Model Context Protocol Demo

import asyncio import json import logging import os from typing import Dict, List, Any, Optional from fastmcp import Client # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # LLM Integration Options - Choose one: # Option 1: OpenAI API # Option 2: Anthropic Claude API # Option 3: Local LLM via Ollama # Option 4: Google Gemini API # Option 5: AWS Bedrock Claude 3.5 LLM_MODELS = ["openai", "anthropic", "ollama", "gemini", "bedrock"] LLM_PROVIDER = "gemini" # --- LLM Query Parser --- class LLMQueryParser: """Use LLM to parse natural language queries into tool calls""" def __init__(self, provider: str = "ollama"): self.provider = provider self.setup_llm_client() def setup_llm_client(self): """Setup the appropriate LLM client based on provider""" if self.provider == "openai": try: import openai self.client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")) logging.info("✅ OpenAI client initialized") except ImportError: logging.error("❌ OpenAI library not installed. Run: pip install openai") self.client = None elif self.provider == "anthropic": try: import anthropic self.client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY")) logging.info("✅ Anthropic client initialized") except ImportError: logging.error("❌ Anthropic library not installed. Run: pip install anthropic") self.client = None elif self.provider == "ollama": try: import requests # Test if Ollama is running response = requests.get("http://localhost:11434/api/tags", timeout=5) if response.status_code == 200: logging.info("✅ Ollama server detected") self.client = "ollama" # Use string indicator else: logging.error("❌ Ollama server not responding") self.client = None except Exception as e: logging.error(f"❌ Ollama connection failed: {e}") logging.error("Make sure Ollama is installed and running: https://ollama.ai") self.client = None elif self.provider == "gemini": try: import google.generativeai as genai api_key = os.getenv("GEMINI_API_KEY") if api_key: genai.configure(api_key=api_key) # https://ai.google.dev/gemini-api/docs/models gemini_models = [ "gemini-2.5-flash-preview-05-20", "gemini-2.5-pro-preview-05-06", "gemini-2.0-flash", "gemini-1.5-flash", "gemini-1.5-pro", ] model_name = gemini_models[3] self.client = genai.GenerativeModel(model_name) logging.info("✅ Gemini client initialized") else: logging.error("❌ GEMINI_API_KEY not set") self.client = None except ImportError: logging.error("❌ Gemini library not installed. Run: pip install google-generativeai") self.client = None elif self.provider == "bedrock": try: import boto3 from botocore.exceptions import NoCredentialsError, PartialCredentialsError # Initialize Bedrock client # AWS credentials should be configured via: # 1. AWS CLI: aws configure # 2. Environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION # 3. IAM roles (if running on EC2) # 4. AWS SSO region = os.getenv("AWS_DEFAULT_REGION", "us-east-1") try: self.client = boto3.client( service_name='bedrock-runtime', region_name=region ) # Test connection with a simple call # Note: We'll test this during the first actual query to avoid unnecessary calls self.bedrock_model_id = "anthropic.claude-3-5-sonnet-20241022-v2:0" # Claude 3.5 Sonnet logging.info(f"✅ AWS Bedrock client initialized (region: {region})") logging.info(f"📋 Using model: {self.bedrock_model_id}") except (NoCredentialsError, PartialCredentialsError) as e: logging.error(f"❌ AWS credentials not configured: {e}") logging.error("Configure AWS credentials using one of these methods:") logging.error(" 1. AWS CLI: aws configure") logging.error(" 2. Environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY") logging.error(" 3. IAM roles (if running on AWS)") self.client = None except ImportError: logging.error("❌ Boto3 library not installed. Run: pip install boto3") self.client = None def get_system_prompt(self, available_tools: List[Dict]) -> str: """Create system prompt with available tools""" tools_desc = "\n".join([ f"- {tool['name']}: {tool.get('description', 'No description')}" for tool in available_tools ]) return f"""You are a tool selection assistant. Given a user query, you must decide which tool to use and extract the necessary parameters. Available tools: {tools_desc} For each user query, respond with ONLY a JSON object in this exact format: {{ "tool": "tool_name", "params": {{ "param1": "value1", "param2": "value2" }}, "confidence": 0.95, "reasoning": "Brief explanation of why this tool was chosen" }} Tool-specific parameter requirements: - calculator: operation (add/subtract/multiply/divide), num1 (number), num2 (number) - trig: operation (sine/cosine/tangent/arc sine/arc cosine/arc tangent), num1 (float), unit(str) - stock_quote: ticker (stock symbol like AAPL, MSFT) - health: no parameters needed - echo: message (text to echo back) If you cannot determine which tool to use or extract parameters, respond with: {{ "tool": null, "params": {{}}, "confidence": 0.0, "reasoning": "Could not parse the query" }} Examples: User: "What is 15 plus 27?" Response: {{"tool": "calculator", "params": {{"operation": "add", "num1": 15, "num2": 27}}, "confidence": 0.98, "reasoning": "Clear arithmetic addition request"}} User: "What is sine 30 degrees?" Response: {{"tool": "trig", "params": {{"operation": "sine", "num1": 30, "unit": "degree"}}, "confidence": 0.98, "reasoning": "Clear trigonometic sine function request"}} User: "Get Apple stock price" Response: {{"tool": "stock_quote", "params": {{"ticker": "AAPL"}}, "confidence": 0.95, "reasoning": "Request for Apple (AAPL) stock information"}} User: "echo hello world" Response: {{"tool": "echo", "params": {{"message": "hello world"}}, "confidence": 0.99, "reasoning": "Direct echo command"}} User: "server status" Response: {{"tool": "health", "params": {{}}, "confidence": 0.90, "reasoning": "Request for server health information"}} Remember: Respond with ONLY the JSON object, no additional text.""" async def parse_query_with_llm(self, query: str, available_tools: List[Dict]) -> Optional[Dict[str, Any]]: """Use LLM to parse the query""" if not self.client: logging.error("❌ LLM client not available, falling back to rule-based parsing") return None system_prompt = self.get_system_prompt(available_tools) try: if self.provider == "openai": response = self.client.chat.completions.create( model="gpt-3.5-turbo", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": query} ], temperature=0.1, max_tokens=200 ) llm_response = response.choices[0].message.content.strip() elif self.provider == "anthropic": response = self.client.messages.create( model="claude-3-haiku-20240307", max_tokens=200, temperature=0.1, system=system_prompt, messages=[{"role": "user", "content": query}] ) llm_response = response.content[0].text.strip() elif self.provider == "ollama": import requests payload = { "model": "llama3.2", # or "mistral", "codellama", etc. "prompt": f"{system_prompt}\n\nUser: {query}\nResponse:", "stream": False, "options": { "temperature": 0.1, "num_predict": 200 } } response = requests.post("http://localhost:11434/api/generate", json=payload, timeout=30) if response.status_code == 200: llm_response = response.json()["response"].strip() else: logging.error(f"Ollama API error: {response.status_code}") return None elif self.provider == "gemini": response = self.client.generate_content( f"{system_prompt}\n\nUser: {query}", generation_config={ "temperature": 0.1, "max_output_tokens": 200 } ) llm_response = response.text.strip() elif self.provider == "bedrock": import json as json_lib # Construct the request for Claude 3.5 Sonnet request_body = { "anthropic_version": "bedrock-2023-05-31", "max_tokens": 200, "temperature": 0.1, "system": system_prompt, "messages": [ { "role": "user", "content": query } ] } # Make the request to Bedrock response = self.client.invoke_model( modelId=self.bedrock_model_id, body=json_lib.dumps(request_body), contentType="application/json", accept="application/json" ) # Parse the response response_body = json_lib.loads(response['body'].read()) llm_response = response_body['content'][0]['text'].strip() logging.debug(f"Bedrock response usage: {response_body.get('usage', {})}") # Parse the JSON response logging.debug(f"LLM Response: {llm_response}") # Clean up the response - sometimes LLMs add extra text if llm_response.startswith("```json"): llm_response = llm_response.replace("```json", "").replace("```", "").strip() parsed_response = json.loads(llm_response) # Validate the response structure if not isinstance(parsed_response, dict): logging.error("LLM response is not a dictionary") return None if parsed_response.get("tool") is None: logging.warning(f"LLM could not parse query: {parsed_response.get('reasoning', 'Unknown reason')}") return None confidence = parsed_response.get("confidence", 0.0) if confidence < 0.5: logging.warning(f"Low confidence ({confidence}) in LLM parsing: {parsed_response.get('reasoning')}") return None logging.info(f"🤖 LLM parsed query: {parsed_response.get('reasoning')} (confidence: {confidence})") return { "tool": parsed_response["tool"], "params": parsed_response.get("params", {}), "confidence": confidence, "reasoning": parsed_response.get("reasoning", "") } except json.JSONDecodeError as e: logging.error(f"Failed to parse LLM JSON response: {e}") logging.error(f"Raw LLM response: {llm_response}") return None except Exception as e: logging.error(f"LLM query parsing error: {e}") return None # --- Fallback Rule-Based Parser (same as before) --- class RuleBasedQueryParser: """Fallback rule-based query parser""" @staticmethod def parse_calculator_query(query: str) -> Optional[Dict[str, Any]]: import re query = query.lower().strip() patterns = [ ("add", ["plus", "add", "+", "sum"]), ("subtract", ["minus", "subtract", "-", "difference"]), ("multiply", ["times", "multiply", "*", "×", "product"]), ("divide", ["divide", "divided by", "/", "÷"]), ] for operation, keywords in patterns: for keyword in keywords: if keyword in query: numbers = re.findall(r'-?\d+(?:\.\d+)?', query) if len(numbers) >= 2: try: return { "tool": "calculator", "params": { "operation": operation, "num1": float(numbers[0]), "num2": float(numbers[1]) } } except ValueError: pass return None @staticmethod def parse_stock_query(query: str) -> Optional[Dict[str, Any]]: import re query_lower = query.lower().strip() stock_keywords = ["stock", "price", "quote", "ticker", "share"] if any(keyword in query_lower for keyword in stock_keywords): tickers = re.findall(r'\b[A-Z]{2,5}\b', query.upper()) if tickers: excluded_words = {"GET", "THE", "FOR", "AND", "BUT", "NOT", "YOU", "ALL", "CAN"} valid_tickers = [t for t in tickers if t not in excluded_words] if valid_tickers: return { "tool": "stock_quote", "params": {"ticker": valid_tickers[0]} } common_tickers = ["AAPL", "GOOGL", "GOOG", "MSFT", "TSLA", "AMZN", "META", "NVDA"] for ticker in common_tickers: if ticker in query.upper(): return { "tool": "stock_quote", "params": {"ticker": ticker} } return None @staticmethod def parse_query(query: str) -> Optional[Dict[str, Any]]: if any(word in query.lower() for word in ["health", "status", "ping"]): return {"tool": "health", "params": {}} if query.lower().startswith("echo "): message = query[5:].strip() return {"tool": "echo", "params": {"message": message}} calc_result = RuleBasedQueryParser.parse_calculator_query(query) if calc_result: return calc_result stock_result = RuleBasedQueryParser.parse_stock_query(query) if stock_result: return stock_result return None # --- Same extraction and formatting functions as before --- def extract_result_data(result): """Extract actual data from FastMCP result object""" try: if isinstance(result, list) and len(result) > 0: content_item = result[0] if hasattr(content_item, 'text'): try: return json.loads(content_item.text) except json.JSONDecodeError: return {"text": content_item.text} else: return {"content": str(content_item)} elif hasattr(result, 'content') and result.content: content_item = result.content[0] if hasattr(content_item, 'text'): try: return json.loads(content_item.text) except json.JSONDecodeError: return {"text": content_item.text} else: return {"content": str(content_item)} else: if isinstance(result, dict): return result else: return {"result": str(result)} except Exception as e: logging.error(f"Error extracting result data: {e}") return {"error": f"Could not parse result: {e}"} def format_result(tool_name: str, result: Dict) -> str: """Format tool results for display""" if isinstance(result, dict) and "error" in result: return f"❌ Error: {result['error']}" if tool_name == "calculator": if "result" in result: expression = result.get('expression') if expression: return f"🧮 {expression}" else: num1 = result.get('num1', '?') num2 = result.get('num2', '?') operation = result.get('operation', '?') calc_result = result.get('result', '?') return f"🧮 {num1} {operation} {num2} = {calc_result}" elif "error" in result: return f"❌ Calculator Error: {result['error']}" elif tool_name == "stock_quote": if "current_price" in result: ticker = result.get('ticker', 'Unknown') name = result.get('company_name', ticker) price = result.get('current_price', 0) currency = result.get('currency', 'USD') extra_info = [] if result.get('volume'): extra_info.append(f"Vol: {result['volume']:,}") if result.get('day_high') and result.get('day_low'): extra_info.append(f"Range: {result['day_low']}-{result['day_high']}") extra = f" ({', '.join(extra_info)})" if extra_info else "" return f"📈 {name} ({ticker}): {currency} {price}{extra}" elif "error" in result: return f"❌ Stock Error: {result['error']}" elif tool_name == "health": if isinstance(result, dict): return f"✅ {result.get('message', 'Server is healthy')}" elif tool_name == "echo": if isinstance(result, dict): return f"🔊 {result.get('echo', result.get('message', str(result)))}" try: return f"✅ Result: {json.dumps(result, indent=2)}" except (TypeError, ValueError): return f"✅ Result: {str(result)}" # --- Main Demo with LLM Integration --- async def run_llm_demo(): print("🚀 LLM-powered MCP client demo starting...") print(f"🤖 Using LLM Provider: {LLM_PROVIDER}") # Initialize parsers llm_parser = LLMQueryParser(LLM_PROVIDER) fallback_parser = RuleBasedQueryParser() server_path = "mcp_server.py" try: print(f"📡 Connecting to MCP server: {server_path}") async with Client(server_path) as client: print("✅ Connected to MCP server (local)!") # Discover available tools for LLM context available_tools = [] try: tools = await client.list_tools() # print(f"tools: {tools}") if tools: available_tools = [ {"name": tool.name, "description": tool.description} for tool in tools ] print(f"✅ Found {len(available_tools)} tools for LLM context") else: # Fallback tool definitions available_tools = [ {"name": "calculator", "description": "Perform arithmetic operations"}, {"name": "trig", "description": "Performs trigonometric operations"}, {"name": "stock_quote", "description": "Get stock price data"}, {"name": "health", "description": "Check server health"}, {"name": "echo", "description": "Echo back messages"} ] except Exception as e: logging.error(f"Tool discovery failed: {e}") print(f"\n{'='*60}") print("🎯 LLM-Powered Interactive Demo Started!") print("\n📝 Try natural language queries:") print(" • 'What's fifteen plus twenty seven?'") print(" • 'Can you multiply 12 by 8?'") print(" • 'Find sine of 30 degrees'") print(" • 'I need the current Apple stock price'") print(" • 'How much is Tesla trading for?'") print(" • 'Please echo back: Hello AI!'") print(" • 'Is the server working properly?'") print("\n💡 Commands:") print(" • 'tools' - List available tools") print(" • 'switch' - Switch parsing mode") print(" • 'exit' - Quit the demo") print(f"{'='*60}") use_llm = llm_parser.client is not None while True: try: user_input = input(f"\n💬 Your query {'🤖' if use_llm else '🔧'}: ").strip() if not user_input: continue if user_input.lower() in ('exit', "quit", "bye", "q"): print("👋 Goodbye!") break if user_input.lower() == 'switch': use_llm = not use_llm mode = "LLM" if use_llm else "Rule-based" print(f"🔄 Switched to {mode} parsing mode") continue if user_input.lower() == 'tools': print("\n🔧 Available tools:") for tool in available_tools: print(f" • {tool['name']}: {tool['description']}") continue # Parse the query parsed_query = None if use_llm and llm_parser.client: parsed_query = await llm_parser.parse_query_with_llm(user_input, available_tools) # Fallback to rule-based if LLM fails if not parsed_query: if use_llm: print("🔄 LLM parsing failed, trying rule-based parser...") parsed_query = fallback_parser.parse_query(user_input) if not parsed_query: print("❓ I couldn't understand your query. Try rephrasing or being more specific.") continue # Execute the tool call tool_name = parsed_query["tool"] parameters = parsed_query["params"] print(f"🔧 Calling tool: {tool_name}") if parsed_query.get("reasoning"): print(f"💭 Reasoning: {parsed_query['reasoning']}") if parameters: print(f"📝 Parameters: {json.dumps(parameters, indent=2)}") try: result = await client.call_tool(tool_name, parameters) result_data = extract_result_data(result) print(format_result(tool_name, result_data)) except Exception as e: print(f"❌ Error calling tool: {e}") logging.error(f"Tool call error: {e}", exc_info=True) except KeyboardInterrupt: print("\n\n👋 Goodbye!") break except Exception as e: logging.error(f"Unexpected error: {e}", exc_info=True) print(f"❌ Unexpected error: {e}") except Exception as e: print(f"❌ Failed to connect to server: {e}") print("\nMake sure the server file exists and FastMCP is installed:") print(" pip install fastmcp yfinance") print(f" Ensure {server_path} exists in the current directory") def main(): """Run the LLM-powered async demo""" print("🤖 LLM-Powered MCP Client") print("=" * 40) print(f"Selected LLM Provider: {LLM_PROVIDER}") if LLM_PROVIDER == "ollama": print("📋 Ollama Setup Instructions:") print("1. Install Ollama: https://ollama.ai") print("2. Run: ollama pull llama3.2") print("3. Start Ollama service") elif LLM_PROVIDER == "openai": print("📋 Set environment variable: OPENAI_API_KEY") elif LLM_PROVIDER == "anthropic": print("📋 Set environment variable: ANTHROPIC_API_KEY") elif LLM_PROVIDER == "gemini": print("📋 Set environment variable: GEMINI_API_KEY") elif LLM_PROVIDER == "bedrock": print("📋 AWS Bedrock Setup Instructions:") print("1. Install boto3: pip install boto3") print("2. Configure AWS credentials (choose one):") print(" • AWS CLI: aws configure") print(" • Environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION") print(" • IAM roles (if running on AWS)") print(" • AWS SSO: aws sso login") print("3. Ensure you have access to Claude 3.5 Sonnet in Bedrock") print("4. Set AWS_DEFAULT_REGION if needed (default: us-east-1)") print("=" * 40) asyncio.run(run_llm_demo()) if __name__ == '__main__': main()