"""
AgentServer class for managing an agent in both client and server modes.
"""
import asyncio
import json
import logging
from typing import Any, Dict, Optional
from src.agent.agent import Agent
from src.capabilities.capability import CapabilityRegistry
from src.llm.base_llm import BaseLLM
from src.llm.groq_llm import GroqLLM
from src.mcp.mcp_connection_manager import MCPConnectionManager
from src.mcp.mcp_server_wrapper import MCPServerWrapper
from src.tools.tool import ToolRegistry
from src.utils.context import cleanup_context, get_context, initialize_context
class AgentServer:
"""
Manages an agent that can operate both as a client and as a server.
"""
def __init__(
self, config: Dict[str, Any], server_mode: bool = False, server_name: str = None
):
self.config = config
self.server_mode = server_mode
self.server_name = server_name or config.get(
"server_name", f"agent-server-{id(self)}"
)
self.agent = None
self.server = None
self.connection_manager = None
self.tool_registry = None
self.capability_registry = None
self.llm_client = None
self.logger = logging.getLogger("agent-server")
async def initialize(self):
"""Initialize the agent and server if in server mode."""
# Initialize context
await initialize_context(self.config)
context = get_context()
# Create connection manager
self.connection_manager = MCPConnectionManager()
# Create registries
self.tool_registry = ToolRegistry()
self.capability_registry = CapabilityRegistry()
# Create LLM client
self.llm_client = self._create_llm_client()
# Create agent
self.agent = Agent(
llm_client=self.llm_client,
connection_manager=self.connection_manager,
tool_registry=self.tool_registry,
capability_registry=self.capability_registry,
name=self.config.get("agent_name", "agent"),
)
# Connect to servers and discover tools
await self._connect_to_servers_and_discover_tools()
# Load capabilities from config
await self._load_capabilities()
# If in server mode, initialize the server wrapper
if self.server_mode:
self.server = MCPServerWrapper(self.agent, self.server_name)
self.logger.info(
f"Initialized agent in server mode with name: {self.server_name}"
)
else:
self.logger.info(f"Initialized agent in client-only mode")
return self
async def run(self):
"""Run the agent in the appropriate mode."""
if self.server_mode:
# Run as a server
self.logger.info(f"Running agent as server with name: {self.server_name}")
await self.server.run_stdio_async()
else:
# Run as a standalone agent
self.logger.info("Running agent in interactive conversation mode")
await self.agent.start_conversation()
async def cleanup(self):
"""Clean up agent resources."""
try:
self.logger.info("Cleaning up agent resources")
if self.connection_manager:
await self.connection_manager.disconnect_all()
await cleanup_context()
self.logger.info("Cleanup complete")
except Exception as e:
self.logger.error(f"Error during cleanup: {e}")
async def _connect_to_servers_and_discover_tools(self):
"""Connect to servers and discover their tools."""
try:
# First establish connections to all servers
servers = self.config.get("servers", {})
for name, server_config in servers.items():
self.logger.info(f"Connecting to server: {name}")
await self.connection_manager.connect_server(name, server_config)
# Then discover tools from all servers
self.logger.info("Discovering tools from servers")
await self.tool_registry.load_from_config(
self.config, self.connection_manager
)
# Log the tools that were discovered
tools = self.tool_registry.list_tools()
self.logger.info(f"Discovered {len(tools)} tools from MCP servers")
except Exception as e:
self.logger.error(f"Error connecting to servers or discovering tools: {e}")
raise
async def _load_capabilities(self):
"""Load capabilities from the configuration."""
try:
self.logger.info("Loading capabilities from configuration")
await self.capability_registry.load_from_config(self.config)
# Log the capabilities that were loaded
capabilities = self.capability_registry.list_capabilities()
self.logger.info(
f"Loaded {len(capabilities)} capabilities from configuration"
)
except Exception as e:
self.logger.error(f"Error loading capabilities: {e}")
raise
def _create_llm_client(self) -> BaseLLM:
"""Create an LLM client based on configuration."""
provider = self.config.get("llm_provider", "groq")
api_key = self.config.get("llm_api_key", "")
self.logger.info(f"Creating LLM client with provider: {provider}")
if provider == "groq":
return GroqLLM(api_key)
# TODO: add more providers (Anthropic, OpenAI...)
else:
raise ValueError(f"Unsupported LLM provider: {provider}")
