Implemented vector index for chat history context and MCP tool for semantic search & summarization

Author: keyurdoshi03

services/chatbot/src/chatbot/chat_service.py

@@ -1,9 +1,9 @@
 from uuid import uuid4
-
+import os
 from langgraph.graph.message import Messages
-
+from .vector_index import build_vector_index_from_chat_history, update_vector_index, retrieval_index_path
 from .extensions import db
-from .langgraph_agent import build_langgraph_agent, execute_langgraph_agent
+from .langgraph_agent import execute_langgraph_agent
 
 
 async def get_chat_history(session_id):
@@ -39,4 +39,8 @@ async def process_user_message(session_id, user_message, api_key, model_name, us
     # Limit chat history to last 20 messages
     history = history[-20:]
     await update_chat_history(session_id, history)
+    if not os.path.exists(retrieval_index_path):
+        await build_vector_index_from_chat_history(api_key)
+    else:
+        await update_vector_index(api_key, session_id, {"user": user_message, "assistant": reply.content})
     return reply.content, response_message_id

services/chatbot/src/chatbot/vector_index.py

@@ -0,0 +1,45 @@
+from langchain_community.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document
+from .extensions import db
+
+retrieval_index_path = "/app/resources/chat_index"
+
+async def build_vector_index_from_chat_history(api_key):
+    docs = []
+    async for chat in db.chat_sessions.find({}):
+        session_id = chat.get("session_id", "unknown")
+        messages = chat.get("messages", [])
+        for msg in messages:
+            role = msg.get("role")
+            content = msg.get("content")
+            if content:
+                doc = Document(
+                    page_content=content,
+                    metadata={"session_id": session_id, "role": role}
+                )
+                docs.append(doc)
+
+    embeddings = OpenAIEmbeddings(api_key=api_key)
+    vectorstore = FAISS.from_documents(docs, embeddings)
+    vectorstore.save_local(retrieval_index_path)
+
+async def update_vector_index(api_key, session_id, new_messages):
+    docs = []
+    for role, content in new_messages.items():
+        if content:
+            doc = Document(
+                page_content=content,
+                metadata={"session_id": session_id, "role": role}
+            )
+            docs.append(doc)
+
+    if docs:
+        embeddings = OpenAIEmbeddings(api_key=api_key)
+        vectorstore = FAISS.load_local(
+            retrieval_index_path,
+            embeddings,
+            allow_dangerous_deserialization=True,
+        )
+        vectorstore.add_documents(docs)
+        vectorstore.save_local(retrieval_index_path)

services/chatbot/src/mcpserver/server.py

@@ -4,6 +4,10 @@
 import os
 import logging
 import time
+from .tool_helpers import (
+    get_any_api_key,
+    get_chat_history_retriever,
+)
 # Configure logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
@@ -77,6 +81,27 @@ def get_http_client():
     name="My crAPI MCP Server"
 )
 
+@mcp.tool(tags={"history", "search", "summary", "context"},)
+async def search_chat_history(question: str) -> str:
+    """Answer questions based on user chat history (summarized and semantically indexed). 
+    Use this when the user asks about prior chats, what they asked earlier, or wants a summary of past conversations.    
+    Answer questions based on the user's prior chat history.
+
+    Use this tool when the user refers to anything mentioned before, asks for a summary of previous messages or sessions, 
+    or references phrases like 'what I said earlier', 'things we discussed', 'my earlier question', 'until now', 'till date', 'all my conversations' or 'previously mentioned'.
+    The chat history is semantically indexed and summarized using vector search."""
+
+    logger.info(f"search_chat_history called with: {question}")
+    api_key=await get_any_api_key()
+    if not api_key:
+        logger.error("API key is not available. Cannot search chat history.")
+        return "OpenAI API key is not available. Cannot search chat history."
+    retriever = await get_chat_history_retriever(api_key=api_key)
+    response = await retriever.ainvoke({"query": question})
+    result = response["result"]
+    logger.info(f"RESULT: {result}")
+    return result
+
 if __name__ == "__main__":
     mcp_server_port = int(os.environ.get("MCP_SERVER_PORT", 5500))
     mcp.run(transport="streamable-http", host="0.0.0.0", port=mcp_server_port,)

services/chatbot/src/mcpserver/tool_helpers.py

@@ -0,0 +1,48 @@
+import os
+from langchain_community.embeddings import OpenAIEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.prompts import PromptTemplate
+from chatbot.extensions import db
+from langchain.chains import RetrievalQA
+from langchain_openai import ChatOpenAI
+
+retrieval_index_path = "/app/resources/chat_index"
+
+async def get_any_api_key():
+    if os.environ.get("CHATBOT_OPENAI_API_KEY"):
+        return os.environ.get("CHATBOT_OPENAI_API_KEY")
+    doc = await db.sessions.find_one(
+        {"openai_api_key": {"$exists": True, "$ne": None}},
+        {"openai_api_key": 1}
+    )
+    if doc and "openai_api_key" in doc:
+        return doc["openai_api_key"]    
+    return None
+
+async def get_chat_history_retriever(api_key: str):
+    prompt_template = PromptTemplate.from_template(
+        """You are an assistant that summarizes chat history across sessions.
+
+        Given the following chat excerpts:
+        {context}
+        Answer the user's question: {question}
+
+        If the user asks for a summary, provide a coherent, high-level summary of the conversations in natural language.
+        If the user asks a specific question, extract and answer it from the chats.
+        Be detailed, accurate, and neutral."""
+    )
+    embeddings = OpenAIEmbeddings(api_key=api_key)
+    vectorstore = FAISS.load_local(
+        retrieval_index_path,
+        embeddings,
+        allow_dangerous_deserialization=True
+    )
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=ChatOpenAI(api_key=api_key, model="gpt-4o"),
+        retriever=retriever,
+        chain_type="stuff",
+        chain_type_kwargs={"prompt": prompt_template, "document_variable_name": "context"},
+        return_source_documents=False,
+    )
+    return qa_chain