from llama_index.callbacks import CallbackManager, LlamaDebugHandler, CBEventType
from llama_index import ListIndex, ServiceContext, SimpleDirectoryReader, VectorStoreIndex

'''
Title of the page: A simple Python implementation of the ReAct pattern for LLMs
Name of the website: LlamaIndex (GPT Index) is a data framework for your LLM application.
URL: https://github.com/jerryjliu/llama_index
'''
docs = SimpleDirectoryReader("../data/paul_graham/").load_data()

from llama_index import ServiceContext, LLMPredictor, TreeIndex
from langchain.chat_models import ChatOpenAI
llm_predictor = LLMPredictor(llm=ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0))

llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])
service_context = ServiceContext.from_defaults(callback_manager=callback_manager, llm_predictor=llm_predictor)

index = VectorStoreIndex.from_documents(docs, service_context=service_context)
query_engine = index.as_query_engine()

response = query_engine.query("What did the author do growing up?")

# Print info on the LLM calls during the list index query
print(llama_debug.get_event_time_info(CBEventType.LLM))

# Print info on llm inputs/outputs - returns start/end events for each LLM call
event_pairs = llama_debug.get_llm_inputs_outputs()
print(event_pairs[0][0])
print(event_pairs[0][1].payload.keys())
print(event_pairs[0][1].payload['response'])

# Get info on any event type
event_pairs = llama_debug.get_event_pairs(CBEventType.CHUNKING)
print(event_pairs[0][0].payload.keys())  # get first chunking start event
print(event_pairs[0][1].payload.keys())  # get first chunking end event

# Clear the currently cached events
llama_debug.flush_event_logs()

langchain.chat_models.ChatOpenAI
[((676, 718), 'llama_index.callbacks.LlamaDebugHandler', 'LlamaDebugHandler', ([], {'print_trace_on_end': '(True)'}), '(print_trace_on_end=True)\n', (693, 718), False, 'from llama_index.callbacks import CallbackManager, LlamaDebugHandler, CBEventType\n'), ((738, 768), 'llama_index.callbacks.CallbackManager', 'CallbackManager', (['[llama_debug]'], {}), '([llama_debug])\n', (753, 768), False, 'from llama_index.callbacks import CallbackManager, LlamaDebugHandler, CBEventType\n'), ((787, 883), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'callback_manager': 'callback_manager', 'llm_predictor': 'llm_predictor'}), '(callback_manager=callback_manager,\n    llm_predictor=llm_predictor)\n', (815, 883), False, 'from llama_index import ServiceContext, LLMPredictor, TreeIndex\n'), ((889, 959), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['docs'], {'service_context': 'service_context'}), '(docs, service_context=service_context)\n', (920, 959), False, 'from llama_index import ListIndex, ServiceContext, SimpleDirectoryReader, VectorStoreIndex\n'), ((405, 450), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (['"""../data/paul_graham/"""'], {}), "('../data/paul_graham/')\n", (426, 450), False, 'from llama_index import ListIndex, ServiceContext, SimpleDirectoryReader, VectorStoreIndex\n'), ((606, 659), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)'}), "(model_name='gpt-3.5-turbo', temperature=0)\n", (616, 659), False, 'from langchain.chat_models import ChatOpenAI\n')]