Spaces:
Paused
Paused
| import os | |
| from typing import List | |
| from langchain.embeddings.openai import OpenAIEmbeddings | |
| from langchain.vectorstores.pinecone import Pinecone | |
| from langchain.chains import ConversationalRetrievalChain | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.memory import ChatMessageHistory, ConversationBufferMemory | |
| from langchain.docstore.document import Document | |
| import pinecone | |
| import chainlit as cl | |
| pinecone.init( | |
| api_key=os.environ.get("PINECONE_API_KEY"), | |
| environment=os.environ.get("PINECONE_ENV"), | |
| ) | |
| index_name = "langchain-demo" | |
| embeddings = OpenAIEmbeddings() | |
| welcome_message = "Welcome to the Chainlit Pinecone demo! Ask anything about Shakespeare's King Lear vectorized documents from Pinecone DB." | |
| async def start(): | |
| await cl.Message(content=welcome_message).send() | |
| docsearch = Pinecone.from_existing_index( | |
| index_name=index_name, embedding=embeddings | |
| ) | |
| message_history = ChatMessageHistory() | |
| memory = ConversationBufferMemory( | |
| memory_key="chat_history", | |
| output_key="answer", | |
| chat_memory=message_history, | |
| return_messages=True, | |
| ) | |
| chain = ConversationalRetrievalChain.from_llm( | |
| ChatOpenAI( | |
| model_name="gpt-3.5-turbo", | |
| temperature=0, | |
| streaming=True), | |
| chain_type="stuff", | |
| retriever=docsearch.as_retriever(search_kwargs={'k': 3}), # I only want maximum of three document back with the highest similarity score | |
| memory=memory, | |
| return_source_documents=True, | |
| ) | |
| cl.user_session.set("chain", chain) | |
| async def main(message: cl.Message): | |
| chain = cl.user_session.get("chain") # type: ConversationalRetrievalChain | |
| cb = cl.AsyncLangchainCallbackHandler() | |
| res = await chain.acall(message.content, callbacks=[cb]) | |
| answer = res["answer"] | |
| source_documents = res["source_documents"] # type: List[Document] | |
| text_elements = [] # type: List[cl.Text] | |
| if source_documents: | |
| for source_idx, source_doc in enumerate(source_documents): | |
| source_name = f"source_{source_idx}" | |
| # Create the text element referenced in the message | |
| text_elements.append( | |
| cl.Text(content=source_doc.page_content, name=source_name) | |
| ) | |
| source_names = [text_el.name for text_el in text_elements] | |
| if source_names: | |
| answer += f"\nSources: {', '.join(source_names)}" | |
| else: | |
| answer += "\nNo sources found" | |
| await cl.Message(content=answer, elements=text_elements).send() | |