Add chromadb

piyushroshan · piyushroshan · commit c6a6e7237979 · 2025-07-27T18:19:27.000+05:30
diff --git a/deploy/docker/docker-compose.yml b/deploy/docker/docker-compose.yml
@@ -167,15 +167,16 @@ services:
       - MONGO_DB_PASSWORD=crapisecretpassword
       - MONGO_DB_NAME=crapi
       - DEFAULT_MODEL=gpt-4o-mini
-      - CHROMA_PERSIST_DIRECTORY=/app/vectorstore
+      - CHROMA_HOST=chromadb
+      - CHROMA_PORT=8000
       # - CHATBOT_OPENAI_API_KEY=
-    volumes:
-      - chatbot-vectors:/app/vectorstore
     depends_on:
       mongodb:
         condition: service_healthy
       crapi-identity:
         condition: service_healthy
+      chromadb:
+        condition: service_healthy
     # ports:
     #  - "${LISTEN_IP:-127.0.0.1}:5002:5002"
 
@@ -257,6 +258,16 @@ services:
           cpus: '0.3'
           memory: 128M
 
+  chromadb:
+    container_name: chromadb
+    image: 'chromadb/chroma:latest'
+    environment:
+      IS_PERSISTENT: 'TRUE'
+    volumes:
+      - chromadb-data:/data
+    # ports:
+    #   - "${LISTEN_IP:-127.0.0.1}:8000:8000"
+
   mailhog:
     user: root
     container_name: mailhog
@@ -298,4 +309,4 @@ services:
 volumes:
   mongodb-data:
   postgresql-data:
-  chatbot-vectors:
+  chromadb-data:
diff --git a/deploy/helm/templates/chatbot/config.yaml b/deploy/helm/templates/chatbot/config.yaml
@@ -22,4 +22,5 @@ data:
     MONGO_DB_NAME: {{ .Values.mongodb.config.mongoDbName }}
     CHATBOT_OPENAI_API_KEY: {{ .Values.openAIApiKey }}
     DEFAULT_MODEL: {{ .Values.chatbot.config.defaultModel | quote }}
-    CHROMA_PERSIST_DIRECTORY: {{ .Values.chatbot.config.chromaPersistDirectory | quote }}
+    CHROMA_HOST: {{ .Values.chromadb.service.name }}
+    CHROMA_PORT: {{ .Values.chromadb.port | quote }}
diff --git a/deploy/helm/templates/chatbot/deployment.yaml b/deploy/helm/templates/chatbot/deployment.yaml
@@ -57,10 +57,3 @@ spec:
               port: {{ .Values.chatbot.port }}
             initialDelaySeconds: 15
             periodSeconds: 10
