feat(compression): U2 HeadroomCompressor with SmartCrusher and CCR cache
Add HeadroomCompressor implementing CompressionStrategy Protocol with content-type routing (JSON→SmartCrusher, code→CodeCompressor), CCR reversible compression cache, and graceful degradation when headroom-ai is not installed.
This commit is contained in:
chiguyong
parent
5d3a5f2bf3
commit
ea705b979b
|
|
@ -1,6 +1,7 @@
|
|||
"""AgentKit Core - 基础组件"""
|
||||
|
||||
from agentkit.core.base import BaseAgent
|
||||
from agentkit.core.compressor import CompressionStrategy, ContextCompressor, create_compressor
|
||||
from agentkit.core.config_driven import AgentConfig, ConfigDrivenAgent
|
||||
from agentkit.core.exceptions import (
|
||||
AgentAlreadyRegisteredError,
|
||||
|
|
@ -36,10 +37,20 @@ from agentkit.core.protocol import (
|
|||
TaskStatus,
|
||||
)
|
||||
|
||||
# Optional: HeadroomCompressor — only available when headroom-ai is installed
|
||||
try:
|
||||
from agentkit.core.headroom_compressor import HeadroomCompressor
|
||||
except ImportError:
|
||||
HeadroomCompressor = None # type: ignore[misc,assignment]
|
||||
|
||||
__all__ = [
|
||||
"BaseAgent",
|
||||
"AgentConfig",
|
||||
"ConfigDrivenAgent",
|
||||
"CompressionStrategy",
|
||||
"ContextCompressor",
|
||||
"create_compressor",
|
||||
"HeadroomCompressor",
|
||||
"AgentCapability",
|
||||
"AgentStatus",
|
||||
"CancellationToken",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,202 @@
|
|||
"""HeadroomCompressor — 基于 headroom-ai 的上下文压缩器
|
||||
|
||||
在工具输出拼装到对话历史前进行智能压缩，减少 60-90% token 消耗。
|
||||
使用 headroom-ai Library 模式集成，支持 SmartCrusher (JSON) 和 CodeCompressor (代码)。
|
||||
CCR 可逆压缩保证原始数据不丢失。
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from agentkit.core.compressor import CompressionStrategy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Optional dependency detection
|
||||
_HEADROOM_AVAILABLE = False
|
||||
headroom_compress = None # type: ignore[misc,assignment]
|
||||
try:
|
||||
from headroom import compress as headroom_compress
|
||||
_HEADROOM_AVAILABLE = True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def _is_json_content(text: str) -> bool:
|
||||
"""检测文本是否为 JSON 内容"""
|
||||
text = text.strip()
|
||||
if text.startswith(("{", "[")):
|
||||
try:
|
||||
json.loads(text)
|
||||
return True
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _is_code_content(text: str) -> bool:
|
||||
"""检测文本是否为代码内容"""
|
||||
# Common code patterns
|
||||
code_indicators = [
|
||||
r"^\s*(def |class |import |from |func |fn |pub |package |#include )", # Python/Go/Rust/Java/C
|
||||
r"^\s*(function |const |let |var |export |import )", # JS/TS
|
||||
r"```[a-z]", # Code blocks
|
||||
r"^\s*(if |for |while |try |catch |switch )", # Control flow
|
||||
]
|
||||
lines = text.split("\n")
|
||||
code_line_count = 0
|
||||
for line in lines[:20]: # Check first 20 lines
|
||||
for pattern in code_indicators:
|
||||
if re.search(pattern, line, re.MULTILINE):
|
||||
code_line_count += 1
|
||||
break
|
||||
# If more than 30% of first 20 lines look like code, treat as code
|
||||
return code_line_count > min(6, len(lines) * 0.3)
|
||||
|
||||
|
||||
class HeadroomCompressor:
|
||||
"""基于 headroom-ai 的上下文压缩器
|
||||
|
||||
支持 SmartCrusher (JSON) 和 CodeCompressor (代码) 两种压缩策略。
|
||||
CCR 可逆压缩保证原始数据可通过 headroom_retrieve 取回。
|
||||
|
||||
配置项:
|
||||
enabled: bool — 开关
|
||||
compressors: list[str] — 启用的压缩器 ["smart_crusher", "code_compressor"]
|
||||
ccr_ttl: int — CCR 缓存 TTL（秒），默认 300
|
||||
min_length: int — 最小压缩长度（字符），默认 500
|
||||
model: str — 传给 headroom 的模型名
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, Any]):
|
||||
self._config = config
|
||||
self._compressors = config.get("compressors", ["smart_crusher", "code_compressor"])
|
||||
self._ccr_ttl = config.get("ccr_ttl", 300)
|
||||
self._min_length = config.get("min_length", 500)
|
||||
self._model = config.get("model", "default")
|
||||
# CCR cache: hash -> original content
|
||||
self._ccr_cache: dict[str, str] = {}
|
||||
|
||||
def is_available(self) -> bool:
|
||||
"""检查 headroom-ai 是否已安装"""
|
||||
return _HEADROOM_AVAILABLE
|
||||
|
||||
async def compress(self, messages: list[dict]) -> list[dict]:
|
||||
"""压缩消息列表中 role=tool 的消息"""
|
||||
if not _HEADROOM_AVAILABLE:
|
||||
return messages
|
||||
|
||||
compressed = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "tool" and len(str(msg.get("content", ""))) >= self._min_length:
|
||||
try:
|
||||
original_content = str(msg.get("content", ""))
|
||||
# Use headroom compress on the tool message
|
||||
result = headroom_compress(
|
||||
[msg],
|
||||
model=self._model,
|
||||
)
|
||||
# result.messages contains the compressed messages
|
||||
if hasattr(result, "messages") and result.messages:
|
||||
compressed_msg = result.messages[0]
|
||||
# Store original in CCR cache
|
||||
ccr_hash = self._store_ccr(original_content)
|
||||
# Append CCR hash to compressed content
|
||||
content = compressed_msg.get("content", original_content)
|
||||
if ccr_hash:
|
||||
content += f"\n<!-- CCR:hash={ccr_hash} -->"
|
||||
compressed.append({**msg, "content": content})
|
||||
else:
|
||||
compressed.append(msg)
|
||||
except Exception as e:
|
||||
logger.warning(f"Headroom compression failed for tool message: {e}")
|
||||
compressed.append(msg)
|
||||
else:
|
||||
compressed.append(msg)
|
||||
|
||||
return compressed
|
||||
|
||||
async def compress_tool_result(self, tool_name: str, result: Any) -> str:
|
||||
"""压缩单个工具输出结果"""
|
||||
content = str(result)
|
||||
|
||||
if not _HEADROOM_AVAILABLE:
|
||||
return content
|
||||
|
||||
if len(content) < self._min_length:
|
||||
return content
|
||||
|
||||
try:
|
||||
# Route by content type
|
||||
content_type = self._detect_content_type(content)
|
||||
|
||||
if content_type == "json" and "smart_crusher" in self._compressors:
|
||||
compressed = self._compress_with_headroom(content, "smart_crusher")
|
||||
elif content_type == "code" and "code_compressor" in self._compressors:
|
||||
compressed = self._compress_with_headroom(content, "code_compressor")
|
||||
else:
|
||||
# No applicable compressor
|
||||
return content
|
||||
|
||||
if compressed and len(compressed) < len(content):
|
||||
ccr_hash = self._store_ccr(content)
|
||||
if ccr_hash:
|
||||
compressed += f"\n<!-- CCR:hash={ccr_hash} -->"
|
||||
return compressed
|
||||
|
||||
return content
|
||||
except Exception as e:
|
||||
logger.warning(f"Tool result compression failed for '{tool_name}': {e}")
|
||||
return content
|
||||
|
||||
def _detect_content_type(self, content: str) -> str:
|
||||
"""检测内容类型"""
|
||||
if _is_json_content(content):
|
||||
return "json"
|
||||
if _is_code_content(content):
|
||||
return "code"
|
||||
return "text"
|
||||
|
||||
def _compress_with_headroom(self, content: str, compressor: str) -> str | None:
|
||||
"""使用 headroom 压缩内容"""
|
||||
try:
|
||||
msg = [{"role": "user", "content": content}]
|
||||
result = headroom_compress(msg, model=self._model)
|
||||
if hasattr(result, "messages") and result.messages:
|
||||
return result.messages[0].get("content", content)
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning(f"Headroom {compressor} compression failed: {e}")
|
||||
return None
|
||||
|
||||
def _store_ccr(self, original: str) -> str | None:
|
||||
"""存储原始内容到 CCR 缓存，返回哈希"""
|
||||
import hashlib
|
||||
ccr_hash = hashlib.sha256(original.encode()).hexdigest()[:16]
|
||||
self._ccr_cache[ccr_hash] = original
|
||||
return ccr_hash
|
||||
|
||||
def retrieve(self, ccr_hash: str | None = None, query: str | None = None) -> dict:
|
||||
"""从 CCR 缓存检索原始数据"""
|
||||
if ccr_hash and ccr_hash in self._ccr_cache:
|
||||
return {
|
||||
"content": self._ccr_cache[ccr_hash],
|
||||
"ccr_hash": ccr_hash,
|
||||
"success": True,
|
||||
}
|
||||
|
||||
if query:
|
||||
# Simple keyword search in cached content
|
||||
results = []
|
||||
for h, content in self._ccr_cache.items():
|
||||
if query.lower() in content.lower():
|
||||
results.append({"ccr_hash": h, "content": content[:500]})
|
||||
if results:
|
||||
return {"results": results, "success": True}
|
||||
|
||||
return {
|
||||
"error": f"CCR hash '{ccr_hash}' not found in cache",
|
||||
"success": False,
|
||||
}
|
||||
|
|
@ -0,0 +1,361 @@
|
|||
"""HeadroomCompressor 单元测试
|
||||
|
||||
所有测试使用 mock headroom 模块，无需安装 headroom-ai。
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agentkit.core.headroom_compressor import (
|
||||
HeadroomCompressor,
|
||||
_is_code_content,
|
||||
_is_json_content,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_headroom_compress_mock(return_content="compressed"):
|
||||
"""创建 mock headroom.compress 函数，返回带有 messages 属性的结果对象"""
|
||||
mock_result = MagicMock()
|
||||
mock_result.messages = [{"role": "user", "content": return_content}]
|
||||
return mock_result
|
||||
|
||||
|
||||
def _long_json_content():
|
||||
"""生成超过 min_length 的 JSON 内容"""
|
||||
import json
|
||||
items = [{"id": i, "name": f"item_{i}", "description": f"description for item {i}"} for i in range(50)]
|
||||
return json.dumps({"items": items})
|
||||
|
||||
|
||||
def _long_code_content():
|
||||
"""生成超过 min_length 的代码内容"""
|
||||
lines = []
|
||||
for i in range(50):
|
||||
lines.append(f"def function_{i}():")
|
||||
lines.append(f" result = process_data({i})")
|
||||
lines.append(f" return result")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _long_text_content():
|
||||
"""生成超过 min_length 的纯文本内容"""
|
||||
return "This is plain text content. " * 100
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestHeadroomAvailability
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestHeadroomAvailability:
|
||||
"""测试 headroom-ai 可用性检测"""
|
||||
|
||||
def test_is_available_false_when_not_installed(self):
|
||||
"""_HEADROOM_AVAILABLE=False 时 is_available() 返回 False"""
|
||||
compressor = HeadroomCompressor({})
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", False):
|
||||
assert compressor.is_available() is False
|
||||
|
||||
def test_is_available_true_when_installed(self):
|
||||
"""_HEADROOM_AVAILABLE=True 时 is_available() 返回 True"""
|
||||
compressor = HeadroomCompressor({})
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True):
|
||||
assert compressor.is_available() is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestContentTypeDetection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestContentTypeDetection:
|
||||
"""测试内容类型检测函数"""
|
||||
|
||||
def test_json_content_detected(self):
|
||||
"""有效 JSON 对象被正确检测"""
|
||||
assert _is_json_content('{"key": "value"}') is True
|
||||
|
||||
def test_json_array_detected(self):
|
||||
"""有效 JSON 数组被正确检测"""
|
||||
assert _is_json_content('[1, 2, 3]') is True
|
||||
|
||||
def test_non_json_content(self):
|
||||
"""普通文本不被识别为 JSON"""
|
||||
assert _is_json_content("hello world") is False
|
||||
|
||||
def test_invalid_json_start(self):
|
||||
"""以 { 开头但无效的 JSON 不被识别"""
|
||||
assert _is_json_content("{invalid") is False
|
||||
|
||||
def test_code_content_detected(self):
|
||||
"""Python 代码（含 def/class 关键字）被正确检测"""
|
||||
code = "def hello():\n pass\n\nclass Foo:\n pass\nimport os\nfrom sys import path"
|
||||
assert _is_code_content(code) is True
|
||||
|
||||
def test_non_code_content(self):
|
||||
"""纯文本不被识别为代码"""
|
||||
text = "This is just a regular paragraph of text with no code keywords at all."
|
||||
assert _is_code_content(text) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestCompressToolResult
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCompressToolResult:
|
||||
"""测试 compress_tool_result 方法"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_content_not_compressed(self):
|
||||
"""短于 min_length 的内容不压缩"""
|
||||
compressor = HeadroomCompressor({"min_length": 500})
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True):
|
||||
result = await compressor.compress_tool_result("test_tool", "short content")
|
||||
assert result == "short content"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_content_compressed_with_smart_crusher(self):
|
||||
"""JSON 内容使用 smart_crusher 压缩"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher", "code_compressor"],
|
||||
})
|
||||
json_content = _long_json_content()
|
||||
mock_fn = MagicMock(return_value=_make_headroom_compress_mock("compressed json"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress_tool_result("json_tool", json_content)
|
||||
assert "compressed json" in result
|
||||
assert "<!-- CCR:hash=" in result
|
||||
mock_fn.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_content_compressed_with_code_compressor(self):
|
||||
"""代码内容使用 code_compressor 压缩"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher", "code_compressor"],
|
||||
})
|
||||
code_content = _long_code_content()
|
||||
mock_fn = MagicMock(return_value=_make_headroom_compress_mock("compressed code"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress_tool_result("code_tool", code_content)
|
||||
assert "compressed code" in result
|
||||
assert "<!-- CCR:hash=" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_text_content_not_compressed(self):
|
||||
"""纯文本内容无适用的压缩器，返回原始内容"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher", "code_compressor"],
|
||||
})
|
||||
text_content = _long_text_content()
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True):
|
||||
result = await compressor.compress_tool_result("text_tool", text_content)
|
||||
assert result == text_content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compression_includes_ccr_hash(self):
|
||||
"""压缩后的输出包含 CCR hash 标记"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher"],
|
||||
})
|
||||
json_content = _long_json_content()
|
||||
mock_fn = MagicMock(return_value=_make_headroom_compress_mock("small"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress_tool_result("tool", json_content)
|
||||
assert "<!-- CCR:hash=" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_headroom_not_available_returns_original(self):
|
||||
"""headroom-ai 不可用时返回原始内容"""
|
||||
compressor = HeadroomCompressor({"min_length": 100})
|
||||
json_content = _long_json_content()
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", False):
|
||||
result = await compressor.compress_tool_result("tool", json_content)
|
||||
assert result == json_content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_compression_failure_falls_back(self):
|
||||
"""headroom_compress 抛出异常时回退到原始内容"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher"],
|
||||
})
|
||||
json_content = _long_json_content()
|
||||
mock_fn = MagicMock(side_effect=RuntimeError("headroom error"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress_tool_result("tool", json_content)
|
||||
assert result == json_content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestCompress
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCompress:
|
||||
"""测试 compress 方法（消息列表压缩）"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_messages_compressed(self):
|
||||
"""role=tool 且内容足够长的消息被压缩"""
|
||||
compressor = HeadroomCompressor({"min_length": 100})
|
||||
long_content = "x" * 500
|
||||
messages = [
|
||||
{"role": "system", "content": "system prompt"},
|
||||
{"role": "tool", "content": long_content, "tool_call_id": "call_1"},
|
||||
]
|
||||
mock_fn = MagicMock(return_value=_make_headroom_compress_mock("compressed"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress(messages)
|
||||
assert len(result) == 2
|
||||
assert result[0]["role"] == "system"
|
||||
assert result[1]["role"] == "tool"
|
||||
assert "compressed" in result[1]["content"]
|
||||
assert "<!-- CCR:hash=" in result[1]["content"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_tool_messages_preserved(self):
|
||||
"""system/user/assistant 消息保持不变"""
|
||||
compressor = HeadroomCompressor({"min_length": 100})
|
||||
messages = [
|
||||
{"role": "system", "content": "system prompt"},
|
||||
{"role": "user", "content": "x" * 500},
|
||||
{"role": "assistant", "content": "y" * 500},
|
||||
]
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True):
|
||||
result = await compressor.compress(messages)
|
||||
assert result == messages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_short_tool_messages_not_compressed(self):
|
||||
"""短于 min_length 的 tool 消息不压缩"""
|
||||
compressor = HeadroomCompressor({"min_length": 500})
|
||||
messages = [
|
||||
{"role": "tool", "content": "short", "tool_call_id": "call_1"},
|
||||
]
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True):
|
||||
result = await compressor.compress(messages)
|
||||
assert result[0]["content"] == "short"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_headroom_not_available_returns_original(self):
|
||||
"""headroom-ai 不可用时返回原始消息列表"""
|
||||
compressor = HeadroomCompressor({"min_length": 100})
|
||||
messages = [
|
||||
{"role": "tool", "content": "x" * 500, "tool_call_id": "call_1"},
|
||||
]
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", False):
|
||||
result = await compressor.compress(messages)
|
||||
assert result == messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestCCRRetrieve
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestCCRRetrieve:
|
||||
"""测试 CCR 缓存检索"""
|
||||
|
||||
def test_retrieve_by_hash(self):
|
||||
"""通过 hash 检索存储的原始内容"""
|
||||
compressor = HeadroomCompressor({})
|
||||
original = "this is the original content"
|
||||
ccr_hash = compressor._store_ccr(original)
|
||||
assert ccr_hash is not None
|
||||
|
||||
result = compressor.retrieve(ccr_hash=ccr_hash)
|
||||
assert result["success"] is True
|
||||
assert result["content"] == original
|
||||
assert result["ccr_hash"] == ccr_hash
|
||||
|
||||
def test_retrieve_by_query(self):
|
||||
"""通过关键词搜索 CCR 缓存"""
|
||||
compressor = HeadroomCompressor({})
|
||||
compressor._store_ccr("Python is a programming language")
|
||||
compressor._store_ccr("Rust is a systems programming language")
|
||||
|
||||
result = compressor.retrieve(query="Python")
|
||||
assert result["success"] is True
|
||||
assert len(result["results"]) >= 1
|
||||
assert any("Python" in r["content"] for r in result["results"])
|
||||
|
||||
def test_retrieve_not_found(self):
|
||||
"""无效 hash 返回错误"""
|
||||
compressor = HeadroomCompressor({})
|
||||
result = compressor.retrieve(ccr_hash="nonexistent_hash")
|
||||
assert result["success"] is False
|
||||
assert "error" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ccr_hash_in_compressed_output(self):
|
||||
"""compress_tool_result 后内容可通过 CCR hash 检索"""
|
||||
compressor = HeadroomCompressor({
|
||||
"min_length": 100,
|
||||
"compressors": ["smart_crusher"],
|
||||
})
|
||||
json_content = _long_json_content()
|
||||
mock_fn = MagicMock(return_value=_make_headroom_compress_mock("compressed"))
|
||||
|
||||
with patch("agentkit.core.headroom_compressor._HEADROOM_AVAILABLE", True), \
|
||||
patch("agentkit.core.headroom_compressor.headroom_compress", mock_fn):
|
||||
result = await compressor.compress_tool_result("tool", json_content)
|
||||
# Extract hash from output
|
||||
import re
|
||||
match = re.search(r"CCR:hash=([a-f0-9]+)", result)
|
||||
assert match is not None
|
||||
ccr_hash = match.group(1)
|
||||
|
||||
# Retrieve original via hash
|
||||
retrieved = compressor.retrieve(ccr_hash=ccr_hash)
|
||||
assert retrieved["success"] is True
|
||||
assert retrieved["content"] == json_content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestHeadroomCompressorConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestHeadroomCompressorConfig:
|
||||
"""测试配置项"""
|
||||
|
||||
def test_default_config(self):
|
||||
"""默认配置值"""
|
||||
compressor = HeadroomCompressor({})
|
||||
assert compressor._compressors == ["smart_crusher", "code_compressor"]
|
||||
assert compressor._ccr_ttl == 300
|
||||
assert compressor._min_length == 500
|
||||
assert compressor._model == "default"
|
||||
|
||||
def test_custom_config(self):
|
||||
"""自定义配置值"""
|
||||
config = {
|
||||
"compressors": ["smart_crusher"],
|
||||
"ccr_ttl": 600,
|
||||
"min_length": 1000,
|
||||
"model": "gpt-4",
|
||||
}
|
||||
compressor = HeadroomCompressor(config)
|
||||
assert compressor._compressors == ["smart_crusher"]
|
||||
assert compressor._ccr_ttl == 600
|
||||
assert compressor._min_length == 1000
|
||||
assert compressor._model == "gpt-4"
|
||||
Loading…
Reference in New Issue