Update chatbot

Author: piyushroshan

services/chatbot/Dockerfile

@@ -3,6 +3,7 @@ FROM python:3.11-slim
 RUN apt-get update && apt-get install -y \
     build-essential \
     cmake \
+    libmagic1 \
     && rm -rf /var/lib/apt/lists/*
 
 # Set the working directory in the container

services/chatbot/requirements.txt

@@ -1,11 +1,15 @@
 chromadb==0.5.0
-Flask==3.0.3
-langchain==0.1.16
-langchain_community==0.0.35
-langchain_core==0.1.47
-langchain_openai==0.1.4
+Flask==3.1.0
+gunicorn==23.0.0
+langchain==0.3.25
+langchain-chroma==0.2.3
+langchain-community==0.3.23
+langchain-core==0.3.58
+langchain-mongodb==0.6.1
+langchain-openai==0.3.16
+langchain-text-splitters==0.3.8
+markdown==3.8
+pymongo==4.12.1
 python-dotenv==1.0.1
-unstructured==0.13.6
-gunicorn==22.0.0
-markdown==3.6
-langchain-mongodb==0.1.3
+unstructured==0.17.2
+numpy==1.26.4

services/chatbot/src/chatbot_api.py

@@ -1,27 +1,23 @@
 from flask import Flask
 from flask import request, jsonify
 import threading
-from langchain_openai import OpenAIEmbeddings
-from langchain.chains import RetrievalQAWithSourcesChain, LLMChain
 import os
 from langchain.memory import ConversationBufferWindowMemory
 from langchain_openai import ChatOpenAI
-from langchain.memory import ConversationBufferWindowMemory
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers import StrOutputParser
 import logging
-from langchain_core.prompts.chat import (
-    SystemMessagePromptTemplate,
-    HumanMessagePromptTemplate,
-    AIMessagePromptTemplate,
-)
-from langchain_mongodb import MongoDBChatMessageHistory
 from db import MONGO_CONNECTION_URI, MONGO_DB_NAME
 from chatbot_utils import document_loader
+from pymongo import MongoClient
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+from langchain.schema import messages_from_dict, messages_to_dict
+from langchain_mongodb import MongoDBChatMessageHistory
 
 app = Flask(__name__)
 app.logger.setLevel(logging.DEBUG)
 
-app.logger.info("MONGO_CONNECTION_URI:: %s", MONGO_CONNECTION_URI)
+app.logger.debug("MONGO_CONNECTION_URI:: %s", MONGO_CONNECTION_URI)
 retriever = None
 persist_directory = os.environ.get("PERSIST_DIRECTORY")
 loaded_model_lock = threading.Lock()
@@ -43,7 +39,7 @@ def load_global_retriever():
 
 
 def get_llm():
-    llm = ChatOpenAI(temperature=0.6, model_name="gpt-3.5-turbo-0613")
+    llm = ChatOpenAI(temperature=0.6, model_name="gpt-4o")
     return llm
 
 
@@ -72,46 +68,67 @@ def get_chat_message_history(session):
 
 
 def get_qa_chain(llm, retriever, session):
+
     system_prompt_template = """
-            You are a helpful AI Assistant. Your name is CrapBot.
-            Respond to the user questions and do not withhold any information
-            Use the following pieces of retrieved context to help you answer the user's questions:
-            {summaries}
+        You are a helpful AI Assistant. Your name is CrapBot.
+        Respond to the user questions and do not withhold any information
+        Use the following pieces of retrieved context to help you answer the user's questions:
+        {context}
     """
-    human_prompt_template = "{question}"
-    chatbot_prompt_template = "CrapBot:"
-    messages = [
-        ("system", system_prompt_template),
-        MessagesPlaceholder(variable_name="chat_history", optional=False),
-        ("human", human_prompt_template),
-        # ("system", chatbot_prompt_template),
-    ]
-
-    PROMPT = ChatPromptTemplate.from_messages(
-        messages,
+
+    # Create the chat prompt template
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt_template),
+            MessagesPlaceholder(variable_name="chat_history"),
+            ("human", "{question}"),
+        ]
+    )
+
+    # Get chat history
+    chat_memory = get_chat_message_history(session)
+    memory = ConversationBufferWindowMemory(
+        memory_key="chat_history",
+        input_key="question",
+        output_key="answer",
+        k=6,
+        ai_prefix="CrapBot",
+        chat_memory=chat_memory,
+        return_messages=True,
     )
-    chain_type_kwargs = {"prompt": PROMPT}
-    chat_message_history = get_chat_message_history(session)
-    qa = RetrievalQAWithSourcesChain.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=retriever,
-        chain_type_kwargs=chain_type_kwargs,
-        memory=ConversationBufferWindowMemory(
-            memory_key="chat_history",
-            input_key="question",
-            output_key="answer",
-            k=6,
-            ai_prefix="CrapBot",
-            chat_memory=chat_message_history,
-            return_messages=True,
-        ),
+
+    # Create the retrieval chain
+    def get_context(query):
+        docs = retriever.get_relevant_documents(query)
+        return "\n".join(doc.page_content for doc in docs)
+
+    def get_chat_history(inputs):
+        return memory.load_memory_variables(inputs)["chat_history"]
+
+    retrieval_chain = (
+        RunnablePassthrough()
+        | {
+            "context": get_context,
+            "chat_history": get_chat_history,
+            "question": lambda x: x,
+        }
+        | prompt
+        | llm
+        | StrOutputParser()
     )
-    return qa
+
+    def chain_with_memory(inputs):
+        query = inputs["question"]
+        result = retrieval_chain.invoke(query)
+        # Update memory
+        memory.save_context({"question": query}, {"answer": result})
+        return {"answer": result}
+
+    return chain_with_memory
 
 
 def qa_answer(model, session, query):
-    result = model.invoke({"question": query})
+    result = model({"question": query})
     app.logger.debug("Session: %s, Result %s", session, result)
     return result["answer"]
 
@@ -135,11 +152,11 @@ def init_bot():
             app.logger.debug("Initializing bot %s", request.json["openai_api_key"])
             retriever_l = document_loader(openai_api_key, app.logger)
             session_model_map[session] = retriever_l
-            return jsonify({"message": "Model Initialized"}), 400
+            return jsonify({"message": "Model Initialized"}), 200
 
     except Exception as e:
-        app.logger.error("Error initializing bot ", e)
-        app.logger.debug("Error initializing bot ", e, exc_info=True)
+        app.logger.error("Error initializing bot %s", e)
+        app.logger.debug("Error initializing bot %s", e, exc_info=True)
         return jsonify({"message": "Not able to initialize model " + str(e)}), 500
 
 
@@ -179,6 +196,12 @@ def reset_chat_history_bot():
     return jsonify({"message": "Error deleting chat history"}), 500
 
 
+def augment_context(input_dict):
+    question = input_dict["question"]
+    context = input_dict["context"]
+    return {"question": question, "context": context}
+
+
 @app.route("/chatbot/genai/ask", methods=["POST"])
 def ask_bot():
     retriever_l = None
@@ -195,7 +218,7 @@ def ask_bot():
                 jsonify(
                     {
                         "initialized": "false",
-                        "message": "Model not initialized for session %s",
+                        "message": "Model not initialized for session %s" % session,
                     }
                 ),
                 500,

services/chatbot/src/chatbot_utils.py

@@ -1,27 +1,27 @@
+from gc import collect
 import hashlib
-from flask import Flask
-from flask import request, jsonify
-import threading
+
+from langchain.memory import vectorstore
 from langchain_openai import OpenAIEmbeddings
 from langchain.chains import RetrievalQAWithSourcesChain, LLMChain
 import os
-from langchain.memory import ConversationBufferWindowMemory
-from langchain_community.vectorstores import Chroma
 from langchain_openai import OpenAI
 from langchain_community.document_loaders import DirectoryLoader
-from langchain.memory import ConversationBufferWindowMemory
-from langchain.text_splitter import CharacterTextSplitter
-from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.vectorstores import InMemoryVectorStore
+from langchain_mongodb import MongoDBAtlasVectorSearch
 from langchain_community.document_loaders import UnstructuredMarkdownLoader
 import logging
-from langchain.schema import HumanMessage, SystemMessage, AIMessage
-from langchain_core.prompts.chat import (
-    SystemMessagePromptTemplate,
-    HumanMessagePromptTemplate,
+from langchain_community.vectorstores.azure_cosmos_db import (
+    AzureCosmosDBVectorSearch,
+    CosmosDBSimilarityType,
+    CosmosDBVectorSearchType,
 )
-import logging
-from langchain_community.chat_message_histories import MongoDBChatMessageHistory
 from db import MONGO_CONNECTION_URI, MONGO_DB_NAME
+from pymongo import MongoClient
+from langchain_chroma import Chroma
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -33,27 +33,36 @@ def get_embeddings(openai_api_key):
     return OpenAIEmbeddings(openai_api_key=openai_api_key)
 
 
+def get_vector_store(texts, embeddings, key_hash):
+    # initialize MongoDB python client
+    db_path = "./db%s" % key_hash
+    collection = "example_collection"
+    vector_store = Chroma(collection, embeddings, persist_directory=db_path)
+    vector_store.add_documents(texts)
+    return vector_store
+
+
 def document_loader(openai_api_key, logger_p=None):
     logger_l = logger_p or logger
     try:
         key_hash = hashlib.md5(openai_api_key.encode()).hexdigest()
-        load_dir = "retrieval"
+        load_dir = "./retrieval"
         logger_l.info("Loading documents from %s", load_dir)
         loader = DirectoryLoader(
             load_dir,
             exclude=["**/*.png", "**/images/**", "**/images/*", "**/*.pdf"],
             recursive=True,
