"""This example shows how to use the ChatGPT API
with LangChain to answer questions about Prefect."""

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import ChatVectorDBChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

from langchain_prefect.loaders import GitHubRepoLoader
from langchain_prefect.plugins import RecordLLMCalls

documents = GitHubRepoLoader("PrefectHQ/prefect", glob="**/*.md").load()

text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
documents = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

system_template = """Use the following pieces of context to answer the users question. 
If you don't know the answer, just say that you don't know, don't make up an answer.
----------------
{context}"""

prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{question}"),
    ]
)

qa = ChatVectorDBChain.from_llm(
    llm=ChatOpenAI(temperature=0),
    vectorstore=Chroma.from_documents(documents, embeddings),
    qa_prompt=prompt,
)

with RecordLLMCalls(
    tags={qa.vectorstore.__class__.__name__}, max_prompt_tokens=int(1e4)
):
    chat_history = []
    query = "What infrastructures does Prefect support?"
    result = qa({"question": query, "chat_history": chat_history})

    print(result["answer"])

    chat_history = [(query, result["answer"])]
    query = "Can I use Prefect with AWS?"
    result = qa({"question": query, "chat_history": chat_history})

    print(result["answer"])

langchain.text_splitter.CharacterTextSplitter langchain_prefect.loaders.GitHubRepoLoader langchain.prompts.chat.SystemMessagePromptTemplate.from_template langchain.chat_models.ChatOpenAI langchain.vectorstores.Chroma.from_documents langchain.prompts.chat.HumanMessagePromptTemplate.from_template langchain.embeddings.openai.OpenAIEmbeddings
[((680, 735), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(0)'}), '(chunk_size=1000, chunk_overlap=0)\n', (701, 735), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((803, 821), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (819, 821), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((602, 655), 'langchain_prefect.loaders.GitHubRepoLoader', 'GitHubRepoLoader', (['"""PrefectHQ/prefect"""'], {'glob': '"""**/*.md"""'}), "('PrefectHQ/prefect', glob='**/*.md')\n", (618, 655), False, 'from langchain_prefect.loaders import GitHubRepoLoader\n'), ((1084, 1142), 'langchain.prompts.chat.SystemMessagePromptTemplate.from_template', 'SystemMessagePromptTemplate.from_template', (['system_template'], {}), '(system_template)\n', (1125, 1142), False, 'from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((1152, 1206), 'langchain.prompts.chat.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['"""{question}"""'], {}), "('{question}')\n", (1192, 1206), False, 'from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((1258, 1283), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (1268, 1283), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1301, 1345), 'langchain.vectorstores.Chroma.from_documents', 'Chroma.from_documents', (['documents', 'embeddings'], {}), '(documents, embeddings)\n', (1322, 1345), False, 'from langchain.vectorstores import Chroma\n')]