mcp_chatbot.py

from dotenv import load_dotenv
from anthropic import Anthropic
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client
from typing import List, Dict, TypedDict
from contextlib import AsyncExitStack
import json
import asyncio

load_dotenv()

class ToolDefinition(TypedDict):
    name: str
    description: str
    input_schema: dict

class PromptDefinition(TypedDict):
    name: str
    description: str
    arguments: list

class MCP_ChatBot:

    def __init__(self):
        # Initialize session and client objects
        self.sessions: List[ClientSession] = [] # new
        self.exit_stack = AsyncExitStack() # new
        self.anthropic = Anthropic()
        self.available_tools: List[ToolDefinition] = []
        self.available_prompts: List[PromptDefinition] = []
        self.available_resources = [] # new
        self.tool_to_session: Dict[str, ClientSession] = {}
        self.prompt_to_session: Dict[str, ClientSession] = {}
        self.resource_to_session: Dict[str, ClientSession] = {} # new


    async def connect_to_server(self, server_name: str, server_config: dict) -> None:
        """Connect to a single MCP server."""
        try:
            server_params = StdioServerParameters(**server_config)
            stdio_transport = await self.exit_stack.enter_async_context(
                stdio_client(server_params)
            ) # new
            read, write = stdio_transport
            session = await self.exit_stack.enter_async_context(
                ClientSession(read, write)
            ) # new
            await session.initialize()
            self.sessions.append(session)

            
            try:
                # List available tools for this session
                response = await session.list_tools()
                tools = response.tools
                print(f"\nConnected to {server_name} with tools:", [t.name for t in tools])

                for tool in tools: # new
                    self.tool_to_session[tool.name] = session
                    self.available_tools.append({
                        "name": tool.name,
                        "description": tool.description,
                        "input_schema": tool.inputSchema
                    })

                # List available prompts
                prompt_response = await session.list_prompts()
                if prompt_response and prompt_response.prompts:
                    for prompt in prompt_response.prompts:
                        self.prompt_to_session[prompt.name] = session
                        self.available_prompts.append({
                            "name": prompt.name,
                            "description": prompt.description,
                            "arguments": prompt.arguments
                        })
                
                # List available resources
                resource_response = await session.list_resources()
                if resource_response and resource_response.resources:
                    for resource in resource_response.resources:
                        resource_uri = str(resource.uri)
                        self.resource_to_session[resource_uri] = session
                        self.available_resources.append(resource_uri)
            
            except Exception as e:
                print(f"Error listing tools, prompts or resources for {server_name}: {e}")


        except Exception as e:
            print(f"Failed to connect to {server_name}: {e}")

    async def connect_to_servers(self): # new
        """Connect to all configured MCP servers."""
        try:
            with open("server_config.json", "r") as file:
                data = json.load(file)

            servers = data.get("mcpServers", {})

            for server_name, server_config in servers.items():
                await self.connect_to_server(server_name, server_config)
        except Exception as e:
            print(f"Error loading server configuration: {e}")
            raise

    async def process_query(self, query):
        messages = [{'role':'user', 'content':query}]
        response = self.anthropic.messages.create(max_tokens = 2024,
                                                  system = """print("Type your queries or 'quit' to exit.")
        print("Use @folders to see available topics")
        print("Use @<topic> to search papers in that topic")
        print("Use /prompts to list available prompts")
        print("Use /prompt <name> <arg1 = value1> to execute a prompt")""",
                                      model = 'claude-3-7-sonnet-20250219', 
                                      tools = self.available_tools,
                                      messages = messages)
        process_query = True
        while process_query:
            assistant_content = []
            has_tool_use = False
            for content in response.content:
                if content.type =='text':
                    print(content.text)
                    assistant_content.append(content)
                    if(len(response.content) == 1):
                        process_query= False
                elif content.type == 'tool_use':
                    has_tool_use = True
                    assistant_content.append(content)
                    messages.append({'role':'assistant', 'content':assistant_content})
                    tool_id = content.id
                    tool_args = content.input
                    tool_name = content.name


                    print(f"Calling tool {tool_name} with args {tool_args}")

                    # Call a tool
                    session = self.tool_to_session[tool_name] # new
                    if not session:
                        print(f"No tool with name {tool_name} found.")
                        break
                    result = await session.call_tool(tool_name, arguments=tool_args)
                    messages.append({"role": "user", 
                                      "content": [
                                          {
                                              "type": "tool_result",
                                              "tool_use_id":tool_id,
                                              "content": result.content
                                          }
                                      ]
                                    })
                    response = self.anthropic.messages.create(max_tokens = 2024,
                                                              system = """print("Type your queries or 'quit' to exit.")
        print("Use @folders to see available topics")
        print("Use @<topic> to search papers in that topic")
        print("Use /prompts to list available prompts")
        print("Use /prompt <name> <arg1 = value1> to execute a prompt")""",
                                      model = 'claude-3-7-sonnet-20250219', 
                                      tools = self.available_tools,
                                      messages = messages) 

                    if(len(response.content) == 1 and response.content[0].type == "text"):
                        print(response.content[0].text)
                        process_query= False

    async def get_resources(self, resource_uri):
        session = self.resource_to_session.get(resource_uri)

        # Fall back for papers URIs - try any papers resource session
        if not session and resource_uri.startswith("papers://"):
            for uri, sess in self.resource_to_session.items():
                if uri.startswith("papers://"):
                    session = sess
                    break

        if not session:
            print(f"Resource {resource_uri} not found in any server")
            return
        
        try:
            result = await session.read_resource(uri=resource_uri)
            if result and result.contents:
                print(f"\nResource: {resource_uri}")
                print("Content:")
                print(result.contents[0].text)
            else:
                print(f"No content available")
        except Exception as e:
            print(f"Error: {e}")
    
    async def list_prompts(self):
        """List all available prompts"""
        if not self.available_prompts:
            print("No prompts available.")
            return
        
        print("\nAvailable Prompts:")
        for prompt in self.available_prompts:
            print(f"- {prompt['name']}: {prompt['description']}")
            if prompt['arguments']:
                for arg in prompt['arguments']:
                    arg_name = arg.name if hasattr(arg, 'name') else arg.get('name', '')
                    print(f"   - {arg_name}")
    
    async def execute_prompt(self, prompt_name, args):
        """Execute a specific prompt with given arguments"""
        session = self.prompt_to_session.get(prompt_name)
        if not session:
            print(f"Prompt {prompt_name} not found in any server")
            return
        
        try:
            result = await session.get_prompt(prompt_name, arguments = args)
            if result and result.messages:
                prompt_content = result.messages[0].content
                # Extract text from content (handles different formats)
                if isinstance(prompt_content, str):
                    text = prompt_content
                elif hasattr(prompt_content, 'text'):
                    text = prompt_content.text
                else:
                    # Handles list of content items
                    text = " ".join(item.text if hasattr(item, 'text') else str(item) for item in prompt_content)
                print(f"\nExecuting Prompt: {prompt_name}...")
                await self.process_query(text)
                
        except Exception as e:
            print(f"Error executing prompt {prompt_name}: {e}")

    async def chat_loop(self):
        # For adding the particular user interface for getting accesss to the prompts and resources in each servers. 
        """Run an interactive chat loop"""
        print("\nMCP Chatbot Started!")
        print("Type your queries or 'quit' to exit.")
        print("Use @folders to see available topics")
        print("Use @<topic> to search papers in that topic")
        print("Use /prompts to list available prompts")
        print("Use /prompt <name> <arg1 = value1> to execute a prompt")

        while True:
            try:
                query = input("\nQuery: ").strip()
                if not query:
                    continue
                if query.lower() == 'quit':
                    break
                
                # Check for @Resource syntax first
                if query.startswith('@'):
                    # Remove @ sign
                    topic = query[1:].strip()
                    if topic == "folders":
                        resource_uri = "papers://folders"
                    else:
                        resource_uri = f"papers://{topic}"
                    await self.get_resources(resource_uri)
                    continue

                # Check for /prompts command
                if query.startswith('/'):
                    parts = query.split()
                    command = parts[0].lower()

                    if command == '/prompts':
                        await self.list_prompts()
                    elif command == '/prompt':
                        if len(parts) < 2:
                            print("Usage: /prompt <name> <arg1=value1> <arg2=value2> ...")
                            continue
                        prompt_name = parts[1]
                        args = {}

                        #parse arguments
                        for arg in parts[2:]:
                            if '=' in arg:
                                key, value = arg.split('=', 1)
                                args[key] = value.strip()
                            else:
                                print(f"Invalid argument format: {arg}")
                                continue
                        await self.execute_prompt(prompt_name, args)
                    else:
                        print(f"Unknown command: {command}")
                    continue

                await self.process_query(query)

            except Exception as e:
                print(f"\nError: {str(e)}")

    async def cleanup(self): # new
        """Cleanly close all resources using AsyncExitStack."""
        await self.exit_stack.aclose()


async def main():
    chatbot = MCP_ChatBot()
    try:
        # the mcp clients and sessions are not initialized using "with"
        # like in the previous lesson
        # so the cleanup should be manually handled
        await chatbot.connect_to_servers() # new! 
        await chatbot.chat_loop()
    finally:
        await chatbot.cleanup() #new! 


if __name__ == "__main__":
    asyncio.run(main())