From 194ff2d28d626a8a300917d0cd0e237053e28441 Mon Sep 17 00:00:00 2001 From: David Hyrule Date: Mon, 29 Jun 2026 04:32:09 +0200 Subject: [PATCH] feat: surface OpenRouter token usage in knowledge enrich CostUsage (Phase 2b) call_openrouter now returns (content, usage) and requests usage.include so OpenRouter returns cost; enrich_target captures usage and emits agent-core CostUsage (input/output/total tokens + usd) via the flag-gated emitter. Emission moved from cmd_enrich into enrich_target so both the CLI and the knowledge-loop batch path emit, and only on the live (non-dry-run) path. Still off by default (HYRULE_KNOWLEDGE_AGENT_CORE_TRACE). ruff + mypy --strict clean; suite 97 passed. Co-Authored-By: Claude Opus 4.8 --- src/hyrule_knowledge/agent_core_trace.py | 19 +++++++++++++++---- src/hyrule_knowledge/cli.py | 4 +--- src/hyrule_knowledge/enrich.py | 9 ++++++++- src/hyrule_knowledge/llm.py | 8 ++++++-- tests/test_agent_core_trace.py | 10 +++++++++- 5 files changed, 39 insertions(+), 11 deletions(-) diff --git a/src/hyrule_knowledge/agent_core_trace.py b/src/hyrule_knowledge/agent_core_trace.py index 64babb9..29374ee 100644 --- a/src/hyrule_knowledge/agent_core_trace.py +++ b/src/hyrule_knowledge/agent_core_trace.py @@ -57,18 +57,29 @@ def emit_context_pack(pack_json: dict[str, Any], *, run_id: str | None = None) - def emit_enrich_cost( - provider: str, model: str, target: str, *, run_id: str | None = None + provider: str, + model: str, + target: str, + *, + usage: dict[str, Any] | None = None, + run_id: str | None = None, ) -> dict[str, Any] | None: - """Emit a model_call TraceEvent carrying CostUsage(provider, model) for an enrichment. + """Emit a model_call TraceEvent carrying CostUsage for an enrichment. - Token/USD fidelity is a follow-up: the enrichment path does not yet surface usage. + OpenRouter ``usage`` (when present) maps prompt/completion/total tokens and ``cost`` -> usd. """ if not enabled(): return None try: models_mod = importlib.import_module("agent_core.contracts.models") tracing_mod = importlib.import_module("agent_core.contracts.tracing") - cost = models_mod.CostUsage(provider=provider, model=model) + cost_fields: dict[str, Any] = {"provider": provider, "model": model} + if usage: + cost_fields["input_tokens"] = usage.get("prompt_tokens") + cost_fields["output_tokens"] = usage.get("completion_tokens") + cost_fields["total_tokens"] = usage.get("total_tokens") + cost_fields["usd"] = usage.get("cost") + cost = models_mod.CostUsage(**cost_fields) event = tracing_mod.TraceEvent( event_type="model_call", node_id="enrich", diff --git a/src/hyrule_knowledge/cli.py b/src/hyrule_knowledge/cli.py index 6864a1c..100102f 100644 --- a/src/hyrule_knowledge/cli.py +++ b/src/hyrule_knowledge/cli.py @@ -9,7 +9,7 @@ from datetime import UTC, datetime from pathlib import Path -from .agent_core_trace import emit_context_pack, emit_enrich_cost +from .agent_core_trace import emit_context_pack from .authority import AuthorityTier from .builder import build_all, source_ref from .config import SourceConfig, load_config @@ -406,8 +406,6 @@ def cmd_enrich(args: argparse.Namespace) -> int: print(f"enrichment failed: {exc}", file=sys.stderr) return 1 print(f"wrote enrichment concept: {path}") - if not args.dry_run: - emit_enrich_cost(args.provider, args.model, args.target) return 0 diff --git a/src/hyrule_knowledge/enrich.py b/src/hyrule_knowledge/enrich.py index 6403b91..c456583 100644 --- a/src/hyrule_knowledge/enrich.py +++ b/src/hyrule_knowledge/enrich.py @@ -5,9 +5,11 @@ import json from datetime import UTC, datetime from pathlib import Path +from typing import Any import yaml +from .agent_core_trace import emit_enrich_cost from .llm import LLMConfig, LLMError, call_openrouter from .models import Concept, EnrichmentMetadata, SourceRef from .okf_writer import dump_concept, slugify @@ -25,12 +27,15 @@ def enrich_target( ) -> Path: source_pack = build_source_pack(bundle_root, target, max_input_chars) generated_at = datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z") + usage: dict[str, Any] | None = None if dry_run: output = _dry_run_output(source_pack) else: if provider != "openrouter": raise LLMError(f"unsupported provider for this tranche: {provider}") - output = call_openrouter(source_pack, LLMConfig(provider=provider, model=model, temperature=temperature)) + output, usage = call_openrouter( + source_pack, LLMConfig(provider=provider, model=model, temperature=temperature) + ) body = _body_from_output(output) output_hash = sha256_text(json.dumps(output, sort_keys=True, ensure_ascii=False)) source_refs = _source_refs_from_pack(source_pack) @@ -60,6 +65,8 @@ def enrich_target( path.parent.mkdir(parents=True, exist_ok=True) path.write_text(dump_concept(concept), encoding="utf-8") _write_run_record(bundle_root.parent / "exports/enrichment-runs.jsonl", concept, source_pack, output) + if not dry_run: + emit_enrich_cost(provider, model, target, usage=usage) return path diff --git a/src/hyrule_knowledge/llm.py b/src/hyrule_knowledge/llm.py index 36c0944..e751ccd 100644 --- a/src/hyrule_knowledge/llm.py +++ b/src/hyrule_knowledge/llm.py @@ -47,7 +47,9 @@ def build_user_prompt(source_pack: dict[str, Any]) -> str: ) -def call_openrouter(source_pack: dict[str, Any], config: LLMConfig) -> dict[str, Any]: +def call_openrouter( + source_pack: dict[str, Any], config: LLMConfig +) -> tuple[dict[str, Any], dict[str, Any] | None]: api_key = os.environ.get(config.api_key_env) if not api_key: raise LLMError(f"missing {config.api_key_env}; LLM enrichment is manual/opt-in") @@ -59,6 +61,7 @@ def call_openrouter(source_pack: dict[str, Any], config: LLMConfig) -> dict[str, {"role": "user", "content": build_user_prompt(source_pack)}, ], "response_format": {"type": "json_object"}, + "usage": {"include": True}, } request = urllib.request.Request( "https://openrouter.ai/api/v1/chat/completions", @@ -83,7 +86,8 @@ def call_openrouter(source_pack: dict[str, Any], config: LLMConfig) -> dict[str, content = data["choices"][0]["message"]["content"] result = parse_json_object_response(content) validate_enrichment_json(result, _allowed_citations(source_pack)) - return result + usage: dict[str, Any] | None = data.get("usage") if isinstance(data, dict) else None + return result, usage def parse_json_object_response(content: Any) -> dict[str, Any]: diff --git a/tests/test_agent_core_trace.py b/tests/test_agent_core_trace.py index c34635a..a4d55de 100644 --- a/tests/test_agent_core_trace.py +++ b/tests/test_agent_core_trace.py @@ -41,7 +41,15 @@ def test_enrich_cost_emits_when_enabled(monkeypatch, tmp_path): sink = tmp_path / "t.jsonl" monkeypatch.setenv("HYRULE_KNOWLEDGE_AGENT_CORE_TRACE", "1") monkeypatch.setenv("HYRULE_KNOWLEDGE_AGENT_CORE_TRACE_PATH", str(sink)) - record = agent_core_trace.emit_enrich_cost("openrouter", "anthropic/claude-sonnet-4.6", "ansible") + record = agent_core_trace.emit_enrich_cost( + "openrouter", + "anthropic/claude-sonnet-4.6", + "ansible", + usage={"prompt_tokens": 1500, "completion_tokens": 300, "total_tokens": 1800, "cost": 0.012}, + ) assert record is not None assert record["event_type"] == "model_call" assert record["cost"]["provider"] == "openrouter" + assert record["cost"]["input_tokens"] == 1500 + assert record["cost"]["output_tokens"] == 300 + assert record["cost"]["usd"] == 0.012