-          volumeMounts:
-            - name: chatbot-vectors
-              mountPath: {{ .Values.chatbot.config.chromaPersistDirectory | quote }}
-      volumes:
-        - name: chatbot-vectors
-          persistentVolumeClaim:
-            claimName: {{ .Values.chatbot.storage.pvc.name }}
diff --git a/deploy/helm/templates/chatbot/storage.yaml b/deploy/helm/templates/chatbot/storage.yaml
diff --git a/deploy/helm/templates/chromadb/config.yaml b/deploy/helm/templates/chromadb/config.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: {{ .Values.chromadb.config.name }}
+  labels:
+    release: {{ .Release.Name }}
+    {{- with .Values.chromadb.config.labels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+data:
+  IS_PERSISTENT: {{ .Values.chromadb.config.isPersistent | quote }}
diff --git a/deploy/helm/templates/chromadb/service.yaml b/deploy/helm/templates/chromadb/service.yaml
@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ .Values.chromadb.service.name }}
+  labels:
+    release: {{ .Release.Name }}
+    {{- with .Values.chromadb.service.labels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+spec:
+  ports:
+  - port: {{ .Values.chromadb.port }}
+    name: chromadb
+  selector:
+    {{- toYaml .Values.chromadb.serviceSelectorLabels | nindent 4 }}
diff --git a/deploy/helm/templates/chromadb/statefulset.yaml b/deploy/helm/templates/chromadb/statefulset.yaml
@@ -0,0 +1,38 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: {{ .Values.chromadb.name }}
+  labels:
+    release: {{ .Release.Name }}
+    {{- with .Values.chromadb.statefulsetLabels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+spec:
+  serviceName: {{ .Values.chromadb.service.name }}
+  replicas: {{ .Values.chromadb.replicaCount }}
+  selector:
+    matchLabels:
+      {{- toYaml .Values.chromadb.statefulsetSelectorMatchLabels | nindent 6 }}
+  template:
+    metadata:
+      labels:
+        release: {{ .Release.Name }}
+        {{- toYaml .Values.chromadb.podLabels | nindent 8 }}
+    spec:
+      containers:
+        - name: {{ .Values.chromadb.name }}
+          image: {{ .Values.chromadb.image }}:{{ .Values.chromadb.version }}
+          imagePullPolicy: {{ .Values.chromadb.imagePullPolicy }}
+          ports:
+            - containerPort: {{ .Values.chromadb.port }}
+          envFrom:
+            - configMapRef:
+                name: {{ .Values.chromadb.config.name }}
+          volumeMounts:
+            - mountPath: /data
+              name: chromadb-data
+      volumes:
+        - name: chromadb-data
+          persistentVolumeClaim:
+            claimName: {{ .Values.chromadb.storage.pvc.name }}
+            
diff --git a/deploy/helm/templates/chromadb/storage.yaml b/deploy/helm/templates/chromadb/storage.yaml
@@ -0,0 +1,34 @@
+{{- if eq .Values.chromadb.storage.type "manual" }}
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: {{ .Values.chromadb.storage.pv.name }}
+  labels:
+    release: {{ .Release.Name }}
+    {{- toYaml .Values.chromadb.storage.pv.labels | nindent 4 }}
+spec:
+  storageClassName: {{ .Values.chromadb.storage.type }}
+  capacity:
+    storage: {{ .Values.chromadb.storage.pv.resources.storage }}
+  accessModes:
+    - ReadWriteOnce
+  hostPath:
+    path: {{ .Values.chromadb.storage.pv.hostPath }}
+---
+{{- end }}
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: {{ .Values.chromadb.storage.pvc.name }}
+  labels:
+    release: {{ .Release.Name }}
+    {{- toYaml .Values.chromadb.storage.pvc.labels | nindent 4 }}
+spec:
+  {{- if ne .Values.chromadb.storage.type "default" }}
+  storageClassName: {{ .Values.chromadb.storage.type }}
+  {{- end }}
+  accessModes:
+    - ReadWriteOnce
+  resources:
+    {{- toYaml .Values.chromadb.storage.pvc.resources | nindent 4 }}
+
diff --git a/services/chatbot/requirements.txt b/services/chatbot/requirements.txt
@@ -19,4 +19,4 @@ faiss-cpu==1.11.0
 psycopg2-binary
 uvicorn==0.35.0
 fastmcp==2.10.2
-chromadb==1.0.15
+chromadb-client==1.0.15
diff --git a/services/chatbot/src/chatbot/chat_api.py b/services/chatbot/src/chatbot/chat_api.py
@@ -10,7 +10,7 @@
     get_or_create_session_id,
     store_api_key,
     store_model_name,
-    get_user_jwt
+    get_user_jwt,
 )
 
 chat_bp = Blueprint("chat", __name__, url_prefix="/genai")
@@ -38,6 +38,7 @@ async def init():
     await store_api_key(session_id, openai_api_key)
     return jsonify({"message": "Initialized"}), 200
 
+
 @chat_bp.route("/model", methods=["POST"])
 async def model():
     session_id = await get_or_create_session_id()
@@ -49,6 +50,7 @@ async def model():
     await store_model_name(session_id, model_name)
     return jsonify({"model_used": model_name}), 200
 
+
 @chat_bp.route("/ask", methods=["POST"])
 async def chat():
     session_id = await get_or_create_session_id()
@@ -62,7 +64,9 @@ async def chat():
     id = data.get("id", uuid4().int & (1 << 63) - 1)
     if not message:
         return jsonify({"message": "Message is required", "id": id}), 400
-    reply, response_id = await process_user_message(session_id, message, openai_api_key, model_name, user_jwt)
+    reply, response_id = await process_user_message(
+        session_id, message, openai_api_key, model_name, user_jwt
+    )
     return jsonify({"id": response_id, "message": reply}), 200
 
 
diff --git a/services/chatbot/src/chatbot/chat_service.py b/services/chatbot/src/chatbot/chat_service.py
@@ -1,6 +1,6 @@
 from uuid import uuid4
 from langgraph.graph.message import Messages
-from .vector_index import update_vector_index
+from services.chatbot.src.chatbot.retrieverutils import add_to_chroma_collection
 from .extensions import db
 from .langgraph_agent import execute_langgraph_agent
 
@@ -27,19 +27,20 @@ async def process_user_message(session_id, user_message, api_key, model_name, us
     source_message_id = uuid4().int & (1 << 63) - 1
     history.append({"id": source_message_id, "role": "user", "content": user_message})
     # Run LangGraph agent
-    response = await execute_langgraph_agent(api_key, model_name, history, user_jwt, session_id)
+    response = await execute_langgraph_agent(
+        api_key, model_name, history, user_jwt, session_id
+    )
     print("Response", response)
     reply: Messages = response.get("messages", [{}])[-1]
     print("Reply", reply.content)
     response_message_id = uuid4().int & (1 << 63) - 1
     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}]
+    )
     # 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
diff --git a/services/chatbot/src/chatbot/config.py b/services/chatbot/src/chatbot/config.py
@@ -2,7 +2,7 @@
 
 from dotenv import load_dotenv
 
-from .dbconnections import MONGO_CONNECTION_URI
+from .dbconnections import MONGO_CONNECTION_URI, CHROMA_HOST, CHROMA_PORT
 
 load_dotenv()
 
@@ -11,4 +11,5 @@ class Config:
     SECRET_KEY = os.getenv("SECRET_KEY", "super-secret")
     MONGO_URI = MONGO_CONNECTION_URI
     DEFAULT_MODEL_NAME = os.getenv("DEFAULT_MODEL", "gpt-4o-mini")
-    CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY", "/app/vectorstore")
+    CHROMA_HOST = CHROMA_HOST
+    CHROMA_PORT = CHROMA_PORT
diff --git a/services/chatbot/src/chatbot/dbconnections.py b/services/chatbot/src/chatbot/dbconnections.py
@@ -32,3 +32,6 @@
     POSTGRES_PORT,
     POSTGRES_DB,
 )
