initial mcp server setup
This commit is contained in:
OwusuBlessing
@@ -0,0 +1,514 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP AI Client with Flexible Transport and LLM Provider Support
|
||||
|
||||
This client can connect to MCP servers using either SSE or stdio transport
|
||||
and use various AI models (OpenAI, Claude, Grok) to process queries with access to MCP tools.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import Any, Dict, List, Optional
|
||||
from enum import Enum
|
||||
|
||||
import nest_asyncio
|
||||
from dotenv import load_dotenv
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp.client.sse import sse_client
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Apply nest_asyncio to allow nested event loops
|
||||
nest_asyncio.apply()
|
||||
|
||||
# Import LLM client factory
|
||||
try:
|
||||
# Try to import from source first
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
|
||||
|
||||
from mcp_template.llm_client.client_factory import AIClientFactory
|
||||
LLM_CLIENT_AVAILABLE = True
|
||||
except ImportError:
|
||||
try:
|
||||
# Fallback to build version
|
||||
from mcp_template.llm_client.client_factory import AIClientFactory
|
||||
LLM_CLIENT_AVAILABLE = True
|
||||
except ImportError:
|
||||
LLM_CLIENT_AVAILABLE = False
|
||||
|
||||
|
||||
class TransportType(Enum):
|
||||
"""Supported MCP transport types"""
|
||||
STDIO = "stdio"
|
||||
SSE = "sse"
|
||||
|
||||
|
||||
class MCPAIClient:
|
||||
"""Client for interacting with AI models using MCP tools with flexible transport and provider support."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "gpt-4o",
|
||||
transport: TransportType = TransportType.SSE,
|
||||
provider: str = "openai",
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 1000,
|
||||
top_k: Optional[int] = None,
|
||||
top_p: Optional[float] = None,
|
||||
**kwargs
|
||||
):
|
||||
"""Initialize the AI MCP client.
|
||||
|
||||
Args:
|
||||
model: The AI model to use.
|
||||
transport: The MCP transport type to use (stdio or sse).
|
||||
provider: The AI provider to use (openai, claude, grok).
|
||||
temperature: Sampling temperature for AI responses.
|
||||
max_tokens: Maximum tokens for AI responses.
|
||||
**kwargs: Additional parameters for the AI client.
|
||||
"""
|
||||
self.transport = transport
|
||||
self.model = model
|
||||
self.provider = provider
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
|
||||
# Initialize session and client objects
|
||||
self.session: Optional[ClientSession] = None
|
||||
self.exit_stack = AsyncExitStack()
|
||||
|
||||
# Initialize AI client using factory
|
||||
if not LLM_CLIENT_AVAILABLE:
|
||||
raise ImportError("LLM client not available. Make sure the LLM client modules are properly installed.")
|
||||
|
||||
# Prepare additional parameters for AI client
|
||||
ai_kwargs = kwargs.copy()
|
||||
if top_k is not None:
|
||||
ai_kwargs["top_k"] = top_k
|
||||
if top_p is not None:
|
||||
ai_kwargs["top_p"] = top_p
|
||||
|
||||
self.ai_client = AIClientFactory.create_client(
|
||||
provider=provider,
|
||||
model_name=model,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
**ai_kwargs
|
||||
)
|
||||
|
||||
self.stdio: Optional[Any] = None
|
||||
self.write: Optional[Any] = None
|
||||
|
||||
# Transport-specific attributes
|
||||
self.read_stream: Optional[Any] = None
|
||||
self.write_stream: Optional[Any] = None
|
||||
|
||||
async def connect_stdio(self, server_script_path: str = "run_mcp_server.py", server_args: List[str] = None):
|
||||
"""Connect to an MCP server using stdio transport.
|
||||
|
||||
Args:
|
||||
server_script_path: Path to the server script.
|
||||
server_args: Additional arguments for the server script.
|
||||
"""
|
||||
if server_args is None:
|
||||
server_args = ["--transport", "stdio"]
|
||||
|
||||
# Server configuration
|
||||
server_params = StdioServerParameters(
|
||||
command="python",
|
||||
args=[server_script_path] + server_args,
|
||||
)
|
||||
|
||||
# Connect to the server
|
||||
stdio_transport = await self.exit_stack.enter_async_context(
|
||||
stdio_client(server_params)
|
||||
)
|
||||
self.stdio, self.write = stdio_transport
|
||||
self.session = await self.exit_stack.enter_async_context(
|
||||
ClientSession(self.stdio, self.write)
|
||||
)
|
||||
|
||||
async def connect_sse(self, server_url: str = "http://localhost:8050/sse"):
|
||||
"""Connect to an MCP server using SSE transport.
|
||||
|
||||
Args:
|
||||
server_url: The SSE endpoint URL of the server.
|
||||
"""
|
||||
# Connect to the server
|
||||
sse_transport = await self.exit_stack.enter_async_context(
|
||||
sse_client(server_url)
|
||||
)
|
||||
self.read_stream, self.write_stream = sse_transport
|
||||
self.session = await self.exit_stack.enter_async_context(
|
||||
ClientSession(self.read_stream, self.write_stream)
|
||||
)
|
||||
|
||||
async def connect_to_server(self, server_url: str = "http://localhost:8050/sse",
|
||||
server_script_path: str = "run_mcp_server.py"):
|
||||
"""Connect to an MCP server using the configured transport.
|
||||
|
||||
Args:
|
||||
server_url: The SSE endpoint URL (used for SSE transport).
|
||||
server_script_path: Path to the server script (used for stdio transport).
|
||||
"""
|
||||
if self.transport == TransportType.SSE:
|
||||
await self.connect_sse(server_url)
|
||||
elif self.transport == TransportType.STDIO:
|
||||
await self.connect_stdio(server_script_path)
|
||||
else:
|
||||
raise ValueError(f"Unsupported transport type: {self.transport}")
|
||||
|
||||
# Initialize the connection
|
||||
await self.session.initialize()
|
||||
|
||||
print(f"✅ Connected to MCP server using {self.transport.value.upper()} transport")
|
||||
|
||||
# List available components
|
||||
await self.list_available_components()
|
||||
|
||||
async def list_available_components(self):
|
||||
"""List all available tools, prompts, and resources."""
|
||||
print("\n" + "="*60)
|
||||
print("📋 AVAILABLE MCP COMPONENTS")
|
||||
print("="*60)
|
||||
|
||||
# List available tools
|
||||
tools_result = await self.session.list_tools()
|
||||
print(f"\n🔧 Tools ({len(tools_result.tools)}):")
|
||||
for tool in tools_result.tools:
|
||||
print(f" • {tool.name}: {tool.description}")
|
||||
|
||||
# List available prompts
|
||||
prompts_result = await self.session.list_prompts()
|
||||
print(f"\n💬 Prompts ({len(prompts_result.prompts)}):")
|
||||
for prompt in prompts_result.prompts:
|
||||
print(f" • {prompt.name}: {prompt.description}")
|
||||
|
||||
# List available resources
|
||||
resources_result = await self.session.list_resources()
|
||||
print(f"\n📄 Resources ({len(resources_result.resources)}):")
|
||||
for resource in resources_result.resources:
|
||||
print(f" • {resource.uri}: {resource.name}")
|
||||
if resource.description:
|
||||
print(f" └─ {resource.description}")
|
||||
|
||||
print("\n" + "="*60)
|
||||
|
||||
|
||||
async def get_mcp_tools(self) -> List[Dict[str, Any]]:
|
||||
"""Get available tools from the MCP server in AI provider format.
|
||||
|
||||
Returns:
|
||||
A list of tools formatted for the specific AI provider.
|
||||
"""
|
||||
tools_result = await self.session.list_tools()
|
||||
|
||||
# Convert MCP tools to standard format
|
||||
standard_tools = [
|
||||
{
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"inputSchema": tool.inputSchema,
|
||||
}
|
||||
for tool in tools_result.tools
|
||||
]
|
||||
|
||||
# Use AI client's formatting method
|
||||
return self.ai_client._format_tools_for_provider(standard_tools)
|
||||
|
||||
async def process_query(self, query: str) -> str:
|
||||
"""Process a query using AI model and available MCP tools.
|
||||
|
||||
Args:
|
||||
query: The user query.
|
||||
|
||||
Returns:
|
||||
The response from the AI model.
|
||||
"""
|
||||
print(f"\n🤔 Processing query: '{query}'")
|
||||
|
||||
# Get available tools
|
||||
tools = await self.get_mcp_tools()
|
||||
|
||||
# Initial AI API call
|
||||
print(f"🧠 Calling {self.provider}/{self.model} with {len(tools)} available tools...")
|
||||
response = await self.ai_client.chat_completion(
|
||||
messages=[{"role": "user", "content": query}],
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
# Get assistant's response
|
||||
assistant_message = response["choices"][0]["message"]
|
||||
|
||||
# Initialize conversation with user query and assistant response
|
||||
messages = [
|
||||
{"role": "user", "content": query},
|
||||
{
|
||||
"role": assistant_message["role"],
|
||||
"content": assistant_message.get("content"),
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc["id"],
|
||||
"type": tc.get("type", "function"),
|
||||
"function": tc["function"]
|
||||
}
|
||||
for tc in assistant_message.get("tool_calls", [])
|
||||
] if assistant_message.get("tool_calls") else None,
|
||||
},
|
||||
]
|
||||
|
||||
# Handle tool calls if present
|
||||
if "tool_calls" in assistant_message and assistant_message["tool_calls"]:
|
||||
print(f"🔧 Assistant wants to use {len(assistant_message['tool_calls'])} tool(s)")
|
||||
|
||||
# Process each tool call
|
||||
for i, tool_call in enumerate(assistant_message["tool_calls"], 1):
|
||||
tool_name = tool_call["function"]["name"]
|
||||
tool_args = json.loads(tool_call["function"]["arguments"])
|
||||
|
||||
print(f" {i}. Calling tool: {tool_name}")
|
||||
print(f" Arguments: {tool_args}")
|
||||
|
||||
# In proper MCP, arguments should be passed directly
|
||||
fastmcp_args = tool_args
|
||||
|
||||
# Execute tool call
|
||||
try:
|
||||
result = await self.session.call_tool(
|
||||
tool_name,
|
||||
arguments=fastmcp_args,
|
||||
)
|
||||
|
||||
tool_result = result.content[0].text if result.content else "No result"
|
||||
print(f" Result: {tool_result}{'...' if len(tool_result) > 100 else ''}")
|
||||
|
||||
# Add tool response to conversation
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call["id"],
|
||||
"content": tool_result,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"Tool execution failed: {e}"
|
||||
print(f" ❌ Error: {error_msg}")
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call["id"],
|
||||
"content": error_msg,
|
||||
}
|
||||
)
|
||||
|
||||
# Get final response from AI with tool results
|
||||
print(f"🧠 Getting final response from {self.provider}/{self.model}...")
|
||||
final_response = await self.ai_client.chat_completion(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="none", # Don't allow more tool calls
|
||||
)
|
||||
|
||||
final_answer = final_response["choices"][0]["message"]["content"]
|
||||
print(f"💡 Final answer: {final_answer[:200]}{'...' if len(final_answer) > 200 else ''}")
|
||||
return final_answer
|
||||
|
||||
# No tool calls, just return the direct response
|
||||
direct_answer = assistant_message["content"]
|
||||
print(f"💡 Direct answer: {direct_answer[:200]}{'...' if len(direct_answer) > 200 else ''}")
|
||||
return direct_answer
|
||||
|
||||
async def interactive_session(self):
|
||||
"""Start an interactive session for querying."""
|
||||
print("🚀 Starting interactive MCP-AI session")
|
||||
print(f"📡 Transport: {self.transport.value.upper()}")
|
||||
print(f"🤖 Model: {self.provider}/{self.model}")
|
||||
print(f"🌡️ Temperature: {self.temperature}")
|
||||
print(f"📏 Max Tokens: {self.max_tokens}")
|
||||
print("💡 Type 'quit' or 'exit' to end the session")
|
||||
print("-" * 50)
|
||||
|
||||
while True:
|
||||
try:
|
||||
query = input("\n❓ Your query: ").strip()
|
||||
|
||||
if query.lower() in ['quit', 'exit', 'q']:
|
||||
print("👋 Goodbye!")
|
||||
break
|
||||
|
||||
if not query:
|
||||
continue
|
||||
|
||||
# Process the query
|
||||
response = await self.process_query(query)
|
||||
print(f"\n🎯 Response: {response}")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n👋 Session interrupted. Goodbye!")
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"❌ Error processing query: {e}")
|
||||
continue
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up resources."""
|
||||
try:
|
||||
await self.exit_stack.aclose()
|
||||
print("🧹 Resources cleaned up successfully")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Cleanup warning: {e}")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Main entry point for the client."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="MCP AI Client with Flexible Transport and LLM Provider Support")
|
||||
parser.add_argument(
|
||||
"--transport",
|
||||
choices=["sse", "stdio"],
|
||||
default="sse",
|
||||
help="MCP transport type (default: sse)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--provider",
|
||||
choices=["openai", "claude", "grok"],
|
||||
default="openai",
|
||||
help="AI provider to use (default: openai)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
help="AI model to use (defaults based on provider: openai=gpt-4o, claude=claude-3-opus-20240229, grok=grok-1)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature",
|
||||
type=float,
|
||||
default=0.7,
|
||||
help="Sampling temperature (default: 0.7)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Maximum tokens for response (default: 1000)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-k",
|
||||
type=int,
|
||||
help="Top-k sampling parameter (provider-specific)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--top-p",
|
||||
type=float,
|
||||
help="Top-p sampling parameter (provider-specific)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server-url",
|
||||
default="http://localhost:8050/sse",
|
||||
help="Server URL for SSE transport (default: http://localhost:8050/sse)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server-script",
|
||||
default="run_mcp_server.py",
|
||||
help="Server script path for stdio transport (default: run_mcp_server.py)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query",
|
||||
help="Single query to process (if not provided, starts interactive mode)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose", "-v",
|
||||
action="store_true",
|
||||
help="Enable verbose output"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Validate API keys based on provider
|
||||
provider = args.provider
|
||||
api_key_env_vars = {
|
||||
"openai": "OPENAI_API_KEY",
|
||||
"claude": "ANTHROPIC_API_KEY",
|
||||
"grok": "GROK_API_KEY"
|
||||
}
|
||||
|
||||
api_key_env = api_key_env_vars.get(provider)
|
||||
if not api_key_env:
|
||||
print(f"❌ Error: Unknown provider: {provider}")
|
||||
return
|
||||
|
||||
api_key = os.getenv(api_key_env)
|
||||
if not api_key:
|
||||
print(f"❌ Error: {api_key_env} environment variable not set")
|
||||
print(f" Please set your {provider.upper()} API key:")
|
||||
print(f" export {api_key_env}='your-api-key-here'")
|
||||
print(f" Or create a .env file with: {api_key_env}=your-api-key-here")
|
||||
return
|
||||
|
||||
# Basic API key validation
|
||||
if provider == "openai" and (not api_key.startswith("sk-") or len(api_key) < 20):
|
||||
print("❌ Error: OPENAI_API_KEY appears to be invalid")
|
||||
print(" API key should start with 'sk-' and be at least 20 characters long")
|
||||
return
|
||||
|
||||
# Set default model if not provided
|
||||
model = args.model
|
||||
if not model:
|
||||
default_models = {
|
||||
"openai": "gpt-4o",
|
||||
"claude": "claude-3-opus-20240229",
|
||||
"grok": "grok-1"
|
||||
}
|
||||
model = default_models.get(provider, "gpt-4o")
|
||||
|
||||
# Create client
|
||||
transport = TransportType(args.transport)
|
||||
client = MCPAIClient(
|
||||
model=model,
|
||||
transport=transport,
|
||||
provider=provider,
|
||||
temperature=args.temperature,
|
||||
max_tokens=args.max_tokens,
|
||||
top_k=args.top_k,
|
||||
top_p=args.top_p
|
||||
)
|
||||
|
||||
try:
|
||||
# Connect to server
|
||||
if transport == TransportType.SSE:
|
||||
await client.connect_to_server(server_url=args.server_url)
|
||||
else:
|
||||
await client.connect_to_server(server_script_path=args.server_script)
|
||||
|
||||
# Handle single query or interactive mode
|
||||
if args.query:
|
||||
response = await client.process_query(args.query)
|
||||
print(f"\n🎯 Response: {response}")
|
||||
else:
|
||||
await client.interactive_session()
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
if "ConnectionRefusedError" in str(e):
|
||||
if transport == TransportType.SSE:
|
||||
print("💡 Make sure the MCP server is running with SSE transport:")
|
||||
print(f" python {args.server_script} --transport sse")
|
||||
else:
|
||||
print("💡 Make sure the server script path is correct")
|
||||
finally:
|
||||
await client.cleanup()
|
||||
|
||||
|
||||
# Backward compatibility alias
|
||||
MCPOpenAIClient = MCPAIClient
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user