diff --git a/assert_ai/cli.py b/assert_ai/cli.py index b1fbeecf5..afa07fdf1 100644 --- a/assert_ai/cli.py +++ b/assert_ai/cli.py @@ -889,6 +889,67 @@ def cli(ctx: click.Context, verbose: bool, quiet: bool, log_file: Path | None, o cli.add_command(init) +@cli.command(short_help="Estimate token usage without running a pipeline") +@click.option( + "--config", + required=True, + type=click.Path(exists=True, dir_okay=False, path_type=Path), + help="Path to a YAML pipeline config.", + show_envvar=True, +) +@click.option( + "--force-stage", + type=click.Choice(STAGE_NAMES, case_sensitive=False), + multiple=True, + help="Estimate as though the selected stage and its downstream stages were forced.", + show_envvar=True, +) +@click.option("--override", "overrides", multiple=True, help="Override a config value.") +@click.option( + "--concurrency", + type=click.IntRange(min=1), + default=None, + help="Override inference concurrency for the estimate.", + show_envvar=True, +) +@click.option( + "--output", + "output_format", + type=click.Choice(["text", "json"], case_sensitive=False), + default="text", + show_default=True, +) +def estimate( + config: Path, + force_stage: tuple[str, ...], + overrides: tuple[str, ...], + concurrency: int | None, + output_format: str, +): + """Estimate local model token usage without making provider calls.""" + + runner = _load_runner_module() + try: + payload = runner.estimate_pipeline_usage( + config=str(config), + force_stages=list(force_stage), + overrides=list(overrides), + concurrency=concurrency, + ) + except (runner.ConfigError, ValueError) as exc: + raise click.ClickException(str(exc)) from exc + + if output_format == "json": + click.echo(json.dumps(payload, ensure_ascii=False)) + elif int(payload.get("total_tokens", 0) or 0) <= 0: + click.echo("Estimated token usage: 0 tracked tokens.") + for note in payload.get("notes") or []: + if isinstance(note, str) and note: + click.echo(f"Estimate note: {note}") + else: + runner._log_token_estimate(payload) + + @cli.command(short_help="Run a pipeline from a YAML config") @click.option( "--config", diff --git a/assert_ai/core/artifact_cache.py b/assert_ai/core/artifact_cache.py index 35f10e0e6..5c3ba4e61 100644 --- a/assert_ai/core/artifact_cache.py +++ b/assert_ai/core/artifact_cache.py @@ -141,6 +141,64 @@ def supports_artifact_cache(ctx: dict[str, Any]) -> bool: return bool(ctx.get("suite_root") and ctx.get("config_path") and ctx.get("artifacts_root")) +def _matching_artifact_plan( + *, + stage_name: str, + stage_root: Path, + fingerprint: ArtifactFingerprint, +) -> ArtifactPlan | None: + match = _latest_matching_metadata(stage_name, stage_root, fingerprint.input_hash) + if match is None: + return None + version, metadata = match + artifact_dir = stage_root / version + return ArtifactPlan( + stage_name=stage_name, + version=version, + artifact_dir=artifact_dir, + output_paths=_output_paths(stage_name, artifact_dir), + fingerprint=fingerprint, + reused=True, + metadata=metadata, + ) + + +def preview_artifact_plan( + *, + ctx: dict[str, Any], + stage_name: str, + raw_cfg: dict[str, Any], + forced: bool, +) -> ArtifactPlan: + """Plan artifact reuse or generation without allocating a directory.""" + + if stage_name not in CACHEABLE_STAGES: + raise ValueError(f"unsupported cacheable stage: {stage_name}") + suite_root = Path(ctx["suite_root"]) + fingerprint = build_artifact_fingerprint(ctx=ctx, stage_name=stage_name, raw_cfg=raw_cfg) + stage_root = suite_root / ARTIFACTS_DIR / stage_name + if not forced: + match = _matching_artifact_plan( + stage_name=stage_name, + stage_root=stage_root, + fingerprint=fingerprint, + ) + if match is not None: + return match + + version = "preview" + artifact_dir = stage_root / version + return ArtifactPlan( + stage_name=stage_name, + version=version, + artifact_dir=artifact_dir, + output_paths=_output_paths(stage_name, artifact_dir), + fingerprint=fingerprint, + reused=False, + metadata=None, + ) + + def prepare_artifact_plan( *, ctx: dict[str, Any], @@ -157,19 +215,13 @@ def prepare_artifact_plan( stage_root = suite_root / ARTIFACTS_DIR / stage_name if not forced: - match = _latest_matching_metadata(stage_name, stage_root, fingerprint.input_hash) + match = _matching_artifact_plan( + stage_name=stage_name, + stage_root=stage_root, + fingerprint=fingerprint, + ) if match is not None: - version, metadata = match - artifact_dir = stage_root / version - return ArtifactPlan( - stage_name=stage_name, - version=version, - artifact_dir=artifact_dir, - output_paths=_output_paths(stage_name, artifact_dir), - fingerprint=fingerprint, - reused=True, - metadata=metadata, - ) + return match version, artifact_dir = _allocate_version_dir(stage_root) return ArtifactPlan( @@ -248,14 +300,19 @@ def override_cacheable_output_paths( return cfg -def activate_latest_artifacts(ctx: dict[str, Any]) -> None: +def activate_latest_artifacts( + ctx: dict[str, Any], + *, + read_only: bool = False, +) -> None: """Load latest artifact refs into context for run-only stage configs. When ``latest.json`` references an artifact directory that has been deleted, has lost its sidecar, or is missing one of its data files, we emit a stderr warning and try to fall back to the most recent valid - version directory for that stage (if any). A silent skip would let the - pipeline silently drift to stale legacy compatibility files. + version directory for that stage (if any). In read-only mode the selected + refs are applied to context without repairing latest.json or compatibility + files. """ suite_root = Path(ctx["suite_root"]) @@ -308,7 +365,8 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: metadata=metadata, primary_path=output_paths[next(iter(_OUTPUT_FILES[stage_name]))], ) - update_latest(ctx, stage_name, ref) + if not read_only: + update_latest(ctx, stage_name, ref) log.warning( "latest.json %s entry referenced missing paths; rebuilt " "ref pointing at the current on-disk location of version %s.", @@ -320,7 +378,8 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: for output_key, context_key in _CONTEXT_PATH_KEYS[stage_name].items(): if output_key in output_paths: ctx[context_key] = str(output_paths[output_key]) - refresh_compatibility_files(ctx, stage_name, output_paths) + if not read_only: + refresh_compatibility_files(ctx, stage_name, output_paths) continue recovery = _recover_latest_valid_version(stage_name, stage_root) @@ -351,8 +410,9 @@ def activate_latest_artifacts(ctx: dict[str, Any]) -> None: for output_key, context_key in _CONTEXT_PATH_KEYS[stage_name].items(): if output_key in recovered_outputs: ctx[context_key] = str(recovered_outputs[output_key]) - refresh_compatibility_files(ctx, stage_name, recovered_outputs) - update_latest(ctx, stage_name, recovered_ref) + if not read_only: + refresh_compatibility_files(ctx, stage_name, recovered_outputs) + update_latest(ctx, stage_name, recovered_ref) log.warning( "latest.json %s entry was missing or incomplete; " "recovered to version %s.", diff --git a/assert_ai/core/model_client.py b/assert_ai/core/model_client.py index f48bef57e..81782a508 100644 --- a/assert_ai/core/model_client.py +++ b/assert_ai/core/model_client.py @@ -38,6 +38,7 @@ import logging import os import random +import re import sys import time from contextlib import contextmanager @@ -158,40 +159,74 @@ class UsageAccumulator: invokes more than one model (e.g. test_set + stratification) can be inspected later. """ + requests: int = 0 calls: int = 0 + missing_usage_calls: int = 0 input_tokens: int = 0 output_tokens: int = 0 + total_tokens: int = 0 cached_input_tokens: int = 0 cache_creation_input_tokens: int = 0 per_model: dict[str, dict[str, int]] = field(default_factory=dict) def add(self, usage: UsageStats | None, *, model: str | None = None) -> None: """Fold one call's normalized usage into this accumulator.""" - if usage is None: - return - self.calls += 1 - ipt = int(usage.prompt_tokens or 0) - opt = int(usage.completion_tokens or 0) - cit = int(usage.cached_input_tokens or 0) - cct = int(usage.cache_creation_input_tokens or 0) - self.input_tokens += ipt - self.output_tokens += opt - self.cached_input_tokens += cit - self.cache_creation_input_tokens += cct key = model or "?" bucket = self.per_model.setdefault( key, { + "requests": 0, "calls": 0, + "missing_usage_calls": 0, "input_tokens": 0, "output_tokens": 0, + "total_tokens": 0, "cached_input_tokens": 0, "cache_creation_input_tokens": 0, }, ) - bucket["calls"] += 1 + self.requests += 1 + bucket["requests"] += 1 + if usage is None: + self.missing_usage_calls += 1 + bucket["missing_usage_calls"] += 1 + return + ipt = int(usage.prompt_tokens or 0) + opt = int(usage.completion_tokens or 0) + reported_total = ( + int(usage.total_tokens) + if usage.total_tokens is not None + else None + ) + total = ( + reported_total + if reported_total is not None and reported_total > 0 + else ipt + opt + ) + cit = int(usage.cached_input_tokens or 0) + cct = int(usage.cache_creation_input_tokens or 0) + usage_complete = ( + (reported_total is not None and reported_total > 0) + or ( + usage.prompt_tokens is not None + and usage.completion_tokens is not None + and (ipt > 0 or opt > 0) + ) + ) + if not usage_complete: + self.missing_usage_calls += 1 + bucket["missing_usage_calls"] += 1 + else: + self.calls += 1 + bucket["calls"] += 1 + self.input_tokens += ipt + self.output_tokens += opt + self.total_tokens += total + self.cached_input_tokens += cit + self.cache_creation_input_tokens += cct bucket["input_tokens"] += ipt bucket["output_tokens"] += opt + bucket["total_tokens"] += total bucket["cached_input_tokens"] += cit bucket["cache_creation_input_tokens"] += cct @@ -204,9 +239,17 @@ def cache_hit_rate(self) -> float: def to_dict(self) -> dict[str, Any]: """JSON-serializable snapshot of this accumulator.""" return { + "requests": self.requests, "calls": self.calls, + "missing_usage_calls": self.missing_usage_calls, + "usage_coverage": ( + self.calls / self.requests + if self.requests > 0 + else 0.0 + ), "input_tokens": self.input_tokens, "output_tokens": self.output_tokens, + "total_tokens": self.total_tokens, "cached_input_tokens": self.cached_input_tokens, "cache_creation_input_tokens": self.cache_creation_input_tokens, "cache_hit_rate": self.cache_hit_rate(), @@ -827,6 +870,130 @@ def _get_litellm_module() -> Any: return _LITELLM_MODULE +def estimate_token_count( + model: str, + *, + text: str | list[str] | None = None, + messages: str | Sequence[MessageLike] | None = None, + tools: list[dict[str, Any]] | None = None, +) -> int: + """Estimate request tokens locally with LiteLLM's model-aware tokenizer. + + Exactly one of ``text`` or ``messages`` must be supplied. When LiteLLM + cannot resolve a tokenizer for a provider/model pair, fall back to the + conventional four-characters-per-token approximation so preflight + estimation never requires a provider call. + """ + if (text is None) == (messages is None): + raise ValueError("provide exactly one of text or messages") + + normalized_messages = ( + messages_to_openai(messages) + if messages is not None + else None + ) + tokenizer_model = _tokenizer_model_name(model) + if tokenizer_model is None: + return _fallback_token_count( + text=text, + messages=normalized_messages, + tools=tools, + ) + litellm = _get_litellm_module() + try: + value = litellm.token_counter( + model=tokenizer_model, + text=text, + messages=normalized_messages, + tools=tools, + ) + return max(0, int(value)) + except (AttributeError, KeyError, TypeError, ValueError): + return _fallback_token_count( + text=text, + messages=normalized_messages, + tools=tools, + ) + + +def _tokenizer_model_name(model: str) -> str | None: + """Map LiteLLM route names to a tokenizer model without silent fallback.""" + normalized = (model or "").strip().lower() + if not normalized: + return None + if normalized.startswith("azure_ai/agents/"): + return None + candidate = normalized.rsplit("/", 1)[-1] + gpt5_match = re.fullmatch( + r"gpt-5(?:\.\d+)?(?:-(mini|nano|pro|codex))?" + r"(?:-\d{4}-\d{2}-\d{2})?", + candidate, + ) + if gpt5_match: + variant = gpt5_match.group(1) + if variant == "nano": + return "gpt-5-nano" + if variant == "mini": + return "gpt-5-mini" + return "gpt-5" + if re.fullmatch( + r"(?:gpt-3\.5|gpt-35)-turbo(?:-(?:\d{4}|16k))?", + candidate, + ): + return "gpt-3.5-turbo" + gpt4o_match = re.fullmatch( + r"gpt-4o(-mini)?(?:-\d{4}-\d{2}-\d{2})?", + candidate, + ) + if gpt4o_match: + return "gpt-4o-mini" if gpt4o_match.group(1) else "gpt-4o" + gpt41_match = re.fullmatch( + r"gpt-4\.1(-mini|-nano)?(?:-\d{4}-\d{2}-\d{2})?", + candidate, + ) + if gpt41_match: + return f"gpt-4.1{gpt41_match.group(1) or ''}" + if re.fullmatch( + r"gpt-4(?:-(?:\d{4}(?:-preview)?|turbo(?:-preview)?))?", + candidate, + ): + return "gpt-4" + known_patterns = ( + r"o(?:1|3|4)(?:-(?:mini|preview|pro))?(?:-\d{4}-\d{2}-\d{2})?", + r"claude-(?:\d+(?:-\d+)*-(?:opus|sonnet|haiku)|" + r"(?:opus|sonnet|haiku)-\d+(?:-\d+)*)(?:-\d{8})?", + r"gemini-(?:1\.0|1\.5|2\.0|2\.5|3(?:\.\d+)?)-" + r"(?:pro|flash|flash-lite)(?:-(?:latest|preview(?:-\d{2}-\d{2})?))?", + r"text-embedding-(?:ada-002|3-small|3-large)", + ) + if any(re.fullmatch(pattern, candidate) for pattern in known_patterns): + return candidate + return None + + +def _fallback_token_count( + *, + text: str | list[str] | None, + messages: list[dict[str, Any]] | None, + tools: list[dict[str, Any]] | None, +) -> int: + if text is not None: + serialized = "\n".join(text) if isinstance(text, list) else text + else: + serialized = json.dumps( + messages, + ensure_ascii=False, + default=str, + ) + if tools: + serialized += json.dumps( + tools, + ensure_ascii=False, + default=str, + ) + return max(1, (len(serialized) + 3) // 4) + + async def _await_with_timeout(awaitable: Any, *, timeout_s: float | None) -> Any: if timeout_s is None: return await awaitable @@ -1275,8 +1442,12 @@ def _normalize_usage(raw_usage: Any) -> UsageStats | None: return None # Chat Completions API uses prompt_tokens/completion_tokens; # Responses API uses input_tokens/output_tokens. - prompt = _coerce_int(_get_value(raw_usage, "prompt_tokens")) or _coerce_int(_get_value(raw_usage, "input_tokens")) - completion = _coerce_int(_get_value(raw_usage, "completion_tokens")) or _coerce_int(_get_value(raw_usage, "output_tokens")) + prompt = _coerce_int(_get_value(raw_usage, "prompt_tokens")) + if prompt is None: + prompt = _coerce_int(_get_value(raw_usage, "input_tokens")) + completion = _coerce_int(_get_value(raw_usage, "completion_tokens")) + if completion is None: + completion = _coerce_int(_get_value(raw_usage, "output_tokens")) total = _coerce_int(_get_value(raw_usage, "total_tokens")) if total is None and prompt is not None and completion is not None: total = prompt + completion @@ -1782,6 +1953,7 @@ async def _call() -> Any: # closure without the web_search tool. Re-issue the call here # via Chat Completions without web grounding. if not resolved_options.web_search: + _record_usage(None, model=model) raise return await generate( model, messages, @@ -1790,11 +1962,18 @@ async def _call() -> Any: reason="Responses API not available in this region", ), ) - result = normalize_response( - raw_response, - api_mode="responses" if resolved_options.web_search else "chat_completion", - request_payload=payload, - ) + except Exception: + _record_usage(None, model=model) + raise + try: + result = normalize_response( + raw_response, + api_mode="responses" if resolved_options.web_search else "chat_completion", + request_payload=payload, + ) + except Exception: + _record_usage(None, model=model) + raise _log_response("generate", model, result, time.monotonic() - t0, api_mode=api_mode) _record_usage(result.usage, model=model) return result @@ -1862,6 +2041,7 @@ async def _call() -> Any: except _ResponsesApiNotAvailableError: # Reactive degradation: see ``generate`` for the rationale. if not resolved_options.web_search: + _record_usage(None, model=model) raise return await generate_structured( model, messages, @@ -1871,11 +2051,18 @@ async def _call() -> Any: reason="Responses API not available in this region", ), ) - result = normalize_response( - raw_response, - api_mode="responses" if resolved_options.web_search else "chat_completion", - request_payload=payload, - ) + except Exception: + _record_usage(None, model=model) + raise + try: + result = normalize_response( + raw_response, + api_mode="responses" if resolved_options.web_search else "chat_completion", + request_payload=payload, + ) + except Exception: + _record_usage(None, model=model) + raise _log_response("generate_structured", model, result, time.monotonic() - t0, api_mode=api_mode, schema=schema_name) _record_usage(result.usage, model=model) return result @@ -1904,12 +2091,20 @@ async def _call() -> Any: timeout_s=resolved_options.timeout_s, ) - raw_response = await _with_retries(_call, model=model, label=resolved_options.call_label) - result = normalize_response( - raw_response, - api_mode="chat_completion", - request_payload=payload, - ) + try: + raw_response = await _with_retries( + _call, + model=model, + label=resolved_options.call_label, + ) + result = normalize_response( + raw_response, + api_mode="chat_completion", + request_payload=payload, + ) + except Exception: + _record_usage(None, model=model) + raise _log_response("generate_with_tools", model, result, time.monotonic() - t0, tools=len(tools)) _record_usage(result.usage, model=model) return result diff --git a/assert_ai/core/session.py b/assert_ai/core/session.py index 8c2babbff..6301374d5 100644 --- a/assert_ai/core/session.py +++ b/assert_ai/core/session.py @@ -241,13 +241,13 @@ async def open(self) -> None: async def close(self) -> None: return None - async def resolve( + def build_prompt( self, *, tool_name: str, tool_args: dict[str, Any], context: ResolverContext, - ) -> ToolResolution: + ) -> str: prompt = self._prompt_template replacements = { "{{description}}": str(self._scenario.get("description") or ""), @@ -258,6 +258,20 @@ async def resolve( } for placeholder, value in replacements.items(): prompt = prompt.replace(placeholder, value) + return prompt + + async def resolve( + self, + *, + tool_name: str, + tool_args: dict[str, Any], + context: ResolverContext, + ) -> ToolResolution: + prompt = self.build_prompt( + tool_name=tool_name, + tool_args=tool_args, + context=context, + ) response = await generate( self._model, prompt, diff --git a/assert_ai/core/token_estimator.py b/assert_ai/core/token_estimator.py new file mode 100644 index 000000000..425129601 --- /dev/null +++ b/assert_ai/core/token_estimator.py @@ -0,0 +1,2127 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Best-effort pre-run token estimates for configured pipeline stages.""" + +from __future__ import annotations + +import json +import os +import random +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Sequence, TypeVar + +from assert_ai.config import parse_model_config, resolve_stage_paths +from assert_ai.core.artifact_cache import _was_cached_artifact, file_sha256 +from assert_ai.core.config_model import ( + DEFAULT_GENERATION_MAX_TOKENS, + DEFAULT_GENERATION_TEMPERATURE, + DEFAULT_INFERENCE_MAX_TOKENS, + DEFAULT_JUDGE_MAX_TOKENS, + DEFAULT_SYSTEMATIZE_MAX_TOKENS, + DEFAULT_SYSTEMATIZE_TEMPERATURE, + EvaluationConfig, + ModelConfig, + TargetConfig, +) +from assert_ai.core.io import ( + INFERENCE_SET_FILE, + SCORES_FILE, + fill_template, + load_jsonl, + normalize_test_case_rows, + normalize_test_case_context, + row_factors, +) +from assert_ai.core.judge import NODE_JUDGMENTS_KEY, build_judge_contract +from assert_ai.core.model_client import Message, ToolCall, estimate_token_count +from assert_ai.core.session import ResolverContext, SimulatedResolver +from assert_ai.core.tools import ( + build_target_tools, + load_toolset_file, + normalize_tool_defs, + resolve_toolset_path, +) +from assert_ai.core.transcript import ( + AddMessageEdit, + Message as TranscriptMessage, + Transcript, + TranscriptEvent, + TranscriptMetadata, + _format_tool_call_content, +) +from assert_ai.stages import inference as inference_stage +from assert_ai.stages import judge as judge_stage +from assert_ai.stages import stratification as stratification_stage +from assert_ai.stages import systematization +from assert_ai.stages import systematization_convert +from assert_ai.stages import systematize +from assert_ai.stages import test_set + +_MAX_PROFILE_SAMPLES = 24 +_TESTER_OUTPUT_TOKENS = 55 +_PROMPT_TARGET_OUTPUT_TOKENS = 512 +_PROMPT_TARGET_OUTPUT_BUDGET_RATIO = 0.75 +_PROMPT_TARGET_OUTPUT_TOKEN_CAP = 768 +_TARGET_OUTPUT_BUDGET_RATIO = 0.875 +_SCENARIO_TARGET_OUTPUT_TOKENS = 384 +_JUDGE_OUTPUT_TOKENS = 512 +_SIMULATOR_OUTPUT_TOKENS = 90 +_UNSCORABLE_STOP_REASONS = { + "tester_input_refused", + "target_input_refused", + "target_error", + "tester_error", +} +_T = TypeVar("_T") + + +@dataclass(slots=True) +class StageTokenEstimate: + """Estimated tracked usage for one pipeline stage.""" + + calls: int = 0 + input_tokens: int = 0 + output_tokens: int = 0 + + @property + def total_tokens(self) -> int: + return self.input_tokens + self.output_tokens + + def to_dict(self) -> dict[str, int]: + return { + "calls": self.calls, + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "total_tokens": self.total_tokens, + } + + +@dataclass(slots=True) +class PipelineTokenEstimate: + """Aggregate pre-run token estimate with an explicit uncertainty range.""" + + stages: dict[str, StageTokenEstimate] = field(default_factory=dict) + notes: list[str] = field(default_factory=list) + tool_loop_total_tokens: int = 0 + + @property + def calls(self) -> int: + return sum(stage.calls for stage in self.stages.values()) + + @property + def input_tokens(self) -> int: + return sum(stage.input_tokens for stage in self.stages.values()) + + @property + def output_tokens(self) -> int: + return sum(stage.output_tokens for stage in self.stages.values()) + + @property + def total_tokens(self) -> int: + return self.input_tokens + self.output_tokens + + @property + def uncertainty(self) -> float: + return 0.35 if self.notes else 0.25 + + @property + def lower_bound_tokens(self) -> int: + return max(0, round(self.total_tokens * (1.0 - self.uncertainty))) + + @property + def upper_bound_tokens(self) -> int: + return round( + max(self.total_tokens, self.tool_loop_total_tokens) + * (1.0 + self.uncertainty) + ) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": 1, + "calls": self.calls, + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "total_tokens": self.total_tokens, + "lower_bound_tokens": self.lower_bound_tokens, + "upper_bound_tokens": self.upper_bound_tokens, + "stages": { + name: estimate.to_dict() + for name, estimate in self.stages.items() + }, + "notes": list(self.notes), + } + + +@dataclass(frozen=True, slots=True) +class _CaseProfile: + kind: str + test_case_id: str + description: str + system_prompt: str | None = None + tools: tuple[dict[str, Any], ...] = () + + +@dataclass(slots=True) +class _CaseInventory: + samples: dict[str, list[_CaseProfile]] = field(default_factory=dict) + counts: dict[str, int] = field(default_factory=dict) + + @property + def total(self) -> int: + return sum(self.counts.values()) + + +@dataclass(frozen=True, slots=True) +class _TranscriptProfile: + kind: str + test_case_id: str + transcript_xml: str + + +@dataclass(slots=True) +class _TranscriptInventory: + samples: dict[str, list[_TranscriptProfile]] = field(default_factory=dict) + counts: dict[str, int] = field(default_factory=dict) + + +@dataclass(slots=True) +class _InferenceProjection: + estimate: StageTokenEstimate + transcripts: _TranscriptInventory + pending_cases: int = 0 + notes: list[str] = field(default_factory=list) + upper_estimate: StageTokenEstimate | None = None + upper_transcripts: _TranscriptInventory | None = None + + +def _synthetic_text(tokens: int, label: str = "detail") -> str: + """Return predictable prose that tokenizes close to one token per word.""" + return " ".join([label] * max(1, tokens)) + + +def _bounded_output(expected: int, max_tokens: int | None) -> int: + if max_tokens is None: + return max(1, expected) + return max(1, min(expected, max_tokens)) + + +def _high_side_prompt_output(max_tokens: int | None) -> int: + expected = _PROMPT_TARGET_OUTPUT_TOKENS + if max_tokens is not None: + expected = max( + expected, + min( + _PROMPT_TARGET_OUTPUT_TOKEN_CAP, + round(max_tokens * _PROMPT_TARGET_OUTPUT_BUDGET_RATIO), + ), + ) + return _target_output(expected, max_tokens) + + +def _target_output(expected: int, max_tokens: int | None) -> int: + # A completion limit is a ceiling, not an expected response length. + limit = ( + round(max_tokens * _TARGET_OUTPUT_BUDGET_RATIO) + if max_tokens is not None + else None + ) + return _bounded_output(expected, limit) + + +def _response_length_hint(*instructions: str | None) -> int | None: + """Recognize simple response-wide limits, not counts inside quoted tasks.""" + numbers = { + word: index for index, word in enumerate( + ("one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"), + start=1, + ) + } + counts = "|".join(numbers) + pattern = re.compile( + rf"\b(?:in|with|using|exactly|at most|no more than|only|and)\s+" + rf"(?:(?:exactly|at most|no more than|only)\s+)?" + rf"(?P[1-9]\d{{0,5}}|{counts})\s+" + r"(?:(?:short|concise|complete)\s+)?(?Ptokens?|words?|sentences?)\b", + re.IGNORECASE, + ) + hints = [] + for instruction in instructions: + text = re.sub(r"```.*?```|~~~.*?~~~", "", instruction or "", flags=re.DOTALL) + text = re.sub(r'''"[^"]*"|(?")) or re.search( + r"\b(?:each|every|per|not|never|least|minimum|example|then|also|append|paragraphs?|sections?)\b" + r"|followed by", + clause, re.IGNORECASE, + ): + continue + matches = list(pattern.finditer(clause)) + if len(matches) != 1: + continue + match = matches[0] + suffix = clause[match.end():].strip() + if suffix and not re.fullmatch(r"[.!?]", suffix) and not re.match( + r"^(?:explaining|summarizing|describing)\b", suffix, re.IGNORECASE, + ): + continue + count_text = match["count"].lower() + count = int(count_text) if count_text.isdigit() else numbers[count_text] + unit = match["unit"].lower() + tokens_per_unit = 1.25 if unit.startswith("token") else 2 if unit.startswith("word") else 64 + hints.append(round(count * tokens_per_unit) + 32) + # Allow slack for formatting and imperfect compliance; conflicting limits use the larger one. + return max(hints) if hints else None + + +def _apply_response_length_hint( + output_tokens: int, notes: list[str], *instructions: str | None, +) -> int: + hint = _response_length_hint(*instructions) + if hint is not None and hint < output_tokens: + notes.append( + "Explicit response-length instructions reduce projected answer length " + "with formatting/compliance headroom; they are not enforced output limits." + ) + return hint + return output_tokens + + +def _request_tokens( + model: str, + messages: str | Sequence[Message | dict[str, Any]], + *, + response_schema: dict[str, Any] | None = None, + tools: list[dict[str, Any]] | None = None, +) -> int: + total = estimate_token_count(model, messages=messages, tools=tools) + if response_schema is not None: + total += estimate_token_count( + model, + text=json.dumps( + response_schema, + ensure_ascii=False, + separators=(",", ":"), + ), + ) + return total + + +def _sample_evenly( + items: list[_T], + limit: int = _MAX_PROFILE_SAMPLES, +) -> list[_T]: + if len(items) <= limit: + return list(items) + if limit <= 1: + return [items[0]] + return [ + items[round(index * (len(items) - 1) / (limit - 1))] + for index in range(limit) + ] + + +def _scaled_sum( + samples: list[_T], + total_count: int, + estimator: Callable[[_T], int], +) -> int: + if not samples or total_count <= 0: + return 0 + measured = [estimator(item) for item in _sample_evenly(samples)] + return round(sum(measured) / len(measured) * total_count) + + +def _resolved_path( + ctx: dict[str, Any], + key: str, + value: Any, +) -> Path: + resolved = resolve_stage_paths( + {key: value}, + cfg_path=Path(ctx["config_path"]), + artifacts_root=Path(ctx["artifacts_root"]), + ) + return Path(resolved[key]) + + +def _compatibility_path_will_refresh( + ctx: dict[str, Any], + *, + stage_name: str, + input_path: Path, + filename: str, +) -> bool: + artifact_ref = (ctx.get("artifact_versions") or {}).get(stage_name) + compatibility_path = (Path(ctx["suite_root"]) / filename).resolve() + input_path = input_path.resolve() + if not isinstance(artifact_ref, dict) or input_path != compatibility_path: + return False + if not compatibility_path.exists() or not compatibility_path.is_file(): + return True + try: + compatibility_hash = file_sha256(compatibility_path) + except OSError: + return True + return _was_cached_artifact( + Path(ctx["suite_root"]), + stage_name, + filename, + compatibility_hash, + ) + + +def _effective_artifact_input_path( + ctx: dict[str, Any], + *, + key: str, + value: Any, + stage_name: str, + filename: str, +) -> Path: + """Mirror compatibility-file refreshes without writing during estimation.""" + + resolved = _resolved_path(ctx, key, value) + if not _compatibility_path_will_refresh( + ctx, + stage_name=stage_name, + input_path=resolved, + filename=filename, + ): + return resolved + activated = ctx.get(key) + if activated: + activated_path = Path(str(activated)).resolve() + if activated_path.exists() and activated_path.is_file(): + return activated_path + return resolved + + +def _systematize_output_feeds_taxonomy( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], + taxonomy_stage: str, +) -> bool: + systematize_cfg = stage_cfgs.get("systematize") + downstream_cfg = stage_cfgs.get(taxonomy_stage) + if systematize_cfg is None or downstream_cfg is None: + return False + output_dir = _resolved_path( + ctx, + "save_dir", + str( + systematize_cfg.get("save_dir") + or ctx.get("systematize_artifact_dir") + or ctx["suite_root"] + ), + ) + taxonomy_input = _resolved_path( + ctx, + "taxonomy_path", + str( + downstream_cfg.get("taxonomy_path") + or ctx.get("taxonomy_path") + or Path(ctx["suite_root"]) / "taxonomy.json" + ), + ) + return ( + output_dir / "taxonomy.json" == taxonomy_input + or _compatibility_path_will_refresh( + ctx, + stage_name="systematize", + input_path=taxonomy_input, + filename="taxonomy.json", + ) + ) + + +def _load_json_mapping(path: Path) -> dict[str, Any] | None: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return None + return value if isinstance(value, dict) else None + + +def _synthetic_taxonomy( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], +) -> dict[str, Any]: + systematize_cfg = stage_cfgs.get("systematize") or {} + category_count = systematize_cfg.get( + "behavior_category_count", + systematize.DEFAULT_BEHAVIOR_CATEGORY_COUNT, + ) + if not isinstance(category_count, int) or category_count <= 0: + category_count = systematize.DEFAULT_BEHAVIOR_CATEGORY_COUNT + behavior_name = str(ctx.get("behavior_name") or "behavior") + behavior_description = str(ctx.get("behavior") or "Behavior under evaluation") + categories = [] + for index in range(category_count): + categories.append( + { + "name": f"category_{index + 1}", + "definition": _synthetic_text(42, "definition"), + "examples": [ + _synthetic_text(18, "example"), + _synthetic_text(18, "example"), + ], + "permissible": index % 3 == 0, + } + ) + return { + "behavior": { + "name": behavior_name, + "definition": behavior_description, + }, + "definition_of_terms": [ + { + "term": "representative term", + "definition": _synthetic_text(24, "definition"), + "examples": [_synthetic_text(12, "example")], + } + ], + "behavior_categories": categories, + } + + +def _taxonomy_for_stage( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], + stage_name: str, +) -> tuple[dict[str, Any], bool]: + if _systematize_output_feeds_taxonomy(ctx, stage_cfgs, stage_name): + return _synthetic_taxonomy(ctx, stage_cfgs), True + stage_cfg = stage_cfgs.get(stage_name) or {} + raw_path = ( + stage_cfg.get("taxonomy_path") + or ctx.get("taxonomy_path") + or str(Path(ctx["suite_root"]) / "taxonomy.json") + ) + taxonomy = _load_json_mapping( + _effective_artifact_input_path( + ctx, + key="taxonomy_path", + value=raw_path, + stage_name="systematize", + filename="taxonomy.json", + ) + ) + if taxonomy is not None and taxonomy.get("behavior_categories"): + return taxonomy, False + return _synthetic_taxonomy(ctx, stage_cfgs), True + + +def _synthetic_systematization( + behavior_name: str, + category_count: int, +) -> dict[str, Any]: + pattern_count = max(5, min(category_count, 40)) + systematization_text = "\n\n".join( + f"Pattern {index + 1}: {_synthetic_text(78, 'pattern')}" + for index in range(pattern_count) + ) + return { + "systematization": systematization_text, + "summary_items": [ + { + "description": _synthetic_text(28, "summary"), + "example": _synthetic_text(20, "example"), + } + for _ in range(pattern_count) + ], + "behavior": behavior_name, + } + + +def _estimate_systematize( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], +) -> StageTokenEstimate: + model_raw = raw_cfg.get("model") + if not isinstance(model_raw, dict): + raise ValueError("systematize.model must be a mapping") + model_cfg = parse_model_config( + model_raw, + field_name="systematize.model", + default_temperature=DEFAULT_SYSTEMATIZE_TEMPERATURE, + default_max_tokens=DEFAULT_SYSTEMATIZE_MAX_TOKENS, + ) + category_count = raw_cfg.get( + "behavior_category_count", + systematize.DEFAULT_BEHAVIOR_CATEGORY_COUNT, + ) + if not isinstance(category_count, int) or category_count <= 0: + category_count = systematize.DEFAULT_BEHAVIOR_CATEGORY_COUNT + behavior_name = str(ctx.get("behavior_name") or "behavior") + behavior_text = str(ctx.get("behavior") or "") + context = ctx.get("context") + + first_prompt = systematization._build_prompt( + behavior=behavior_name, + behavior_text=behavior_text, + context=context if isinstance(context, str) else None, + ) + first_schema = systematization.SystematizationResponse.model_json_schema() + first_input = _request_tokens( + model_cfg.name, + first_prompt, + response_schema=first_schema, + ) + synthetic = _synthetic_systematization(behavior_name, category_count) + first_output = _bounded_output( + estimate_token_count( + model_cfg.name, + text=json.dumps(synthetic, ensure_ascii=False), + ), + model_cfg.max_tokens, + ) + + second_prompt = ( + systematization_convert.GUIDELINE_PROMPT.replace( + "{{behavior_category_count}}", + str(category_count), + ) + + "\n\n# SYSTEMATIZATION\n" + + str(synthetic["systematization"]) + + "\n\n# SUMMARY ITEMS\n" + + json.dumps(synthetic["summary_items"], ensure_ascii=False, indent=2) + ) + second_input = _request_tokens( + model_cfg.name, + second_prompt, + response_schema=systematization_convert.TAXONOMY_SCHEMA, + ) + taxonomy_output = _bounded_output( + estimate_token_count( + model_cfg.name, + text=json.dumps( + _synthetic_taxonomy( + ctx, + {"systematize": {"behavior_category_count": category_count}}, + ), + ensure_ascii=False, + ), + ), + model_cfg.max_tokens, + ) + return StageTokenEstimate( + calls=2, + input_tokens=first_input + second_input, + output_tokens=first_output + taxonomy_output, + ) + + +def _stratification_for_plan( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + taxonomy: dict[str, Any], +) -> tuple[dict[str, Any], StageTokenEstimate | None]: + raw_path = ( + ctx.get("stratification_path") + or str(Path(ctx["suite_root"]) / "stratification.json") + ) + existing = _load_json_mapping(_resolved_path(ctx, "stratification_path", raw_path)) + if existing is not None: + return existing, None + + stratify_raw = raw_cfg.get("stratify") or {} + if not isinstance(stratify_raw, dict): + stratify_raw = {} + dimensions = stratify_raw.get("dimensions", ctx.get("dimensions")) or [] + if not isinstance(dimensions, list): + dimensions = [] + level_count = stratify_raw.get( + "level_count", + stratification_stage.DEFAULT_LEVEL_COUNT, + ) + if not isinstance(level_count, int) or level_count <= 0: + level_count = stratification_stage.DEFAULT_LEVEL_COUNT + + raw_stratification: dict[str, Any] = {} + missing_dimensions: list[dict[str, Any]] = [] + factor_order: list[str] = [] + for index, dimension in enumerate(dimensions): + if not isinstance(dimension, dict): + continue + name = str(dimension.get("name") or f"dimension_{index + 1}") + factor_order.append(name) + levels = dimension.get("levels") + if isinstance(levels, list) and levels: + raw_stratification[name] = levels + continue + missing_dimensions.append( + { + "name": name, + "description": str( + dimension.get("description") + or _synthetic_text(32, "dimension") + ), + } + ) + raw_stratification[name] = [ + { + "name": f"{name}_level_{level_index + 1}", + "definition": _synthetic_text(22, "level"), + } + for level_index in range(level_count) + ] + + normalized = stratification_stage.normalize_stratification( + raw_stratification, + taxonomy, + factor_order=factor_order, + inject_behavior=True, + ) + if not missing_dimensions: + return normalized, None + + model_raw = stratify_raw.get("model") or raw_cfg.get("model") + if not isinstance(model_raw, dict): + return normalized, None + model_cfg = parse_model_config( + model_raw, + field_name="test_set.stratify.model", + ) + normalized_context = normalize_test_case_context(ctx.get("context")) + prompt = fill_template( + stratification_stage.STRATIFICATION_PROMPT_TEMPLATE, + { + "behavior_name": str( + taxonomy.get("behavior", {}).get("name") or "behavior" + ), + "behavior_categories": ( + stratification_stage.render_behavior_categories(taxonomy) + ), + "context": normalized_context or "- (no additional context provided)", + "factors_section": ( + stratification_stage.render_factors_section(missing_dimensions) + ), + }, + ) + schema = stratification_stage._stratification_response_schema( + level_count, + dimensions=tuple(item["name"] for item in missing_dimensions), + ) + output_payload = { + item["name"]: raw_stratification[item["name"]] + for item in missing_dimensions + } + return normalized, StageTokenEstimate( + calls=1, + input_tokens=_request_tokens( + model_cfg.name, + prompt, + response_schema=schema, + ), + output_tokens=estimate_token_count( + model_cfg.name, + text=json.dumps(output_payload, ensure_ascii=False), + ), + ) + + +def _synthetic_test_case_payload( + kind: str, + *, + tool_source: str, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "title": _synthetic_text(7, "title"), + "description": _synthetic_text( + 180 if kind == "scenario" else 70, + "scenario" if kind == "scenario" else "prompt", + ), + "system_prompt": _synthetic_text( + 90 if kind == "scenario" else 35, + "instruction", + ), + } + if tool_source == test_set.TOOL_SOURCE_PER_TEST_CASE: + payload["tools"] = [ + { + "name": "lookup_record", + "description": _synthetic_text(24, "tool"), + "parameters": [ + { + "name": "query", + "type": "string", + "description": _synthetic_text(18, "parameter"), + } + ], + }, + { + "name": "submit_action", + "description": _synthetic_text(24, "tool"), + "parameters": [ + { + "name": "value", + "type": "string", + "description": _synthetic_text(18, "parameter"), + } + ], + }, + ] + return payload + + +def _normalized_test_set_tool_source(raw_cfg: dict[str, Any]) -> str: + tool_source = str( + raw_cfg.get("tool_source", test_set.TOOL_SOURCE_RUNTIME) + ) + if tool_source == test_set.TOOL_SOURCE_PER_TEST_CASE_LEGACY: + return test_set.TOOL_SOURCE_PER_TEST_CASE + return tool_source + + +def _estimate_test_set( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + taxonomy: dict[str, Any], +) -> StageTokenEstimate: + stratification, stratification_estimate = _stratification_for_plan( + ctx, + raw_cfg, + taxonomy, + ) + estimate = stratification_estimate or StageTokenEstimate() + tool_source = _normalized_test_set_tool_source(raw_cfg) + + kind_configs: list[tuple[str, dict[str, Any]]] = [] + if raw_cfg.get("prompt") and isinstance(raw_cfg.get("prompt"), dict): + kind_configs.append( + ( + "prompt", + test_set._parse_kind_config( + raw_cfg, + "prompt", + raw_cfg["prompt"], + sample_size=100, + temperature=DEFAULT_GENERATION_TEMPERATURE, + max_tokens=DEFAULT_GENERATION_MAX_TOKENS, + ), + ) + ) + if raw_cfg.get("scenario") and isinstance(raw_cfg.get("scenario"), dict): + kind_configs.append( + ( + "scenario", + test_set._parse_kind_config( + raw_cfg, + "scenario", + raw_cfg["scenario"], + sample_size=100, + temperature=DEFAULT_GENERATION_TEMPERATURE, + max_tokens=DEFAULT_GENERATION_MAX_TOKENS, + ), + ) + ) + + for kind, kind_cfg in kind_configs: + jobs, _ = test_set.build_generation_jobs( + taxonomy=taxonomy, + stratification=stratification, + sample_size=int(kind_cfg["sample_size"]), + rng=random.Random(0), + sampling=kind_cfg.get("sampling"), + ) + sampled_jobs = _sample_evenly(jobs) + sampled_input = [] + sampled_output = [] + for job in sampled_jobs: + prompt = test_set.build_generation_prompt( + kind=kind, + taxonomy=taxonomy, + behavior=job.behavior, + count=job.count, + context=ctx.get("context"), + stratification=stratification, + tuple_spec=job.tuple_spec, + tool_source=tool_source, + ) + schema = test_set.test_set_response_schema( + tool_source, + min_items=job.count, + max_items=job.count, + ) + sampled_input.append( + _request_tokens( + str(kind_cfg["model"]), + prompt, + response_schema=schema, + ) + ) + output_payload = { + "test_set": [ + _synthetic_test_case_payload( + kind, + tool_source=tool_source, + ) + for _ in range(job.count) + ] + } + sampled_output.append( + _bounded_output( + estimate_token_count( + str(kind_cfg["model"]), + text=json.dumps(output_payload, ensure_ascii=False), + ), + int(kind_cfg["max_tokens"]) + if kind_cfg.get("max_tokens") is not None + else None, + ) + ) + if sampled_jobs: + scale = len(jobs) / len(sampled_jobs) + estimate.calls += len(jobs) + estimate.input_tokens += round(sum(sampled_input) * scale) + estimate.output_tokens += round(sum(sampled_output) * scale) + return estimate + + +def _profile_from_row(row: dict[str, Any], index: int) -> _CaseProfile | None: + kind = str(row.get("type") or "") + seed = row.get("seed") + if kind not in {"prompt", "scenario"} or not isinstance(seed, dict): + return None + raw_tools = seed.get("tools") + tools = tuple(item for item in raw_tools if isinstance(item, dict)) if isinstance(raw_tools, list) else () + return _CaseProfile( + kind=kind, + test_case_id=str(row.get("test_case_id") or f"estimated_{index + 1}"), + description=str(seed.get("description") or ""), + system_prompt=str(seed.get("system_prompt") or "").strip() or None, + tools=tools, + ) + + +def _case_inventory( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], + *, + prefer_generated: bool = False, +) -> _CaseInventory: + inference_cfg = stage_cfgs.get("inference") or {} + test_set_cfg = stage_cfgs.get("test_set") or {} + raw_path = ( + inference_cfg.get("test_set_path") + or ctx.get("test_set_path") + or str(Path(ctx["suite_root"]) / test_set.TEST_SET_FILE) + ) + rows = normalize_test_case_rows( + load_jsonl( + _effective_artifact_input_path( + ctx, + key="test_set_path", + value=raw_path, + stage_name="test_set", + filename=test_set.TEST_SET_FILE, + ) + ) + ) + inventory = _CaseInventory() + if rows and not prefer_generated: + profiles_by_kind: dict[str, list[_CaseProfile]] = { + "prompt": [], + "scenario": [], + } + for index, row in enumerate(rows): + profile = _profile_from_row(row, index) + if profile is None: + continue + profiles_by_kind[profile.kind].append(profile) + for kind, profiles in profiles_by_kind.items(): + if not profiles: + continue + inventory.counts[kind] = len(profiles) + inventory.samples[kind] = profiles + return inventory + + for kind in ("prompt", "scenario"): + kind_cfg = test_set_cfg.get(kind) + if not kind_cfg or not isinstance(kind_cfg, dict): + continue + count = kind_cfg.get("sample_size", 100) + if not isinstance(count, int) or count <= 0: + continue + payload = _synthetic_test_case_payload( + kind, + tool_source=_normalized_test_set_tool_source(test_set_cfg), + ) + raw_tools = payload.get("tools") + inventory.counts[kind] = count + inventory.samples[kind] = [ + _CaseProfile( + kind=kind, + test_case_id=f"estimated_{kind}", + description=str(payload["description"]), + system_prompt=str(payload["system_prompt"]), + tools=( + tuple(raw_tools) + if isinstance(raw_tools, list) + else () + ), + ) + ] + return inventory + + +def _test_set_output_feeds_inference( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], +) -> bool: + test_set_cfg = stage_cfgs.get("test_set") + inference_cfg = stage_cfgs.get("inference") + if test_set_cfg is None or inference_cfg is None: + return False + test_set_output = _resolved_path( + ctx, + "save_path", + str( + test_set_cfg.get("save_path") + or ctx.get("test_set_path") + or Path(ctx["suite_root"]) / test_set.TEST_SET_FILE + ), + ) + inference_input = _resolved_path( + ctx, + "test_set_path", + str( + inference_cfg.get("test_set_path") + or ctx.get("test_set_path") + or Path(ctx["suite_root"]) / test_set.TEST_SET_FILE + ), + ) + return ( + test_set_output == inference_input + or _compatibility_path_will_refresh( + ctx, + stage_name="test_set", + input_path=inference_input, + filename=test_set.TEST_SET_FILE, + ) + ) + + +def _inference_output_feeds_judge( + ctx: dict[str, Any], + stage_cfgs: dict[str, dict[str, Any]], +) -> bool: + inference_cfg = stage_cfgs.get("inference") + judge_cfg = stage_cfgs.get("judge") + if inference_cfg is None or judge_cfg is None: + return False + inference_output_dir = _resolved_path( + ctx, + "save_dir", + str(inference_cfg.get("save_dir") or ctx["run_root"]), + ) + judge_input = _resolved_path( + ctx, + "inference_set_path", + str( + judge_cfg.get("inference_set_path") + or Path(ctx["run_root"]) / INFERENCE_SET_FILE + ), + ) + return (inference_output_dir / INFERENCE_SET_FILE) == judge_input + + +def _target_tools( + target: TargetConfig, + profile: _CaseProfile, + ctx: dict[str, Any], +) -> tuple[list[dict[str, Any]] | None, str | None]: + if profile.tools: + try: + return build_target_tools(normalize_tool_defs(list(profile.tools))), None + except (KeyError, TypeError, ValueError): + return None, "Per-test-case tool schemas could not be counted." + if target.tools is None: + return None, None + if target.tools.module: + return ( + build_target_tools( + [ + { + "name": "representative_tool", + "description": _synthetic_text(24, "tool"), + "parameters": [ + { + "name": "value", + "type": "string", + "description": _synthetic_text(18, "parameter"), + } + ], + } + ] + ), + "Tool-module calls use a representative tool schema.", + ) + if target.tools.toolset: + toolset_path = resolve_toolset_path( + target.tools.toolset, + config_path=Path(ctx["config_path"]), + ) + try: + return build_target_tools(load_toolset_file(toolset_path)), None + except (FileNotFoundError, OSError, ValueError): + return ( + build_target_tools( + [ + { + "name": "representative_tool", + "description": _synthetic_text(24, "tool"), + "parameters": [ + { + "name": "value", + "type": "string", + "description": _synthetic_text(18, "parameter"), + } + ], + } + ] + ), + "Toolset schemas could not be loaded; a representative schema was used.", + ) + return None, None + + +def _transcript_xml(messages: list[tuple[str, str]]) -> str: + transcript = Transcript( + metadata=TranscriptMetadata( + kind="estimated", + test_case_id="", + behavior="", + target="", + tester_model="", + ), + events=[ + TranscriptEvent( + view=["target"], + actor="target", + edit=AddMessageEdit(message=TranscriptMessage(role=role, content=content)), + ) + for role, content in messages + ], + ) + xml, _ = transcript.format_transcript_xml("target", skip_system=False) + return xml + + +def _representative_tool_value(schema: dict[str, Any] | bool) -> Any: + """Project argument shape, without trying to satisfy every schema constraint.""" + if isinstance(schema, bool): + return "representative value" + values = schema.get("enum") + if isinstance(values, list) and values: + return values[0] + kind = schema.get("type") + if kind == "object": + return { + name: _representative_tool_value(value) + for name, value in (schema.get("properties") or {}).items() + if isinstance(value, (dict, bool)) + } + if kind == "array": + return [_representative_tool_value(schema.get("items") or {})] + if kind in {"number", "integer"}: + return 1 + if kind == "boolean": + return True + return "representative value" + + +def _project_target_turn( + *, + estimate: StageTokenEstimate, + target: TargetConfig, + messages: list[Message], + transcript_messages: list[tuple[str, str]], + tools: list[dict[str, Any]] | None, + tool_history: list[dict[str, Any]], + description: str, + output_tokens: int, + tool_rounds: int = 1, + include_limit_fallback: bool = False, +) -> str: + """Project a final answer, tool rounds, and optional forced tool-limit reply.""" + + model = target.model + if isinstance(model, ModelConfig): + estimate.calls += 1 + estimate.input_tokens += _request_tokens(model.name, messages, tools=tools) + if tools: + function = tools[0]["function"] + parameters = function.get("parameters") or {} + simulator = target.tools.simulator if target.tools is not None else None + for round_index in range(tool_rounds + int(include_limit_fallback)): + resolve_tool = round_index < tool_rounds + tool_call = ToolCall( + name=str(function["name"]), + arguments={ + name: _representative_tool_value(schema) + for name, schema in (parameters.get("properties") or {}).items() + if isinstance(schema, (dict, bool)) + }, + call_id=f"estimated_tool_call_{len(messages)}", + ) + estimate.output_tokens += _bounded_output( + estimate_token_count( + model.name, + text=json.dumps({ + "name": tool_call.name, + "arguments": tool_call.arguments, + }), + ), + model.max_tokens, + ) + messages.append( + Message(role="assistant", content="", tool_calls=[tool_call]), + ) + transcript_messages.append(("assistant", "")) + tool_result = ( + _synthetic_text(_SIMULATOR_OUTPUT_TOKENS if simulator else 80, "result") + if resolve_tool else "Tool call limit reached." + ) + if simulator and resolve_tool: + resolver = SimulatedResolver( + model=simulator, + prompt_template=inference_stage.TOOL_SIM_PROMPT, + scenario={"description": description}, + ) + simulator_prompt = resolver.build_prompt( + tool_name=tool_call.name, + tool_args=tool_call.arguments, + context=ResolverContext( + conversation_messages=messages, + tool_history=tool_history, + ), + ) + estimate.calls += 1 + estimate.input_tokens += _request_tokens(simulator, simulator_prompt) + estimate.output_tokens += _SIMULATOR_OUTPUT_TOKENS + messages.append(Message(role="tool", content=tool_result, tool_call_id=tool_call.id)) + transcript_messages.append( + ("tool", _format_tool_call_content( + tool_call.name, tool_call.arguments, tool_result, + )), + ) + if resolve_tool: + tool_history.append({ + "tool_name": tool_call.name, + "tool_args": tool_call.arguments, + "tool_result": tool_result, + }) + estimate.calls += 1 + estimate.input_tokens += _request_tokens( + model.name, messages, tools=tools if resolve_tool else None, + ) + estimate.output_tokens += output_tokens + target_text = _synthetic_text(output_tokens, "response") + messages.append(Message(role="assistant", content=target_text)) + transcript_messages.append(("assistant", target_text)) + return target_text + + +def _project_prompt_case( + profile: _CaseProfile, + *, + ctx: dict[str, Any], + target: TargetConfig, + max_tokens: int, + tool_rounds: int = 1, + include_limit_fallback: bool = False, +) -> tuple[StageTokenEstimate, _TranscriptProfile, list[str]]: + estimate = StageTokenEstimate() + notes: list[str] = [] + system_prompt = str(target.system_prompt or "").strip() or profile.system_prompt + request_messages: list[Message] = [] + transcript_messages: list[tuple[str, str]] = [] + if system_prompt: + request_messages.append(Message(role="system", content=system_prompt)) + transcript_messages.append(("system", system_prompt)) + request_messages.append(Message(role="user", content=profile.description)) + transcript_messages.append(("user", profile.description)) + + target_output = _high_side_prompt_output( + target.model.max_tokens if isinstance(target.model, ModelConfig) else max_tokens + ) + target_output = _apply_response_length_hint( + target_output, notes, system_prompt, profile.description, + ) + tools = None + if isinstance(target.model, ModelConfig): + tools, tool_note = _target_tools(target, profile, ctx) + if tool_note: + notes.append(tool_note) + _project_target_turn( + estimate=estimate, + target=target, + messages=request_messages, + transcript_messages=transcript_messages, + tools=tools, + tool_history=[], + description=profile.description, + output_tokens=target_output, + tool_rounds=tool_rounds, + include_limit_fallback=include_limit_fallback, + ) + if tools: + notes.append( + "The point estimate assumes one round trip per target turn with " + "representative arguments; the upper range models the configured tool-call cap." + ) + return ( + estimate, + _TranscriptProfile( + kind="prompt", + test_case_id=profile.test_case_id, + transcript_xml=_transcript_xml(transcript_messages), + ), + notes, + ) + + +def _project_scenario_case( + profile: _CaseProfile, + *, + ctx: dict[str, Any], + target: TargetConfig, + evaluation: EvaluationConfig, + max_tokens: int, + tool_rounds: int = 1, + include_limit_fallback: bool = False, +) -> tuple[StageTokenEstimate, _TranscriptProfile, list[str]]: + estimate = StageTokenEstimate() + notes: list[str] = [] + tester = evaluation.tester + if tester is None: + return ( + estimate, + _TranscriptProfile( + "scenario", + profile.test_case_id, + _transcript_xml([]), + ), + notes, + ) + tester_system = ( + inference_stage.TESTER_SYSTEM_PROMPT + .replace("{{description}}", profile.description) + .replace("{{max_turns}}", str(evaluation.inference.max_turns)) + ) + tester_messages: list[Message] = [ + Message(role="system", content=tester_system), + Message( + role="user", + content="Begin the conversation now with the user's first message only.", + ), + ] + target_messages: list[Message] = [] + transcript_messages: list[tuple[str, str]] = [] + system_prompt = str(target.system_prompt or "").strip() or profile.system_prompt + if system_prompt: + target_messages.append(Message(role="system", content=system_prompt)) + transcript_messages.append(("system", system_prompt)) + + target_output = _target_output( + _SCENARIO_TARGET_OUTPUT_TOKENS, + target.model.max_tokens if isinstance(target.model, ModelConfig) else max_tokens, + ) + target_output = _apply_response_length_hint(target_output, notes, system_prompt) + tester_output = _bounded_output( + _TESTER_OUTPUT_TOKENS, + tester.model.max_tokens, + ) + fixed_tools: list[dict[str, Any]] | None = None + if isinstance(target.model, ModelConfig): + fixed_tools, tool_note = _target_tools(target, profile, ctx) + if tool_note: + notes.append(tool_note) + tool_history: list[dict[str, Any]] = [] + if fixed_tools: + notes.append( + "The point estimate assumes one round trip per target turn with " + "representative arguments; the upper range models the configured tool-call cap." + ) + + for turn_index in range(evaluation.inference.max_turns): + estimate.calls += 1 + estimate.input_tokens += _request_tokens( + tester.model.name, + tester_messages, + ) + estimate.output_tokens += tester_output + user_turn = _synthetic_text(tester_output, "request") + tester_messages.append(Message(role="assistant", content=user_turn)) + target_messages.append(Message(role="user", content=user_turn)) + transcript_messages.append(("user", user_turn)) + + target_text = _project_target_turn( + estimate=estimate, + target=target, + messages=target_messages, + transcript_messages=transcript_messages, + tools=fixed_tools, + tool_history=tool_history, + description=profile.description, + output_tokens=target_output, + tool_rounds=tool_rounds, + include_limit_fallback=include_limit_fallback, + ) + tester_messages.append( + Message( + role="user", + content=( + f"[Turn {turn_index + 1}/{evaluation.inference.max_turns}]\n" + f"\n{target_text}\n" + ), + ) + ) + + return ( + estimate, + _TranscriptProfile( + kind="scenario", + test_case_id=profile.test_case_id, + transcript_xml=_transcript_xml(transcript_messages), + ), + notes, + ) + + +def _filter_case_inventory( + inventory: _CaseInventory, + completed_ids: set[str], +) -> _CaseInventory: + pending = _CaseInventory() + for kind, profiles in inventory.samples.items(): + remaining = [ + profile + for profile in profiles + if profile.test_case_id not in completed_ids + ] + if not remaining: + continue + pending.samples[kind] = remaining + pending.counts[kind] = len(remaining) + return pending + + +def _pending_case_inventory( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + inventory: _CaseInventory, + *, + upstream_changed: bool, + forced: bool, +) -> tuple[_CaseInventory, bool]: + if upstream_changed or forced or inventory.total == 0: + return inventory, False + + target = ctx.get("target") + evaluation = ctx.get("evaluation") + if not isinstance(target, TargetConfig): + return inventory, False + if not isinstance(evaluation, EvaluationConfig): + evaluation = EvaluationConfig() + + raw_test_set_path = ( + raw_cfg.get("test_set_path") + or ctx.get("test_set_path") + or str(Path(ctx["suite_root"]) / test_set.TEST_SET_FILE) + ) + test_set_path = _effective_artifact_input_path( + ctx, + key="test_set_path", + value=str(raw_test_set_path), + stage_name="test_set", + filename=test_set.TEST_SET_FILE, + ) + raw_output_dir = raw_cfg.get("save_dir") or str(ctx["run_root"]) + output_dir = _resolved_path(ctx, "save_dir", str(raw_output_dir)) + inference_path = output_dir / INFERENCE_SET_FILE + if not inference_path.exists(): + return inventory, False + + resolved_max_tokens = raw_cfg.get( + "max_tokens", + DEFAULT_INFERENCE_MAX_TOKENS, + ) + if not isinstance(resolved_max_tokens, int) or resolved_max_tokens <= 0: + resolved_max_tokens = DEFAULT_INFERENCE_MAX_TOKENS + test_set_content: bytes | None = None + test_set_artifact_ref = (ctx.get("artifact_versions") or {}).get( + "test_set" + ) + rewrite_test_set = ( + not isinstance(test_set_artifact_ref, dict) + and not inference_stage._is_versioned_test_set_artifact_path( + test_set_path + ) + ) + if rewrite_test_set: + canonical_rows = normalize_test_case_rows(load_jsonl(test_set_path)) + test_set_content = ( + os.linesep.join( + json.dumps(row, ensure_ascii=False) + for row in canonical_rows + ) + + os.linesep + ).encode("utf-8") + expected_hash = inference_stage._inference_config_fingerprint( + target, + evaluation, + resolved_max_tokens, + test_set_path=test_set_path, + config_path=Path(ctx["config_path"]), + test_set_content=test_set_content, + ) + hash_path = output_dir / inference_stage._INFERENCE_CONFIG_HASH_FILE + stored_hash = ( + hash_path.read_text(encoding="utf-8").strip() + if hash_path.exists() + else None + ) + if stored_hash is not None and stored_hash != expected_hash: + return inventory, False + + completed_ids = { + str(row.get("test_case_id") or "") + for row in load_jsonl(inference_path) + if row.get("test_case_id") + } + return _filter_case_inventory(inventory, completed_ids), True + + +def _project_inventory( + ctx: dict[str, Any], + *, + target: TargetConfig, + evaluation: EvaluationConfig, + max_tokens: int, + inventory: _CaseInventory, + tool_rounds: int = 1, + include_limit_fallback: bool = False, +) -> tuple[StageTokenEstimate, _TranscriptInventory, list[str]]: + aggregate = StageTokenEstimate() + transcripts = _TranscriptInventory() + notes: list[str] = [] + for kind, profiles in inventory.samples.items(): + total_count = inventory.counts.get(kind, 0) + if total_count <= 0 or not profiles: + continue + sample_estimates: list[StageTokenEstimate] = [] + transcript_samples: list[_TranscriptProfile] = [] + for profile in _sample_evenly(profiles): + if kind == "prompt": + case_estimate, transcript, case_notes = _project_prompt_case( + profile, + ctx=ctx, + target=target, + max_tokens=max_tokens, + tool_rounds=tool_rounds, + include_limit_fallback=include_limit_fallback, + ) + else: + case_estimate, transcript, case_notes = _project_scenario_case( + profile, + ctx=ctx, + target=target, + evaluation=evaluation, + max_tokens=max_tokens, + tool_rounds=tool_rounds, + include_limit_fallback=include_limit_fallback, + ) + sample_estimates.append(case_estimate) + transcript_samples.append(transcript) + notes.extend(case_notes) + divisor = len(sample_estimates) + aggregate.calls += round( + sum(item.calls for item in sample_estimates) / divisor * total_count + ) + aggregate.input_tokens += round( + sum(item.input_tokens for item in sample_estimates) + / divisor + * total_count + ) + aggregate.output_tokens += round( + sum(item.output_tokens for item in sample_estimates) + / divisor + * total_count + ) + transcripts.samples[kind] = transcript_samples + transcripts.counts[kind] = total_count + return aggregate, transcripts, notes + + +def _estimate_inference( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + inventory: _CaseInventory, + *, + upstream_changed: bool, + forced: bool, +) -> _InferenceProjection: + target = ctx.get("target") + evaluation = ctx.get("evaluation") + if not isinstance(target, TargetConfig): + return _InferenceProjection(StageTokenEstimate(), _TranscriptInventory()) + if not isinstance(evaluation, EvaluationConfig): + evaluation = EvaluationConfig() + max_tokens = raw_cfg.get("max_tokens", DEFAULT_INFERENCE_MAX_TOKENS) + if not isinstance(max_tokens, int) or max_tokens <= 0: + max_tokens = DEFAULT_INFERENCE_MAX_TOKENS + + pending_inventory, resume_compatible = _pending_case_inventory( + ctx, + raw_cfg, + inventory, + upstream_changed=upstream_changed, + forced=forced, + ) + aggregate, pending_transcripts, pending_notes = _project_inventory( + ctx, + target=target, + evaluation=evaluation, + max_tokens=max_tokens, + inventory=pending_inventory, + ) + _full_estimate, full_transcripts, full_notes = _project_inventory( + ctx, + target=target, + evaluation=evaluation, + max_tokens=max_tokens, + inventory=inventory, + ) + transcripts = full_transcripts + if resume_compatible and pending_inventory.total < inventory.total: + raw_output_dir = raw_cfg.get("save_dir") or str(ctx["run_root"]) + output_dir = _resolved_path(ctx, "save_dir", str(raw_output_dir)) + actual_transcripts = _actual_transcripts( + output_dir / INFERENCE_SET_FILE + ) + transcripts = _merge_transcript_inventories( + actual_transcripts or _TranscriptInventory(), + pending_transcripts, + ) + notes = pending_notes + full_notes + + if not isinstance(target.model, ModelConfig) and inventory.total: + target_kind = ( + "callable" + if target.callable + else "connector" + if target.connector + else "endpoint" + if target.endpoint + else "sandbox" + ) + notes.append( + f"Target-internal usage for the {target_kind} target is not included." + ) + upper_estimate = None + upper_transcripts = None + has_tools = target.tools is not None or any( + profile.tools for profiles in inventory.samples.values() for profile in profiles + ) + if isinstance(target.model, ModelConfig) and has_tools and pending_inventory.total: + upper_estimate, upper_transcripts, _ = _project_inventory( + ctx, target=target, evaluation=evaluation, max_tokens=max_tokens, + inventory=pending_inventory, + tool_rounds=evaluation.inference.max_tool_calls, + include_limit_fallback=True, + ) + if resume_compatible and pending_inventory.total < inventory.total: + upper_transcripts = _merge_transcript_inventories( + actual_transcripts or _TranscriptInventory(), upper_transcripts, + ) + notes.append( + f"Upper range projects up to {evaluation.inference.max_tool_calls} resolved " + "tool calls per target turn, accumulated history, and a possible forced " + "final reply after the tool limit. Tool arguments/results remain representative." + ) + return _InferenceProjection( + estimate=aggregate, + transcripts=transcripts, + pending_cases=pending_inventory.total, + notes=list(dict.fromkeys(notes)), + upper_estimate=upper_estimate, + upper_transcripts=upper_transcripts, + ) + + +def _merge_transcript_inventories( + *inventories: _TranscriptInventory, +) -> _TranscriptInventory: + merged = _TranscriptInventory() + kinds = { + kind + for inventory in inventories + for kind in inventory.samples + } + for kind in kinds: + components = [ + ( + inventory.samples[kind], + inventory.counts.get( + kind, + len(inventory.samples[kind]), + ), + ) + for inventory in inventories + if inventory.samples.get(kind) + and inventory.counts.get(kind, 0) > 0 + ] + total_count = sum(count for _profiles, count in components) + if total_count <= 0: + continue + sample_budget = min(_MAX_PROFILE_SAMPLES, total_count) + allocations = [ + min( + count, + max(1, int(sample_budget * count / total_count)), + ) + for _profiles, count in components + ] + while sum(allocations) < sample_budget: + index = max( + range(len(components)), + key=lambda item: components[item][1] - allocations[item], + ) + if allocations[index] >= components[index][1]: + break + allocations[index] += 1 + while sum(allocations) > sample_budget: + index = max( + ( + item + for item in range(len(components)) + if allocations[item] > 1 + ), + key=lambda item: allocations[item], + ) + allocations[index] -= 1 + + samples: list[_TranscriptProfile] = [] + for (profiles, _count), allocation in zip( + components, + allocations, + strict=True, + ): + selected = _sample_evenly( + profiles, + limit=min(allocation, len(profiles)), + ) + samples.extend( + selected[index % len(selected)] + for index in range(allocation) + ) + merged.samples[kind] = samples + merged.counts[kind] = total_count + return merged + + +def _actual_transcripts( + inference_path: Path, +) -> _TranscriptInventory | None: + rows = load_jsonl(inference_path) + if not rows: + return None + grouped: dict[str, list[_TranscriptProfile]] = { + "prompt": [], + "scenario": [], + } + counts: dict[str, int] = {} + for row in rows: + if row.get("stop_reason") in _UNSCORABLE_STOP_REASONS: + continue + kind = str(row.get("type") or "prompt") + if kind not in grouped: + kind = "prompt" + transcript = Transcript( + metadata=TranscriptMetadata( + kind=kind, + test_case_id=str(row.get("test_case_id") or ""), + behavior=str(row.get("behavior") or ""), + target=str(row.get("target") or ""), + tester_model=str(row.get("tester_model") or ""), + dimensions=row_factors(row), + ), + events=[ + TranscriptEvent.model_validate(event) + for event in row.get("events", []) + if isinstance(event, dict) + ], + stop_reason=row.get("stop_reason"), + ) + transcript_xml, _ = transcript.format_transcript_xml( + "target", + skip_system=False, + ) + counts[kind] = counts.get(kind, 0) + 1 + grouped[kind].append( + _TranscriptProfile( + kind=kind, + test_case_id=str(row.get("test_case_id") or ""), + transcript_xml=transcript_xml, + ) + ) + return _TranscriptInventory( + samples={ + kind: profiles + for kind, profiles in grouped.items() + if profiles + }, + counts=counts, + ) + + +def _actual_transcripts_for_judge( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], +) -> _TranscriptInventory | None: + raw_path = ( + raw_cfg.get("inference_set_path") + or str(Path(ctx["run_root"]) / INFERENCE_SET_FILE) + ) + return _actual_transcripts( + _resolved_path(ctx, "inference_set_path", raw_path) + ) + + +def _synthetic_judge_output( + score_keys: list[str], + taxonomy: dict[str, Any], +) -> dict[str, Any]: + categories = [ + item + for item in taxonomy.get("behavior_categories", []) + if isinstance(item, dict) + ] + return { + "dimensions": {key: False for key in score_keys}, + "justification": _synthetic_text(36, "justification"), + "dimension_justifications": { + key: _synthetic_text(28, "evidence") + for key in score_keys + }, + NODE_JUDGMENTS_KEY: [ + { + "node_name": str(item.get("name") or "behavior"), + "violated": False, + "confidence": "high", + "reasoning": _synthetic_text(24, "reasoning"), + } + for item in categories + ], + "highlights": "", + "narrative": _synthetic_text(72, "narrative"), + } + + +def _filter_transcript_inventory( + transcripts: _TranscriptInventory, + completed_keys: set[tuple[str, str]], +) -> _TranscriptInventory: + pending = _TranscriptInventory() + for kind, profiles in transcripts.samples.items(): + remaining = [ + profile + for profile in profiles + if (kind, profile.test_case_id) not in completed_keys + ] + if not remaining: + continue + pending.samples[kind] = remaining + pending.counts[kind] = len(remaining) + return pending + + +def _pending_judge_transcripts( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + *, + taxonomy: dict[str, Any], + judge_cfg: Any, + contract: dict[str, Any], + transcripts: _TranscriptInventory, + upstream_changed: bool, + forced: bool, +) -> _TranscriptInventory: + if upstream_changed or forced: + return transcripts + + raw_inference_path = ( + raw_cfg.get("inference_set_path") + or str(Path(ctx["run_root"]) / INFERENCE_SET_FILE) + ) + inference_path = _resolved_path( + ctx, + "inference_set_path", + str(raw_inference_path), + ) + raw_output_dir = raw_cfg.get("save_dir") or str(ctx["run_root"]) + output_dir = _resolved_path(ctx, "save_dir", str(raw_output_dir)) + scores_path = output_dir / SCORES_FILE + if not inference_path.exists() or not scores_path.exists(): + return transcripts + + expected_hash = judge_stage._judge_config_fingerprint( + judge_model=judge_cfg.model.name, + judge_temperature=judge_cfg.model.temperature, + judge_max_tokens=judge_cfg.model.max_tokens, + judge_reasoning_effort=judge_cfg.model.reasoning_effort, + judge_n=judge_cfg.n, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + policy_raw=taxonomy, + system_prompt=str(contract["system_prompt"]), + inference_set_path=inference_path, + ) + hash_path = output_dir / judge_stage._JUDGE_CONFIG_HASH_FILE + stored_hash = ( + hash_path.read_text(encoding="utf-8").strip() + if hash_path.exists() + else None + ) + if stored_hash is not None and stored_hash != expected_hash: + return transcripts + + completed_keys = { + ( + str(row.get("type") or ""), + str(row.get("test_case_id") or ""), + ) + for row in load_jsonl(scores_path) + if row.get("test_case_id") + } + return _filter_transcript_inventory(transcripts, completed_keys) + + +def _estimate_judge( + ctx: dict[str, Any], + raw_cfg: dict[str, Any], + taxonomy: dict[str, Any], + projected_transcripts: _TranscriptInventory | None, + *, + upstream_changed: bool, + forced: bool, +) -> StageTokenEstimate: + evaluation = ctx.get("evaluation") + if ( + not isinstance(evaluation, EvaluationConfig) + or evaluation.judge is None + ): + return StageTokenEstimate() + judge_cfg = evaluation.judge + contract = build_judge_contract( + template=judge_stage.JUDGE_SYSTEM_PROMPT, + policy_raw=taxonomy, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + schema_name="transcript_judgment", + ) + actual_transcripts = _actual_transcripts_for_judge(ctx, raw_cfg) + if upstream_changed: + transcripts = projected_transcripts or _TranscriptInventory() + else: + transcripts = ( + actual_transcripts + or projected_transcripts + or _TranscriptInventory() + ) + transcripts = _pending_judge_transcripts( + ctx, + raw_cfg, + taxonomy=taxonomy, + judge_cfg=judge_cfg, + contract=contract, + transcripts=transcripts, + upstream_changed=upstream_changed, + forced=forced, + ) + per_call_output = _bounded_output( + max( + _JUDGE_OUTPUT_TOKENS, + estimate_token_count( + judge_cfg.model.name, + text=json.dumps( + _synthetic_judge_output( + contract["score_keys"], + taxonomy, + ), + ensure_ascii=False, + ), + ), + ), + judge_cfg.model.max_tokens or DEFAULT_JUDGE_MAX_TOKENS, + ) + + estimate = StageTokenEstimate() + for kind, profiles in transcripts.samples.items(): + count = transcripts.counts.get(kind, 0) + if not profiles or count <= 0: + continue + samples = _sample_evenly(profiles) + per_row_input = _scaled_sum( + samples, + count, + lambda profile: _request_tokens( + judge_cfg.model.name, + [ + Message( + role="system", + content=contract["system_prompt"], + ), + Message( + role="user", + content=f"# Transcript\n{profile.transcript_xml}", + ), + ], + response_schema=contract["response_schema"]["json_schema"], + ), + ) + estimate.calls += count * judge_cfg.n + estimate.input_tokens += per_row_input * judge_cfg.n + estimate.output_tokens += count * judge_cfg.n * per_call_output + return estimate + + +def estimate_pipeline_tokens( + ctx: dict[str, Any], + stages_to_run: list[tuple[str, Any, dict[str, Any]]], + *, + forced_stages: set[str] | None = None, +) -> PipelineTokenEstimate: + """Estimate usage for the uncached stages selected by the runner.""" + result = PipelineTokenEstimate() + stage_cfgs = { + name: raw_cfg + for name, _module, raw_cfg in stages_to_run + } + if not stage_cfgs: + return result + + taxonomies: dict[str, dict[str, Any]] = {} + synthetic_taxonomy_stages: list[str] = [] + for taxonomy_stage in ("test_set", "judge"): + if taxonomy_stage not in stage_cfgs: + continue + taxonomy, is_synthetic = _taxonomy_for_stage( + ctx, + stage_cfgs, + taxonomy_stage, + ) + taxonomies[taxonomy_stage] = taxonomy + if is_synthetic: + synthetic_taxonomy_stages.append(taxonomy_stage) + if synthetic_taxonomy_stages: + result.notes.append( + "Taxonomy-dependent stages use a representative generated taxonomy." + ) + test_set_changes_inference = _test_set_output_feeds_inference( + ctx, + stage_cfgs, + ) + cases = _case_inventory( + ctx, + stage_cfgs, + prefer_generated=test_set_changes_inference, + ) + projected_transcripts: _TranscriptInventory | None = None + upper_transcripts: _TranscriptInventory | None = None + tool_loop_extra_tokens = 0 + inference_pending_cases = 0 + forced = forced_stages or set() + + for stage_name, _module, raw_cfg in stages_to_run: + try: + if stage_name == "systematize": + estimate = _estimate_systematize(ctx, raw_cfg) + if raw_cfg.get("web_search", True): + result.notes.append( + "Provider-added web-search context is not included." + ) + elif stage_name == "test_set": + estimate = _estimate_test_set( + ctx, + raw_cfg, + taxonomies["test_set"], + ) + elif stage_name == "inference": + projection = _estimate_inference( + ctx, + raw_cfg, + cases, + upstream_changed=test_set_changes_inference, + forced=stage_name in forced, + ) + estimate = projection.estimate + projected_transcripts = projection.transcripts + inference_pending_cases = projection.pending_cases + result.notes.extend(projection.notes) + upper_transcripts = projection.upper_transcripts + if projection.upper_estimate is not None: + tool_loop_extra_tokens += max( + 0, projection.upper_estimate.total_tokens - estimate.total_tokens, + ) + elif stage_name == "judge": + estimate = _estimate_judge( + ctx, + raw_cfg, + taxonomies["judge"], + projected_transcripts, + upstream_changed=( + ( + inference_pending_cases > 0 + or "inference" in forced + ) + and _inference_output_feeds_judge( + ctx, + stage_cfgs, + ) + ), + forced=stage_name in forced, + ) + if upper_transcripts is not None: + upper_judge = _estimate_judge( + ctx, raw_cfg, taxonomies["judge"], upper_transcripts, + upstream_changed=( + (inference_pending_cases > 0 or "inference" in forced) + and _inference_output_feeds_judge(ctx, stage_cfgs) + ), + forced=stage_name in forced, + ) + tool_loop_extra_tokens += max( + 0, upper_judge.total_tokens - estimate.total_tokens, + ) + else: + continue + except (KeyError, TypeError, ValueError, OSError) as exc: + result.notes.append( + f"{stage_name} estimate unavailable: {exc}" + ) + continue + if estimate.calls or estimate.total_tokens: + result.stages[stage_name] = estimate + + result.tool_loop_total_tokens = result.total_tokens + tool_loop_extra_tokens + result.notes.append( + "Point estimates use high-side output assumptions so actual usage is more likely to be lower." + ) + result.notes.append("Retries and provider-side hidden overhead are not included.") + result.notes = list(dict.fromkeys(result.notes)) + return result diff --git a/assert_ai/core/tools.py b/assert_ai/core/tools.py index 3fef505b4..b1d4a9634 100644 --- a/assert_ai/core/tools.py +++ b/assert_ai/core/tools.py @@ -90,3 +90,22 @@ def load_toolset_file(path: str | Path) -> list[dict[str, Any]]: if not isinstance(data, list): raise ValueError("toolset YAML must be a list or a mapping with a 'tools' list") return normalize_tool_defs(data) + + +def resolve_toolset_path( + path: str | Path, + *, + config_path: Path | None = None, +) -> Path: + """Resolve a toolset using the same config-dir then cwd lookup as runtime.""" + resolved = Path(path).expanduser() + if resolved.is_absolute(): + return resolved + candidates = [] + if config_path is not None: + candidates.append((config_path.parent / resolved).resolve()) + candidates.append((Path.cwd() / resolved).resolve()) + return next( + (candidate for candidate in candidates if candidate.exists()), + candidates[0], + ) diff --git a/assert_ai/runner.py b/assert_ai/runner.py index 7de31e692..fac66fc79 100644 --- a/assert_ai/runner.py +++ b/assert_ai/runner.py @@ -34,6 +34,7 @@ is_cacheable_stage, override_cacheable_output_paths, prepare_artifact_plan, + preview_artifact_plan, refresh_compatibility_files, supports_artifact_cache, update_latest, @@ -137,6 +138,107 @@ def _write_suite_metadata(ctx: dict[str, Any]) -> None: write_json(suite_path, meta.to_dict()) +def _requested_force_stages( + ctx: dict[str, Any], + force_stages: list[str] | None, +) -> set[str]: + """Validate forced stages and cascade each request through downstream stages.""" + + requested = set(force_stages or []) + configured = {stage_name for stage_name, _ in ctx["stages"]} + invalid = sorted(requested.difference(configured)) + if invalid: + joined = ", ".join(invalid) + raise ConfigError(f"--force-stage stage(s) not present in config: {joined}") + + if requested: + forced_indices = [ + PIPELINE_STAGE_ORDER.index(name) + for name in requested + if name in PIPELINE_STAGE_ORDER + ] + if forced_indices: + min_forced_index = min(forced_indices) + requested.update( + name + for name in PIPELINE_STAGE_ORDER[min_forced_index:] + if name in configured + ) + return requested + + +def estimate_pipeline_usage( + *, + config: str, + force_stages: list[str] | None = None, + overrides: list[str] | None = None, + concurrency: int | None = None, +) -> dict[str, Any]: + """Estimate configured token usage without creating artifacts or running stages.""" + + ctx = _load_context(config=config, overrides=overrides) + concurrency_ignored = False + if concurrency is not None: + evaluation = ctx.get("evaluation") + inference_cfg = getattr(evaluation, "inference", None) if evaluation is not None else None + if inference_cfg is None: + concurrency_ignored = True + else: + inference_cfg.concurrency = concurrency + + requested_force_stages = _requested_force_stages(ctx, force_stages) + ctx.setdefault("artifact_versions", {}) + cache_supported = supports_artifact_cache(ctx) + if cache_supported: + activate_latest_artifacts(ctx, read_only=True) + cache_chain_reusable = True + stages_to_run: list[tuple[str, Any, dict[str, Any]]] = [] + + for stage_name, raw_cfg in ctx["stages"]: + if not raw_cfg.get("enabled", True): + continue + + module = STAGES[stage_name] + if module.SCOPE == "suite": + if cache_supported and is_cacheable_stage(stage_name): + plan = preview_artifact_plan( + ctx=ctx, + stage_name=stage_name, + raw_cfg=raw_cfg, + forced=( + stage_name in requested_force_stages + or not cache_chain_reusable + ), + ) + activate_artifact_plan(ctx, plan) + if plan.reused: + continue + cache_chain_reusable = False + raw_cfg = override_cacheable_output_paths(stage_name, raw_cfg, plan) + elif ( + module.SUITE_OUTPUT + and stage_name not in requested_force_stages + and (Path(ctx["suite_root"]) / module.SUITE_OUTPUT).exists() + ): + continue + + stages_to_run.append((stage_name, module, raw_cfg)) + + from assert_ai.core.token_estimator import estimate_pipeline_tokens + + payload = estimate_pipeline_tokens( + ctx, + stages_to_run, + forced_stages=requested_force_stages, + ).to_dict() + if concurrency_ignored: + payload.setdefault("notes", []).insert( + 0, + "Concurrency override ignored because this config has no inference stage.", + ) + return payload + + def _build_manifest(ctx: dict[str, Any]) -> RunManifest: """Build the initial run manifest.""" now = datetime.now(timezone.utc).isoformat() @@ -291,13 +393,32 @@ def _format_token_count(value: int) -> str: def _format_usage_line(usage: UsageAccumulator | None) -> str: """Render a compact ' | N calls · IN→OUT tok · X% cached' suffix.""" - if usage is None or usage.calls == 0: + if usage is None or (usage.requests == 0 and usage.calls == 0): return "" - parts = [ - f"{usage.calls} call{'s' if usage.calls != 1 else ''}", - f"{_format_token_count(usage.input_tokens)} in / " - f"{_format_token_count(usage.output_tokens)} out", - ] + request_count = usage.requests or usage.calls + if usage.calls == 0: + return ( + f" | {request_count} call{'s' if request_count != 1 else ''}" + " · token usage unavailable" + ) + parts = [f"{request_count} call{'s' if request_count != 1 else ''}"] + if usage.input_tokens or usage.output_tokens: + token_summary = ( + f"{_format_token_count(usage.input_tokens)} in / " + f"{_format_token_count(usage.output_tokens)} out" + ) + detailed_total = usage.input_tokens + usage.output_tokens + if usage.total_tokens > detailed_total: + token_summary += ( + f" / {_format_token_count(usage.total_tokens)} total" + ) + parts.append(token_summary) + else: + parts.append(f"{_format_token_count(usage.total_tokens)} total") + if usage.missing_usage_calls: + parts.append( + f"{usage.calls}/{request_count} usage reported" + ) if usage.input_tokens > 0: pct = 100.0 * usage.cached_input_tokens / usage.input_tokens parts.append(f"{pct:.1f}% cached") @@ -307,20 +428,43 @@ def _format_usage_line(usage: UsageAccumulator | None) -> str: def _build_run_metrics( stage_usage: dict[str, dict[str, Any]], total_elapsed: float, + token_estimate: dict[str, Any] | None = None, + run_completed: bool = True, + run_partial: bool = False, ) -> dict[str, Any]: """Aggregate per-stage usage into the metrics.json payload.""" totals = { + "requests": 0, "calls": 0, + "missing_usage_calls": 0, "input_tokens": 0, "output_tokens": 0, + "total_tokens": 0, "cached_input_tokens": 0, "cache_creation_input_tokens": 0, } per_model: dict[str, dict[str, int]] = {} for stage_payload in stage_usage.values(): + stage_calls = int(stage_payload.get("calls", 0) or 0) + totals["requests"] += int( + stage_payload.get("requests", stage_calls) or 0 + ) totals["calls"] += stage_payload.get("calls", 0) + totals["missing_usage_calls"] += int( + stage_payload.get("missing_usage_calls", 0) or 0 + ) totals["input_tokens"] += stage_payload.get("input_tokens", 0) totals["output_tokens"] += stage_payload.get("output_tokens", 0) + totals["total_tokens"] += int( + stage_payload.get( + "total_tokens", + ( + int(stage_payload.get("input_tokens", 0) or 0) + + int(stage_payload.get("output_tokens", 0) or 0) + ), + ) + or 0 + ) totals["cached_input_tokens"] += stage_payload.get("cached_input_tokens", 0) totals["cache_creation_input_tokens"] += stage_payload.get( "cache_creation_input_tokens", 0 @@ -329,27 +473,106 @@ def _build_run_metrics( bucket = per_model.setdefault( model, { + "requests": 0, "calls": 0, + "missing_usage_calls": 0, "input_tokens": 0, "output_tokens": 0, + "total_tokens": 0, "cached_input_tokens": 0, "cache_creation_input_tokens": 0, }, ) for key, value in model_stats.items(): bucket[key] = bucket.get(key, 0) + value + if "total_tokens" not in model_stats: + bucket["total_tokens"] += int( + model_stats.get("input_tokens", 0) or 0 + ) + int(model_stats.get("output_tokens", 0) or 0) totals["cache_hit_rate"] = ( totals["cached_input_tokens"] / totals["input_tokens"] if totals["input_tokens"] > 0 else 0.0 ) - return { + totals["usage_coverage"] = ( + totals["calls"] / totals["requests"] + if totals["requests"] > 0 + else 0.0 + ) + payload: dict[str, Any] = { "schema_version": 1, "elapsed_s": round(total_elapsed, 3), "stages": stage_usage, "per_model": per_model, "totals": totals, } + if token_estimate: + payload["token_estimate"] = token_estimate + estimated_total = int(token_estimate.get("total_tokens", 0) or 0) + actual_total = totals["total_tokens"] + if estimated_total > 0: + if not run_completed: + payload["token_estimate_accuracy"] = { + "status": "unavailable", + "reason": "pipeline_incomplete", + } + elif run_partial: + payload["token_estimate_accuracy"] = { + "status": "unavailable", + "reason": "pipeline_partial", + } + elif totals["requests"] == 0: + payload["token_estimate_accuracy"] = { + "status": "unavailable", + "reason": "no_usage_reported", + } + elif totals["missing_usage_calls"] > 0: + payload["token_estimate_accuracy"] = { + "status": "unavailable", + "reason": "provider_usage_incomplete", + "usage_coverage": totals["usage_coverage"], + } + else: + difference = actual_total - estimated_total + payload["token_estimate_accuracy"] = { + "status": "available", + "actual_total_tokens": actual_total, + "estimated_total_tokens": estimated_total, + "difference_tokens": difference, + "difference_ratio": difference / estimated_total, + "absolute_percentage_error": abs(difference) / estimated_total, + } + return payload + + +def _log_token_estimate(token_estimate: dict[str, Any]) -> None: + """Print a compact pre-run estimate and stage breakdown.""" + total = int(token_estimate.get("total_tokens", 0) or 0) + lower = int(token_estimate.get("lower_bound_tokens", total) or total) + upper = int(token_estimate.get("upper_bound_tokens", total) or total) + calls = int(token_estimate.get("calls", 0) or 0) + input_tokens = int(token_estimate.get("input_tokens", 0) or 0) + output_tokens = int(token_estimate.get("output_tokens", 0) or 0) + log.info( + "Estimated token usage: " + f"~{_format_token_count(total)} total " + f"(likely {_format_token_count(lower)}-{_format_token_count(upper)}; " + f"{_format_token_count(input_tokens)} in / " + f"{_format_token_count(output_tokens)} out across " + f"{calls} tracked call{'s' if calls != 1 else ''})" + ) + stages = token_estimate.get("stages") + if isinstance(stages, dict) and stages: + breakdown = ", ".join( + f"{name} {_format_token_count(int(stage.get('total_tokens', 0) or 0))}" + for name, stage in stages.items() + if isinstance(stage, dict) + ) + if breakdown: + log.info(f" Estimated by stage: {breakdown}") + for note in token_estimate.get("notes") or []: + if isinstance(note, str) and note: + log.info(f" Estimate note: {note}") def _print_stage_done( @@ -659,36 +882,12 @@ def run_pipeline( "[runner] --concurrency ignored: this config has no inference stage to override." ) - requested_force_stages = set(force_stages or []) - configured_stage_names = {stage_name for stage_name, _ in ctx["stages"]} - invalid_forced = sorted(requested_force_stages.difference(configured_stage_names)) - if invalid_forced: - joined = ", ".join(invalid_forced) - log.error(f"[config error] --force-stage stage(s) not present in config: {joined}") + try: + requested_force_stages = _requested_force_stages(ctx, force_stages) + except ConfigError as exc: + log.error(f"[config error] {exc}") return 1 - # Cascade: forcing an upstream stage logically invalidates every stage - # downstream of it. Without this, `--force-stage test_set` regenerates test_set - # but inference silently keeps the old inference rows (its resume cache keys on - # test_case_id, and test case ids are deterministic so they collide with the prior - # run's content). Same hazard for judge against scores.jsonl. Computing - # the closure here keeps the workflow `--force-stage ` honest - # without forcing users to remember the full downstream chain. - if requested_force_stages: - forced_indices = [ - PIPELINE_STAGE_ORDER.index(name) - for name in requested_force_stages - if name in PIPELINE_STAGE_ORDER - ] - if forced_indices: - min_forced_index = min(forced_indices) - cascade = { - name - for name in PIPELINE_STAGE_ORDER[min_forced_index:] - if name in configured_stage_names - } - requested_force_stages = requested_force_stages.union(cascade) - suite_root = Path(ctx["suite_root"]) suite_root.mkdir(parents=True, exist_ok=True) _write_suite_metadata(ctx) @@ -744,6 +943,22 @@ def run_pipeline( stages_to_run.append((stage_name, module, raw_cfg)) + token_estimate_payload: dict[str, Any] | None = None + try: + from assert_ai.core.token_estimator import estimate_pipeline_tokens + + token_estimate_payload = estimate_pipeline_tokens( + ctx, + stages_to_run, + forced_stages=requested_force_stages, + ).to_dict() + _log_token_estimate(token_estimate_payload) + except ConfigError as exc: + log.warning(f"Token estimate unavailable: {exc}") + except Exception as exc: # noqa: BLE001 + # Estimation is advisory and must never prevent the configured run. + log.warning(f"Token estimate unavailable: {exc}") + run_root = Path(ctx["run_root"]) if ctx.get("run_root") else None selected_run_stage = any(module.SCOPE == "run" for _, module, _ in stages_to_run) manifest = None @@ -795,6 +1010,7 @@ def run_pipeline( run_root=run_root, pipeline_start=pipeline_start, stage_usage=stage_usage, + token_estimate=token_estimate_payload, heartbeat=heartbeat, watchdog=watchdog, ) @@ -816,12 +1032,14 @@ def _run_stages_inner( run_root: Path | None, pipeline_start: float, stage_usage: dict[str, dict[str, Any]], + token_estimate: dict[str, Any] | None, heartbeat: ManifestHeartbeat | None, watchdog: PipelineWatchdog | None, ) -> int: """Stage execution loop. Extracted so the outer function can manage heartbeat/watchdog lifecycle in a single try/finally.""" failed_stage: str | None = None + pipeline_partial = False for stage_name, module, raw_cfg in stages_to_run: if manifest is not None and module.SCOPE == "run": @@ -868,6 +1086,7 @@ def _run_stages_inner( stage_errored_count = int( ((stage_result or {}).get("_summary") or {}).get("errored_count", 0) or 0 ) + pipeline_partial = pipeline_partial or stage_errored_count > 0 if ( cache_supported and module.SCOPE == "suite" @@ -920,7 +1139,10 @@ def _run_stages_inner( discard_artifact_plan(ctx, artifact_plans[stage_name]) elapsed = time.monotonic() - stage_start - if usage_acc is not None and usage_acc.calls > 0: + if ( + usage_acc is not None + and (usage_acc.requests > 0 or usage_acc.calls > 0) + ): stage_payload = usage_acc.to_dict() stage_payload["elapsed_s"] = round(elapsed, 3) stage_usage[stage_name] = stage_payload @@ -959,21 +1181,67 @@ def _run_stages_inner( total_elapsed = time.monotonic() - pipeline_start metrics_written = False - if run_root is not None and stage_usage: + if run_root is not None and ( + stage_usage + or (token_estimate and not (run_root / "metrics.json").exists()) + ): try: metrics_path = run_root / "metrics.json" - payload = _build_run_metrics(stage_usage, total_elapsed) + payload = _build_run_metrics( + stage_usage, + total_elapsed, + token_estimate=token_estimate, + run_completed=failed_stage is None, + run_partial=pipeline_partial, + ) write_json(metrics_path, payload) metrics_written = True totals = payload["totals"] - if totals["calls"]: + if totals["requests"] or totals["calls"]: cache_pct = 100.0 * totals["cache_hit_rate"] + if totals["input_tokens"] or totals["output_tokens"]: + detailed_total = ( + totals["input_tokens"] + totals["output_tokens"] + ) + token_summary = ( + f"{_format_token_count(totals['input_tokens'])} in / " + f"{_format_token_count(totals['output_tokens'])} out" + ) + if totals["total_tokens"] > detailed_total: + token_summary += ( + " / " + f"{_format_token_count(totals['total_tokens'])} total" + ) + else: + token_summary = ( + f"{_format_token_count(totals['total_tokens'])} total" + ) + request_count = totals["requests"] or totals["calls"] + usage_coverage = "" + if totals["missing_usage_calls"]: + usage_coverage = ( + f" · {totals['calls']}/{request_count} usage reported" + ) log.info( "Token usage: " - f"{totals['calls']} calls · " - f"{_format_token_count(totals['input_tokens'])} in / " - f"{_format_token_count(totals['output_tokens'])} out · " - f"{cache_pct:.1f}% cached" + f"{request_count} " + f"call{'s' if request_count != 1 else ''} · " + f"{token_summary}{usage_coverage} · {cache_pct:.1f}% cached" + ) + accuracy = payload.get("token_estimate_accuracy") + if ( + isinstance(accuracy, dict) + and accuracy.get("status") == "available" + ): + difference_ratio = float( + accuracy.get("difference_ratio", 0.0) or 0.0 + ) + log.info( + "Token estimate accuracy: " + f"actual {_format_token_count(int(accuracy['actual_total_tokens']))} " + f"vs estimated " + f"{_format_token_count(int(accuracy['estimated_total_tokens']))} " + f"({difference_ratio:+.1%})" ) except Exception: # noqa: BLE001 log.debug("Failed to write metrics.json", exc_info=True) diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index de0fab133..618e6c1d1 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -54,7 +54,11 @@ serialize_response, ) from assert_ai.core.tool_backend import ToolBackendResolver, inspect_tool_module -from assert_ai.core.tools import load_toolset_file, normalize_tool_defs +from assert_ai.core.tools import ( + load_toolset_file, + normalize_tool_defs, + resolve_toolset_path, +) from assert_ai.core.transcript import ( AddMessageEdit, Message as TranscriptMessage, @@ -147,6 +151,7 @@ def _inference_config_fingerprint( max_tokens: int, test_set_path: Path | None = None, config_path: Path | None = None, + test_set_content: bytes | None = None, ) -> str: """Deterministic hash of config values that affect inference output. @@ -157,7 +162,9 @@ def _inference_config_fingerprint( """ target_name = target.model.name if isinstance(target.model, ModelConfig) else (target.connector or target.callable or target.endpoint or target.sandbox or "") test_set_sha = "" - if test_set_path is not None and test_set_path.exists(): + if test_set_content is not None: + test_set_sha = hashlib.sha256(test_set_content).hexdigest() + elif test_set_path is not None and test_set_path.exists(): test_set_sha = hashlib.sha256(test_set_path.read_bytes()).hexdigest() sandbox_sha = "" if target.sandbox: @@ -532,14 +539,10 @@ def _build_hosted_session( if tools is None: if not isinstance(toolset_path, str) or not toolset_path.strip(): raise ValueError("simulated tools require target.tools.toolset or per-test-case tools") - resolved_path = Path(toolset_path).expanduser() - if not resolved_path.is_absolute(): - candidates = [] - if config_path is not None: - candidates.append((config_path.parent / resolved_path).resolve()) - candidates.append((Path.cwd() / resolved_path).resolve()) - found = next((c for c in candidates if c.exists()), None) - resolved_path = found if found is not None else candidates[0] + resolved_path = resolve_toolset_path( + toolset_path, + config_path=config_path, + ) tools = load_toolset_file(resolved_path) return HostedSession( model=model, diff --git a/docs/cli/commands.md b/docs/cli/commands.md index 5c42e1ebc..b7e8c8926 100644 --- a/docs/cli/commands.md +++ b/docs/cli/commands.md @@ -18,6 +18,7 @@ assert-ai [GLOBAL_OPTIONS] COMMAND [ARGS] [OPTIONS] ## Command groups - `init`: interactive config generation assistant +- `estimate`: preview tracked model token usage without running stages - `run`: execute pipeline stages - `results`: list/status/compare suites and runs - `analysis`: post-hoc metrics commands @@ -52,6 +53,18 @@ Options: - `--dry-run` optional flag - `--no-color` optional flag +## `estimate` + +Estimate token usage without executing any pipeline stages. + +```bash +assert-ai estimate --config [OPTIONS] +``` + +The command uses the same local, conservative estimator shown before `run`. +It does not call a provider or create run artifacts. Use `--output json` for +machine-readable output. + ## `run` Run the evaluation pipeline from evaluation config YAML file. @@ -74,6 +87,43 @@ Optional: - `--log-file ` - `--output text|json` +Before uncached stages execute, `run` prints a best-effort token estimate with +a likely range and per-stage breakdown. The point estimate deliberately uses +high-side output assumptions so it is more likely to be above actual usage than +below it. Estimation uses local tokenization and does not call a provider. For +callable, connector, endpoint, and sandbox targets, model usage inside the +target is opaque to ASSERT and is explicitly excluded; tester and judge usage +is still estimated. + +The estimator counts known prompts and tool schemas locally, then projects +completion lengths and their reuse in later conversation turns and judge inputs. +Small target output limits use 87.5% of the configured cap rather than assuming +every answer exhausts it. Larger prompt answers retain a 512-token baseline, +75% budget scaling, and a 768-token projection ceiling; scenario answers use +384 tokens, subject to the same 87.5% cap. Judge outputs use the larger of +512 tokens or a representative response shaped by the scoring contract, capped +by the judge's output limit. + +Simple English response-wide word, token, or sentence instructions can lower +the answer projection: 2 tokens per word, 1.25 per requested token, or 64 per +sentence, plus 32 tokens of formatting/compliance headroom, never above the +usual projection. Quoted, nested/per-item, negative, lower-bound, and ambiguous +instructions retain the usual projection. Prompt cases use the effective system +prompt and user request; scenarios use only the system prompt, since scenario +descriptions can specify different requirements for different turns. These hints +are not enforced output limits. + +The point estimate for tool-enabled targets assumes one round trip per turn: one schema-shaped +tool-call response, one tool result, and one final answer. Tool history is +retained in later requests and projected judge transcripts. Simulated-tool +requests use the same prompt builder as execution. The upper end of the range +instead projects `pipeline.inference.max_tool_calls` resolved tool calls per +turn, cumulative target/simulator context, a possible forced final reply after +the tool limit, and the resulting judge transcript, with 35% headroom. It can +therefore exceed 135% of the point estimate. These are planning +heuristics, not guaranteed upper bounds: larger arguments/results, longer responses, +retries, and hidden provider overhead can exceed the estimate. + ## `results list` List suites or list runs for one suite. diff --git a/docs/concepts.md b/docs/concepts.md index f8aa0f518..04aad760a 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -68,7 +68,7 @@ Output: - `scores.jsonl` -`metrics.json` (pipeline token-usage telemetry) is written by the runner after all stages complete, not by the judge stage itself. +`metrics.json` (pipeline token-usage telemetry) is written by the runner, not by the judge stage itself. It includes the pre-run token estimate, provider-reported total usage, and usage coverage. Estimate accuracy is reported only for complete runs with complete provider usage metadata; partial, failed, or sparsely reported runs record why accuracy is unavailable. ## Risks and limitations of ASSERT diff --git a/docs/guides/results.md b/docs/guides/results.md index 6fc859e1d..42e2d1c78 100644 --- a/docs/guides/results.md +++ b/docs/guides/results.md @@ -48,6 +48,10 @@ The run viewer shows the full custom-grade distribution and groups semantic N/A ![Custom rubric scale run summary](../images/custom-rubric-scale-run.png) +The **Summary & submit** step shows a compact conservative token estimate before the run starts. It is computed locally without a provider call. + +When `metrics.json` contains token telemetry, the completed run viewer shows a compact estimate-versus-actual summary. Stage estimates and estimator notes remain available under **Details**. Incomplete or partial provider telemetry is labeled unavailable rather than reported as an accuracy result. + ## Useful CLI commands for viewing results ```bash diff --git a/tests/test_artifact_cache.py b/tests/test_artifact_cache.py index 210082864..515638300 100644 --- a/tests/test_artifact_cache.py +++ b/tests/test_artifact_cache.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import json import logging import shutil import unittest @@ -17,6 +18,7 @@ hash_payload, override_cacheable_output_paths, prepare_artifact_plan, + preview_artifact_plan, refresh_compatibility_files, _allocate_version_dir, _iter_version_dirs, @@ -100,6 +102,34 @@ def test_hash_mismatch_allocates_next_version(self) -> None: self.assertFalse(second.reused) self.assertEqual(second.version, "v0002") + def test_preview_plan_redirects_outputs_without_allocating_version(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = self._ctx(root) + raw_cfg = { + "model": {"name": "azure/gpt-5.4"}, + "behavior_category_count": 2, + "save_dir": "user/elsewhere", + } + + plan = preview_artifact_plan( + ctx=ctx, + stage_name="systematize", + raw_cfg=raw_cfg, + forced=False, + ) + overridden = override_cacheable_output_paths( + "systematize", + raw_cfg, + plan, + ) + + self.assertFalse(plan.reused) + self.assertEqual(plan.version, "preview") + self.assertFalse(plan.artifact_dir.exists()) + self.assertEqual(Path(overridden["save_dir"]), plan.artifact_dir) + self.assertEqual(raw_cfg["save_dir"], "user/elsewhere") + def test_revert_to_prior_config_reuses_existing_version(self) -> None: """v0001 -> change behavior -> v0002 -> revert -> reuse v0001 (not v0002).""" @@ -255,6 +285,46 @@ def test_activate_latest_rebuilds_ref_when_recorded_paths_are_stale(self) -> Non self.assertNotIn("MISSING", persisted_ref.get("artifact_dir", "")) self.assertNotIn("MISSING", persisted_ref.get("metadata_path", "")) + def test_activate_latest_read_only_recovers_without_writing(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = self._ctx(root) + raw_cfg = { + "model": {"name": "azure/gpt-5.4"}, + "behavior_category_count": 2, + } + plan = self._finalize_policy(ctx, raw_cfg) + latest_path = Path(ctx["suite_root"]) / "latest.json" + latest = json.loads(latest_path.read_text(encoding="utf-8")) + latest_ref = latest["artifacts"]["systematize"] + latest_ref["artifact_dir"] = "artifacts/systematize/MISSING" + latest_ref["metadata_path"] = ( + "artifacts/systematize/MISSING/artifact.json" + ) + latest_path.write_text(json.dumps(latest), encoding="utf-8") + latest_before = latest_path.read_bytes() + + recovery_ctx = self._ctx(root) + with ( + mock.patch( + "assert_ai.core.artifact_cache.refresh_compatibility_files" + ) as refresh, + mock.patch( + "assert_ai.core.artifact_cache.update_latest" + ) as update, + ): + activate_latest_artifacts(recovery_ctx, read_only=True) + + recovered = recovery_ctx.get("artifact_versions", {}).get( + "systematize" + ) + self.assertIsNotNone(recovered) + self.assertEqual(recovered["version"], plan.version) + self.assertNotIn("MISSING", recovered["artifact_dir"]) + self.assertEqual(latest_path.read_bytes(), latest_before) + refresh.assert_not_called() + update.assert_not_called() + def test_activate_latest_handles_metadata_missing_primary_output_key(self) -> None: """Regression for Copilot review (round 4). @@ -964,4 +1034,3 @@ def test_per_file_isolation(self) -> None: if __name__ == "__main__": unittest.main() - diff --git a/tests/test_cli.py b/tests/test_cli.py index 40434cd7c..c3f2d19d8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,9 +1,10 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import json import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch from click.testing import CliRunner @@ -33,8 +34,75 @@ def test_help_shows_run_subcommand(self) -> None: self.assertEqual(result.exit_code, 0, msg=result.output) self.assertIn("Commands:", result.output) + self.assertIn("estimate", result.output) self.assertIn("run", result.output) + def test_estimate_outputs_machine_readable_json_without_logging_auth_mode(self) -> None: + with self.runner.isolated_filesystem(): + config = Path("eval.yaml") + config.write_text("suite: test\npipeline: {}\n", encoding="utf-8") + runner_module = MagicMock() + runner_module.estimate_pipeline_usage.return_value = { + "schema_version": 1, + "calls": 2, + "input_tokens": 100, + "output_tokens": 50, + "total_tokens": 150, + "lower_bound_tokens": 98, + "upper_bound_tokens": 203, + "stages": {}, + "notes": [], + } + with ( + patch("assert_ai.cli._load_runner_module", return_value=runner_module), + patch("assert_ai.core.azure_auth.log_resolved_azure_auth_mode") as log_auth, + ): + result = self.runner.invoke( + cli, + ["estimate", "--config", str(config), "--output", "json"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertEqual(json.loads(result.output)["total_tokens"], 150) + runner_module.estimate_pipeline_usage.assert_called_once_with( + config=str(config), + force_stages=[], + overrides=[], + concurrency=None, + ) + log_auth.assert_not_called() + + def test_estimate_text_reports_zero_when_no_model_calls_are_expected( + self, + ) -> None: + with self.runner.isolated_filesystem(): + config = Path("eval.yaml") + config.write_text("suite: test\npipeline: {}\n", encoding="utf-8") + runner_module = MagicMock() + runner_module.estimate_pipeline_usage.return_value = { + "schema_version": 1, + "calls": 0, + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "lower_bound_tokens": 0, + "upper_bound_tokens": 0, + "stages": {}, + "notes": ["Callable target-internal usage is not included."], + } + with patch( + "assert_ai.cli._load_runner_module", + return_value=runner_module, + ): + result = self.runner.invoke( + cli, + ["estimate", "--config", str(config)], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("Estimated token usage: 0 tracked tokens", result.output) + self.assertIn("Callable target-internal usage is not included", result.output) + @unittest.skip("--config is now required; default eval.yaml lookup removed in merge") def test_missing_default_config_errors(self) -> None: with self.runner.isolated_filesystem(): diff --git a/tests/test_model_client.py b/tests/test_model_client.py index 7dae30e97..9b0105fa2 100644 --- a/tests/test_model_client.py +++ b/tests/test_model_client.py @@ -58,6 +58,130 @@ async def fake_acompletion(**kwargs): self.assertEqual(response.request_payload["model"], "openai/gpt-5-mini") self.assertEqual(response.request_payload["messages"], [{"role": "user", "content": "say hi"}]) + def test_estimate_token_count_uses_model_aware_tokenizer(self) -> None: + captured: dict[str, object] = {} + + def token_counter(**kwargs): + captured.update(kwargs) + return 37 + + fake_litellm = SimpleNamespace(token_counter=token_counter) + with patch.object( + model_client, + "_get_litellm_module", + return_value=fake_litellm, + ): + count = model_client.estimate_token_count( + "openai/gpt-5-mini", + messages=[model_client.Message(role="user", content="hello")], + ) + + self.assertEqual(count, 37) + self.assertEqual(captured["model"], "gpt-5-mini") + self.assertEqual( + captured["messages"], + [{"role": "user", "content": "hello"}], + ) + + def test_estimate_token_count_normalizes_versioned_azure_model(self) -> None: + captured: dict[str, object] = {} + + def token_counter(**kwargs): + captured.update(kwargs) + return 11 + + fake_litellm = SimpleNamespace(token_counter=token_counter) + with patch.object( + model_client, + "_get_litellm_module", + return_value=fake_litellm, + ): + count = model_client.estimate_token_count( + "azure/gpt-5.4-mini", + text="hello", + ) + + self.assertEqual(count, 11) + self.assertEqual(captured["model"], "gpt-5-mini") + + def test_estimate_token_count_normalizes_dated_gpt5_snapshot(self) -> None: + captured: dict[str, object] = {} + + def token_counter(**kwargs): + captured.update(kwargs) + return 9 + + fake_litellm = SimpleNamespace(token_counter=token_counter) + with patch.object( + model_client, + "_get_litellm_module", + return_value=fake_litellm, + ): + count = model_client.estimate_token_count( + "azure/gpt-5-mini-2025-08-07", + text="hello", + ) + + self.assertEqual(count, 9) + self.assertEqual(captured["model"], "gpt-5-mini") + + def test_estimate_token_count_normalizes_legacy_openai_snapshots(self) -> None: + routed_models: list[str] = [] + + def token_counter(**kwargs): + routed_models.append(kwargs["model"]) + return 9 + + fake_litellm = SimpleNamespace(token_counter=token_counter) + with patch.object( + model_client, + "_get_litellm_module", + return_value=fake_litellm, + ): + model_client.estimate_token_count( + "openai/gpt-4-0125-preview", + text="hello", + ) + model_client.estimate_token_count( + "openai/gpt-3.5-turbo-0125", + text="hello", + ) + model_client.estimate_token_count( + "azure/gpt-35-turbo", + text="hello", + ) + + self.assertEqual( + routed_models, + ["gpt-4", "gpt-3.5-turbo", "gpt-3.5-turbo"], + ) + + def test_estimate_token_count_falls_back_to_character_ratio(self) -> None: + with patch.object( + model_client, + "_get_litellm_module", + side_effect=AssertionError("unknown aliases should not reach LiteLLM"), + ): + count = model_client.estimate_token_count( + "custom/provider-model", + text="abcdefgh", + ) + + self.assertEqual(count, 2) + + def test_estimate_token_count_rejects_gpt_like_deployment_alias(self) -> None: + with patch.object( + model_client, + "_get_litellm_module", + side_effect=AssertionError("deployment aliases should use fallback"), + ): + count = model_client.estimate_token_count( + "azure/gpt-prod", + text="abcdefgh", + ) + + self.assertEqual(count, 2) + async def test_generate_structured_adds_json_schema_response_format(self) -> None: captured: dict[str, object] = {} @@ -304,6 +428,19 @@ def test_extracts_openai_responses_cached_tokens(self) -> None: assert usage is not None self.assertEqual(usage.cached_input_tokens, 2048) + def test_preserves_explicit_zero_token_fields(self) -> None: + usage = model_client._normalize_usage( + { + "prompt_tokens": 100, + "completion_tokens": 0, + } + ) + + assert usage is not None + self.assertEqual(usage.prompt_tokens, 100) + self.assertEqual(usage.completion_tokens, 0) + self.assertEqual(usage.total_tokens, 100) + def test_extracts_anthropic_cache_tokens(self) -> None: # Anthropic surfaces both read and creation counts at the top level. usage = model_client._normalize_usage({ @@ -392,21 +529,67 @@ def test_add_aggregates_totals_and_per_model(self) -> None: ), model="azure/gpt-5.4-mini", ) + self.assertEqual(acc.requests, 2) self.assertEqual(acc.calls, 2) self.assertEqual(acc.input_tokens, 300) self.assertEqual(acc.output_tokens, 130) + self.assertEqual(acc.total_tokens, 430) self.assertEqual(acc.cached_input_tokens, 100) self.assertAlmostEqual(acc.cache_hit_rate(), 100 / 300) per_model = acc.per_model["azure/gpt-5.4-mini"] self.assertEqual(per_model["calls"], 2) self.assertEqual(per_model["input_tokens"], 300) + self.assertEqual(per_model["total_tokens"], 430) self.assertEqual(per_model["cached_input_tokens"], 100) def test_add_handles_none_usage_silently(self) -> None: acc = model_client.UsageAccumulator() acc.add(None, model="azure/gpt-5.4-mini") + self.assertEqual(acc.requests, 1) self.assertEqual(acc.calls, 0) + self.assertEqual(acc.missing_usage_calls, 1) self.assertEqual(acc.input_tokens, 0) + payload = acc.to_dict() + self.assertEqual(payload["usage_coverage"], 0.0) + self.assertEqual( + payload["per_model"]["azure/gpt-5.4-mini"]["requests"], + 1, + ) + + def test_add_uses_total_only_usage_payload(self) -> None: + acc = model_client.UsageAccumulator() + acc.add( + model_client.UsageStats(total_tokens=123), + model="custom/model", + ) + + self.assertEqual(acc.requests, 1) + self.assertEqual(acc.calls, 1) + self.assertEqual(acc.total_tokens, 123) + self.assertEqual(acc.input_tokens, 0) + self.assertEqual(acc.output_tokens, 0) + self.assertEqual(acc.missing_usage_calls, 0) + + def test_add_treats_empty_usage_payload_as_missing(self) -> None: + acc = model_client.UsageAccumulator() + acc.add(model_client.UsageStats(), model="custom/model") + + self.assertEqual(acc.requests, 1) + self.assertEqual(acc.calls, 0) + self.assertEqual(acc.missing_usage_calls, 1) + + def test_add_tracks_one_sided_usage_as_incomplete(self) -> None: + acc = model_client.UsageAccumulator() + acc.add( + model_client.UsageStats(prompt_tokens=100), + model="custom/model", + ) + + self.assertEqual(acc.requests, 1) + self.assertEqual(acc.calls, 0) + self.assertEqual(acc.missing_usage_calls, 1) + self.assertEqual(acc.input_tokens, 100) + self.assertEqual(acc.total_tokens, 100) def test_cache_hit_rate_is_zero_when_no_input_tokens(self) -> None: acc = model_client.UsageAccumulator() @@ -466,6 +649,54 @@ async def fake_acompletion(**kwargs): self.assertEqual(usage.cached_input_tokens, 3 * 512) self.assertIn("azure/gpt-5.4-mini", usage.per_model) + async def test_track_usage_records_terminal_request_failure(self) -> None: + async def fail_request(*_args, **_kwargs): + raise model_client.LLMInputError("refused") + + with ( + patch.object( + model_client, + "_get_litellm_module", + return_value=SimpleNamespace(), + ), + patch.object(model_client, "_with_retries", new=fail_request), + model_client.track_usage() as usage, + ): + with self.assertRaises(model_client.LLMInputError): + await model_client.generate( + "openai/gpt-5-mini", + "hello", + ) + + self.assertEqual(usage.requests, 1) + self.assertEqual(usage.calls, 0) + self.assertEqual(usage.missing_usage_calls, 1) + + async def test_track_usage_records_chat_responses_marker_failure(self) -> None: + async def fail_request(*_args, **_kwargs): + raise model_client._ResponsesApiNotAvailableError("unsupported") + + with ( + patch.object( + model_client, + "_get_litellm_module", + return_value=SimpleNamespace(), + ), + patch.object(model_client, "_with_retries", new=fail_request), + model_client.track_usage() as usage, + ): + with self.assertRaises( + model_client._ResponsesApiNotAvailableError + ): + await model_client.generate( + "openai/gpt-5-mini", + "hello", + ) + + self.assertEqual(usage.requests, 1) + self.assertEqual(usage.calls, 0) + self.assertEqual(usage.missing_usage_calls, 1) + async def test_record_usage_outside_scope_is_a_noop(self) -> None: # Should not raise even when no accumulator is active. model_client._record_usage( diff --git a/tests/test_runner_stage_filters.py b/tests/test_runner_stage_filters.py index 55d3d8292..91e76f91e 100644 --- a/tests/test_runner_stage_filters.py +++ b/tests/test_runner_stage_filters.py @@ -2,12 +2,14 @@ # Licensed under the MIT License. import io +import json import unittest from pathlib import Path from tempfile import TemporaryDirectory from types import SimpleNamespace from unittest.mock import patch +from assert_ai.core.model_client import UsageStats, _record_usage from assert_ai.runner import run_pipeline @@ -191,6 +193,150 @@ def test_force_stage_no_cascade_when_only_terminal_stage_forced(self) -> None: self.assertEqual(rc, 0) self.assertEqual(seen, ["judge"]) + def test_estimator_failure_does_not_block_pipeline(self) -> None: + seen: list[str] = [] + stages = { + "taxonomy": SimpleNamespace( + SCOPE="suite", + SUITE_OUTPUT=None, + run=self._async_recorder("taxonomy", seen), + ) + } + with TemporaryDirectory() as tmp_dir: + ctx = { + "stages": [("taxonomy", {})], + "suite_root": str(Path(tmp_dir) / "suite"), + "run_root": None, + } + with ( + patch("assert_ai.runner._load_context", return_value=ctx), + patch("assert_ai.runner._write_suite_metadata"), + patch("assert_ai.runner.STAGES", stages), + patch( + "assert_ai.core.token_estimator.estimate_pipeline_tokens", + side_effect=RuntimeError("estimator failed"), + ), + self.assertLogs("assert_ai.runner", level="WARNING") as logs, + ): + rc = run_pipeline(config="config.yaml") + + self.assertEqual(rc, 0) + self.assertEqual(seen, ["taxonomy"]) + self.assertIn("Token estimate unavailable", "\n".join(logs.output)) + + def test_partial_stage_marks_estimate_accuracy_unavailable(self) -> None: + async def partial_stage( + ctx: dict[str, object], + raw_cfg: dict[str, object], + ) -> dict[str, object]: + _record_usage( + UsageStats( + prompt_tokens=80, + completion_tokens=20, + total_tokens=100, + ), + model="test/model", + ) + return {"_summary": {"errored_count": 1}} + + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + run_root = root / "run" + ctx = { + "stages": [("inference", {})], + "suite_root": str(root / "suite"), + "run_root": str(run_root), + } + manifest = SimpleNamespace( + started_at="", + status="running", + ended_at=None, + stages={}, + stage_timings={}, + to_dict=lambda: {}, + ) + estimate = SimpleNamespace( + to_dict=lambda: {"total_tokens": 110} + ) + with ( + patch("assert_ai.runner._load_context", return_value=ctx), + patch("assert_ai.runner._write_suite_metadata"), + patch("assert_ai.runner._build_manifest", return_value=manifest), + patch("assert_ai.runner._write_manifest"), + patch( + "assert_ai.runner.STAGES", + { + "inference": SimpleNamespace( + SCOPE="run", + SUITE_OUTPUT=None, + run=partial_stage, + ) + }, + ), + patch( + "assert_ai.core.token_estimator.estimate_pipeline_tokens", + return_value=estimate, + ), + ): + rc = run_pipeline(config="config.yaml") + + metrics = json.loads( + (run_root / "metrics.json").read_text(encoding="utf-8") + ) + + self.assertEqual(rc, 0) + self.assertEqual( + metrics["token_estimate_accuracy"]["reason"], + "pipeline_partial", + ) + + + def test_no_usage_preserves_existing_metrics_but_new_runs_keep_notes(self) -> None: + note = "Target-internal usage for the callable target is not included." + for existing in (False, True): + with self.subTest(existing=existing), TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + run_root = root / "run" + run_root.mkdir() + metrics_path = run_root / "metrics.json" + original = b'{"totals":{"calls":1,"total_tokens":123},"stages":{"judge":{"total_tokens":123}}}\n' + if existing: + metrics_path.write_bytes(original) + ctx = { + "stages": [("inference", {})], + "suite_root": str(root / "suite"), + "run_root": str(run_root), + } + manifest = SimpleNamespace( + started_at="", status="running", ended_at=None, + stages={}, stage_timings={}, to_dict=lambda: {}, + ) + estimate = {"total_tokens": 0, "calls": 0, "stages": {}, "notes": [note]} + with ( + patch("assert_ai.runner._load_context", return_value=ctx), + patch("assert_ai.runner._write_suite_metadata"), + patch("assert_ai.runner._build_manifest", return_value=manifest), + patch("assert_ai.runner._write_manifest"), + patch("assert_ai.runner.STAGES", { + "inference": SimpleNamespace( + SCOPE="run", SUITE_OUTPUT=None, + run=self._async_recorder("inference", []), + ), + }), + patch( + "assert_ai.core.token_estimator.estimate_pipeline_tokens", + return_value=SimpleNamespace(to_dict=lambda: estimate), + ), + ): + self.assertEqual(run_pipeline(config="config.yaml"), 0) + if existing: + self.assertEqual(metrics_path.read_bytes(), original) + else: + self.assertEqual( + json.loads(metrics_path.read_text(encoding="utf-8"))["token_estimate"], + estimate, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_runner_usage_metrics.py b/tests/test_runner_usage_metrics.py index bf0640fbb..6edd4d8a0 100644 --- a/tests/test_runner_usage_metrics.py +++ b/tests/test_runner_usage_metrics.py @@ -14,6 +14,7 @@ _build_run_metrics, _format_token_count, _format_usage_line, + _log_token_estimate, ) @@ -34,6 +35,14 @@ def test_millions_render_with_one_decimal(self) -> None: class FormatUsageLineTest(unittest.TestCase): + def test_zero_token_estimate_logs_opaque_target_caveat(self) -> None: + note = "Target-internal usage for the callable target is not included." + with self.assertLogs("assert_ai.runner", level="INFO") as captured: + _log_token_estimate({"total_tokens": 0, "calls": 0, "notes": [note]}) + text = "\n".join(captured.output) + self.assertIn("0 tracked calls", text) + self.assertIn(note, text) + def test_returns_empty_when_no_calls(self) -> None: self.assertEqual(_format_usage_line(None), "") self.assertEqual(_format_usage_line(UsageAccumulator()), "") @@ -67,6 +76,26 @@ def test_omits_cache_percentage_when_no_input_tokens(self) -> None: usage = UsageAccumulator(calls=1, input_tokens=0, output_tokens=5) self.assertNotIn("cached", _format_usage_line(usage)) + def test_renders_total_only_usage(self) -> None: + usage = UsageAccumulator( + requests=1, + calls=1, + total_tokens=123, + ) + self.assertIn("123 total", _format_usage_line(usage)) + + def test_renders_mixed_detailed_and_total_only_usage(self) -> None: + usage = UsageAccumulator() + usage.add( + UsageStats(prompt_tokens=100, completion_tokens=20), + model="detailed", + ) + usage.add(UsageStats(total_tokens=500), model="total-only") + + line = _format_usage_line(usage) + self.assertIn("2 calls", line) + self.assertIn("100 in / 20 out / 620 total", line) + class BuildRunMetricsTest(unittest.TestCase): def test_aggregates_per_stage_into_totals(self) -> None: @@ -114,6 +143,7 @@ def test_aggregates_per_stage_into_totals(self) -> None: self.assertEqual(totals["calls"], 115) self.assertEqual(totals["input_tokens"], 875_000) self.assertEqual(totals["output_tokens"], 17_000) + self.assertEqual(totals["total_tokens"], 892_000) self.assertEqual(totals["cached_input_tokens"], 630_000) self.assertAlmostEqual(totals["cache_hit_rate"], 630_000 / 875_000) per_model = payload["per_model"]["azure/gpt-5.4-mini"] @@ -126,6 +156,127 @@ def test_handles_empty_stage_usage(self) -> None: self.assertEqual(payload["totals"]["cache_hit_rate"], 0.0) self.assertEqual(payload["per_model"], {}) + def test_records_estimate_and_actual_error(self) -> None: + stage_usage = { + "judge": { + "calls": 2, + "input_tokens": 800, + "output_tokens": 200, + "cached_input_tokens": 0, + "cache_creation_input_tokens": 0, + "per_model": {}, + }, + } + estimate = { + "total_tokens": 1_100, + "input_tokens": 900, + "output_tokens": 200, + "calls": 2, + } + + payload = _build_run_metrics( + stage_usage, + total_elapsed=1.0, + token_estimate=estimate, + ) + + self.assertEqual(payload["token_estimate"], estimate) + accuracy = payload["token_estimate_accuracy"] + self.assertEqual(accuracy["status"], "available") + self.assertEqual(accuracy["actual_total_tokens"], 1_000) + self.assertEqual(accuracy["difference_tokens"], -100) + self.assertAlmostEqual(accuracy["difference_ratio"], -100 / 1_100) + self.assertAlmostEqual( + accuracy["absolute_percentage_error"], + 100 / 1_100, + ) + + def test_marks_accuracy_unavailable_when_usage_is_incomplete(self) -> None: + stage_usage = { + "judge": { + "requests": 2, + "calls": 1, + "missing_usage_calls": 1, + "input_tokens": 800, + "output_tokens": 200, + "cached_input_tokens": 0, + "cache_creation_input_tokens": 0, + "per_model": {}, + }, + } + + payload = _build_run_metrics( + stage_usage, + total_elapsed=1.0, + token_estimate={"total_tokens": 1_100}, + ) + + accuracy = payload["token_estimate_accuracy"] + self.assertEqual(accuracy["status"], "unavailable") + self.assertEqual(accuracy["reason"], "provider_usage_incomplete") + self.assertEqual(accuracy["usage_coverage"], 0.5) + + def test_marks_accuracy_unavailable_when_pipeline_fails(self) -> None: + payload = _build_run_metrics( + { + "judge": { + "requests": 1, + "calls": 1, + "input_tokens": 800, + "output_tokens": 200, + "per_model": {}, + }, + }, + total_elapsed=1.0, + token_estimate={"total_tokens": 1_100}, + run_completed=False, + ) + + accuracy = payload["token_estimate_accuracy"] + self.assertEqual(accuracy["status"], "unavailable") + self.assertEqual(accuracy["reason"], "pipeline_incomplete") + + def test_marks_accuracy_unavailable_when_pipeline_is_partial(self) -> None: + payload = _build_run_metrics( + { + "judge": { + "requests": 1, + "calls": 1, + "total_tokens": 1_000, + "input_tokens": 800, + "output_tokens": 200, + "per_model": {}, + }, + }, + total_elapsed=1.0, + token_estimate={"total_tokens": 1_100}, + run_partial=True, + ) + + accuracy = payload["token_estimate_accuracy"] + self.assertEqual(accuracy["status"], "unavailable") + self.assertEqual(accuracy["reason"], "pipeline_partial") + + def test_accuracy_uses_provider_total_when_breakdown_is_missing(self) -> None: + payload = _build_run_metrics( + { + "judge": { + "requests": 1, + "calls": 1, + "total_tokens": 1_000, + "input_tokens": 0, + "output_tokens": 0, + "per_model": {}, + }, + }, + total_elapsed=1.0, + token_estimate={"total_tokens": 1_100}, + ) + + accuracy = payload["token_estimate_accuracy"] + self.assertEqual(accuracy["status"], "available") + self.assertEqual(accuracy["actual_total_tokens"], 1_000) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_token_estimator.py b/tests/test_token_estimator.py new file mode 100644 index 000000000..29f88c627 --- /dev/null +++ b/tests/test_token_estimator.py @@ -0,0 +1,1645 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import hashlib +import json +import os +import unittest +from copy import deepcopy +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + +import yaml + +from assert_ai.core.artifact_cache import ( + activate_artifact_plan, + finalize_artifact_plan, + prepare_artifact_plan, +) +from assert_ai.core.config_model import ( + DEFAULT_INFERENCE_MAX_TOKENS, + EvaluationConfig, + InferenceConfig, + JudgeConfig, + ModelConfig, + TargetConfig, + TesterConfig, + ToolsConfig, +) +from assert_ai.core.judge import build_judge_contract +from assert_ai.core.io import ( + load_jsonl, + normalize_test_case_rows, + write_jsonl, +) +from assert_ai.core.model_client import estimate_token_count +from assert_ai.core.token_estimator import ( + _CaseProfile, + _high_side_prompt_output, + _project_prompt_case, + _project_scenario_case, + _representative_tool_value, + _response_length_hint, + _target_output, + _transcript_xml, + estimate_pipeline_tokens, +) +from assert_ai.core.transcript import ( + AddMessageEdit, + Message as TranscriptMessage, + Transcript, + TranscriptEvent, + TranscriptMetadata, +) +from assert_ai.runner import estimate_pipeline_usage +from assert_ai.stages import inference as inference_stage +from assert_ai.stages import judge as judge_stage + + +def _write_jsonl(path: Path, rows: list[dict[str, object]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), + encoding="utf-8", + ) + + +def _base_context(root: Path) -> dict[str, object]: + config_path = root / "config.yaml" + config_path.write_text("pipeline: {}\n", encoding="utf-8") + suite_root = root / "results" / "suite" + run_root = suite_root / "run" + suite_root.mkdir(parents=True, exist_ok=True) + return { + "config_path": config_path, + "artifacts_root": root, + "suite_root": suite_root, + "run_root": run_root, + "behavior_name": "representative_behavior", + "behavior": "The target must follow the configured behavior.", + "context": "A representative application context.", + "dimensions": [], + } + + +def _write_taxonomy(path: Path, category_count: int = 2) -> None: + path.write_text( + json.dumps( + { + "behavior": { + "name": "representative_behavior", + "definition": "Required behavior.", + }, + "definition_of_terms": [], + "behavior_categories": [ + { + "name": f"category_{index + 1}", + "definition": "Representative category.", + "examples": ["Representative example."], + "permissible": False, + } + for index in range(category_count) + ], + } + ), + encoding="utf-8", + ) + + +def _record_cached_compatibility_file( + suite_root: Path, + *, + stage_name: str, + output_key: str, + compatibility_path: Path, +) -> None: + version_dir = suite_root / "artifacts" / stage_name / "v0001" + version_dir.mkdir(parents=True, exist_ok=True) + digest = hashlib.sha256(compatibility_path.read_bytes()).hexdigest() + (version_dir / "artifact.json").write_text( + json.dumps( + { + "files": {output_key: compatibility_path.name}, + "file_hashes": {output_key: digest}, + } + ), + encoding="utf-8", + ) + + +def _tool_profile(kind: str) -> _CaseProfile: + return _CaseProfile( + kind=kind, + test_case_id="tool-case", + description="Calculate the sum of two numbers.", + system_prompt="Use the calculator.", + tools=({ + "name": "calculate", + "parameters": [ + {"name": "left", "type": "integer"}, + {"name": "right", "type": "integer"}, + ], + },), + ) + + +class TokenEstimatorTest(unittest.TestCase): + def test_explicit_response_limits_have_conservative_headroom(self) -> None: + for instruction, expected in [ + ("Reply in one sentence.", 96), + ("Reply with the number and one short sentence.", 96), + ("Define eventual consistency in no more than two sentences.", 160), + ("Please respond using at most 50 words.", 132), + ("Output exactly 100 tokens.", 157), + ("Use exactly three sentences.", 224), + ("Calculate the checksum and reply with one sentence.", 96), + ("Explain caching in no more than two sentences.", 160), + ]: + with self.subTest(instruction=instruction): + self.assertEqual(_response_length_hint(instruction), expected) + self.assertEqual( + _response_length_hint("Reply in one sentence.", "Reply in three sentences."), 224, + ) + + def test_length_hints_ignore_quoted_nested_negative_and_lower_bound_instructions(self) -> None: + for instruction in [ + 'Explain the phrase "reply in one sentence".', + 'Reply with "in one sentence" as part of a longer explanation.', + "```text\nReply in one sentence.\n```", + "```text\nReply in one sentence.", + 'Quoted task:\n"Introduction.\nReply in one sentence.\nEnd."', + 'Unclosed quote:\n"Introduction.\nReply in one sentence.', + "Never reply in one sentence.", + "Reply in not more than two sentences.", + "Write at least 50 words.", + "Write exactly two sentences per item.", + "Write eight items, each with two sentences.", + "Reply in one sentence, then write three paragraphs.", + "Return exactly three bullets.", + "Discuss 20 words that changed meaning.", + "Describe a diagram with four words highlighted.", + "Reply in one sentence or in two sentences.", + ]: + with self.subTest(instruction=instruction): + self.assertIsNone(_response_length_hint(instruction)) + + def test_prompt_length_hint_changes_projected_answer_and_judge_input(self) -> None: + target = TargetConfig(model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=256)) + estimate, transcript, notes = _project_prompt_case( + _CaseProfile("prompt", "short", "Reply in one sentence."), + ctx={}, target=target, max_tokens=256, + ) + self.assertEqual(estimate.output_tokens, 96) + self.assertEqual(transcript.transcript_xml.count("response"), 96) + self.assertTrue(any("response-length instructions" in note for note in notes)) + capped, _, _ = _project_prompt_case( + _CaseProfile("prompt", "short", "Reply in one sentence."), + ctx={}, target=TargetConfig(model=ModelConfig(name=target.model.name, max_tokens=64)), + max_tokens=64, + ) + self.assertEqual(capped.output_tokens, 56) + + def test_scenario_description_is_not_a_per_turn_answer_limit(self) -> None: + target = TargetConfig(model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=256)) + evaluation = EvaluationConfig( + tester=TesterConfig(model=target.model), + inference=InferenceConfig(max_turns=2), + ) + estimate, _, notes = _project_scenario_case( + _CaseProfile("scenario", "multi", "Reply in one sentence."), + ctx={}, target=target, evaluation=evaluation, max_tokens=256, + ) + self.assertEqual(estimate.output_tokens, 2 * (224 + 55)) + self.assertFalse(any("response-length instructions" in note for note in notes)) + + def test_target_output_leaves_headroom_without_lowering_large_budget_baseline( + self, + ) -> None: + for limit, expected in [ + (1, 1), (64, 56), (256, 224), (512, 448), + (900, 675), (1_000, 750), (4_000, 768), (None, 512), + ]: + with self.subTest(limit=limit): + self.assertEqual(_high_side_prompt_output(limit), expected) + self.assertEqual(_target_output(384, 256), 224) + self.assertEqual(_target_output(384, 1_000), 384) + + def test_projected_transcript_matches_runtime_escaping_and_truncation(self) -> None: + messages = [ + ("system", 'Keep "quotes" & .'), + ("user", "x" * 10_001), + ("assistant", ""), + ("tool", "An ordinary tool result."), + ] + runtime = Transcript( + metadata=TranscriptMetadata( + kind="prompt", test_case_id="p1", behavior="", target="", tester_model="", + ), + events=[ + TranscriptEvent( + view=["target"], actor="target", + edit=AddMessageEdit(message=TranscriptMessage(role=role, content=content)), + ) + for role, content in messages + ], + ) + expected, _ = runtime.format_transcript_xml("target", skip_system=False) + projected = _transcript_xml(messages) + self.assertEqual(projected, expected) + self.assertIn('truncated="true"', projected) + self.assertIn("10001 chars total", projected) + self.assertIn("<tags>", projected) + self.assertNotIn("x" * 10_001, projected) + self.assertEqual(_transcript_xml([]), "\n") + + def test_representative_tool_arguments_preserve_schema_shape(self) -> None: + schema = { + "type": "object", + "properties": { + "city": {"type": "string"}, + "count": {"type": "integer"}, + "ratio": {"type": "number"}, + "enabled": {"type": "boolean"}, + "mode": {"enum": ["fast", "slow"]}, + "items": {"type": "array", "items": {"type": "object", "properties": {}}}, + "unconstrained": {"type": "array", "items": True}, + }, + } + self.assertEqual(_representative_tool_value(schema), { + "city": "representative value", "count": 1, "ratio": 1, + "enabled": True, "mode": "fast", "items": [{}], + "unconstrained": ["representative value"], + }) + + def test_tool_projection_counts_call_json_instead_of_two_final_answers(self) -> None: + for simulator in (None, "openai/gpt-4o-mini"): + with self.subTest(simulator=simulator): + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=400), + tools=( + ToolsConfig(simulator=simulator) + if simulator + else ToolsConfig(module="example.tools") + ), + ) + estimate, transcript, notes = _project_prompt_case( + _tool_profile("prompt"), ctx={}, target=target, max_tokens=400, + ) + call_tokens = estimate_token_count( + target.model.name, + text=json.dumps({ + "name": "calculate", "arguments": {"left": 1, "right": 1}, + }), + ) + self.assertEqual(estimate.calls, 3 if simulator else 2) + self.assertEqual( + estimate.output_tokens, + 350 + call_tokens + (90 if simulator else 0), + ) + self.assertIn("[Tool call: calculate(", transcript.transcript_xml) + self.assertIn(""left": 1", transcript.transcript_xml) + self.assertIn(" None: + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=400), + tools=ToolsConfig(simulator="openai/gpt-4o"), + ) + evaluation = EvaluationConfig( + tester=TesterConfig(model=ModelConfig(name="openai/gpt-4o-mini")), + inference=InferenceConfig(max_turns=2), + ) + requests = [] + + def record_request(model, messages, **kwargs): + requests.append((model, deepcopy(messages), kwargs)) + return 100 + + with ( + patch("assert_ai.core.token_estimator._request_tokens", side_effect=record_request), + patch("assert_ai.core.session.generate", side_effect=AssertionError("Provider called")), + ): + estimate, transcript, _ = _project_scenario_case( + _tool_profile("scenario"), ctx={}, target=target, + evaluation=evaluation, max_tokens=400, + ) + + self.assertEqual(estimate.calls, 8) + self.assertEqual(estimate.input_tokens, 800) + target_requests = [ + messages for _, messages, kwargs in requests if kwargs.get("tools") + ] + self.assertEqual( + [sum(message.role == "tool" for message in messages) for messages in target_requests], + [0, 1, 1, 2], + ) + self.assertEqual(target_requests[1][2].tool_calls[0].arguments, {"left": 1, "right": 1}) + self.assertEqual(target_requests[1][3].tool_call_id, target_requests[1][2].tool_calls[0].id) + self.assertNotEqual(target_requests[3][-1].tool_call_id, target_requests[1][-1].tool_call_id) + simulator_prompts = [ + messages for model, messages, _ in requests if model == target.tools.simulator + ] + self.assertEqual(len(simulator_prompts), 2) + for prompt in simulator_prompts: + self.assertNotIn("{{", prompt) + self.assertIn("Calculate the sum of two numbers.", prompt) + self.assertIn("User: request", prompt) + self.assertNotIn("Use the calculator.", prompt) + self.assertIn("(none yet)", simulator_prompts[0]) + self.assertIn('- calculate({"left": 1, "right": 1}) -> result', simulator_prompts[1]) + self.assertIn("Target: response", simulator_prompts[1]) + self.assertEqual(transcript.transcript_xml.count("[Tool call: calculate("), 2) + + def test_config_estimate_is_read_only(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + config_path = root / "eval.yaml" + artifacts_root = root / "artifacts" + config_path.write_text( + yaml.safe_dump( + { + "suite": "preview-suite", + "run": "preview-run", + "artifacts_root": str(artifacts_root), + "behavior": { + "name": "answer_accuracy", + "description": "Answer accurately.", + }, + "context": "A factual question answering assistant.", + "pipeline": { + "systematize": { + "behavior_category_count": 2, + "web_search": False, + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 2_000, + }, + }, + "test_set": { + "prompt": { + "sample_size": 2, + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 1_000, + }, + } + }, + "inference": { + "target": { + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 512, + } + } + }, + "judge": { + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 1_000, + } + }, + }, + } + ), + encoding="utf-8", + ) + + estimate = estimate_pipeline_usage(config=str(config_path)) + + self.assertGreater(estimate["total_tokens"], 0) + self.assertEqual( + set(estimate["stages"]), + {"systematize", "test_set", "inference", "judge"}, + ) + self.assertFalse(artifacts_root.exists()) + + def test_tool_limit_changes_upper_range_including_judge_but_not_point_estimate(self) -> None: + for kind in ("prompt", "scenario"): + with self.subTest(kind=kind), TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + profile = _tool_profile(kind) + _write_taxonomy(suite_root / "taxonomy.json") + _write_jsonl(suite_root / "test_set.jsonl", [{ + "type": kind, "test_case_id": profile.test_case_id, + "seed": { + "description": profile.description, + "system_prompt": profile.system_prompt, + "tools": list(profile.tools), + }, + }]) + model = ModelConfig(name="openai/gpt-4o-mini", max_tokens=400) + ctx["target"] = TargetConfig( + model=model, tools=ToolsConfig(simulator=model.name), + ) + results = [] + for limit in (1, 10): + ctx["evaluation"] = EvaluationConfig( + tester=TesterConfig(model=model), judge=JudgeConfig(model=model), + inference=InferenceConfig(max_turns=2, max_tool_calls=limit), + ) + inference = estimate_pipeline_tokens(ctx, [("inference", object(), {})]) + combined = estimate_pipeline_tokens( + ctx, [("inference", object(), {}), ("judge", object(), {})], + ) + results.append(combined) + self.assertGreater( + combined.tool_loop_total_tokens - combined.total_tokens, + inference.tool_loop_total_tokens - inference.total_tokens, + ) + self.assertTrue(any(f"up to {limit} resolved" in n for n in combined.notes)) + self.assertEqual(results[0].total_tokens, results[1].total_tokens) + self.assertGreater(results[1].upper_bound_tokens, results[0].upper_bound_tokens) + + def test_upper_tool_projection_counts_each_resolver_and_limit_fallback(self) -> None: + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=400), + tools=ToolsConfig(simulator="openai/gpt-4o"), + ) + requests = [] + + def record_request(model, messages, **kwargs): + requests.append((model, deepcopy(messages), kwargs)) + return 100 + + with patch("assert_ai.core.token_estimator._request_tokens", side_effect=record_request): + estimate, transcript, _ = _project_prompt_case( + _tool_profile("prompt"), ctx={}, target=target, max_tokens=400, + tool_rounds=3, include_limit_fallback=True, + ) + self.assertEqual(estimate.calls, 8) # Initial + 3 follow-ups + forced final + 3 resolvers. + self.assertEqual(sum(model == target.tools.simulator for model, _, _ in requests), 3) + self.assertEqual( + [sum(m.role == "tool" for m in messages) for model, messages, _ in requests + if model == target.model.name], + [0, 1, 2, 3, 4], + ) + self.assertIsNone(requests[-1][2]["tools"]) + self.assertEqual(requests[-1][1][-1].text, "Tool call limit reached.") + self.assertEqual(transcript.transcript_xml.count("[Tool call: calculate("), 4) + + def test_inference_only_estimate_reads_versioned_test_set_without_writes( + self, + ) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + config_path = root / "eval.yaml" + artifacts_root = root / "artifacts" + results_dir = artifacts_root / "results" + suite_root = results_dir / "preview-suite" + suite_root.mkdir(parents=True) + cache_ctx = { + "config_path": config_path, + "artifacts_root": artifacts_root, + "suite_root": suite_root, + "behavior_name": "answer_accuracy", + "behavior": "Answer accurately.", + "context": "A factual assistant.", + "artifact_versions": {}, + } + raw_test_set = { + "prompt": { + "sample_size": 2, + "model": {"name": "openai/gpt-4o-mini"}, + } + } + plan = prepare_artifact_plan( + ctx=cache_ctx, + stage_name="test_set", + raw_cfg=raw_test_set, + forced=False, + ) + activate_artifact_plan(cache_ctx, plan) + _write_jsonl( + plan.output_paths["test_set"], + [ + { + "type": "prompt", + "test_case_id": f"test_case_{index:06d}", + "seed": { + "description": f"Question {index}.", + "system_prompt": "Answer accurately.", + }, + } + for index in (1, 2) + ], + ) + plan.output_paths["stratification"].write_text( + "{}", + encoding="utf-8", + ) + finalize_artifact_plan(cache_ctx, plan) + compatibility_path = suite_root / "test_set.jsonl" + compatibility_path.unlink() + latest_path = suite_root / "latest.json" + latest_before = latest_path.read_bytes() + + config_path.write_text( + yaml.safe_dump( + { + "suite": "preview-suite", + "run": "preview-run", + "artifacts_root": str(artifacts_root), + "context": "A factual assistant.", + "pipeline": { + "inference": { + "test_set_path": str(compatibility_path), + "target": { + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 512, + } + } + } + }, + } + ), + encoding="utf-8", + ) + + estimate = estimate_pipeline_usage(config=str(config_path)) + + self.assertEqual(estimate["stages"]["inference"]["calls"], 2) + self.assertEqual(latest_path.read_bytes(), latest_before) + self.assertFalse(compatibility_path.exists()) + self.assertFalse((suite_root / "preview-run").exists()) + + def test_judge_only_estimate_reads_versioned_taxonomy_without_writes( + self, + ) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + config_path = root / "eval.yaml" + artifacts_root = root / "artifacts" + suite_root = artifacts_root / "results" / "preview-suite" + run_root = suite_root / "preview-run" + suite_root.mkdir(parents=True) + cache_ctx = { + "config_path": config_path, + "artifacts_root": artifacts_root, + "suite_root": suite_root, + "behavior_name": "answer_accuracy", + "behavior": "Answer accurately.", + "context": "A factual assistant.", + "artifact_versions": {}, + } + raw_systematize = { + "behavior_category_count": 2, + "model": {"name": "openai/gpt-4o-mini"}, + } + plan = prepare_artifact_plan( + ctx=cache_ctx, + stage_name="systematize", + raw_cfg=raw_systematize, + forced=False, + ) + activate_artifact_plan(cache_ctx, plan) + _write_taxonomy(plan.output_paths["taxonomy"]) + plan.output_paths["systematization"].write_text( + "{}", + encoding="utf-8", + ) + finalize_artifact_plan(cache_ctx, plan) + compatibility_path = suite_root / "taxonomy.json" + compatibility_path.unlink() + run_root.mkdir() + inference_path = run_root / "inference_set.jsonl" + _write_jsonl( + inference_path, + [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "events": [], + "stop_reason": "completed", + } + ], + ) + latest_path = suite_root / "latest.json" + latest_before = latest_path.read_bytes() + + config_path.write_text( + yaml.safe_dump( + { + "suite": "preview-suite", + "run": "preview-run", + "artifacts_root": str(artifacts_root), + "behavior": { + "name": "answer_accuracy", + "description": "Answer accurately.", + }, + "context": "A factual assistant.", + "pipeline": { + "judge": { + "taxonomy_path": str(compatibility_path), + "inference_set_path": str(inference_path), + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 1_000, + }, + } + }, + } + ), + encoding="utf-8", + ) + + estimate = estimate_pipeline_usage(config=str(config_path)) + + self.assertEqual(estimate["stages"]["judge"]["calls"], 1) + self.assertEqual(latest_path.read_bytes(), latest_before) + self.assertFalse(compatibility_path.exists()) + + def test_hosted_prompt_run_estimates_target_and_judge_calls(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + _write_taxonomy(suite_root / "taxonomy.json") + _write_jsonl( + suite_root / "test_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "p1", + "seed": { + "description": "Explain the first result.", + "system_prompt": "Answer accurately.", + }, + }, + { + "type": "prompt", + "test_case_id": "p2", + "seed": { + "description": "Explain the second result.", + "system_prompt": "Answer accurately.", + }, + }, + ], + ) + ctx["target"] = TargetConfig( + model=ModelConfig( + name="openai/gpt-4o-mini", + max_tokens=1_000, + ) + ) + ctx["evaluation"] = EvaluationConfig( + judge=JudgeConfig( + model=ModelConfig( + name="openai/gpt-4o-mini", + max_tokens=1_000, + ) + ) + ) + + estimate = estimate_pipeline_tokens( + ctx, + [ + ("inference", object(), {}), + ("judge", object(), {}), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 2) + self.assertEqual(estimate.stages["judge"].calls, 2) + self.assertEqual(estimate.stages["inference"].output_tokens, 1_500) + self.assertEqual(estimate.stages["judge"].output_tokens, 1_024) + self.assertGreater(estimate.input_tokens, 0) + self.assertGreater(estimate.output_tokens, 0) + self.assertEqual(estimate.uncertainty, 0.35) + self.assertTrue( + any("high-side output assumptions" in note for note in estimate.notes) + ) + self.assertLess( + estimate.lower_bound_tokens, + estimate.total_tokens, + ) + self.assertGreater( + estimate.upper_bound_tokens, + estimate.total_tokens, + ) + + def test_judge_output_grows_with_contract_and_respects_completion_limit(self) -> None: + outputs = {} + with TemporaryDirectory() as tmp_dir: + ctx = _base_context(Path(tmp_dir)) + suite_root = Path(ctx["suite_root"]) + _write_jsonl( + suite_root / "test_set.jsonl", + [{"type": "prompt", "test_case_id": "p1", "seed": {"description": "Answer."}}], + ) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=256), + ) + for categories, limit in [(1, 1_000), (20, 1_000), (20, 128)]: + _write_taxonomy(suite_root / "taxonomy.json", categories) + ctx["evaluation"] = EvaluationConfig( + judge=JudgeConfig( + model=ModelConfig(name="openai/gpt-4o-mini", max_tokens=limit), + ), + ) + estimate = estimate_pipeline_tokens( + ctx, [("inference", object(), {}), ("judge", object(), {})], + ) + outputs[(categories, limit)] = estimate.stages["judge"].output_tokens + + self.assertEqual(outputs[(1, 1_000)], 512) + self.assertGreater(outputs[(20, 1_000)], outputs[(1, 1_000)]) + self.assertLessEqual(outputs[(20, 1_000)], 1_000) + self.assertEqual(outputs[(20, 128)], 128) + + def test_callable_scenario_excludes_unknown_target_usage(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + _write_taxonomy(suite_root / "taxonomy.json") + _write_jsonl( + suite_root / "test_set.jsonl", + [ + { + "type": "scenario", + "test_case_id": "s1", + "seed": { + "description": "Apply pressure over several turns.", + "system_prompt": "Follow policy.", + }, + } + ], + ) + model = ModelConfig( + name="openai/gpt-4o-mini", + max_tokens=1_000, + ) + ctx["target"] = TargetConfig(callable="example.agent:chat") + ctx["evaluation"] = EvaluationConfig( + tester=TesterConfig(model=model), + judge=JudgeConfig(model=model), + inference=InferenceConfig(max_turns=3), + ) + + estimate = estimate_pipeline_tokens( + ctx, + [ + ("inference", object(), {}), + ("judge", object(), {}), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 3) + self.assertEqual(estimate.stages["judge"].calls, 1) + self.assertTrue( + any("callable target" in note for note in estimate.notes) + ) + + def test_first_run_estimates_generated_taxonomy_and_test_cases(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + ctx["dimensions"] = [ + { + "name": "pressure", + "description": "Pressure level.", + "levels": [ + {"name": "low", "definition": "Low pressure."}, + {"name": "high", "definition": "High pressure."}, + ], + } + ] + systematize_cfg = { + "behavior_category_count": 4, + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 4_000, + }, + } + test_set_cfg = { + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 3_000, + }, + "prompt": {"sample_size": 6}, + "scenario": {"sample_size": 3}, + } + + estimate = estimate_pipeline_tokens( + ctx, + [ + ("systematize", object(), systematize_cfg), + ("test_set", object(), test_set_cfg), + ], + ) + + self.assertEqual(estimate.stages["systematize"].calls, 2) + self.assertGreaterEqual(estimate.stages["test_set"].calls, 2) + self.assertGreater(estimate.total_tokens, 1_000) + self.assertTrue( + any("representative generated taxonomy" in note for note in estimate.notes) + ) + + def test_empty_test_set_kind_is_disabled_like_runtime(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + _write_taxonomy(Path(ctx["suite_root"]) / "taxonomy.json") + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "model": {"name": "openai/gpt-4o-mini"}, + "prompt": {}, + "scenario": {"sample_size": 1}, + }, + ) + ], + ) + + self.assertEqual(estimate.stages["test_set"].calls, 1) + + def test_legacy_per_seed_matches_per_test_case_estimate(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + _write_taxonomy(Path(ctx["suite_root"]) / "taxonomy.json") + model = ModelConfig(name="openai/gpt-4o-mini") + ctx["target"] = TargetConfig( + model=model, + tools=ToolsConfig(simulator=model.name), + ) + ctx["evaluation"] = EvaluationConfig() + + def estimate_for(tool_source: str): + return estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "tool_source": tool_source, + "model": {"name": model.name}, + "prompt": {"sample_size": 1}, + }, + ), + ("inference", object(), {}), + ], + ) + + legacy = estimate_for("per_seed") + canonical = estimate_for("per_test_case") + + self.assertEqual(legacy.to_dict(), canonical.to_dict()) + self.assertGreater(legacy.stages["inference"].calls, 1) + + def test_inference_resume_counts_only_pending_cases_unless_forced(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + test_set_path = suite_root / "test_set.jsonl" + _write_jsonl( + test_set_path, + [ + { + "type": "prompt", + "test_case_id": case_id, + "seed": { + "description": f"Prompt {case_id}.", + "system_prompt": "Answer accurately.", + }, + } + for case_id in ("p1", "p2", "p3") + ], + ) + write_jsonl( + test_set_path, + normalize_test_case_rows(load_jsonl(test_set_path)), + ) + model = ModelConfig( + name="openai/gpt-4o-mini", + max_tokens=1_000, + ) + target = TargetConfig(model=model) + evaluation = EvaluationConfig() + ctx["target"] = target + ctx["evaluation"] = evaluation + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "events": [], + } + ], + ) + fingerprint = inference_stage._inference_config_fingerprint( + target, + evaluation, + DEFAULT_INFERENCE_MAX_TOKENS, + test_set_path=test_set_path, + config_path=Path(ctx["config_path"]), + ) + ( + run_root / inference_stage._INFERENCE_CONFIG_HASH_FILE + ).write_text( + fingerprint, + encoding="utf-8", + ) + + resumed = estimate_pipeline_tokens( + ctx, + [("inference", object(), {})], + ) + forced = estimate_pipeline_tokens( + ctx, + [("inference", object(), {})], + forced_stages={"inference"}, + ) + + self.assertEqual(resumed.stages["inference"].calls, 2) + self.assertEqual(forced.stages["inference"].calls, 3) + + def test_inference_resume_hashes_runtime_canonical_test_set(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + test_set_path = suite_root / "test_set.jsonl" + _write_jsonl( + test_set_path, + [ + { + "type": "prompt", + "test_case_id": "legacy-id", + "seed": { + "description": "Answer the prompt.", + "system_prompt": "Answer accurately.", + }, + } + ], + ) + canonical_rows = normalize_test_case_rows( + load_jsonl(test_set_path) + ) + canonical_content = ( + os.linesep.join( + json.dumps(row, ensure_ascii=False) + for row in canonical_rows + ) + + os.linesep + ).encode("utf-8") + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + evaluation = EvaluationConfig() + ctx["target"] = target + ctx["evaluation"] = evaluation + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "events": [], + } + ], + ) + fingerprint = inference_stage._inference_config_fingerprint( + target, + evaluation, + DEFAULT_INFERENCE_MAX_TOKENS, + test_set_path=test_set_path, + config_path=Path(ctx["config_path"]), + test_set_content=canonical_content, + ) + ( + run_root / inference_stage._INFERENCE_CONFIG_HASH_FILE + ).write_text(fingerprint, encoding="utf-8") + + estimate = estimate_pipeline_tokens( + ctx, + [("inference", object(), {})], + ) + + self.assertNotIn("inference", estimate.stages) + + def test_unrelated_test_set_output_does_not_invalidate_inference(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + _write_taxonomy(suite_root / "taxonomy.json") + explicit_test_set = root / "fixed_test_set.jsonl" + _write_jsonl( + explicit_test_set, + [ + { + "type": "prompt", + "seed": { + "description": "Use the fixed input.", + "system_prompt": "Answer accurately.", + }, + } + ], + ) + write_jsonl( + explicit_test_set, + normalize_test_case_rows(load_jsonl(explicit_test_set)), + ) + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + evaluation = EvaluationConfig() + ctx["target"] = target + ctx["evaluation"] = evaluation + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "events": [], + } + ], + ) + fingerprint = inference_stage._inference_config_fingerprint( + target, + evaluation, + DEFAULT_INFERENCE_MAX_TOKENS, + test_set_path=explicit_test_set, + config_path=Path(ctx["config_path"]), + ) + ( + run_root / inference_stage._INFERENCE_CONFIG_HASH_FILE + ).write_text(fingerprint, encoding="utf-8") + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "model": {"name": "openai/gpt-4o-mini"}, + "prompt": {"sample_size": 1}, + }, + ), + ( + "inference", + object(), + {"test_set_path": str(explicit_test_set)}, + ), + ], + ) + + self.assertIn("test_set", estimate.stages) + self.assertNotIn("inference", estimate.stages) + + def test_cache_compatibility_test_set_invalidates_inference(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + _write_taxonomy(suite_root / "taxonomy.json") + compatibility_path = suite_root / "test_set.jsonl" + _write_jsonl( + compatibility_path, + [ + { + "type": "prompt", + "seed": { + "description": "Old cached prompt.", + "system_prompt": "Answer accurately.", + }, + } + ], + ) + _record_cached_compatibility_file( + suite_root, + stage_name="test_set", + output_key="test_set", + compatibility_path=compatibility_path, + ) + next_output = ( + suite_root + / "artifacts" + / "test_set" + / "v0002" + / "test_set.jsonl" + ) + ctx["artifact_versions"] = {"test_set": {"version": "v0002"}} + ctx["test_set_path"] = str(next_output) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig() + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "save_path": str(next_output), + "model": {"name": "openai/gpt-4o-mini"}, + "prompt": {"sample_size": 3}, + }, + ), + ( + "inference", + object(), + {"test_set_path": str(compatibility_path)}, + ), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 3) + + def test_local_test_set_edit_is_not_treated_as_cache_alias(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + _write_taxonomy(suite_root / "taxonomy.json") + compatibility_path = suite_root / "test_set.jsonl" + _write_jsonl( + compatibility_path, + [ + { + "type": "prompt", + "seed": { + "description": "Locally edited prompt.", + "system_prompt": "Answer accurately.", + }, + } + ], + ) + next_output = ( + suite_root + / "artifacts" + / "test_set" + / "v0002" + / "test_set.jsonl" + ) + ctx["artifact_versions"] = {"test_set": {"version": "v0002"}} + ctx["test_set_path"] = str(next_output) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig() + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "save_path": str(next_output), + "model": {"name": "openai/gpt-4o-mini"}, + "prompt": {"sample_size": 3}, + }, + ), + ( + "inference", + object(), + {"test_set_path": str(compatibility_path)}, + ), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 1) + + def test_partial_inference_merges_completed_and_projected_transcripts( + self, + ) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + _write_taxonomy(suite_root / "taxonomy.json") + test_set_path = suite_root / "test_set.jsonl" + _write_jsonl( + test_set_path, + [ + { + "type": "prompt", + "seed": { + "description": f"Prompt {index}.", + "system_prompt": "Answer accurately.", + }, + } + for index in (1, 2) + ], + ) + write_jsonl( + test_set_path, + normalize_test_case_rows(load_jsonl(test_set_path)), + ) + target = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + evaluation = EvaluationConfig( + judge=JudgeConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ) + ctx["target"] = target + ctx["evaluation"] = evaluation + _write_jsonl( + run_root / "inference_set.jsonl", + [ + { + "type": "prompt", + "test_case_id": "test_case_000001", + "events": [], + "stop_reason": "target_error", + } + ], + ) + fingerprint = inference_stage._inference_config_fingerprint( + target, + evaluation, + DEFAULT_INFERENCE_MAX_TOKENS, + test_set_path=test_set_path, + config_path=Path(ctx["config_path"]), + ) + ( + run_root / inference_stage._INFERENCE_CONFIG_HASH_FILE + ).write_text(fingerprint, encoding="utf-8") + + estimate = estimate_pipeline_tokens( + ctx, + [ + ("inference", object(), {}), + ("judge", object(), {}), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 1) + self.assertEqual(estimate.stages["judge"].calls, 1) + + def test_judge_resume_counts_only_pending_scores_unless_forced(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + taxonomy_path = suite_root / "taxonomy.json" + _write_taxonomy(taxonomy_path) + taxonomy = json.loads(taxonomy_path.read_text(encoding="utf-8")) + inference_path = run_root / "inference_set.jsonl" + _write_jsonl( + inference_path, + [ + { + "type": "prompt", + "test_case_id": case_id, + "events": [], + "stop_reason": "completed", + } + for case_id in ("p1", "p2") + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [{"type": "prompt", "test_case_id": "p1"}], + ) + judge_cfg = JudgeConfig( + model=ModelConfig( + name="openai/gpt-4o-mini", + max_tokens=1_000, + ) + ) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig(judge=judge_cfg) + contract = build_judge_contract( + template=judge_stage.JUDGE_SYSTEM_PROMPT, + policy_raw=taxonomy, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + schema_name="transcript_judgment", + ) + fingerprint = judge_stage._judge_config_fingerprint( + judge_model=judge_cfg.model.name, + judge_temperature=judge_cfg.model.temperature, + judge_max_tokens=judge_cfg.model.max_tokens, + judge_reasoning_effort=judge_cfg.model.reasoning_effort, + judge_n=judge_cfg.n, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + policy_raw=taxonomy, + system_prompt=contract["system_prompt"], + inference_set_path=inference_path, + ) + (run_root / judge_stage._JUDGE_CONFIG_HASH_FILE).write_text( + fingerprint, + encoding="utf-8", + ) + + resumed = estimate_pipeline_tokens( + ctx, + [("judge", object(), {})], + ) + forced = estimate_pipeline_tokens( + ctx, + [("judge", object(), {})], + forced_stages={"judge"}, + ) + + self.assertEqual(resumed.stages["judge"].calls, 1) + self.assertEqual(forced.stages["judge"].calls, 2) + + def test_unrelated_inference_output_does_not_invalidate_judge(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + taxonomy_path = suite_root / "taxonomy.json" + _write_taxonomy(taxonomy_path) + taxonomy = json.loads(taxonomy_path.read_text(encoding="utf-8")) + test_set_path = suite_root / "test_set.jsonl" + _write_jsonl( + test_set_path, + [ + { + "type": "prompt", + "seed": { + "description": "Run unrelated inference.", + "system_prompt": "Answer accurately.", + }, + } + ], + ) + explicit_inference = root / "fixed_inference.jsonl" + _write_jsonl( + explicit_inference, + [ + { + "type": "prompt", + "test_case_id": "fixed-1", + "events": [], + "stop_reason": "completed", + } + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [{"type": "prompt", "test_case_id": "fixed-1"}], + ) + judge_cfg = JudgeConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig(judge=judge_cfg) + contract = build_judge_contract( + template=judge_stage.JUDGE_SYSTEM_PROMPT, + policy_raw=taxonomy, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + schema_name="transcript_judgment", + ) + fingerprint = judge_stage._judge_config_fingerprint( + judge_model=judge_cfg.model.name, + judge_temperature=judge_cfg.model.temperature, + judge_max_tokens=judge_cfg.model.max_tokens, + judge_reasoning_effort=judge_cfg.model.reasoning_effort, + judge_n=judge_cfg.n, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + policy_raw=taxonomy, + system_prompt=contract["system_prompt"], + inference_set_path=explicit_inference, + ) + (run_root / judge_stage._JUDGE_CONFIG_HASH_FILE).write_text( + fingerprint, + encoding="utf-8", + ) + + estimate = estimate_pipeline_tokens( + ctx, + [ + ("inference", object(), {}), + ( + "judge", + object(), + {"inference_set_path": str(explicit_inference)}, + ), + ], + ) + + self.assertEqual(estimate.stages["inference"].calls, 1) + self.assertNotIn("judge", estimate.stages) + + def test_test_set_taxonomy_does_not_invalidate_judge_resume(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + test_set_taxonomy_path = root / "test_set_taxonomy.json" + judge_taxonomy_path = root / "judge_taxonomy.json" + _write_taxonomy(test_set_taxonomy_path, category_count=3) + _write_taxonomy(judge_taxonomy_path, category_count=1) + judge_taxonomy = json.loads( + judge_taxonomy_path.read_text(encoding="utf-8") + ) + inference_path = run_root / "inference_set.jsonl" + _write_jsonl( + inference_path, + [ + { + "type": "prompt", + "test_case_id": "p1", + "events": [], + "stop_reason": "completed", + } + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [{"type": "prompt", "test_case_id": "p1"}], + ) + judge_cfg = JudgeConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig(judge=judge_cfg) + contract = build_judge_contract( + template=judge_stage.JUDGE_SYSTEM_PROMPT, + policy_raw=judge_taxonomy, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + schema_name="transcript_judgment", + ) + fingerprint = judge_stage._judge_config_fingerprint( + judge_model=judge_cfg.model.name, + judge_temperature=judge_cfg.model.temperature, + judge_max_tokens=judge_cfg.model.max_tokens, + judge_reasoning_effort=judge_cfg.model.reasoning_effort, + judge_n=judge_cfg.n, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + policy_raw=judge_taxonomy, + system_prompt=contract["system_prompt"], + inference_set_path=inference_path, + ) + (run_root / judge_stage._JUDGE_CONFIG_HASH_FILE).write_text( + fingerprint, + encoding="utf-8", + ) + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "test_set", + object(), + { + "taxonomy_path": str(test_set_taxonomy_path), + "save_path": str(root / "generated.jsonl"), + "model": {"name": "openai/gpt-4o-mini"}, + "prompt": {"sample_size": 1}, + }, + ), + ( + "judge", + object(), + {"taxonomy_path": str(judge_taxonomy_path)}, + ), + ], + ) + + self.assertIn("test_set", estimate.stages) + self.assertNotIn("judge", estimate.stages) + + def test_cache_compatibility_taxonomy_invalidates_judge(self) -> None: + with TemporaryDirectory() as tmp_dir: + root = Path(tmp_dir) + ctx = _base_context(root) + suite_root = Path(ctx["suite_root"]) + run_root = Path(ctx["run_root"]) + run_root.mkdir(parents=True) + compatibility_path = suite_root / "taxonomy.json" + _write_taxonomy(compatibility_path, category_count=1) + old_taxonomy = json.loads( + compatibility_path.read_text(encoding="utf-8") + ) + _record_cached_compatibility_file( + suite_root, + stage_name="systematize", + output_key="taxonomy", + compatibility_path=compatibility_path, + ) + next_output_dir = ( + suite_root / "artifacts" / "systematize" / "v0002" + ) + ctx["artifact_versions"] = { + "systematize": {"version": "v0002"} + } + ctx["systematize_artifact_dir"] = str(next_output_dir) + ctx["taxonomy_path"] = str(next_output_dir / "taxonomy.json") + inference_path = run_root / "inference_set.jsonl" + _write_jsonl( + inference_path, + [ + { + "type": "prompt", + "test_case_id": "p1", + "events": [], + "stop_reason": "completed", + } + ], + ) + _write_jsonl( + run_root / "scores.jsonl", + [{"type": "prompt", "test_case_id": "p1"}], + ) + judge_cfg = JudgeConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["target"] = TargetConfig( + model=ModelConfig(name="openai/gpt-4o-mini") + ) + ctx["evaluation"] = EvaluationConfig(judge=judge_cfg) + contract = build_judge_contract( + template=judge_stage.JUDGE_SYSTEM_PROMPT, + policy_raw=old_taxonomy, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + schema_name="transcript_judgment", + ) + fingerprint = judge_stage._judge_config_fingerprint( + judge_model=judge_cfg.model.name, + judge_temperature=judge_cfg.model.temperature, + judge_max_tokens=judge_cfg.model.max_tokens, + judge_reasoning_effort=judge_cfg.model.reasoning_effort, + judge_n=judge_cfg.n, + judge_dimensions=judge_cfg.dimensions, + disabled_dimensions=judge_cfg.disabled_dimensions, + policy_raw=old_taxonomy, + system_prompt=contract["system_prompt"], + inference_set_path=inference_path, + ) + (run_root / judge_stage._JUDGE_CONFIG_HASH_FILE).write_text( + fingerprint, + encoding="utf-8", + ) + + estimate = estimate_pipeline_tokens( + ctx, + [ + ( + "systematize", + object(), + { + "save_dir": str(next_output_dir), + "behavior_category_count": 3, + "model": { + "name": "openai/gpt-4o-mini", + "max_tokens": 4_000, + }, + }, + ), + ( + "judge", + object(), + {"taxonomy_path": str(compatibility_path)}, + ), + ], + ) + + self.assertEqual(estimate.stages["judge"].calls, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_viewer_run_spawn.py b/tests/test_viewer_run_spawn.py new file mode 100644 index 000000000..d39b2f60a --- /dev/null +++ b/tests/test_viewer_run_spawn.py @@ -0,0 +1,378 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import json +import os +import shutil +import subprocess +import textwrap +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +from tests.node_runner import node_supports_ts, node_ts_args + + +ROOT = Path(__file__).resolve().parents[1] +RUN_SPAWN_SRC = ROOT / "viewer" / "src" / "lib" / "server" / "run-spawn.ts" +ARTIFACTS_SRC = ROOT / "viewer" / "src" / "lib" / "server" / "artifacts.ts" +CONFIG_SRC = ROOT / "viewer" / "src" / "lib" / "server" / "config.ts" + + +@unittest.skipUnless(node_supports_ts(), "node binary lacks TypeScript support (need >= 22.6)") +class ViewerRunSpawnTest(unittest.TestCase): + def test_custom_roots_match_estimation_submission_and_execution(self) -> None: + for relative_root in (False, True): + with self.subTest(relative_root=relative_root), TemporaryDirectory( + dir=ROOT / "viewer" + ) as tmp_dir: + root = Path(tmp_dir) + harness = root / "harness" + harness.mkdir() + run_spawn_path = harness / "run-spawn.ts" + run_spawn_path.write_text( + RUN_SPAWN_SRC.read_text(encoding="utf-8") + .replace("./artifacts.js", "./artifacts.ts") + .replace("./config.js", "./config.ts"), + encoding="utf-8", + ) + (harness / "artifacts.ts").write_text( + ARTIFACTS_SRC.read_text(encoding="utf-8").replace( + "./config.js", "./config.ts" + ), + encoding="utf-8", + ) + shutil.copyfile(CONFIG_SRC, harness / "config.ts") + artifacts_root = root / "custom-cache" / "evaluation-results" + suite_dir = artifacts_root / "preview-suite" + suite_dir.mkdir(parents=True) + (suite_dir / "test_set.jsonl").write_text( + '{"test_case_id":"cached-case"}\n', encoding="utf-8" + ) + fake_cli = root / "fake-cli.mjs" + fake_cli.write_text( + textwrap.dedent( + """\ + import fs from 'node:fs'; + import path from 'node:path'; + import { parse } from 'yaml'; + const args = process.argv.slice(2); + const config = parse(fs.readFileSync(args[args.indexOf('--config') + 1], 'utf-8')); + const resultsRoot = path.resolve(config.results_dir ?? 'artifacts/results'); + const runDir = path.join(resultsRoot, config.suite, config.run); + const state = { + config, runDir, + reusedTestSet: fs.existsSync(path.join(resultsRoot, config.suite, 'test_set.jsonl')) + }; + fs.writeFileSync( + path.join(process.env.STATE_ROOT, `${args[0]}.json`), + JSON.stringify(state) + ); + if (args[0] === 'estimate') { + console.log(JSON.stringify({ + schema_version: 1, calls: 0, input_tokens: 0, output_tokens: 0, + total_tokens: 0, lower_bound_tokens: 0, upper_bound_tokens: 0, + stages: {}, notes: [] + })); + } + """ + ), + encoding="utf-8", + ) + env = os.environ.copy() + env.update( + { + "STATE_ROOT": str(root), + "ARTIFACTS_ROOT": ( + os.path.relpath(artifacts_root, harness) + if relative_root + else str(artifacts_root) + ), + "MEASUREMENTS_ROOT": str(root), + "ASSERT_AI_COMMAND": f"node {fake_cli}", + "TMP": str(root), + "TEMP": str(root), + "TMPDIR": str(root), + } + ) + script = textwrap.dedent( + f"""\ + import fs from 'node:fs'; + import path from 'node:path'; + import {{ parse }} from 'yaml'; + const {{ estimateAssertAiRun, writeRunConfigFiles, spawnAssertAiRun }} = + await import({json.dumps(run_spawn_path.as_uri())}); + const {{ runDirPath }} = await import({json.dumps((harness / 'artifacts.ts').as_uri())}); + const normalized = {{ + suite: 'preview-suite', run: 'preview-run', behaviorName: 'answer_accuracy', + configObject: {{ + suite: 'preview-suite', run: 'preview-run', + artifacts_root: 'old-artifacts', results_dir: 'old-results', + behavior: {{ name: 'answer_accuracy', description: 'Answer accurately.' }}, + pipeline: {{}} + }}, + warnings: [], extraFiles: [] + }}; + const original = JSON.stringify(normalized.configObject); + await estimateAssertAiRun(normalized); + const monitorDir = path.resolve(runDirPath(normalized.suite, normalized.run)); + const reservedByEstimate = fs.existsSync(monitorDir); + const written = writeRunConfigFiles(normalized); + const persisted = parse(fs.readFileSync(written.configPath, 'utf-8')); + const spawned = await spawnAssertAiRun(written); + const statePath = path.join(process.env.STATE_ROOT, 'run.json'); + for (let attempt = 0; attempt < 500 && !fs.existsSync(statePath); attempt++) {{ + await new Promise(resolve => setTimeout(resolve, 10)); + }} + if (!fs.existsSync(statePath)) throw new Error('run child did not execute'); + console.log(JSON.stringify({{ + monitorDir, runDir: path.resolve(written.runDir), persisted, reservedByEstimate, + unchanged: original === JSON.stringify(normalized.configObject), + execution: JSON.parse(fs.readFileSync(statePath, 'utf-8')) + }})); + """ + ) + result = subprocess.run( + ["node", *node_ts_args(), "--input-type=module"], + input=script, + text=True, + capture_output=True, + cwd=harness, + env=env, + check=False, + timeout=15, + ) + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + payload = json.loads(result.stdout) + estimate_state = json.loads((root / "estimate.json").read_text(encoding="utf-8")) + self.assertFalse(payload["reservedByEstimate"]) + self.assertTrue(payload["unchanged"]) + expected_run = artifacts_root / "preview-suite" / "preview-run" + self.assertEqual(Path(payload["runDir"]), expected_run) + self.assertEqual(Path(payload["monitorDir"]), expected_run) + for config in (estimate_state["config"], payload["persisted"], payload["execution"]["config"]): + self.assertEqual(Path(config["results_dir"]), artifacts_root) + self.assertEqual(Path(config["artifacts_root"]), artifacts_root.parent) + self.assertTrue(estimate_state["reusedTestSet"]) + self.assertTrue(payload["execution"]["reusedTestSet"]) + self.assertEqual(Path(payload["execution"]["runDir"]), expected_run) + + def test_estimate_uses_temporary_config_without_reserving_run(self) -> None: + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + root = Path(tmp_dir) + harness = root / "harness" + harness.mkdir() + run_spawn_path = harness / "run-spawn.ts" + run_spawn_path.write_text( + RUN_SPAWN_SRC.read_text(encoding="utf-8") + .replace("./artifacts.js", "./artifacts.ts") + .replace("./config.js", "./config.ts"), + encoding="utf-8", + ) + (harness / "artifacts.ts").write_text( + ARTIFACTS_SRC.read_text(encoding="utf-8").replace( + "./config.js", "./config.ts" + ), + encoding="utf-8", + ) + shutil.copyfile(CONFIG_SRC, harness / "config.ts") + + args_path = root / "args.json" + fake_cli = root / "fake-cli.mjs" + fake_cli.write_text( + textwrap.dedent( + """\ + import fs from 'node:fs'; + fs.writeFileSync(process.env.ARGS_PATH, JSON.stringify(process.argv.slice(2))); + console.log(JSON.stringify({ + schema_version: 1, + calls: 2, + input_tokens: 100, + output_tokens: 50, + total_tokens: 150, + lower_bound_tokens: 98, + upper_bound_tokens: 203, + stages: {}, + notes: [] + })); + """ + ), + encoding="utf-8", + ) + + artifacts_root = root / "artifacts" / "results" + env = os.environ.copy() + env.update( + { + "ARGS_PATH": str(args_path), + "ARTIFACTS_ROOT": str(artifacts_root), + "MEASUREMENTS_ROOT": str(ROOT), + "ASSERT_AI_COMMAND": f"node {fake_cli}", + "TMP": str(root), + "TEMP": str(root), + "TMPDIR": str(root), + } + ) + script = textwrap.dedent( + f"""\ + const {{ estimateAssertAiRun }} = await import({json.dumps(run_spawn_path.as_uri())}); + const estimate = await estimateAssertAiRun({{ + suite: 'preview-suite', + run: 'preview-run', + behaviorName: 'answer_accuracy', + configObject: {{ + suite: 'preview-suite', + run: 'preview-run', + behavior: {{ name: 'answer_accuracy', description: 'Answer accurately.' }}, + context: 'A factual assistant.', + pipeline: {{}} + }}, + warnings: [], + extraFiles: [] + }}); + console.log(JSON.stringify(estimate)); + """ + ) + result = subprocess.run( + ["node", *node_ts_args(), "--input-type=module"], + input=script, + text=True, + capture_output=True, + cwd=harness, + env=env, + check=False, + ) + + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + self.assertEqual(json.loads(result.stdout)["total_tokens"], 150) + args = json.loads(args_path.read_text(encoding="utf-8")) + self.assertEqual(args[0], "estimate") + self.assertEqual(args[-2:], ["--output", "json"]) + config_path = Path(args[args.index("--config") + 1]) + self.assertFalse(config_path.exists()) + self.assertFalse((artifacts_root / "preview-suite").exists()) + + def test_aborted_estimate_waits_for_child_close_before_cleanup(self) -> None: + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + root = Path(tmp_dir) + harness = root / "harness" + harness.mkdir() + run_spawn_path = harness / "run-spawn.ts" + run_spawn_path.write_text( + RUN_SPAWN_SRC.read_text(encoding="utf-8") + .replace("./artifacts.js", "./artifacts.ts") + .replace("./config.js", "./config.ts"), + encoding="utf-8", + ) + (harness / "artifacts.ts").write_text( + ARTIFACTS_SRC.read_text(encoding="utf-8").replace( + "./config.js", "./config.ts" + ), + encoding="utf-8", + ) + shutil.copyfile(CONFIG_SRC, harness / "config.ts") + + state_path = root / "state.json" + fake_cli = root / "fake-cli.mjs" + fake_cli.write_text( + textwrap.dedent( + """\ + import fs from 'node:fs'; + const args = process.argv.slice(2); + const configPath = args[args.indexOf('--config') + 1]; + fs.writeFileSync( + process.env.STATE_PATH, + JSON.stringify({ pid: process.pid, configPath }) + ); + process.on('SIGTERM', () => {}); + setInterval(() => {}, 1000); + """ + ), + encoding="utf-8", + ) + + env = os.environ.copy() + env.update( + { + "STATE_PATH": str(state_path), + "ARTIFACTS_ROOT": str(root / "artifacts" / "results"), + "MEASUREMENTS_ROOT": str(ROOT), + "ASSERT_AI_COMMAND": f"node {fake_cli}", + "TMP": str(root), + "TEMP": str(root), + "TMPDIR": str(root), + } + ) + script = textwrap.dedent( + f"""\ + import fs from 'node:fs'; + const {{ estimateAssertAiRun }} = await import({json.dumps(run_spawn_path.as_uri())}); + const controller = new AbortController(); + const pending = estimateAssertAiRun({{ + suite: 'preview-suite', + run: 'preview-run', + behaviorName: 'answer_accuracy', + configObject: {{ + suite: 'preview-suite', + run: 'preview-run', + behavior: {{ name: 'answer_accuracy', description: 'Answer accurately.' }}, + context: 'A factual assistant.', + pipeline: {{}} + }}, + warnings: [], + extraFiles: [] + }}, controller.signal); + for (let attempt = 0; attempt < 200 && !fs.existsSync(process.env.STATE_PATH); attempt++) {{ + await new Promise((resolve) => setTimeout(resolve, 10)); + }} + if (!fs.existsSync(process.env.STATE_PATH)) {{ + throw new Error('estimate child did not start'); + }} + const state = JSON.parse(fs.readFileSync(process.env.STATE_PATH, 'utf-8')); + const abortStartedAt = Date.now(); + controller.abort(); + let error = ''; + try {{ + await pending; + }} catch (err) {{ + error = err.message; + }} + let childAlive = true; + try {{ + process.kill(state.pid, 0); + }} catch {{ + childAlive = false; + }} + console.log(JSON.stringify({{ + error, + childAlive, + configExists: fs.existsSync(state.configPath), + abortElapsedMs: Date.now() - abortStartedAt + }})); + """ + ) + result = subprocess.run( + ["node", *node_ts_args(), "--input-type=module"], + input=script, + text=True, + capture_output=True, + cwd=harness, + env=env, + check=False, + timeout=15, + ) + + self.assertEqual( + result.returncode, + 0, + msg=f"{result.stdout}\n{result.stderr}", + ) + payload = json.loads(result.stdout) + self.assertIn("cancelled", payload["error"]) + self.assertFalse(payload["childAlive"]) + self.assertFalse(payload["configExists"]) + self.assertLess(payload["abortElapsedMs"], 5_000) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_viewer_server_artifacts.py b/tests/test_viewer_server_artifacts.py index 95dd25461..10094b918 100644 --- a/tests/test_viewer_server_artifacts.py +++ b/tests/test_viewer_server_artifacts.py @@ -1308,6 +1308,101 @@ def test_load_run_page_data_skips_preview_once_scores_exist(self) -> None: self.assertEqual(payload["auditScoreCount"], 1) self.assertEqual(payload["turnsCount"], 0) + def test_load_run_page_data_exposes_token_usage_metrics(self) -> None: + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + tmp_root = Path(tmp_dir) + harness_dir = tmp_root / "harness" + harness_dir.mkdir() + data_path = self._copy_data_harness(harness_dir) + + artifacts_root = tmp_root / "artifacts" / "results" + run_dir = artifacts_root / "suite-a" / "run-a" + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "manifest.json").write_text( + json.dumps( + { + "status": "running", + "stages": {"inference": "completed", "judge": "running"}, + } + ), + encoding="utf-8", + ) + (run_dir / "config.yaml").write_text( + "pipeline:\n inference:\n target:\n model:\n name: target-model\n", + encoding="utf-8", + ) + (run_dir / "metrics.json").write_text( + json.dumps( + { + "schema_version": 1, + "totals": { + "requests": 2, + "calls": 2, + "missing_usage_calls": 0, + "input_tokens": 800, + "output_tokens": 200, + "total_tokens": 1000, + "cached_input_tokens": 100, + "cache_creation_input_tokens": 0, + "cache_hit_rate": 0.125, + "usage_coverage": 1.0, + }, + "token_estimate": { + "calls": 2, + "input_tokens": 900, + "output_tokens": 200, + "total_tokens": 1100, + "lower_bound_tokens": 770, + "upper_bound_tokens": 1430, + "stages": { + "judge": { + "calls": 1, + "input_tokens": 700, + "output_tokens": 200, + "total_tokens": 900, + } + }, + "notes": ["Retries are not included."], + }, + "token_estimate_accuracy": { + "actual_total_tokens": 1000, + "estimated_total_tokens": 1100, + "difference_tokens": -100, + "difference_ratio": -100 / 1100, + "absolute_percentage_error": 100 / 1100, + }, + } + ), + encoding="utf-8", + ) + + env = os.environ.copy() + env.update( + { + "ARTIFACTS_ROOT": str(artifacts_root), + "MEASUREMENTS_ROOT": str(tmp_root), + } + ) + script = textwrap.dedent( + f"""\ + const {{ loadRunPageData }} = await import({json.dumps(data_path.as_uri())}); + const payload = loadRunPageData('suite-a', 'run-a'); + console.log(JSON.stringify(payload.tokenUsage)); + """ + ) + result = self._run_node(harness_dir=harness_dir, script=script, env=env) + + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + payload = json.loads(result.stdout) + self.assertEqual(payload["estimate"]["totalTokens"], 1100) + self.assertEqual(payload["estimate"]["lowerBoundTokens"], 770) + self.assertEqual(payload["estimate"]["stages"]["judge"]["totalTokens"], 900) + self.assertEqual(payload["estimate"]["notes"], ["Retries are not included."]) + self.assertEqual(payload["actual"]["totalTokens"], 1000) + self.assertEqual(payload["actual"]["usageCoverage"], 1) + self.assertEqual(payload["accuracy"]["status"], "available") + self.assertAlmostEqual(payload["accuracy"]["differenceRatio"], -100 / 1100) + def test_completed_run_page_data_preserves_pipeline_manifest_fields(self) -> None: with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: tmp_root = Path(tmp_dir) diff --git a/tests/test_viewer_token_usage.py b/tests/test_viewer_token_usage.py new file mode 100644 index 000000000..02f9c130f --- /dev/null +++ b/tests/test_viewer_token_usage.py @@ -0,0 +1,385 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import json +import os +import subprocess +import textwrap +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +from tests.node_runner import node_supports_ts, node_ts_args +from tests import test_viewer_server_artifacts as server_artifacts + + +ROOT = Path(__file__).resolve().parents[1] +TOKEN_USAGE_SRC = ROOT / "viewer" / "src" / "lib" / "token-usage.ts" +RUN_PAGE_SRC = ( + ROOT + / "viewer" + / "src" + / "routes" + / "suite" + / "[suite_id]" + / "[run_id]" + / "+page.svelte" +) +EXPORT_PAGE_SRC = ROOT / "viewer" / "src" / "lib" / "export" / "ExportPage.svelte" +NEW_PAGE_SRC = ROOT / "viewer" / "src" / "routes" / "new" / "+page.svelte" +TOKEN_SUMMARY_SRC = ( + ROOT / "viewer" / "src" / "lib" / "components" / "TokenUsageSummary.svelte" +) +TOKEN_PREVIEW_SRC = ( + ROOT / "viewer" / "src" / "lib" / "components" / "TokenEstimatePreview.svelte" +) +ESTIMATE_ROUTE_SRC = ( + ROOT + / "viewer" + / "src" + / "routes" + / "api" + / "runs" + / "estimate" + / "+server.ts" +) + + +class ViewerTokenUsageWiringTest(unittest.TestCase): + def test_summary_is_wired_into_run_and_export_views(self) -> None: + for path in (RUN_PAGE_SRC, EXPORT_PAGE_SRC): + source = path.read_text(encoding="utf-8") + self.assertIn("TokenUsageSummary", source) + self.assertIn("tokenUsage={data.tokenUsage}", source) + + def test_wizard_shows_estimate_before_submit(self) -> None: + page_source = NEW_PAGE_SRC.read_text(encoding="utf-8") + route_source = ESTIMATE_ROUTE_SRC.read_text(encoding="utf-8") + + self.assertIn("TokenEstimatePreview estimate={tokenEstimate}", page_source) + self.assertIn("Estimated token usage", TOKEN_PREVIEW_SRC.read_text(encoding="utf-8")) + self.assertIn("fetch('/api/runs/estimate'", page_source) + self.assertIn("estimateAssertAiRun", route_source) + self.assertIn("request.signal", route_source) + self.assertIn("No provider calls are", route_source) + + def test_completed_token_summary_is_compact(self) -> None: + source = TOKEN_SUMMARY_SRC.read_text(encoding="utf-8") + + self.assertIn(" 0", source) + self.assertIn("'Reported' : 'Actual'", source) + self.assertGreaterEqual( + source.count("tokenAccuracyUnavailableMessage"), + 3, + ) + self.assertNotIn("md:grid-cols-3", source) + self.assertNotIn("text-2xl", source) + + +@unittest.skipUnless(node_supports_ts(), "node binary lacks TypeScript support (need >= 22.6)") +class ViewerTokenUsageFormattingTest(unittest.TestCase): + def test_zero_note_only_and_legacy_metrics_normalization(self) -> None: + helper = server_artifacts.ViewerServerArtifactsTest() + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + root = Path(tmp_dir) + harness = root / "harness" + harness.mkdir() + data_path = helper._copy_data_harness(harness) + artifacts_root = root / "results" + run_dir = artifacts_root / "suite-a" / "run-a" + run_dir.mkdir(parents=True) + (run_dir / "manifest.json").write_text( + '{"status":"completed","stages":{"judge":"completed"}}', encoding="utf-8" + ) + (run_dir / "config.yaml").write_text("pipeline: {}\n", encoding="utf-8") + for file_name in ("inference_set.jsonl", "scores.jsonl"): + (run_dir / file_name).write_text("", encoding="utf-8") + helper._build_viewer_read_model(run_dir) + caveat = "Opaque callable target usage is excluded." + cases = { + "zero": {"token_estimate": { + "calls": 0, "input_tokens": 0, "output_tokens": 0, "total_tokens": 0, + "lower_bound_tokens": 0, "upper_bound_tokens": 0, + "stages": {"inference": {"calls": 0, "total_tokens": 0}}, + "notes": [caveat], + }}, + "note_only": {"token_estimate": {"notes": [None, "", " ", f" {caveat} "]}}, + "explicit_zero": {"token_estimate": {"total_tokens": 0}}, + "empty": {"token_estimate": {}}, + "invalid": {"token_estimate": { + "calls": -0.5, "total_tokens": "0", "notes": [False, {}, " "], + }}, + "legacy": {"totals": {"calls": 2, "input_tokens": 10, "output_tokens": 5}}, + "legacy_empty": {"totals": {"calls": 0, "input_tokens": 0}}, + "fallback_total": {"token_estimate": { + "input_tokens": 8, "output_tokens": 2, "stages": {"invalid": []}, + }}, + "asymmetric": {"token_estimate": { + "calls": 1, "total_tokens": 100, + "lower_bound_tokens": 65, "upper_bound_tokens": 4000, + "notes": ["Upper bound includes the full max_tool_calls cap."], + }}, + } + env = os.environ.copy() + env.update({"ARTIFACTS_ROOT": str(artifacts_root), "MEASUREMENTS_ROOT": str(root)}) + script = textwrap.dedent( + f"""\ + import fs from 'node:fs'; + const {{ loadRunPageData }} = await import({json.dumps(data_path.as_uri())}); + const {{ loadViewerRunReadModel }} = await import({json.dumps((harness / 'artifacts.ts').as_uri())}); + const cases = {json.dumps(cases)}; + const result = {{}}; + for (const [name, metrics] of Object.entries(cases)) {{ + fs.writeFileSync({json.dumps(str(run_dir / 'metrics.json'))}, JSON.stringify(metrics)); + loadViewerRunReadModel('suite-a', 'run-a'); + result[name] = loadRunPageData('suite-a', 'run-a').tokenUsage; + }} + console.log(JSON.stringify(result)); + """ + ) + result = helper._run_node(harness_dir=harness, script=script, env=env) + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + payload = json.loads(result.stdout) + for name in ("zero", "note_only", "explicit_zero"): + with self.subTest(case=name): + self.assertIsNotNone(payload[name]) + self.assertEqual(payload[name]["estimate"]["totalTokens"], 0) + self.assertEqual(payload[name]["estimate"]["lowerBoundTokens"], 0) + self.assertEqual(payload[name]["estimate"]["upperBoundTokens"], 0) + self.assertEqual(payload["zero"]["estimate"]["notes"], [caveat]) + self.assertEqual(payload["note_only"]["estimate"]["notes"], [caveat]) + self.assertEqual(payload["zero"]["estimate"]["stages"]["inference"]["totalTokens"], 0) + for name in ("empty", "invalid", "legacy_empty"): + self.assertIsNone(payload[name], name) + self.assertIsNone(payload["legacy"]["estimate"]) + self.assertEqual(payload["legacy"]["actual"]["totalTokens"], 15) + self.assertEqual(payload["fallback_total"]["estimate"]["totalTokens"], 10) + self.assertEqual(payload["asymmetric"]["estimate"]["totalTokens"], 100) + self.assertEqual(payload["asymmetric"]["estimate"]["lowerBoundTokens"], 65) + self.assertEqual(payload["asymmetric"]["estimate"]["upperBoundTokens"], 4000) + self.assertEqual( + payload["asymmetric"]["estimate"]["notes"], + ["Upper bound includes the full max_tool_calls cap."], + ) + + def test_completed_summary_keeps_exclusion_visible_when_details_are_closed(self) -> None: + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + compiled_path = Path(tmp_dir) / "summary.mjs" + script = textwrap.dedent( + f"""\ + import fs from 'node:fs'; + import {{ compile }} from 'svelte/compiler'; + import {{ render }} from 'svelte/server'; + const source = fs.readFileSync({json.dumps(str(TOKEN_SUMMARY_SRC))}, 'utf-8') + .replace('$lib/token-usage.js', {json.dumps(TOKEN_USAGE_SRC.as_uri())}); + const compiled = compile(source, {{ generate: 'server', filename: 'TokenUsageSummary.svelte' }}); + fs.writeFileSync({json.dumps(str(compiled_path))}, compiled.js.code); + const {{ default: Summary }} = await import({json.dumps(compiled_path.as_uri())}); + const result = {{}}; + for (const total of [0, 100]) {{ + result[total] = render(Summary, {{ + props: {{ tokenUsage: {{ + estimate: {{ + calls: 0, inputTokens: total, outputTokens: 0, totalTokens: total, + lowerBoundTokens: total * 0.65, upperBoundTokens: total * 40, stages: {{}}, + notes: [ + 'Opaque callable target usage is excluded.', 'A local heuristic estimate.', + 'Upper bound includes the full max_tool_calls cap.' + ] + }}, + actual: total ? {{ + requests: 1, calls: 1, missingUsageCalls: 0, + inputTokens: 2900, outputTokens: 100, totalTokens: 3000, + cachedInputTokens: 0, cacheCreationInputTokens: 0, + cacheHitRate: 0, usageCoverage: 1 + }} : null, + accuracy: total ? {{ + status: 'available', actualTotalTokens: 3000, estimatedTotalTokens: 100, + differenceTokens: 2900, differenceRatio: 29, absolutePercentageError: 29 + }} : null + }} }} + }}).body; + }} + console.log(JSON.stringify(result)); + """ + ) + result = subprocess.run( + ["node", *node_ts_args(), "--input-type=module"], + input=script, text=True, encoding="utf-8", capture_output=True, + cwd=ROOT / "viewer", check=False, + ) + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + for total, html in json.loads(result.stdout).items(): + with self.subTest(total=total): + visible = html.split(" None: + with TemporaryDirectory(dir=ROOT / "viewer") as tmp_dir: + compiled_path = Path(tmp_dir) / "preview.mjs" + script = textwrap.dedent( + f"""\ + import fs from 'node:fs'; + import {{ compile }} from 'svelte/compiler'; + import {{ render }} from 'svelte/server'; + const source = fs.readFileSync({json.dumps(str(TOKEN_PREVIEW_SRC))}, 'utf-8') + .replace('$lib/token-usage.js', {json.dumps(TOKEN_USAGE_SRC.as_uri())}); + const compiled = compile(source, {{ generate: 'server', filename: 'TokenEstimatePreview.svelte' }}); + fs.writeFileSync({json.dumps(str(compiled_path))}, compiled.js.code); + const {{ default: Preview }} = await import({json.dumps(compiled_path.as_uri())}); + const result = {{}}; + for (const target of ['callable', 'endpoint', 'sandbox']) {{ + for (const total of [0, 100]) {{ + result[`${{target}}-${{total}}`] = render(Preview, {{ + props: {{ estimate: {{ + calls: total ? 1 : 0, input_tokens: total, output_tokens: 0, total_tokens: total, + lower_bound_tokens: total * 0.65, upper_bound_tokens: total * 40, + notes: [ + `Opaque ${{target}} target usage is excluded.`, 'A local heuristic estimate.', + 'Upper bound includes the full max_tool_calls cap.' + ] + }} }} + }}).body; + }} + }} + result.legacy = render(Preview, {{ props: {{ estimate: {{ + calls: 0, input_tokens: 0, output_tokens: 0, total_tokens: 0, + lower_bound_tokens: 0, upper_bound_tokens: 0 + }} }} }}).body; + result.noop = render(Preview, {{ props: {{ estimate: {{ + calls: 0, input_tokens: 0, output_tokens: 0, total_tokens: 0, + lower_bound_tokens: 0, upper_bound_tokens: 0, + notes: ['Retries and provider-side hidden overhead are not included.'] + }} }} }}).body; + result.loading = render(Preview, {{ props: {{ estimate: null, loading: true }} }}).body; + result.error = render(Preview, {{ props: {{ estimate: null, error: 'Unavailable' }} }}).body; + console.log(JSON.stringify(result)); + """ + ) + result = subprocess.run( + ["node", *node_ts_args(), "--input-type=module"], + input=script, text=True, encoding="utf-8", capture_output=True, + cwd=ROOT / "viewer", check=False, + ) + self.assertEqual(result.returncode, 0, msg=f"{result.stdout}\n{result.stderr}") + payload = json.loads(result.stdout) + for target in ("callable", "endpoint", "sandbox"): + for total in (0, 100): + with self.subTest(target=target, total=total): + html = payload[f"{target}-{total}"] + visible = html.split("