+
+CHROMA_HOST = os.environ.get("CHROMA_HOST", "chromadb")
+CHROMA_PORT = os.environ.get("CHROMA_PORT", "8000")
diff --git a/services/chatbot/src/chatbot/langgraph_agent.py b/services/chatbot/src/chatbot/langgraph_agent.py
@@ -16,34 +16,13 @@
 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
-
-
-async def get_retriever_tool(api_key):
-    embeddings = OpenAIEmbeddings(api_key=api_key)
-    if os.path.exists("faiss_index"):
-        vectorstore = FAISS.load_local(
-            "faiss_index", embeddings, allow_dangerous_deserialization=True
-        )
-    else:
-        retrival_dir = os.path.join(os.path.dirname(__file__), "../../retrieval")
-        loader = DirectoryLoader(retrival_dir)  # or PDF, Markdown, etc.
-        docs = loader.load()
-        vectorstore = FAISS.from_documents(docs, embeddings)
-        vectorstore.save_local("faiss_index")
-    retriever = vectorstore.as_retriever(
-        search_type="similarity", search_kwargs={"k": 3}
-    )
-
-    # ✅ Create RAG tool
-    retriever_tool = create_retriever_tool(
-        retriever,
-        name="crapi_rag",
-        description="Use this to answer questions about crAPI, its endpoints, flows, vulnerabilities, and APIs.",
-    )
-    return retriever_tool
+import chromadb
 
 
 async def build_langgraph_agent(api_key, model_name, user_jwt):
@@ -90,13 +69,15 @@ 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)
+    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
 
 
-async def execute_langgraph_agent(api_key, model_name, messages, user_jwt, session_id=None):
+async def execute_langgraph_agent(
+    api_key, model_name, messages, user_jwt, session_id=None
+):
     agent = await build_langgraph_agent(api_key, model_name, user_jwt)
     print("messages", messages)
     print("Session ID", session_id)
diff --git a/services/chatbot/src/chatbot/mcp_client.py b/services/chatbot/src/chatbot/mcp_client.py
@@ -1,5 +1,6 @@
 from langchain_mcp_adapters.client import MultiServerMCPClient
 
+
 def get_mcp_client(user_jwt: str | None) -> MultiServerMCPClient:
     headers = {}
     if user_jwt:
@@ -13,4 +14,4 @@ def get_mcp_client(user_jwt: str | None) -> MultiServerMCPClient:
                 "headers": headers,
             }
         }
-    )
+    )
diff --git a/services/chatbot/src/chatbot/retrieverutils.py b/services/chatbot/src/chatbot/retrieverutils.py
diff --git a/services/chatbot/src/chatbot/session_service.py b/services/chatbot/src/chatbot/session_service.py
diff --git a/services/chatbot/src/chatbot/vector_index.py b/services/chatbot/src/chatbot/vector_index.py
diff --git a/services/chatbot/src/mcpserver/__main__.py b/services/chatbot/src/mcpserver/__main__.py
diff --git a/services/chatbot/src/mcpserver/server.py b/services/chatbot/src/mcpserver/server.py
diff --git a/services/chatbot/src/mcpserver/tool_helpers.py b/services/chatbot/src/mcpserver/tool_helpers.py
