fischer
/
fischer-agentkit
1
0
Fork 0
Code Issues Pull Requests 1 Actions 91 Packages Projects Releases Wiki Activity

feat(compression): U4 ServerConfig compression field and Agent injection 

Add compression config to ServerConfig (following telemetry pattern),
create compressor in create_app, pass through AgentPool to
ConfigDrivenAgent, and inject into ReActEngine.execute() calls.
This commit is contained in:
chiguyong 2026-06-07 18:20:05 +08:00
parent fcb4fb33f3
commit 286804792d
5 changed files with 297 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/agentkit/core/agent_pool.py
View File

@ -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

5
src/agentkit/core/config_driven.py
View File

@ -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

16
src/agentkit/server/app.py
View File

@ -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()

7
src/agentkit/server/config.py
View File

@ -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}")

262
tests/unit/test_compression_config.py Normal file
View File

@ -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