import os
import pinecone

from rich.console import Console
from rich.markdown import Markdown

import langchain
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.embeddings import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.vectorstores import Pinecone

# langchain.debug = True

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_KEY")
PINECONE_ENVIRONMENT = os.getenv("PINECONE_ENVIRONMENT")
PINECONE_INDEX = os.getenv("PINECONE_INDEX")

embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)
index = pinecone.Index(PINECONE_INDEX)
vector_store = Pinecone(index, embeddings, "text")

prompt_template = """
You are a question-answering bot for Airbyte company employees and will be provided
relevant context from Notion pages on the Airbyte company knowlege base.

Whenever you are asked a question you answer with a helpful answer if you can, along
with the links to those relevant pages for further information. If you are not sure, you
will say that you are not sure but still provide links if anything might be helpful to the
questioner.

Only use the provided context. Do not guess and do not use prior knowlege.

Please provide your response in markdown format, starting with a level 2 header that describes
the answer under a reasonable summary header.

Notion context for this question:
{context}

Question: {question}

Please provide a helpful answer along one or more URLs that would be helpful for finding additional information:
"""

prompt = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

qa = RetrievalQA.from_chain_type(
    llm=OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY),
    chain_type="stuff",
    retriever=vector_store.as_retriever(),
    chain_type_kwargs={"prompt": prompt},
)

console = Console()

console.print(Markdown("\n------\n> What do you want to know?"))
console.print("")
while True:
    try:
        query = input("")
    except KeyboardInterrupt:
        console.print("\n")
        console.print(Markdown("_Goodbye!_ 👋"))
        exit(0)

    answer = qa.run(query)
    console.print(Markdown(answer))

    console.print(Markdown("\n------\n> What else do you want to know?\n"))
    console.print("\n")

langchain.llms.OpenAI langchain.embeddings.OpenAIEmbeddings langchain.prompts.PromptTemplate langchain.vectorstores.Pinecone
[((371, 398), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (380, 398), False, 'import os\n'), ((418, 443), 'os.getenv', 'os.getenv', (['"""PINECONE_KEY"""'], {}), "('PINECONE_KEY')\n", (427, 443), False, 'import os\n'), ((467, 500), 'os.getenv', 'os.getenv', (['"""PINECONE_ENVIRONMENT"""'], {}), "('PINECONE_ENVIRONMENT')\n", (476, 500), False, 'import os\n'), ((518, 545), 'os.getenv', 'os.getenv', (['"""PINECONE_INDEX"""'], {}), "('PINECONE_INDEX')\n", (527, 545), False, 'import os\n'), ((560, 607), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'openai_api_key': 'OPENAI_API_KEY'}), '(openai_api_key=OPENAI_API_KEY)\n', (576, 607), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((608, 681), 'pinecone.init', 'pinecone.init', ([], {'api_key': 'PINECONE_API_KEY', 'environment': 'PINECONE_ENVIRONMENT'}), '(api_key=PINECONE_API_KEY, environment=PINECONE_ENVIRONMENT)\n', (621, 681), False, 'import pinecone\n'), ((690, 720), 'pinecone.Index', 'pinecone.Index', (['PINECONE_INDEX'], {}), '(PINECONE_INDEX)\n', (704, 720), False, 'import pinecone\n'), ((736, 771), 'langchain.vectorstores.Pinecone', 'Pinecone', (['index', 'embeddings', '"""text"""'], {}), "(index, embeddings, 'text')\n", (744, 771), False, 'from langchain.vectorstores import Pinecone\n'), ((1643, 1728), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['context', 'question']"}), "(template=prompt_template, input_variables=['context',\n    'question'])\n", (1657, 1728), False, 'from langchain.prompts import PromptTemplate\n'), ((1950, 1959), 'rich.console.Console', 'Console', ([], {}), '()\n', (1957, 1959), False, 'from rich.console import Console\n'), ((1975, 2026), 'rich.markdown.Markdown', 'Markdown', (['"""\n------\n> What do you want to know?"""'], {}), '("""\n------\n> What do you want to know?""")\n', (1983, 2026), False, 'from rich.markdown import Markdown\n'), ((1774, 1826), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'openai_api_key': 'OPENAI_API_KEY'}), '(temperature=0, openai_api_key=OPENAI_API_KEY)\n', (1780, 1826), False, 'from langchain.llms import OpenAI\n'), ((2259, 2275), 'rich.markdown.Markdown', 'Markdown', (['answer'], {}), '(answer)\n', (2267, 2275), False, 'from rich.markdown import Markdown\n'), ((2296, 2353), 'rich.markdown.Markdown', 'Markdown', (['"""\n------\n> What else do you want to know?\n"""'], {}), '("""\n------\n> What else do you want to know?\n""")\n', (2304, 2353), False, 'from rich.markdown import Markdown\n'), ((2171, 2195), 'rich.markdown.Markdown', 'Markdown', (['"""_Goodbye!_ 👋"""'], {}), "('_Goodbye!_ 👋')\n", (2179, 2195), False, 'from rich.markdown import Markdown\n')]