Skip to content

feat(checks): aggregate LLM token usage on SuiteResult #2671

Description

@davidberenstein1957

Feature Request

Roll up token usage from LLM-backed steps (judges, attackers, input generators) onto SuiteResult, and print a one-line summary in the suite report when totals are non-zero.

Motivation

giskard-llm already returns Usage on completions. Suite/scan reports do not expose aggregates, so “how many tokens did this cost?” needs custom instrumentation.

Implementation plan

  1. Define a small aggregate model (reuse or wrap giskard.llm Usage fields):

    class SuiteUsage(BaseModel):
        prompt_tokens: int = 0
        completion_tokens: int = 0
        total_tokens: int = 0
  2. Collect during run (preferred): mutable collector on run context incremented wherever acompletion / generator usage is known. Fallback: sum from check details if usage is already stored there.

  3. Attach to SuiteResult as usage: SuiteUsage (default zeros).

  4. Extend print_report() footer when total_tokens > 0.

  5. Tests with a fake generator that records known usage → aggregate matches; pure string-check suites stay at zero / omit footer line.

API usage

from giskard.checks import Scenario, Suite, Equals

suite = Suite().append(
    Scenario("no-llm")
    .interact(inputs="hi", outputs="hi")
    .check(Equals(expected="hi"))  # after key-default issue; or key=...
)
result = await suite.run()
assert result.usage.total_tokens == 0

# After an LLM-judge / scan run:
result = await vulnerability_scan(target=my_agent, description="...", languages=["en"])
print(result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens)

# Programmatic gate (optional user pattern — not built-in budget)
if result.usage.total_tokens > 500_000:
    raise RuntimeError("eval exceeded token budget")

Example report footer

SUMMARY
  36 scenarios  pass=30  fail=5  error=1  pass_rate=83.3%
  tokens: in=128400  out=39200  total=167600

Example JSON (suite dump fragment)

{
  "usage": {
    "prompt_tokens": 128400,
    "completion_tokens": 39200,
    "total_tokens": 167600
  }
}

Dollars / pricing maps are out of scope for v1.

Acceptance criteria

  • LLM-backed runs expose aggregated token counts on SuiteResult when usage was recorded
  • Report includes a one-line usage summary when total_tokens > 0
  • Unit test: fake recorded usage → aggregate matches
  • Non-LLM suites: zeros and no misleading footer

Out of scope

  • Hard spend budgets / early-stop
  • Multi-provider $ tables
  • Generation progress UI (separate)

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions