from llama_index import (
    VectorStoreIndex,
    SummaryIndex,
    SimpleKeywordTableIndex,
    SimpleDirectoryReader,
    ServiceContext,
    StorageContext,
    load_index_from_storage
)
from llama_index.schema import IndexNode
from llama_index.tools import QueryEngineTool, ToolMetadata
from llama_index.llms import OpenAI

import os
os.environ["OPENAI_API_KEY"]
llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
service_context = ServiceContext.from_defaults(llm=llm)

PERSIST_DIR = "clm_guidance_metadata"
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
query_engine = index.as_query_engine(similarity_top_k=3, llm=OpenAI(model="gpt-3.5-turbo"))

import gradio as gr

def clm(question: str, conversation_history: list[str]):
    
    context = " ".join([item["user"] + " " + item["chatbot"] for item in conversation_history])
    response = query_engine.query("the user is asking questions about community led monitoring, or CLM. The user previously asked and received the following: " +
                                  context +
                                  question)

    conversation_history.append({"user": question, "chatbot": response.response})   

    source1 = ("File Name: " +
      response.source_nodes[0].metadata["file_name"] +
      "\nPage Number: " +
      response.source_nodes[0].metadata["page_label"] +
      "\n Source Text: " +
     response.source_nodes[0].text)
    
    source2 = ("File Name: " +
      response.source_nodes[1].metadata["file_name"] +
      "\nPage Number: " +
      response.source_nodes[1].metadata["page_label"] +
      "\n Source Text: " +
     response.source_nodes[1].text)
        
    source3 = ("File Name: " +
      response.source_nodes[2].metadata["file_name"] +
      "\nPage Number: " +
      response.source_nodes[2].metadata["page_label"] +
      "\n Source Text: " +
     response.source_nodes[2].text)

    return response, source1, source2, source3, conversation_history

inputs = [gr.Textbox(lines=10, label="Question"),
          gr.State(value=[])]
outputs = [
    gr.Textbox(label="Chatbot Response", type="text"),
    gr.Textbox(label="Source 1", max_lines = 10, autoscroll = False, type="text"),
    gr.Textbox(label="Source 2", max_lines = 10, autoscroll = False, type="text"),
    gr.Textbox(label="Source 3", max_lines = 10, autoscroll = False, type="text"),
    gr.State()
]

gr.Interface(fn=clm, inputs=inputs, outputs=outputs, title="CLM Chatbot", 
             description="Enter a question and see the processed outputs in collapsible boxes.").launch()