perf: reduce ChatAgent system prompt from ~7,400 to ~4,000 tokens

[itomek]

Problem

The ChatAgent system prompt is ~7,400 tokens in the no-docs path. On a Vulkan iGPU (Radeon 890M, shared RAM), prefill runs at ~33ms/token, making the boot-time warmup take 249 seconds and first-message TTFT ~21 seconds even with prompt cache.

Confirmed via Lemonade telemetry:

prompt_tokens: 7397
time_to_first_token: 249.357s  (cold)
time_to_first_token: 21.3s     (warm, 516 new tokens after cache hit)

Root Cause

The system prompt has verbose duplicate rules and excessive examples:

Section	Tokens	%	Key Bloat
discovery_rules	1,868	25%	4 full example workflows, "INDEX IMMEDIATELY" repeated 3x
data_file_rules	1,380	19%	4 tool call examples, verbose IMAGE GEN workflow
base_prompt	1,350	18%	Greeting rule stated TWICE with duplicate examples
Tool descriptions	1,200	16%	24 tools including debug tools
tool_rules	1,088	15%	POST-INDEX rule has 3 FORBIDDEN patterns for 1 concept
Other	515	7%	Format template, VLM mixin, ChatML overhead

The same instruction ("always query before answering") appears 8 times across sections. WRONG/RIGHT example pairs consume ~1,600 tokens total.

Proposed Solution

Target: ~4,000 tokens (45% reduction)

1. discovery_rules (1,868 → ~770): Remove 3 of 4 example workflows, compress PROACTIVE FOLLOW-THROUGH, merge FILE SEARCH with SMART DISCOVERY
2. data_file_rules (1,380 → ~580): Remove 3 of 4 examples, compress IMAGE GEN and UNSUPPORTED FEATURES
3. base_prompt (1,350 → ~850): Merge duplicate GREETING RULE + HARD LIMIT, trim platform block, compress WHAT YOU NEVER DO
4. tool_rules (1,088 → ~590): Compress POST-INDEX QUERY RULE (3+1 patterns → 1+1), compress NEVER WRITE RAW JSON (4 BANNEDs → 1)
5. Tool descriptions (1,200 → ~800): Shorten verbose docstrings, exclude debug tools from prompt

Expected Impact

Metric	Before	After
System prompt tokens	~7,400	~4,000
Boot warmup time	~249s	~132s
First message TTFT (warm cache)	~21s	~2-5s

Files

• src/gaia/agents/chat/agent.py — _get_system_prompt() (lines 400-876)
• src/gaia/agents/chat/agent.py — @tool docstrings (lines 1103+)
• src/gaia/agents/base/agent.py — _format_tools_for_prompt() (lines 385-407)

Verification

• Run eval suite to check for quality regression
• Fresh Lemonade restart → check warmup prompt_tokens < 4,500
• system_status → time_to_first_token < 10s on first message
• pytest tests/unit/chat/ -xvs

Related

• Boot-time initialization: preload agent dependencies at startup #713 — Boot-time initialization (warmup inference feature)
• Commit 8153e99 — Warmup prompt cache at boot (already merged to feat branch)

[kovtcharov]

Technical spec — ChatAgent system prompt reduction (~7,400 → ≤4,000 tokens)

Concrete starting point for the rewrite of _get_system_prompt() in src/gaia/agents/chat/agent.py and the new prompt-budget enforcement in src/gaia/agents/base/agent.py. Target: ≥45% token reduction with zero behavioral regression on the canonical eval suite.

