mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Add Query Echo: re-inject user question after compressed tool outputs
Query Echo addresses attention decay in compressed contexts. After SmartCrusher compresses tool outputs, the user's question may be thousands of tokens away. Echo appends a brief reminder after the last compressed block. - New: headroom/transforms/query_echo.py - Compression-ratio-proportional: only triggers when >30% compressed - Cache-safe: appended at end (after KV cache boundary) - Provider-agnostic: Anthropic, OpenAI, Gemini - 18 tests (15 unit + 3 integration with real API) - Fix _crush_array 4-tuple unpack in test_critical_fixes.py
This commit is contained in:
chopratejas
parent
729cc035a4
commit
d4d8dd0c26
5 changed files with 466 additions and 34 deletions
|
|
@ -1777,6 +1777,15 @@ class HeadroomProxy:
|
|||
f"[{request_id}] Memory: Added beta header: {key}={headers[key]}"
|
||||
)
|
||||
|
||||
# Query Echo: re-inject user's question after compressed tool outputs
|
||||
# Helps LLM attend to the question after reading dense compressed data
|
||||
if tokens_saved > 0:
|
||||
from headroom.transforms.query_echo import extract_user_query, inject_query_echo
|
||||
|
||||
user_query = extract_user_query(messages) # From original messages
|
||||
if inject_query_echo(optimized_messages, user_query, tokens_saved, original_tokens):
|
||||
logger.debug(f"[{request_id}] Query echo injected after compression")
|
||||
|
||||
# Update body
|
||||
body["messages"] = optimized_messages
|
||||
if tools is not None:
|
||||
|
|
@ -4103,6 +4112,14 @@ class HeadroomProxy:
|
|||
f"[{request_id}] CCR: Tool already present (MCP?), skipped injection for hashes: {injector.detected_hashes}"
|
||||
)
|
||||
|
||||
# Query Echo: re-inject user's question after compressed tool outputs
|
||||
if tokens_saved > 0:
|
||||
from headroom.transforms.query_echo import extract_user_query, inject_query_echo
|
||||
|
||||
user_query = extract_user_query(messages)
|
||||
if inject_query_echo(optimized_messages, user_query, tokens_saved, original_tokens):
|
||||
logger.debug(f"[{request_id}] Query echo injected after compression")
|
||||
|
||||
body["messages"] = optimized_messages
|
||||
if tools is not None:
|
||||
body["tools"] = tools
|
||||
|
|
@ -5201,6 +5218,14 @@ class HeadroomProxy:
|
|||
tokens_saved = original_tokens - optimized_tokens
|
||||
optimization_latency = (time.time() - start_time) * 1000
|
||||
|
||||
# Query Echo: re-inject user's question after compressed tool outputs
|
||||
if tokens_saved > 0:
|
||||
from headroom.transforms.query_echo import extract_user_query, inject_query_echo
|
||||
|
||||
user_query = extract_user_query(messages)
|
||||
if inject_query_echo(optimized_messages, user_query, tokens_saved, original_tokens):
|
||||
logger.debug(f"[{request_id}] Query echo injected after Gemini compression")
|
||||
|
||||
# Convert back to Gemini format if optimized
|
||||
if optimized_messages != messages:
|
||||
optimized_contents, optimized_system = self._messages_to_gemini_contents(
|
||||
|
|
|
|||
123
headroom/transforms/query_echo.py
Normal file
123
headroom/transforms/query_echo.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
"""Query Echo — re-injects the user's question after compressed tool outputs.
|
||||
|
||||
After Headroom compresses tool outputs, the user's original question may be
|
||||
thousands of tokens away from where the LLM generates its answer. Attention
|
||||
to the question decays with distance (the "lost in the middle" problem).
|
||||
|
||||
Query Echo solves this by appending a brief reminder of the user's question
|
||||
after the last compressed content block. This:
|
||||
|
||||
1. Improves answer quality from compressed data (fresh attention on the question)
|
||||
2. Makes CCR retrieval triggers more accurate (LLM is more aware of what's missing)
|
||||
3. Feeds cleaner signals to TOIN (more accurate retrievals → better learning)
|
||||
|
||||
The echo is:
|
||||
- Cache-safe: always placed after the cache boundary (in the "new tokens" region)
|
||||
- Compression-ratio-proportional: only triggers when compression was >30%
|
||||
- Provider-agnostic: appends to content strings, works for all providers
|
||||
- Cheap: ~50 tokens overhead, vs thousands of tokens saved by compression
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def extract_user_query(messages: list[dict]) -> str:
|
||||
"""Extract the last substantive user question from messages.
|
||||
|
||||
Scans backwards through messages to find the last user message
|
||||
that contains a real question (not just "yes", "continue", etc.).
|
||||
|
||||
Works for both Anthropic and OpenAI message formats.
|
||||
"""
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") != "user":
|
||||
continue
|
||||
|
||||
content = msg.get("content", "")
|
||||
|
||||
# Anthropic: content can be a list of blocks
|
||||
if isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text = str(block.get("text", "")).strip()
|
||||
if len(text) > 10: # Skip trivial messages
|
||||
return text
|
||||
continue
|
||||
|
||||
# OpenAI/generic: content is a string
|
||||
if isinstance(content, str):
|
||||
text = content.strip()
|
||||
if len(text) > 10:
|
||||
return text
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
def inject_query_echo(
|
||||
messages: list[dict],
|
||||
query: str,
|
||||
tokens_saved: int,
|
||||
original_tokens: int,
|
||||
min_compression_ratio: float = 0.3,
|
||||
max_query_chars: int = 200,
|
||||
) -> bool:
|
||||
"""Inject a query reminder after the last compressed tool output.
|
||||
|
||||
Args:
|
||||
messages: The optimized messages list (will be modified in-place).
|
||||
query: The user's question to echo.
|
||||
tokens_saved: Tokens saved by compression.
|
||||
original_tokens: Original token count before compression.
|
||||
min_compression_ratio: Minimum compression ratio to trigger echo (0.0–1.0).
|
||||
max_query_chars: Maximum characters of query to include in echo.
|
||||
|
||||
Returns:
|
||||
True if echo was injected, False if skipped.
|
||||
"""
|
||||
if not query or original_tokens == 0:
|
||||
return False
|
||||
|
||||
compression_ratio = tokens_saved / original_tokens
|
||||
if compression_ratio < min_compression_ratio:
|
||||
return False
|
||||
|
||||
# Truncate and sanitize query
|
||||
safe_query = query[:max_query_chars].replace("\n", " ").replace('"', "'").strip()
|
||||
if not safe_query:
|
||||
return False
|
||||
|
||||
echo = f"\n\n[Recall: User asked: '{safe_query}']"
|
||||
|
||||
# Walk backwards through messages, find last compressed content
|
||||
for msg in reversed(messages):
|
||||
content = msg.get("content", "")
|
||||
|
||||
# String content (OpenAI tool messages, Anthropic string content)
|
||||
if isinstance(content, str) and "compressed" in content:
|
||||
msg["content"] = content + echo
|
||||
logger.debug(
|
||||
"Query echo injected (ratio=%.1f%%, query=%s...)",
|
||||
compression_ratio * 100,
|
||||
safe_query[:50],
|
||||
)
|
||||
return True
|
||||
|
||||
# List of content blocks (Anthropic format)
|
||||
if isinstance(content, list):
|
||||
for block in reversed(content):
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
block_content = block.get("content", "")
|
||||
if isinstance(block_content, str) and "compressed" in block_content:
|
||||
block["content"] = block_content + echo
|
||||
logger.debug(
|
||||
"Query echo injected in content block (ratio=%.1f%%)",
|
||||
compression_ratio * 100,
|
||||
)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
@ -165,36 +165,6 @@ class TestSummaryHelpfulness:
|
|||
|
||||
def test_code_summary_helps_identify_functions(self):
|
||||
"""LLM can identify which functions were removed from compressed code."""
|
||||
original_code = '''
|
||||
class PaymentProcessor:
|
||||
"""Processes payments via Stripe."""
|
||||
|
||||
def __init__(self, api_key: str):
|
||||
self.stripe = stripe.Client(api_key)
|
||||
self.retry_count = 3
|
||||
|
||||
def charge(self, amount: float, currency: str, token: str) -> dict:
|
||||
for attempt in range(self.retry_count):
|
||||
try:
|
||||
return self.stripe.charges.create(
|
||||
amount=int(amount * 100),
|
||||
currency=currency,
|
||||
source=token,
|
||||
)
|
||||
except stripe.RateLimitError:
|
||||
time.sleep(2 ** attempt)
|
||||
raise PaymentError("Max retries exceeded")
|
||||
|
||||
def refund(self, charge_id: str, amount: float = None) -> dict:
|
||||
params = {"charge": charge_id}
|
||||
if amount:
|
||||
params["amount"] = int(amount * 100)
|
||||
return self.stripe.refunds.create(**params)
|
||||
|
||||
def get_balance(self) -> float:
|
||||
balance = self.stripe.balance.retrieve()
|
||||
return balance.available[0].amount / 100
|
||||
'''
|
||||
compressed_code = '''
|
||||
class PaymentProcessor:
|
||||
"""Processes payments via Stripe."""
|
||||
|
|
@ -215,9 +185,15 @@ class PaymentProcessor:
|
|||
# [2 lines omitted]
|
||||
pass
|
||||
'''
|
||||
from headroom.transforms.compression_summary import summarize_removed_code
|
||||
from headroom.transforms.compression_summary import summarize_compressed_code
|
||||
|
||||
code_summary = summarize_removed_code(original_code, compressed_code)
|
||||
# Use AST-based summary (language-agnostic)
|
||||
bodies = [
|
||||
("def charge(self, amount: float, currency: str, token: str) -> dict:", "...", 10),
|
||||
("def refund(self, charge_id: str, amount: float = None) -> dict:", "...", 20),
|
||||
("def get_balance(self) -> float:", "...", 30),
|
||||
]
|
||||
code_summary = summarize_compressed_code(bodies, 3)
|
||||
|
||||
prompt = f"Here is a compressed Python file:\n\n```python\n{compressed_code}\n```\n\n"
|
||||
if code_summary:
|
||||
|
|
|
|||
|
|
@ -349,7 +349,7 @@ class TestSmartCrusherTOINIntegration:
|
|||
len(toin._patterns)
|
||||
|
||||
# Crush the array
|
||||
result, info, markers = crusher._crush_array(
|
||||
result, info, markers, _summary = crusher._crush_array(
|
||||
items, query_context="test query", tool_name="test_tool"
|
||||
)
|
||||
|
||||
|
|
@ -418,7 +418,7 @@ class TestAllFixesIntegrated:
|
|||
]
|
||||
|
||||
# Step 1: Compress
|
||||
result, info, markers = crusher._crush_array(
|
||||
result, info, markers, _summary = crusher._crush_array(
|
||||
items, query_context="find status", tool_name="api_response"
|
||||
)
|
||||
|
||||
|
|
|
|||
308
tests/test_query_echo.py
Normal file
308
tests/test_query_echo.py
Normal file
|
|
@ -0,0 +1,308 @@
|
|||
"""Tests for Query Echo — unit tests + integration eval with real API.
|
||||
|
||||
Unit tests: verify echo injection logic across message formats.
|
||||
Integration: verify echo actually improves LLM answers on compressed data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.transforms.query_echo import extract_user_query, inject_query_echo
|
||||
|
||||
# Load .env for integration tests
|
||||
env_path = Path(__file__).parent.parent / ".env"
|
||||
if env_path.exists():
|
||||
for line in env_path.read_text().splitlines():
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#") and "=" in line:
|
||||
key, _, value = line.partition("=")
|
||||
os.environ.setdefault(key.strip(), value.strip())
|
||||
|
||||
ANTHROPIC_KEY = os.environ.get("ANTHROPIC_API_KEY", "")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unit Tests: extract_user_query
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestExtractUserQuery:
|
||||
def test_simple_string_message(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "What are the test failures?"},
|
||||
{"role": "assistant", "content": "Let me check."},
|
||||
]
|
||||
assert extract_user_query(messages) == "What are the test failures?"
|
||||
|
||||
def test_anthropic_content_blocks(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Find the auth bug in these results"},
|
||||
{"type": "tool_result", "tool_use_id": "x", "content": "data..."},
|
||||
],
|
||||
},
|
||||
]
|
||||
assert extract_user_query(messages) == "Find the auth bug in these results"
|
||||
|
||||
def test_skips_trivial_messages(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "What are the top 3 errors in this log output?"},
|
||||
{"role": "assistant", "content": "Looking..."},
|
||||
{"role": "user", "content": "yes"},
|
||||
]
|
||||
# "yes" is too short (<10 chars), should get the earlier question
|
||||
assert "errors" in extract_user_query(messages)
|
||||
|
||||
def test_empty_messages(self):
|
||||
assert extract_user_query([]) == ""
|
||||
|
||||
def test_no_user_messages(self):
|
||||
messages = [{"role": "assistant", "content": "Hello"}]
|
||||
assert extract_user_query(messages) == ""
|
||||
|
||||
def test_multi_turn_gets_last_substantive(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "Tell me about the weather"},
|
||||
{"role": "assistant", "content": "It's sunny."},
|
||||
{"role": "user", "content": "What about the authentication failures in CI?"},
|
||||
{"role": "assistant", "content": "Let me check."},
|
||||
]
|
||||
assert "authentication" in extract_user_query(messages)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unit Tests: inject_query_echo
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestInjectQueryEcho:
|
||||
def test_injects_after_compressed_content(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "What are the failures?"},
|
||||
{
|
||||
"role": "tool",
|
||||
"content": '[{"status":"pass"}]\n[50 items compressed to 5. Retrieve more: hash=abc123]',
|
||||
},
|
||||
]
|
||||
result = inject_query_echo(
|
||||
messages, "What are the failures?", tokens_saved=500, original_tokens=1000
|
||||
)
|
||||
assert result is True
|
||||
assert "[Recall:" in messages[-1]["content"]
|
||||
assert "failures" in messages[-1]["content"]
|
||||
|
||||
def test_skips_when_no_compression(self):
|
||||
messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "tool", "content": "some data"},
|
||||
]
|
||||
result = inject_query_echo(messages, "hello", tokens_saved=0, original_tokens=100)
|
||||
assert result is False
|
||||
|
||||
def test_skips_when_low_compression(self):
|
||||
messages = [
|
||||
{"role": "tool", "content": "data\n[10 items compressed to 8]"},
|
||||
]
|
||||
result = inject_query_echo(messages, "query", tokens_saved=20, original_tokens=100)
|
||||
# 20% compression — below 30% threshold
|
||||
assert result is False
|
||||
|
||||
def test_triggers_at_high_compression(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "tool",
|
||||
"content": "data\n[500 items compressed to 20. Retrieve more: hash=abc]",
|
||||
},
|
||||
]
|
||||
result = inject_query_echo(
|
||||
messages, "What are the errors?", tokens_saved=800, original_tokens=1000
|
||||
)
|
||||
assert result is True
|
||||
assert "[Recall:" in messages[0]["content"]
|
||||
|
||||
def test_anthropic_content_blocks(self):
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "1",
|
||||
"content": '[{"x":1}]\n[100 items compressed to 10. hash=abc123]',
|
||||
},
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "2",
|
||||
"content": '[{"y":2}]\n[200 items compressed to 15. hash=def456]',
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
result = inject_query_echo(
|
||||
messages, "Find the errors", tokens_saved=600, original_tokens=1000
|
||||
)
|
||||
assert result is True
|
||||
# Echo should be on the LAST compressed block
|
||||
last_block = messages[0]["content"][-1]
|
||||
assert "[Recall:" in last_block["content"]
|
||||
|
||||
def test_truncates_long_query(self):
|
||||
long_query = "x" * 500
|
||||
messages = [
|
||||
{"role": "tool", "content": "data\n[100 items compressed to 5. hash=abc]"},
|
||||
]
|
||||
inject_query_echo(messages, long_query, tokens_saved=500, original_tokens=1000)
|
||||
# Should be truncated to 200 chars
|
||||
assert len(messages[0]["content"]) < 800
|
||||
|
||||
def test_sanitizes_quotes(self):
|
||||
messages = [
|
||||
{"role": "tool", "content": "data\n[100 items compressed to 5. hash=abc]"},
|
||||
]
|
||||
inject_query_echo(
|
||||
messages, 'Find "auth_error" in logs', tokens_saved=500, original_tokens=1000
|
||||
)
|
||||
# Double quotes replaced with single quotes
|
||||
assert '"auth_error"' not in messages[0]["content"]
|
||||
assert "'auth_error'" in messages[0]["content"]
|
||||
|
||||
def test_no_echo_without_compressed_marker(self):
|
||||
"""If no message contains 'compressed', no echo injected."""
|
||||
messages = [
|
||||
{"role": "tool", "content": "just plain tool output without compression"},
|
||||
]
|
||||
result = inject_query_echo(messages, "query", tokens_saved=500, original_tokens=1000)
|
||||
assert result is False
|
||||
|
||||
def test_empty_query_skipped(self):
|
||||
messages = [
|
||||
{"role": "tool", "content": "data\n[100 items compressed]"},
|
||||
]
|
||||
result = inject_query_echo(messages, "", tokens_saved=500, original_tokens=1000)
|
||||
assert result is False
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Integration Eval: Does echo improve answer quality?
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _call_claude(messages, max_tokens=200):
|
||||
import httpx
|
||||
|
||||
resp = httpx.post(
|
||||
"https://api.anthropic.com/v1/messages",
|
||||
headers={
|
||||
"X-Api-Key": ANTHROPIC_KEY,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"model": "claude-sonnet-4-5-20250929",
|
||||
"max_tokens": max_tokens,
|
||||
"messages": messages,
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _make_compressed_test_data():
|
||||
"""Create realistic compressed test results where answer is in omitted data."""
|
||||
# Visible: 10 passing tests
|
||||
visible = [
|
||||
{"test": f"test_module_{i}", "status": "passed", "duration_ms": 50 + i} for i in range(10)
|
||||
]
|
||||
compressed_output = json.dumps(visible, indent=2)
|
||||
|
||||
# The important info is in the COMPRESSED summary
|
||||
summary = "87 passed, 2 failed, 1 error; notable: test_auth (fail); test_db (timeout)"
|
||||
compressed_output += (
|
||||
f"\n[90 items compressed to 10. Omitted: {summary}."
|
||||
f" Retrieve more: hash=abc123def456789012345678."
|
||||
f" Expires in 5m.]"
|
||||
)
|
||||
return compressed_output
|
||||
|
||||
|
||||
@pytest.mark.skipif(not ANTHROPIC_KEY, reason="ANTHROPIC_API_KEY not set")
|
||||
class TestQueryEchoIntegration:
|
||||
def test_with_echo_finds_failures(self):
|
||||
"""WITH echo → LLM should identify failures from compressed data."""
|
||||
compressed = _make_compressed_test_data()
|
||||
question = "What are the test failures and what went wrong?"
|
||||
|
||||
# Add echo at the end
|
||||
echoed_content = compressed + f"\n\n[Recall: User asked: '{question}']"
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Here are the CI test results:\n\n{echoed_content}\n\nAnswer my question.",
|
||||
},
|
||||
]
|
||||
resp = _call_claude(messages)
|
||||
text = resp.get("content", [{}])[0].get("text", "").lower()
|
||||
|
||||
print(f"\n WITH echo response: {text[:200]}")
|
||||
|
||||
has_failure_info = any(w in text for w in ["fail", "error", "timeout", "auth", "db"])
|
||||
assert has_failure_info, f"LLM missed failures with echo: {text[:300]}"
|
||||
|
||||
def test_without_echo_baseline(self):
|
||||
"""WITHOUT echo → baseline for comparison."""
|
||||
compressed = _make_compressed_test_data()
|
||||
question = "What are the test failures and what went wrong?"
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Here are the CI test results:\n\n{compressed}\n\n{question}",
|
||||
},
|
||||
]
|
||||
resp = _call_claude(messages)
|
||||
text = resp.get("content", [{}])[0].get("text", "").lower()
|
||||
|
||||
print(f"\n WITHOUT echo response: {text[:200]}")
|
||||
|
||||
# Not asserting — this is the baseline. May or may not find failures.
|
||||
has_failure_info = any(w in text for w in ["fail", "error", "timeout"])
|
||||
print(f" Found failure info without echo: {has_failure_info}")
|
||||
|
||||
def test_echo_with_specific_lookup(self):
|
||||
"""Echo helps LLM find specific data in compressed results."""
|
||||
# Compressed config data — visible items are all 'production'
|
||||
visible = [
|
||||
{"env": "production", "key": "DB_URL", "value": "postgres://prod:5432/app"},
|
||||
{"env": "production", "key": "REDIS_URL", "value": "redis://prod:6379"},
|
||||
{"env": "production", "key": "API_KEY", "value": "sk-prod-xxx"},
|
||||
]
|
||||
compressed = json.dumps(visible, indent=2)
|
||||
compressed += (
|
||||
"\n[45 items compressed to 3. Omitted: 30 staging, 12 development."
|
||||
" Retrieve more: hash=cfg123def456789012345678."
|
||||
" Expires in 5m.]"
|
||||
)
|
||||
|
||||
question = "What is the staging database URL?"
|
||||
echoed = compressed + f"\n\n[Recall: User asked: '{question}']"
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": f"Application configs:\n\n{echoed}\n\nAnswer precisely."},
|
||||
]
|
||||
resp = _call_claude(messages)
|
||||
text = resp.get("content", [{}])[0].get("text", "").lower()
|
||||
|
||||
print(f"\n Staging lookup with echo: {text[:200]}")
|
||||
|
||||
# With echo + summary mentioning "30 staging", LLM should know
|
||||
# staging data exists in compressed items
|
||||
knows_staging = "staging" in text or "compressed" in text or "retrieve" in text
|
||||
assert knows_staging, f"LLM didn't reference staging data: {text[:300]}"
|
||||
Loading…
Add table
Add a link
Reference in a new issue