From c6a6e72379792c27b96147b08d370cf47e739858 Mon Sep 17 00:00:00 2001 From: Roshan Piyush Date: Sun, 27 Jul 2025 18:19:27 +0530 Subject: [PATCH 1/4] Add chromadb --- deploy/docker/docker-compose.yml | 19 ++++- deploy/helm/templates/chatbot/config.yaml | 3 +- deploy/helm/templates/chatbot/deployment.yaml | 7 -- deploy/helm/templates/chatbot/storage.yaml | 33 -------- deploy/helm/templates/chromadb/config.yaml | 11 +++ deploy/helm/templates/chromadb/service.yaml | 15 ++++ .../helm/templates/chromadb/statefulset.yaml | 38 +++++++++ deploy/helm/templates/chromadb/storage.yaml | 34 ++++++++ services/chatbot/requirements.txt | 2 +- services/chatbot/src/chatbot/chat_api.py | 8 +- services/chatbot/src/chatbot/chat_service.py | 13 +-- services/chatbot/src/chatbot/config.py | 5 +- services/chatbot/src/chatbot/dbconnections.py | 3 + .../chatbot/src/chatbot/langgraph_agent.py | 37 +++------ services/chatbot/src/chatbot/mcp_client.py | 3 +- .../chatbot/src/chatbot/retrieverutils.py | 80 +++++++++++++++++++ .../chatbot/src/chatbot/session_service.py | 3 + services/chatbot/src/chatbot/vector_index.py | 23 ------ services/chatbot/src/mcpserver/__main__.py | 2 +- services/chatbot/src/mcpserver/server.py | 48 +++++------ .../chatbot/src/mcpserver/tool_helpers.py | 33 +------- 21 files changed, 251 insertions(+), 169 deletions(-) delete mode 100644 deploy/helm/templates/chatbot/storage.yaml create mode 100644 deploy/helm/templates/chromadb/config.yaml create mode 100644 deploy/helm/templates/chromadb/service.yaml create mode 100644 deploy/helm/templates/chromadb/statefulset.yaml create mode 100644 deploy/helm/templates/chromadb/storage.yaml create mode 100644 services/chatbot/src/chatbot/retrieverutils.py delete mode 100644 services/chatbot/src/chatbot/vector_index.py diff --git a/deploy/docker/docker-compose.yml b/deploy/docker/docker-compose.yml index 50e2d85a..5d1771df 100755 --- 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 index 215d9af8..1645d6ce 100644 --- 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 }} \ No newline at end of file diff --git a/deploy/helm/templates/chatbot/deployment.yaml b/deploy/helm/templates/chatbot/deployment.yaml index f58c047f..692dfa0c 100644 --- 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 deleted file mode 100644 index bace224f..00000000 --- a/deploy/helm/templates/chatbot/storage.yaml +++ /dev/null @@ -1,33 +0,0 @@ -{{- if eq .Values.chatbot.storage.type "manual" }} -apiVersion: v1 -kind: PersistentVolume -metadata: - name: {{ .Values.chatbot.storage.pv.name }} - labels: - release: {{ .Release.Name }} - {{- toYaml .Values.chatbot.storage.pv.labels | nindent 4 }} -spec: - storageClassName: {{ .Values.chatbot.storage.type }} - capacity: - storage: {{ .Values.chatbot.storage.pv.resources.storage }} - accessModes: - - ReadWriteOnce - hostPath: - path: {{ .Values.chatbot.storage.pv.hostPath }} ---- -{{- end }} -apiVersion: v1 -kind: PersistentVolumeClaim -metadata: - name: {{ .Values.chatbot.storage.pvc.name }} - labels: - release: {{ .Release.Name }} - {{- toYaml .Values.chatbot.storage.pvc.labels | nindent 4 }} -spec: - {{- if ne .Values.chatbot.storage.type "default" }} - storageClassName: {{ .Values.chatbot.storage.type }} - {{- end }} - accessModes: - - ReadWriteOnce - resources: - {{- toYaml .Values.chatbot.storage.pvc.resources | nindent 4 }} diff --git a/deploy/helm/templates/chromadb/config.yaml b/deploy/helm/templates/chromadb/config.yaml new file mode 100644 index 00000000..30b8d4a1 --- /dev/null +++ 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 new file mode 100644 index 00000000..f9eca312 --- /dev/null +++ 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 new file mode 100644 index 00000000..1af046ee --- /dev/null +++ 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 }} + \ No newline at end of file diff --git a/deploy/helm/templates/chromadb/storage.yaml b/deploy/helm/templates/chromadb/storage.yaml new file mode 100644 index 00000000..b3a1873d --- /dev/null +++ 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 index d4cbc666..89fc448a 100644 --- 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 \ No newline at end of file +chromadb-client==1.0.15 \ No newline at end of file diff --git a/services/chatbot/src/chatbot/chat_api.py b/services/chatbot/src/chatbot/chat_api.py index ed403157..798f1e84 100644 --- 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 index 8a0491f6..f94429e3 100644 --- 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,7 +27,9 @@ 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) @@ -35,11 +37,10 @@ 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}] + ) # 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 index eb7d81f7..7f3d2ecf 100644 --- 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 index b283d79c..385546f9 100644 --- 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 index 4ae5e450..d942cc48 100644 --- 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 index c2632184..f9d8678e 100644 --- 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, } } - ) \ No newline at end of file + ) diff --git a/services/chatbot/src/chatbot/retrieverutils.py b/services/chatbot/src/chatbot/retrieverutils.py new file mode 100644 index 00000000..4bea1727 --- /dev/null +++ b/services/chatbot/src/chatbot/retrieverutils.py @@ -0,0 +1,80 @@ +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 + +from langchain_community.embeddings import OpenAIEmbeddings +from langchain_community.vectorstores import Chroma +from langchain_core.documents import Document +from .config import Config +from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction + + +async def get_chroma_collection(api_key): + chroma_client = chromadb.AsyncHttpClient( + host=Config.CHROMA_HOST, + port=Config.CHROMA_PORT, + ssl=False, + headers=None, + settings=Settings(), + tenant=DEFAULT_TENANT, + database=DEFAULT_DATABASE, + ) + + collection = chroma_client.get_or_create_collection( + name="chats", + embedding_function=OpenAIEmbeddingFunction( + api_key=api_key, + model="text-embedding-3-large", + ), + ) + return collection + + +async def add_to_chroma_collection(api_key, session_id, new_messages): + collection = await get_chroma_collection(api_key) + collection.add( + documents=[ + {"content": content, "metadata": {"session_id": session_id, "role": role}} + for role, content in new_messages.items() + ] + ) + + +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 diff --git a/services/chatbot/src/chatbot/session_service.py b/services/chatbot/src/chatbot/session_service.py index 9216d68c..53612c21 100644 --- a/services/chatbot/src/chatbot/session_service.py +++ b/services/chatbot/src/chatbot/session_service.py @@ -45,11 +45,13 @@ async def delete_api_key(session_id): {"session_id": session_id}, {"$unset": {"openai_api_key": ""}} ) + async def store_model_name(session_id, model_name): await db.sessions.update_one( {"session_id": session_id}, {"$set": {"model_name": model_name}}, upsert=True ) + async def get_model_name(session_id): doc = await db.sessions.find_one({"session_id": session_id}) if not doc: @@ -58,6 +60,7 @@ async def get_model_name(session_id): return Config.DEFAULT_MODEL_NAME return doc["model_name"] + async def get_user_jwt() -> str | None: auth = request.headers.get("Authorization", "") if auth.startswith("Bearer "): diff --git a/services/chatbot/src/chatbot/vector_index.py b/services/chatbot/src/chatbot/vector_index.py deleted file mode 100644 index 8397ea9a..00000000 --- a/services/chatbot/src/chatbot/vector_index.py +++ /dev/null @@ -1,23 +0,0 @@ -from langchain_community.embeddings import OpenAIEmbeddings -from langchain_community.vectorstores import Chroma -from langchain_core.documents import Document -from .config import Config - -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, model="text-embedding-3-large") - vectorstore = Chroma( - embedding_function=embeddings, - persist_directory=Config.CHROMA_PERSIST_DIRECTORY - ) - vectorstore.add_documents(docs) - vectorstore.persist() \ No newline at end of file diff --git a/services/chatbot/src/mcpserver/__main__.py b/services/chatbot/src/mcpserver/__main__.py index 1e67dc86..ffd62d3f 100644 --- a/services/chatbot/src/mcpserver/__main__.py +++ b/services/chatbot/src/mcpserver/__main__.py @@ -12,4 +12,4 @@ if __name__ == "__main__": logger.info("Starting MCP server...") mcp_server_port = int(os.environ.get("MCP_SERVER_PORT", 5500)) - app.run(transport="streamable-http", host="0.0.0.0", port=mcp_server_port) + app.run(transport="streamable-http", host="0.0.0.0", port=mcp_server_port) diff --git a/services/chatbot/src/mcpserver/server.py b/services/chatbot/src/mcpserver/server.py index 78760626..58aec289 100644 --- a/services/chatbot/src/mcpserver/server.py +++ b/services/chatbot/src/mcpserver/server.py @@ -8,6 +8,7 @@ get_any_api_key, get_chat_history_retriever, ) + # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" @@ -28,6 +29,7 @@ API_KEY = None API_AUTH_TYPE = "ApiKey" + def get_api_key(): global API_KEY # Try 5 times to get API key @@ -47,9 +49,15 @@ def get_api_key(): response = client.post(apikey_url, json=login_body) if response.status_code != 200: if i == MAX_ATTEMPTS - 1: - logger.error(f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}") - raise Exception(f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}") - logger.error(f"Failed to get API key in attempt {i+1}: {response.status_code} {response.text}. Sleeping for {i} seconds...") + logger.error( + f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}" + ) + raise Exception( + f"Failed to get API key after {i+1} attempts: {response.status_code} {response.text}" + ) + logger.error( + f"Failed to get API key in attempt {i+1}: {response.status_code} {response.text}. Sleeping for {i} seconds..." + ) time.sleep(i) response_json = response.json() logger.info(f"Response: {response_json}") @@ -70,38 +78,20 @@ def get_http_client(): headers=headers, ) -# Load your OpenAPI spec + +# Load your OpenAPI spec with open("/app/resources/crapi-openapi-spec.json", "r") as f: openapi_spec = json.load(f) # Create the MCP server mcp = FastMCP.from_openapi( - openapi_spec=openapi_spec, - client=get_http_client(), - name="My crAPI MCP Server" + openapi_spec=openapi_spec, client=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,) + mcp.run( + transport="streamable-http", + host="0.0.0.0", + port=mcp_server_port, + ) diff --git a/services/chatbot/src/mcpserver/tool_helpers.py b/services/chatbot/src/mcpserver/tool_helpers.py index d78066ce..b760e489 100644 --- a/services/chatbot/src/mcpserver/tool_helpers.py +++ b/services/chatbot/src/mcpserver/tool_helpers.py @@ -7,42 +7,13 @@ 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} + {"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 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, model="text-embedding-3-large") - vectorstore = Chroma( - embedding_function=embeddings, - persist_directory=Config.CHROMA_PERSIST_DIRECTORY - ) - 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 From 73657bbf0aa4114f9be539c1a164f424e96145c5 Mon Sep 17 00:00:00 2001 From: Roshan Piyush Date: Wed, 30 Jul 2025 00:08:17 +0530 Subject: [PATCH 2/4] Update values.yaml --- deploy/helm/values.yaml | 42 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/deploy/helm/values.yaml b/deploy/helm/values.yaml index 94864f11..3146869e 100644 --- a/deploy/helm/values.yaml +++ b/deploy/helm/values.yaml @@ -326,3 +326,45 @@ apiGatewayService: app: gateway-service serviceSelectorLabels: app: gateway-service + + +chromadb: + name: chromadb + image: chromadb/chroma + version: latest + imagePullPolicy: IfNotPresent + port: 8000 + replicaCount: 1 + service: + name: chromadb + labels: + app: chromadb + config: + name: chromadb-config + labels: + app: chromadb + storage: + # type: "manual" + # pv: + # name: chromadb-pv + # labels: + # app: chromadb + # resources: + # storage: 1Gi + # hostPath: /mnt/chromadb + type: "default" + pvc: + name: chromadb-pv-claim + labels: + app: chromadb + resources: + requests: + storage: 2Gi + serviceSelectorLabels: + app: chromadb + podLabels: + app: chromadb + statefulsetLabels: + app: chromadb + statefulsetSelectorMatchLabels: + app: chromadb \ No newline at end of file From ec804d6cdc64dcf8962ab51762a2b008d1e544c8 Mon Sep 17 00:00:00 2001 From: Keyur Doshi Date: Thu, 31 Jul 2025 00:18:11 +0530 Subject: [PATCH 3/4] Chromadb client-server flow --- deploy/docker/docker-compose.yml | 6 ++ deploy/helm/values.yaml | 1 - services/chatbot/requirements.txt | 3 +- services/chatbot/src/chatbot/chat_service.py | 6 +- .../chatbot/src/chatbot/langgraph_agent.py | 21 ----- .../chatbot/src/chatbot/retrieverutils.py | 80 ------------------- services/chatbot/src/chatbot/vectordb.py | 56 +++++++++++++ services/chatbot/src/mcpserver/config.py | 4 +- services/chatbot/src/mcpserver/server.py | 25 +++++- .../chatbot/src/mcpserver/tool_helpers.py | 37 ++++++++- 10 files changed, 128 insertions(+), 111 deletions(-) delete mode 100644 services/chatbot/src/chatbot/retrieverutils.py create mode 100644 services/chatbot/src/chatbot/vectordb.py diff --git a/deploy/docker/docker-compose.yml b/deploy/docker/docker-compose.yml index 2caab5c0..6fc2c725 100755 --- a/deploy/docker/docker-compose.yml +++ b/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: diff --git a/deploy/helm/values.yaml b/deploy/helm/values.yaml index 3146869e..4a7921d8 100644 --- a/deploy/helm/values.yaml +++ b/deploy/helm/values.yaml @@ -327,7 +327,6 @@ apiGatewayService: serviceSelectorLabels: app: gateway-service - chromadb: name: chromadb image: chromadb/chroma diff --git a/services/chatbot/requirements.txt b/services/chatbot/requirements.txt index 89fc448a..d85c39ba 100644 --- a/services/chatbot/requirements.txt +++ b/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 \ No newline at end of file +chromadb-client==1.0.15 +langchain-chroma==0.2.4 \ No newline at end of file diff --git a/services/chatbot/src/chatbot/chat_service.py b/services/chatbot/src/chatbot/chat_service.py index f94429e3..a3bd1f73 100644 --- 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 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) diff --git a/services/chatbot/src/chatbot/langgraph_agent.py b/services/chatbot/src/chatbot/langgraph_agent.py index d942cc48..e5a4fcf4 100644 --- a/services/chatbot/src/chatbot/langgraph_agent.py +++ b/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 diff --git a/services/chatbot/src/chatbot/retrieverutils.py b/services/chatbot/src/chatbot/retrieverutils.py deleted file mode 100644 index 4bea1727..00000000 --- a/services/chatbot/src/chatbot/retrieverutils.py +++ /dev/null @@ -1,80 +0,0 @@ -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 - -from langchain_community.embeddings import OpenAIEmbeddings -from langchain_community.vectorstores import Chroma -from langchain_core.documents import Document -from .config import Config -from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction - - -async def get_chroma_collection(api_key): - chroma_client = chromadb.AsyncHttpClient( - host=Config.CHROMA_HOST, - port=Config.CHROMA_PORT, - ssl=False, - headers=None, - settings=Settings(), - tenant=DEFAULT_TENANT, - database=DEFAULT_DATABASE, - ) - - collection = chroma_client.get_or_create_collection( - name="chats", - embedding_function=OpenAIEmbeddingFunction( - api_key=api_key, - model="text-embedding-3-large", - ), - ) - return collection - - -async def add_to_chroma_collection(api_key, session_id, new_messages): - collection = await get_chroma_collection(api_key) - collection.add( - documents=[ - {"content": content, "metadata": {"session_id": session_id, "role": role}} - for role, content in new_messages.items() - ] - ) - - -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 diff --git a/services/chatbot/src/chatbot/vectordb.py b/services/chatbot/src/chatbot/vectordb.py new file mode 100644 index 00000000..f01f469b --- /dev/null +++ b/services/chatbot/src/chatbot/vectordb.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 diff --git a/services/chatbot/src/mcpserver/config.py b/services/chatbot/src/mcpserver/config.py index 49ea9fe8..34259a1d 100644 --- a/services/chatbot/src/mcpserver/config.py +++ b/services/chatbot/src/mcpserver/config.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", "admin@example.com") - API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123") \ No newline at end of file + API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123") + CHROMA_HOST = os.environ.get("CHROMA_HOST", "chromadb") + CHROMA_PORT = os.environ.get("CHROMA_PORT", "8000") \ No newline at end of file diff --git a/services/chatbot/src/mcpserver/server.py b/services/chatbot/src/mcpserver/server.py index 05630ed5..57d204ac 100644 --- a/services/chatbot/src/mcpserver/server.py +++ b/services/chatbot/src/mcpserver/server.py @@ -1,5 +1,5 @@ import httpx -from fastmcp import FastMCP, settings +from fastmcp import FastMCP import json import os from .config import Config @@ -7,7 +7,7 @@ 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( diff --git a/services/chatbot/src/mcpserver/tool_helpers.py b/services/chatbot/src/mcpserver/tool_helpers.py index 87783ed3..ecd84f60 100644 --- a/services/chatbot/src/mcpserver/tool_helpers.py +++ b/services/chatbot/src/mcpserver/tool_helpers.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, + ) \ No newline at end of file From 3efe47c821cc7c6faa4e776bcff7376b3eab06e6 Mon Sep 17 00:00:00 2001 From: Keyur Doshi Date: Thu, 31 Jul 2025 03:30:07 +0530 Subject: [PATCH 4/4] Retriever provided at agent creation --- services/chatbot/src/chatbot/chat_service.py | 2 +- .../chatbot/src/chatbot/langgraph_agent.py | 5 +- services/chatbot/src/chatbot/retriever.py | 67 +++++++++++++++++++ services/chatbot/src/chatbot/vectordb.py | 56 ---------------- services/chatbot/src/mcpserver/config.py | 4 +- services/chatbot/src/mcpserver/server.py | 25 ------- .../chatbot/src/mcpserver/tool_helpers.py | 53 --------------- 7 files changed, 72 insertions(+), 140 deletions(-) create mode 100644 services/chatbot/src/chatbot/retriever.py delete mode 100644 services/chatbot/src/chatbot/vectordb.py delete mode 100644 services/chatbot/src/mcpserver/tool_helpers.py diff --git a/services/chatbot/src/chatbot/chat_service.py b/services/chatbot/src/chatbot/chat_service.py index a3bd1f73..ee5dea1f 100644 --- 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 .vectordb import update_collection +from .retriever import update_collection from .extensions import db from .langgraph_agent import execute_langgraph_agent diff --git a/services/chatbot/src/chatbot/langgraph_agent.py b/services/chatbot/src/chatbot/langgraph_agent.py index e5a4fcf4..f76a8819 100644 --- a/services/chatbot/src/chatbot/langgraph_agent.py +++ b/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 diff --git a/services/chatbot/src/chatbot/retriever.py b/services/chatbot/src/chatbot/retriever.py new file mode 100644 index 00000000..c367d0b2 --- /dev/null +++ b/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. + """, + ) diff --git a/services/chatbot/src/chatbot/vectordb.py b/services/chatbot/src/chatbot/vectordb.py deleted file mode 100644 index f01f469b..00000000 --- a/services/chatbot/src/chatbot/vectordb.py +++ /dev/null @@ -1,56 +0,0 @@ -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 diff --git a/services/chatbot/src/mcpserver/config.py b/services/chatbot/src/mcpserver/config.py index 34259a1d..49ea9fe8 100644 --- a/services/chatbot/src/mcpserver/config.py +++ b/services/chatbot/src/mcpserver/config.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", "admin@example.com") - API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123") - CHROMA_HOST = os.environ.get("CHROMA_HOST", "chromadb") - CHROMA_PORT = os.environ.get("CHROMA_PORT", "8000") \ No newline at end of file + API_PASSWORD = os.getenv("API_PASSWORD", "Admin!123") \ No newline at end of file diff --git a/services/chatbot/src/mcpserver/server.py b/services/chatbot/src/mcpserver/server.py index 57d204ac..05059282 100644 --- a/services/chatbot/src/mcpserver/server.py +++ b/services/chatbot/src/mcpserver/server.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( diff --git a/services/chatbot/src/mcpserver/tool_helpers.py b/services/chatbot/src/mcpserver/tool_helpers.py deleted file mode 100644 index ecd84f60..00000000 --- a/services/chatbot/src/mcpserver/tool_helpers.py +++ /dev/null @@ -1,53 +0,0 @@ -import os -import chromadb -from langchain_community.embeddings import OpenAIEmbeddings -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"): - 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 - - -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, - ) \ No newline at end of file