from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.docstore.document import Document
from langchain.prompts import PromptTemplate
from langchain.indexes.vectorstore import VectorstoreIndexCreator
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from langchain.vectorstores import Pinecone
import langchain
import json
import pandas as pd
import pinecone
import openai
import os

# Clear the terminal
os.system('cls' if os.name == 'nt' else 'clear')

## Set local environment variables
embeddings = OpenAIEmbeddings()
OPENAI_API_KEY=os.getenv("OPEN_API_KEY")
pinecone.init(api_key=os.getenv("PINECONE_API_KEY"),
              environment=os.getenv("PINECONE_ENVIRONMENT_KEY"))

# Create a Pinecone index object
index_name = "llm-demo"
index = pinecone.Index(index_name=index_name)

## Langchain setup
model = langchain.OpenAI(temperature=0, model_name="gpt-4")

## Create documents to send to QA Chain
def get_documents(response):
    # Create lists
    ids = []
    scores = [] 
    contents = []
    docs= []
    
    # Create docs list for langchain Qa Chain
    for match in response['matches']:
        ids.append(match['metadata']['embedding_id'])
        scores.append(match['score'])
        contents.append(match['metadata']['embedding_content'])
        content=match['metadata']['embedding_content']
  
        # Create Document object
        doc = Document(
            page_content=content
        )
        docs.append(doc)
    print(docs)

    get_response_from_llm(docs)

    # Create a dataframe (THIS IS NOT USED)
    search_results_df = pd.DataFrame({
        'id': ids,
        'score': scores,
        'page_content': contents
    })

## Get response from langchain Qa Chain   
def get_response_from_llm(docs):
    # Load QA Chain
    qa_chain = load_qa_chain(model, chain_type="stuff")
    response = qa_chain.run(
        question=question, 
        input_documents=docs
    )  
    print(response)


## Generate the query embedding
def answer_question(question):
 
    question_emb = embeddings.embed_query(question)

    # Perform the query
    response = index.query([question_emb], top_k=20, include_metadata=True, include_values=False)

    get_documents(response)

###########################-MAIN-############################################## 
#question =  "What did the president say about Justice Breyer" 
#question =  "What did the president say about Ukraine"
#question = "What did the president say about immigration. Provide 5 as bullets.  be concise"
question =  "What did the president say about southern border. Provide 3 as bullets.  be concise"
#question = "What is the president' birthday"

answer = answer_question(question)
langchain.chains.question_answering.load_qa_chain langchain.docstore.document.Document langchain.OpenAI langchain.embeddings.openai.OpenAIEmbeddings
[((568, 616), 'os.system', 'os.system', (["('cls' if os.name == 'nt' else 'clear')"], {}), "('cls' if os.name == 'nt' else 'clear')\n", (577, 616), False, 'import os\n'), ((666, 684), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (682, 684), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((700, 725), 'os.getenv', 'os.getenv', (['"""OPEN_API_KEY"""'], {}), "('OPEN_API_KEY')\n", (709, 725), False, 'import os\n'), ((910, 947), 'pinecone.Index', 'pinecone.Index', ([], {'index_name': 'index_name'}), '(index_name=index_name)\n', (924, 947), False, 'import pinecone\n'), ((976, 1027), 'langchain.OpenAI', 'langchain.OpenAI', ([], {'temperature': '(0)', 'model_name': '"""gpt-4"""'}), "(temperature=0, model_name='gpt-4')\n", (992, 1027), False, 'import langchain\n'), ((1724, 1792), 'pandas.DataFrame', 'pd.DataFrame', (["{'id': ids, 'score': scores, 'page_content': contents}"], {}), "({'id': ids, 'score': scores, 'page_content': contents})\n", (1736, 1792), True, 'import pandas as pd\n'), ((1935, 1975), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['model'], {'chain_type': '"""stuff"""'}), "(model, chain_type='stuff')\n", (1948, 1975), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((748, 777), 'os.getenv', 'os.getenv', (['"""PINECONE_API_KEY"""'], {}), "('PINECONE_API_KEY')\n", (757, 777), False, 'import os\n'), ((805, 842), 'os.getenv', 'os.getenv', (['"""PINECONE_ENVIRONMENT_KEY"""'], {}), "('PINECONE_ENVIRONMENT_KEY')\n", (814, 842), False, 'import os\n'), ((1528, 1558), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': 'content'}), '(page_content=content)\n', (1536, 1558), False, 'from langchain.docstore.document import Document\n')]