From 4652d32bc70e5821443669444ecb9db412aa2c73 Mon Sep 17 00:00:00 2001 From: Keyur Doshi Date: Thu, 17 Jul 2025 15:04:30 +0530 Subject: [PATCH 1/3] Model selection implemented --- services/chatbot/.env | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/services/chatbot/.env b/services/chatbot/.env index 349ae2ec..8a1be487 100644 --- a/services/chatbot/.env +++ b/services/chatbot/.env @@ -1,2 +1,3 @@ PERSIST_DIRECTORY=db -TARGET_SOURCE_CHUNKS=4 \ No newline at end of file +TARGET_SOURCE_CHUNKS=4 +DEFAULT_MODEL=gpt-4o-mini \ No newline at end of file From e07c90904768a8cdea8e4fd88eb9e16cecb145c9 Mon Sep 17 00:00:00 2001 From: Keyur Doshi Date: Thu, 17 Jul 2025 16:11:35 +0530 Subject: [PATCH 2/3] Refactor: moved default model env variable to correct files --- services/chatbot/.env | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/services/chatbot/.env b/services/chatbot/.env index 8a1be487..349ae2ec 100644 --- a/services/chatbot/.env +++ b/services/chatbot/.env @@ -1,3 +1,2 @@ PERSIST_DIRECTORY=db -TARGET_SOURCE_CHUNKS=4 -DEFAULT_MODEL=gpt-4o-mini \ No newline at end of file +TARGET_SOURCE_CHUNKS=4 \ No newline at end of file From 26f5a81c844430eae278552d9b4824d013a77ffb Mon Sep 17 00:00:00 2001 From: Keyur Doshi Date: Mon, 21 Jul 2025 18:23:40 +0530 Subject: [PATCH 3/3] User context provided --- services/chatbot/src/chatbot/chat_api.py | 4 ++- services/chatbot/src/chatbot/chat_service.py | 4 +-- .../chatbot/src/chatbot/langgraph_agent.py | 9 ++++--- services/chatbot/src/chatbot/mcp_client.py | 26 ++++++++++--------- .../chatbot/src/chatbot/session_service.py | 6 +++++ 5 files changed, 30 insertions(+), 19 deletions(-) diff --git a/services/chatbot/src/chatbot/chat_api.py b/services/chatbot/src/chatbot/chat_api.py index 1c8b0211..ed403157 100644 --- a/services/chatbot/src/chatbot/chat_api.py +++ b/services/chatbot/src/chatbot/chat_api.py @@ -10,6 +10,7 @@ get_or_create_session_id, store_api_key, store_model_name, + get_user_jwt ) chat_bp = Blueprint("chat", __name__, url_prefix="/genai") @@ -53,6 +54,7 @@ async def chat(): session_id = await get_or_create_session_id() openai_api_key = await get_api_key(session_id) model_name = await get_model_name(session_id) + user_jwt = await get_user_jwt() if not openai_api_key: return jsonify({"message": "Missing OpenAI API key. Please authenticate."}), 400 data = await request.get_json() @@ -60,7 +62,7 @@ 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) + 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 7f38907e..41846c4f 100644 --- a/services/chatbot/src/chatbot/chat_service.py +++ b/services/chatbot/src/chatbot/chat_service.py @@ -22,13 +22,13 @@ async def delete_chat_history(session_id): await db.chat_sessions.delete_one({"session_id": session_id}) -async def process_user_message(session_id, user_message, api_key, model_name): +async def process_user_message(session_id, user_message, api_key, model_name, user_jwt): history = await get_chat_history(session_id) # generate a unique numeric id for the message that is random but unique 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, 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) diff --git a/services/chatbot/src/chatbot/langgraph_agent.py b/services/chatbot/src/chatbot/langgraph_agent.py index ed164a4f..4ae5e450 100644 --- a/services/chatbot/src/chatbot/langgraph_agent.py +++ b/services/chatbot/src/chatbot/langgraph_agent.py @@ -18,7 +18,7 @@ from langgraph.prebuilt import create_react_agent from .extensions import postgresdb -from .mcp_client import mcp_client +from .mcp_client import get_mcp_client async def get_retriever_tool(api_key): @@ -46,7 +46,7 @@ async def get_retriever_tool(api_key): return retriever_tool -async def build_langgraph_agent(api_key, model_name): +async def build_langgraph_agent(api_key, model_name, user_jwt): system_prompt = textwrap.dedent( """ You are crAPI Assistant — an expert agent that helps users explore and test the Completely Ridiculous API (crAPI), a vulnerable-by-design application for learning and evaluating modern API security issues. @@ -86,6 +86,7 @@ async def build_langgraph_agent(api_key, model_name): ) llm = ChatOpenAI(api_key=api_key, model=model_name) toolkit = SQLDatabaseToolkit(db=postgresdb, llm=llm) + mcp_client = get_mcp_client(user_jwt) mcp_tools = await mcp_client.get_tools() db_tools = toolkit.get_tools() tools = mcp_tools + db_tools @@ -95,8 +96,8 @@ async def build_langgraph_agent(api_key, model_name): return agent_node -async def execute_langgraph_agent(api_key, model_name, messages, session_id=None): - agent = await build_langgraph_agent(api_key, model_name) +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) response = await agent.ainvoke({"messages": messages}) diff --git a/services/chatbot/src/chatbot/mcp_client.py b/services/chatbot/src/chatbot/mcp_client.py index 38d98e44..c2632184 100644 --- a/services/chatbot/src/chatbot/mcp_client.py +++ b/services/chatbot/src/chatbot/mcp_client.py @@ -1,14 +1,16 @@ -import asyncio -import os - from langchain_mcp_adapters.client import MultiServerMCPClient -mcp_client = MultiServerMCPClient( - { - "crapi": { - "transport": "streamable_http", - "url": "http://localhost:5500/mcp/", - "headers": {}, - }, - } -) +def get_mcp_client(user_jwt: str | None) -> MultiServerMCPClient: + headers = {} + if user_jwt: + headers["Authorization"] = f"Bearer {user_jwt}" + + return MultiServerMCPClient( + { + "crapi": { + "transport": "streamable_http", + "url": "http://localhost:5500/mcp/", + "headers": headers, + } + } + ) \ No newline at end of file diff --git a/services/chatbot/src/chatbot/session_service.py b/services/chatbot/src/chatbot/session_service.py index 54902cb8..9216d68c 100644 --- a/services/chatbot/src/chatbot/session_service.py +++ b/services/chatbot/src/chatbot/session_service.py @@ -57,3 +57,9 @@ async def get_model_name(session_id): if "model_name" not in doc: 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 "): + return auth.replace("Bearer ", "") + return None