Chromadb client-server flow

Author: keyurdoshi03

deploy/docker/docker-compose.yml

@@ -268,6 +268,12 @@ services:
     image: 'chromadb/chroma:latest'
     environment:
       IS_PERSISTENT: 'TRUE'
+    healthcheck:
+      test: [ "CMD", "/bin/bash", "-c", "cat < /dev/null > /dev/tcp/localhost/8000" ]
+      interval: 15s
+      timeout: 15s
+      retries: 15
+      start_period: 20s
     volumes:
       - chromadb-data:/data
     # ports:

deploy/helm/values.yaml

@@ -327,7 +327,6 @@ apiGatewayService:
   serviceSelectorLabels:
     app: gateway-service
 
-
 chromadb:
   name: chromadb
   image: chromadb/chroma

services/chatbot/requirements.txt

@@ -19,4 +19,5 @@ faiss-cpu==1.11.0
 psycopg2-binary
 uvicorn==0.35.0
 fastmcp==2.10.2
-chromadb-client==1.0.15
+chromadb-client==1.0.15
+langchain-chroma==0.2.4

services/chatbot/src/chatbot/chat_service.py

@@ -1,6 +1,6 @@
 from uuid import uuid4
 from langgraph.graph.message import Messages
-from services.chatbot.src.chatbot.retrieverutils import add_to_chroma_collection
+from .vectordb import update_collection
 from .extensions import db
 from .langgraph_agent import execute_langgraph_agent
 
@@ -37,9 +37,7 @@ async def process_user_message(session_id, user_message, api_key, model_name, us
     history.append(
         {"id": response_message_id, "role": "assistant", "content": reply.content}
     )
-    await add_to_chroma_collection(
-        api_key, session_id, [{"user": user_message}, {"assistant": reply.content}]
-    )
+    await update_collection(api_key, session_id, {"user": user_message, "assistant": reply.content})
     # Limit chat history to last 20 messages
     history = history[-20:]
     await update_chat_history(session_id, history)

services/chatbot/src/chatbot/langgraph_agent.py

@@ -1,28 +1,9 @@
-import os
 import textwrap
-from typing import Annotated, Sequence, TypedDict
-
-from langchain.agents.agent_toolkits import create_retriever_tool
-from langchain.chains import LLMChain, RetrievalQA
-from langchain.prompts import PromptTemplate
-from langchain.schema import BaseMessage
-from langchain.tools import Tool
 from langchain_community.agent_toolkits import SQLDatabaseToolkit
-from langchain_community.agent_toolkits.sql.base import create_sql_agent
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_community.embeddings import OpenAIEmbeddings
-from langchain_community.vectorstores import FAISS  # or Chroma, Weaviate, etc.
 from langchain_openai import ChatOpenAI
-from langgraph.graph import MessageGraph, StateGraph
-from langgraph.graph.message import add_messages
 from langgraph.prebuilt import create_react_agent
-from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE, Settings
-
-
 from .extensions import postgresdb
-from .config import Config
 from .mcp_client import get_mcp_client
-import chromadb
 
 
 async def build_langgraph_agent(api_key, model_name, user_jwt):
@@ -69,8 +50,6 @@ async def build_langgraph_agent(api_key, model_name, user_jwt):
     mcp_tools = await mcp_client.get_tools()
     db_tools = toolkit.get_tools()
     tools = mcp_tools + db_tools
-    retriever_tool = await get_retriever_tool(api_key)
-    tools.append(retriever_tool)
     agent_node = create_react_agent(model=llm, tools=tools, prompt=system_prompt)
     return agent_node
 

services/chatbot/src/chatbot/retrieverutils.py

@@ -0,0 +1,56 @@
+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
+
+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
+        )
+
+# async def get_retriever_tool(api_key):
+#     collection = await get_chroma_collection(api_key)
+#     retriever = collection.as_retriever()
+#     retriever_tool = create_retriever_tool(
+#         retriever,
+#         name="chat_rag",
+#         description="""
+#         Use this to 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.
+        
+#         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.
+#         """,
+#     )
+#     return retriever_tool

services/chatbot/src/chatbot/vectordb.py

@@ -10,4 +10,6 @@ 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")
+    API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123")
+    CHROMA_HOST = os.environ.get("CHROMA_HOST", "chromadb")
+    CHROMA_PORT = os.environ.get("CHROMA_PORT", "8000")

services/chatbot/src/mcpserver/config.py

@@ -1,13 +1,13 @@
 import httpx
-from fastmcp import FastMCP, settings
+from fastmcp import FastMCP
 import json
 import os
 from .config import Config
 import logging
 import time
 from .tool_helpers import (
     get_any_api_key,
-    get_chat_history_retriever,
+    build_retrieverQA,
 )
 
 # Configure logging
@@ -78,6 +78,27 @@ 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(

services/chatbot/src/mcpserver/server.py

@@ -1,11 +1,13 @@
 import os
+import chromadb
 from langchain_community.embeddings import OpenAIEmbeddings
-from langchain_community.vectorstores import Chroma
 from langchain.prompts import PromptTemplate
 from chatbot.extensions import db
 from .config import Config
 from langchain.chains import RetrievalQA
 from langchain_openai import ChatOpenAI
+from langchain_chroma import Chroma
+
 
 async def get_any_api_key():
     if os.environ.get("CHATBOT_OPENAI_API_KEY"):
@@ -16,3 +18,36 @@ async def get_any_api_key():
     if doc and "openai_api_key" in doc:
         return doc["openai_api_key"]
     return None
+
+
+def build_retrieverQA(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."""
+    )
+
+    chroma_client = chromadb.HttpClient(
+        host=Config.CHROMA_HOST,
+        port=Config.CHROMA_PORT,
+    )
+    embeddings = OpenAIEmbeddings(api_key=api_key, model="text-embedding-3-large")
+    vectorstore = Chroma(
+        collection_name="chats",
+        embedding_function=embeddings,
+        client=chroma_client,
+    )
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    return 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,
+    )