context-compression¶
第2章 · 上下文工程 · 配套项目
chapter2/context-compression
项目说明¶
Context Compression Strategies Experiment¶
This project demonstrates and compares different context compression strategies for LLM agents, using the task of researching OpenAI co-founders' current affiliations as a test case.
Overview¶
As LLM context windows grow larger (128K+ tokens), managing context efficiently becomes crucial for: - Cost optimization - Reducing token usage - Performance - Faster response times - Reliability - Avoiding context overflow errors - Relevance - Maintaining focus on important information
This experiment implements and compares 6 different context compression strategies to understand their trade-offs.
Compression Strategies¶
1. No Compression¶
- Description: Puts all original webpage content directly into agent context
- Expected Result: Fails after a few tool calls due to context overflow
- Purpose: Demonstrates the baseline problem
2. Non-Context-Aware: Individual Summaries¶
- Description: Summarizes each webpage independently using LLM, then concatenates all summaries
- Expected Result: Preserves page-specific details but may lose cross-page relationships
- Trade-off: Multiple LLM calls (one per page) but maintains page boundaries
- Best for: When each source should be treated independently
3. Non-Context-Aware: Combined Summary¶
- Description: Concatenates all webpage content first, then creates a single comprehensive summary
- Expected Result: Better understanding of overall content but may lose page-specific attribution
- Trade-off: Single LLM call but might hit token limits with many pages
- Best for: When looking for overarching themes across multiple sources
4. Context-Aware Summarization¶
- Description: Combines all search results and creates query-focused summary
- Expected Result: Better relevance preservation
- Trade-off: Requires additional LLM call for summarization
5. Context-Aware with Citations¶
- Description: Similar to #4 but includes citations and source links
- Expected Result: Enables follow-up questions with source tracking
- Trade-off: Slightly larger context but maintains traceability
6. Windowed Context¶
- Description: Keeps full content for last tool call, compresses older history
- Expected Result: Balance between detail and efficiency
- Trade-off: Recent detail vs. historical compression
- Smart Compression: Only compresses uncompressed messages (marked with [COMPRESSED] to prevent re-compression)
Installation¶
-
Navigate to the project:
-
Install dependencies:
-
Set up environment variables:
Required API keys:
- MOONSHOT_API_KEY: For the Kimi (Moonshot) model (required). The book's 实验 2-9 uses Kimi K3 (a reasoning
model whose real context window is ~1M tokens; the demo deliberately caps context at a 128K budget via
CONTEXT_WINDOW_SIZE so the overflow/compression behavior is observable). The model name is configurable via
MODEL_NAME in .env or the -m/--model CLI flag (e.g. kimi-k2.5, kimi-k3, moonshot-v1-128k).
- OPENROUTER_API_KEY: 通用回退。未设置 MOONSHOT_API_KEY 时，只要配置了
OPENROUTER_API_KEY，实验会自动改走 OpenRouter（kimi-* 映射为
moonshotai/kimi-k2）。设置了 MOONSHOT_API_KEY 时行为完全不变。
- SERPER_API_KEY: For web search (optional, will use mock data if not provided)
Get API keys: - Moonshot/Kimi: https://platform.moonshot.cn/ - Serper (free tier): https://serper.dev/
Scripts Overview¶
| Script | Purpose | Output |
|---|---|---|
main.py |
Interactive demo / single strategy runner | Console output |
experiment.py |
Automated comparison of strategies (token / compression / success table) | Results in results/ |
run_all_strategies.py |
Run strategies with detailed per-round logging | Logs in logs/ |
quickstart.py |
Menu wrapper that checks env and launches the above | Console output |
All three main entrypoints ship an argparse CLI (Chinese --help). Run any of them with -h
to see the full option list. The three most useful flags are shared:
-s/--strategy— 选择要运行的一种或多种策略（默认全部 6 种）；取值见下方“Compression Strategies”或运行--list-strategies-m/--model— 覆盖模型名（默认读取环境变量MODEL_NAME）-n/--max-iterations— 每个策略允许的最大工具调用轮数
Strategy aliases accepted by --strategy: no_compression, individual, combined,
context_aware, citations, windowed.
Usage¶
Run Full Experiment (comparison table + JSON)¶
Compare all 6 strategies (default), or a subset:
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 # 查看可选策略名
This will:
- Test each selected compression strategy sequentially
- Research OpenAI co-founders' affiliations
- Print a comparison table (Success / Time / Tokens / Compression / Overflows)
- Save results to results/experiment_TIMESTAMP.json (or the path given by -o/--output)
Key flags: -s/--strategy, -m/--model, -o/--output, -n/--max-iterations, --streaming, --list-strategies.
Run All Strategies with Logging¶
Run strategies with detailed logging and compression output:
python run_all_strategies.py # 全部 6 种策略
python run_all_strategies.py -s windowed # 只跑自适应窗口化
python run_all_strategies.py --log-dir logs/k2 -m kimi-k2.5
Features:
- Runs the selected compression strategies sequentially
- Logs all compression summaries to file
- Shows streaming output in real-time
- Saves detailed logs to <log-dir>/strategy_run_TIMESTAMP.log
- Saves JSON results to <log-dir>/strategy_results_TIMESTAMP.json
- Generates comparison summary at the end
Key flags: -s/--strategy, -m/--model, --log-dir, -n/--max-iterations, --list-strategies.
Interactive Demo¶
Test individual strategies with streaming output:
python main.py # Interactive: choose a strategy at the prompt
python main.py -s citations # Run a specific strategy non-interactively
python main.py -s windowed --no-streaming # Disable streaming output
Features:
- Choose any compression strategy (interactively, or with -s/--strategy)
- Streaming responses enabled by default (--no-streaming to disable)
- See real-time execution
- Try follow-up questions (for citation strategy)
Custom Usage¶
from agent import ResearchAgent
from compression_strategies import CompressionStrategy
# Create agent with specific strategy
agent = ResearchAgent(
api_key="your_api_key",
compression_strategy=CompressionStrategy.CONTEXT_AWARE_CITATIONS,
enable_streaming=True
)
# Execute research
result = agent.execute_research()
# Access results
if result['success']:
print(result['final_answer'])
print(f"Tool calls: {len(result['trajectory'].tool_calls)}")
Project Structure¶
context-compression/
├── config.py # Configuration management
├── web_tools.py # Web search and fetch tools
├── compression_strategies.py # Compression strategy implementations
├── agent.py # Main research agent with streaming
├── experiment.py # Experiment runner for comparisons (CLI)
├── run_all_strategies.py # Detailed per-round logging runner (CLI)
├── main.py # Interactive demo / single strategy runner (CLI)
├── quickstart.py # Menu wrapper (env check + launcher)
├── requirements.txt # Python dependencies
├── env.example # Environment variables template
├── logs/ # Detailed logs (created by run_all_strategies.py)
└── results/ # Experiment results (created on run)
Key Components¶
Web Tools (web_tools.py)¶
- search_web: Searches using Serper API, crawls each result
- fetch_webpage: Fetches and converts HTML to clean text
- Mock data: Provides sample data when API key unavailable
Compression Strategies (compression_strategies.py)¶
- ContextCompressor: Implements all 6 strategies
- CompressedContent: Data class for compressed results
- Dynamic compression: Based on query and context
Research Agent (agent.py)¶
- Streaming support: Real-time response streaming
- Tool integration: Web search and fetch capabilities
- Message management: Handles conversation history
- Windowed compression: Dynamic history compression
Experiment Runner (experiment.py)¶
- Automated testing: Runs all strategies
- Metrics collection: Execution time, compression ratio, success rate
- Comparison report: Visual comparison table
- Results persistence: JSON output for analysis
Metrics Collected¶
- Success Rate: Whether task completed successfully
- Execution Time: Total time to complete research
- Compression Ratio: Compressed size / original size
- Context Overflows: Number of times context limit approached
- Tool Calls: Number of web searches performed
- Final Answer Length: Size of generated report
Expected Results¶
Based on the compression strategies:
- No Compression: ❌ Fails with context overflow
- Non-Context-Aware: ⚠️ Completes but may miss details
- Context-Aware: ✅ Good balance of size and relevance
- With Citations: ✅ Best for follow-ups, slightly larger
- Windowed Context: ✅ Most efficient for long conversations
Measured Results (real run)¶
The numbers below are from a real end-to-end run — no mock data. Every strategy used live Serper web search (real 2026 web pages fetched and crawled) and the current Moonshot reasoning model.
- Model:
kimi-k3(Moonshot reasoning model; real window ~1M tokens, but the demo caps the compression/overflow budget atCONTEXT_WINDOW_SIZE = 128000) - Search: real Serper (
google.serper.dev) web search + page crawling - Task: 识别并追踪 OpenAI 联合创始人的职业状态 (track current affiliations of the ~11 OpenAI co-founders)
- Run date: 2026-07-18 ·
MAX_ITERATIONS=15· raw JSON:results/kimi_k3_real_20260718.json
Columns: Tokens = cumulative Kimi API token usage (prompt + completion across all iterations); Compress = compressed chars / original chars (smaller = compressed harder); Overflows = times the 128K budget was approached/exceeded.
| # | 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 exactly as designed: context overflows the 128K budget (165,227 tokens used) around the 5th iteration. - Context-aware summary (#4) is the most token-efficient successful strategy (40,157 tokens, 3.0% char compression) — it compresses hardest while still solving the task. - Individual summaries (#2) are by far the slowest (~50 min) because each fetched page is summarized separately by the reasoning model; token usage is also the highest. - Windowed context (#6) only compresses when prompt usage crosses the 80% threshold (≈102,400 tokens), then batch-compresses all uncompressed tool messages at once; because it keeps recent full content, its char-level "compression ratio" stays ~100% while it still completes the task fastest among the compressing strategies. - These are single-run measurements with a reasoning model and live web search, so absolute numbers will vary run to run; the relative ordering is the takeaway.
Configuration¶
Edit .env or config.py for:
MODEL_NAME: LLM model to use (default: kimi-k3)MODEL_TEMPERATURE: Response randomness (default: 0.3)MAX_ITERATIONS: Maximum tool calls (default: 50)MAX_WEBPAGE_LENGTH: Max chars per webpage (default: 50000)SUMMARY_MAX_TOKENS: Max tokens for summaries (default: 500)CONTEXT_WINDOW_SIZE: Context budget for the demo's overflow/compression trigger (default: 128000; note Kimi K3's real window is ~1M tokens — this is an intentionally smaller budget so compression is exercised)
Troubleshooting¶
No API Keys¶
The system will use mock data if SERPER_API_KEY is not set, allowing you to test the compression strategies without web search.
Context Overflow¶
If you encounter context overflow with strategies other than "No Compression", try:
- Reducing MAX_WEBPAGE_LENGTH
- Decreasing SUMMARY_MAX_TOKENS
- Limiting search results with num_results
Slow Execution¶
- Disable streaming in demo for faster output
- Reduce
MAX_ITERATIONSfor quicker experiments - Use mock data instead of real web search
Research Task¶
The experiment uses a specific research task:
"Find the current affiliations of all OpenAI co-founders"
This task is ideal because it: - Requires multiple searches (one per co-founder) - Generates substantial content (biographical information) - Tests context management (accumulating information) - Has verifiable results (known affiliations)
Extending the Project¶
To add new compression strategies:
- Add strategy to
CompressionStrategyenum - Implement in
ContextCompressorclass - Add handling in
compress_search_results() - Update experiment runner if needed
To change the research task:
- Modify system prompt in
agent.py - Update mock data in
web_tools.py - Adjust tool descriptions as needed
License¶
This project is part of the AI Agent practical training course and is for educational purposes.
源代码¶
agent.py¶
"""
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())
@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
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 tool_name == "search_web":
result = self.web_tools.search_web(**arguments)
# 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":
result = self.web_tools.fetch_webpage(**arguments)
# 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)
context_threshold = Config.CONTEXT_WINDOW_SIZE * 0.8
if self.trajectory.prompt_tokens_used <= context_threshold:
logger.debug(f"Windowed compression: Context usage below threshold ({self.trajectory.prompt_tokens_used:,}/{context_threshold:.0f} tokens)")
return messages # No compression needed yet
logger.info(f"⚠️ Context usage exceeds 80% threshold ({self.trajectory.prompt_tokens_used:,}/{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 hasattr(call, 'id') 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.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.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
if self.trajectory.prompt_tokens_used > Config.CONTEXT_WINDOW_SIZE * 0.8:
logger.warning(f"Approaching context limit: {self.trajectory.prompt_tokens_used:,} prompt tokens used")
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.prompt_tokens_used:,} tokens used (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']
function_args = json.loads(tool_call['function']['arguments'])
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
)
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")
compression_strategies.py¶
"""
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:
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:
# 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]} # Limit for history compression
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:
content = chunk.choices[0].delta.content
print(content, end="", flush=True)
summary_parts.append(content)
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', ''))
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]} # Limit to prevent token overflow
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} # Already limited per page
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} # Already limited per page
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
config.py¶
"""
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
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", "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
"""
if not cls.MOONSHOT_API_KEY and not cls.OPENROUTER_API_KEY:
print("ERROR: neither MOONSHOT_API_KEY nor OPENROUTER_API_KEY is set")
print("Please set MOONSHOT_API_KEY (primary) or OPENROUTER_API_KEY "
"(universal fallback) in .env or as an environment variable")
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)`` honoring the MOONSHOT->OpenRouter fallback.
Computed at call time so a runtime override of ``Config.MODEL_NAME``
(e.g. via ``--model``) is respected.
"""
from openrouter_fallback import resolve_llm as _resolve
return _resolve(
model=cls.MODEL_NAME,
primary_keys=("MOONSHOT_API_KEY", "KIMI_API_KEY"),
primary_base_url=cls.MOONSHOT_BASE_URL,
)
@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"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")
experiment.py¶
#!/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-9)
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-9）。\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()
main.py¶
#!/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-9)
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(" 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.MOONSHOT_API_KEY,
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()
openrouter_fallback.py¶
"""Universal OpenRouter fallback helper (Chapter 2 experiments).
Goal: every experiment keeps working when the direct provider key is missing
but ``OPENROUTER_API_KEY`` is present. Default behavior is fully preserved:
* If a primary provider key is present -> use the primary provider unchanged.
* Else if ``OPENROUTER_API_KEY`` is present -> route through OpenRouter
(base_url=https://openrouter.ai/api/v1) and translate the model id.
* Else raise a clear error listing every accepted key.
Model translation (only applied when the OpenRouter fallback is active):
* ids already containing "/" are passed through unchanged;
* gpt-* / o1* / o3* / o4* / chatgpt* -> "openai/<id>";
* claude-* -> "anthropic/claude-opus-4.8";
* kimi-* -> "moonshotai/kimi-k2.6";
* anything else -> passed through unchanged.
"""
import os
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
def is_openrouter_key(api_key):
"""OpenRouter keys reliably start with ``sk-or-``."""
return bool(api_key) and api_key.startswith("sk-or-")
def map_model_to_openrouter(model):
"""Translate a bare provider model id into an OpenRouter-qualified id."""
if not model or "/" in model:
return model
low = model.lower()
if low.startswith(("gpt-", "gpt5", "o1", "o3", "o4", "chatgpt")):
return "openai/" + model
if low.startswith("claude"):
return "anthropic/claude-opus-4.8"
if low.startswith("kimi"):
return "moonshotai/kimi-k2.6"
return model
def resolve_llm(model=None, primary_keys=("OPENAI_API_KEY",), primary_base_url=None):
"""Resolve ``(api_key, base_url, model)`` honoring the primary->OpenRouter fallback.
Args:
model: requested model id (may be remapped when the fallback activates).
primary_keys: env var names checked in order for the primary provider.
primary_base_url: base_url used when a primary key is found
(``None`` means the OpenAI SDK default / official endpoint).
"""
for env in primary_keys:
key = os.getenv(env)
if key:
return key, primary_base_url, model
or_key = os.getenv("OPENROUTER_API_KEY")
if or_key:
return or_key, OPENROUTER_BASE_URL, map_model_to_openrouter(model)
accepted = ", ".join(list(primary_keys) + ["OPENROUTER_API_KEY"])
raise RuntimeError(
"No LLM API key found. Set one of the primary keys or the universal "
"fallback. Accepted keys: " + accepted + ". See env.example."
)
quickstart.py¶
#!/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"""
moonshot_key = os.getenv("MOONSHOT_API_KEY")
serper_key = os.getenv("SERPER_API_KEY")
print("🔍 Checking environment configuration...")
print("-" * 40)
if moonshot_key:
print("✅ MOONSHOT_API_KEY is set")
else:
print("❌ MOONSHOT_API_KEY is NOT set")
print(" Please add it to your .env file")
print(" Get API key at: https://platform.moonshot.cn/")
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.MOONSHOT_API_KEY,
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()
run_all_strategies.py¶
#!/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-9)
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()
test_strategies.py¶
test_streaming_logging.py¶
test_token_tracking.py¶
web_tools.py¶
"""
Web tools for searching and fetching web pages
"""
import json
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', '')),
'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...]"
result = {
'url': url,
'title': soup.title.string if soup.title else 'No 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")