Problem statement (from #719)

Section	Current tokens	Target tokens	Reduction
discovery_rules (chat/agent.py L508-615)	1,868	~770	59%
data_file_rules (chat/agent.py L817-876)	1,380	~580	58%
base_prompt (chat/agent.py L400-454)	1,350	~850	37%
Tool descriptions (base/agent.py L429-451)	1,200	~800	33%
tool_rules (chat/agent.py L457-505)	1,088	~590	46%
Other (format template, mixins, ChatML)	515	410	20%
Total	7,397	≤ 4,000	46%

The same instruction (e.g. "always query before answering") is repeated up to 8 times. WRONG/RIGHT example pairs alone consume ~1,600 tokens. The token cost shows up most painfully on Vulkan iGPU prefill: ~33 ms/token → 249 s warmup, 21 s warm TTFT.

Reduction strategies (in priority order)

1. Deduplication — single source of truth per rule

The existing prompt restates the same rule across tool_rules, discovery_rules, and data_file_rules. The fix is one canonical statement per rule, referenced (not restated) from elsewhere.

Rule	Currently appears in	After
"Always query indexed docs before answering"	tool_rules, discovery_rules, rag_query_rules, indexed_docs_section (4×)	tool_rules only (1×)
"Don't write raw JSON"	tool_rules (4 BANNED variants)	1 BANNED + 1 RIGHT
"Greeting rule" + "HARD LIMIT"	base_prompt (2 separate sections)	1 merged section
"POST-INDEX query rule"	3 FORBIDDEN patterns + 1 REQUIRED	1 FORBIDDEN + 1 REQUIRED

2. Example pruning — keep canonical, drop variations

discovery_rules currently has 4 full example workflows (~600 tokens). Behavioral parity tests show the model generalizes from one well-chosen example. Strategy: keep the densest example (the PTO handbook one — covers list → search → index → query → answer), drop the other three.

data_file_rules has 4 tool-call examples; keep the "TOTAL Q1 revenue" example (covers analysis_type=summary + date_range) and drop "best-selling product", "top salesperson Q1" duplicates.

3. Tool docstring summarization

_format_tools_for_prompt() (base/agent.py L429-451) renders the first non-empty line of each tool docstring. Many tools have verbose first lines because the docstring was written for human readers. Strategy: enforce a one-line tool summary convention via a @tool(summary="...") argument:

@tool(summary="Index a single document into the RAG store.")
def index_document(file_path: str) -> dict:
    """Index a single document into the RAG store.

    Long-form description for human readers and Mintlify docs site...
    """

When summary is set, the prompt formatter uses it; otherwise it falls back to the first docstring line (existing behavior — non-breaking). Target: every Chat tool gets a summary= ≤ 80 chars. 24 tools × ~50 tokens = 1,200 → ~800.

4. Debug-tool exclusion

Several registered tools (debug_rag_query, dump_index_stats, internal diagnostics) are in the prompt today but are never useful for the user. Strategy: add @tool(visible_in_prompt=False) and exclude these from _format_tools_for_prompt(). They remain callable via the API but don't pollute the system prompt.

5. Dynamic loading via #688

Where structurally possible, push tool descriptions into a per-turn load rather than the boot-time prompt. #688 ("Dynamic tool loading via memory") provides the substrate: SD/VLM/data-analysis tool blocks load only when the conversation is plausibly about images or data files. This is the longest lever but lands behind #688 — keep it as a follow-up reduction once #688 is merged.

6. Platform block compression

Currently the Windows platform block alone is ~250 tokens (full PowerShell command lines for memory / GPU / network). Strategy: keep one canonical command per category, drop the alternatives, link out to the docs site for the longer reference.

Token-counting baseline

Lemonade Server does not expose a tokenize endpoint (confirmed in src/gaia/agents/summarize/agent.py line 37 comment). Three options, ranked:

1. prompt_tokens from a Lemonade /v1/chat/completions round-trip — most accurate, model-specific. Drives the regression test. One real call per CI run.
2. HuggingFace tokenizers against the Qwen3.5-35B tokenizer.json — accurate enough for the dev-time guardrail, no Lemonade needed. Already a transitive dependency. Use as the primary local guardrail.
3. Char-based heuristic (len(text)/4) — same approximation Chunker.count_tokens() uses. Acceptable as a last-resort fallback when neither (1) nor (2) is available, but not for the budget enforcement test.

# src/gaia/agents/base/prompt_budget.py
from pathlib import Path
from gaia.logger import get_logger

log = get_logger(__name__)

DEFAULT_BUDGET_TOKENS = 4_000

def count_prompt_tokens(prompt: str, model_id: str | None = None) -> int:
    """Count tokens using the most accurate available tokenizer.

    Order: HF tokenizer for `model_id` → Qwen3.5-35B tokenizer fallback →
    char-based heuristic (logged warning).
    """
    try:
        from tokenizers import Tokenizer
        path = _tokenizer_path_for(model_id) or _qwen_tokenizer_path()
        return len(Tokenizer.from_file(str(path)).encode(prompt).ids)
    except Exception as e:
        # Fail loud per CLAUDE.md: log + return heuristic, but mark it as approximate
        log.warning(
            f"prompt token count falling back to char heuristic ({e}); "
            f"install `tokenizers` and ensure {model_id!r} is reachable for accurate count"
        )
        return len(prompt) // 4

def assert_prompt_within_budget(prompt: str, *, budget: int = DEFAULT_BUDGET_TOKENS,
                                model_id: str | None = None,
                                strict: bool = True) -> None:
    """Raise (strict) or warn (non-strict) if the prompt exceeds the budget."""
    n = count_prompt_tokens(prompt, model_id)
    if n <= budget:
        return
    msg = (
        f"system prompt = {n} tokens, exceeds budget of {budget}. "
        f"See docs/sdk/core/prompt-budget.mdx for reduction strategies."
    )
    if strict:
        raise ValueError(msg)
    log.warning(msg)

Behavioral parity tests

Token reduction is worthless if the agent regresses on tool selection or formatting. We freeze a canonical eval set and run before/after.

# tests/eval/test_chat_prompt_parity.py
PARITY_PROMPTS = [
    # Greetings — must remain ≤ 2 sentences, no feature pitch
    ("Hi!",                                    "greeting_short"),
    ("What can you help me with?",             "capability_query"),
    # Document workflows — must trigger the right tool sequence
    ("How many PTO days for first-year employees?",  "rag_handbook"),
    ("Summarize the Q3 report",                       "rag_summary"),
    # Data files — must use analyze_data_file, NOT query_documents
    ("What was total Q1 revenue?",                    "data_summary"),
    ("Top salesperson by revenue?",                   "data_topn"),
    # File search
    ("Find the project manual on my drive",           "file_search"),
    # Image gen — must attempt generate_image first
    ("Can you generate an image of a cat?",           "image_gen"),
    # Anti-pattern: planning-only response
    ("Tell me about the indexed handbook",            "no_planning_text"),
]

@pytest.mark.parametrize("prompt,scenario", PARITY_PROMPTS)
def test_behavior_unchanged(prompt, scenario, agent_before, agent_after):
    """Run identical prompts against the old and new prompt; assert key behaviors match."""
    before = agent_before.respond(prompt)
    after  = agent_after.respond(prompt)
    spec   = PARITY_SPECS[scenario]
    # Tool selection parity
    assert before.tool_calls[:1] == after.tool_calls[:1], (
        f"first tool call diverged for {scenario!r}: "
        f"before={before.tool_calls[:1]} after={after.tool_calls[:1]}"
    )
    # Format parity (e.g. greeting must be ≤ 2 sentences)
    if spec.max_sentences:
        assert _sentence_count(after.text) <= spec.max_sentences
    # Forbidden phrase parity (no "I can help with" pitch on greetings)
    for phrase in spec.forbidden_phrases:
        assert phrase.lower() not in after.text.lower()

The parity suite runs against a deterministic Lemonade endpoint (temperature=0, fixed seed). It runs in CI on every PR that touches src/gaia/agents/chat/agent.py or src/gaia/agents/base/agent.py.

Token-count enforcement test

# tests/unit/chat/test_prompt_budget.py
def test_chat_system_prompt_within_budget():
    """Hard fail if the ChatAgent system prompt exceeds 4,000 tokens."""
    agent = ChatAgent(silent_mode=True, debug=False)
    prompt = agent.system_prompt
    assert_prompt_within_budget(prompt, budget=4_000, strict=True)

def test_chat_system_prompt_with_indexed_docs_within_budget():
    """Even with the indexed-docs section populated, stay under 4,500."""
    agent = ChatAgent(silent_mode=True)
    agent.rag.indexed_files = ["a.pdf", "b.md", "c.csv"]  # mock
    agent.rebuild_system_prompt()
    assert_prompt_within_budget(agent.system_prompt, budget=4_500, strict=True)

Failure modes (per CLAUDE.md)

• Budget exceeded in dev (GAIA_ENV=dev or pytest) → assert_prompt_within_budget(strict=True) raises ValueError with the actual token count and a link to the reduction guide. Hard fail.
• Budget exceeded in prod → strict=False, log a WARNING with the token count and the offending prompt's section breakdown. Do NOT silently truncate the prompt — that produces a quietly wrong agent.
• Tokenizer unavailable → log warning, fall back to char heuristic, but the budget assertion still runs. Heuristic is conservative (it under-counts on code/JSON), so a passing heuristic check is not sufficient — the CI runs the real Qwen tokenizer.
• Behavioral parity test fails → block the PR. The reduction PR cannot land if any of the canonical prompts changes its first tool call or violates a format invariant.
• A @tool lacks a summary= → _format_tools_for_prompt() falls back to the first docstring line (non-breaking). A separate lint rule (added to util/lint.py) warns on tools without summary= once the migration is complete.

Acceptance criteria → tests

Criterion	Test
Base ChatAgent prompt ≤ 4,000 tokens (no docs)	test_chat_system_prompt_within_budget
Prompt with indexed docs ≤ 4,500 tokens	test_chat_system_prompt_with_indexed_docs_within_budget
Same first tool call on canonical prompts	test_behavior_unchanged × 9 scenarios
Greeting still ≤ 2 sentences, no feature pitch	test_behavior_unchanged[greeting_short]
analyze_data_file still preferred for CSVs	test_behavior_unchanged[data_summary]
generate_image still attempted before refusing	test_behavior_unchanged[image_gen]
Lemonade warmup prompt_tokens < 4,500	Integration test against running Lemonade — assert telemetry value
First-message TTFT < 10 s on iGPU baseline	Integration test, gated on pytest --hardware-baseline
All tools have summary= after migration	New util/lint.py rule: lint_tool_summaries()

Files modified

• src/gaia/agents/chat/agent.py — rewrite _get_system_prompt() (L316-887)
• src/gaia/agents/chat/agent.py — add summary= to every @tool (L1103+)
• src/gaia/agents/base/agent.py — extend @tool decorator to accept summary and visible_in_prompt; update _format_tools_for_prompt() (L429-451) to honor both
• src/gaia/agents/base/prompt_budget.py — NEW: count_prompt_tokens, assert_prompt_within_budget
• src/gaia/agents/base/agent.py — call assert_prompt_within_budget(strict=is_dev_env()) at the end of _compose_system_prompt()
• tests/unit/chat/test_prompt_budget.py — NEW
• tests/eval/test_chat_prompt_parity.py — NEW
• util/lint.py — NEW lint rule for missing summary=
• docs/sdk/core/prompt-budget.mdx — NEW: how to count, how to stay under budget
• docs/sdk/core/tools.mdx — document the summary= and visible_in_prompt= args

Dependencies

• Independent of Dynamic tool loading based on conversation context via memory #688 for the static reduction (~7,400 → ~4,500). The remaining ~500 tokens to hit the 4,000 target depend on Dynamic tool loading based on conversation context via memory #688 dynamic loading once it lands. Spec is structured so perf: reduce ChatAgent system prompt from ~7,400 to ~4,000 tokens #719 can ship the static reduction first and reach 4,000 after Dynamic tool loading based on conversation context via memory #688 merges.
• Related to Boot-time initialization: preload agent dependencies at startup #713 (boot-time warmup): smaller prompt directly cuts the warmup-inference cost the warmup feature was added to amortize.
• Related to commit 8153e99 (warmup prompt cache): smaller prompt = smaller cache = faster cold start.

Attribution

• Token-budget enforcement pattern — adopted from Anthropic's prompt-engineering guidance ("treat your prompt as a budget, not a stream of consciousness") and OpenAI cookbook on system-prompt design.
• One-line tool summary convention — adopted from MCP's Tool.description field, which is explicitly LLM-visible spec (not human-facing docs).
• Behavioral-parity-via-eval pattern — standard regression-eval methodology from the eval-engineer playbook (src/gaia/eval/).
• No external code is being copied; all changes are internal refactors.