MCP init

.chainlit/config.toml

@@ -62,7 +62,7 @@ edit_message = true
     # Only the executables in the allow list can be used for MCP stdio server.
     # Only need the base name of the executable, e.g. "npx", not "/usr/bin/npx".
     # Please don't comment this line for now, we need it to parse the executable name.
-    allowed_executables = [ "npx", "uvx" ]
+    allowed_executables = [ "npx", "python", "node" ]
 
 [UI]
 # Name of the assistant.

app.py

@@ -1,9 +1,11 @@
 import os
-from typing import Optional, Dict
+from typing import Optional, Dict, List, Any
 import asyncio
+import json
 import chainlit as cl
 import google.generativeai as genai
 from openai import AsyncOpenAI
+from mcp import ClientSession
 
 
 # --- Logging setup ---
@@ -39,51 +41,207 @@ def oauth_callback(
     return None
 
 
+@cl.on_mcp_connect
+async def on_mcp_connect(connection, session: ClientSession):
+    """Called when a new MCP connection is established."""
+    await cl.Message(
+        f"Establishing connection with MCP server: `{connection.name}`..."
+    ).send()
+
+    try:
+        result = await session.list_tools()
+        tools_for_llm = [
+            {
+                "type": "function",
+                "function": {
+                    "name": t.name,
+                    "description": t.description,
+                    "parameters": t.inputSchema,
+                },
+            }
+            for t in result.tools
+        ]
+
+        all_mcp_tools = cl.user_session.get("mcp_tools", {})
+        all_mcp_tools[connection.name] = tools_for_llm
+        cl.user_session.set("mcp_tools", all_mcp_tools)
+
+        tool_names = [t.name for t in result.tools]
+        await cl.Message(
+            content=f"✅ Connection to `{connection.name}` successful. Found tools:\n`{', '.join(tool_names)}`"
+        ).send()
+    except Exception as e:
+        await cl.ErrorMessage(
+            f"Error discovering tools for `{connection.name}`: {e}"
+        ).send()
+
+
+@cl.on_mcp_disconnect
+async def on_mcp_disconnect(name: str, session: ClientSession):
+    """Called when an MCP connection is terminated."""
+    await cl.Message(f"MCP connection `{name}` has been disconnected.").send()
+    all_mcp_tools = cl.user_session.get("mcp_tools", {})
+    if name in all_mcp_tools:
+        del all_mcp_tools[name]
+        cl.user_session.set("mcp_tools", all_mcp_tools)
+
+
 @cl.on_chat_start
 async def start():
     log_info("Chat session started.")
-    await cl.Message(content="# Welcome!").send()
+    await cl.Message(content="# Welcome to Naked Insurance! How can I help?").send()
     cl.user_session.set("message_history", [])
+    cl.user_session.set("mcp_tools", {})
+
+
+async def execute_tool_call(tool_call: Any):
+    """Helper function to find the correct MCP session and execute a tool call."""
+    tool_name = tool_call.function.name
+    tool_args_str = tool_call.function.arguments
+
+    history = cl.user_session.get("message_history")
+
+    mcp_connection_name = None
+    all_mcp_tools = cl.user_session.get("mcp_tools", {})
+    for conn_name, tools in all_mcp_tools.items():
+        if any(t["function"]["name"] == tool_name for t in tools):
+            mcp_connection_name = conn_name
+            break
+
+    if not mcp_connection_name:
+        error_msg = f"Error: Tool `{tool_name}` not found in any active MCP connection."
+        history.append(
+            {
+                "tool_call_id": tool_call.id,
+                "role": "tool",
+                "name": tool_name,
+                "content": error_msg,
+            }
+        )
+        return
+
+    async with cl.Step(type="tool", name=tool_name, input=tool_args_str) as step:
+        try:
+            mcp_session, _ = cl.context.session.mcp_sessions.get(mcp_connection_name)
+            tool_args = json.loads(tool_args_str)
+            tool_result = await mcp_session.call_tool(tool_name, tool_args)
+            step.output = tool_result.content
+            history.append(
+                {
+                    "tool_call_id": tool_call.id,
+                    "role": "tool",
+                    "name": tool_name,
+                    "content": tool_result.content,
+                }
+            )
+        except Exception as e:
+            error_msg = f"Error executing tool `{tool_name}`: {e}"
+            step.error = error_msg
+            history.append(
+                {
+                    "tool_call_id": tool_call.id,
+                    "role": "tool",
+                    "name": tool_name,
+                    "content": error_msg,
+                }
+            )
+
+    cl.user_session.set("message_history", history)
 
 
 @cl.on_message
 async def main(message: cl.Message):
-    log_info("=" * 20)
-    log_info(f"on_message triggered for: '{message.content}'")
-
-    model = "gpt-4o"  # Hardcoding for simplicity during debug
-    temperature = 0.7
-    message_history = cl.user_session.get("message_history", [])
-    message_history.append({"role": "user", "content": message.content})
-
-    try:
-        # Create an empty message for streaming
-        msg = cl.Message(content="", author="Assistant")
-
-        log_info(f"Calling OpenAI model '{model}' with streaming...")
-        if not openai_client:
-            raise ValueError("OpenAI API key not configured")
-
-        # Stream the response from OpenAI
-        stream = await openai_client.chat.completions.create(
-            model=model, messages=message_history, temperature=temperature, stream=True
+    history = cl.user_session.get("message_history")
+    history.append({"role": "user", "content": message.content})
+
+    # Aggregate tools from all connections
+    all_mcp_tools = cl.user_session.get("mcp_tools", {})
+    aggregated_tools = [
+        tool for conn_tools in all_mcp_tools.values() for tool in conn_tools
+    ]
+
+    # First API call to get assistant response or tool calls
+    stream = await openai_client.chat.completions.create(
+        model="gpt-4o",
+        messages=history,
+        tools=aggregated_tools if aggregated_tools else None,
+        stream=True,
+    )
+
+    # Handle the streaming response
+    output_message = cl.Message(content="", author="Assistant")
+    tool_calls = []
+    tool_calls_buffer = {}  # Buffer to reconstruct tool calls that arrive in chunks
+
+    async for part in stream:
+        delta = part.choices[0].delta
+        if delta.content:
+            # Stream content tokens
+            await output_message.stream_token(delta.content)
+
+        if delta.tool_calls:
+            for tool_call_chunk in delta.tool_calls:
+                index = tool_call_chunk.index
+                if index not in tool_calls_buffer:
+                    # First time we see this tool call, initialize it
+                    tool_calls_buffer[index] = {
+                        "id": tool_call_chunk.id or "",
+                        "type": "function",
+                        "function": {
+                            "name": tool_call_chunk.function.name or "",
+                            "arguments": "",
+                        },
+                    }
+
+                # Append argument chunks
+                if tool_call_chunk.function.arguments:
+                    tool_calls_buffer[index]["function"][
+                        "arguments"
+                    ] += tool_call_chunk.function.arguments
+
+    # Finalize message if it has content
+    if output_message.content:
+        await output_message.update()
+
+    # If there were tool calls, process them
+    if tool_calls_buffer:
+        # Reconstruct the full tool calls from the buffer
+        assistant_message = {"role": "assistant", "content": None, "tool_calls": []}
+        for index in sorted(tool_calls_buffer.keys()):
+            assistant_message["tool_calls"].append(tool_calls_buffer[index])
+
+        history.append(assistant_message)  # Add assistant's decision to use tools
+
+        # Execute all tool calls
+        for tool_call_dict in assistant_message["tool_calls"]:
+            # The OpenAI library expects a Pydantic model, not a dict, so we create a mock one
+            from pydantic import BaseModel
+
+            class Func(BaseModel):
+                name: str
+                arguments: str
+
+            class ToolCall(BaseModel):
+                id: str
+                function: Func
+                type: str
+
+            tool_call_obj = ToolCall(**tool_call_dict)
+            await execute_tool_call(tool_call_obj)
+
+        # Second API call with tool results to get the final natural language response
+        final_stream = await openai_client.chat.completions.create(
+            model="gpt-4o", messages=history, stream=True
         )
-
-        log_info("Starting to stream response...")
-        async for part in stream:
+        final_output_message = cl.Message(content="", author="Assistant")
+        async for part in final_stream:
             if token := part.choices[0].delta.content or "":
-                await msg.stream_token(token)
+                await final_output_message.stream_token(token)
 
-        log_info(f"Streaming complete. Final content: '{msg.content}'")
+        await final_output_message.update()
+        history.append({"role": "assistant", "content": final_output_message.content})
 
-        # Update the message to finalize it
-        await msg.update()
+    else:  # No tool calls, just a simple response
+        history.append({"role": "assistant", "content": output_message.content})
 
-        # Update message history with the complete response
-        message_history.append({"role": "assistant", "content": msg.content})
-        cl.user_session.set("message_history", message_history)
-        log_info("on_message finished successfully with streaming.")
-
-    except Exception as e:
-        log_error(f"An exception occurred: {str(e)}")
-        await cl.Message(content=f"An error occurred: {str(e)}").send()
+    cl.user_session.set("message_history", history)

requirements.txt

@@ -3,4 +3,5 @@ openai
 google-generativeai
 python-dotenv
 requests
-websockets
+websockets
+mcp