363 lines
14 KiB
Python
363 lines
14 KiB
Python
"""ContextCompressor - 上下文压缩与 Prompt 缓存
|
||
|
||
长会话自动压缩历史消息，保持 Token 在预算内；
|
||
会话内 Prompt 不重复渲染。
|
||
"""
|
||
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
from typing import Any, Protocol, runtime_checkable
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _is_cjk(char: str) -> bool:
|
||
"""Check if a character is CJK (1 token ≈ 1 char).
|
||
|
||
Covers CJK Unified Ideographs, Hiragana, Katakana, and Hangul Syllables.
|
||
"""
|
||
cp = ord(char)
|
||
return (
|
||
0x4E00 <= cp <= 0x9FFF # CJK Unified Ideographs
|
||
or 0x3040 <= cp <= 0x30FF # Hiragana + Katakana
|
||
or 0xAC00 <= cp <= 0xD7AF # Hangul Syllables
|
||
)
|
||
|
||
|
||
def estimate_text_tokens(text: str) -> int:
|
||
"""Estimate token count: CJK 1:1, other characters 4:1.
|
||
|
||
CJK characters typically tokenize to ~1 token per character, while
|
||
ASCII/Latin text averages ~4 chars per token. Avoids the 4x
|
||
underestimation that ``len(text) // 4`` produces for CJK conversations.
|
||
|
||
ponytail ceiling: pure CJK may still underestimate ~10-20%, but
|
||
headroom_threshold=0.8 absorbs this. Upgrade path: litellm.token_counter
|
||
or provider-specific tokenizer.
|
||
"""
|
||
cjk_count = 0
|
||
non_cjk_count = 0
|
||
for char in text:
|
||
if _is_cjk(char):
|
||
cjk_count += 1
|
||
else:
|
||
non_cjk_count += 1
|
||
return cjk_count + non_cjk_count // 4
|
||
|
||
|
||
@runtime_checkable
|
||
class CompressionStrategy(Protocol):
|
||
"""压缩策略协议 — 所有压缩器必须实现此接口"""
|
||
|
||
async def compress(self, messages: list[dict]) -> list[dict]:
|
||
"""压缩消息列表"""
|
||
...
|
||
|
||
async def compress_tool_result(self, tool_name: str, result: Any) -> str:
|
||
"""压缩单个工具输出结果，返回压缩后的字符串"""
|
||
...
|
||
|
||
def is_available(self) -> bool:
|
||
"""检查压缩器是否可用"""
|
||
...
|
||
|
||
|
||
class ContextCompressor:
|
||
"""Compress long conversation histories to stay within token budgets"""
|
||
|
||
def __init__(
|
||
self,
|
||
llm_gateway: Any = None,
|
||
max_tokens: int = 4000,
|
||
keep_recent: int = 3,
|
||
model: str = "default",
|
||
model_context_limit: int = 128_000,
|
||
headroom_threshold: float = 0.8,
|
||
min_tokens: int = 8_000,
|
||
auxiliary_model: str | None = None,
|
||
):
|
||
self._llm_gateway = llm_gateway
|
||
self._max_tokens = max_tokens
|
||
self._keep_recent = keep_recent
|
||
self._model = model
|
||
# U3: Headroom-based compression trigger — predict context overflow
|
||
# before the single-request limit is hit.
|
||
self._model_context_limit = model_context_limit
|
||
self._headroom_threshold = headroom_threshold
|
||
self._min_tokens = min_tokens
|
||
# G4/U1: Auxiliary model for cost-sensitive summarization (e.g. "fast" alias).
|
||
# When set and differs from main model, _summarize tries auxiliary first,
|
||
# falls back to main model on failure OR empty content (Finding 4 anti-pattern).
|
||
# ponytail: ceiling — auxiliary is best-effort; main model is authoritative fallback.
|
||
self._auxiliary_model = auxiliary_model
|
||
|
||
def should_compress(self, messages: list[dict]) -> bool:
|
||
"""Check if compression should be triggered based on headroom ratio.
|
||
|
||
Triggers when either:
|
||
1. estimated_tokens / model_context_limit > headroom_threshold (headroom)
|
||
2. estimated_tokens > min_tokens (fixed fallback, preserves old behavior)
|
||
"""
|
||
estimated = self.estimate_tokens(messages)
|
||
if estimated / self._model_context_limit > self._headroom_threshold:
|
||
return True
|
||
if estimated > self._min_tokens:
|
||
return True
|
||
return False
|
||
|
||
def estimate_tokens(self, messages: list[dict]) -> int:
|
||
"""Estimate total tokens in message list (CJK 1:1, ASCII 4:1)"""
|
||
return sum(estimate_text_tokens(str(m.get("content", ""))) for m in messages)
|
||
|
||
async def compress(self, messages: list[dict]) -> list[dict]:
|
||
"""Compress messages if they exceed token budget.
|
||
|
||
Linear flow: summarize -> aggressive -> truncate.
|
||
Each step only fires if the previous didn't bring tokens under budget.
|
||
"""
|
||
tokens_before = self.estimate_tokens(messages)
|
||
if tokens_before <= self._max_tokens:
|
||
return messages
|
||
|
||
# Separate system messages, old messages, and recent messages
|
||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||
non_system = [m for m in messages if m.get("role") != "system"]
|
||
|
||
if len(non_system) <= self._keep_recent:
|
||
return messages # Not enough messages to compress
|
||
|
||
old_msgs = non_system[: -self._keep_recent]
|
||
recent_msgs = non_system[-self._keep_recent :]
|
||
|
||
# Step 1: Summarize old messages
|
||
summary = await self._summarize(old_msgs)
|
||
compressed = list(system_msgs)
|
||
if summary:
|
||
compressed.append(
|
||
{
|
||
"role": "system",
|
||
"content": f"## Conversation Summary\n{summary}",
|
||
}
|
||
)
|
||
compressed.extend(recent_msgs)
|
||
|
||
# Step 2: If still over budget, aggressive compress
|
||
# F-010: pass original `messages` (not `compressed`) to avoid summary-of-summary
|
||
strategy = "summary"
|
||
if self.estimate_tokens(compressed) > self._max_tokens:
|
||
compressed = await self._compress_aggressive(messages)
|
||
strategy = "aggressive"
|
||
|
||
# Step 3: If still over budget, truncate as last resort
|
||
if self.estimate_tokens(compressed) > self._max_tokens:
|
||
compressed = self._truncate(compressed)
|
||
strategy = "truncate"
|
||
|
||
# Step 4: Log compression result
|
||
tokens_after = self.estimate_tokens(compressed)
|
||
self._log_compression(tokens_before, tokens_after, len(messages), len(compressed), strategy)
|
||
|
||
return compressed
|
||
|
||
def _log_compression(
|
||
self,
|
||
tokens_before: int,
|
||
tokens_after: int,
|
||
msg_count_before: int,
|
||
msg_count_after: int,
|
||
strategy: str,
|
||
) -> None:
|
||
"""Log structured compression info (tokens_before/after/ratio/msg_count)."""
|
||
ratio = tokens_after / tokens_before if tokens_before > 0 else 0.0
|
||
logger.info(
|
||
"context compressed: %d -> %d tokens (%.1f%%), messages: %d -> %d, strategy: %s",
|
||
tokens_before,
|
||
tokens_after,
|
||
ratio * 100,
|
||
msg_count_before,
|
||
msg_count_after,
|
||
strategy,
|
||
)
|
||
|
||
async def _summarize(self, messages: list[dict], max_input_tokens: int = 3200) -> str:
|
||
"""Summarize a list of messages using LLM.
|
||
|
||
G4/U1: When ``auxiliary_model`` is configured and differs from the main
|
||
model, try auxiliary first (cost-optimization). On auxiliary failure OR
|
||
empty content (Finding 4 anti-pattern — "did not throw is not succeeded"),
|
||
fall back to main model. Existing ``_simple_summary`` degradation
|
||
preserved as the final tier when main model also fails.
|
||
"""
|
||
if not self._llm_gateway:
|
||
# No LLM available, do simple truncation
|
||
return self._simple_summary(messages)
|
||
|
||
# Build summary prompt
|
||
conversation_text = "\n".join(
|
||
f"[{m.get('role', 'unknown')}]: {m.get('content', '')}" for m in messages
|
||
)
|
||
|
||
# Pre-truncate if conversation_text exceeds safe token threshold
|
||
estimated_tokens = estimate_text_tokens(conversation_text)
|
||
if estimated_tokens > max_input_tokens:
|
||
# CJK-aware char limit: max_input_tokens chars is exact for CJK (1:1),
|
||
# conservative for ASCII (4:1, truncates to 1/4 budget but safe).
|
||
# Review fix #1: old `* 4` allowed 4x token budget for CJK text.
|
||
max_chars = max_input_tokens
|
||
conversation_text = conversation_text[:max_chars] + "\n...[truncated]"
|
||
|
||
prompt = (
|
||
"Summarize the following conversation history concisely, "
|
||
"preserving key facts, decisions, and context. "
|
||
"Focus on information that would be needed for continuing the conversation.\n\n"
|
||
f"{conversation_text}"
|
||
)
|
||
|
||
# G4: Try auxiliary model first when configured (cheap route).
|
||
if self._auxiliary_model and self._auxiliary_model != self._model:
|
||
try:
|
||
response = await self._llm_gateway.chat(
|
||
messages=[{"role": "user", "content": prompt}],
|
||
model=self._auxiliary_model,
|
||
agent_name="compressor",
|
||
task_type="summarization",
|
||
)
|
||
# Finding 4: empty content is a failure, not a success.
|
||
if response.content and response.content.strip():
|
||
return response.content
|
||
logger.info("Auxiliary model returned empty content, falling back to main model")
|
||
except Exception as e:
|
||
logger.info(
|
||
f"Auxiliary model summarization failed, falling back to main model: {e}"
|
||
)
|
||
|
||
# Main model path (or auxiliary fallback).
|
||
try:
|
||
response = await self._llm_gateway.chat(
|
||
messages=[{"role": "user", "content": prompt}],
|
||
model=self._model,
|
||
agent_name="compressor",
|
||
task_type="summarization",
|
||
)
|
||
return response.content
|
||
except Exception as e:
|
||
logger.warning(f"LLM summarization failed, using simple summary: {e}")
|
||
return self._simple_summary(messages)
|
||
|
||
def _simple_summary(self, messages: list[dict]) -> str:
|
||
"""Simple truncation-based summary when LLM is unavailable"""
|
||
parts = []
|
||
for msg in messages:
|
||
role = msg.get("role", "unknown")
|
||
content = str(msg.get("content", ""))[:200]
|
||
parts.append(f"[{role}]: {content}...")
|
||
return "\n".join(parts)
|
||
|
||
async def _compress_aggressive(self, messages: list[dict]) -> list[dict]:
|
||
"""Aggressive compression: keep only last message + summary of the rest."""
|
||
system_msgs = [m for m in messages if m.get("role") == "system"]
|
||
non_system = [m for m in messages if m.get("role") != "system"]
|
||
|
||
# Keep only the last message
|
||
if non_system:
|
||
summary = await self._summarize(non_system[:-1])
|
||
compressed = list(system_msgs)
|
||
if summary:
|
||
compressed.append(
|
||
{
|
||
"role": "system",
|
||
"content": f"## Conversation Summary\n{summary}",
|
||
}
|
||
)
|
||
compressed.append(non_system[-1])
|
||
return compressed
|
||
|
||
return messages
|
||
|
||
def _truncate(self, messages: list[dict]) -> list[dict]:
|
||
"""Last resort: truncate long messages"""
|
||
result = []
|
||
for msg in messages:
|
||
content = str(msg.get("content", ""))
|
||
if len(content) > self._max_tokens * 4:
|
||
msg = {**msg, "content": content[: self._max_tokens * 4] + "...[truncated]"}
|
||
result.append(msg)
|
||
return result
|
||
|
||
async def compress_tool_result(self, tool_name: str, result: Any) -> str:
|
||
"""默认实现：不做压缩，直接返回字符串表示"""
|
||
return str(result)
|
||
|
||
def is_available(self) -> bool:
|
||
"""ContextCompressor 始终可用"""
|
||
return True
|
||
|
||
|
||
def create_compressor(config: dict[str, Any] | None = None) -> CompressionStrategy | None:
|
||
"""根据配置创建压缩器实例
|
||
|
||
Args:
|
||
config: 压缩配置字典，支持以下字段：
|
||
- enabled: bool, 是否启用压缩（默认 False）
|
||
- provider: "headroom" | "summary", 压缩提供者
|
||
- max_tokens: int, token 预算（summary 模式）
|
||
- keep_recent: int, 保留最近 N 条消息（summary 模式）
|
||
- 其他 provider 特定配置
|
||
|
||
Returns:
|
||
CompressionStrategy 实例，或 None（未启用时）
|
||
"""
|
||
if not config or not config.get("enabled", False):
|
||
return None
|
||
|
||
provider = config.get("provider", "summary")
|
||
|
||
if provider == "headroom":
|
||
try:
|
||
from agentkit.core.headroom_compressor import HeadroomCompressor
|
||
|
||
compressor = HeadroomCompressor(config)
|
||
if compressor.is_available():
|
||
return compressor
|
||
logger.warning(
|
||
"HeadroomCompressor not available (headroom-ai not installed?). "
|
||
"Falling back to ContextCompressor."
|
||
)
|
||
except ImportError:
|
||
logger.warning(
|
||
"HeadroomCompressor module not available. Falling back to ContextCompressor."
|
||
)
|
||
# Fallback to summary compressor
|
||
return ContextCompressor(
|
||
max_tokens=config.get("max_tokens", 4000),
|
||
keep_recent=config.get("keep_recent", 3),
|
||
)
|
||
|
||
# Default: summary-based compression
|
||
return ContextCompressor(
|
||
max_tokens=config.get("max_tokens", 4000),
|
||
keep_recent=config.get("keep_recent", 3),
|
||
)
|
||
|
||
|
||
def render_cached(template, variables: dict[str, Any] | None = None) -> list[dict[str, str]]:
|
||
"""Render PromptTemplate with caching - returns cached result for same variables"""
|
||
cache_key = hashlib.md5(json.dumps(variables or {}, sort_keys=True).encode()).hexdigest()
|
||
|
||
if not hasattr(template, "_render_cache"):
|
||
template._render_cache = {}
|
||
|
||
if cache_key in template._render_cache:
|
||
return template._render_cache[cache_key]
|
||
|
||
result = template.render(variables=variables)
|
||
template._render_cache[cache_key] = result
|
||
return result
|
||
|
||
|
||
def clear_cache(template) -> None:
|
||
"""Clear the render cache on a PromptTemplate instance"""
|
||
if hasattr(template, "_render_cache"):
|
||
template._render_cache.clear()
|