dotku/fastapi-columbus

1import os2import pinecone3 4from fastapi import FastAPI5from fastapi.middleware.cors import CORSMiddleware6from langchain.chains import RetrievalQA7from langchain.embeddings.openai import OpenAIEmbeddings8from langchain.llms import OpenAI9from langchain.vectorstores import Pinecone10 11PINECONE_API_KEY = os.getenv('PINECONE_API_KEY')12PINECONE_ENV = os.getenv('PINECONE_ENV')13OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')14PINECONE_INDEX_NAME = os.getenv('PINECONE_INDEX_NAME')15 16def parse_response(response):17    result = response['result']18    result += '\n\nSources:  \n'19    for source_name in response["source_documents"]:20        result += ''.join((source_name.metadata['source'],21              " page #:", str(source_name.metadata['page']), '  \n'))22    return result23 24app = FastAPI()25 26app.add_middleware(27    CORSMiddleware,28    allow_origins=['*']29)30 31@app.get("/")32def read_root():33    return {"message": "Hello World"}34    35@app.get("/api/python")36def hello_python():37    return {"message": "Hello Python"}38 39@app.get("/prompt")40def read_root(p: str='According to HQ H303140, what is "Country  of  origin"  means?'):41    pinecone.init(42        api_key=PINECONE_API_KEY,43        environment=PINECONE_ENV44    )45    index = pinecone.Index(PINECONE_INDEX_NAME)46    index.describe_index_stats()47    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)48    docsearch = Pinecone.from_existing_index(PINECONE_INDEX_NAME, embeddings)49    retriever = docsearch.as_retriever(50        include_metadata=True, 51        metadata_key='source'52    )53    llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)54    qa_chain = RetrievalQA.from_chain_type(llm=llm,55                                       chain_type="stuff",56                                       retriever=retriever,57                                       return_source_documents=True)58    response = qa_chain(p)59    return {60        "prompt": p,61        "response": parse_response(response)62    }