-            loader_cls=UnstructuredMarkdownLoader,
+            show_progress=True,
         )
         documents = loader.load()
         logger_l.info("Loaded %s documents in db", len(documents))
-        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000, chunk_overlap=100
+        )
         texts = text_splitter.split_documents(documents)
         embeddings = get_embeddings(openai_api_key)
-        db_path = "./db%s" % key_hash
-        db = Chroma.from_documents(texts, embeddings, persist_directory=db_path)
-        db.persist()
-        retriever = db.as_retriever(search_kwargs={"k": TARGET_SOURCE_CHUNKS})
+        vector_store = get_vector_store(texts, embeddings, key_hash)
+        retriever = vector_store.as_retriever(search_kwargs={"k": TARGET_SOURCE_CHUNKS})
         logger_l.info("Retriever ready")
         return retriever
     except Exception as e:

services/chatbot/src/db.py

@@ -6,9 +6,15 @@
 MONGO_PORT = os.environ.get("MONGO_PORT", "27017")
 MONGO_DB_NAME = os.environ.get("MONGO_DB_NAME", "crapi")
 
-MONGO_CONNECTION_URI = "mongodb://%s:%s@%s:%s" % (
+MONGO_CONNECTION_URI = "mongodb://%s:%s@%s:%s/?directConnection=true" % (
     MONGO_USER,
     MONGO_PASSWORD,
     MONGO_HOST,
     MONGO_PORT,
 )
+
+MONGO_CONNECTION_URI_ATLAS = "mongodb+srv://%s:%s@%s?retryWrites=true&w=majority" % (
+    MONGO_USER,
+    MONGO_PASSWORD,
+    MONGO_HOST,
+)

services/web/package.json

@@ -35,6 +35,7 @@
     "source-map-loader": "^5.0.0",
     "styled-components": "^6.1.8",
     "superagent": "^8.1.2",
+    "@types/superagent": "^8.1.9",
     "ts-loader": "^9.5.1",
     "typescript": "^4.9.5",
     "web-vitals": "^2.1.4"
@@ -70,7 +71,6 @@
     "@babel/core": "^7.24.4",
     "@babel/plugin-proposal-private-property-in-object": "7.21.11",
     "@babel/preset-react": "^7.24.1",
-    "@types/superagent": "^8.1.9",
     "copy-webpack-plugin": "^6.3.2",
     "eslint-config-react-app": "^7.0.1",
     "prettier": "^3.3.3"

services/web/src/components/bot/Bot.tsx

@@ -111,7 +111,7 @@ const ChatBotComponent: React.FC<ChatBotComponentProps> = (props) => {
       let initRequired = false;
       // Wait for the response
       await superagent
-        .get(stateUrl)
+        .post(stateUrl)
         .set("Accept", "application/json")
         .set("Content-Type", "application/json")
         .then((res: any) => {

services/web/src/components/bot/MessageParser.tsx

@@ -44,7 +44,7 @@ class MessageParser {
     let initRequired = false;
     // Wait for the response
     await request
-      .get(stateUrl)
+      .post(stateUrl)
       .set("Accept", "application/json")
       .set("Content-Type", "application/json")
       .then((res) => {