Update main.py

main.py

@@ -1,5 +1,23 @@
+import os
+
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from langchain.vectorstores import Pinecone
+from langchain.llms import OpenAI
+from langchain.chains import RetrievalQA
+
+PINECONE_API_KEY = os.getenv('PINECONE_API_KEY')
+PINECONE_ENV = os.getenv('PINECONE_ENV')
+OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
+PINECONE_INDEX_NAME = os.getenv('PINECONE_INDEX_NAME')
+
+def parse_response(response):
+    result = response['result']
+    result += '\n\nSources:'
+    for source_name in response["source_documents"]:
+        result += ''.join((source_name.metadata['source'],
+              "page #:", str(source_name.metadata['page'])))
+    return result
 
 app = FastAPI()
 
@@ -14,4 +32,29 @@ def read_root():
     
 @app.get("/api/python")
 def hello_python():
-    return {"message": "Hello Python"}
+    return {"message": "Hello Python"}
+
+@app.get("/prompt")
+def read_root(p: str='According to HQ H303140, what is "Country  of  origin"  means?'):
+    pinecone.init(
+        api_key=PINECONE_API_KEY,
+        environment=PINECONE_ENV
+    )
+    index = pinecone.Index(PINECONE_INDEX_NAME)
+    index.describe_index_stats()
+    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
+    docsearch = Pinecone.from_existing_index(PINECONE_INDEX_NAME, embeddings)
+    retriever = docsearch.as_retriever(
+        include_metadata=True, 
+        metadata_key='source'
+    )
+    llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)
+    qa_chain = RetrievalQA.from_chain_type(llm=llm,
+                                       chain_type="stuff",
+                                       retriever=retriever,
+                                       return_source_documents=True)
+    response = qa_chain(p)
+    return {
+        "prompt": p,
+        "response": parse_response(response)
+    }