diff --git a/src/agentkit/core/agent_pool.py b/src/agentkit/core/agent_pool.py index 141cae4..200ac77 100644 --- a/src/agentkit/core/agent_pool.py +++ b/src/agentkit/core/agent_pool.py @@ -1,6 +1,7 @@ """AgentPool - 运行时 Agent 实例池""" import logging +from typing import TYPE_CHECKING from agentkit.core.config_driven import ConfigDrivenAgent from agentkit.core.protocol import AgentStatus @@ -8,6 +9,9 @@ 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__) @@ -19,11 +23,13 @@ class AgentPool: llm_gateway: LLMGateway, skill_registry: SkillRegistry, tool_registry: ToolRegistry | None = None, + compressor: "CompressionStrategy | None" = 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 async def create_agent(self, config) -> ConfigDrivenAgent: """Create and start an Agent instance @@ -42,6 +48,7 @@ class AgentPool: config=config, tool_registry=self._tool_registry, llm_gateway=self._llm_gateway, + compressor=self._compressor, ) await agent.start() self._agents[config.name] = agent diff --git a/src/agentkit/core/config_driven.py b/src/agentkit/core/config_driven.py index e723b8c..df75233 100644 --- a/src/agentkit/core/config_driven.py +++ b/src/agentkit/core/config_driven.py @@ -201,6 +201,7 @@ class ConfigDrivenAgent(BaseAgent, EvolutionMixin): custom_handlers: dict[str, Callable[..., Coroutine]] | None = None, llm_gateway: Any = None, # NEW v2 param: LLMGateway mcp_servers: dict[str, str] | None = None, # NEW v2 param: MCP server URLs + compressor: Any = None, # CompressionStrategy | None ): # v2: If SkillConfig, extract skill info from agentkit.skills.base import SkillConfig, Skill @@ -277,6 +278,9 @@ class ConfigDrivenAgent(BaseAgent, EvolutionMixin): from agentkit.quality.output import OutputStandardizer self._output_standardizer = OutputStandardizer() + # v2: Store compressor for ReAct engine + self._compressor = compressor + # 从配置构建 Prompt 模板 if config.prompt: sections = PromptSection( @@ -656,6 +660,7 @@ class ConfigDrivenAgent(BaseAgent, EvolutionMixin): retrieval_config=retrieval_config or None, cancellation_token=cancellation_token, timeout_seconds=timeout_seconds, + compressor=self._compressor, ) # Parse result diff --git a/src/agentkit/server/app.py b/src/agentkit/server/app.py index d980108..499a49b 100644 --- a/src/agentkit/server/app.py +++ b/src/agentkit/server/app.py @@ -234,10 +234,26 @@ def create_app( app.state.mcp_manager = mcp_manager else: app.state.mcp_manager = None + # Initialize compressor if compression is configured + from agentkit.core.compressor import create_compressor + compressor = create_compressor(server_config.compression) if server_config else None + app.state.compressor = compressor + # Register headroom_retrieve tool if HeadroomCompressor is active + if compressor is not None: + try: + from agentkit.core.headroom_compressor import HeadroomCompressor + if isinstance(compressor, HeadroomCompressor) and compressor.is_available(): + from agentkit.tools.headroom_retrieve import HeadroomRetrieveTool + retrieve_tool = HeadroomRetrieveTool(compressor=compressor) + app.state.tool_registry.register(retrieve_tool) + logger.info("HeadroomRetrieveTool registered (CCR retrieval enabled)") + except ImportError: + pass app.state.agent_pool = AgentPool( llm_gateway=app.state.llm_gateway, skill_registry=app.state.skill_registry, tool_registry=app.state.tool_registry, + compressor=compressor, ) app.state.intent_router = IntentRouter(llm_gateway=app.state.llm_gateway) app.state.quality_gate = QualityGate() diff --git a/src/agentkit/server/config.py b/src/agentkit/server/config.py index 8900671..be7b66a 100644 --- a/src/agentkit/server/config.py +++ b/src/agentkit/server/config.py @@ -105,6 +105,7 @@ class ServerConfig: memory: dict[str, Any] | None = None, mcp_servers: dict[str, MCPServerConfig] | None = None, telemetry: dict[str, Any] | None = None, + compression: dict[str, Any] | None = None, on_change: Callable[["ServerConfig"], None] | None = None, ): self.host = host @@ -122,6 +123,7 @@ class ServerConfig: self.memory = memory or {} self.mcp_servers = mcp_servers or {} self.telemetry = telemetry or {} + self.compression = compression or {} self.on_change = on_change # Config watching state @@ -167,6 +169,9 @@ class ServerConfig: # Telemetry config telemetry_data = data.get("telemetry", {}) + # Compression config + compression_data = data.get("compression", {}) + return cls( host=server.get("host", "0.0.0.0"), port=server.get("port", 8001), @@ -183,6 +188,7 @@ class ServerConfig: memory=memory_data, mcp_servers=mcp_servers, telemetry=telemetry_data, + compression=compression_data, ) @staticmethod @@ -373,6 +379,7 @@ class ServerConfig: self.memory = new_config.memory self.mcp_servers = new_config.mcp_servers self.telemetry = new_config.telemetry + self.compression = new_config.compression self._last_mtime = new_config._last_mtime logger.info(f"Config reloaded from {path}") diff --git a/tests/unit/test_compression_config.py b/tests/unit/test_compression_config.py new file mode 100644 index 0000000..af384c3 --- /dev/null +++ b/tests/unit/test_compression_config.py @@ -0,0 +1,262 @@ +"""Tests for compression config integration (U4) + +Covers: +1. ServerConfig.compression field +2. create_app compression setup +3. ConfigDrivenAgent compressor passthrough +""" + +import tempfile +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agentkit.server.config import ServerConfig + + +# --------------------------------------------------------------------------- +# 1. ServerConfig.compression +# --------------------------------------------------------------------------- + + +class TestServerConfigCompression: + """Test compression field on ServerConfig""" + + def test_default_compression_is_empty_dict(self): + config = ServerConfig() + assert config.compression == {} + + def test_compression_from_dict(self): + data = { + "compression": { + "enabled": True, + "provider": "headroom", + "compressors": ["smart_crusher"], + } + } + config = ServerConfig.from_dict(data) + assert config.compression["enabled"] is True + assert config.compression["provider"] == "headroom" + assert config.compression["compressors"] == ["smart_crusher"] + + def test_compression_none_when_not_in_yaml(self): + data = {"server": {"host": "0.0.0.0"}} + config = ServerConfig.from_dict(data) + assert config.compression == {} + + def test_compression_hot_reload(self): + """_try_reload_config should update compression""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".yaml", delete=False + ) as f: + f.write( + "server:\n host: 0.0.0.0\n port: 8001\n" + "compression:\n enabled: false\n" + ) + f.flush() + config = ServerConfig.from_yaml(f.name) + assert config.compression == {"enabled": False} + + # Write new content + f2 = open(f.name, "w") + f2.write( + "server:\n host: 0.0.0.0\n port: 8001\n" + "compression:\n enabled: true\n provider: headroom\n" + ) + f2.close() + + config._try_reload_config(f.name) + assert config.compression["enabled"] is True + assert config.compression["provider"] == "headroom" + + +# --------------------------------------------------------------------------- +# 2. create_app compression setup +# --------------------------------------------------------------------------- + + +class TestCreateAppCompression: + """Test compression setup in create_app""" + + def test_compressor_created_when_enabled(self): + from agentkit.server.app import create_app + + with patch("agentkit.core.compressor.create_compressor") as mock_create: + mock_compressor = MagicMock() + mock_create.return_value = mock_compressor + + server_config = ServerConfig( + compression={"enabled": True, "provider": "summary"} + ) + app = create_app(server_config=server_config) + + mock_create.assert_called_once_with({"enabled": True, "provider": "summary"}) + assert app.state.compressor is mock_compressor + + def test_compressor_none_when_disabled(self): + from agentkit.server.app import create_app + + with patch("agentkit.core.compressor.create_compressor") as mock_create: + # create_compressor returns None when disabled + mock_create.return_value = None + + server_config = ServerConfig( + compression={"enabled": False} + ) + app = create_app(server_config=server_config) + + mock_create.assert_called_once_with({"enabled": False}) + assert app.state.compressor is None + + def test_compressor_none_when_no_config(self): + from agentkit.server.app import create_app + + with patch("agentkit.core.compressor.create_compressor") as mock_create: + mock_create.return_value = None + + # No server_config at all + app = create_app() + + # create_compressor should not be called (no server_config) + mock_create.assert_not_called() + assert app.state.compressor is None + + +# --------------------------------------------------------------------------- +# 3. ConfigDrivenAgent compressor passthrough +# --------------------------------------------------------------------------- + + +class TestConfigDrivenAgentCompression: + """Test compressor passthrough from ConfigDrivenAgent to ReActEngine""" + + @pytest.fixture + def agent_config(self): + from agentkit.core.config_driven import AgentConfig + + return AgentConfig( + name="test_agent", + agent_type="test", + task_mode="llm_generate", + prompt={"identity": "test", "instructions": "test"}, + ) + + @pytest.fixture + def skill_config(self): + from agentkit.skills.base import SkillConfig + + return SkillConfig( + name="test_skill", + agent_type="test", + description="test", + prompt={"identity": "test", "instructions": "test instructions"}, + execution_mode="react", + ) + + def test_compressor_stored_on_agent(self, agent_config): + from agentkit.core.config_driven import ConfigDrivenAgent + + mock_compressor = MagicMock() + agent = ConfigDrivenAgent( + config=agent_config, + compressor=mock_compressor, + ) + assert agent._compressor is mock_compressor + + def test_no_compressor_backward_compatible(self, agent_config): + from agentkit.core.config_driven import ConfigDrivenAgent + + agent = ConfigDrivenAgent(config=agent_config) + assert agent._compressor is None + + @pytest.mark.asyncio + async def test_compressor_passed_to_react_engine(self, skill_config): + from agentkit.core.config_driven import ConfigDrivenAgent + + mock_compressor = MagicMock() + + with patch.object( + ConfigDrivenAgent, "__init__", return_value=None + ) as mock_init: + # We need to test _handle_react directly, so set up the agent manually + agent = object.__new__(ConfigDrivenAgent) + agent._config = skill_config + agent._skill_config = skill_config + agent._prompt_template = None + agent._tools = [] + agent._memory_retriever = None + agent._compressor = mock_compressor + agent._evolution_enabled = False + agent._current_module = None + agent._active_tokens = {} + agent.name = "test_agent" + + # Mock the ReActEngine + mock_engine = MagicMock() + mock_result = MagicMock() + mock_result.output = '{"result": "ok"}' + mock_engine.execute = AsyncMock(return_value=mock_result) + agent._react_engine = mock_engine + + from agentkit.core.protocol import TaskMessage + from datetime import datetime, timezone + + task = TaskMessage( + task_id="t1", + agent_name="test_agent", + task_type="test", + input_data={"query": "hello"}, + priority=1, + created_at=datetime.now(timezone.utc), + callback_url=None, + ) + + await agent._handle_react(task) + + # Verify compressor was passed to execute + mock_engine.execute.assert_called_once() + call_kwargs = mock_engine.execute.call_args.kwargs + assert call_kwargs["compressor"] is mock_compressor + + @pytest.mark.asyncio + async def test_no_compressor_backward_compatible_react(self, skill_config): + from agentkit.core.config_driven import ConfigDrivenAgent + + with patch.object( + ConfigDrivenAgent, "__init__", return_value=None + ): + agent = object.__new__(ConfigDrivenAgent) + agent._config = skill_config + agent._skill_config = skill_config + agent._prompt_template = None + agent._tools = [] + agent._memory_retriever = None + agent._compressor = None + agent._evolution_enabled = False + agent._current_module = None + agent._active_tokens = {} + agent.name = "test_agent" + + mock_engine = MagicMock() + mock_result = MagicMock() + mock_result.output = '{"result": "ok"}' + mock_engine.execute = AsyncMock(return_value=mock_result) + agent._react_engine = mock_engine + + from agentkit.core.protocol import TaskMessage + from datetime import datetime, timezone + + task = TaskMessage( + task_id="t2", + agent_name="test_agent", + task_type="test", + input_data={"query": "hello"}, + priority=1, + created_at=datetime.now(timezone.utc), + callback_url=None, + ) + + await agent._handle_react(task) + + call_kwargs = mock_engine.execute.call_args.kwargs + assert call_kwargs["compressor"] is None