liqiang/ai-agent-book
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）
Build latest book artifacts / build (push) Canceled after 0s
Details
dependency resolution / resolve (3.11) (push) Canceled after 0s
Details
dependency resolution / resolve (3.13) (push) Canceled after 0s
Details
deploy-pages / build (push) Canceled after 0s
Details
deploy-pages / deploy (push) Canceled after 0s
Details
i18n consistency check / check (push) Canceled after 0s
Details
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
Details
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
Details
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
Details
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
Details
web-search-agent tests / test (push) Canceled after 0s
Details
web-search-agent tests / agentbook (push) Canceled after 0s
Details

Browse Source
This commit is contained in:
liqiang
2026-08-20 13:12:50 +00:00
commit b119135836
10275 changed files with 3284984 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
chapter2/context-compression/README.md
+525
View File
@@ -0,0 +1,525 @@
# Context Compression Strategies / 上下文压缩策略对比
> Companion material for *AI Agents in Depth*, Chapter 2 — **Experiment 2-10 ★★★: Comparison of context compression strategies**.
> 配套《深入理解 AI Agent》第 2 章 **实验 2-10 ★★★：上下文压缩策略对比**。
← [Chapter 2 index / 返回第 2 章目录](../README.md)
---
## Code map
- **Run first:** python experiment.py -s context_aware (or python quickstart.py for the menu).
- **Start here:** experiment.py::ExperimentRunner controls one strategy comparison.
- **Core behavior:** agent.py::ResearchAgent records the tool trajectory; compression_strategies.py::ContextCompressor applies the policy.
- **State / protocol:** AgentTrajectory, ToolCall and CompressionStrategy.
- **Verifier:** results/ JSON plus token/overflow counters; tests cover malformed tool results.
- **Experiment variable:** six compression strategies and the context-window budget.
- **Skip on first pass:** web search provider, streaming UI and plotting helpers.
## English
### Overview
Demonstrates and compares context compression strategies for LLM agents, using research on OpenAI co-founders’ current affiliations as the test task.
As context windows grow (128K+), efficient context management matters for:
- **Cost** — fewer tokens
- **Performance** — lower latency
- **Reliability** — fewer overflow errors
- **Relevance** — keep what matters
This lab implements and compares **6** strategies and their trade-offs.
### Compression strategies
#### 1. No compression
- Full webpage content into context
- Expected: fails after a few tool calls (overflow)
- Purpose: baseline problem
#### 2. Non-context-aware: individual summaries
- Summarize each page with LLM, then concatenate
- Preserves page-specific detail; may lose cross-page links
- Multiple LLM calls; good when sources are independent
#### 3. Non-context-aware: combined summary
- Concatenate all pages, then one summary
- Better overall picture; may lose per-page attribution
- One LLM call; may hit limits with many pages
#### 4. Context-aware summarization
- Query-focused summary over all search results
- Better relevance; extra LLM call
#### 5. Context-aware with citations
- Like #4 plus citations / source links
- Better for follow-ups; slightly larger
#### 6. Windowed context
- Full content for latest tool call; compress older history
- Balance detail vs efficiency
- Only compresses messages not already marked `[COMPRESSED]`
### Installation
```bash
# From the repository root: use the shared Chapter 2 environment
uv sync --locked --python 3.12 --extra ch2
# Activate it before changing directories:
# macOS/Linux:
source .venv/bin/activate
# Windows PowerShell: .venv\Scripts\Activate.ps1
# Windows cmd: .venv\Scripts\activate.bat
# pip fallback when uv is not installed:
# python -m pip install -e ".[ch2]"
cd chapter2/context-compression
# Single-project compatibility path, still supported during migration:
# python -m pip install -r requirements.txt
cp env.example .env
# Edit .env with your API keys
```
**API keys:**
- `LLM_PROVIDER` — `kimi` (default), `dashscope`/`qwen`/`bailian`, or `openrouter`.
- `DASHSCOPE_API_KEY` — Alibaba Cloud Model Studio / Bailian key when using DashScope; default model is `qwen3.7-plus` (set `DASHSCOPE_BASE_URL` for international keys).
- `MOONSHOT_API_KEY` — Kimi/Moonshot for live runs. Book 实验 2-10 uses Kimi K3 (~1M real window); the demo **caps** the compression/overflow budget at `CONTEXT_WINDOW_SIZE` (default 128K) so overflow/compression is observable. Override model via `MODEL_NAME` or `-m/--model` (e.g. `kimi-k2.5`, `kimi-k3`, `moonshot-v1-128k`).
- `OPENROUTER_API_KEY` — fallback if Moonshot key unset (`kimi-*` → `moonshotai/kimi-k2`). Unchanged if `MOONSHOT_API_KEY` is set.
- `SERPER_API_KEY` — web search (optional; mock data if missing)
Keys: [Moonshot](https://platform.moonshot.cn/), [Serper free tier](https://serper.dev/)
### Scripts overview
| Script | Purpose | Output |
|--------|---------|--------|
| `main.py` | Interactive demo / single strategy | Console |
| `experiment.py` | Automated comparison (token / compression / success table) | `results/` |
| `run_all_strategies.py` | Strategies with detailed per-round logs | `logs/` |
| `quickstart.py` | Menu wrapper (env check + launcher) | Console |
CLIs use Chinese `--help`. Shared useful flags:
- `-s/--strategy` — one or more strategies (default all 6); see list or `--list-strategies`
- `-m/--model` — override `MODEL_NAME`
- `-n/--max-iterations` — max tool-call rounds per strategy
Strategy aliases: `no_compression`, `individual`, `combined`, `context_aware`, `citations`, `windowed`.
### Usage
#### Full experiment (comparison table + JSON)
```bash
python experiment.py # all 6 strategies + comparison table
python experiment.py -s context_aware # one strategy
python experiment.py -s individual combined # two non-task-aware strategies
python experiment.py -m moonshot-v1-128k -o results/run.json
python experiment.py --list-strategies
```
Runs selected strategies sequentially, researches co-founder affiliations, prints Success / Time / **Tokens** / Compression / Overflows, saves `results/experiment_TIMESTAMP.json` (or `-o`).
Key flags: `-s/--strategy`, `-m/--model`, `-o/--output`, `-n/--max-iterations`, `--streaming`, `--list-strategies`.
#### All strategies with logging
```bash
python run_all_strategies.py
python run_all_strategies.py -s windowed
python run_all_strategies.py --log-dir logs/k2 -m kimi-k2.5
```
- Sequential strategies
- Compression summaries to log file
- Streaming by default
- Logs: `<log-dir>/strategy_run_TIMESTAMP.log`
- JSON: `<log-dir>/strategy_results_TIMESTAMP.json`
- End comparison summary
Flags: `-s/--strategy`, `-m/--model`, `--log-dir`, `-n/--max-iterations`, `--list-strategies`.
#### Interactive demo
```bash
python main.py # choose strategy at prompt
python main.py -s citations
python main.py -s windowed --no-streaming
```
Streaming on by default; follow-ups useful for citation strategy.
#### Custom usage
```python
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
agent = ResearchAgent(
api_key="your_api_key",
compression_strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS,
enable_streaming=True
)
result = agent.execute_research()
if result['success']:
print(result['final_answer'])
print(f"Tool calls: {len(result['trajectory'].tool_calls)}")
```
### Project structure
```
context-compression/
├── config.py
├── web_tools.py
├── compression_strategies.py
├── agent.py
├── experiment.py
├── run_all_strategies.py
├── main.py
├── quickstart.py
├── requirements.txt
├── env.example
├── logs/ # from run_all_strategies.py
└── results/ # experiment JSON
```
### Key components
- **web_tools.py:** `search_web` (Serper + crawl), `fetch_webpage`, mock data without key
- **compression_strategies.py:** `ContextCompressor`, `CompressedContent`, dynamic compression
- **agent.py:** streaming, tools, history, windowed compression
- **experiment.py:** automated runs, metrics, comparison table, JSON
### Metrics
Success rate, execution time, compression ratio (compressed/original size), context overflows, tool calls, final answer length.
### Expected results (qualitative)
1. No compression → overflow fail
2. Non-context-aware → may complete, miss detail
3. Context-aware → good size/relevance balance
4. With citations → best for follow-ups
5. Windowed → efficient for long multi-turn
### Measured results (real run)
Real end-to-end run (no mock): live Serper + Moonshot reasoning model.
- **Model:** `kimi-k3` (real window ~1M; demo budget `CONTEXT_WINDOW_SIZE = 128000`)
- **Search:** real Serper + page crawl
- **Task:** track current affiliations of ~11 OpenAI co-founders
- **Date:** 2026-07-18 · `MAX_ITERATIONS=15` · raw: `results/kimi_k3_real_20260718.json`
| # | Strategy | Success | Iterations | Tokens | Compress | Overflows | Time |
|---|----------|---------|-----------|--------|----------|-----------|------|
| 1 | `no_compression` | ❌ (overflow at 165,227 tok > 128K) | 5 | 166,043 | 102.1% | 1 | 107s |
| 2 | `non_context_aware_individual_summary` | ✅ | 12 | 276,608 | 10.9% | 4 | 2980s |
| 3 | `non_context_aware_combined_summary` | ✅ | 10 | 93,449 | 4.3% | 0 | 1189s |
| 4 | `context_aware_summary` | ✅ | 7 | 40,157 | 3.0% | 0 | 967s |
| 5 | `context_aware_with_citations` | ✅ | 10 | 222,992 | 4.1% | 3 | 1235s |
| 6 | `windowed_context` | ✅ | 7 | 174,601 | 102.4% | 4 | 867s |
Notes:
- **No compression** fails as designed past 128K (~5th iteration).
- **Context-aware summary (#4)** most token-efficient success (40,157 tokens, 3.0% char compression).
- **Individual summaries (#2)** slowest (~50 min): per-page summaries on a reasoning model.
- **Windowed (#6)** compresses only when usage crosses ~80% of budget; keeps recent full content → char “compression ratio” ~100% while still finishing fastest among compressing strategies.
- Single-run numbers vary; relative ordering is the takeaway.
### Configuration
`.env` or `config.py`:
- `MODEL_NAME` (default kimi-k3)
- `MODEL_TEMPERATURE` (default 0.3)
- `MAX_ITERATIONS` (default 50)
- `MAX_WEBPAGE_LENGTH` (default 50000)
- `SUMMARY_MAX_TOKENS` (default 500)
- `CONTEXT_WINDOW_SIZE` (default 128000; intentional cap vs K3’s real ~1M window)
### Troubleshooting
- **No Serper key:** mock data still exercises compression logic
- **Overflow on non-baseline strategies:** lower `MAX_WEBPAGE_LENGTH` / `SUMMARY_MAX_TOKENS` / search `num_results`
- **Slow:** `--no-streaming`, lower `-n/--max-iterations`, mock search
### Research task
> “Find the current affiliations of all OpenAI co-founders”
Good because it needs many searches, accumulates text, stresses context management, and has checkable outcomes.
### Extending
New strategy: enum → `ContextCompressor` → `compress_search_results()` → experiment runner.
New task: system prompt in `agent.py`, mock data in `web_tools.py`, tool descriptions as needed.
---
## 中文
### 概述
演示并对比 LLM Agent 的多种上下文压缩策略，测试任务为调研 OpenAI 联合创始人当前职业归属。
上下文窗口越来越大（128K+）时，高效管理上下文关乎：
- **成本** — 减少 token
- **性能** — 更低延迟
- **可靠性** — 减少溢出错误
- **相关性** — 保留关键信息
本实验实现并对比 **6** 种策略及其取舍。
### 压缩策略
#### 1. 无压缩
- 网页原文直接进入上下文
- 预期：几次工具调用后溢出失败
- 目的：展示基线问题
#### 2. 非任务感知：逐页摘要
- 每页单独 LLM 摘要再拼接
- 保留页内细节，可能丢跨页关系
- 多次 LLM 调用；适合来源彼此独立
#### 3. 非任务感知：合并摘要
- 先拼接全部网页再做一次总摘要
- 更利把握全局，可能丢页级归属
- 单次 LLM 调用；页多时可能撞限
#### 4. 上下文感知摘要
- 结合查询对全部搜索结果做聚焦摘要
- 相关性更好；多一次 LLM 调用
#### 5. 带引用的上下文感知摘要
- 在 #4 基础上加引用与来源链接
- 利于追问；上下文略大
#### 6. 窗口化上下文
- 最近一次工具调用保留全文，更早历史压缩
- 细节与效率折中
- 只压缩尚未标记 `[COMPRESSED]` 的消息
### 安装
```bash
# 在仓库根目录使用统一的第 2 章环境
uv sync --locked --python 3.12 --extra ch2
# 切换目录前先激活环境：
# macOS/Linux：
source .venv/bin/activate
# Windows PowerShell：.venv\Scripts\Activate.ps1
# Windows cmd：.venv\Scripts\activate.bat
# 未安装 uv 时可用 pip 兜底：
# python -m pip install -e ".[ch2]"
cd chapter2/context-compression
# 迁移期间仍支持单项目兼容路径：
# python -m pip install -r requirements.txt
cp env.example .env
# 编辑 .env 填入 API Key
```
**所需 Key：**
- `LLM_PROVIDER`：`kimi`（默认）、`dashscope`/`qwen`/`bailian` 或 `openrouter`。
- `DASHSCOPE_API_KEY`：使用阿里云百炼 / Model Studio 时的 Key，默认模型 `qwen3.7-plus`（国际区 Key 可设置 `DASHSCOPE_BASE_URL`）。
- `MOONSHOT_API_KEY`：Kimi/Moonshot（在线跑必需）。书中实验 2-10 使用 Kimi K3（真实窗口约 1M）；演示通过 `CONTEXT_WINDOW_SIZE`（默认 128K）**故意收紧**溢出/压缩预算以便观察。可用 `MODEL_NAME` 或 `-m/--model` 覆盖（如 `kimi-k2.5`、`kimi-k3`、`moonshot-v1-128k`）。
- `OPENROUTER_API_KEY`：未设置 Moonshot key 时的通用回退（`kimi-*` → `moonshotai/kimi-k2`）。设了 `MOONSHOT_API_KEY` 时行为不变。
- `SERPER_API_KEY`：联网搜索（可选；缺失则用 mock 数据）
获取：[Moonshot](https://platform.moonshot.cn/)、[Serper 免费档](https://serper.dev/)
### 脚本一览
| 脚本 | 作用 | 输出 |
|------|------|------|
| `main.py` | 交互演示 / 单策略 | 控制台 |
| `experiment.py` | 自动对比（token / 压缩 / 成功表） | `results/` |
| `run_all_strategies.py` | 带逐轮详细日志 | `logs/` |
| `quickstart.py` | 菜单封装（检查环境并启动） | 控制台 |
均提供中文 `--help`。共用常用参数：
- `-s/--strategy` — 一种或多种策略（默认全部 6 种）；见列表或 `--list-strategies`
- `-m/--model` — 覆盖 `MODEL_NAME`
- `-n/--max-iterations` — 每策略最大工具调用轮数
策略别名：`no_compression`、`individual`、`combined`、`context_aware`、`citations`、`windowed`。
### 用法
#### 完整实验（对比表 + JSON）
```bash
python experiment.py # 运行全部 6 种策略并生成对比表
python experiment.py -s context_aware # 只运行「上下文感知压缩」
python experiment.py -s individual combined # 只对比两种非任务感知策略
python experiment.py -m moonshot-v1-128k -o results/run.json
python experiment.py --list-strategies
```
依次测试所选策略、调研联合创始人归属、打印 Success / Time / **Tokens** / Compression / Overflows，保存到 `results/experiment_TIMESTAMP.json`（或 `-o`）。
主要参数：`-s/--strategy`、`-m/--model`、`-o/--output`、`-n/--max-iterations`、`--streaming`、`--list-strategies`。
#### 带日志跑全部策略
```bash
python run_all_strategies.py
python run_all_strategies.py -s windowed
python run_all_strategies.py --log-dir logs/k2 -m kimi-k2.5
```
- 顺序跑所选策略
- 压缩摘要写入日志
- 默认流式
- 日志：`<log-dir>/strategy_run_TIMESTAMP.log`
- JSON：`<log-dir>/strategy_results_TIMESTAMP.json`
- 末尾对比摘要
参数：`-s/--strategy`、`-m/--model`、`--log-dir`、`-n/--max-iterations`、`--list-strategies`。
#### 交互演示
```bash
python main.py # 提示选择策略
python main.py -s citations
python main.py -s windowed --no-streaming
```
默认开启流式；引用策略适合追问。
#### 编程调用
```python
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
agent = ResearchAgent(
api_key="your_api_key",
compression_strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS,
enable_streaming=True
)
result = agent.execute_research()
if result['success']:
print(result['final_answer'])
print(f"Tool calls: {len(result['trajectory'].tool_calls)}")
```
### 项目结构
```
context-compression/
├── config.py
├── web_tools.py
├── compression_strategies.py
├── agent.py
├── experiment.py
├── run_all_strategies.py
├── main.py
├── quickstart.py
├── requirements.txt
├── env.example
├── logs/
└── results/
```
### 关键组件
- **web_tools.py：** `search_web`（Serper + 抓取）、`fetch_webpage`、无 Key 时 mock
- **compression_strategies.py：** `ContextCompressor`、`CompressedContent`、动态压缩
- **agent.py：** 流式、工具、历史、窗口化压缩
- **experiment.py：** 自动跑、指标、对比表、JSON
### 采集指标
成功率、执行时间、压缩比（压缩后/原始）、上下文溢出次数、工具调用次数、最终答案长度。
### 定性预期
1. 无压缩 → 溢出失败
2. 非任务感知 → 可能完成但丢细节
3. 上下文感知 → 体积与相关性较均衡
4. 带引用 → 最利于追问
5. 窗口化 → 长对话更高效
### 实测结果（真实运行）
真实端到端（无 mock）：实时 Serper + Moonshot 推理模型。
- **模型：** `kimi-k3`（真实窗口约 1M；演示预算 `CONTEXT_WINDOW_SIZE = 128000`）
- **搜索：** 真实 Serper + 页面抓取
- **任务：** 识别并追踪约 11 位 OpenAI 联合创始人的职业状态
- **日期：** 2026-07-18 · `MAX_ITERATIONS=15` · 原始 JSON：`results/kimi_k3_real_20260718.json`
| # | Strategy | Success | Iterations | Tokens | Compress | Overflows | Time |
|---|----------|---------|-----------|--------|----------|-----------|------|
| 1 | `no_compression` | ❌ (overflow at 165,227 tok > 128K) | 5 | 166,043 | 102.1% | 1 | 107s |
| 2 | `non_context_aware_individual_summary` | ✅ | 12 | 276,608 | 10.9% | 4 | 2980s |
| 3 | `non_context_aware_combined_summary` | ✅ | 10 | 93,449 | 4.3% | 0 | 1189s |
| 4 | `context_aware_summary` | ✅ | 7 | 40,157 | 3.0% | 0 | 967s |
| 5 | `context_aware_with_citations` | ✅ | 10 | 222,992 | 4.1% | 3 | 1235s |
| 6 | `windowed_context` | ✅ | 7 | 174,601 | 102.4% | 4 | 867s |
说明：
- **无压缩**按设计在超过 128K 时失败（约第 5 轮）。
- **上下文感知摘要（#4）** token 最省（40,157 tokens，字符压缩 3.0%）。
- **逐页摘要（#2）**最慢（约 50 分钟）：推理模型对每页单独摘要。
- **窗口化（#6）**仅在用量跨过约 80% 预算时批量压缩未压缩工具消息；保留近期全文，字符「压缩比」约 100%，但在可完成任务的策略中总时间最短。
- 单次运行绝对值会波动；相对排序是关键 takeaway。
### 配置
`.env` 或 `config.py`：
- `MODEL_NAME`（默认 kimi-k3）
- `MODEL_TEMPERATURE`（默认 0.3）
- `MAX_ITERATIONS`（默认 50）
- `MAX_WEBPAGE_LENGTH`（默认 50000）
- `SUMMARY_MAX_TOKENS`（默认 500）
- `CONTEXT_WINDOW_SIZE`（默认 128000；相对 K3 真实 ~1M 的故意收紧）
### 故障排除
- **无 Serper Key：** mock 仍可验证压缩逻辑
- **非基线策略仍溢出：** 降低 `MAX_WEBPAGE_LENGTH` / `SUMMARY_MAX_TOKENS` / 搜索 `num_results`
- **偏慢：** `--no-streaming`、减小 `-n/--max-iterations`、改用 mock 搜索
### 研究任务
> 「查找所有 OpenAI 联合创始人的当前职业归属」
适合原因：需多次搜索、内容量大、考验上下文累积管理，结果可核对。
### 扩展
新策略：枚举 → `ContextCompressor` → `compress_search_results()` → 实验 runner。
新任务：改 `agent.py` 系统提示、`web_tools.py` mock、工具描述。
---
## Notes / 说明
- The 128K budget is intentional so compression/overflow behavior is visible even on models with larger real windows.
- 128K 预算是故意收紧的，以便在真实窗口更大的模型上仍能观察到压缩与溢出行为。
chapter2/context-compression/agent.py
+675
View File
@@ -0,0 +1,675 @@
"""
Context Compression Research Agent with Streaming Support
"""
import json
import logging
import time
import sys
from typing import List, Dict, Any, Optional, Generator, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from openai import OpenAI
from config import Config
from web_tools import WebTools
from compression_strategies import (
CompressionStrategy,
ContextCompressor,
CompressedContent
)
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
# Configure logging
logging.basicConfig(level=logging.INFO, format=Config.LOG_FORMAT)
logger = logging.getLogger(__name__)
@dataclass
class ToolCall:
"""Represents a single tool call"""
tool_name: str
arguments: Dict[str, Any]
result: Optional[Any] = None
compressed_result: Optional[CompressedContent] = None
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
# Provider-side tool_call id, so a tool message in the history can be
# matched back to the call that produced it (used by windowed compression
# to recover the original query).
id: Optional[str] = None
@dataclass
class AgentTrajectory:
"""Tracks the agent's execution trajectory"""
tool_calls: List[ToolCall] = field(default_factory=list)
total_tokens_used: int = 0
prompt_tokens_used: int = 0
completion_tokens_used: int = 0
# Prompt tokens of the most recent API call = the current context size.
# prompt_tokens_used above is a cumulative COST counter (each call's
# prompt re-counts the shared prefix), so it must not be compared
# against the per-request context window.
last_prompt_tokens: int = 0
context_overflows: int = 0
compression_strategy: CompressionStrategy = CompressionStrategy.NO_COMPRESSION
start_time: float = field(default_factory=time.time)
end_time: Optional[float] = None
class ResearchAgent:
"""
AI Agent for researching with context compression
"""
def __init__(
self,
api_key: str,
compression_strategy: CompressionStrategy = CompressionStrategy.NO_COMPRESSION,
verbose: bool = False,
enable_streaming: bool = True
):
"""
Initialize the research agent
Args:
api_key: API key for Moonshot/Kimi
compression_strategy: Strategy for context compression
verbose: Enable verbose logging
enable_streaming: Enable streaming responses
"""
# Moonshot 官方 key 存在则直连；否则回退 OpenRouter（见 Config.resolve_llm）。
resolved_key, resolved_base_url, resolved_model = Config.resolve_llm()
self.client = OpenAI(
api_key=resolved_key,
base_url=resolved_base_url
)
self.model = resolved_model
self.compression_strategy = compression_strategy
self.verbose = verbose
self.enable_streaming = enable_streaming
# Initialize tools
self.web_tools = WebTools()
self.compressor = ContextCompressor(compression_strategy, api_key, enable_streaming)
# Initialize trajectory
self.trajectory = AgentTrajectory(compression_strategy=compression_strategy)
# Initialize conversation history
self.conversation_history = []
self._init_system_prompt()
logger.info(f"Agent initialized with compression strategy: {compression_strategy.value}")
def _init_system_prompt(self):
"""Initialize the system prompt for OpenAI co-founders research"""
# Get current date dynamically
from datetime import datetime
today = datetime.now()
date_string = today.strftime("%A, %B %d, %Y")
self.conversation_history = [
{
"role": "system",
"content": f"""You are a research assistant tasked with finding information about OpenAI co-founders.
Your task is to:
1. First, search for and identify ALL OpenAI co-founders
2. Then, search for EACH co-founder individually to find their CURRENT affiliations
3. Compile a comprehensive report with current status for each co-founder
Important instructions:
- Be thorough and systematic - search for each person individually
- Focus on CURRENT affiliations, not historical roles
- Include company names, positions, and any recent changes
- If someone left a position, note where they went
- When you have gathered all information, provide a FINAL ANSWER with a complete list
Available tools:
- search_web: Search the web for information
- fetch_webpage: Fetch specific webpage content
Start by searching for the complete list of OpenAI co-founders.
TODAY'S DATE: {date_string}"""
}
]
def _get_tools_description(self) -> List[Dict[str, Any]]:
"""Get tool descriptions for the model"""
return [
{
"type": "function",
"function": {
"name": "search_web",
"description": "Search the web for information. Returns multiple search results with content from each webpage.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query"
},
"num_results": {
"type": "integer",
"description": "Number of results to return (default: 5)",
"default": 5
}
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "fetch_webpage",
"description": "Fetch and extract text content from a specific webpage URL",
"parameters": {
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "The URL of the webpage to fetch"
}
},
"required": ["url"]
}
}
}
]
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Tuple[Any, Optional[CompressedContent]]:
"""
Execute a tool and return the result with optional compression
Args:
tool_name: Name of the tool to execute
arguments: Arguments for the tool
Returns:
Tuple of (tool result, compressed content if applicable)
"""
if not isinstance(arguments, dict):
arguments = {}
if tool_name == "search_web":
if "query" not in arguments:
return {"error": "Missing required argument 'query' for search_web"}, None
try:
result = self.web_tools.search_web(**arguments)
except Exception as e:
logger.error(f"Failed to execute search_web: {e}")
return {"error": f"Failed to execute search_web: {e}"}, None
# Apply compression strategy
query = arguments.get('query', '')
current_context = self._get_current_context_summary()
compressed = self.compressor.compress_search_results(
result,
query,
current_context
)
return result, compressed
elif tool_name == "fetch_webpage":
if "url" not in arguments:
return {"error": "Missing required argument 'url' for fetch_webpage"}, None
try:
result = self.web_tools.fetch_webpage(**arguments)
except Exception as e:
logger.error(f"Failed to execute fetch_webpage: {e}")
return {"error": f"Failed to execute fetch_webpage: {e}"}, None
# For fetch, we typically don't compress (used for follow-ups)
return result, None
else:
return {"error": f"Unknown tool: {tool_name}"}, None
def _get_current_context_summary(self) -> str:
"""Get a summary of current context for context-aware compression"""
if not self.trajectory.tool_calls:
return ""
# Get last few tool calls for context
recent_calls = self.trajectory.tool_calls[-3:]
context_parts = []
for call in recent_calls:
context_parts.append(f"Previous search: {call.arguments.get('query', 'N/A')}")
return " | ".join(context_parts)
def _handle_windowed_compression(self, messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Apply windowed compression strategy to message history
Only compresses when context usage exceeds 80% threshold
Args:
messages: Current message history
Returns:
Messages with compressed history when needed
"""
if self.compression_strategy != CompressionStrategy.WINDOWED_CONTEXT:
return messages
# Check if we should start compressing (80% context usage).
# Use the LAST call's prompt size (current context), not the
# cumulative cost counter, which grows quadratically and would
# trigger compression long before the window is actually near full.
context_threshold = Config.CONTEXT_WINDOW_SIZE * 0.8
if self.trajectory.last_prompt_tokens <= context_threshold:
logger.debug(f"Windowed compression: Context usage below threshold ({self.trajectory.last_prompt_tokens:,}/{context_threshold:.0f} tokens)")
return messages # No compression needed yet
logger.info(f"⚠️ Context usage exceeds 80% threshold ({self.trajectory.last_prompt_tokens:,}/{Config.CONTEXT_WINDOW_SIZE} tokens) - Starting compression")
# Compression marker to identify already-compressed messages
COMPRESSION_MARKER = "[COMPRESSED]"
# First, count how many tool messages we have and how many need compression
tool_messages_to_compress = []
already_compressed_count = 0
for i, msg in enumerate(messages):
if msg.get('role') == 'tool':
original_content = msg.get('content', '')
if original_content.startswith(COMPRESSION_MARKER):
already_compressed_count += 1
else:
tool_messages_to_compress.append((i, msg))
total_tool_messages = already_compressed_count + len(tool_messages_to_compress)
if not tool_messages_to_compress:
logger.debug(f"Windowed compression: All {total_tool_messages} tool messages already compressed")
return messages # All tool messages already compressed
logger.info(f"📊 Compressing {len(tool_messages_to_compress)} uncompressed tool messages (out of {total_tool_messages} total)")
# Build the result with compression for all uncompressed tool messages
compressed_messages = []
compressed_in_this_pass = 0
for i, msg in enumerate(messages):
if msg.get('role') == 'tool':
original_content = msg.get('content', '')
# Check if already compressed
if original_content.startswith(COMPRESSION_MARKER):
# Already compressed, keep as is
compressed_messages.append(msg)
else:
# Compress this tool result
compressed_in_this_pass += 1
# Find the corresponding tool call to get context
tool_call_id = msg.get('tool_call_id')
query = "Information search" # Default
# Try to find the query from the tool call
for call in self.trajectory.tool_calls:
if call.id is not None and call.id == tool_call_id:
query = call.arguments.get('query', query)
break
logger.debug(f"Compressing tool message {compressed_in_this_pass}/{len(tool_messages_to_compress)} at index {i} (query: {query[:50]}...)")
compressed = self.compressor.compress_for_history(
original_content,
'search_web',
query,
preserve_citations=True
)
logger.debug(f"Compressed: {compressed.original_length:,} → {compressed.compressed_length:,} chars")
# Mark as compressed with clear marker
compressed_content = (
f"{COMPRESSION_MARKER} "
f"[Original: {compressed.original_length:,} chars → Compressed: {compressed.compressed_length:,} chars]\n"
f"{compressed.content}"
)
compressed_messages.append({
**msg,
'content': compressed_content
})
else:
compressed_messages.append(msg)
logger.info(f"✅ Compressed {compressed_in_this_pass} tool messages in this pass")
return compressed_messages
def _stream_response(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
"""
Stream response from the model
Args:
messages: Conversation messages
Returns:
Complete message object with token usage
"""
try:
stream = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=self._get_tools_description(),
tool_choice="auto",
temperature=_reasoning_safe_temperature(self.model, Config.MODEL_TEMPERATURE),
max_tokens=Config.MODEL_MAX_TOKENS,
stream=True,
stream_options={"include_usage": True} # Request token usage in stream
)
collected_chunks = []
collected_messages = []
current_tool_calls = []
usage_data = None
print("\n🤖 Assistant: ", end="", flush=True)
for chunk in stream:
collected_chunks.append(chunk)
# Capture usage data if present (might be in a chunk without choices)
if hasattr(chunk, 'usage') and chunk.usage is not None:
usage_data = chunk.usage
# Check if chunk has choices before accessing
if hasattr(chunk, 'choices') and chunk.choices and len(chunk.choices) > 0:
delta = chunk.choices[0].delta
# Handle content
if hasattr(delta, 'content') and delta.content:
content = delta.content
print(content, end="", flush=True)
collected_messages.append(content)
# Handle tool calls in streaming
if hasattr(delta, 'tool_calls') and delta.tool_calls:
for tool_call_delta in delta.tool_calls:
if tool_call_delta.index is not None:
# Ensure we have enough tool calls in the list
while len(current_tool_calls) <= tool_call_delta.index:
current_tool_calls.append({
"id": "",
"type": "function",
"function": {"name": "", "arguments": ""}
})
if tool_call_delta.id:
current_tool_calls[tool_call_delta.index]["id"] = tool_call_delta.id
if tool_call_delta.function:
if tool_call_delta.function.name:
current_tool_calls[tool_call_delta.index]["function"]["name"] = tool_call_delta.function.name
if tool_call_delta.function.arguments:
current_tool_calls[tool_call_delta.index]["function"]["arguments"] += tool_call_delta.function.arguments
print("\n", flush=True)
# Log token usage if available
if usage_data:
prompt_tokens = usage_data.prompt_tokens if hasattr(usage_data, 'prompt_tokens') else 0
completion_tokens = usage_data.completion_tokens if hasattr(usage_data, 'completion_tokens') else 0
total_tokens = usage_data.total_tokens if hasattr(usage_data, 'total_tokens') else 0
logger.info(f"🔢 Kimi API Token Usage - Prompt: {prompt_tokens}, Completion: {completion_tokens}, Total: {total_tokens}")
# Update trajectory
self.trajectory.last_prompt_tokens = prompt_tokens
self.trajectory.prompt_tokens_used += prompt_tokens
self.trajectory.completion_tokens_used += completion_tokens
self.trajectory.total_tokens_used += total_tokens
# Construct the complete message
complete_message = {
"role": "assistant",
"content": "".join(collected_messages) if collected_messages else None
}
if current_tool_calls:
complete_message["tool_calls"] = current_tool_calls
return complete_message
except Exception as e:
logger.error(f"Error in streaming response: {str(e)}")
raise
def _non_streaming_response(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
"""
Get non-streaming response from the model
Args:
messages: Conversation messages
Returns:
Complete message object with token usage
"""
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=self._get_tools_description(),
tool_choice="auto",
temperature=_reasoning_safe_temperature(self.model, Config.MODEL_TEMPERATURE),
max_tokens=Config.MODEL_MAX_TOKENS,
stream=False
)
message = response.choices[0].message
# Log token usage
if hasattr(response, 'usage') and response.usage:
prompt_tokens = response.usage.prompt_tokens
completion_tokens = response.usage.completion_tokens
total_tokens = response.usage.total_tokens
logger.info(f"🔢 Kimi API Token Usage - Prompt: {prompt_tokens}, Completion: {completion_tokens}, Total: {total_tokens}")
# Update trajectory
self.trajectory.last_prompt_tokens = prompt_tokens
self.trajectory.prompt_tokens_used += prompt_tokens
self.trajectory.completion_tokens_used += completion_tokens
self.trajectory.total_tokens_used += total_tokens
# Convert to dict format
message_dict = {
"role": "assistant",
"content": message.content
}
if hasattr(message, 'tool_calls') and message.tool_calls:
message_dict["tool_calls"] = [
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
}
for tc in message.tool_calls
]
# Display the response
if message.content:
print(f"\n🤖 Assistant: {message.content}\n")
return message_dict
def execute_research(self, max_iterations: int = 15) -> Dict[str, Any]:
"""
Execute the research task
Args:
max_iterations: Maximum number of tool calls
Returns:
Research results
"""
# Add initial user message
self.conversation_history.append({
"role": "user",
"content": "Please research and find the current affiliations of all OpenAI co-founders."
})
messages = self.conversation_history.copy()
iteration = 0
final_answer = None
print("\n" + "="*60)
print(f"Starting research with {self.compression_strategy.value} strategy")
print("="*60)
while iteration < max_iterations:
iteration += 1
print(f"\n📍 Iteration {iteration}/{max_iterations}")
try:
# Apply windowed compression if needed
if self.compression_strategy == CompressionStrategy.WINDOWED_CONTEXT:
messages = self._handle_windowed_compression(messages)
# Display current token usage from trajectory
print(f"📊 Cumulative Token Usage - Prompt: {self.trajectory.prompt_tokens_used:,}, Completion: {self.trajectory.completion_tokens_used:,}, Total: {self.trajectory.total_tokens_used:,}")
# Check if we're approaching token limit based on actual usage
if self.trajectory.total_tokens_used > 0: # Only check after first call
# Compression demo uses a 128k context budget. Compare the
# LAST call's prompt size (the actual context) against the
# window — the cumulative counter re-counts the shared
# prefix every call and overstates usage quadratically.
if self.trajectory.last_prompt_tokens > Config.CONTEXT_WINDOW_SIZE * 0.8:
logger.warning(f"Approaching context limit: {self.trajectory.last_prompt_tokens:,} prompt tokens in last request")
self.trajectory.context_overflows += 1
if self.compression_strategy == CompressionStrategy.NO_COMPRESSION:
print("\n⚠️ Context overflow detected! This demonstrates the limitation of no compression.")
return {
"error": f"Context window exceeded - {self.trajectory.last_prompt_tokens:,} tokens in last request (limit: {Config.CONTEXT_WINDOW_SIZE})",
"trajectory": self.trajectory,
"iterations": iteration
}
# Get response from model
if self.enable_streaming:
message = self._stream_response(messages)
else:
message = self._non_streaming_response(messages)
# Handle tool calls
if message.get('tool_calls'):
messages.append(message)
if message.get('content'):
print(f"\n🤖 Assistant: {message['content']}")
for tool_call in message['tool_calls']:
function_name = tool_call['function']['name']
raw_args = tool_call['function'].get('arguments') or "{}"
try:
if isinstance(raw_args, dict):
function_args = raw_args
elif isinstance(raw_args, (bytes, bytearray)):
function_args = json.loads(raw_args.decode("utf-8"))
elif isinstance(raw_args, str):
function_args = json.loads(raw_args)
else:
function_args = json.loads(str(raw_args))
if not isinstance(function_args, dict):
logger.warning(
"Tool argument JSON is not an object, proceeding with empty object: %r",
raw_args,
)
function_args = {}
except (json.JSONDecodeError, TypeError, UnicodeDecodeError):
# Tolerate bad tool-arg JSON; keep the loop alive.
function_args = {}
logger.warning(
"Tool argument is not valid JSON, proceeding with empty object: %r",
raw_args,
)
print(f"\n🔧 Executing: {function_name}")
print(f" Args: {function_args}")
# Execute the tool
result, compressed = self._execute_tool(function_name, function_args)
# Record the tool call
tool_call_record = ToolCall(
tool_name=function_name,
arguments=function_args,
result=result,
compressed_result=compressed,
id=tool_call['id']
)
self.trajectory.tool_calls.append(tool_call_record)
# Determine what content to add to messages
if compressed and self.compression_strategy != CompressionStrategy.NO_COMPRESSION:
# Use compressed content
tool_content = compressed.content
print(f" ✂️ Compressed: {compressed.original_length:,} → {compressed.compressed_length:,} chars")
else:
# Use original content (for no compression or last message in windowed)
if function_name == "search_web":
# Format search results
tool_content = json.dumps(result, indent=2)
else:
tool_content = json.dumps(result)
# Add tool result to messages
tool_msg = {
"role": "tool",
"tool_call_id": tool_call['id'],
"content": tool_content
}
messages.append(tool_msg)
print(f" 📄 Result size: {len(tool_content):,} characters")
elif message.get('content'):
# No tool calls, just content
messages.append(message)
final_answer = message['content']
logger.info("Final answer found")
break
except Exception as e:
logger.error(f"Error during research: {str(e)}")
return {
"error": str(e),
"trajectory": self.trajectory,
"iterations": iteration
}
# Set end time
self.trajectory.end_time = time.time()
return {
"final_answer": final_answer,
"trajectory": self.trajectory,
"iterations": iteration,
"success": final_answer is not None,
"execution_time": self.trajectory.end_time - self.trajectory.start_time
}
def reset(self):
"""Reset the agent's state"""
self.trajectory = AgentTrajectory(compression_strategy=self.compression_strategy)
self._init_system_prompt()
self.web_tools.clear_cache()
logger.info("Agent state reset")
chapter2/context-compression/benchmark_compression.py
+339
View File
@@ -0,0 +1,339 @@
"""
Context Compression Benchmark Module.
Systematically benchmarks Summary, Truncation, Key-Sentence, and Observation-Filtering
compression strategies on long-context tasks. Measures compression ratio, Time-to-First-Token (TTFT),
token cost savings, and downstream QA retention accuracy.
"""
import math
import re
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union, Tuple
def count_tokens(text: str) -> int:
"""Estimate token count for a given text string.
Uses tiktoken if available, with a reliable character/word-based fallback.
"""
if not text:
return 0
try:
import tiktoken
try:
encoding = tiktoken.encoding_for_model("gpt-4")
except Exception:
encoding = tiktoken.get_encoding("cl100k_base")
return len(encoding.encode(text))
except Exception:
# Fallback estimation: ~4 chars per token or ~0.75 words per token
words = len(text.split())
chars = len(text)
return max(1, int((words * 1.3 + chars / 4) / 2))
@dataclass
class StrategyMetrics:
"""Performance metrics for a context compression strategy."""
strategy: str
original_tokens: int
compressed_tokens: int
compression_ratio: float # compressed_tokens / original_tokens
ttft_ms: float # Time to first token in milliseconds
token_cost_savings: float # Cost savings ratio (0.0 to 1.0)
qa_retention_accuracy: float # Downstream QA accuracy (0.0 to 1.0)
def to_dict(self) -> Dict[str, Any]:
"""Convert metrics to a standard dictionary representation."""
return {
"strategy": self.strategy,
"original_tokens": self.original_tokens,
"compressed_tokens": self.compressed_tokens,
"compression_ratio": self.compression_ratio,
"ttft_ms": self.ttft_ms,
"token_cost_savings": self.token_cost_savings,
"qa_retention_accuracy": self.qa_retention_accuracy,
}
class ContextCompressionBenchmark:
"""Benchmark harness for evaluating context compression strategies."""
STRATEGIES = ["summary", "truncation", "key_sentence", "observation_filtering"]
def __init__(
self,
base_ttft_ms: float = 50.0,
per_token_ttft_ms: float = 0.05,
token_cost_per_1k: float = 0.0015,
target_max_tokens: int = 500,
):
"""Initialize the benchmark suite with configurable performance parameters."""
self.base_ttft_ms = base_ttft_ms
self.per_token_ttft_ms = per_token_ttft_ms
self.token_cost_per_1k = token_cost_per_1k
self.target_max_tokens = target_max_tokens
def compress_summary(self, context: str, query: str = "") -> str:
"""Summary Strategy: Condenses context into key abstract points."""
if not context:
return ""
sentences = [s.strip() for s in re.split(r'(?<=[.!?])\s+', context) if s.strip()]
if not sentences:
return context
if len(sentences) <= 3:
return context
# Extract beginning, middle, and end sentences to form a concise summary
step = max(1, len(sentences) // 3)
summary_sentences = [sentences[0]]
if step < len(sentences):
summary_sentences.append(sentences[step])
if len(sentences) - 1 > step:
summary_sentences.append(sentences[-1])
return " ".join(summary_sentences)
def compress_truncation(self, context: str, max_tokens: Optional[int] = None) -> str:
"""Truncation Strategy: Slices context to fit within strict token limits."""
if not context:
return ""
limit = self.target_max_tokens if max_tokens is None else max_tokens
if limit <= 0:
return ""
words = context.split()
if not words:
# No whitespace-separated words (e.g. CJK text): truncate by characters.
# CJK characters are roughly 1-2 tokens each, so use a conservative 1:1 ratio.
return context[:limit]
# Estimate max words corresponding to limit tokens (~0.75 words per token)
max_words = max(1, int(limit * 0.75))
truncated_words = words[:max_words]
return " ".join(truncated_words)
def compress_key_sentence(self, context: str, query: str = "") -> str:
"""Key-Sentence Strategy: Retains sentences with high query term match/relevance."""
if not context:
return ""
query = query or ""
sentences = [s.strip() for s in re.split(r'(?<=[.!?])\s+', context) if s.strip()]
if not sentences:
return context
if not query:
# Fallback to sentence length / position scoring if query is empty
scored = sorted(enumerate(sentences), key=lambda x: len(x[1]), reverse=True)
top_indices = sorted([idx for idx, _ in scored[:max(1, len(sentences) // 2)]])
return " ".join([sentences[i] for i in top_indices])
query_terms = set(re.findall(r'\w+', query.lower()))
scored_sentences = []
for idx, sentence in enumerate(sentences):
sentence_terms = set(re.findall(r'\w+', sentence.lower()))
overlap = len(query_terms.intersection(sentence_terms))
scored_sentences.append((overlap, idx, sentence))
# Sort by overlap descending, then by original position
scored_sentences.sort(key=lambda x: (-x[0], x[1]))
# Keep top half of sentences or those with overlap > 0
keep_count = max(1, math.ceil(len(sentences) * 0.5))
selected = scored_sentences[:keep_count]
# Sort selected back into original context order
selected.sort(key=lambda x: x[1])
return " ".join([s[2] for s in selected])
def compress_observation_filtering(self, context: str) -> str:
"""Observation-Filtering Strategy: Removes verbose system output, logs, hex, and JSON blobs."""
if not context:
return ""
lines = context.splitlines()
filtered_lines = []
for line in lines:
stripped = line.strip()
# Filter out JSON-like blobs, long hex hashes, trace logs, or repetitive debug markers
if (
re.match(r'^\s*[\{\[\}\]].*$', stripped) or
re.search(r'\b[0-9a-fA-F]{32,64}\b', stripped) or
re.search(r'^\s*(DEBUG|TRACE|INFO|VERBOSE)\b', stripped, re.IGNORECASE) or
re.search(r'^\s*<.*?>\s*$', stripped)
):
continue
filtered_lines.append(line)
result = "\n".join(filtered_lines).strip()
return result if result else context
def compress(self, strategy: str, context: str, query: str = "") -> str:
"""Apply a specific compression strategy to a given context string."""
strat = strategy.lower().replace("-", "_")
if strat == "summary":
return self.compress_summary(context, query)
elif strat == "truncation":
return self.compress_truncation(context)
elif strat in ("key_sentence", "keysentence"):
return self.compress_key_sentence(context, query)
elif strat in ("observation_filtering", "observationfiltering"):
return self.compress_observation_filtering(context)
else:
raise ValueError(f"Unknown compression strategy: {strategy}")
def evaluate_retention(self, compressed_text: str, task: Union[str, Dict[str, Any]]) -> Optional[float]:
"""Evaluate downstream QA retention accuracy on compressed context."""
compressed_text = compressed_text or ""
task = task or ""
query = task if isinstance(task, str) else (task.get("query", "") if isinstance(task, dict) else "")
expected = task.get("expected_answer", "") if isinstance(task, dict) else ""
if query is None:
query = ""
if expected is None:
expected = ""
# Only score against the expected answer, not the query.
# Using query words as fallback inflates scores because the question
# text often survives compression even when the answer is deleted.
target_text = expected.strip()
target_tokens = set(re.findall(r'\w+', target_text.lower()))
if not target_tokens:
# No expected answer to check against: cannot evaluate retention.
return None
compressed_tokens = set(re.findall(r'\w+', compressed_text.lower()))
matched = target_tokens.intersection(compressed_tokens)
# Calculate recall accuracy
accuracy = len(matched) / len(target_tokens)
return min(1.0, max(0.0, accuracy))
def evaluate_strategy(
self,
strategy: str,
contexts: List[str],
tasks: List[Union[str, Dict[str, Any]]],
) -> StrategyMetrics:
"""Benchmark a single compression strategy over multiple contexts and tasks."""
total_orig_tokens = 0
total_comp_tokens = 0
total_retention_acc = 0.0
retention_count = 0
sample_count = 0
start_time = time.perf_counter()
for idx, ctx in enumerate(contexts):
task = tasks[idx % len(tasks)] if tasks else ""
if task is None:
task = ""
query = task if isinstance(task, str) else (task.get("query", "") if isinstance(task, dict) else "")
query = query or ""
orig_tokens = count_tokens(ctx)
compressed_ctx = self.compress(strategy, ctx, query=query)
comp_tokens = count_tokens(compressed_ctx)
retention_acc = self.evaluate_retention(compressed_ctx, task)
total_orig_tokens += orig_tokens
total_comp_tokens += comp_tokens
if retention_acc is not None:
total_retention_acc += retention_acc
retention_count += 1
sample_count += 1
elapsed_ms = (time.perf_counter() - start_time) * 1000
avg_orig_tokens = total_orig_tokens / max(1, sample_count)
avg_comp_tokens = total_comp_tokens / max(1, sample_count)
avg_retention_acc = total_retention_acc / max(1, retention_count)
if avg_orig_tokens == 0:
ratio = 0.0
savings = 0.0
else:
ratio = avg_comp_tokens / avg_orig_tokens
savings = max(0.0, 1.0 - ratio)
# Simulate TTFT: Base TTFT + processing time + prefill latency based on compressed tokens
simulated_ttft = self.base_ttft_ms + (avg_comp_tokens * self.per_token_ttft_ms) + (elapsed_ms / max(1, sample_count))
# Format normalized strategy key
strat_key = strategy.lower().replace("-", "_")
return StrategyMetrics(
strategy=strat_key,
original_tokens=int(avg_orig_tokens),
compressed_tokens=int(avg_comp_tokens),
compression_ratio=round(ratio, 4),
ttft_ms=round(simulated_ttft, 2),
token_cost_savings=round(savings, 4),
qa_retention_accuracy=round(avg_retention_acc, 4),
)
def run_benchmark(
self,
contexts: Union[str, List[Union[str, Dict[str, Any]]]],
tasks: Union[str, List[Union[str, Dict[str, Any]]]],
) -> Dict[str, Any]:
"""Run systematic benchmark across all compression strategies.
Args:
contexts: Single context string, dict, or list of context strings/dicts.
tasks: Single task/query string, dict, or list of tasks/queries.
Returns:
Comparative metrics dictionary mapping strategy names to performance metrics dicts.
"""
# Standardize contexts into list of text strings
if isinstance(contexts, (str, dict)):
raw_contexts = [contexts]
else:
raw_contexts = list(contexts)
normalized_contexts = []
for c in raw_contexts:
if isinstance(c, str):
normalized_contexts.append(c)
elif isinstance(c, dict):
content = c.get("content")
if content is None:
content = c.get("text")
# Use the extracted content, or empty string if none found.
# Falling back to str(c) would treat the raw dict repr as
# context text, producing nonsensical benchmark metrics.
normalized_contexts.append(content if content is not None else "")
else:
normalized_contexts.append(str(c))
# Standardize tasks into list of queries/task objects
if isinstance(tasks, (str, dict)):
normalized_tasks = [tasks]
else:
normalized_tasks = list(tasks)
results: Dict[str, Any] = {}
for strategy in self.STRATEGIES:
metrics = self.evaluate_strategy(strategy, normalized_contexts, normalized_tasks)
metrics_dict = metrics.to_dict()
display_name = {
"summary": "Summary",
"truncation": "Truncation",
"key_sentence": "Key-Sentence",
"observation_filtering": "Observation-Filtering",
}.get(strategy, strategy)
metrics_dict["display_name"] = display_name
results[strategy] = metrics_dict
return results
def run_benchmark(
contexts: Union[str, List[Union[str, Dict[str, Any]]]],
tasks: Union[str, List[Union[str, Dict[str, Any]]]],
) -> Dict[str, Any]:
"""Module-level entrypoint for executing the compression benchmark.
Args:
contexts: Input contexts (strings or dicts).
tasks: Downstream QA tasks or queries.
Returns:
Dictionary of comparative performance metrics per compression strategy.
"""
benchmark = ContextCompressionBenchmark()
return benchmark.run_benchmark(contexts, tasks)
chapter2/context-compression/compression_strategies.py
+694
View File
@@ -0,0 +1,694 @@
"""
Context Compression Strategies for the experiment
"""
import json
import logging
from typing import List, Dict, Any, Optional, Tuple
from enum import Enum
from dataclasses import dataclass, field
from datetime import datetime
from openai import OpenAI
import tiktoken
from config import Config
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
def _reasoning_safe_max_tokens(model, requested, reasoning_budget=2048):
"""Reasoning models (Kimi K3, GPT-5, ...) spend part of the max_tokens
budget on reasoning_content *before* emitting the visible answer. If we
pass only the summary budget (e.g. 300-500), the reasoning trace can eat
into it and the summary comes back truncated or empty. Give reasoning
models extra headroom so the requested output budget is fully available
for the summary itself; non-reasoning models are unchanged."""
m = str(model or "").lower().replace("/", "-")
if "kimi-k3" in m or "gpt-5" in m:
return requested + reasoning_budget
return requested
# Configure logging
logging.basicConfig(level=logging.INFO, format=Config.LOG_FORMAT)
logger = logging.getLogger(__name__)
class CompressionStrategy(Enum):
"""Different context compression strategies"""
NO_COMPRESSION = "no_compression"
NON_CONTEXT_AWARE_INDIVIDUAL = "non_context_aware_individual_summary" # Summarize each page individually then concat
NON_CONTEXT_AWARE_COMBINED = "non_context_aware_combined_summary" # Concat all pages then summarize once
CONTEXT_AWARE = "context_aware_summary"
CONTEXT_AWARE_CITATIONS = "context_aware_with_citations"
WINDOWED_CONTEXT = "windowed_context"
@dataclass
class CompressedContent:
"""Represents compressed content"""
original_length: int
compressed_length: int
content: str
citations: List[Dict[str, str]] = field(default_factory=list)
strategy: CompressionStrategy = CompressionStrategy.NO_COMPRESSION
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
class ContextCompressor:
"""Handles different context compression strategies"""
def __init__(self, strategy: CompressionStrategy, api_key: str, enable_streaming: bool = True):
"""
Initialize the context compressor
Args:
strategy: Compression strategy to use
api_key: API key for LLM
enable_streaming: Whether to enable streaming for summarization
"""
self.strategy = strategy
self.enable_streaming = enable_streaming
# Moonshot 官方 key 存在则直连；否则回退 OpenRouter（见 Config.resolve_llm）。
resolved_key, resolved_base_url, resolved_model = Config.resolve_llm()
self.client = OpenAI(
api_key=resolved_key,
base_url=resolved_base_url
)
self.model = resolved_model
# Initialize tokenizer for token counting
try:
self.encoding = tiktoken.encoding_for_model("gpt-4")
except Exception:
self.encoding = tiktoken.get_encoding("cl100k_base")
logger.info(f"Context compressor initialized with strategy: {strategy.value}, streaming: {enable_streaming}")
def count_tokens(self, text: str) -> int:
"""Count the number of tokens in a text string."""
try:
return len(self.encoding.encode(text))
except Exception:
# Fallback to character-based estimation (1 token ≈ 4 chars)
return len(text) // 4
def compress_search_results(
self,
search_results: Dict[str, Any],
query: str,
current_context: Optional[str] = None
) -> CompressedContent:
"""
Compress search results based on the selected strategy
Args:
search_results: Raw search results from web tool
query: The original search query
current_context: Current conversation context (for context-aware strategies)
Returns:
Compressed content
"""
if self.strategy == CompressionStrategy.NO_COMPRESSION:
return self._no_compression(search_results)
elif self.strategy == CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL:
return self._non_context_aware_individual_summary(search_results)
elif self.strategy == CompressionStrategy.NON_CONTEXT_AWARE_COMBINED:
return self._non_context_aware_combined_summary(search_results)
elif self.strategy == CompressionStrategy.CONTEXT_AWARE:
return self._context_aware_summary(search_results, query, current_context)
elif self.strategy == CompressionStrategy.CONTEXT_AWARE_CITATIONS:
return self._context_aware_with_citations(search_results, query, current_context)
elif self.strategy == CompressionStrategy.WINDOWED_CONTEXT:
# For windowed context, return full content (compression happens later)
return self._no_compression(search_results)
else:
raise ValueError(f"Unknown compression strategy: {self.strategy}")
def compress_for_history(
self,
content: str,
tool_name: str,
query: str,
preserve_citations: bool = True
) -> CompressedContent:
"""
Compress content for message history (used in windowed context strategy)
Args:
content: Content to compress
tool_name: Name of the tool that generated the content
query: The query that triggered the tool call
preserve_citations: Whether to preserve citations
Returns:
Compressed content for history
"""
original_length = len(content)
try:
prompt = f"""Compress the following {tool_name} results into a concise summary that preserves key information.
Focus on information relevant to: {query}
Original content:
{content[:10000]}
Requirements:
1. Keep all important facts, names, dates, and affiliations
2. Remove redundant information
3. Maintain clarity and coherence
{"4. Include [Source: URL] citations for important facts" if preserve_citations else ""}
5. Maximum length: {Config.SUMMARY_MAX_TOKENS} tokens
Provide a focused summary:"""
# Log prompt length
prompt_tokens = self.count_tokens(prompt)
logger.info(f"Simple summary request - Prompt tokens: {prompt_tokens}, Prompt length: {len(prompt)} chars")
if self.enable_streaming:
# Stream the summary to console
print(f"\n📝 Creating simple summary...\n", flush=True)
stream = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates concise summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS),
stream=True
)
summary_parts = []
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
# NB: do not name this `content` — that shadows the
# `content` parameter (the original tool output) and
# breaks the truncation fallback below when the stream
# fails part-way through.
delta_text = chunk.choices[0].delta.content
print(delta_text, end="", flush=True)
summary_parts.append(delta_text)
print("\n") # New lines after streaming
compressed = "".join(summary_parts)
else:
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates concise summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS)
)
compressed = response.choices[0].message.content
return CompressedContent(
original_length=original_length,
compressed_length=len(compressed),
content=compressed,
strategy=CompressionStrategy.WINDOWED_CONTEXT
)
except Exception as e:
logger.error(f"Error compressing for history: {str(e)}")
# Fallback to truncation
truncated = content[:2000] + "\n\n[Content truncated for history...]"
return CompressedContent(
original_length=original_length,
compressed_length=len(truncated),
content=truncated,
strategy=CompressionStrategy.WINDOWED_CONTEXT
)
def _no_compression(self, search_results: Dict[str, Any]) -> CompressedContent:
"""
Strategy 1: No compression - return all original content
"""
all_content = []
total_length = 0
for result in search_results.get('results', []):
content = f"""
===== Search Result =====
Title: {result.get('title', 'N/A')}
URL: {result.get('url', 'N/A')}
Snippet: {result.get('snippet', 'N/A')}
Full Content:
{result.get('content', 'No content available')}
========================
"""
all_content.append(content)
total_length += len(result.get('content') or '')
full_content = "\n\n".join(all_content)
return CompressedContent(
original_length=total_length,
compressed_length=len(full_content),
content=full_content,
strategy=CompressionStrategy.NO_COMPRESSION
)
def _non_context_aware_individual_summary(self, search_results: Dict[str, Any]) -> CompressedContent:
"""
Strategy 2A: Non-context-aware summarization - Summarize each page individually then concatenate
"""
summaries = []
total_original = 0
for result in search_results.get('results', []):
if not result.get('content'):
continue
original_content = result.get('content', '')
total_original += len(original_content)
try:
# Summarize each page independently
prompt = f"""Summarize the following webpage content in 2-3 paragraphs:
Title: {result.get('title', 'N/A')}
URL: {result.get('url', 'N/A')}
Content:
{original_content[:5000]}
Provide a concise summary:"""
# Log prompt length
prompt_tokens = self.count_tokens(prompt)
logger.info(f"Non-context-aware summary - Prompt tokens: {prompt_tokens}, Prompt length: {len(prompt)} chars")
if self.enable_streaming:
# Stream the summary to console
print(f"\n📝 Summarizing: {result.get('title', 'N/A')[:50]}...", end=" ", flush=True)
stream = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates concise summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, 300),
stream=True
)
summary_parts = []
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
content = chunk.choices[0].delta.content
print(content, end="", flush=True)
summary_parts.append(content)
print() # New line after streaming
summary = "".join(summary_parts)
else:
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates concise summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, 300)
)
summary = response.choices[0].message.content
summaries.append(f"""
Source: {result.get('title', 'N/A')}
URL: {result.get('url', 'N/A')}
Summary: {summary}
""")
except Exception as e:
logger.error(f"Error summarizing page: {str(e)}")
# Fallback to snippet
summaries.append(f"""
Source: {result.get('title', 'N/A')}
URL: {result.get('url', 'N/A')}
Summary: {result.get('snippet', 'No summary available')}
""")
compressed_content = "\n".join(summaries)
return CompressedContent(
original_length=total_original,
compressed_length=len(compressed_content),
content=compressed_content,
strategy=CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL
)
def _non_context_aware_combined_summary(self, search_results: Dict[str, Any]) -> CompressedContent:
"""
Strategy 2B: Non-context-aware summarization - Concatenate all pages then summarize once
"""
# Combine all content first
all_content = []
total_original = 0
max_chars_per_page = 5000 # Limit each page to prevent token overflow
for result in search_results.get('results', []):
if result.get('content'):
original_content = result.get('content', '')
total_original += len(original_content)
# Limit each page's content
limited_content = original_content[:max_chars_per_page]
all_content.append(f"""
===== Page: {result.get('title', 'N/A')} =====
URL: {result.get('url', 'N/A')}
Content: {limited_content}
""")
if not all_content:
return CompressedContent(
original_length=0,
compressed_length=0,
content="No content available",
strategy=CompressionStrategy.NON_CONTEXT_AWARE_COMBINED
)
combined_content = "\n\n".join(all_content)
try:
# Create a single summary for all combined content
prompt = f"""Summarize the following combined webpage content comprehensively:
{combined_content}
Requirements:
1. Create a comprehensive summary covering all pages
2. Include key information from each source
3. Maintain factual accuracy
4. Maximum length: {Config.SUMMARY_MAX_TOKENS} tokens
Provide a comprehensive summary:"""
# Log prompt length
prompt_tokens = self.count_tokens(prompt)
logger.info(f"Non-context-aware combined summary - Prompt tokens: {prompt_tokens}, Prompt length: {len(prompt)} chars")
if self.enable_streaming:
# Stream the summary to console
print(f"\n📄 Creating combined summary for all {len(search_results.get('results', []))} pages...\n", flush=True)
stream = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates comprehensive summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS),
stream=True
)
summary_parts = []
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
content = chunk.choices[0].delta.content
print(content, end="", flush=True)
summary_parts.append(content)
print("\n") # New lines after streaming
summary = "".join(summary_parts)
else:
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates comprehensive summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS)
)
summary = response.choices[0].message.content
return CompressedContent(
original_length=total_original,
compressed_length=len(summary),
content=summary,
strategy=CompressionStrategy.NON_CONTEXT_AWARE_COMBINED
)
except Exception as e:
logger.error(f"Error creating combined summary: {str(e)}")
# Fallback to concatenated snippets
fallback = "\n\n".join([
f"{r.get('title', 'N/A')}: {r.get('snippet', 'No summary available')}"
for r in search_results.get('results', [])
])
return CompressedContent(
original_length=total_original,
compressed_length=len(fallback),
content=fallback,
strategy=CompressionStrategy.NON_CONTEXT_AWARE_COMBINED
)
def _context_aware_summary(
self,
search_results: Dict[str, Any],
query: str,
current_context: Optional[str] = None
) -> CompressedContent:
"""
Strategy 3: Context-aware summarization considering the query
"""
# Combine all content with per-page limits
all_content = []
total_original = 0
max_chars_per_page = 5000 # Limit each page to prevent token overflow
for result in search_results.get('results', []):
if result.get('content'):
original_content = result.get('content', '')
total_original += len(original_content)
# Limit each page's content
limited_content = original_content[:max_chars_per_page]
all_content.append(f"""
Title: {result.get('title', 'N/A')}
URL: {result.get('url', 'N/A')}
Content: {limited_content}
""")
combined_content = "\n\n".join(all_content)
try:
# Create context-aware summary
prompt = f"""Given the search query: "{query}"
{f"Current context: {current_context[:1000]}" if current_context else ""}
Analyze the following search results and provide a focused summary that directly addresses the query.
Focus on extracting information most relevant to answering: {query}
Search Results:
{combined_content}
Requirements:
1. Focus only on information relevant to the query
2. Prioritize current/recent information
3. Include specific names, dates, and affiliations
4. Maximum length: {Config.SUMMARY_MAX_TOKENS} tokens
Provide a query-focused summary:"""
# Log prompt length
prompt_tokens = self.count_tokens(prompt)
logger.info(f"Context-aware summary - Prompt tokens: {prompt_tokens}, Prompt length: {len(prompt)} chars")
if self.enable_streaming:
# Stream the summary to console
print(f"\n🎯 Creating context-aware summary for query: '{query[:50]}...'\n", flush=True)
stream = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates focused, context-aware summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS),
stream=True
)
summary_parts = []
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
content = chunk.choices[0].delta.content
print(content, end="", flush=True)
summary_parts.append(content)
print("\n") # New lines after streaming
summary = "".join(summary_parts)
else:
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates focused, context-aware summaries."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS)
)
summary = response.choices[0].message.content
return CompressedContent(
original_length=total_original,
compressed_length=len(summary),
content=summary,
strategy=CompressionStrategy.CONTEXT_AWARE
)
except Exception as e:
logger.error(f"Error creating context-aware summary: {str(e)}")
# Fallback to simple concatenation
fallback = "\n\n".join([r.get('snippet', '') for r in search_results.get('results', [])])
return CompressedContent(
original_length=total_original,
compressed_length=len(fallback),
content=fallback,
strategy=CompressionStrategy.CONTEXT_AWARE
)
def _context_aware_with_citations(
self,
search_results: Dict[str, Any],
query: str,
current_context: Optional[str] = None
) -> CompressedContent:
"""
Strategy 4: Context-aware summarization with citations
"""
# Track sources with per-page limits
sources = []
all_content = []
total_original = 0
max_chars_per_page = 5000 # Limit each page to prevent token overflow
for i, result in enumerate(search_results.get('results', [])):
if result.get('content'):
source_id = f"[{i+1}]"
original_content = result.get('content', '')
total_original += len(original_content)
# Limit each page's content
limited_content = original_content[:max_chars_per_page]
sources.append({
'id': source_id,
'title': result.get('title', 'N/A'),
'url': result.get('url', 'N/A')
})
all_content.append(f"""
{source_id} Title: {result.get('title', 'N/A')}
Content: {limited_content}
""")
combined_content = "\n\n".join(all_content)
try:
# Create context-aware summary with citations
prompt = f"""Given the search query: "{query}"
{f"Current context: {current_context[:1000]}" if current_context else ""}
Analyze the following search results and provide a focused summary with citations.
Search Results (with source IDs):
{combined_content}
Requirements:
1. Focus on information relevant to: {query}
2. Include inline citations using [1], [2], etc. for each fact
3. Prioritize current/recent information
4. Include specific names, dates, and affiliations with citations
5. Maximum length: {Config.SUMMARY_MAX_TOKENS} tokens
Provide a query-focused summary with citations:"""
# Log prompt length
prompt_tokens = self.count_tokens(prompt)
logger.info(f"Citation-based summary - Prompt tokens: {prompt_tokens}, Prompt length: {len(prompt)} chars")
if self.enable_streaming:
# Stream the summary to console
print(f"\n📚 Creating summary with citations for: '{query[:50]}...'\n", flush=True)
stream = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates focused summaries with proper citations."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS),
stream=True
)
summary_parts = []
for chunk in stream:
if chunk.choices and chunk.choices[0].delta.content:
content = chunk.choices[0].delta.content
print(content, end="", flush=True)
summary_parts.append(content)
print("\n") # New lines after streaming
summary = "".join(summary_parts)
else:
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "system", "content": "You are a helpful assistant that creates focused summaries with proper citations."},
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=_reasoning_safe_max_tokens(self.model, Config.SUMMARY_MAX_TOKENS)
)
summary = response.choices[0].message.content
# Append source list
source_list = "\n\nSources:\n"
for source in sources:
source_list += f"{source['id']} {source['title']} - {source['url']}\n"
final_content = summary + source_list
return CompressedContent(
original_length=total_original,
compressed_length=len(final_content),
content=final_content,
citations=sources,
strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS
)
except Exception as e:
logger.error(f"Error creating summary with citations: {str(e)}")
# Fallback
fallback = "\n\n".join([
f"[{i+1}] {r.get('title', '')}: {r.get('snippet', '')}"
for i, r in enumerate(search_results.get('results', []))
])
return CompressedContent(
original_length=total_original,
compressed_length=len(fallback),
content=fallback,
citations=sources,
strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS
)
def estimate_tokens(self, text: str) -> int:
"""
Estimate token count for text (rough approximation)
Args:
text: Text to estimate tokens for
Returns:
Estimated token count
"""
# Rough approximation: 1 token ≈ 4 characters
return len(text) // 4
chapter2/context-compression/config.py
+121
View File
@@ -0,0 +1,121 @@
"""
Configuration module for Context Compression Experiment
"""
import os
from typing import Optional
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
class Config:
"""Configuration settings for the context compression experiment"""
# API Configuration
LLM_PROVIDER: str = os.getenv("LLM_PROVIDER", "kimi").lower()
LLM_PROVIDER = {"qwen": "dashscope", "bailian": "dashscope"}.get(
LLM_PROVIDER, LLM_PROVIDER
)
DASHSCOPE_API_KEY: str = os.getenv("DASHSCOPE_API_KEY", "")
DASHSCOPE_BASE_URL: str = os.getenv(
"DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1"
)
MOONSHOT_API_KEY: str = os.getenv("MOONSHOT_API_KEY", "")
MOONSHOT_BASE_URL: str = "https://api.moonshot.cn/v1"
# Universal fallback: 当 MOONSHOT_API_KEY 缺失但设置了 OPENROUTER_API_KEY 时，
# 自动改走 OpenRouter（kimi-* 模型名映射为 moonshotai/kimi-k2）。
OPENROUTER_API_KEY: str = os.getenv("OPENROUTER_API_KEY", "")
SERPER_API_KEY: str = os.getenv("SERPER_API_KEY", "")
SERPER_BASE_URL: str = "https://google.serper.dev"
# Model Configuration
MODEL_NAME: str = os.getenv(
"MODEL_NAME", "qwen3.7-plus" if LLM_PROVIDER == "dashscope" else "kimi-k3"
)
MODEL_TEMPERATURE: float = float(os.getenv("MODEL_TEMPERATURE", "0.3"))
MODEL_MAX_TOKENS: int = int(os.getenv("MODEL_MAX_TOKENS", "8192"))
# Agent Configuration
MAX_ITERATIONS: int = int(os.getenv("MAX_ITERATIONS", "50"))
ENABLE_VERBOSE: bool = os.getenv("ENABLE_VERBOSE", "false").lower() == "true"
# Compression Configuration
MAX_WEBPAGE_LENGTH: int = int(os.getenv("MAX_WEBPAGE_LENGTH", "50000"))
SUMMARY_MAX_TOKENS: int = int(os.getenv("SUMMARY_MAX_TOKENS", "500"))
# Context Window Configuration
CONTEXT_WINDOW_SIZE: int = 128000 # 128K context budget for the compression demo (K3 supports up to 1M)
# Logging Configuration
LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO")
LOG_FORMAT: str = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
# File paths
RESULTS_DIR: str = "results"
CACHE_DIR: str = "cache"
@classmethod
def validate(cls) -> bool:
"""
Validate required configuration
Returns:
True if configuration is valid
"""
try:
cls.resolve_llm()
except ValueError as exc:
print(f"ERROR: {exc}")
return False
if not cls.SERPER_API_KEY:
print("WARNING: SERPER_API_KEY is not set")
print("Web search functionality will be limited")
print("Get a free API key at: https://serper.dev")
return True
@classmethod
def resolve_llm(cls):
"""Return ``(api_key, base_url, model)`` for the configured provider.
Computed at call time so a runtime override of ``Config.MODEL_NAME``
(e.g. via ``--model``) is respected.
端点、接受的 key 变量与模型名映射由 agentbook 的 provider 注册表统一
维护。此处保持三元组返回值：调用方按 3 个字段解包，测试也按这个形状
打桩。
"""
from agentbook.providers import resolve_backend
backend = resolve_backend(cls.LLM_PROVIDER, model=cls.MODEL_NAME)
return backend.api_key, backend.base_url, backend.model
@classmethod
def create_directories(cls):
"""Create necessary directories if they don't exist"""
os.makedirs(cls.RESULTS_DIR, exist_ok=True)
os.makedirs(cls.CACHE_DIR, exist_ok=True)
@classmethod
def print_config(cls):
"""Print current configuration (hiding sensitive data)"""
print("\n" + "="*50)
print("CONFIGURATION")
print("="*50)
print(f"Model: {cls.MODEL_NAME}")
print(f"Temperature: {cls.MODEL_TEMPERATURE}")
print(f"Max Tokens: {cls.MODEL_MAX_TOKENS}")
print(f"Max Iterations: {cls.MAX_ITERATIONS}")
print(f"Context Window: {cls.CONTEXT_WINDOW_SIZE:,} tokens")
print(f"Max Webpage Length: {cls.MAX_WEBPAGE_LENGTH:,} chars")
print(f"Summary Max Tokens: {cls.SUMMARY_MAX_TOKENS}")
print(f"Provider: {cls.LLM_PROVIDER}")
print(f"DashScope API Key Set: {'Yes' if cls.DASHSCOPE_API_KEY else 'No'}")
print(f"Kimi API Key Set: {'Yes' if cls.MOONSHOT_API_KEY else 'No'}")
print(f"Serper API Key Set: {'Yes' if cls.SERPER_API_KEY else 'No'}")
print("="*50 + "\n")
chapter2/context-compression/env.example
+29
View File
@@ -0,0 +1,29 @@
# LLM provider: kimi (default), dashscope/qwen/bailian, or openrouter
LLM_PROVIDER=kimi
# Kimi API Configuration
MOONSHOT_API_KEY=your_moonshot_api_key_here
# Alibaba Cloud Model Studio / Bailian (Qwen)
# DASHSCOPE_API_KEY=your_dashscope_api_key_here
# DASHSCOPE_BASE_URL=https://dashscope-intl.aliyuncs.com/compatible-mode/v1
# 通用回退：未设置 MOONSHOT_API_KEY 时，若配置了 OPENROUTER_API_KEY，则自动改走
# OpenRouter（kimi-* 会映射为 moonshotai/kimi-k2）。
# OPENROUTER_API_KEY=your-openrouter-api-key
# Search API Configuration (using Serper API - free tier available)
SERPER_API_KEY=your_serper_api_key_here
# Model Configuration
MODEL_NAME=kimi-k3
MODEL_TEMPERATURE=0.3
MODEL_MAX_TOKENS=8192
# Agent Configuration
MAX_ITERATIONS=15
ENABLE_VERBOSE=false
# Compression Configuration
MAX_WEBPAGE_LENGTH=50000
SUMMARY_MAX_TOKENS=500
chapter2/context-compression/experiment.py
Executable
+382
View File
@@ -0,0 +1,382 @@
#!/usr/bin/env python3
"""
Context Compression Strategies Comparison Experiment
"""
import os
import sys
import json
import time
import argparse
from typing import Dict, Any, List, Optional
from datetime import datetime
from dataclasses import asdict
from colorama import init, Fore, Style
from tqdm import tqdm
from config import Config
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
# Initialize colorama for colored output
init(autoreset=True)
# Short CLI aliases -> compression strategy (order matches the book's 实验 2-10)
STRATEGY_CHOICES = {
"no_compression": CompressionStrategy.NO_COMPRESSION,
"individual": CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
"combined": CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
"context_aware": CompressionStrategy.CONTEXT_AWARE,
"citations": CompressionStrategy.CONTEXT_AWARE_CITATIONS,
"windowed": CompressionStrategy.WINDOWED_CONTEXT,
}
ALL_STRATEGIES = list(STRATEGY_CHOICES.values())
class ExperimentRunner:
"""Runs experiments comparing different compression strategies"""
def __init__(self, api_key: str, results_file: Optional[str] = None,
enable_streaming: bool = False):
"""
Initialize the experiment runner
Args:
api_key: API key for Kimi/Moonshot
results_file: Optional explicit path for the results JSON (default: results/experiment_TIMESTAMP.json)
enable_streaming: Stream compression/model output to the console during the run
"""
self.api_key = api_key
self.results = []
self.enable_streaming = enable_streaming
# Create results directory
Config.create_directories()
# Results file
if results_file:
self.results_file = results_file
parent = os.path.dirname(self.results_file)
if parent:
os.makedirs(parent, exist_ok=True)
else:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.results_file = os.path.join(Config.RESULTS_DIR, f"experiment_{timestamp}.json")
def run_single_strategy(self, strategy: CompressionStrategy, verbose: bool = False) -> Dict[str, Any]:
"""
Run experiment with a single compression strategy
Args:
strategy: Compression strategy to test
verbose: Enable verbose output
Returns:
Experiment results
"""
print(f"\n{Fore.CYAN}{'='*70}")
print(f"{Fore.CYAN}Testing Strategy: {Fore.YELLOW}{strategy.value}")
print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
# Create agent with the strategy
agent = ResearchAgent(
api_key=self.api_key,
compression_strategy=strategy,
verbose=verbose,
enable_streaming=self.enable_streaming # Off by default for cleaner experiment output
)
start_time = time.time()
try:
# Execute the research task
result = agent.execute_research(max_iterations=Config.MAX_ITERATIONS)
end_time = time.time()
execution_time = end_time - start_time
# Analyze results
trajectory = result.get('trajectory')
# Calculate metrics
metrics = {
'strategy': strategy.value,
'success': result.get('success', False),
'iterations': result.get('iterations', 0),
'tool_calls': len(trajectory.tool_calls) if trajectory else 0,
'context_overflows': trajectory.context_overflows if trajectory else 0,
'execution_time': execution_time,
'total_tokens': trajectory.total_tokens_used if trajectory else 0,
'error': result.get('error'),
'final_answer_length': len(result.get('final_answer', '')) if result.get('final_answer') else 0
}
# Calculate compression ratios
if trajectory and trajectory.tool_calls:
total_original = 0
total_compressed = 0
for call in trajectory.tool_calls:
if call.compressed_result:
total_original += call.compressed_result.original_length
total_compressed += call.compressed_result.compressed_length
elif call.result and call.tool_name == 'search_web':
# No compression - count full size
content = json.dumps(call.result)
total_original += len(content)
total_compressed += len(content)
if total_original > 0:
metrics['compression_ratio'] = round(total_compressed / total_original, 3)
metrics['total_original_size'] = total_original
metrics['total_compressed_size'] = total_compressed
else:
metrics['compression_ratio'] = 1.0
metrics['total_original_size'] = 0
metrics['total_compressed_size'] = 0
# Print summary
self._print_summary(metrics)
# Store full result
full_result = {
'metrics': metrics,
'final_answer': result.get('final_answer'),
'timestamp': datetime.now().isoformat()
}
return full_result
except Exception as e:
print(f"{Fore.RED}Error during experiment: {str(e)}{Style.RESET_ALL}")
return {
'metrics': {
'strategy': strategy.value,
'success': False,
'error': str(e),
'execution_time': time.time() - start_time
},
'timestamp': datetime.now().isoformat()
}
def _print_summary(self, metrics: Dict[str, Any]):
"""Print a summary of the metrics"""
print(f"\n{Fore.GREEN}📊 Results Summary:{Style.RESET_ALL}")
print(f" Success: {self._format_bool(metrics['success'])}")
print(f" Iterations: {metrics['iterations']}")
print(f" Tool Calls: {metrics['tool_calls']}")
print(f" Execution Time: {metrics['execution_time']:.2f}s")
print(f" Total Tokens: {metrics.get('total_tokens', 0):,}")
if 'compression_ratio' in metrics:
print(f" Compression Ratio: {metrics['compression_ratio']:.1%}")
print(f" Original Size: {metrics['total_original_size']:,} chars")
print(f" Compressed Size: {metrics['total_compressed_size']:,} chars")
if metrics.get('context_overflows', 0) > 0:
print(f" {Fore.YELLOW}Context Overflows: {metrics['context_overflows']}{Style.RESET_ALL}")
if metrics.get('error'):
print(f" {Fore.RED}Error: {metrics['error'][:100]}...{Style.RESET_ALL}")
def _format_bool(self, value: bool) -> str:
"""Format boolean value with color"""
if value:
return f"{Fore.GREEN}✓ Yes{Style.RESET_ALL}"
else:
return f"{Fore.RED}✗ No{Style.RESET_ALL}"
def run_all_strategies(self, strategies: Optional[List[CompressionStrategy]] = None) -> None:
"""Run experiments for the given compression strategies (default: all six)"""
if strategies is None:
strategies = list(ALL_STRATEGIES)
print(f"\n{Fore.MAGENTA}{'='*70}")
print(f"{Fore.MAGENTA}CONTEXT COMPRESSION STRATEGIES COMPARISON EXPERIMENT")
print(f"{Fore.MAGENTA}{'='*70}{Style.RESET_ALL}")
print(f"\nTesting {len(strategies)} compression strategies...")
print(f"Task: Research current affiliations of OpenAI co-founders")
# Run each strategy
for strategy in tqdm(strategies, desc="Running experiments"):
result = self.run_single_strategy(strategy)
self.results.append(result)
# Save intermediate results
self._save_results()
# Small delay between experiments
time.sleep(2)
# Print final comparison
self._print_comparison()
def _save_results(self):
"""Save results to JSON file"""
with open(self.results_file, 'w') as f:
json.dump(self.results, f, indent=2, default=str)
print(f"\n💾 Results saved to: {self.results_file}")
def _print_comparison(self):
"""Print comparison table of all strategies"""
print(f"\n{Fore.MAGENTA}{'='*70}")
print(f"{Fore.MAGENTA}FINAL COMPARISON")
print(f"{Fore.MAGENTA}{'='*70}{Style.RESET_ALL}")
# Create comparison table
print(f"\n{'Strategy':<38} {'Success':<9} {'Time':<9} {'Tokens':<11} {'Compress':<10} {'Overflows':<10}")
print("-" * 90)
for result in self.results:
metrics = result['metrics']
strategy = metrics['strategy'][:36]
success = "✓" if metrics['success'] else "✗"
time_str = f"{metrics.get('execution_time', 0):.1f}s"
tokens = f"{metrics.get('total_tokens', 0):,}" if metrics.get('total_tokens') else "N/A"
compress = f"{metrics.get('compression_ratio', 1.0):.1%}" if 'compression_ratio' in metrics else "N/A"
overflows = str(metrics.get('context_overflows', 0))
# Color code success
color = Fore.GREEN if metrics['success'] else Fore.RED
print(f"{color}{strategy:<38} {success:<9} {time_str:<9} {tokens:<11} {compress:<10} {overflows:<10}{Style.RESET_ALL}")
print("\n" + "="*90)
# Analysis summary
self._print_analysis()
def _print_analysis(self):
"""Print analysis of the results"""
print(f"\n{Fore.CYAN}📈 Analysis:{Style.RESET_ALL}")
successful = [r for r in self.results if r['metrics']['success']]
failed = [r for r in self.results if not r['metrics']['success']]
print(f"\n Successful Strategies: {len(successful)}/{len(self.results)}")
if successful:
# Find best performing
fastest = min(successful, key=lambda x: x['metrics']['execution_time'])
most_efficient = min(successful, key=lambda x: x['metrics'].get('total_compressed_size', float('inf')))
print(f" Fastest: {fastest['metrics']['strategy']} ({fastest['metrics']['execution_time']:.1f}s)")
print(f" Most Efficient: {most_efficient['metrics']['strategy']} ({most_efficient['metrics'].get('total_compressed_size', 0):,} chars)")
if failed:
print(f"\n Failed Strategies:")
for r in failed:
# error may be present-but-None when a strategy fails by hitting the
# iteration cap (rather than raising), so coalesce before slicing.
err = r['metrics'].get('error') or 'No final answer within max iterations'
print(f" - {r['metrics']['strategy']}: {err[:50]}...")
# Key findings
print(f"\n{Fore.CYAN}🔍 Key Findings:{Style.RESET_ALL}")
print(" 1. No Compression: Expected to fail with context overflow ✓")
print(" 2. Non-Context-Aware: May lose important context details")
print(" 3. Context-Aware: Better relevance preservation")
print(" 4. With Citations: Enables follow-up questions")
print(" 5. Windowed Context: Balance between detail and efficiency")
def build_parser() -> argparse.ArgumentParser:
"""构建命令行参数解析器"""
parser = argparse.ArgumentParser(
prog="experiment.py",
description="上下文压缩策略对比实验（对应《深入理解 AI Agent》实验 2-10）。\n"
"对同一个研究任务（追踪 OpenAI 联合创始人的现状）分别运行多种压缩策略，"
"输出 token 用量 / 压缩率 / 成功率对比表，并保存 JSON 结果。",
epilog="示例：\n"
" python experiment.py # 运行全部 6 种策略并对比\n"
" python experiment.py -s context_aware # 只运行“上下文感知压缩”\n"
" python experiment.py -s individual combined # 只对比两种非任务感知策略\n"
" python experiment.py --model kimi-k3 -o results/k2.json\n"
" python experiment.py --list-strategies # 查看可选策略名",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"-s", "--strategy", nargs="+", choices=list(STRATEGY_CHOICES.keys()), metavar="NAME",
help="要运行的压缩策略（可指定多个，默认运行全部 6 种）。可选值："
+ ", ".join(STRATEGY_CHOICES.keys()),
)
parser.add_argument(
"-m", "--model", default=None,
help=f"覆盖使用的模型名称（默认读取环境变量 MODEL_NAME，当前为 {Config.MODEL_NAME}）",
)
parser.add_argument(
"-o", "--output", default=None, metavar="PATH",
help="结果 JSON 的保存路径（默认 results/experiment_<时间戳>.json）",
)
parser.add_argument(
"-n", "--max-iterations", type=int, default=None, metavar="N",
help=f"每个策略允许的最大迭代（工具调用轮数），默认 {Config.MAX_ITERATIONS}",
)
parser.add_argument(
"--streaming", action="store_true",
help="实时流式打印模型与压缩过程的输出（默认关闭，以获得更整洁的对比输出）",
)
parser.add_argument(
"--list-strategies", action="store_true",
help="列出所有可选的压缩策略名称后退出",
)
return parser
def main():
"""Main entry point"""
parser = build_parser()
args = parser.parse_args()
if args.list_strategies:
print("可选的压缩策略（--strategy 的取值）：")
for alias, strat in STRATEGY_CHOICES.items():
print(f" {alias:<16} -> {strat.value}")
return
# Apply CLI overrides onto the shared Config
if args.model:
Config.MODEL_NAME = args.model
if args.max_iterations is not None:
Config.MAX_ITERATIONS = args.max_iterations
# Resolve which strategies to run
if args.strategy:
strategies = [STRATEGY_CHOICES[name] for name in args.strategy]
else:
strategies = list(ALL_STRATEGIES)
# Check configuration
if not Config.validate():
print(f"\n{Fore.RED}Configuration validation failed!{Style.RESET_ALL}")
print("\nPlease set up your .env file with:")
print(" MOONSHOT_API_KEY=your_api_key_here")
print(" SERPER_API_KEY=your_api_key_here (optional)")
sys.exit(1)
# Print configuration
Config.print_config()
# Create runner
runner = ExperimentRunner(
Config.MOONSHOT_API_KEY,
results_file=args.output,
enable_streaming=args.streaming,
)
# Run experiments
try:
runner.run_all_strategies(strategies)
print(f"\n{Fore.GREEN}✅ Experiment completed successfully!{Style.RESET_ALL}")
except KeyboardInterrupt:
print(f"\n{Fore.YELLOW}⚠️ Experiment interrupted by user{Style.RESET_ALL}")
except Exception as e:
print(f"\n{Fore.RED}❌ Experiment failed: {str(e)}{Style.RESET_ALL}")
sys.exit(1)
if __name__ == "__main__":
main()
chapter2/context-compression/main.py
Executable
+231
View File
@@ -0,0 +1,231 @@
#!/usr/bin/env python3
"""
Interactive demo for context compression strategies
"""
import os
import sys
import argparse
from colorama import init, Fore, Style
from config import Config
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
# Initialize colorama
init(autoreset=True)
# Short CLI aliases -> compression strategy (order matches the book's 实验 2-10)
STRATEGY_CHOICES = {
"no_compression": CompressionStrategy.NO_COMPRESSION,
"individual": CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
"combined": CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
"context_aware": CompressionStrategy.CONTEXT_AWARE,
"citations": CompressionStrategy.CONTEXT_AWARE_CITATIONS,
"windowed": CompressionStrategy.WINDOWED_CONTEXT,
}
def print_banner():
"""Print demo banner"""
print(f"\n{Fore.CYAN}{'='*70}")
print(f"{Fore.CYAN}CONTEXT COMPRESSION RESEARCH AGENT - INTERACTIVE DEMO")
print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
print("\nThis demo allows you to test different compression strategies")
print("for researching OpenAI co-founders' current affiliations.\n")
def select_strategy() -> CompressionStrategy:
"""Let user select a compression strategy"""
print(f"{Fore.YELLOW}Available Compression Strategies:{Style.RESET_ALL}")
print("1. No Compression (expected to fail with large contexts)")
print("2. Non-Context-Aware: Individual Summaries (summarize each page, then concatenate)")
print("3. Non-Context-Aware: Combined Summary (concatenate all pages, then summarize once)")
print("4. Context-Aware Summarization")
print("5. Context-Aware with Citations")
print("6. Windowed Context (only compress when approaching context limit)")
while True:
try:
choice = input(f"\n{Fore.GREEN}Select strategy (1-6): {Style.RESET_ALL}")
strategies = [
CompressionStrategy.NO_COMPRESSION,
CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
CompressionStrategy.CONTEXT_AWARE,
CompressionStrategy.CONTEXT_AWARE_CITATIONS,
CompressionStrategy.WINDOWED_CONTEXT
]
return strategies[int(choice) - 1]
except (ValueError, IndexError):
print(f"{Fore.RED}Invalid choice. Please enter 1-6.{Style.RESET_ALL}")
def run_demo(enable_streaming=True, strategy: CompressionStrategy = None):
"""Run the interactive demo
Args:
enable_streaming: Whether to enable streaming output (default: True)
strategy: Preselected compression strategy; if None, prompt the user interactively
"""
print_banner()
# Check configuration
if not Config.validate():
print(f"\n{Fore.RED}Configuration validation failed!{Style.RESET_ALL}")
print("\nPlease set up your .env file with:")
print(" DASHSCOPE_API_KEY=your_api_key_here (for LLM_PROVIDER=dashscope/qwen/bailian)")
print(" MOONSHOT_API_KEY=your_api_key_here")
print(" SERPER_API_KEY=your_api_key_here (optional, will use mock data)")
sys.exit(1)
# Select strategy (interactively unless one was passed on the command line)
if strategy is None:
strategy = select_strategy()
print(f"\n{Fore.CYAN}Selected: {strategy.value}{Style.RESET_ALL}")
# Display streaming status
streaming_status = "ENABLED" if enable_streaming else "DISABLED"
print(f"{Fore.YELLOW}Streaming output: {streaming_status}{Style.RESET_ALL}")
# Create agent
print(f"\n{Fore.YELLOW}Initializing agent...{Style.RESET_ALL}")
agent = ResearchAgent(
api_key=Config.resolve_llm()[0],
compression_strategy=strategy,
verbose=False,
enable_streaming=enable_streaming
)
print(f"\n{Fore.CYAN}Starting research task...{Style.RESET_ALL}")
print("Task: Find current affiliations of all OpenAI co-founders\n")
print("-" * 70)
try:
# Execute research
result = agent.execute_research(max_iterations=Config.MAX_ITERATIONS)
# Print results
print("\n" + "="*70)
print(f"{Fore.GREEN}RESEARCH COMPLETE{Style.RESET_ALL}")
print("="*70)
if result.get('success'):
print(f"\n{Fore.GREEN}✅ Success!{Style.RESET_ALL}")
print(f"\nFinal Answer:\n{result.get('final_answer', 'No answer found')}")
else:
print(f"\n{Fore.RED}❌ Failed{Style.RESET_ALL}")
if result.get('error'):
print(f"Error: {result['error']}")
# Print statistics
trajectory = result.get('trajectory')
if trajectory:
print(f"\n{Fore.CYAN}📊 Statistics:{Style.RESET_ALL}")
print(f" Tool Calls: {len(trajectory.tool_calls)}")
print(f" Context Overflows: {trajectory.context_overflows}")
print(f" Execution Time: {result.get('execution_time', 0):.2f}s")
print(f" Total Tokens Used: {trajectory.total_tokens_used:,}")
print(f" - Prompt Tokens: {trajectory.prompt_tokens_used:,}")
print(f" - Completion Tokens: {trajectory.completion_tokens_used:,}")
# Calculate compression stats
if trajectory.tool_calls:
total_original = 0
total_compressed = 0
for call in trajectory.tool_calls:
if call.compressed_result:
total_original += call.compressed_result.original_length
total_compressed += call.compressed_result.compressed_length
if total_original > 0:
ratio = total_compressed / total_original
print(f" Compression Ratio: {ratio:.1%}")
print(f" Space Saved: {total_original - total_compressed:,} chars")
# Follow-up question demo (for citation strategy)
if strategy == CompressionStrategy.CONTEXT_AWARE_CITATIONS and result.get('success'):
print(f"\n{Fore.YELLOW}This strategy supports follow-up questions!{Style.RESET_ALL}")
follow_up = input("\nAsk a follow-up question (or press Enter to skip): ")
if follow_up:
print(f"\n{Fore.CYAN}Processing follow-up...{Style.RESET_ALL}")
# Add follow-up to conversation
agent.conversation_history.append({"role": "user", "content": follow_up})
# Get response (simplified for demo)
messages = agent.conversation_history.copy()
if enable_streaming:
message = agent._stream_response(messages)
else:
message = agent._non_streaming_response(messages)
if message.get('content'):
print(f"\n{Fore.GREEN}Follow-up Answer:{Style.RESET_ALL}")
print(message['content'])
except KeyboardInterrupt:
print(f"\n\n{Fore.YELLOW}Demo interrupted by user{Style.RESET_ALL}")
except Exception as e:
print(f"\n{Fore.RED}Error: {str(e)}{Style.RESET_ALL}")
def main():
"""Main entry point"""
# Parse command line arguments
parser = argparse.ArgumentParser(
prog="main.py",
description="上下文压缩策略交互式演示：针对“追踪 OpenAI 联合创始人现状”这一研究任务，"
"单独运行某一种压缩策略并实时观察其执行与压缩过程。",
epilog="示例：\n"
" python main.py # 交互式选择策略\n"
" python main.py -s citations # 直接运行“带引用的上下文感知”策略\n"
" python main.py -s windowed --no-streaming\n"
"如需批量对比全部策略并生成对比表，请使用 experiment.py。",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
'-s', '--strategy', choices=list(STRATEGY_CHOICES.keys()), metavar="NAME",
help="直接指定压缩策略（跳过交互式选择）。可选值：" + ", ".join(STRATEGY_CHOICES.keys()),
)
parser.add_argument(
'-m', '--model', default=None,
help=f"覆盖使用的模型名称（默认读取环境变量 MODEL_NAME，当前为 {Config.MODEL_NAME}）",
)
parser.add_argument(
'--no-streaming',
action='store_true',
help='关闭流式输出（默认开启流式）'
)
args = parser.parse_args()
if args.model:
Config.MODEL_NAME = args.model
# Determine streaming preference
enable_streaming = not args.no_streaming
preset_strategy = STRATEGY_CHOICES[args.strategy] if args.strategy else None
try:
run_demo(enable_streaming=enable_streaming, strategy=preset_strategy)
# Ask if user wants to try another strategy
while True:
again = input(f"\n{Fore.GREEN}Try another strategy? (y/n): {Style.RESET_ALL}")
if again.lower() == 'y':
run_demo(enable_streaming=enable_streaming)
else:
print(f"\n{Fore.CYAN}Thank you for using the demo!{Style.RESET_ALL}")
break
except KeyboardInterrupt:
print(f"\n\n{Fore.YELLOW}Goodbye!{Style.RESET_ALL}")
sys.exit(0)
if __name__ == "__main__":
main()
chapter2/context-compression/quickstart.py
Executable
+133
View File
@@ -0,0 +1,133 @@
#!/usr/bin/env python3
"""
Quick start script to test the context compression experiment
"""
import os
import sys
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def check_environment():
"""Check if environment is properly configured"""
provider = os.getenv("LLM_PROVIDER", "kimi").lower()
provider_key = os.getenv("DASHSCOPE_API_KEY") if provider in {"dashscope", "qwen", "bailian"} else os.getenv("MOONSHOT_API_KEY")
serper_key = os.getenv("SERPER_API_KEY")
print("🔍 Checking environment configuration...")
print("-" * 40)
if provider_key:
print(f"✅ API key for {provider} is set")
else:
print(f"❌ API key for {provider} is NOT set")
print(" Please add it to your .env file")
print(" Set DASHSCOPE_API_KEY for dashscope/qwen/bailian or MOONSHOT_API_KEY for kimi")
return False
if serper_key:
print("✅ SERPER_API_KEY is set")
else:
print("⚠️ SERPER_API_KEY is NOT set")
print(" Web search will use mock data")
print(" Get free API key at: https://serper.dev/")
print("-" * 40)
return True
def quick_test():
"""Run a quick test with context-aware citations strategy"""
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
from config import Config
print("\n🚀 Running quick test with Context-Aware Citations strategy...")
print("Task: Research OpenAI co-founders' current affiliations\n")
# Create agent
agent = ResearchAgent(
api_key=Config.resolve_llm()[0],
compression_strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS,
verbose=False,
enable_streaming=True
)
# Execute research
result = agent.execute_research(max_iterations=10)
# Print results
print("\n" + "="*60)
if result.get('success'):
print("✅ SUCCESS!")
print("\nFinal Answer:")
print(result.get('final_answer', 'No answer found'))
# Statistics
trajectory = result.get('trajectory')
if trajectory:
print(f"\n📊 Statistics:")
print(f" - Tool calls: {len(trajectory.tool_calls)}")
print(f" - Execution time: {result.get('execution_time', 0):.2f}s")
else:
print("❌ FAILED")
if result.get('error'):
print(f"Error: {result['error']}")
print("="*60)
def main():
"""Main entry point"""
print("\n" + "="*60)
print("CONTEXT COMPRESSION EXPERIMENT - QUICK START")
print("="*60 + "\n")
# Check environment
if not check_environment():
print("\n❌ Please configure your environment first!")
print("\n1. Copy env.example to .env:")
print(" cp env.example .env")
print("\n2. Edit .env and add your API keys")
print("\n3. Run this script again")
sys.exit(1)
# Menu
print("\n📋 What would you like to do?")
print("1. Run quick test (Context-Aware Citations)")
print("2. Run full experiment (all 6 strategies)")
print("3. Interactive demo (choose strategy)")
print("4. Exit")
try:
choice = input("\nSelect option (1-4): ")
if choice == "1":
quick_test()
elif choice == "2":
print("\n🔬 Starting full experiment...")
import experiment
experiment.main()
elif choice == "3":
print("\n🎮 Starting interactive demo...")
import main as demo
demo.main()
elif choice == "4":
print("\n👋 Goodbye!")
sys.exit(0)
else:
print("\n❌ Invalid choice")
sys.exit(1)
except KeyboardInterrupt:
print("\n\n⚠️ Interrupted by user")
sys.exit(0)
except Exception as e:
print(f"\n❌ Error: {str(e)}")
sys.exit(1)
if __name__ == "__main__":
main()
chapter2/context-compression/requirements.txt
+13
View File
@@ -0,0 +1,13 @@
# Shared provider resolver from the repository root. Run this requirements file
# from the experiment directory, as shown in the README.
-e ../..
openai>=1.0.0
requests>=2.31.0
beautifulsoup4>=4.12.0
html2text>=2020.1.16
python-dotenv>=1.0.0
lxml>=4.9.0
colorama>=0.4.6
tqdm>=4.66.0
tiktoken>=0.5.0
chapter2/context-compression/run_all_strategies.py
+458
View File
@@ -0,0 +1,458 @@
#!/usr/bin/env python3
"""
Script to run all compression strategies sequentially and save results to log
"""
import os
import sys
import json
import time
import argparse
import logging
from datetime import datetime
from typing import Dict, Any, List, Optional
from agent import ResearchAgent
from compression_strategies import CompressionStrategy, ContextCompressor
from config import Config
from colorama import init, Fore, Style
# Initialize colorama
init(autoreset=True)
# Short CLI aliases -> compression strategy (order matches the book's 实验 2-10)
STRATEGY_CHOICES = {
"no_compression": CompressionStrategy.NO_COMPRESSION,
"individual": CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
"combined": CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
"context_aware": CompressionStrategy.CONTEXT_AWARE,
"citations": CompressionStrategy.CONTEXT_AWARE_CITATIONS,
"windowed": CompressionStrategy.WINDOWED_CONTEXT,
}
ALL_STRATEGIES = list(STRATEGY_CHOICES.values())
class StrategyRunner:
"""Runs all compression strategies and logs results"""
def __init__(self, log_dir: str = "logs"):
"""
Initialize the strategy runner
Args:
log_dir: Directory to save log files
"""
self.log_dir = log_dir
os.makedirs(log_dir, exist_ok=True)
# Create log file with timestamp
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.log_file = os.path.join(log_dir, f"strategy_run_{timestamp}.log")
self.json_file = os.path.join(log_dir, f"strategy_results_{timestamp}.json")
# Configure logging
self.setup_logging()
# Results storage
self.results = []
def setup_logging(self):
"""Configure logging to file and console"""
# Create formatter
formatter = logging.Formatter(
'%(asctime)s - %(levelname)s - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
)
# File handler
file_handler = logging.FileHandler(self.log_file)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(formatter)
# Console handler - with custom filter for cleaner output
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
# Use a simpler format for console
console_formatter = logging.Formatter(
'%(asctime)s - %(levelname)s - %(message)s',
datefmt='%H:%M:%S'
)
console_handler.setFormatter(console_formatter)
# Configure root logger
self.logger = logging.getLogger('StrategyRunner')
self.logger.setLevel(logging.DEBUG)
self.logger.addHandler(file_handler)
self.logger.addHandler(console_handler)
def log_banner(self, message: str, char: str = "=", width: int = 70):
"""Log a banner message"""
border = char * width
self.logger.info(border)
self.logger.info(message.center(width))
self.logger.info(border)
def run_strategy(self, strategy: CompressionStrategy) -> Dict[str, Any]:
"""
Run a single compression strategy
Args:
strategy: The compression strategy to test
Returns:
Dictionary with results
"""
self.log_banner(f"Testing: {strategy.value}", char="-")
self.logger.info(f"Strategy: {strategy.value}")
result = {
'strategy': strategy.value,
'start_time': datetime.now().isoformat(),
'success': False,
'error': None,
'metrics': {}
}
try:
# Create agent with the strategy
self.logger.info("Creating agent...")
agent = ResearchAgent(
api_key=Config.MOONSHOT_API_KEY,
compression_strategy=strategy,
verbose=False,
enable_streaming=True # Enable streaming to see compressions
)
# Execute research task
self.logger.info("Starting research task...")
start_time = time.time()
# Custom stream handler to capture and log streaming output
class StreamCapture:
def __init__(self, logger, original_stdout):
self.logger = logger
self.original_stdout = original_stdout
self.buffer = []
self.current_line = []
def write(self, text):
# Accumulate text
self.current_line.append(text)
# If we have a newline, log the complete line
if '\n' in text:
full_line = ''.join(self.current_line)
lines = full_line.split('\n')
# Log all complete lines through logger (will go to both console and file)
for line in lines[:-1]:
if line.strip():
# Use INFO level for important summaries, DEBUG for other output
if any(keyword in line for keyword in ['📝', '🎯', '📚', '📄', 'Summarizing:', 'Creating']):
self.logger.info(f"[COMPRESSION] {line}")
else:
self.logger.debug(f"[AGENT] {line}")
self.buffer.append(line)
# Keep any partial line for next write
self.current_line = [lines[-1]] if lines[-1] else []
def flush(self):
# Flush any remaining partial line
if self.current_line:
remaining = ''.join(self.current_line)
if remaining.strip():
self.logger.debug(f"[AGENT] {remaining}")
self.buffer.append(remaining)
self.current_line = []
def get_output(self):
# Ensure any remaining content is flushed
self.flush()
return '\n'.join(self.buffer)
# Capture streaming output
original_stdout = sys.stdout
stream_capture = StreamCapture(self.logger, original_stdout)
try:
sys.stdout = stream_capture
research_result = agent.execute_research(max_iterations=Config.MAX_ITERATIONS)
finally:
sys.stdout = original_stdout
execution_time = time.time() - start_time
# Get the complete captured output for storage
output = stream_capture.get_output()
# Store output in result for later analysis
result['agent_output'] = output
# Process results
trajectory = research_result.get('trajectory')
if research_result.get('success'):
result['success'] = True
result['final_answer'] = research_result.get('final_answer', 'No answer found')
self.logger.info("✅ Strategy completed successfully")
else:
result['error'] = research_result.get('error', 'Unknown error')
self.logger.warning(f"⚠️ Strategy failed: {result['error']}")
# Collect metrics
if trajectory:
result['metrics'] = {
'execution_time': execution_time,
'tool_calls': len(trajectory.tool_calls),
'context_overflows': trajectory.context_overflows,
'total_tokens': trajectory.total_tokens_used,
'prompt_tokens': trajectory.prompt_tokens_used,
'completion_tokens': trajectory.completion_tokens_used
}
# Calculate compression statistics
total_original = 0
total_compressed = 0
for call in trajectory.tool_calls:
if call.compressed_result:
total_original += call.compressed_result.original_length
total_compressed += call.compressed_result.compressed_length
if total_original > 0:
compression_ratio = total_compressed / total_original
result['metrics']['compression_ratio'] = compression_ratio
result['metrics']['total_original_size'] = total_original
result['metrics']['total_compressed_size'] = total_compressed
result['metrics']['space_saved'] = total_original - total_compressed
# Log metrics
self.logger.info(f"Execution time: {execution_time:.2f}s")
self.logger.info(f"Tool calls: {result['metrics']['tool_calls']}")
self.logger.info(f"Context overflows: {result['metrics']['context_overflows']}")
self.logger.info(f"Total tokens: {result['metrics']['total_tokens']:,}")
if 'compression_ratio' in result['metrics']:
self.logger.info(f"Compression ratio: {result['metrics']['compression_ratio']:.1%}")
self.logger.info(f"Space saved: {result['metrics']['space_saved']:,} chars")
# Log compression details for each tool call
self.logger.debug("\nCompression details by tool call:")
for i, call in enumerate(trajectory.tool_calls, 1):
if call.compressed_result:
self.logger.debug(f" Tool call {i}: {call.tool_name}")
self.logger.debug(f" - Original: {call.compressed_result.original_length:,} chars")
self.logger.debug(f" - Compressed: {call.compressed_result.compressed_length:,} chars")
self.logger.debug(f" - Strategy: {call.compressed_result.strategy.value}")
except Exception as e:
result['error'] = str(e)
self.logger.error(f"❌ Error running strategy: {e}", exc_info=True)
result['end_time'] = datetime.now().isoformat()
return result
def run_all_strategies(self, strategies: Optional[List[CompressionStrategy]] = None):
"""Run the given compression strategies (default: all six)"""
if strategies is None:
strategies = list(ALL_STRATEGIES)
self.log_banner("COMPRESSION STRATEGIES TEST RUN", char="=")
self.logger.info(f"Testing {len(strategies)} strategies")
self.logger.info(f"Log file: {self.log_file}")
self.logger.info(f"JSON results: {self.json_file}")
# Run each strategy
for i, strategy in enumerate(strategies, 1):
self.logger.info(f"\n[{i}/{len(strategies)}] Running {strategy.value}")
result = self.run_strategy(strategy)
self.results.append(result)
# Small delay between strategies
if i < len(strategies):
time.sleep(2)
# Generate summary
self.generate_summary()
# Save results to JSON
self.save_json_results()
self.log_banner("TEST RUN COMPLETE", char="=")
self.logger.info(f"Results saved to:")
self.logger.info(f" - Log: {self.log_file}")
self.logger.info(f" - JSON: {self.json_file}")
def generate_summary(self):
"""Generate and log a summary of all results"""
self.log_banner("RESULTS SUMMARY", char="=")
# Create comparison table
self.logger.info("\nStrategy Comparison:")
self.logger.info("-" * 100)
self.logger.info(f"{'Strategy':<40} {'Success':<10} {'Time(s)':<10} {'Tokens':<12} {'Compression':<12} {'Overflows':<10}")
self.logger.info("-" * 100)
for result in self.results:
strategy = result['strategy'][:38] # Truncate if too long
success = "✅ Yes" if result['success'] else "❌ No"
metrics = result.get('metrics', {})
exec_time = f"{metrics.get('execution_time', 0):.2f}" if metrics else "N/A"
tokens = f"{metrics.get('total_tokens', 0):,}" if metrics else "N/A"
compression = f"{metrics.get('compression_ratio', 0):.1%}" if metrics.get('compression_ratio') else "N/A"
overflows = str(metrics.get('context_overflows', 0)) if metrics else "N/A"
self.logger.info(f"{strategy:<40} {success:<10} {exec_time:<10} {tokens:<12} {compression:<12} {overflows:<10}")
self.logger.info("-" * 100)
# Summary statistics
successful = sum(1 for r in self.results if r['success'])
failed = len(self.results) - successful
self.logger.info(f"\nOverall Results:")
self.logger.info(f" - Successful: {successful}/{len(self.results)}")
self.logger.info(f" - Failed: {failed}/{len(self.results)}")
# Find best performers
if successful > 0:
# Best compression ratio
compressed_results = [r for r in self.results if r.get('metrics', {}).get('compression_ratio')]
if compressed_results:
best_compression = min(compressed_results, key=lambda r: r['metrics']['compression_ratio'])
self.logger.info(f" - Best compression: {best_compression['strategy']} ({best_compression['metrics']['compression_ratio']:.1%})")
# Fastest execution
timed_results = [r for r in self.results if r.get('metrics', {}).get('execution_time')]
if timed_results:
fastest = min(timed_results, key=lambda r: r['metrics']['execution_time'])
self.logger.info(f" - Fastest: {fastest['strategy']} ({fastest['metrics']['execution_time']:.2f}s)")
# Most tokens used
token_results = [r for r in self.results if r.get('metrics', {}).get('total_tokens')]
if token_results:
most_tokens = max(token_results, key=lambda r: r['metrics']['total_tokens'])
least_tokens = min(token_results, key=lambda r: r['metrics']['total_tokens'])
self.logger.info(f" - Most tokens: {most_tokens['strategy']} ({most_tokens['metrics']['total_tokens']:,})")
self.logger.info(f" - Least tokens: {least_tokens['strategy']} ({least_tokens['metrics']['total_tokens']:,})")
def save_json_results(self):
"""Save results to JSON file"""
try:
with open(self.json_file, 'w') as f:
json.dump({
'run_date': datetime.now().isoformat(),
'config': {
'model': Config.MODEL_NAME,
'max_iterations': Config.MAX_ITERATIONS,
'context_window': Config.CONTEXT_WINDOW_SIZE,
'summary_max_tokens': Config.SUMMARY_MAX_TOKENS
},
'results': self.results
}, f, indent=2)
self.logger.info(f"JSON results saved to {self.json_file}")
except Exception as e:
self.logger.error(f"Failed to save JSON results: {e}")
def build_parser() -> argparse.ArgumentParser:
"""构建命令行参数解析器"""
parser = argparse.ArgumentParser(
prog="run_all_strategies.py",
description="逐个运行压缩策略并将完整过程（含流式压缩摘要）写入日志。\n"
"与 experiment.py 相比，本脚本侧重“可复盘的详细日志”：每次运行都会生成 "
".log 文本日志和 .json 结果文件，便于逐轮检查压缩效果。",
epilog="示例：\n"
" python run_all_strategies.py # 运行全部 6 种策略\n"
" python run_all_strategies.py -s windowed # 只跑自适应窗口化策略\n"
" python run_all_strategies.py --model kimi-k3 --log-dir logs/k2\n"
" python run_all_strategies.py --list-strategies",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"-s", "--strategy", nargs="+", choices=list(STRATEGY_CHOICES.keys()), metavar="NAME",
help="要运行的压缩策略（可指定多个，默认运行全部 6 种）。可选值："
+ ", ".join(STRATEGY_CHOICES.keys()),
)
parser.add_argument(
"-m", "--model", default=None,
help=f"覆盖使用的模型名称（默认读取环境变量 MODEL_NAME，当前为 {Config.MODEL_NAME}）",
)
parser.add_argument(
"--log-dir", default="logs", metavar="DIR",
help="日志与 JSON 结果的输出目录（默认 logs/）",
)
parser.add_argument(
"-n", "--max-iterations", type=int, default=None, metavar="N",
help=f"每个策略允许的最大迭代（工具调用轮数），默认 {Config.MAX_ITERATIONS}",
)
parser.add_argument(
"--list-strategies", action="store_true",
help="列出所有可选的压缩策略名称后退出",
)
return parser
def main():
"""Main entry point"""
parser = build_parser()
args = parser.parse_args()
if args.list_strategies:
print("可选的压缩策略（--strategy 的取值）：")
for alias, strat in STRATEGY_CHOICES.items():
print(f" {alias:<16} -> {strat.value}")
return
# Apply CLI overrides onto the shared Config
if args.model:
Config.MODEL_NAME = args.model
if args.max_iterations is not None:
Config.MAX_ITERATIONS = args.max_iterations
strategies = ([STRATEGY_CHOICES[name] for name in args.strategy]
if args.strategy else list(ALL_STRATEGIES))
print(f"\n{Fore.CYAN}{'='*70}")
print(f"{Fore.CYAN}COMPRESSION STRATEGIES AUTOMATED TEST RUNNER")
print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}\n")
# Validate configuration
if not Config.validate():
print(f"{Fore.RED}Configuration validation failed!{Style.RESET_ALL}")
print("\nPlease set up your .env file with:")
print(" MOONSHOT_API_KEY=your_api_key_here")
print(" SERPER_API_KEY=your_api_key_here (optional)")
sys.exit(1)
# Create directories
Config.create_directories()
# Run all strategies
runner = StrategyRunner(log_dir=args.log_dir)
try:
print(f"{Fore.YELLOW}Starting test run...{Style.RESET_ALL}")
print(f"Log file: {runner.log_file}\n")
runner.run_all_strategies(strategies)
print(f"\n{Fore.GREEN}✅ Test run complete!{Style.RESET_ALL}")
print(f"\nResults saved to:")
print(f" 📄 Log: {runner.log_file}")
print(f" 📊 JSON: {runner.json_file}")
except KeyboardInterrupt:
print(f"\n{Fore.YELLOW}Test run interrupted by user{Style.RESET_ALL}")
runner.logger.warning("Test run interrupted by user")
except Exception as e:
print(f"\n{Fore.RED}Fatal error: {e}{Style.RESET_ALL}")
runner.logger.error(f"Fatal error: {e}", exc_info=True)
sys.exit(1)
if __name__ == "__main__":
main()
chapter2/context-compression/test_malformed_tool_json.py
+141
View File
@@ -0,0 +1,141 @@
"""
Malformed tool-argument JSON must not abort execute_research or cause real dispatch to raise TypeError.
"""
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock, patch
# Optional deps used at import time by web_tools.
sys.modules.setdefault("html2text", types.ModuleType("html2text"))
sys.modules.setdefault("dotenv", types.SimpleNamespace(load_dotenv=lambda: None))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from compression_strategies import CompressionStrategy
from agent import ResearchAgent
def test_execute_research_survives_malformed_tool_arguments_json():
with patch("agent.Config.resolve_llm", return_value=("k", "http://x", "m")), \
patch("agent.OpenAI"), \
patch("agent.WebTools") as mock_web_tools_cls, \
patch("agent.ContextCompressor"):
mock_web_tools = MagicMock()
mock_web_tools_cls.return_value = mock_web_tools
agent = ResearchAgent(
api_key="k",
compression_strategy=CompressionStrategy.NO_COMPRESSION,
verbose=False,
enable_streaming=False,
)
bad_call_search = {
"id": "call-bad-search",
"type": "function",
"function": {
"name": "search_web",
"arguments": '{"query": "openai",}', # trailing comma
},
}
bad_call_fetch = {
"id": "call-bad-fetch",
"type": "function",
"function": {
"name": "fetch_webpage",
"arguments": '{"url": "https://example.com",}', # malformed JSON
},
}
tool_msg = {"role": "assistant", "content": "searching", "tool_calls": [bad_call_search, bad_call_fetch]}
final_msg = {"role": "assistant", "content": "FINAL ANSWER: ok", "tool_calls": None}
agent._non_streaming_response = MagicMock(side_effect=[tool_msg, final_msg])
# Do NOT mock _execute_tool, let real dispatch run over missing query/url arguments
result = agent.execute_research(max_iterations=3)
assert result.get("error") is None
assert len(agent.trajectory.tool_calls) == 2
assert agent.trajectory.tool_calls[0].result == {"error": "Missing required argument 'query' for search_web"}
assert agent.trajectory.tool_calls[1].result == {"error": "Missing required argument 'url' for fetch_webpage"}
def test_execute_research_survives_non_dict_and_invalid_bytes_tool_arguments():
"""
Non-dict JSON structures (lists, numbers) and invalid UTF-8 bytes must normalize to {} and not raise.
"""
with patch("agent.Config.resolve_llm", return_value=("k", "http://x", "m")), \
patch("agent.OpenAI"), \
patch("agent.WebTools") as mock_web_tools_cls, \
patch("agent.ContextCompressor"):
mock_web_tools = MagicMock()
mock_web_tools_cls.return_value = mock_web_tools
agent = ResearchAgent(
api_key="k",
compression_strategy=CompressionStrategy.NO_COMPRESSION,
verbose=False,
enable_streaming=False,
)
non_dict_calls = [
{"id": "c1", "type": "function", "function": {"name": "search_web", "arguments": "[]"}},
{"id": "c2", "type": "function", "function": {"name": "fetch_webpage", "arguments": "123"}},
{"id": "c3", "type": "function", "function": {"name": "search_web", "arguments": b"\x80\xff"}},
{"id": "c4", "type": "function", "function": {"name": "fetch_webpage", "arguments": {"invalid": 1}}},
]
tool_msg = {"role": "assistant", "content": "searching", "tool_calls": non_dict_calls}
final_msg = {"role": "assistant", "content": "FINAL ANSWER: ok", "tool_calls": None}
agent._non_streaming_response = MagicMock(side_effect=[tool_msg, final_msg])
result = agent.execute_research(max_iterations=3)
assert result.get("error") is None
assert len(agent.trajectory.tool_calls) == 4
assert agent.trajectory.tool_calls[0].arguments == {}
assert agent.trajectory.tool_calls[1].arguments == {}
assert agent.trajectory.tool_calls[2].arguments == {}
assert agent.trajectory.tool_calls[3].arguments == {"invalid": 1}
assert agent.trajectory.tool_calls[0].result == {"error": "Missing required argument 'query' for search_web"}
assert agent.trajectory.tool_calls[1].result == {"error": "Missing required argument 'url' for fetch_webpage"}
def test_execute_research_survives_tool_execution_exceptions():
"""
Tool execution exceptions in search_web or fetch_webpage must return error dicts with compressed=None and not abort loop.
"""
with patch("agent.Config.resolve_llm", return_value=("k", "http://x", "m")), \
patch("agent.OpenAI"), \
patch("agent.WebTools") as mock_web_tools_cls, \
patch("agent.ContextCompressor"):
mock_web_tools = MagicMock()
mock_web_tools.search_web.side_effect = RuntimeError("network connection failed")
mock_web_tools.fetch_webpage.side_effect = RuntimeError("http 500 error")
mock_web_tools_cls.return_value = mock_web_tools
agent = ResearchAgent(
api_key="k",
compression_strategy=CompressionStrategy.NO_COMPRESSION,
verbose=False,
enable_streaming=False,
)
call_search = {"id": "c1", "type": "function", "function": {"name": "search_web", "arguments": '{"query": "test"}'}}
call_fetch = {"id": "c2", "type": "function", "function": {"name": "fetch_webpage", "arguments": '{"url": "http://test.com"}'}}
tool_msg = {"role": "assistant", "content": "searching", "tool_calls": [call_search, call_fetch]}
final_msg = {"role": "assistant", "content": "FINAL ANSWER: ok", "tool_calls": None}
agent._non_streaming_response = MagicMock(side_effect=[tool_msg, final_msg])
result = agent.execute_research(max_iterations=3)
assert result.get("error") is None
assert len(agent.trajectory.tool_calls) == 2
assert agent.trajectory.tool_calls[0].result == {"error": "Failed to execute search_web: network connection failed"}
assert agent.trajectory.tool_calls[0].compressed_result is None
assert agent.trajectory.tool_calls[1].result == {"error": "Failed to execute fetch_webpage: http 500 error"}
assert agent.trajectory.tool_calls[1].compressed_result is None
chapter2/context-compression/test_null_content_compression.py
+26
View File
@@ -0,0 +1,26 @@
"""
Test suite locking out TypeError in ContextCompressor._no_compression
when a search result dictionary contains 'content': None.
"""
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
from compression_strategies import ContextCompressor
def test_no_compression_handles_null_content():
"""
Ensure _no_compression does not raise TypeError when result['content'] is None.
"""
compressor = ContextCompressor.__new__(ContextCompressor)
search_results = {
'results': [
{'title': 'Test', 'url': 'http://example.com', 'snippet': 'snippet', 'content': None}
]
}
compressed = compressor._no_compression(search_results)
assert compressed.original_length == 0
assert "Full Content:" in compressed.content
chapter2/context-compression/test_strategies.py
+1
View File
@@ -0,0 +1 @@
chapter2/context-compression/test_streaming_logging.py
+1
View File
@@ -0,0 +1 @@
chapter2/context-compression/test_token_tracking.py
+1
View File
@@ -0,0 +1 @@
chapter2/context-compression/test_web_tools_null_content.py
+59
View File
@@ -0,0 +1,59 @@
"""
Test suite locking out TypeError in WebTools.search_web
when fetch_webpage returns a dictionary containing 'content': None.
"""
import os
import sys
from unittest.mock import MagicMock
sys.modules['html2text'] = MagicMock()
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
from web_tools import WebTools
def test_search_web_handles_null_page_content():
"""
Ensure search_web calculates content_length correctly when page_content['content'] is None.
"""
tool = WebTools.__new__(WebTools)
tool.serper_api_key = "dummy"
tool.fetch_webpage = MagicMock(return_value={'content': None, 'success': True})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.json.return_value = {
'organic': [
{'title': 'Example', 'link': 'http://example.com', 'snippet': 'Test snippet'}
]
}
import requests
requests.post = MagicMock(return_value=mock_resp)
res = tool.search_web("test query", num_results=1)
assert res['num_results'] == 1
assert res['results'][0]['content_length'] == 0
def test_fetch_webpage_title_with_nested_elements():
"""
Ensure fetch_webpage extracts title text correctly when title tag has nested HTML elements.
"""
html_content = "<html><head><title>Report <span>2026</span></title></head><body><p>Test content</p></body></html>"
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.text = html_content
mock_resp.raise_for_status = MagicMock()
import requests
requests.get = MagicMock(return_value=mock_resp)
tool = WebTools.__new__(WebTools)
tool.page_cache = {}
tool.html_converter = MagicMock()
tool.html_converter.handle.return_value = "Test content"
res = tool.fetch_webpage("http://example.com/report")
assert res['title'] == "Report 2026"
chapter2/context-compression/web_tools.py
+306
View File
@@ -0,0 +1,306 @@
"""
Web tools for searching and fetching web pages
"""
import json
import html
import re
import logging
import requests
from typing import List, Dict, Any, Optional
from bs4 import BeautifulSoup
import html2text
from urllib.parse import urlparse, urljoin
import time
from config import Config
# Configure logging
logging.basicConfig(level=logging.INFO, format=Config.LOG_FORMAT)
logger = logging.getLogger(__name__)
class WebTools:
"""Tools for web search and page fetching"""
def __init__(self):
"""Initialize web tools"""
self.serper_api_key = Config.SERPER_API_KEY
self.html_converter = html2text.HTML2Text()
self.html_converter.ignore_links = False
self.html_converter.ignore_images = True
self.html_converter.ignore_emphasis = False
self.html_converter.body_width = 0 # Don't wrap lines
self.html_converter.single_line_break = True
# Cache for fetched pages to avoid redundant fetches
self.page_cache = {}
def search_web(self, query: str, num_results: int = 5) -> Dict[str, Any]:
"""
Search the web using Serper API
Args:
query: Search query
num_results: Number of results to return
Returns:
Dictionary containing search results with crawled content
"""
try:
if not self.serper_api_key:
# Fallback to mock results for demo
logger.warning("No Serper API key, using mock results")
return self._get_mock_search_results(query)
logger.info(f"Searching web for: {query}")
# Call Serper API
headers = {
'X-API-KEY': self.serper_api_key,
'Content-Type': 'application/json'
}
payload = {
'q': query,
'num': num_results
}
response = requests.post(
f"{Config.SERPER_BASE_URL}/search",
headers=headers,
json=payload,
timeout=10
)
if response.status_code != 200:
logger.error(f"Serper API error: {response.status_code}")
return self._get_mock_search_results(query)
data = response.json()
# Process organic results
results = []
organic_results = data.get('organic', [])[:num_results]
for result in organic_results:
# Fetch and convert each page
url = result.get('link', '')
if url:
page_content = self.fetch_webpage(url)
results.append({
'title': result.get('title', ''),
'url': url,
'snippet': result.get('snippet', ''),
'content': page_content.get('content', ''),
'content_length': len(page_content.get('content') or ''),
'fetch_success': page_content.get('success', False)
})
# Small delay to be respectful
time.sleep(0.5)
return {
'query': query,
'num_results': len(results),
'results': results,
'timestamp': time.time()
}
except Exception as e:
logger.error(f"Error searching web: {str(e)}")
return self._get_mock_search_results(query)
def fetch_webpage(self, url: str) -> Dict[str, Any]:
"""
Fetch a webpage and convert HTML to text
Args:
url: URL of the webpage to fetch
Returns:
Dictionary containing the converted text content
"""
try:
# Check cache first
if url in self.page_cache:
logger.info(f"Using cached content for: {url}")
return self.page_cache[url]
logger.info(f"Fetching webpage: {url}")
# Fetch the page
headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
}
response = requests.get(url, headers=headers, timeout=10)
response.raise_for_status()
# Parse HTML
soup = BeautifulSoup(response.text, 'lxml')
# Remove script and style elements
for script in soup(["script", "style", "nav", "footer", "header"]):
script.decompose()
# Convert to text
text_content = self.html_converter.handle(str(soup))
# Clean up the text
lines = text_content.split('\n')
cleaned_lines = []
for line in lines:
line = line.strip()
if line and not line.startswith('#'): # Remove empty lines and navigation markers
cleaned_lines.append(line)
cleaned_text = '\n'.join(cleaned_lines)
# Truncate if too long
if len(cleaned_text) > Config.MAX_WEBPAGE_LENGTH:
cleaned_text = cleaned_text[:Config.MAX_WEBPAGE_LENGTH] + "\n\n[Content truncated...]"
title = 'No title'
if soup.title:
raw_title = soup.title.get_text()
cleaned_title = html.unescape(re.sub(r'<[^>]+>', '', raw_title)).strip()
if cleaned_title:
title = cleaned_title
result = {
'url': url,
'title': title,
'content': cleaned_text,
'content_length': len(cleaned_text),
'success': True,
'timestamp': time.time()
}
# Cache the result
self.page_cache[url] = result
return result
except Exception as e:
logger.error(f"Error fetching webpage {url}: {str(e)}")
error_result = {
'url': url,
'title': 'Error',
'content': f"Failed to fetch webpage: {str(e)}",
'content_length': 0,
'success': False,
'error': str(e),
'timestamp': time.time()
}
# Cache even failed results to avoid retrying
self.page_cache[url] = error_result
return error_result
def _get_mock_search_results(self, query: str) -> Dict[str, Any]:
"""
Get mock search results for testing without API key
Args:
query: Search query
Returns:
Mock search results
"""
# Mock results for OpenAI co-founders
mock_data = {
"openai": [
{
'title': 'OpenAI - Wikipedia',
'url': 'https://en.wikipedia.org/wiki/OpenAI',
'snippet': 'OpenAI was founded in 2015 by Sam Altman, Elon Musk, Ilya Sutskever, Greg Brockman, Wojciech Zaremba, and John Schulman...',
'content': '''OpenAI was founded in December 2015 by Sam Altman, Elon Musk, Ilya Sutskever, Greg Brockman, Wojciech Zaremba, and John Schulman.
The organization was founded with the goal of advancing digital intelligence in a way that benefits humanity.
Current Status of Co-founders (as of 2024):
- Sam Altman: CEO of OpenAI (returned after brief departure in November 2023)
- Elon Musk: Left OpenAI board in 2018, founded xAI in 2023
- Ilya Sutskever: Former Chief Scientist, left OpenAI in May 2024, co-founded Safe Superintelligence Inc.
- Greg Brockman: President and Chairman of OpenAI
- Wojciech Zaremba: Head of Language and Code Generation at OpenAI
- John Schulman: Co-founder, left OpenAI in August 2024 to join Anthropic
Additional early members:
- Andrej Karpathy: Former Director of AI at Tesla, briefly returned to OpenAI, now independent
- Dario Amodei: Left to co-found Anthropic in 2021
- Daniela Amodei: Left to co-found Anthropic in 2021'''
}
],
"sam altman": [
{
'title': 'Sam Altman - CEO of OpenAI',
'url': 'https://example.com/sam-altman',
'snippet': 'Sam Altman is the CEO of OpenAI...',
'content': 'Sam Altman is currently the CEO of OpenAI. He briefly left the company in November 2023 but returned after employee protests. He is also known for his work at Y Combinator and various investments in startups.'
}
],
"elon musk": [
{
'title': 'Elon Musk launches xAI',
'url': 'https://example.com/elon-musk-ai',
'snippet': 'Elon Musk founded xAI in 2023...',
'content': 'Elon Musk, who co-founded OpenAI in 2015, left the board in 2018 citing conflicts of interest with Tesla\'s AI development. In 2023, he founded xAI, a new AI company focused on understanding the universe. He is also CEO of Tesla, SpaceX, and owner of X (formerly Twitter).'
}
],
"ilya sutskever": [
{
'title': 'Ilya Sutskever launches Safe Superintelligence',
'url': 'https://example.com/ilya-sutskever',
'snippet': 'Ilya Sutskever left OpenAI to start SSI...',
'content': 'Ilya Sutskever, former Chief Scientist at OpenAI, left the company in May 2024 after nearly a decade. He co-founded Safe Superintelligence Inc. (SSI) with Daniel Gross and Daniel Levy, focusing on building safe AGI.'
}
]
}
# Find matching mock data
query_lower = query.lower()
for key in mock_data:
if key in query_lower:
results = []
for item in mock_data[key]:
results.append({
'title': item['title'],
'url': item['url'],
'snippet': item['snippet'],
'content': item['content'],
'content_length': len(item['content']),
'fetch_success': True
})
return {
'query': query,
'num_results': len(results),
'results': results,
'timestamp': time.time(),
'mock': True
}
# Default mock result
return {
'query': query,
'num_results': 1,
'results': [{
'title': 'Mock Search Result',
'url': 'https://example.com',
'snippet': 'This is a mock search result for testing',
'content': 'Mock content for testing when no API key is available.',
'content_length': 50,
'fetch_success': True
}],
'timestamp': time.time(),
'mock': True
}
def clear_cache(self):
"""Clear the page cache"""
self.page_cache.clear()
logger.info("Page cache cleared")