109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
"""AgentPool - 运行时 Agent 实例池"""
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from agentkit.core.config_driven import ConfigDrivenAgent
|
|
from agentkit.core.protocol import AgentStatus
|
|
from agentkit.llm.gateway import LLMGateway
|
|
from agentkit.skills.registry import SkillRegistry
|
|
from agentkit.tools.registry import ToolRegistry
|
|
|
|
if TYPE_CHECKING:
|
|
from agentkit.core.compressor import CompressionStrategy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AgentPool:
|
|
"""运行时 Agent 实例池，管理 Agent 的创建、获取、删除"""
|
|
|
|
def __init__(
|
|
self,
|
|
llm_gateway: LLMGateway,
|
|
skill_registry: SkillRegistry,
|
|
tool_registry: ToolRegistry | None = None,
|
|
compressor: "CompressionStrategy | None" = None,
|
|
message_bus: Any = None,
|
|
):
|
|
self._agents: dict[str, ConfigDrivenAgent] = {}
|
|
self._llm_gateway = llm_gateway
|
|
self._skill_registry = skill_registry
|
|
self._tool_registry = tool_registry or ToolRegistry()
|
|
self._compressor = compressor
|
|
self._message_bus = message_bus
|
|
|
|
async def create_agent(self, config) -> ConfigDrivenAgent:
|
|
"""Create and start an Agent instance
|
|
|
|
Args:
|
|
config: AgentConfig or SkillConfig instance
|
|
|
|
Returns:
|
|
The created ConfigDrivenAgent
|
|
"""
|
|
# If agent with same name exists, stop it first
|
|
if config.name in self._agents:
|
|
await self.remove_agent(config.name)
|
|
|
|
agent = ConfigDrivenAgent(
|
|
config=config,
|
|
tool_registry=self._tool_registry,
|
|
llm_gateway=self._llm_gateway,
|
|
compressor=self._compressor,
|
|
)
|
|
await agent.start()
|
|
self._agents[config.name] = agent
|
|
logger.info(f"Agent '{config.name}' created and started in pool")
|
|
|
|
# Register agent to MessageBus if available
|
|
if self._message_bus is not None:
|
|
try:
|
|
async def _handle_bus_message(msg):
|
|
"""Handle incoming bus messages for this agent."""
|
|
logger.debug(f"Agent '{config.name}' received bus message: {msg.topic}")
|
|
|
|
await self._message_bus.subscribe(config.name, _handle_bus_message)
|
|
logger.info(f"Agent '{config.name}' registered to MessageBus")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to register agent '{config.name}' to MessageBus: {e}")
|
|
|
|
return agent
|
|
|
|
async def remove_agent(self, name: str) -> None:
|
|
"""Stop and remove an Agent"""
|
|
agent = self._agents.pop(name, None)
|
|
if agent:
|
|
await agent.stop()
|
|
|
|
# Unregister from MessageBus if available
|
|
if self._message_bus is not None:
|
|
try:
|
|
await self._message_bus.unsubscribe(name)
|
|
logger.info(f"Agent '{name}' unregistered from MessageBus")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to unregister agent '{name}' from MessageBus: {e}")
|
|
|
|
logger.info(f"Agent '{name}' stopped and removed from pool")
|
|
|
|
def get_agent(self, name: str) -> ConfigDrivenAgent | None:
|
|
"""Get agent by name"""
|
|
return self._agents.get(name)
|
|
|
|
def list_agents(self) -> list[dict]:
|
|
"""List all agents with info"""
|
|
return [
|
|
{
|
|
"name": agent.name,
|
|
"agent_type": agent.agent_type,
|
|
"version": agent.version,
|
|
"state": agent.status.value,
|
|
}
|
|
for agent in self._agents.values()
|
|
]
|
|
|
|
async def create_agent_from_skill(self, skill_name: str) -> ConfigDrivenAgent:
|
|
"""Create agent from a registered skill"""
|
|
skill = self._skill_registry.get(skill_name)
|
|
return await self.create_agent(skill.config)
|