Retriever provided at agent creation

Author: keyurdoshi03

services/chatbot/src/chatbot/chat_service.py

@@ -1,6 +1,6 @@
 from uuid import uuid4
 from langgraph.graph.message import Messages
-from .vectordb import update_collection
+from .retriever import update_collection
 from .extensions import db
 from .langgraph_agent import execute_langgraph_agent
 

services/chatbot/src/chatbot/langgraph_agent.py

@@ -4,7 +4,7 @@
 from langgraph.prebuilt import create_react_agent
 from .extensions import postgresdb
 from .mcp_client import get_mcp_client
-
+from .retriever import get_retriever_tool
 
 async def build_langgraph_agent(api_key, model_name, user_jwt):
     system_prompt = textwrap.dedent(
@@ -49,7 +49,8 @@ async def build_langgraph_agent(api_key, model_name, user_jwt):
     mcp_client = get_mcp_client(user_jwt)
     mcp_tools = await mcp_client.get_tools()
     db_tools = toolkit.get_tools()
-    tools = mcp_tools + db_tools
+    retriever_tool = get_retriever_tool(api_key)
+    tools = mcp_tools + db_tools + [retriever_tool]
     agent_node = create_react_agent(model=llm, tools=tools, prompt=system_prompt)
     return agent_node
 

services/chatbot/src/chatbot/retriever.py

@@ -0,0 +1,67 @@
+from langchain.agents.agent_toolkits import create_retriever_tool
+from .config import Config
+import chromadb
+from uuid import uuid4
+from langchain_core.documents import Document
+from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
+from langchain_openai import OpenAIEmbeddings
+from langchain_chroma import Chroma
+
+async def get_collection(api_key):
+    chroma_client = await chromadb.AsyncHttpClient(
+        host=Config.CHROMA_HOST,
+        port=Config.CHROMA_PORT,
+    )
+    embeddings = OpenAIEmbeddingFunction(api_key=api_key, model_name="text-embedding-3-large")
+    collection = await chroma_client.get_or_create_collection(
+        name="chats",
+        embedding_function=embeddings,
+    )
+    return collection
+
+async def update_collection(api_key, session_id, new_messages):
+    docs = []
+    ids = []
+    for role, content in new_messages.items():
+        if content:
+            doc_id = str(uuid4())
+            doc = Document(
+                page_content=content,
+                metadata={"session_id": session_id, "role": role}
+            )
+            docs.append(doc)
+            ids.append(doc_id)
+
+    if docs:
+        collection = await get_collection(api_key)
+        await collection.add(
+            documents=[d.page_content for d in docs],
+            metadatas=[d.metadata for d in docs],
+            ids=ids
+        )
+
+def get_retriever_tool(api_key):
+    embeddings = OpenAIEmbeddings(api_key=api_key, model="text-embedding-3-large")
+    chroma_client = chromadb.HttpClient(
+        host=Config.CHROMA_HOST,
+        port=Config.CHROMA_PORT,
+    )
+    collection = Chroma(
+        collection_name="chats",
+        embedding_function=embeddings,
+        client=chroma_client,
+    )
+    retriever = collection.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    return create_retriever_tool(
+        retriever=retriever,
+        name="retriever_tool",
+        description="""
+        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.
+        """,
+    )

services/chatbot/src/chatbot/vectordb.py

@@ -10,6 +10,4 @@ class Config:
     CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY", "/app/vectorstore")
     OPENAPI_SPEC = os.getenv("OPENAPI_SPEC", "/app/resources/crapi-openapi-spec.json")
     API_USER = os.getenv("API_USER", "[email protected]")
-    API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123")
-    CHROMA_HOST = os.environ.get("CHROMA_HOST", "chromadb")
-    CHROMA_PORT = os.environ.get("CHROMA_PORT", "8000")
+    API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123")

services/chatbot/src/mcpserver/config.py

@@ -5,10 +5,6 @@
 from .config import Config
 import logging
 import time
-from .tool_helpers import (
-    get_any_api_key,
-    build_retrieverQA,
-)
 
 # Configure logging
 logging.basicConfig(
@@ -78,27 +74,6 @@ def get_http_client():
     openapi_spec=openapi_spec, client=get_http_client(), name="My crAPI MCP Server"
 )
 
-@mcp.tool(tags={"history", "search", "summary", "context"},)
-async def retriever_tool(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."
-    retrieverQA = build_retrieverQA(api_key=api_key)
-    response = await retrieverQA.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(