import tiktoken
from llama_index.core import MockEmbedding, VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from llama_index.core.llms.mock import MockLLM

embed_model = MockEmbedding(embed_dim=1536)
llm = MockLLM(max_tokens=256)
token_counter = TokenCountingHandler(
    tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode
)
callback_manager = CallbackManager([token_counter])

Settings.embed_model=embed_model
Settings.llm=llm
Settings.callback_manager=callback_manager

documents = SimpleDirectoryReader("przewidywanie_kosztów_próbki").load_data()
index = VectorStoreIndex.from_documents(
    documents=documents,
    show_progress=True)
print("Liczba tokenów osadzenia:", token_counter.total_embedding_token_count)

query_engine = index.as_query_engine()  
response = query_engine.query("Jak ma na imię kot?")   
print("Liczba tokenów zapytania do LLM:", token_counter.total_llm_token_count)
print("Liczba tokenów osadzenia zapytania:",token_counter.total_embedding_token_count)

