import os
import openai
from dotenv import load_dotenv
load_dotenv()

print(os.getenv("OPENAI_API_KEY"))
openai.api_key = os.getenv("OPENAI_API_KEY")

from langchain import PromptTemplate
from langchain.chains import RetrievalQA

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI
from langchain.text_splitter import CharacterTextSplitter
from langchain.text_splitter import MarkdownTextSplitter
from langchain.vectorstores import Chroma
from langchain.document_loaders import TextLoader

from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI
from langchain.utilities import SerpAPIWrapper
from langchain.agents import initialize_agent
from langchain.agents import AgentType
# from langchain.vectorstores import FAISS

from langchain.chains.router import MultiRetrievalQAChain
from langchain.llms import OpenAI

bulevar_restaurant_texts = [
    "Bulevar is open Sunday through Wednesday from 5-9pm, and Thursday through Satruday from 4-10pm",
    "Bulevar is located in the Arboretum at 360 and Mopac, next to Eddie V's",
    "Bulevar offers tasty Mexican Cuisine, with a laid back style to fine-dining.",
]
bulevar_details_retriever = Chroma.from_texts(bulevar_restaurant_texts, embedding_function=  OpenAIEmbeddings())

loader = TextLoader('raw_text/food_guide.md')
documents = loader.load()

text_splitter = MarkdownTextSplitter(chunk_size=750, chunk_overlap=35)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

docs_retriever = Chroma.from_documents(docs, embeddings)

query = "Who is Jonny?"
docs = docs_retriever.similarity_search(query)

print(docs[0].page_content)

prompt = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

Context: {context}

Question: {question}
Helpful Answer:"""

prompt_template = PromptTemplate(
            template=prompt, input_variables=["context", "question"]
        )
qa_chain = RetrievalQA.from_chain_type(llm= ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.0),  
                                      chain_type='stuff', 
                                      retriever=docs_retriever.as_retriever(), # the doc search
                                    #   k=4, # takes the top 4 documents from the search
                                        return_source_documents=True,
                                        chain_type_kwargs={"prompt": prompt_template, "verbose": True},
                                      )
result = qa_chain({"query": query})
print(result["result"])

retriever_infos = [
    {
        "name": "Food Guide", 
        "description": "Good for answering questions about Bulevar's dinner menu", 
        "retriever": docs_retriever.as_retriever()
    },
    {
        "name": "Bulevar Restaurant Details", 
        "description": "Good for answering questions about Bulevar's hours, and restaurant details such as its mission, history, and owners.", 
        "retriever": bulevar_details_retriever
    }
]

chain = MultiRetrievalQAChain.from_retrievers(OpenAI(), retriever_infos, verbose=True)

print(chain.run("What is Bulevar?"))
print(chain.run("What kinds of meats are offered?"))


# print(result["source_documents"])

# question answer

# How is Dialogues different from others?
# Dialogues is different from others because it is a well-thought-out game with several players, built-in trust, a simple interface, and a feature that allows the user to select and align with the values that matter to them. It is designed to offer value to the individual and is not dictated by Silicon Valley. The founder has put everything they like about reading books, watching movies, and having conversations into building the app. The app also uses a measurement of normality based on duration of time to respond to abnormal and normal signals in the data.

# How is the product coming along?
# The product is nearly half a million lines of code and has an incredibly well-developed and defined product idea, but it is not yet ready for business as the creator is still seeking funding partners and wants to perfect more of it before opening up.

# Who is Jonny?
# Jonny is a person who has created a product related to AI, specifically LLMs and chatbot games, and is seeking investment from potential investors. He believes that his product can increase the valuation of existing companies and is confident in his engineering abilities. He also believes that the future of LLMs will be with writers and is focused on crafting writing to service each person at a personal level.