Update app.py

app.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 from websockets.server import serve
 import os
-from langchain_community.vectorstores import Chroma
+from langchain_chroma import Chroma
 from langchain_huggingface.embeddings import HuggingFaceEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface.llms import HuggingFaceEndpoint
@@ -18,24 +18,22 @@ from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_community.chat_message_histories import ChatMessageHistory
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
 
 if not os.path.isdir('database'):
     os.system("unzip database.zip")
 
-clean_up_tokenization_spaces = True
-
 loader = DirectoryLoader('./database', glob="./*.txt", loader_cls=TextLoader)
 
 documents = loader.load()
 
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-texts = text_splitter.split_documents(documents)
+splits = text_splitter.split_documents(documents)
 
 print()
 print("-------")
 print("TextSplitter, DirectoryLoader")
 print("-------")
-print("--")
 
 persist_directory = 'db'
 
@@ -45,34 +43,14 @@ print()
 print("-------")
 print("Embeddings")
 print("-------")
-print("--")
-
-vectordb = Chroma.from_documents(documents=texts,
-                                 embedding=embedding,
-                                 persist_directory=persist_directory)
-
-vectordb.persist()
-vectordb = None
-
-print()
-print("-------")
-print("Chroma1")
-print("-------")
-print("--")
 
-vectordb = Chroma(persist_directory=persist_directory,
-                  embedding_function=embedding)
-
-print()
-print("-------")
-print("Chroma2")
-print("-------")
-print("--")
+vectorstore = Chroma.from_documents(documents=splits, embedding=embedding)
 
 def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
-retriever = vectordb.as_retriever()
+retriever = vectorstore.as_retriever()
+
 prompt = hub.pull("rlm/rag-prompt")
 llm = HuggingFaceEndpoint(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1")
 rag_chain = (
@@ -86,8 +64,8 @@ print()
 print("-------")
 print("Retriever, Prompt, LLM, Rag_Chain")
 print("-------")
-print("--")
 
+### Contextualize question ###
 contextualize_q_system_prompt = """Given a chat history and the latest user question \
 which might reference context in the chat history, formulate a standalone question \
 which can be understood without the chat history. Do NOT answer the question, \
@@ -103,6 +81,8 @@ history_aware_retriever = create_history_aware_retriever(
     llm, retriever, contextualize_q_prompt
 )
 
+
+### Answer question ###
 qa_system_prompt = """You are an assistant for question-answering tasks. \
 Use the following pieces of retrieved context to answer the question. \
 If you don't know the answer, just say that you don't know. \
@@ -116,16 +96,20 @@ qa_prompt = ChatPromptTemplate.from_messages(
         ("human", "{input}"),
     ]
 )
+question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
+
+rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
+
 
+### Statefully manage chat history ###
 store = {}
 
+
 def get_session_history(session_id: str) -> BaseChatMessageHistory:
     if session_id not in store:
         store[session_id] = ChatMessageHistory()
     return store[session_id]
 
-question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
-rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)
 
 conversational_rag_chain = RunnableWithMessageHistory(
     rag_chain,
@@ -140,11 +124,20 @@ print("started")
 print("-------")
 
 response = conversational_rag_chain.invoke(
-            {"input": "who is the math teacher"},
+            {"input": input()},
+            config={
+                "configurable": {"session_id": "test"}
+            },
+)["answer"]
+print(response)
+
+response = conversational_rag_chain.invoke(
+            {"input": input()},
             config={
                 "configurable": {"session_id": "test"}
             },
 )["answer"]
+print(response)
 
 async def echo(websocket):
     async for message in websocket:
@@ -170,4 +163,4 @@ async def main():
     async with serve(echo, "0.0.0.0", 7860):
         await asyncio.Future()
 
-asyncio.run(main())
+asyncio.run(main())