diff --git a/.gitignore b/.gitignore index 09d9a25e..b9757b08 100644 --- a/.gitignore +++ b/.gitignore @@ -91,6 +91,8 @@ Thumbs.db !tests/test_updater_write_guard.py !tests/test_updater_orchestrator.py !tests/test_usage_ledger.py +!tests/test_pricing.py +!tests/smoke_test_costs.py # Mermaid regression tests and development-only parser harness !tests/test_mermaid_validation.py diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py index 4e5fcfa3..f9e318be 100644 --- a/codewiki/cli/adapters/doc_generator.py +++ b/codewiki/cli/adapters/doc_generator.py @@ -163,6 +163,7 @@ def generate(self) -> DocumentationJob: agent_instructions=self.config.get("agent_instructions"), use_gitignore=self.config.get("use_gitignore", True), prompt_caching=self.config.get("prompt_caching", True), + pricing=self.config.get("pricing", {}), artifacts_enabled=self.config.get("artifacts_enabled", True), artifact_token_budget=self.config.get("artifact_token_budget", 200_000), with_prose=self.config.get("with_prose", False), @@ -181,15 +182,18 @@ def generate(self) -> DocumentationJob: self._finalize_job() # Complete job + self._finish_usage() self.job.complete() return self.job except APIError as e: self.job.fail(str(e)) + self._finish_usage(incomplete=True) raise except Exception as e: self.job.fail(str(e)) + self._finish_usage(incomplete=True) raise async def _run_backend_generation(self, backend_config: BackendConfig): @@ -448,6 +452,7 @@ async def _run_incremental_update( updater = IncrementalUpdater( backend_config, doc_generator.backend, doc_generator, update_opts ) + self._last_update_record = updater.record with doc_generator.usage_ledger.in_phase("update"): record = await updater.run(prev_graph_path, components, leaf_nodes, revision) self._last_update_record = record @@ -484,6 +489,8 @@ def _move_docs_aside(self) -> None: """Keep the previous docs (and the failed attempt's record) next to the output dir.""" import shutil + self._prior_update_history = self._read_update_history(str(self.output_dir)) + suffix = (self.commit_id or "unknown")[:8] target = self.output_dir.parent / f"{self.output_dir.name}.prev-{suffix}" n = 2 @@ -515,11 +522,17 @@ def _merge_update_summary(self, working_dir: str, prior_history: list | None = N return from codewiki.src.be.updater.record import merge_into_metadata - summary = record.summary() preserved = getattr(self, "_preserved_docs_dir", None) + if self._usage_ledger is not None: + record.usage_report = self._usage_ledger.summary() + summary = record.summary() if preserved: summary["previous_docs"] = preserved - merge_into_metadata(working_dir, summary, prior_history=prior_history) + merge_into_metadata( + working_dir, + summary, + prior_history=prior_history or getattr(self, "_prior_update_history", None), + ) try: record.save(working_dir) except OSError: @@ -563,6 +576,7 @@ def _record_usage_statistics(self): summary = self._usage_ledger.summary() totals = summary["totals"] stats = self.job.statistics + stats.usage_report = summary stats.usage_totals = totals stats.total_tokens_used = int( sum( @@ -570,8 +584,6 @@ def _record_usage_statistics(self): for k in ( "input_tokens", "output_tokens", - "cache_read_tokens", - "cache_write_tokens", ) ) or totals.get("total_tokens", 0) @@ -579,6 +591,51 @@ def _record_usage_statistics(self): stats.total_cost_usd = totals.get("cost_usd") stats.unpriced_models = summary["unpriced_models"] + def _finish_usage(self, incomplete: bool = False): + """Persist the final report even when validation or a later stage fails. + + Reporting errors must never hide the original generation failure. + """ + if self._usage_ledger is None: + return + self._usage_ledger.incomplete = incomplete + self._record_usage_statistics() + report = self.job.statistics.usage_report + try: + self.output_dir.mkdir(parents=True, exist_ok=True) + path = self.output_dir / "metadata.json" + try: + metadata = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + metadata = {} + if not isinstance(metadata, dict): + metadata = {} + metadata["usage"] = report + record = getattr(self, "_last_update_record", None) + if record is not None: + record.usage_report = report + record.save(str(self.output_dir)) + summary = record.summary() + if getattr(self, "_preserved_docs_dir", None): + summary["previous_docs"] = self._preserved_docs_dir + history = metadata.get("update_history", []) + if not isinstance(history, list): + history = [] + if history and history[-1].get("run_id") == record.run_id: + history[-1] = summary + else: + history.append(summary) + metadata.update(last_update=summary, update_history=history) + from codewiki.cli.utils.fs import safe_write + + safe_write(path, json.dumps(metadata, indent=2, ensure_ascii=False)) + except Exception: + logging.getLogger(__name__).warning("Could not save the run usage report") + if incomplete: + from codewiki.cli.utils.instructions import _echo_usage + + _echo_usage({"usage_report": report, "provider": self.config.get("provider", "")}) + def _finalize_job(self): """Finalize the job (metadata already created by backend).""" # Just verify metadata exists diff --git a/codewiki/cli/commands/config.py b/codewiki/cli/commands/config.py index e81ea3bd..19d5c92d 100644 --- a/codewiki/cli/commands/config.py +++ b/codewiki/cli/commands/config.py @@ -34,6 +34,110 @@ def config_group(): pass +@config_group.group(name="pricing") +def pricing_group(): + """Manage API prices in USD per million tokens.""" + + +def _pricing_manager(): + manager = ConfigManager() + try: + manager.load() + except ConfigurationError as exc: + raise click.ClickException(str(exc)) from exc + return manager + + +def _pricing_provider(manager, provider): + provider = provider or ( + manager.get_config().provider if manager.get_config() else "openai-compatible" + ) + supported = { + "openai-compatible", + "atlas-cloud", + "anthropic", + "bedrock", + "azure-openai", + "google", + "google-vertex", + } + if provider not in supported: + raise click.BadParameter( + "Choose an API provider; Claude Code and Codex use their provider cost results.", + param_hint="--provider", + ) + return provider + + +@pricing_group.command(name="set") +@click.argument("model") +@click.option( + "--input", "input_rate", required=True, type=float, help="USD per million input tokens." +) +@click.option( + "--output", "output_rate", required=True, type=float, help="USD per million output tokens." +) +@click.option("--cache-read", type=float, help="Defaults to the input rate.") +@click.option("--cache-write", type=float, help="Defaults to the input rate.") +@click.option("--provider", help="Defaults to the active API provider.") +def pricing_set(model, input_rate, output_rate, cache_read, cache_write, provider): + """Set rates for MODEL. Explicit zero rates are supported.""" + from copy import deepcopy + from codewiki.src.be.pricing import validate_pricing + + manager = _pricing_manager() + provider = _pricing_provider(manager, provider) + pricing = deepcopy(manager.get_config().pricing) if manager.get_config() else {} + rates = {"input": input_rate, "output": output_rate} + if cache_read is not None: + rates["cache_read"] = cache_read + if cache_write is not None: + rates["cache_write"] = cache_write + pricing.setdefault(provider, {})[model] = rates + try: + validate_pricing(pricing) + manager.save(pricing=pricing) + except (ValueError, ConfigurationError) as exc: + raise click.ClickException(str(exc)) from exc + click.echo(f"API pricing saved for {provider}/{model} (USD per million tokens).") + + +@pricing_group.command(name="show") +@click.option("--json", "output_json", is_flag=True) +def pricing_show(output_json): + """Show configured rates for all API providers.""" + manager = _pricing_manager() + pricing = manager.get_config().pricing if manager.get_config() else {} + if output_json: + click.echo(json.dumps(pricing, indent=2, sort_keys=True)) + elif not pricing: + click.echo("No API prices configured. Use 'codewiki config pricing set --help'.") + else: + click.echo("API pricing (USD per million tokens; omitted cache rates use input):") + for provider, models in sorted(pricing.items()): + for model, rates in sorted(models.items()): + click.echo(f" {provider}/{model}: {json.dumps(rates, sort_keys=True)}") + + +@pricing_group.command(name="remove") +@click.argument("model") +@click.option("--provider", help="Defaults to the active API provider.") +def pricing_remove(model, provider): + """Remove configured rates for MODEL.""" + from copy import deepcopy + + manager = _pricing_manager() + provider = _pricing_provider(manager, provider) + pricing = deepcopy(manager.get_config().pricing) if manager.get_config() else {} + if model not in pricing.get(provider, {}): + raise click.ClickException(f"No price configured for {provider}/{model}") + del pricing[provider][model] + if not pricing[provider]: + del pricing[provider] + manager.save(pricing=pricing) + click.echo(f"API pricing removed for {provider}/{model}.") + + @config_group.command(name="set") @click.option("--api-key", type=str, help="LLM API key (stored securely in system keychain)") @click.option("--base-url", type=str, help="LLM API base URL (e.g., https://api.anthropic.com)") @@ -473,6 +577,7 @@ def config_show(output_json: bool): "agent_retries": config.agent_retries if config else 3, "use_gitignore": config.use_gitignore if config else True, "prompt_caching": config.prompt_caching if config else True, + "pricing": config.pricing if config else {}, "agent_instructions": config.agent_instructions.to_dict() if config and config.agent_instructions else {}, @@ -542,6 +647,9 @@ def config_show(output_json: bool): click.echo(f" Max Token/Module: {config.max_token_per_module}") click.echo(f" Max Token/Leaf Module: {config.max_token_per_leaf_module}") click.echo(f" Prompt Caching: {config.prompt_caching}") + click.echo( + f" API Pricing: {json.dumps(config.pricing, sort_keys=True)}" + ) click.echo() click.secho("Decomposition Settings", fg="cyan", bold=True) diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index 5d8b74c9..b9002da4 100644 --- a/codewiki/cli/commands/generate.py +++ b/codewiki/cli/commands/generate.py @@ -633,10 +633,7 @@ def generate_command( repo_path, output_dir, logger, verbose, compare_to=compare_to ) if changed_files is not None and len(changed_files) == 0: - logger.success( - "No changes detected since last generation. Documentation is up to date." - ) - sys.exit(EXIT_SUCCESS) + logger.info("No committed changes detected; checking documentation state.") if changed_files is not None and update_rung == "0": logger.info( f" Detected {len(changed_files)} changed files — regenerating affected modules." @@ -792,6 +789,7 @@ def generate_command( "base_url": config.base_url, "api_key": api_key, "provider": getattr(config, "provider", "openai-compatible"), + "pricing": getattr(config, "pricing", {}), "aws_region": getattr(config, "aws_region", "us-east-1"), "google_project": getattr(config, "google_project", ""), "google_location": getattr(config, "google_location", ""), @@ -878,6 +876,7 @@ def generate_command( "total_files_analyzed": job.statistics.total_files_analyzed, "generation_time": generation_time, "total_tokens_used": job.statistics.total_tokens_used, + "usage_report": job.statistics.usage_report, "usage_totals": job.statistics.usage_totals, "total_cost_usd": job.statistics.total_cost_usd, "unpriced_models": job.statistics.unpriced_models, diff --git a/codewiki/cli/config_manager.py b/codewiki/cli/config_manager.py index 1881a7b9..59d48625 100644 --- a/codewiki/cli/config_manager.py +++ b/codewiki/cli/config_manager.py @@ -146,6 +146,7 @@ def save( google_location: Optional[str] = None, use_gitignore: Optional[bool] = None, prompt_caching: Optional[bool] = None, + pricing: Optional[dict] = None, ): """ Save configuration to file and keyring. @@ -199,6 +200,14 @@ def save( was_google = is_google_provider(self._config.provider) # Update fields if provided + if pricing is not None: + from codewiki.src.be.pricing import validate_pricing + + try: + validate_pricing(pricing) + except ValueError as exc: + raise ConfigurationError(str(exc)) from exc + self._config.pricing = pricing if base_url is not None: self._config.base_url = base_url if main_model is not None: diff --git a/codewiki/cli/models/config.py b/codewiki/cli/models/config.py index 0d6be21a..6539e4c1 100644 --- a/codewiki/cli/models/config.py +++ b/codewiki/cli/models/config.py @@ -183,8 +183,18 @@ class Configuration: agent_retries: int = 3 use_gitignore: bool = True prompt_caching: bool = True + pricing: dict = field(default_factory=dict) agent_instructions: AgentInstructions = field(default_factory=AgentInstructions) + def __post_init__(self): + from codewiki.src.be.pricing import validate_pricing + from codewiki.cli.utils.errors import ConfigurationError + + try: + validate_pricing(self.pricing) + except ValueError as exc: + raise ConfigurationError(str(exc)) from exc + def validate(self): """ Validate all configuration fields. @@ -227,6 +237,7 @@ def to_dict(self) -> dict: "agent_retries": self.agent_retries, "use_gitignore": self.use_gitignore, "prompt_caching": self.prompt_caching, + "pricing": self.pricing, "fallback_model": self.fallback_model, } if self.agent_instructions and not self.agent_instructions.is_empty(): @@ -268,6 +279,7 @@ def from_dict(cls, data: dict) -> "Configuration": agent_retries=data.get("agent_retries", 3), use_gitignore=data.get("use_gitignore", True), prompt_caching=data.get("prompt_caching", True), + pricing=data.get("pricing", {}), agent_instructions=agent_instructions, ) @@ -338,4 +350,5 @@ def to_backend_config( agent_instructions=final_instructions.to_dict() if final_instructions else None, use_gitignore=self.use_gitignore, prompt_caching=self.prompt_caching, + pricing=self.pricing, ) diff --git a/codewiki/cli/models/job.py b/codewiki/cli/models/job.py index 11db6e8e..5bbeb083 100644 --- a/codewiki/cli/models/job.py +++ b/codewiki/cli/models/job.py @@ -42,6 +42,7 @@ class JobStatistics: total_cost_usd: Optional[float] = None # Models whose calls reported tokens but no price unpriced_models: List[str] = field(default_factory=list) + usage_report: Dict[str, Any] = field(default_factory=dict) @dataclass diff --git a/codewiki/cli/utils/instructions.py b/codewiki/cli/utils/instructions.py index d292d7fe..c9bacc98 100644 --- a/codewiki/cli/utils/instructions.py +++ b/codewiki/cli/utils/instructions.py @@ -55,6 +55,10 @@ def _compact(n: float) -> str: def _echo_usage(statistics: dict) -> None: """Print the run's token usage and estimated cost, if any was recorded.""" + report = statistics.get("usage_report") + if report: + _echo_usage_report(report, statistics.get("provider", "")) + return tokens = statistics.get("total_tokens_used") or 0 if not tokens: return @@ -91,6 +95,69 @@ def _echo_usage(statistics: dict) -> None: ) +def _usd(cost: float) -> str: + if cost == 0 or cost >= 0.01: + return f"${cost:,.2f}" + if cost < 0.00000001: + return "<$0.00000001" + return f"${cost:.8f}".rstrip("0") + + +def _echo_usage_report(report: dict, provider: str) -> None: + totals = report["totals"] + tokens = int(totals.get("input_tokens", 0) + totals.get("output_tokens", 0)) + click.echo(f" LLM calls: {report['n_calls']:,} (this run)") + click.echo( + f" Tokens used: {tokens:,} (in {_compact(totals.get('input_tokens', 0))}" + f" / out {_compact(totals.get('output_tokens', 0))})" + ) + status = report["cost_status"] + note = ( + " (API-equivalent; billed to your subscription)" + if provider in ("claude-code", "codex") + else "" + ) + if status == "unknown": + click.echo(" Est. cost: unknown") + else: + partial = " (partial; unpriced calls excluded)" if status == "partial" else "" + click.echo(f" Est. cost: {_usd(totals.get('cost_usd', 0))}{partial}{note}") + if report.get("incomplete"): + click.secho(" Run incomplete: available usage only.", fg="yellow") + if report.get("unknown_usage_calls"): + click.secho( + f" Usage unavailable for {report['unknown_usage_calls']} call(s).", fg="yellow" + ) + if report.get("unpriced_models"): + click.secho(f" Cost unknown for: {', '.join(report['unpriced_models'])}", fg="yellow") + if provider == "codex": + click.echo( + " set a price: caw pricing set codex --input --output " + ) + elif provider != "claude-code": + click.echo( + " set a price: codewiki config pricing set MODEL --input USD_PER_1M --output USD_PER_1M" + ) + for title, groups in (("By model", report["by_model"]), ("By phase", report["by_phase"])): + if not groups: + continue + width = max(len(title), *(len(name) for name in groups)) + click.echo( + f"\n {title:<{width}} {'Calls':>5} {'Input':>9} {'Output':>9} {'Cache read/write':>19} Est. USD" + ) + for name, group in sorted(groups.items()): + cost = _usd(group["cost_usd"]) if "cost_usd" in group else "unknown" + if group.get("unpriced_calls") == group["calls"]: + cost = "unknown" + elif group.get("unpriced_calls") and "cost_usd" in group: + cost += " (partial)" + cached = f"{int(group.get('cache_read_tokens', 0)):,}/{int(group.get('cache_write_tokens', 0)):,}" + click.echo( + f" {name:<{width}} {group['calls']:>5} {int(group.get('input_tokens', 0)):>9,}" + f" {int(group.get('output_tokens', 0)):>9,} {cached:>19} {cost}" + ) + + def display_post_generation_instructions( output_dir: Path, repo_name: str, diff --git a/codewiki/mcp/server.py b/codewiki/mcp/server.py index e5c1c5ab..8b6afadc 100644 --- a/codewiki/mcp/server.py +++ b/codewiki/mcp/server.py @@ -545,6 +545,7 @@ async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: cluster_model=config.cluster_model, fallback_model=config.fallback_model, provider=getattr(config, "provider", "openai-compatible"), + pricing=getattr(config, "pricing", {}), aws_region=getattr(config, "aws_region", "us-east-1"), google_project=getattr(config, "google_project", ""), google_location=getattr(config, "google_location", ""), diff --git a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py index fabc3131..492ac45b 100644 --- a/codewiki/src/be/agent_tools/generate_sub_module_documentations.py +++ b/codewiki/src/be/agent_tools/generate_sub_module_documentations.py @@ -16,6 +16,7 @@ from codewiki.src.be.cluster_modules import format_potential_core_components import logging +from codewiki.src.be.api_usage import usage_target logger = logging.getLogger(__name__) @@ -124,22 +125,24 @@ async def generate_sub_module_documentation( # log the current module tree # print(f"Current module tree: {json.dumps(deps.module_tree, indent=4)}") - await sub_agent.run( - format_user_prompt( - module_name=deps.current_module_name, - core_component_ids=core_component_ids, - components=ctx.deps.components, - module_tree=ctx.deps.module_tree, - module_path=list(deps.path_to_current_module), - layout=layout, - ), - deps=ctx.deps, - usage_limits=UsageLimits(request_limit=ctx.deps.config.request_limit), - ) - - # remove the sub-module name from the path to current module and the module tree - deps.path_to_current_module.pop() - deps.current_depth -= 1 + try: + with usage_target(sub_module_name): + await sub_agent.run( + format_user_prompt( + module_name=deps.current_module_name, + core_component_ids=core_component_ids, + components=ctx.deps.components, + module_tree=ctx.deps.module_tree, + module_path=list(deps.path_to_current_module), + layout=layout, + ), + deps=ctx.deps, + usage_limits=UsageLimits(request_limit=ctx.deps.config.request_limit), + ) + finally: + deps.path_to_current_module.pop() + deps.current_depth -= 1 + deps.current_module_name = previous_module_name # restore the previous module name deps.current_module_name = previous_module_name diff --git a/codewiki/src/be/api_usage.py b/codewiki/src/be/api_usage.py new file mode 100644 index 00000000..30329337 --- /dev/null +++ b/codewiki/src/be/api_usage.py @@ -0,0 +1,161 @@ +"""Request-level API accounting, shared with nested agents via context variables.""" + +from __future__ import annotations + +import contextlib +import time +from contextvars import ContextVar +from dataclasses import dataclass, field, replace +from typing import Any + +from pydantic_ai.models.wrapper import WrapperModel + +from codewiki.src.be.backend import UsageLedger, usage_to_dict +from codewiki.src.be.pricing import price_usage + + +def _value(obj: Any, name: str, default=None): + return obj.get(name, default) if isinstance(obj, dict) else getattr(obj, name, default) + + +def normalize_api_usage(usage: Any) -> dict[str, Any] | None: + """Normalize OpenAI aliases and cache details to inclusive input counts. + + pydantic-ai and Google's adapter already provide inclusive input. OpenAI + prompt_tokens also includes cache reads. Never sum aliases or cache subsets. + """ + data = usage_to_dict(usage) + if not data or not any( + k in data for k in ("input_tokens", "prompt_tokens", "output_tokens", "completion_tokens") + ): + return None + details = _value(usage, "prompt_tokens_details") + out = { + "input_tokens": data.get("input_tokens", data.get("prompt_tokens", 0)), + "output_tokens": data.get("output_tokens", data.get("completion_tokens", 0)), + "cache_read_tokens": data.get( + "cache_read_tokens", _value(details, "cached_tokens", 0) or 0 + ), + "cache_write_tokens": data.get("cache_write_tokens", 0), + "requests": 1, + } + out["total_tokens"] = out["input_tokens"] + out["output_tokens"] + return out + + +@dataclass +class UsageContext: + ledger: UsageLedger + config: Any + kind: str + target: str + entries: list[dict[str, Any]] = field(default_factory=list) + + def totals(self) -> dict[str, Any] | None: + totals: dict[str, Any] = {} + for entry in self.entries: + for key in ( + "input_tokens", + "output_tokens", + "cache_read_tokens", + "cache_write_tokens", + "total_tokens", + "requests", + "cost_usd", + ): + if key in entry: + totals[key] = totals.get(key, 0) + entry[key] + return totals or None + + +_CONTEXT: ContextVar[UsageContext | None] = ContextVar("codewiki_api_usage", default=None) + + +@contextlib.contextmanager +def track_usage(ledger: UsageLedger, config: Any, kind: str, target: str = ""): + context = UsageContext(ledger, config, kind, target) + token = _CONTEXT.set(context) + try: + yield context + finally: + _CONTEXT.reset(token) + + +@contextlib.contextmanager +def usage_target(target: str): + """Nested agents share accounting but label their own page.""" + parent = _CONTEXT.get() + token = _CONTEXT.set(replace(parent, target=target) if parent else None) + try: + yield + finally: + _CONTEXT.reset(token) + + +def record_response( + usage: Any, + requested_model: str, + *, + model: str | None = None, + seconds: float = 0.0, + error: str | None = None, +) -> dict[str, Any] | None: + context = _CONTEXT.get() + if context is None: + return normalize_api_usage(usage) + normalized = normalize_api_usage(usage) + actual_model = model or requested_model + priced = price_usage( + normalized, + getattr(context.config, "pricing", {}), + context.config.provider, + actual_model, + requested_model, + ) + entry = context.ledger.add( + context.kind, + context.target, + normalized, + model=actual_model, + seconds=seconds, + provider=context.config.provider, + requested_model=requested_model, + **priced, + **({"error": error} if error else {}), + ) + context.entries.append(entry) + return { + **(normalized or {}), + **({"cost_usd": priced["cost_usd"]} if "cost_usd" in priced else {}), + } or None + + +class UsageModel(WrapperModel): + """Observe each model before FallbackModel can combine its results.""" + + def __init__(self, wrapped, requested_model: str): + super().__init__(wrapped) + self.requested_model = requested_model + + async def request(self, messages, model_settings, model_request_parameters): + started = time.monotonic() + try: + response = await self.wrapped.request( + messages, model_settings, model_request_parameters + ) + except Exception as exc: + # Save the type only: provider exceptions can contain credentials and prompts. + record_response( + None, + self.requested_model, + seconds=time.monotonic() - started, + error=type(exc).__name__, + ) + raise + record_response( + response.usage, + self.requested_model, + model=response.model_name, + seconds=time.monotonic() - started, + ) + return response diff --git a/codewiki/src/be/backend.py b/codewiki/src/be/backend.py index d2c20cd0..7a0a8ef9 100644 --- a/codewiki/src/be/backend.py +++ b/codewiki/src/be/backend.py @@ -23,6 +23,7 @@ import contextlib import threading from collections.abc import Iterator +from copy import deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List @@ -46,7 +47,11 @@ def usage_to_dict(usage: Any) -> dict[str, Any] | None: if usage is None: return None if isinstance(usage, dict): - return {k: v for k, v in usage.items() if isinstance(v, (int, float))} or None + return { + k: v + for k, v in usage.items() + if isinstance(v, (int, float)) and not isinstance(v, bool) + } or None out: dict[str, Any] = {} for key in ( "input_tokens", @@ -88,6 +93,7 @@ class UsageLedger: phase: str = "generate" target: str = "" entries: list[dict[str, Any]] = field(default_factory=list) + incomplete: bool = False _lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False) @contextlib.contextmanager @@ -107,7 +113,8 @@ def add( *, seconds: float = 0.0, model: str | None = None, - ) -> None: + **metadata: Any, + ) -> dict[str, Any]: entry: dict[str, Any] = { "phase": self.phase, "kind": kind, @@ -115,37 +122,75 @@ def add( "model": model or "", "seconds": round(seconds, 2), **(usage or {}), + "usage_reported": usage is not None, + **metadata, } with self._lock: self.entries.append(entry) + return entry def summary(self) -> dict[str, Any]: - """Totals overall and per phase, plus the per-call list. - - A call with tokens but no ``cost_usd`` is *unpriced*: its model is - listed under ``unpriced_models`` so a $0 total is never mistaken for - a free run. - """ + """One accounting result, grouped for display without charging twice.""" with self._lock: - entries = [dict(e) for e in self.entries] + entries = deepcopy(self.entries) totals: dict[str, float] = {} by_phase: dict[str, dict[str, float]] = {} + by_model: dict[str, dict[str, Any]] = {} unpriced: set[str] = set() + priced_calls = 0 + unknown_usage = 0 for e in entries: phase_totals = by_phase.setdefault(e["phase"], {"calls": 0}) + model_totals = by_model.setdefault(e["model"] or "", {"calls": 0}) phase_totals["calls"] += 1 + model_totals["calls"] += 1 for k, v in e.items(): - if k == "seconds" or not isinstance(v, (int, float)) or isinstance(v, bool): + if k not in ( + *_TOKEN_KEYS, + "total_tokens", + "cache_read_tokens", + "cache_write_tokens", + "requests", + "cost_usd", + ): + continue + if not isinstance(v, (int, float)) or isinstance(v, bool): continue totals[k] = totals.get(k, 0) + v phase_totals[k] = phase_totals.get(k, 0) + v - if not e.get("cost_usd") and any(e.get(k) for k in _TOKEN_KEYS): + model_totals[k] = model_totals.get(k, 0) + v + # Legacy integrations use zero for "no rate"; explicitly priced zero is valid. + priced = e.get("cost_usd") is not None and ( + e.get("cost_usd") != 0 + or e.get("cost_source") in ("configured", "provider") + or not any(e.get(k) for k in _TOKEN_KEYS) + ) + if priced: + priced_calls += 1 + else: unpriced.add(e["model"] or "") + for group in (phase_totals, model_totals): + group["unpriced_calls"] = group.get("unpriced_calls", 0) + 1 + if not e.get("usage_reported", True): + unknown_usage += 1 + if not entries: + totals["cost_usd"] = 0.0 + status = ( + "complete" + if priced_calls == len(entries) + else ("partial" if priced_calls else "unknown") + ) return { + "schema_version": 2, "scope": "this run only (modules resumed from disk cost nothing)", + "currency": "USD", + "cost_status": status, + "incomplete": self.incomplete, + "unknown_usage_calls": unknown_usage, "n_calls": len(entries), "totals": totals, "by_phase": by_phase, + "by_model": by_model, "unpriced_models": sorted(unpriced), "calls": entries, } @@ -210,8 +255,31 @@ def _record_usage( ) -> dict[str, Any] | None: """Set ``last_usage`` from a provider usage object and log it in the ledger.""" usage_dict = usage_to_dict(usage) + provider = getattr(getattr(self, "_config", None), "provider", "") + if not provider: + provider = {"claude_code": "claude-code"}.get( + getattr(self, "_caw_provider", ""), getattr(self, "_caw_provider", "") + ) + if usage_dict is not None and is_caw_provider(provider): + # caw exposes uncached input for both CLI providers. + usage_dict["input_tokens"] = sum( + usage_dict.get(k, 0) + for k in ("input_tokens", "cache_read_tokens", "cache_write_tokens") + ) + usage_dict["total_tokens"] = usage_dict["input_tokens"] + usage_dict.get( + "output_tokens", 0 + ) self.last_usage = usage_dict - self.usage_ledger.add(kind, target, usage_dict, seconds=seconds, model=model) + source = "provider" if usage_dict and usage_dict.get("cost_usd") else "unknown" + self.usage_ledger.add( + kind, + target, + usage_dict, + seconds=seconds, + model=model, + provider=provider, + cost_source=source, + ) return usage_dict @abc.abstractmethod diff --git a/codewiki/src/be/documentation_generator.py b/codewiki/src/be/documentation_generator.py index 3e9f8b0c..eb5238f7 100644 --- a/codewiki/src/be/documentation_generator.py +++ b/codewiki/src/be/documentation_generator.py @@ -555,4 +555,15 @@ async def run(self) -> None: except Exception as e: logger.error(f"Documentation generation failed: {e!s}") logger.error(f"Traceback: {traceback.format_exc()}") + self.usage_ledger.incomplete = True + try: + path = os.path.join(self.config.docs_dir, "metadata.json") + metadata = file_manager.load_json(path) if os.path.exists(path) else {} + if not isinstance(metadata, dict): + metadata = {} + metadata["usage"] = self.usage_ledger.summary() + file_manager.ensure_directory(self.config.docs_dir) + file_manager.save_json(metadata, path) + except Exception: + logger.warning("Could not save the incomplete run usage report") raise diff --git a/codewiki/src/be/google_provider.py b/codewiki/src/be/google_provider.py index 1cbca525..937ad20f 100644 --- a/codewiki/src/be/google_provider.py +++ b/codewiki/src/be/google_provider.py @@ -15,14 +15,12 @@ import logging import os -from typing import Any, Iterable, Optional +from typing import Any, Optional from google import genai from google.genai import types as genai_types -from pydantic_ai.messages import ModelResponse from pydantic_ai.models.google import GoogleModel, GoogleModelSettings from pydantic_ai.providers.google import GoogleProvider -from pydantic_ai.usage import RequestUsage logger = logging.getLogger(__name__) @@ -123,36 +121,6 @@ def create_google_model(config, model_name: str) -> GoogleModel: ) -def _price(usage: dict[str, Any], model: str) -> Optional[float]: - """USD price of one call from pydantic-ai's genai-prices data, or None if unknown.""" - response = ModelResponse( - parts=[], - usage=RequestUsage( - input_tokens=usage.get("input_tokens", 0), - output_tokens=usage.get("output_tokens", 0), - cache_read_tokens=usage.get("cache_read_tokens", 0), - ), - model_name=model, - provider_name="google", - ) - return _response_cost(response) - - -def _response_cost(response: ModelResponse) -> Optional[float]: - try: - return float(response.cost().total_price) - except Exception: # noqa: BLE001 - unknown model or old genai-prices; cost is optional - return None - - -def messages_cost(messages: Iterable[Any]) -> Optional[float]: - """Summed price of the ModelResponses of an agent run, or None if none could be priced.""" - prices = [ - p for m in messages if isinstance(m, ModelResponse) if (p := _response_cost(m)) is not None - ] - return sum(prices) if prices else None - - def _usage_dict(metadata) -> dict[str, Any]: if metadata is None: return {"requests": 1} @@ -187,9 +155,20 @@ def call_google( ), ) usage = _usage_dict(response.usage_metadata) - price = _price(usage, response.model_version or name) - if price is not None: - usage["cost_usd"] = price + from codewiki.src.be.api_usage import normalize_api_usage, record_response + from codewiki.src.be.pricing import price_usage + + actual_model = response.model_version or name + priced = price_usage( + normalize_api_usage(usage), + getattr(config, "pricing", {}), + config.provider, + actual_model, + model, + ) + if "cost_usd" in priced: + usage["cost_usd"] = priced["cost_usd"] + record_response(usage, model, model=actual_model) candidate = response.candidates[0] if response.candidates else None if candidate is None: diff --git a/codewiki/src/be/llm_services.py b/codewiki/src/be/llm_services.py index de189aaf..a58fd511 100644 --- a/codewiki/src/be/llm_services.py +++ b/codewiki/src/be/llm_services.py @@ -268,8 +268,10 @@ def create_fallback_model(config: Config) -> Model: def create_fallback_models(config: Config) -> FallbackModel: """Create fallback models chain from configuration.""" - main = create_main_model(config) - fallback = create_fallback_model(config) + from codewiki.src.be.api_usage import UsageModel + + main = UsageModel(create_main_model(config), config.main_model) + fallback = UsageModel(create_fallback_model(config), config.fallback_model) # The default fallback_on (ModelAPIError, or ModelHTTPError on older pydantic-ai) # misses UnexpectedModelBehavior, which pydantic-ai raises for a 200 response whose # body does not parse (e.g. choices=None from an OpenAI-compatible gateway); without @@ -309,6 +311,9 @@ def _extract_content(response, model: str) -> Optional[str]: instead of surfacing later as an opaque NoneType error. """ _remember_usage(response) + from codewiki.src.be.api_usage import record_response + + record_response(getattr(response, "usage", None), model, model=getattr(response, "model", None)) choice = response.choices[0] content = choice.message.content finish_reason = getattr(choice, "finish_reason", None) @@ -449,7 +454,7 @@ def _call_llm_via_litellm( max_tokens=config.max_tokens, api_key=config.llm_api_key if config.provider != "bedrock" else None, ) - return _extract_content(response, litellm_model) + return _extract_content(response, model) def _call_llm_via_azure( @@ -479,4 +484,4 @@ def _call_llm_via_azure( messages=_messages(prompt, system_prompt), max_tokens=config.max_tokens, ) - return _extract_content(response, deployment) + return _extract_content(response, model) diff --git a/codewiki/src/be/pricing.py b/codewiki/src/be/pricing.py new file mode 100644 index 00000000..32474bc1 --- /dev/null +++ b/codewiki/src/be/pricing.py @@ -0,0 +1,67 @@ +"""User-maintained API prices, in USD per million tokens.""" + +from __future__ import annotations + +import math +from decimal import Decimal +from typing import Any + + +def validate_pricing(pricing: dict) -> None: + """Reject malformed persisted prices as well as invalid CLI inputs.""" + if not isinstance(pricing, dict): + raise ValueError("pricing must be a provider/model mapping") + for provider, models in pricing.items(): + if not isinstance(provider, str) or not provider or not isinstance(models, dict): + raise ValueError("pricing must be a provider/model mapping") + for model, rates in models.items(): + if not isinstance(model, str) or not model or not isinstance(rates, dict): + raise ValueError("pricing models must have rate mappings") + if not {"input", "output"} <= rates.keys(): + raise ValueError("pricing requires input and output rates") + for key, rate in rates.items(): + if key not in {"input", "output", "cache_read", "cache_write"}: + raise ValueError(f"Unknown pricing rate: {key}") + if ( + isinstance(rate, bool) + or not isinstance(rate, (int, float)) + or not math.isfinite(rate) + or rate < 0 + ): + raise ValueError("Pricing rates must be finite, nonnegative numbers") + + +def price_usage( + usage: dict[str, Any] | None, + pricing: dict, + provider: str, + model: str, + requested_model: str, +) -> dict[str, Any]: + """Price inclusive input tokens; unknown usage/prices never mean free.""" + models = pricing.get(provider, {}) + pricing_model = model if model in models else requested_model + configured = models.get(pricing_model) + if usage is None or configured is None: + return {"cost_source": "unknown"} + rates = dict(configured) + rates.setdefault("cache_read", rates["input"]) + rates.setdefault("cache_write", rates["input"]) + read = usage.get("cache_read_tokens", 0) + write = usage.get("cache_write_tokens", 0) + uncached = usage.get("input_tokens", 0) - read - write + if uncached < 0: + return {"cost_source": "unknown"} + buckets = { + "input": uncached, + "output": usage.get("output_tokens", 0), + "cache_read": read, + "cache_write": write, + } + cost = sum(Decimal(str(rates[k])) * Decimal(str(n)) for k, n in buckets.items()) + return { + "cost_usd": float(cost / 1_000_000), + "cost_source": "configured", + "pricing_model": pricing_model, + "rates": rates, + } diff --git a/codewiki/src/be/pydantic_ai_backend.py b/codewiki/src/be/pydantic_ai_backend.py index a945b382..704b7ee2 100644 --- a/codewiki/src/be/pydantic_ai_backend.py +++ b/codewiki/src/be/pydantic_ai_backend.py @@ -23,7 +23,8 @@ ) from codewiki.src.be.agent_tools.read_code_components import read_code_components_tool from codewiki.src.be.agent_tools.str_replace_editor import str_replace_editor_tool -from codewiki.src.be.backend import AgentReply, LLMBackend, is_google_provider, usage_to_dict +from codewiki.src.be.api_usage import track_usage, record_response +from codewiki.src.be.backend import AgentReply, LLMBackend, UsageLedger, usage_to_dict from codewiki.src.be.dependency_analyzer.models.core import Node from codewiki.src.be.llm_services import call_llm, create_fallback_models, pop_last_usage from codewiki.src.be.prompt_template import ( @@ -51,27 +52,12 @@ def _run_usage(result: Any) -> dict[str, Any] | None: return None -def _run_usage_with_cost(result: Any, provider: str) -> dict[str, Any] | None: - """`_run_usage`, plus ``cost_usd`` for Google providers (priced from genai-prices).""" - usage = _run_usage(result) - if usage is None or not is_google_provider(provider): - return usage - from codewiki.src.be.google_provider import messages_cost - - try: - cost = messages_cost(result.all_messages()) - except Exception: # noqa: BLE001 — cost is optional telemetry - cost = None - if cost is not None: - usage["cost_usd"] = cost - return usage - - class PydanticAIBackend(LLMBackend): """API-key based backend using pydantic-ai + openai/litellm/google-genai clients.""" def __init__(self, config: Config) -> None: self._config = config + self._usage_ledger = UsageLedger() self._fallback_models = create_fallback_models(config) self._custom_instructions = config.get_prompt_addition() self.last_usage: dict[str, Any] | None = None @@ -83,17 +69,21 @@ def complete( model: str | None = None, system_prompt: str | None = None, ) -> str: - pop_last_usage() - started = time.time() - result = call_llm(prompt, self._config, model=model, system_prompt=system_prompt) - self._record_usage( - "complete", - "", - pop_last_usage(), - seconds=time.time() - started, - model=model or self._config.main_model, - ) - return result + started = time.monotonic() + with track_usage(self.usage_ledger, self._config, "complete") as tracked: + try: + return call_llm(prompt, self._config, model=model, system_prompt=system_prompt) + except Exception as exc: + if not tracked.entries: + record_response( + None, model or self._config.main_model, error=type(exc).__name__ + ) + raise + finally: + if len(tracked.entries) == 1: + tracked.entries[0]["seconds"] = round(time.monotonic() - started, 2) + self.last_usage = tracked.totals() + pop_last_usage() async def run_update_agent( self, @@ -110,19 +100,19 @@ async def run_update_agent( retries=self._config.agent_retries, ) started = time.time() - result = await agent.run( - user_prompt, - deps=deps, - usage_limits=UsageLimits(request_limit=self._config.request_limit), - ) + with track_usage( + self.usage_ledger, self._config, "update_agent", deps.current_module_name + ) as tracked: + try: + result = await agent.run( + user_prompt, + deps=deps, + usage_limits=UsageLimits(request_limit=self._config.request_limit), + ) + finally: + self.last_usage = tracked.totals() seconds = time.time() - started - usage = self._record_usage( - "update_agent", - deps.current_module_name, - _run_usage_with_cost(result, self._config.provider), - seconds=seconds, - model=self._config.main_model, - ) + usage = self.last_usage text = result.output if isinstance(result.output, str) else str(result.output) return AgentReply(text=text, usage=usage, seconds=seconds) @@ -195,29 +185,23 @@ async def run_module_agent( custom_instructions=self._custom_instructions, ) - started = time.time() try: - result = await agent.run( - format_user_prompt( - module_name=module_name, - core_component_ids=core_component_ids, - components=components, - module_tree=deps.module_tree, - module_path=module_path, - layout=config_layout(config), - ), - deps=deps, - usage_limits=UsageLimits(request_limit=config.request_limit), - ) - # Covers this module's own run; pydantic-ai sub-module agents run - # with their own usage and are not counted here. - self._record_usage( - "module_agent", - module_name, - _run_usage_with_cost(result, config.provider), - seconds=time.time() - started, - model=config.main_model, - ) + with track_usage(self.usage_ledger, config, "module_agent", module_name) as tracked: + try: + await agent.run( + format_user_prompt( + module_name=module_name, + core_component_ids=core_component_ids, + components=components, + module_tree=deps.module_tree, + module_path=module_path, + layout=config_layout(config), + ), + deps=deps, + usage_limits=UsageLimits(request_limit=config.request_limit), + ) + finally: + self.last_usage = tracked.totals() file_manager.save_json(deps.module_tree, module_tree_path) return deps.module_tree except Exception as e: diff --git a/codewiki/src/be/updater/record.py b/codewiki/src/be/updater/record.py index 0dc6efa1..f83d70e4 100644 --- a/codewiki/src/be/updater/record.py +++ b/codewiki/src/be/updater/record.py @@ -6,6 +6,7 @@ import json import os import time +import uuid from dataclasses import asdict, dataclass, field from typing import Any @@ -60,6 +61,8 @@ class UpdateRecord: detector_notes: list[str] = field(default_factory=list) errors: list[str] = field(default_factory=list) wall_seconds: float | None = None + usage_report: dict[str, Any] | None = None + run_id: str = field(default_factory=lambda: str(uuid.uuid4())) def add_call(self, cost: CallCost) -> None: self.calls.append(asdict(cost)) @@ -74,6 +77,8 @@ def summary(self) -> dict[str, Any]: if isinstance(v, (int, float)): usage_total[k] = usage_total.get(k, 0) + v return { + "run_id": self.run_id, + "started_at": self.started_at, "outcome": self.outcome, "finished_at": self.finished_at, "rung": self.options.get("rung"), @@ -85,8 +90,9 @@ def summary(self) -> dict[str, Any]: 1 for a in self.ownership if str(a.get("rule", "")).startswith("4:") ), "fallback": self.fallback, - "n_calls": len(self.calls), - "usage_total": usage_total, + "n_calls": self.usage_report["n_calls"] if self.usage_report else len(self.calls), + "usage_total": self.usage_report["totals"] if self.usage_report else usage_total, + "usage": self.usage_report, "pages_written": sorted(set(self.pages_written)), "pages_removed": sorted(set(self.pages_removed)), "wall_seconds": self.wall_seconds, diff --git a/codewiki/src/config.py b/codewiki/src/config.py index 0f0d88a6..338a062b 100644 --- a/codewiki/src/config.py +++ b/codewiki/src/config.py @@ -1,6 +1,6 @@ import argparse import os -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any from dotenv import load_dotenv @@ -115,6 +115,7 @@ class Config: # Prompt caching for agentic/multi-turn calls (auto-disables per model if # the provider rejects cache_control markers) prompt_caching: bool = True + pricing: dict = field(default_factory=dict) # Agent instructions for customization agent_instructions: dict[str, Any] | None = None # Apply Git ignore rules before dependency analysis @@ -263,6 +264,7 @@ def from_cli( artifact_token_budget: int = DEFAULT_ARTIFACT_TOKEN_BUDGET, with_prose: bool = False, layout: str = DEFAULT_LAYOUT, + pricing: dict | None = None, ) -> "Config": """ Create configuration for CLI context. @@ -306,6 +308,9 @@ def from_cli( Config instance """ base_output_dir = os.path.join(output_dir, "temp") + from codewiki.src.be.pricing import validate_pricing + + validate_pricing(pricing or {}) return cls( repo_path=repo_path, @@ -338,4 +343,5 @@ def from_cli( artifact_token_budget=artifact_token_budget, with_prose=with_prose, layout=layout, + pricing=pricing or {}, ) diff --git a/guides/cli-reference.md b/guides/cli-reference.md index 6f2043ae..3c499237 100644 --- a/guides/cli-reference.md +++ b/guides/cli-reference.md @@ -9,6 +9,7 @@ Commands: | --- | --- | | `codewiki generate` | Build or update the documentation for the current directory | | `codewiki config set` | Store provider, model, and token settings | +| `codewiki config pricing set/show/remove` | Manage API model rates in USD per million tokens | | `codewiki config agent` | Store default include/exclude/focus/doc-type/instructions | | `codewiki config show` | Print the stored settings (`--json` for machine-readable output) | | `codewiki config validate` | Check the settings and test the provider connection (`--skip-api-test`, `--verbose`) | @@ -189,6 +190,22 @@ Where things are stored: file. - **Settings and agent defaults**: `~/.codewiki/config.json`. +## `codewiki config pricing` + +```bash +codewiki config pricing set MODEL --input RATE --output RATE \ + [--cache-read RATE] [--cache-write RATE] [--provider PROVIDER] +codewiki config pricing show [--json] +codewiki config pricing remove MODEL [--provider PROVIDER] +``` + +Prices are stored alongside provider settings and appear in `config show`. +Input and output rates are required, finite, nonnegative USD amounts per million +tokens; omitted cache rates use the input rate. The provider defaults to your +active API provider. Claude Code and Codex keep their existing provider costs. +Missing API prices produce unknown or partial reports without stopping the run. +See [Cost tracking](providers.md#cost-tracking) for report details. + ## `codewiki config agent` Stores default analysis settings so you do not repeat them on every run. diff --git a/guides/incremental-updates.md b/guides/incremental-updates.md index be7f3a9a..4ae00dad 100644 --- a/guides/incremental-updates.md +++ b/guides/incremental-updates.md @@ -117,6 +117,14 @@ agent call. Every page written in an update can be explained after the fact. ## Cost +Each update prints tokens and estimated USD cost by model and phase. Its +`metadata.json` usage report is also attached to the update record and the +corresponding update-history entry. Full-build fallback includes both the +update attempt and rebuild in the same run total. An unchanged update records +zero calls and zero cost. Configure API rates with +`codewiki config pricing set MODEL --input RATE --output RATE`; see +[cost tracking](providers.md#cost-tracking) for cache rates and unknown prices. + Measured on svelte (about 125k lines of JavaScript) with the same model: | Run | Cost | Time | diff --git a/guides/providers.md b/guides/providers.md index 56f0c333..3afd9b47 100644 --- a/guides/providers.md +++ b/guides/providers.md @@ -93,8 +93,8 @@ Notes: 65536. - **Privacy.** Google may use prompts sent with a free-tier AI Studio key for product improvement. Use a paid key, or Vertex AI, for private code. -- **Cost.** The cost summary prices Gemini runs from the installed - `genai-prices` data. A model it does not know is reported as unpriced. +- **Cost.** Configure Gemini API rates with `codewiki config pricing set`. + Calls without configured rates are reported as unpriced. ### Atlas Cloud @@ -150,14 +150,39 @@ Things to know: ### Cost tracking -At the end of `codewiki generate`, CodeWiki prints the tokens used and an -estimated cost for the run. `metadata.json` gets a `usage` block with totals, -a per-phase breakdown (`clustering`, `modules`, `parent_modules`, `overview`, -`update`) and one entry per LLM call. Each nested sub-module agent is counted -once. Modules resumed from an earlier run cost nothing and don't appear. +At the end of `codewiki generate` and `codewiki generate --update`, CodeWiki +prints tokens and estimated USD cost, grouped by model and phase. The same +report is saved in `metadata.json` under `usage`; updates also attach it to +`update_record.json` and metadata's update history. Totals cover the current +invocation, including nested agents, retries with reported usage, and fallback +models. Resumed pages add no usage; unchanged updates show zero calls and cost. -The cost is what the same tokens would cost on the API. Your subscription -is billed as usual. +For **API providers**, set your own rates in USD per million tokens. These +example numbers are illustrative; use the rates charged by your endpoint: + +```bash +codewiki config pricing set YOUR_MODEL --input 1 --output 2 \ + --cache-read 0.25 --cache-write 1.25 +codewiki config pricing show --json +codewiki config pricing remove YOUR_MODEL +``` + +Rates are stored by provider and model in `~/.codewiki/config.json`. +`--provider` defaults to the active provider. Configure main, cluster, and +fallback models separately when their names differ. Cache rates default to +the input rate when omitted, and explicit zero rates are supported. The report +records the effective rates and both requested and returned model names; +returned-model rates take precedence over the requested alias's rates. + +Cached tokens are included in input totals, not added again. Costs apply each +cache rate to its subset and the input rate to the remaining input. Unknown +prices or unavailable usage never block generation: estimates are labeled +`partial` or `unknown`, with the affected models listed. Failed runs save and +print available usage with an incomplete marker. Estimates are not invoices. + +**Subscription providers** retain the cost returned through their integration; +CodeWiki's API pricing settings do not override it. The displayed cost is an +API-equivalent estimate; your subscription is billed as usual. - **Claude Code** reports the cost itself. - **Codex** is priced from caw's rate table using the model Codex actually @@ -169,6 +194,13 @@ is billed as usual. caw pricing set codex --input --cached --output ``` +For an opt-in live check, run +`.venv/bin/python tests/smoke_test_costs.py --env-file /absolute/path/.env` +from the project root. It maps `BASE_URL`, `API_KEY`, `MAIN_MODEL`, and +`FALLBACK_MODEL` into an isolated temporary configuration, generates a tiny +repository's docs, updates them, and checks an unchanged update. Its synthetic +rates validate arithmetic rather than provider billing; logs redact credentials. + ## Which model to use - The **main model** writes the pages. It needs strong code understanding diff --git a/tests/smoke_test_costs.py b/tests/smoke_test_costs.py new file mode 100644 index 00000000..4299a236 --- /dev/null +++ b/tests/smoke_test_costs.py @@ -0,0 +1,211 @@ +"""Opt-in live generation/update cost check. Never collected by pytest. + +Run: .venv/bin/python tests/smoke_test_costs.py --env-file /absolute/path/.env +Rates below are synthetic test data, not a provider price quote. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from click.testing import CliRunner +from dotenv import dotenv_values + +import codewiki.cli.config_manager as config_module +from codewiki.cli.main import cli + + +def git(repo, *args): + return subprocess.check_output( + ["git", "-C", str(repo), *args], stderr=subprocess.STDOUT, text=True + ) + + +def verify(report): + assert report["cost_status"] == "complete", report["unpriced_models"] + expected = 0 + for call in report["calls"]: + assert call["usage_reported"] + rates = call["rates"] + read, write = call["cache_read_tokens"], call["cache_write_tokens"] + expected += ( + (call["input_tokens"] - read - write) * rates["input"] + + read * rates["cache_read"] + + write * rates["cache_write"] + + call["output_tokens"] * rates["output"] + ) / 1_000_000 + assert math.isclose(expected, report["totals"]["cost_usd"], abs_tol=1e-12) + for grouping in ("by_model", "by_phase"): + assert math.isclose( + sum(g.get("cost_usd", 0) for g in report[grouping].values()), + expected, + abs_tol=1e-12, + ) + + +def verify_display(report, output): + calls = re.search(r"LLM calls:\s+([\d,]+)", output) + tokens = re.search(r"Tokens used:\s+([\d,]+)", output) + cost = re.search(r"Est\. cost:\s+\$([\d,.]+)", output) + assert calls and int(calls[1].replace(",", "")) == report["n_calls"] + expected_tokens = report["totals"].get("input_tokens", 0) + report["totals"].get( + "output_tokens", 0 + ) + assert tokens and int(tokens[1].replace(",", "")) == expected_tokens + expected_cost = report["totals"]["cost_usd"] + assert cost and math.isclose( + float(cost[1].replace(",", "")), + expected_cost, + abs_tol=0.005 if expected_cost >= 0.01 else 5e-9, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--env-file", required=True, type=Path) + args = parser.parse_args() + env = dotenv_values(args.env_file) + for name in ("BASE_URL", "API_KEY", "MAIN_MODEL", "FALLBACK_MODEL"): + if not env.get(name): + parser.error(f"Missing {name} in env file") + + root = Path(tempfile.mkdtemp(prefix="codewiki-cost-smoke-")) + repo, output, config_dir = root / "repo", root / "docs", root / "config" + repo.mkdir() + config_dir.mkdir() + config_module.CONFIG_DIR = config_dir + config_module.CONFIG_FILE = config_dir / "config.json" + config_module.CREDENTIALS_FILE = config_dir / "credentials.json" + os.environ["CODEWIKI_NO_KEYRING"] = "1" + manager = config_module.ConfigManager() + manager.save( + base_url=env["BASE_URL"], + api_key=env["API_KEY"], + main_model=env["MAIN_MODEL"], + fallback_model=env["FALLBACK_MODEL"], + cluster_model=env["MAIN_MODEL"], + provider="openai-compatible", + ) + runner = CliRunner() + for model in {env["MAIN_MODEL"], env["FALLBACK_MODEL"]}: + result = runner.invoke( + cli, + [ + "config", + "pricing", + "set", + model, + "--input", + "1", + "--output", + "2", + "--cache-read", + "0.25", + ], + ) + assert result.exit_code == 0, result.output + + source = repo / "calculator.py" + source.write_text( + '"""Small integer calculator."""\n\n' + "def add(a: int, b: int) -> int:\n return a + b\n\n" + "def multiply(a: int, b: int) -> int:\n return a * b\n" + ) + git(repo, "init", "-q") + git(repo, "add", ".") + git( + repo, + "-c", + "user.name=Cost Smoke", + "-c", + "user.email=smoke@example.invalid", + "commit", + "-qm", + "initial", + ) + previous_cwd = Path.cwd() + print(f"Smoke artifacts: {root}", flush=True) + try: + os.chdir(repo) + base = [ + "generate", + "--output", + str(output), + "--no-artifacts", + "--max-tokens", + "8192", + "--request-limit", + "24", + "--instructions", + "Keep documentation concise, under 300 words per page. Do not include Mermaid diagrams.", + ] + reports = {} + for stage in ("generate", "update", "no-change"): + if stage == "update": + source.write_text( + source.read_text().replace("return a + b", "return max(-100, min(100, a + b))") + ) + git(repo, "add", ".") + git( + repo, + "-c", + "user.name=Cost Smoke", + "-c", + "user.email=smoke@example.invalid", + "commit", + "-qm", + "clamp addition", + ) + flags = ( + [] if stage == "generate" else ["--update", "--tau-full", "1", "--tau-tree", "1"] + ) + print(f"Running {stage}...", flush=True) + result = runner.invoke(cli, base + flags) + log = result.output + for secret in (env["API_KEY"], env["BASE_URL"]): + log = log.replace(secret, "[redacted]") + (root / f"{stage}.log").write_text(log) + if result.exit_code: + raise RuntimeError( + f"{stage} failed with exit {result.exit_code}; see {root / (stage + '.log')}" + ) + metadata = json.loads((output / "metadata.json").read_text()) + report = metadata["usage"] + verify(report) + verify_display(report, result.output) + assert "By model" in result.output or report["n_calls"] == 0 + assert "By phase" in result.output or report["n_calls"] == 0 + if stage != "generate": + record = json.loads((output / "update_record.json").read_text()) + assert record["usage_report"] == report == metadata["last_update"]["usage"] + assert len(metadata["update_history"]) == (1 if stage == "update" else 2) + assert record["outcome"] == ("incremental" if stage == "update" else "no_change") + if stage == "no-change": + assert report["n_calls"] == 0 and report["totals"]["cost_usd"] == 0 + else: + assert report["n_calls"] > 0 and report["totals"]["cost_usd"] > 0 + reports[stage] = report + (root / f"{stage}.json").write_text(json.dumps(report, indent=2)) + print( + f"{stage}: {report['n_calls']} calls, ${report['totals']['cost_usd']:.8f} at synthetic test rates", + flush=True, + ) + (root / "reports.json").write_text(json.dumps(reports, indent=2)) + finally: + os.chdir(previous_cwd) + shutil.rmtree(config_dir) + + +if __name__ == "__main__": + main() diff --git a/tests/test_agent_limits.py b/tests/test_agent_limits.py index fffa1fd8..eaa1eb63 100644 --- a/tests/test_agent_limits.py +++ b/tests/test_agent_limits.py @@ -47,13 +47,21 @@ def test_backend_config_defaults(): def test_fallback_covers_malformed_responses(monkeypatch): + from pydantic_ai.models.function import FunctionModel + captured = {} - monkeypatch.setattr(llm_services, "create_main_model", lambda config: "main") - monkeypatch.setattr(llm_services, "create_fallback_model", lambda config: "fallback") + monkeypatch.setattr( + llm_services, "create_main_model", lambda config: FunctionModel(lambda m, i: None) + ) + monkeypatch.setattr( + llm_services, "create_fallback_model", lambda config: FunctionModel(lambda m, i: None) + ) monkeypatch.setattr( llm_services, "FallbackModel", lambda *models, **kwargs: captured.update(kwargs) ) - llm_services.create_fallback_models(SimpleNamespace()) + llm_services.create_fallback_models( + SimpleNamespace(main_model="main", fallback_model="fallback") + ) # a 200 response with an unparseable body must also switch to the fallback model assert set(captured["fallback_on"]) == {llm_services.ModelAPIError, UnexpectedModelBehavior} diff --git a/tests/test_google_provider.py b/tests/test_google_provider.py index 8a56812a..2c05b05c 100644 --- a/tests/test_google_provider.py +++ b/tests/test_google_provider.py @@ -1,15 +1,12 @@ """Native Google provider: model factories, one-shot calls, usage and cost.""" import logging -from decimal import Decimal from types import SimpleNamespace import pytest from google.genai import types as gt -from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart from pydantic_ai.models.fallback import FallbackModel from pydantic_ai.models.google import GoogleModel -from pydantic_ai.usage import RequestUsage from codewiki.src.be import google_provider from codewiki.src.be.llm_services import ( @@ -101,7 +98,7 @@ def test_google_fallback_models_are_google_models(clean_env): chain = create_fallback_models(config) assert isinstance(chain, FallbackModel) - assert all(isinstance(m, GoogleModel) for m in chain.models) + assert all(isinstance(m.wrapped, GoogleModel) for m in chain.models) assert chain.models[0].model_name == "gemini-3.8-flash" @@ -148,12 +145,16 @@ def install(response): return install -def test_call_llm_google_text_and_usage(fake_models, monkeypatch): +def test_call_llm_google_text_and_usage(fake_models): models = fake_models(_response([_text_candidate("hello")])) - monkeypatch.setattr(google_provider, "_price", lambda usage, model: 0.25) pop_last_usage() - text = call_llm("p", _config(provider="google", llm_api_key="test-key"), system_prompt="s") + config = _config( + provider="google", + llm_api_key="test-key", + pricing={"google": {MODEL: {"input": 1000, "output": 3000}}}, + ) + text = call_llm("p", config, system_prompt="s") assert text == "hello" call = models.calls[0] @@ -200,7 +201,7 @@ def test_call_llm_google_empty_returns_none(fake_models, caplog, response, reaso assert pop_last_usage()["output_tokens"] == 50 -def test_cost_unknown_model_is_omitted(fake_models, monkeypatch): +def test_cost_unknown_model_is_omitted(fake_models): response = _response([_text_candidate("hello")]) response.model_version = "not-a-real-model" fake_models(response) @@ -209,13 +210,6 @@ def test_cost_unknown_model_is_omitted(fake_models, monkeypatch): assert call_llm("p", _config(provider="google", llm_api_key="test-key")) == "hello" assert "cost_usd" not in pop_last_usage() - priced = ModelResponse(parts=[], usage=RequestUsage(), model_name=MODEL) - monkeypatch.setattr( - ModelResponse, "cost", lambda self: SimpleNamespace(total_price=Decimal("0.5")) - ) - messages = [priced, ModelRequest(parts=[UserPromptPart("x")]), priced] - assert google_provider.messages_cost(messages) == 1.0 - def test_call_google_reuses_sync_client(monkeypatch): built = [] diff --git a/tests/test_pricing.py b/tests/test_pricing.py new file mode 100644 index 00000000..99d4e3b2 --- /dev/null +++ b/tests/test_pricing.py @@ -0,0 +1,475 @@ +"""Pricing configuration, request accounting, and persisted run reports.""" + +import asyncio +import json +from types import SimpleNamespace + +import pytest +from click.testing import CliRunner +from pydantic_ai import Agent +from pydantic_ai.exceptions import ModelHTTPError +from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.usage import RequestUsage + +import codewiki.cli.config_manager as cm +from codewiki.cli.adapters.doc_generator import CLIDocumentationGenerator +from codewiki.cli.commands.config import config_group +from codewiki.cli.models.config import Configuration +from codewiki.cli.models.job import JobStatistics +from codewiki.cli.utils.errors import APIError +from codewiki.cli.utils.instructions import _echo_usage +from codewiki.src.be.api_usage import UsageModel, normalize_api_usage, track_usage, usage_target +from codewiki.src.be.backend import UsageLedger +from codewiki.src.be.pricing import price_usage, validate_pricing +from codewiki.src.be.pydantic_ai_backend import PydanticAIBackend +from codewiki.src.be.updater.record import UpdateRecord +from codewiki.src.config import Config + + +RATES = {"input": 1, "output": 2, "cache_read": 0.25, "cache_write": 1.5} +PRICING = {"openai-compatible": {"main": RATES}} + + +@pytest.fixture +def config_store(tmp_path, monkeypatch): + monkeypatch.setattr(cm, "CONFIG_DIR", tmp_path) + monkeypatch.setattr(cm, "CONFIG_FILE", tmp_path / "config.json") + monkeypatch.setattr(cm, "CREDENTIALS_FILE", tmp_path / "credentials.json") + monkeypatch.setenv("CODEWIKI_NO_KEYRING", "1") + return tmp_path + + +def config(tmp_path, **kwargs): + settings = dict( + repo_path=str(tmp_path), + output_dir=str(tmp_path / "docs"), + llm_base_url="https://example.invalid/v1", + llm_api_key="test-key", + main_model="main", + cluster_model="cluster", + fallback_model="backup", + pricing=PRICING, + ) + settings.update(kwargs) + return Config.from_cli(**settings) + + +def response(name="main", input_tokens=100, output_tokens=20, cache_read_tokens=40): + return ModelResponse( + parts=[TextPart("ok")], + model_name=name, + usage=RequestUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_read_tokens=cache_read_tokens, + ), + ) + + +def test_pricing_commands_preserve_other_settings(config_store): + manager = cm.ConfigManager() + manager.save( + provider="openai-compatible", + base_url="https://example.invalid/v1", + main_model="main", + cluster_model="main", + fallback_model="main", + ) + runner = CliRunner() + result = runner.invoke( + config_group, + ["pricing", "set", "main", "--input", "1", "--output", "2", "--cache-read", "0.25"], + ) + assert result.exit_code == 0, result.output + manager = cm.ConfigManager() + manager.load() + cfg = manager.get_config() + assert cfg.main_model == "main" + assert cfg.pricing["openai-compatible"]["main"]["cache_read"] == 0.25 + manager.save(max_tokens=1234) + again = cm.ConfigManager() + again.load() + assert again.get_config().pricing == cfg.pricing + shown = runner.invoke(config_group, ["pricing", "show", "--json"]) + assert json.loads(shown.output) == cfg.pricing + shown = runner.invoke(config_group, ["show", "--json"]) + assert json.loads(shown.output)["pricing"] == cfg.pricing + removed = runner.invoke(config_group, ["pricing", "remove", "main"]) + assert removed.exit_code == 0, removed.output + assert json.loads(runner.invoke(config_group, ["pricing", "show", "--json"]).output) == {} + + +@pytest.mark.parametrize("bad", ["-1", "nan", "inf", "-inf"]) +def test_invalid_cli_prices_do_not_save(config_store, bad): + result = CliRunner().invoke( + config_group, ["pricing", "set", "main", "--input", bad, "--output", "1"] + ) + assert result.exit_code != 0 + assert not cm.CONFIG_FILE.exists() + + +def test_pricing_before_provider_setup_and_provider_isolation(config_store): + runner = CliRunner() + result = runner.invoke( + config_group, + ["pricing", "set", "main", "--input", "0", "--output", "0", "--provider", "google"], + ) + assert result.exit_code == 0, result.output + manager = cm.ConfigManager() + manager.load() + assert manager.get_config().pricing == {"google": {"main": {"input": 0.0, "output": 0.0}}} + result = runner.invoke( + config_group, + ["pricing", "set", "main", "--input", "1", "--output", "1", "--provider", "codex"], + ) + assert result.exit_code != 0 + assert "provider cost results" in result.output + + +def test_pricing_config_compatibility_and_backend_propagation(tmp_path): + cfg = Configuration.from_dict({"main_model": "main", "pricing": PRICING}) + restored = Configuration.from_dict(cfg.to_dict()) + assert ( + restored.to_backend_config(str(tmp_path), str(tmp_path / "docs"), "key").pricing == PRICING + ) + assert Configuration.from_dict({}).pricing == {} + with pytest.raises(ValueError): + validate_pricing({"api": {"m": {"input": True, "output": 1}}}) + + +def test_cache_arithmetic_alias_and_rate_snapshot(): + usage = { + "input_tokens": 1000, + "output_tokens": 200, + "cache_read_tokens": 400, + "cache_write_tokens": 100, + } + priced = price_usage(usage, PRICING, "openai-compatible", "response-alias", "main") + assert priced["cost_usd"] == pytest.approx((500 + 400 + 100 + 150) / 1_000_000) + assert priced["pricing_model"] == "main" + assert priced["rates"] == RATES + assert price_usage(usage, PRICING, "google", "main", "main") == {"cost_source": "unknown"} + assert price_usage(None, PRICING, "openai-compatible", "main", "main") == { + "cost_source": "unknown" + } + override = {"openai-compatible": {"main": RATES, "response-alias": {"input": 0, "output": 0}}} + assert ( + price_usage(usage, override, "openai-compatible", "response-alias", "main")["cost_usd"] == 0 + ) + + +def test_cache_rates_default_to_input(): + priced = price_usage( + {"input_tokens": 100, "cache_read_tokens": 50, "output_tokens": 10}, + {"api": {"m": {"input": 2, "output": 3}}}, + "api", + "m", + "m", + ) + assert priced["cost_usd"] == pytest.approx(230 / 1_000_000) + + +def test_openai_usage_aliases_and_cached_details(): + normalized = normalize_api_usage( + SimpleNamespace( + prompt_tokens=100, + completion_tokens=20, + total_tokens=120, + prompt_tokens_details=SimpleNamespace(cached_tokens=40), + ) + ) + assert normalized["total_tokens"] == 120 + assert normalized["input_tokens"] == 100 + assert normalized["cache_read_tokens"] == 40 + assert "prompt_tokens" not in normalized + assert normalize_api_usage({"requests": 1}) is None + + +def test_complete_records_actual_model_and_configured_cost(tmp_path, monkeypatch): + from codewiki.src.be import llm_services + + backend = PydanticAIBackend(config(tmp_path)) + reply = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"), finish_reason="stop")], + usage=SimpleNamespace( + prompt_tokens=100, completion_tokens=20, prompt_tokens_details={"cached_tokens": 40} + ), + model="alias", + ) + monkeypatch.setattr( + llm_services, + "create_openai_client", + lambda cfg: SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **kw: reply)) + ), + ) + with backend.usage_ledger.in_phase("overview", "repo"): + assert backend.complete("hello") == "ok" + report = backend.usage_ledger.summary() + assert report["n_calls"] == 1 + assert report["calls"][0]["model"] == "alias" + assert report["calls"][0]["requested_model"] == "main" + assert report["totals"]["cost_usd"] == pytest.approx(110 / 1_000_000) + assert backend.last_usage == report["totals"] + + +def test_nested_agent_requests_counted_once_with_own_targets(tmp_path): + ledger = UsageLedger() + cfg = config(tmp_path) + child = Agent(UsageModel(FunctionModel(lambda m, i: response(), model_name="main"), "main")) + + async def delegate() -> str: + with usage_target("child"): + return (await child.run("child request")).output + + def parent(messages, info): + if len(messages) == 1: + result = response() + result.parts = [ToolCallPart("delegate", {})] + return result + return response() + + agent = Agent(UsageModel(FunctionModel(parent, model_name="main"), "main"), tools=[delegate]) + + async def run(): + with track_usage(ledger, cfg, "module_agent", "parent") as scope: + await agent.run("parent request") + assert scope.totals()["input_tokens"] == 300 + + asyncio.run(run()) + report = ledger.summary() + assert report["n_calls"] == 3 + assert [c["target"] for c in report["calls"]] == ["parent", "child", "parent"] + assert report["totals"]["total_tokens"] == 360 + assert report["cost_status"] == "complete" + + +def test_fallback_attribution_and_unknown_usage(tmp_path, monkeypatch): + from codewiki.src.be import llm_services + + def fail(messages, info): + raise ModelHTTPError(503, "main", "failed") + + monkeypatch.setattr( + llm_services, "create_main_model", lambda c: FunctionModel(fail, model_name="main") + ) + monkeypatch.setattr( + llm_services, + "create_fallback_model", + lambda c: FunctionModel(lambda m, i: response("backup"), model_name="backup"), + ) + cfg = config(tmp_path, pricing={"openai-compatible": {"backup": RATES}}) + ledger = UsageLedger() + + async def run(): + with track_usage(ledger, cfg, "update_agent", "page"): + result = await Agent(llm_services.create_fallback_models(cfg)).run("hi") + assert result.output == "ok" + + asyncio.run(run()) + report = ledger.summary() + assert report["n_calls"] == 2 + assert report["cost_status"] == "partial" + assert report["unknown_usage_calls"] == 1 + assert report["unpriced_models"] == ["main"] + assert report["calls"][1]["pricing_model"] == "backup" + + +def test_tool_retry_requests_and_failure_after_billed_response(tmp_path): + from pydantic_ai import ModelRetry + + ledger = UsageLedger() + cfg = config(tmp_path, pricing={}) # missing prices must not block agent runs + requests, attempts = 0, 0 + + async def sometimes_fails() -> str: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise ModelRetry("try again") + return "done" + + def model(messages, info): + nonlocal requests + requests += 1 + result = response() + if requests <= 2: + result.parts = [ToolCallPart("sometimes_fails", {})] + return result + + agent = Agent( + UsageModel(FunctionModel(model, model_name="main"), "main"), + tools=[sometimes_fails], + retries=1, + ) + + async def run(): + with track_usage(ledger, cfg, "module_agent", "page"): + result = await agent.run("hi") + assert result.output == "ok" + + asyncio.run(run()) + assert attempts == 2 and requests == 3 + report = ledger.summary() + assert report["n_calls"] == 3 + assert report["totals"]["input_tokens"] == 300 + assert report["cost_status"] == "unknown" + + async def fail() -> str: + raise RuntimeError("tool failed after billing") + + def failing_model(messages, info): + result = response() + result.parts = [ToolCallPart("fail", {})] + return result + + ledger = UsageLedger() + agent = Agent(UsageModel(FunctionModel(failing_model, model_name="main"), "main"), tools=[fail]) + + async def failing_run(): + with track_usage(ledger, config(tmp_path), "module_agent", "page"): + await agent.run("hi") + + with pytest.raises(RuntimeError, match="tool failed after billing"): + asyncio.run(failing_run()) + assert ledger.summary()["totals"]["cost_usd"] > 0 + assert ledger.summary()["n_calls"] == 1 + + +def test_zero_unknown_partial_and_small_cost_output(capsys): + ledger = UsageLedger() + _echo_usage({"usage_report": ledger.summary()}) + assert "$0.00" in capsys.readouterr().out + ledger.add( + "request", + "a", + {"input_tokens": 1, "cost_usd": 0.000001}, + model="priced", + cost_source="configured", + ) + ledger.add("request", "b", {"input_tokens": 1}, model="unpriced") + _echo_usage({"usage_report": ledger.summary()}) + output = capsys.readouterr().out + assert "$0.000001" in output and "partial" in output + assert "By model" in output and "By phase" in output + assert "codewiki config pricing set" in output + zero = UsageLedger() + zero.add("request", "", {"input_tokens": 10, "cost_usd": 0}, cost_source="configured") + assert zero.summary()["cost_status"] == "complete" + assert not zero.summary()["unpriced_models"] + + +def adapter(tmp_path, ledger): + instance = object.__new__(CLIDocumentationGenerator) + instance._usage_ledger = ledger + instance.output_dir = tmp_path + instance.config = {"provider": "openai-compatible"} + instance.job = SimpleNamespace(statistics=JobStatistics()) + return instance + + +def test_failure_report_preserves_metadata_and_raises_original(tmp_path, monkeypatch, capsys): + ledger = UsageLedger() + ledger.add( + "request", "", {"input_tokens": 10, "output_tokens": 2, "cost_usd": 0.001}, model="main" + ) + (tmp_path / "metadata.json").write_text(json.dumps({"generation_info": {"commit_id": "old"}})) + instance = CLIDocumentationGenerator(tmp_path, tmp_path, config(tmp_path).__dict__) + + async def fail(cfg): + instance._usage_ledger = ledger + raise APIError("original failure") + + monkeypatch.setattr(instance, "_run_backend_generation", fail) + with pytest.raises(APIError, match="original failure"): + instance.generate() + metadata = json.loads((tmp_path / "metadata.json").read_text()) + assert metadata["generation_info"]["commit_id"] == "old" + assert metadata["usage"]["incomplete"] is True + assert metadata["usage"]["totals"]["cost_usd"] == 0.001 + assert "Run incomplete" in capsys.readouterr().out + + +def test_final_report_and_update_history_are_one_result_per_run(tmp_path): + instance = adapter(tmp_path, UsageLedger()) + first = UpdateRecord() + instance._last_update_record = first + instance._finish_usage() + instance._finish_usage() + second = UpdateRecord() + instance._last_update_record = second + instance._finish_usage() + metadata = json.loads((tmp_path / "metadata.json").read_text()) + record = json.loads((tmp_path / "update_record.json").read_text()) + assert len(metadata["update_history"]) == 2 + assert first.run_id != second.run_id + assert record["usage_report"] == metadata["usage"] == metadata["last_update"]["usage"] + assert instance.job.statistics.total_tokens_used == 0 + + +def test_inclusive_token_statistics_and_resumed_run(tmp_path): + ledger = UsageLedger() + ledger.add( + "request", + "", + { + "input_tokens": 100, + "output_tokens": 20, + "cache_read_tokens": 40, + "cache_write_tokens": 10, + }, + ) + instance = adapter(tmp_path, ledger) + instance._record_usage_statistics() + assert instance.job.statistics.total_tokens_used == 120 + instance._usage_ledger = UsageLedger() + instance._record_usage_statistics() + assert instance.job.statistics.total_tokens_used == 0 + assert instance.job.statistics.total_cost_usd == 0 + + +def test_full_fallback_keeps_attempt_costs_and_prior_history(tmp_path): + output = tmp_path / "docs" + output.mkdir() + (output / "metadata.json").write_text(json.dumps({"update_history": [{"run_id": "old"}]})) + ledger = UsageLedger() + instance = adapter(output, ledger) + instance._last_update_record = UpdateRecord(outcome="full_fallback") + instance.commit_id = "new" + instance.verbose = False + with ledger.in_phase("update"): + ledger.add("routing", "", {"input_tokens": 100, "cost_usd": 0.001}, model="cluster") + instance._move_docs_aside() + with ledger.in_phase("modules"): + ledger.add("module_agent", "", {"input_tokens": 200, "cost_usd": 0.002}, model="main") + instance._merge_update_summary(str(output)) + instance._finish_usage() + metadata = json.loads((output / "metadata.json").read_text()) + assert metadata["usage"]["totals"]["cost_usd"] == pytest.approx(0.003) + assert metadata["usage"]["n_calls"] == 2 + assert len(metadata["update_history"]) == 2 + assert metadata["update_history"][0]["run_id"] == "old" + assert metadata["last_update"]["usage"] == metadata["usage"] + assert metadata["last_update"]["previous_docs"] == instance._preserved_docs_dir + + +def test_direct_generation_failure_saves_available_usage(tmp_path): + from codewiki.src.be.documentation_generator import DocumentationGenerator + + cfg = config(tmp_path) + backend = PydanticAIBackend(cfg) + backend.usage_ledger.add("request", "", {"input_tokens": 10, "cost_usd": 0.0001}) + generator = object.__new__(DocumentationGenerator) + generator.config = cfg + generator.backend = backend + + def fail(): + raise RuntimeError("analysis failed") + + generator.graph_builder = SimpleNamespace(build_dependency_graph=fail) + with pytest.raises(RuntimeError, match="analysis failed"): + asyncio.run(generator.run()) + metadata = json.loads((tmp_path / "docs" / "metadata.json").read_text()) + assert metadata["usage"]["incomplete"] is True + assert metadata["usage"]["totals"]["cost_usd"] == 0.0001 diff --git a/tests/test_usage_ledger.py b/tests/test_usage_ledger.py index afe809c1..2fe678c4 100644 --- a/tests/test_usage_ledger.py +++ b/tests/test_usage_ledger.py @@ -62,8 +62,9 @@ def test_tokens_without_cost_are_unpriced(): ledger = UsageLedger() ledger.add("module_agent", "a", {"input_tokens": 10, "cost_usd": 0.0}, model="gpt-x") ledger.add("module_agent", "b", {"input_tokens": 10}, model="") - ledger.add("module_agent", "c", None, model="ignored") # no usage reported at all - assert ledger.summary()["unpriced_models"] == ["", "gpt-x"] + ledger.add("module_agent", "c", None, model="missing-usage") + assert ledger.summary()["unpriced_models"] == ["", "gpt-x", "missing-usage"] + assert ledger.summary()["unknown_usage_calls"] == 1 def test_add_is_thread_safe(): @@ -121,7 +122,8 @@ def test_caw_complete_records_usage(monkeypatch): assert call["target"] == "repo" assert call["kind"] == "complete" assert call["model"] == "opus" - assert call["input_tokens"] == 1000 + assert call["input_tokens"] == 1050 # normalized inclusive input, including cache reads + assert call["total_tokens"] == 1250 assert call["cost_usd"] == 0.12