diff --git a/demos/06_chat_transcript_cost.py b/demos/06_chat_transcript_cost.py new file mode 100644 index 0000000..f7fb028 --- /dev/null +++ b/demos/06_chat_transcript_cost.py @@ -0,0 +1,45 @@ +"""Scenario 6 - conversational / chatbot products. + +A chat product pays for the *whole* transcript on every turn, because each new +user message re-sends the accumulated history. This demo measures a real +multi-turn conversation, then shows the compounding cost of carrying full +history vs. a windowed / summarized context. + +Audience: teams shipping chat assistants who need to reason about history cost. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate + + +def main() -> None: + rule("CHAT TRANSCRIPT COST - you re-pay for history on every turn") + + convo = read_fixture("05-chat-transcript", "conversation.txt") + model = "claude-haiku" + + # Full transcript resent each turn, ~150-token reply. + full = estimate(convo, model=model, output_tokens=150) + d = full.to_dict() + print(f"\nFull transcript on {model} (150-token reply):") + print(f" history_tokens={d['input_tokens']} per_turn={usd(d['total_cost_usd'])}") + + # A 10-turn session where history grows linearly: turn k pays ~k/N of the + # full history. Approximate the running cost. + turns = 10 + running = sum( + estimate(convo, model=model, input_tokens=int(full.input_tokens * (k + 1) / turns), + output_tokens=150).total_cost + for k in range(turns) + ) + print(f" {turns}-turn session (history grows each turn): {usd(running)}") + + # Windowed alternative: cap history at a third of the transcript. + windowed = estimate(model=model, input_tokens=full.input_tokens // 3, output_tokens=150) + saved = full.total_cost - windowed.total_cost + print(f"\nCap history to 1/3 the transcript: {usd(windowed.total_cost)}/turn " + f"({usd(saved)} saved/turn).") + print("Windowing or summarizing history is the biggest lever on chat cost.") + + +if __name__ == "__main__": + main() diff --git a/demos/07_context_window_guard.py b/demos/07_context_window_guard.py new file mode 100644 index 0000000..cf40337 --- /dev/null +++ b/demos/07_context_window_guard.py @@ -0,0 +1,35 @@ +"""Scenario 7 - platform guardrails. + +Before you send an assembled prompt you want to know it will *fit* the model's +context window — a request that overflows is rejected by the API after you have +already paid to build it. This demo takes a document dump and checks it against +each model's window, showing which models can take it and which cannot. + +Audience: platform engineers adding pre-flight window checks. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, check_budget, list_models + + +def main() -> None: + rule("CONTEXT WINDOW GUARD - will this prompt even fit?") + + dump = read_fixture("06-context-window-guard", "document_dump.txt") + print("\nA large document dump checked against every model's window:\n") + print(f" {'model':<14}{'window':>10}{'used_%':>10} fits?") + for m in list_models(): + est = estimate(dump, model=m.name, output_tokens=500) + res = check_budget(est) + fits = "yes" if res.ok else "NO" + d = est.to_dict() + print(f" {m.name:<14}{m.context_window:>10}{d['context_used_pct']:>10} {fits}") + + smallest = min(list_models(), key=lambda m: m.context_window) + est = estimate(dump, model=smallest.name, output_tokens=500) + print(f"\nSmallest window ({smallest.name}, {smallest.context_window}) " + f"needs {est.input_tokens + est.output_tokens} tokens — " + f"pre-flight the check, don't pay to be rejected. cost would be {usd(est.total_cost)}.") + + +if __name__ == "__main__": + main() diff --git a/demos/08_diff_review_cost.py b/demos/08_diff_review_cost.py new file mode 100644 index 0000000..9d38405 --- /dev/null +++ b/demos/08_diff_review_cost.py @@ -0,0 +1,35 @@ +"""Scenario 8 - AI code review in CI. + +An automated PR reviewer sends the diff to an LLM on every push. Big diffs cost +more and can blow the window. This demo prices a real diff and projects the +monthly bill for a team's PR volume, then shows the cheap-model alternative. + +Audience: dev-tooling teams wiring LLM code review into pipelines. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, compare_models + + +def main() -> None: + rule("DIFF REVIEW COST - what does an AI PR reviewer cost per month?") + + diff = read_fixture("09-stdin-pipeline", "sample_diff.txt") + print("\nOne PR diff reviewed (~600-token review comment):\n") + + ranking = compare_models(diff, output_tokens=600, + models=["claude-haiku", "claude-sonnet", "gpt-4o", "claude-opus"]) + for e in ranking: + d = e.to_dict() + print(f" {d['model']:<14} per_review={usd(d['total_cost_usd'])} in={d['input_tokens']}") + + prs_per_day, workdays = 40, 21 + monthly = prs_per_day * workdays + cheap, dear = ranking[0], ranking[-1] + print(f"\nAt {prs_per_day} PRs/day x {workdays} workdays = {monthly} reviews/month:") + print(f" {cheap.model}: {usd(cheap.total_cost * monthly)}/mo " + f"{dear.model}: {usd(dear.total_cost * monthly)}/mo") + print("Route routine reviews to the cheap model; reserve the expensive one for flagged PRs.") + + +if __name__ == "__main__": + main() diff --git a/demos/09_fewshot_tax.py b/demos/09_fewshot_tax.py new file mode 100644 index 0000000..9e289ce --- /dev/null +++ b/demos/09_fewshot_tax.py @@ -0,0 +1,36 @@ +"""Scenario 9 - prompt engineering trade-offs. + +Few-shot examples improve quality but you pay for them on *every* call, forever. +This demo quantifies the recurring "few-shot tax" against the zero-shot baseline +and shows the break-even volume where trimming examples pays for itself. + +Audience: prompt engineers deciding how many examples to carry. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate + + +def main() -> None: + rule("FEW-SHOT TAX - the recurring cost of carrying examples") + + fewshot = read_fixture("07-fewshot-vs-zeroshot", "fewshot.txt") + zeroshot = read_fixture("07-fewshot-vs-zeroshot", "zeroshot.txt") + model = "claude-sonnet" + + fs = estimate(fewshot, model=model, output_tokens=80) + zs = estimate(zeroshot, model=model, output_tokens=80) + delta = fs.total_cost - zs.total_cost + + print(f"\nOn {model} (80-token answer):") + print(f" zero-shot: {zs.input_tokens:>4} tok {usd(zs.total_cost)}") + print(f" few-shot: {fs.input_tokens:>4} tok {usd(fs.total_cost)}") + print(f" per-call few-shot premium: {usd(delta)}") + + for volume in (10_000, 100_000, 1_000_000): + print(f" at {volume:>9,} calls: ${delta * volume:,.2f} extra") + print("If few-shot doesn't lift quality enough to justify that tax, move examples " + "to a fine-tune or a retrieval step.") + + +if __name__ == "__main__": + main() diff --git a/demos/10_classifier_router.py b/demos/10_classifier_router.py new file mode 100644 index 0000000..6d57aa7 --- /dev/null +++ b/demos/10_classifier_router.py @@ -0,0 +1,39 @@ +"""Scenario 10 - model routing / cascades. + +Many workloads (classification, routing, intent detection) are cheap enough to +run on a small model, escalating to a big model only on low confidence. This +demo prices a classifier prompt on a small vs large model and computes the +blended cost of a cascade that escalates 15% of the time. + +Audience: engineers designing model cascades / routers. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate + + +def main() -> None: + rule("CLASSIFIER ROUTER - small model first, escalate the hard 15%") + + prompt = read_fixture("03-model-selection", "classifier_prompt.txt") + small = estimate(prompt, model="gpt-4o-mini", output_tokens=20) + large = estimate(prompt, model="claude-opus", output_tokens=20) + + print(f"\nClassifier prompt (20-token label):") + print(f" small (gpt-4o-mini): {usd(small.total_cost)}") + print(f" large (claude-opus): {usd(large.total_cost)}") + + escalate = 0.15 + # cascade pays the small model always, and the large model on escalations + blended = small.total_cost + escalate * large.total_cost + always_large = large.total_cost + saved_pct = (1 - blended / always_large) * 100 if always_large else 0 + print(f"\nCascade (small always + large on {escalate:.0%}): {usd(blended)}/call") + print(f" vs always-large {usd(always_large)}/call -> {saved_pct:.0f}% cheaper") + + calls = 5_000_000 + print(f" over {calls:,} calls: ${blended * calls:,.2f} vs ${always_large * calls:,.2f}") + print("Cascades win when the small model handles the easy majority correctly.") + + +if __name__ == "__main__": + main() diff --git a/demos/11_agent_step_budget.py b/demos/11_agent_step_budget.py new file mode 100644 index 0000000..036d7a9 --- /dev/null +++ b/demos/11_agent_step_budget.py @@ -0,0 +1,38 @@ +"""Scenario 11 - autonomous agents. + +An agent loops: think -> act -> observe, many LLM calls per task. Runaway loops +are a real cost incident. This demo prices one agent step and enforces a +per-task budget gate that would halt the agent before it overspends. + +Audience: teams running autonomous / tool-using agents. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, check_budget + + +def main() -> None: + rule("AGENT STEP BUDGET - cap the loop before it runs away") + + step = read_fixture("08-csv-finops", "agent_step_prompt.txt") + model = "claude-sonnet" + per_step = estimate(step, model=model, output_tokens=120) + print(f"\nOne agent step on {model}: {usd(per_step.total_cost)} " + f"({per_step.input_tokens} in / {per_step.output_tokens} out)") + + max_task_cost = 0.05 + max_steps = int(max_task_cost / per_step.total_cost) if per_step.total_cost else 0 + print(f"\nPer-task budget {usd(max_task_cost)} => at most ~{max_steps} steps.") + + for n_steps in (10, max_steps, max_steps + 20): + projected = estimate(model=model, + input_tokens=per_step.input_tokens * n_steps, + output_tokens=per_step.output_tokens * n_steps) + res = check_budget(projected, max_cost_usd=max_task_cost) + status = "OK" if res.ok else "HALT" + print(f" {n_steps:>3} steps -> {usd(projected.total_cost)} [{status}]") + + print("Gate cumulative task cost each step; halt the agent when the ceiling is hit.") + + +if __name__ == "__main__": + main() diff --git a/demos/12_prompt_diet.py b/demos/12_prompt_diet.py new file mode 100644 index 0000000..9dc4dac --- /dev/null +++ b/demos/12_prompt_diet.py @@ -0,0 +1,38 @@ +"""Scenario 12 - prompt optimization ("prompt diet"). + +Trimming boilerplate from a system prompt cuts cost on every single call. This +demo measures a verbose prompt, a trimmed version, and reports the savings both +per-call and at scale — the ROI of the editing work. + +Audience: anyone optimizing a high-volume prompt. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, count_tokens + + +def main() -> None: + rule("PROMPT DIET - measure the ROI of trimming a prompt") + + verbose = read_fixture("01-basic", "prompt.txt") + # A trimmed variant: drop the trailing half (simulates removing boilerplate). + lines = verbose.splitlines() + trimmed = "\n".join(lines[: max(1, len(lines) // 2)]) + + model = "claude-sonnet" + v = estimate(verbose, model=model, output_tokens=200) + t = estimate(trimmed, model=model, output_tokens=200) + saved = v.total_cost - t.total_cost + + print(f"\nOn {model} (200-token answer):") + print(f" verbose: {count_tokens(verbose):>4} tok {usd(v.total_cost)}") + print(f" trimmed: {count_tokens(trimmed):>4} tok {usd(t.total_cost)}") + print(f" saved per call: {usd(saved)}") + + daily = 200_000 + print(f"\nAt {daily:,} calls/day the trim saves ${saved * daily:,.2f}/day " + f"(${saved * daily * 30:,.2f}/mo).") + print("A 10-minute edit can pay for itself many times over at volume.") + + +if __name__ == "__main__": + main() diff --git a/demos/13_rag_chunk_tuning.py b/demos/13_rag_chunk_tuning.py new file mode 100644 index 0000000..2fe01dd --- /dev/null +++ b/demos/13_rag_chunk_tuning.py @@ -0,0 +1,33 @@ +"""Scenario 13 - RAG retriever tuning (top-k sweep). + +The retriever's top-k directly drives context size and cost. This demo sweeps k +by measuring what a growing slice of retrieved context costs and where it starts +eating dangerous fractions of the window. + +Audience: RAG engineers picking a top-k / chunk budget. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, get_pricing + + +def main() -> None: + rule("RAG CHUNK TUNING - sweep top-k against cost and window") + + context = read_fixture("02-rag-context", "retrieved_context.txt") + model = "claude-sonnet" + full_tokens = estimate(context, model=model).input_tokens + window = get_pricing(model).context_window + + print(f"\nRetrieved-context budget on {model} (window {window:,}), 300-token answer:\n") + print(f" {'k (fraction)':<14}{'in_tok':>8}{'cost':>14}{'window_%':>12}") + for frac in (0.25, 0.5, 0.75, 1.0, 1.5): + in_tok = int(full_tokens * frac) + est = estimate(model=model, input_tokens=in_tok, output_tokens=300) + d = est.to_dict() + print(f" {frac:<14}{in_tok:>8}{usd(d['total_cost_usd']):>14}{d['context_used_pct']:>12}") + + print("\nPick the smallest k that answers well; each extra chunk is paid on every query.") + + +if __name__ == "__main__": + main() diff --git a/demos/14_multi_model_portfolio.py b/demos/14_multi_model_portfolio.py new file mode 100644 index 0000000..166cf5e --- /dev/null +++ b/demos/14_multi_model_portfolio.py @@ -0,0 +1,39 @@ +"""Scenario 14 - portfolio cost allocation. + +A platform runs several workloads on different models. This demo builds a +portfolio (each workload = a prompt fixture + a model + a daily volume) and rolls +up the projected daily and monthly spend, so finance sees one defensible number. + +Audience: platform owners doing chargeback / showback. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate + + +WORKLOADS = [ + ("support-triage", ("01-basic", "prompt.txt"), "claude-haiku", 200_000, 300), + ("rag-answers", ("02-rag-context", "retrieved_context.txt"), "claude-sonnet", 50_000, 300), + ("agent-planning", ("08-csv-finops", "agent_step_prompt.txt"), "gpt-4o", 400_000, 120), + ("code-review", ("09-stdin-pipeline", "sample_diff.txt"), "claude-opus", 5_000, 600), +] + + +def main() -> None: + rule("MULTI-MODEL PORTFOLIO - one rollup for finance") + + print(f"\n {'workload':<18}{'model':<14}{'calls/day':>11}{'$/day':>14}") + total_day = 0.0 + for name, fx, model, volume, out in WORKLOADS: + est = estimate(read_fixture(*fx), model=model, output_tokens=out) + day = est.total_cost * volume + total_day += day + print(f" {name:<18}{model:<14}{volume:>11,}{usd(day):>14}") + + print(f" {'-' * 55}") + print(f" {'TOTAL':<18}{'':<14}{'':>11}{usd(total_day):>14}") + print(f"\nProjected spend: ${total_day:,.2f}/day -> ${total_day * 30:,.2f}/mo. " + "Re-run when a prompt or volume changes to catch cost drift early.") + + +if __name__ == "__main__": + main() diff --git a/demos/15_streaming_output_forecast.py b/demos/15_streaming_output_forecast.py new file mode 100644 index 0000000..49f8121 --- /dev/null +++ b/demos/15_streaming_output_forecast.py @@ -0,0 +1,36 @@ +"""Scenario 15 - output-length sensitivity. + +Output tokens are usually the most expensive tokens (output price > input price). +This demo holds the prompt fixed and sweeps expected output length, showing how +much a chatty model or an unbounded max_tokens costs you. + +Audience: engineers setting max_tokens / controlling verbosity. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, get_pricing + + +def main() -> None: + rule("OUTPUT LENGTH FORECAST - output tokens are the pricey ones") + + prompt = read_fixture("01-basic", "prompt.txt") + model = "claude-opus" + p = get_pricing(model) + print(f"\n{model}: input ${p.input_per_1k}/1k, output ${p.output_per_1k}/1k " + f"(output is {p.output_per_1k / p.input_per_1k:.0f}x input)\n") + + print(f" {'max output tok':<16}{'total cost':>14}") + base = None + for out in (100, 500, 1000, 2000, 4000): + est = estimate(prompt, model=model, output_tokens=out) + if base is None: + base = est.total_cost + print(f" {out:<16}{usd(est.total_cost):>14}") + + big = estimate(prompt, model=model, output_tokens=4000) + print(f"\nGoing from 100 to 4000 output tokens is {big.total_cost / base:.1f}x the cost. " + "Cap max_tokens and ask for concise answers where you can.") + + +if __name__ == "__main__": + main() diff --git a/demos/16_batch_library_audit.py b/demos/16_batch_library_audit.py new file mode 100644 index 0000000..bfbea78 --- /dev/null +++ b/demos/16_batch_library_audit.py @@ -0,0 +1,46 @@ +"""Scenario 16 - prompt-library audit with a per-prompt ceiling. + +An engineering org enforces a per-prompt token ceiling in review. This demo +audits the whole prompt library, flags any prompt over the ceiling, and returns +the same aggregate the CLI `batch` command produces — a governance gate. + +Audience: platform / governance teams enforcing prompt hygiene. +""" +import os + +from _common import rule, fixture, usd +from tokenmeter.core import estimate, aggregate, check_budget + + +def main() -> None: + rule("PROMPT LIBRARY AUDIT - flag prompts over the token ceiling") + + lib = fixture("04-batch-prompt-library", "prompts") + paths = sorted(os.path.join(lib, f) for f in os.listdir(lib) if f.endswith(".txt")) + model = "claude-sonnet" + ceiling = 60 # per-prompt input-token ceiling for this fixture set + + ests = [] + flagged = [] + print(f"\nAuditing {len(paths)} prompts on {model} (ceiling {ceiling} tok):\n") + for p in paths: + with open(p, "r", encoding="utf-8", errors="replace") as fh: + est = estimate(fh.read(), model=model) + ests.append(est) + res = check_budget(est, max_tokens=ceiling) + mark = "OK " if res.ok else "OVER" + if not res.ok: + flagged.append(os.path.basename(p)) + print(f" [{mark}] {os.path.basename(p):<24} {est.input_tokens:>4} tok {usd(est.total_cost)}") + + roll = aggregate(ests) + print(f"\n library total: {roll['total_tokens']} tok {usd(roll['total_cost_usd'])}") + if flagged: + print(f" FLAGGED (trim before merge): {', '.join(flagged)}") + else: + print(" all prompts within ceiling.") + print("Wire this as a review gate so oversized prompts never land silently.") + + +if __name__ == "__main__": + main() diff --git a/demos/17_custom_model_pricing.py b/demos/17_custom_model_pricing.py new file mode 100644 index 0000000..bf026c1 --- /dev/null +++ b/demos/17_custom_model_pricing.py @@ -0,0 +1,39 @@ +"""Scenario 17 - self-hosted / custom model pricing. + +Running your own model (or a vendor not in the table) means you set the price. +This demo registers a custom "$/1k" for a self-hosted deployment via add_model, +then compares its effective cost against the hosted options — the classic +build-vs-buy break-even. + +Audience: infra teams pricing a self-hosted model against APIs. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import add_model, estimate, compare_models + + +def main() -> None: + rule("CUSTOM MODEL PRICING - price your self-hosted model in the same table") + + # Amortized GPU cost expressed as $/1k tokens (illustrative). + add_model("self-hosted-8b", input_per_1k=0.0002, output_per_1k=0.0002, context_window=32_000) + print("\nRegistered 'self-hosted-8b' at $0.0002/1k in+out, 32k window.") + + prompt = read_fixture("02-rag-context", "retrieved_context.txt") + ranking = compare_models( + prompt, output_tokens=300, + models=["self-hosted-8b", "gpt-4o-mini", "claude-haiku", "claude-sonnet"], + ) + print("\nSame RAG workload, cheapest first:\n") + for e in ranking: + d = e.to_dict() + print(f" {d['model']:<18} {usd(d['total_cost_usd'])} (ctx {d['context_used_pct']}%)") + + sh = next(e for e in ranking if e.model == "self-hosted-8b") + hosted = next(e for e in ranking if e.model == "claude-haiku") + per_call_saving = hosted.total_cost - sh.total_cost + print(f"\nPer-call saving vs claude-haiku: {usd(per_call_saving)}. " + "Divide your fixed GPU spend by that to get the break-even call volume.") + + +if __name__ == "__main__": + main() diff --git a/demos/18_json_output_pipeline.py b/demos/18_json_output_pipeline.py new file mode 100644 index 0000000..a481397 --- /dev/null +++ b/demos/18_json_output_pipeline.py @@ -0,0 +1,36 @@ +"""Scenario 18 - machine-readable output for pipelines. + +tokenmeter's estimates are plain dicts, so they drop into any pipeline as JSON. +This demo drives the same API the CLI's `--format json` path uses and prints a +parseable record other tools can consume (dashboards, alerting, chargeback). + +Audience: engineers integrating cost telemetry into other systems. +""" +import json + +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, check_budget + + +def main() -> None: + rule("JSON OUTPUT PIPELINE - emit a parseable cost record") + + prompt = read_fixture("01-basic", "prompt.txt") + est = estimate(prompt, model="claude-sonnet", output_tokens=400) + res = check_budget(est, max_cost_usd=0.02, max_tokens=5000) + + record = res.to_dict() + blob = json.dumps(record, indent=2, sort_keys=True) + print("\nEmitted record (pipe to jq, a dashboard, or an alerting rule):\n") + print(blob) + + # Round-trip proves it is valid, parseable JSON. + parsed = json.loads(blob) + assert parsed["estimate"]["model"] == "claude-sonnet" + print(f"\nParsed back OK: cost {usd(parsed['estimate']['total_cost_usd'])}, " + f"budget_ok={parsed['ok']}.") + print("Everything tokenmeter returns is a dict -> JSON, so it composes anywhere.") + + +if __name__ == "__main__": + main() diff --git a/demos/19_regression_guard.py b/demos/19_regression_guard.py new file mode 100644 index 0000000..7fa2301 --- /dev/null +++ b/demos/19_regression_guard.py @@ -0,0 +1,39 @@ +"""Scenario 19 - prompt cost-regression guard. + +Cost creep is silent: a prompt grows a little each PR until the bill jumps. This +demo compares a stored baseline cost against the current prompt and fails (like a +regression test) if the increase exceeds a tolerance — a diff-time cost gate. + +Audience: teams treating prompt cost like a performance budget. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate + + +def main() -> None: + rule("COST REGRESSION GUARD - fail the PR when a prompt gets fatter") + + prompt = read_fixture("01-basic", "prompt.txt") + model = "claude-sonnet" + current = estimate(prompt, model=model, output_tokens=400) + + # Simulate a stored baseline (what the prompt cost last release). + baseline_tokens = int(current.input_tokens * 0.85) # it grew ~18% since + baseline = estimate(model=model, input_tokens=baseline_tokens, output_tokens=400) + + tolerance = 0.10 # allow 10% growth + growth = (current.total_cost - baseline.total_cost) / baseline.total_cost + print(f"\nOn {model}:") + print(f" baseline: {baseline.input_tokens} tok {usd(baseline.total_cost)}") + print(f" current: {current.input_tokens} tok {usd(current.total_cost)}") + print(f" growth: {growth:+.1%} (tolerance {tolerance:.0%})") + + if growth > tolerance: + print(" -> would FAIL the gate: prompt cost regressed beyond tolerance.") + else: + print(" -> within tolerance.") + print("Store the baseline in the repo; compare on every PR to stop silent drift.") + + +if __name__ == "__main__": + main() diff --git a/demos/20_full_month_forecast.py b/demos/20_full_month_forecast.py new file mode 100644 index 0000000..c4f679f --- /dev/null +++ b/demos/20_full_month_forecast.py @@ -0,0 +1,40 @@ +"""Scenario 20 - end-to-end monthly forecast. + +Ties the pieces together: take a real prompt, price it on a chosen model, apply a +traffic profile (weekday vs weekend), and produce a defensible monthly forecast +with a per-model comparison — the number you take to a budget meeting. + +Audience: anyone owning an LLM cost line and forecasting spend. +""" +from _common import rule, read_fixture, usd +from tokenmeter.core import estimate, compare_models + + +def main() -> None: + rule("FULL MONTH FORECAST - the number for the budget meeting") + + prompt = read_fixture("02-rag-context", "retrieved_context.txt") + weekday_calls, weekend_calls = 300_000, 90_000 + weekdays, weekend_days = 22, 8 + monthly_calls = weekday_calls * weekdays + weekend_calls * weekend_days + + print(f"\nTraffic: {weekday_calls:,}/weekday x {weekdays}d + " + f"{weekend_calls:,}/weekend x {weekend_days}d = {monthly_calls:,} calls/mo\n") + + ranking = compare_models( + prompt, output_tokens=300, + models=["claude-haiku", "gpt-4o-mini", "claude-sonnet", "gpt-4o", "claude-opus"], + ) + print(f" {'model':<14}{'$/call':>12}{'$/month':>16}") + for e in ranking: + monthly = e.total_cost * monthly_calls + print(f" {e.model:<14}{usd(e.total_cost):>12}{('$%s' % format(monthly, ',.2f')):>16}") + + cheap, dear = ranking[0], ranking[-1] + spread = dear.total_cost * monthly_calls - cheap.total_cost * monthly_calls + print(f"\nModel choice alone swings this budget by ${spread:,.2f}/month " + f"({cheap.model} vs {dear.model}). Forecast before you commit.") + + +if __name__ == "__main__": + main() diff --git a/demos/run_all.py b/demos/run_all.py index 4089bc1..986bb69 100644 --- a/demos/run_all.py +++ b/demos/run_all.py @@ -18,6 +18,21 @@ "03_ci_budget_gate", "04_eng_manager_prompt_library", "05_rag_context_budget", + "06_chat_transcript_cost", + "07_context_window_guard", + "08_diff_review_cost", + "09_fewshot_tax", + "10_classifier_router", + "11_agent_step_budget", + "12_prompt_diet", + "13_rag_chunk_tuning", + "14_multi_model_portfolio", + "15_streaming_output_forecast", + "16_batch_library_audit", + "17_custom_model_pricing", + "18_json_output_pipeline", + "19_regression_guard", + "20_full_month_forecast", ] diff --git a/docs/DEMOS.md b/docs/DEMOS.md index a4fa786..95c67bd 100644 --- a/docs/DEMOS.md +++ b/docs/DEMOS.md @@ -1,14 +1,14 @@ # Demos -Five runnable scenarios in [`../demos/`](../demos/), each targeting a different -audience. Every scenario drives the **real** `tokenmeter` API (`tokenmeter.core`) -against the bundled, offline fixtures under +**Twenty** runnable scenarios in [`../demos/`](../demos/), each targeting a +different audience. Every scenario drives the **real** `tokenmeter` API +(`tokenmeter.core`) against the bundled, offline fixtures under [`../tokenmeter/demos/`](../tokenmeter/demos) — no network, no API keys, no heavy deps. Run them in any order or on their own; each exits 0, so they double as smoke tests. ```bash -python demos/run_all.py # all five, end to end +python demos/run_all.py # all twenty, end to end python demos/03_ci_budget_gate.py # or just one ``` @@ -44,8 +44,83 @@ Size the real assembled RAG artifact and watch `context_used_pct`, then quantify the recurring few-shot tax — the per-call premium of carrying examples, scaled to a million calls. +## 6. Chat transcript cost — *you re-pay for history every turn* +**Audience:** teams shipping chat assistants. +Measure a real multi-turn transcript, model the compounding cost of a 10-turn +session as history grows, and show the saving from windowing history to a third. + +## 7. Context window guard — *will this prompt even fit?* +**Audience:** platform engineers adding pre-flight checks. +Check a large document dump against every model's window and print which models +can take it — fail fast before the API rejects a prompt you already paid to build. + +## 8. Diff review cost — *what does an AI PR reviewer cost per month?* +**Audience:** dev-tooling teams. +Price one PR diff review across four models and project the monthly bill at 40 +PRs/day, then show the cheap-model-first routing that keeps it affordable. + +## 9. Few-shot tax — *the recurring cost of carrying examples* +**Audience:** prompt engineers. +Quantify the per-call premium of few-shot vs zero-shot and scale it to 10k / 100k +/ 1M calls, so you can decide whether the quality lift earns the tax. + +## 10. Classifier router — *small model first, escalate the hard 15%* +**Audience:** engineers designing model cascades. +Price a classifier prompt on a small vs large model and compute the blended cost +of a cascade that escalates only low-confidence cases. + +## 11. Agent step budget — *cap the loop before it runs away* +**Audience:** teams running autonomous agents. +Price one agent step, derive the max steps a per-task budget allows, and gate a +projected multi-step run — halting the agent before it overspends. + +## 12. Prompt diet — *the ROI of trimming a prompt* +**Audience:** anyone optimizing a high-volume prompt. +Measure a verbose prompt vs a trimmed variant and show the per-call and monthly +savings — the payback on a few minutes of editing. + +## 13. RAG chunk tuning — *sweep top-k against cost and window* +**Audience:** RAG engineers picking a top-k / chunk budget. +Sweep the retrieved-context size and watch cost and `window_%` climb, to pick the +smallest k that still answers well. + +## 14. Multi-model portfolio — *one rollup for finance* +**Audience:** platform owners doing chargeback / showback. +Build a portfolio of workloads (prompt + model + daily volume) and roll up the +projected daily and monthly spend into one defensible number. + +## 15. Output length forecast — *output tokens are the pricey ones* +**Audience:** engineers setting `max_tokens` / controlling verbosity. +Hold the prompt fixed and sweep expected output length to show how a chatty model +or an unbounded `max_tokens` multiplies cost. + +## 16. Prompt library audit — *flag prompts over the token ceiling* +**Audience:** platform / governance teams. +Audit the whole prompt library against a per-prompt token ceiling and flag any +prompt that would need trimming before merge — a governance gate. + +## 17. Custom model pricing — *price your self-hosted model in the same table* +**Audience:** infra teams pricing self-hosted models. +Register a custom `$/1k` via `add_model`, compare it against hosted options, and +derive the build-vs-buy break-even call volume. + +## 18. JSON output pipeline — *emit a parseable cost record* +**Audience:** engineers integrating cost telemetry. +Produce the same JSON record the CLI's `--format json` path emits and round-trip +it, proving everything tokenmeter returns composes into other tools. + +## 19. Cost regression guard — *fail the PR when a prompt gets fatter* +**Audience:** teams treating prompt cost like a performance budget. +Compare a stored baseline cost against the current prompt and fail if growth +exceeds tolerance — a diff-time cost gate that stops silent creep. + +## 20. Full month forecast — *the number for the budget meeting* +**Audience:** anyone owning an LLM cost line. +Apply a weekday/weekend traffic profile to a real prompt and produce a monthly +forecast per model, showing how much model choice alone swings the budget. + --- -Each demo prints clear, narrated output and exits 0. `tests/test_demos.py` -imports and runs every scenario under `pytest`, so the demos are also covered as -tests. +Each demo prints clear, narrated output and exits 0. `tests/test_demos.py` and +`tests/test_demos_behavior.py` import and run every scenario under `pytest`, so +the demos are also covered as tests. diff --git a/tests/test_cli_edge.py b/tests/test_cli_edge.py new file mode 100644 index 0000000..0179459 --- /dev/null +++ b/tests/test_cli_edge.py @@ -0,0 +1,272 @@ +"""CLI edge cases, exit codes, formats, stdin and error paths for tokenmeter.cli. + +Drives ``cli.main(argv)`` directly (no subprocess) and asserts on exit codes and +emitted text. Covers all subcommands (count / budget / models / batch / compare), +all three output formats (table / json / csv), stdin piping, malformed input, and +the documented non-zero exit codes (1 = over budget, 2 = usage/bad input). + +Stdlib + package only. No network. +""" +from __future__ import annotations + +import contextlib +import io +import json +import os +import sys +import tempfile + +import pytest + +from tokenmeter import cli + + +def run(argv, stdin_text=None): + """Invoke the CLI, capturing stdout/stderr and any SystemExit code.""" + out, err = io.StringIO(), io.StringIO() + old_stdin = sys.stdin + if stdin_text is not None: + sys.stdin = io.StringIO(stdin_text) + try: + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + try: + code = cli.main(argv) + except SystemExit as e: # argparse errors raise SystemExit + code = e.code if isinstance(e.code, int) else 2 + finally: + sys.stdin = old_stdin + return code, out.getvalue(), err.getvalue() + + +@pytest.fixture +def tmpfile(): + created = [] + + def _make(content): + fd, path = tempfile.mkstemp(suffix=".txt") + with os.fdopen(fd, "w", encoding="utf-8") as fh: + fh.write(content) + created.append(path) + return path + + yield _make + for p in created: + try: + os.remove(p) + except OSError: + pass + + +# --------------------------------------------------------------------------- # +# count # +# --------------------------------------------------------------------------- # +class TestCount: + def test_count_table_default(self): + code, out, _ = run(["count", "-t", "hello world"]) + assert code == 0 + assert "input_tokens" in out and "total_cost_usd" in out + + def test_count_json(self): + code, out, _ = run(["--format", "json", "count", "-t", "hello", "-m", "gpt-4o"]) + assert code == 0 + d = json.loads(out) + assert d["model"] == "gpt-4o" and d["input_tokens"] > 0 + + def test_count_csv(self): + code, out, _ = run(["count", "-t", "hi", "--format", "csv"]) + assert code == 0 + assert out.startswith("metric,value") + + def test_count_from_file(self, tmpfile): + p = tmpfile("some prompt text here") + code, out, _ = run(["--format", "json", "count", "-f", p]) + assert code == 0 + assert json.loads(out)["input_tokens"] > 0 + + def test_count_from_stdin(self): + code, out, _ = run(["--format", "json", "count"], stdin_text="piped prompt text") + assert code == 0 + assert json.loads(out)["input_tokens"] > 0 + + def test_count_empty_stdin_is_zero(self): + code, out, _ = run(["--format", "json", "count"], stdin_text="") + assert code == 0 + assert json.loads(out)["input_tokens"] == 0 + + def test_count_unknown_model_exit_2(self): + code, _, err = run(["count", "-t", "x", "-m", "nope"]) + assert code == 2 + assert "error" in err.lower() + + def test_count_negative_output_tokens_exit_2(self): + code, _, err = run(["count", "-t", "x", "-o", "-5"]) + assert code == 2 + assert "output_tokens" in err + + def test_count_missing_file_exit_2(self): + code, _, err = run(["count", "-f", os.path.join(tempfile.gettempdir(), "no_such_tm.txt")]) + assert code == 2 + assert "error" in err.lower() + + def test_text_and_file_mutually_exclusive(self, tmpfile): + p = tmpfile("x") + code, _, _ = run(["count", "-t", "hi", "-f", p]) + assert code == 2 # argparse mutually-exclusive error + + +# --------------------------------------------------------------------------- # +# budget # +# --------------------------------------------------------------------------- # +class TestBudget: + def test_pass_exit_0(self): + code, out, _ = run(["budget", "-t", "hi", "-m", "gpt-4o-mini", "--max-cost", "1.0"]) + assert code == 0 + assert "OK" in out + + def test_over_cost_exit_1(self): + code, out, _ = run( + ["budget", "-m", "claude-opus", "-t", "x", "-o", "1", "--max-cost", "0.0"] + ) + assert code == 1 + assert "OVER BUDGET" in out + + def test_over_tokens_exit_1(self): + code, out, _ = run(["budget", "-t", "a b c d e f g", "--max-tokens", "1"]) + assert code == 1 + + def test_json_output(self): + code, out, _ = run(["--format", "json", "budget", "-t", "hi", "--max-cost", "1.0"]) + assert code == 0 + d = json.loads(out) + assert d["ok"] is True and "estimate" in d + + def test_context_overflow_exit_1(self): + # 8k window model, force > window via explicit big text is hard; use a + # large file instead. + big = " ".join(["word"] * 10000) + code, out, _ = run(["budget", "-t", big, "-m", "generic-1k"]) + assert code == 1 + assert "context window" in out + + def test_unknown_model_exit_2(self): + code, _, err = run(["budget", "-t", "x", "-m", "ghost", "--max-cost", "1"]) + assert code == 2 + + +# --------------------------------------------------------------------------- # +# models # +# --------------------------------------------------------------------------- # +class TestModels: + def test_table(self): + code, out, _ = run(["models"]) + assert code == 0 + assert "claude-sonnet" in out + + def test_json(self): + code, out, _ = run(["--format", "json", "models"]) + assert code == 0 + d = json.loads(out) + assert any(m["name"] == "claude-opus" for m in d["models"]) + + def test_csv_header(self): + code, out, _ = run(["models", "--format", "csv"]) + assert code == 0 + assert out.startswith("model,in/1k,out/1k,context") + + +# --------------------------------------------------------------------------- # +# batch # +# --------------------------------------------------------------------------- # +class TestBatch: + def test_two_files_rollup(self, tmpfile): + a, b = tmpfile("hello world"), tmpfile("another prompt here") + code, out, _ = run(["--format", "json", "batch", a, b]) + assert code == 0 + d = json.loads(out) + assert d["rollup"]["files"] == 2 + assert len(d["files"]) == 2 + + def test_batch_table_has_total(self, tmpfile): + a = tmpfile("hi") + code, out, _ = run(["batch", a]) + assert code == 0 + assert "TOTAL" in out + + def test_batch_max_cost_exceeded_exit_1(self, tmpfile): + a = tmpfile(" ".join(["word"] * 500)) + code, _, err = run(["batch", a, "--max-cost", "0.0", "-m", "claude-opus"]) + assert code == 1 + assert "exceeds" in err + + def test_batch_missing_file_exit_2(self): + missing = os.path.join(tempfile.gettempdir(), "definitely_missing_tm.txt") + code, _, err = run(["batch", missing]) + assert code == 2 + assert "cannot read" in err + + def test_batch_requires_at_least_one_file(self): + code, _, _ = run(["batch"]) + assert code == 2 # argparse: nargs="+" requires one + + +# --------------------------------------------------------------------------- # +# compare # +# --------------------------------------------------------------------------- # +class TestCompare: + def test_json_ranked(self): + code, out, _ = run(["compare", "-t", "hello world", "-o", "100", "--format", "json"]) + assert code == 0 + costs = [r["total_cost_usd"] for r in json.loads(out)["ranking"]] + assert costs == sorted(costs) + + def test_subset_csv(self): + code, out, _ = run( + ["compare", "-t", "hi", "-o", "50", "--models", + "claude-opus,gpt-4o-mini", "--format", "csv"] + ) + assert code == 0 + lines = [l for l in out.strip().splitlines() if l] + assert lines[0].startswith("model,in_tok") + assert len(lines) == 3 # header + 2 + + def test_unknown_model_exit_2(self): + code, _, err = run(["compare", "-t", "x", "--models", "nope"]) + assert code == 2 + + def test_from_stdin(self): + code, out, _ = run(["compare", "--format", "json"], stdin_text="piped text here") + assert code == 0 + assert json.loads(out)["ranking"] + + def test_table_output(self): + code, out, _ = run(["compare", "-t", "hello", "-o", "10"]) + assert code == 0 + assert "claude-opus" in out and "gpt-4o-mini" in out + + +# --------------------------------------------------------------------------- # +# global / parser # +# --------------------------------------------------------------------------- # +class TestParser: + def test_no_command_errors(self): + code, _, _ = run([]) + assert code == 2 # subparser required + + def test_version(self): + code, out, _ = run(["--version"]) + assert code == 0 + assert "tokenmeter" in out + + def test_help_exits_zero(self): + code, _, _ = run(["--help"]) + assert code == 0 + + def test_format_before_and_after_subcommand(self): + # sub-level --format wins over top-level + code, out, _ = run(["--format", "csv", "count", "-t", "hi", "--format", "json"]) + assert code == 0 + json.loads(out) # must be JSON, not CSV + + def test_bad_format_choice_exit_2(self): + code, _, _ = run(["count", "-t", "hi", "--format", "xml"]) + assert code == 2 diff --git a/tests/test_core_edge.py b/tests/test_core_edge.py new file mode 100644 index 0000000..29c311d --- /dev/null +++ b/tests/test_core_edge.py @@ -0,0 +1,357 @@ +"""Edge-case and error-path coverage for tokenmeter.core. + +Complements the happy-path smoke suite: these probe the boundaries of the +estimator, the pricing table, budget gates, model comparison, and aggregation +— including the malformed / adversarial inputs that a budgeting tool has to +reject cleanly rather than silently produce a wrong (e.g. negative) number. + +Stdlib + package only. No network. +""" +from __future__ import annotations + +import math + +import pytest + +from tokenmeter.core import ( + MODELS, + ModelPricing, + add_model, + aggregate, + check_budget, + compare_models, + count_tokens, + estimate, + get_pricing, + list_models, +) + + +# --------------------------------------------------------------------------- # +# count_tokens # +# --------------------------------------------------------------------------- # +class TestCountTokens: + def test_empty_is_zero(self): + assert count_tokens("") == 0 + + def test_none_is_zero(self): + # None is tolerated (behaves like empty) — the CLI can hand us None. + assert count_tokens(None) == 0 + + def test_non_string_raises_type_error(self): + with pytest.raises(TypeError): + count_tokens(12345) + + def test_pure_whitespace_no_newline_is_zero(self): + # Spaces/tabs merge into adjacent words; alone they cost nothing. + assert count_tokens(" ") == 0 + assert count_tokens("\t\t") == 0 + + def test_single_newline_counts_one(self): + assert count_tokens("\n") == 1 + + def test_multiple_newlines_count_each(self): + assert count_tokens("\n\n\n") == 3 + + def test_windows_newlines_count(self): + # \r\n contains one \n each -> two newline tokens for two lines. + assert count_tokens("a\r\nb\r\nc") >= 2 + + def test_single_char_word_is_one_token(self): + assert count_tokens("a") == 1 + + def test_four_char_word_is_one_token(self): + assert count_tokens("abcd") == 1 + + def test_five_char_word_is_two_tokens(self): + # ceil(5/4) == 2 + assert count_tokens("abcde") == 2 + + def test_long_word_ceildiv(self): + assert count_tokens("a" * 16) == 4 # ceil(16/4) + assert count_tokens("a" * 17) == 5 # ceil(17/4) + + def test_digits_grouped_by_three(self): + assert count_tokens("1") == 1 + assert count_tokens("123") == 1 + assert count_tokens("1234") == 2 # ceil(4/3) + assert count_tokens("123456") == 2 + + def test_each_symbol_is_a_token(self): + assert count_tokens("!@#") == 3 + + def test_punctuation_spacing(self): + # "a, b. c!" -> a , b . c ! => at least 5 tokens (words + punctuation) + assert count_tokens("a, b. c!") >= 5 + + def test_monotonic_with_length(self): + short = count_tokens("hello") + long = count_tokens("hello world this is a longer sentence") + assert long > short + + def test_unicode_letters_and_emoji(self): + # Should not crash on non-ASCII; produces a positive count. + assert count_tokens("café résumé 北京 🚀") > 0 + + def test_prose_ratio_is_reasonable(self): + text = " ".join(["word"] * 200) + toks = count_tokens(text) + assert 100 <= toks <= 400 + + def test_mixed_alnum_symbols_split(self): + # A code-ish token: identifiers, digits and symbols split on boundaries. + n = count_tokens("foo_bar123 = baz(42);") + assert n > 5 + + +# --------------------------------------------------------------------------- # +# pricing table # +# --------------------------------------------------------------------------- # +class TestPricing: + def test_all_known_models_have_positive_context(self): + for m in list_models(): + assert m.context_window > 0 + assert m.input_per_1k >= 0 + assert m.output_per_1k >= 0 + + def test_list_models_sorted(self): + names = [m.name for m in list_models()] + assert names == sorted(names) + + def test_get_unknown_raises_keyerror_with_hint(self): + with pytest.raises(KeyError) as ei: + get_pricing("nope-not-real") + # the message should list the known models to help the user + assert "known" in str(ei.value) + + def test_pricing_cost_math(self): + p = ModelPricing("x", 0.01, 0.02, 1000) + # 2000 in, 1000 out => 0.02 + 0.02 = 0.04 + assert math.isclose(p.cost(2000, 1000), 0.04, rel_tol=1e-9) + + def test_pricing_zero_tokens_is_zero(self): + p = get_pricing("claude-opus") + assert p.cost(0, 0) == 0.0 + + +class TestAddModel: + def test_add_and_override(self): + add_model("edge-tmp", 0.001, 0.002, 4096) + assert get_pricing("edge-tmp").context_window == 4096 + add_model("edge-tmp", 0.005, 0.006, 8192) + assert get_pricing("edge-tmp").input_per_1k == 0.005 + + def test_blank_name_rejected(self): + with pytest.raises(ValueError): + add_model("", 0.001, 0.002, 4096) + with pytest.raises(ValueError): + add_model(" ", 0.001, 0.002, 4096) + + def test_non_numeric_price_rejected(self): + with pytest.raises(ValueError): + add_model("bad", "abc", 0.002, 4096) + + def test_negative_price_rejected(self): + with pytest.raises(ValueError): + add_model("neg", -0.001, 0.002, 4096) + with pytest.raises(ValueError): + add_model("neg", 0.001, -0.002, 4096) + + def test_nonpositive_context_rejected(self): + with pytest.raises(ValueError): + add_model("zero-ctx", 0.001, 0.002, 0) + with pytest.raises(ValueError): + add_model("neg-ctx", 0.001, 0.002, -100) + + def test_zero_price_allowed(self): + p = add_model("free", 0.0, 0.0, 1024) + assert p.input_per_1k == 0.0 + + +# --------------------------------------------------------------------------- # +# estimate # +# --------------------------------------------------------------------------- # +class TestEstimate: + def test_explicit_input_tokens_override_text(self): + e = estimate("this text is ignored", model="gpt-4o", input_tokens=1000) + assert e.input_tokens == 1000 + + def test_negative_input_tokens_rejected(self): + with pytest.raises(ValueError): + estimate(model="gpt-4o", input_tokens=-1) + + def test_negative_output_tokens_rejected(self): + with pytest.raises(ValueError): + estimate("hi", model="gpt-4o", output_tokens=-5) + + def test_zero_tokens_zero_cost(self): + e = estimate("", model="claude-opus", output_tokens=0) + assert e.total_cost == 0.0 + assert e.input_tokens == 0 + + def test_unknown_model_raises(self): + with pytest.raises(KeyError): + estimate("hi", model="ghost-model") + + def test_context_pct_computed(self): + e = estimate(model="generic-1k", input_tokens=4096, output_tokens=0) + # 4096 / 8192 == 50% + assert math.isclose(e.context_used_pct, 50.0, rel_tol=1e-6) + + def test_context_pct_over_100_when_overflowing(self): + e = estimate(model="generic-1k", input_tokens=16384, output_tokens=0) + assert e.context_used_pct > 100 + + def test_to_dict_keys_and_rounding(self): + d = estimate(model="claude-opus", input_tokens=1000, output_tokens=1000).to_dict() + for k in ( + "model", "input_tokens", "output_tokens", "total_tokens", + "input_cost_usd", "output_cost_usd", "total_cost_usd", + "context_window", "context_used_pct", + ): + assert k in d + assert d["total_tokens"] == 2000 + # opus: 0.015 in + 0.075 out per 1k => 0.015 + 0.075 = 0.09 + assert math.isclose(d["total_cost_usd"], 0.09, rel_tol=1e-6) + + def test_input_cost_matches_pricing(self): + e = estimate(model="gpt-4o", input_tokens=1000, output_tokens=0) + assert math.isclose(e.input_cost, 0.0025, rel_tol=1e-9) + + def test_output_cost_scales(self): + lo = estimate("x", model="gpt-4o", output_tokens=100) + hi = estimate("x", model="gpt-4o", output_tokens=1000) + assert hi.output_cost > lo.output_cost + assert lo.input_tokens == hi.input_tokens + + def test_float_tokens_coerced_to_int(self): + e = estimate(model="gpt-4o", input_tokens=100, output_tokens=0) + assert isinstance(e.input_tokens, int) + + +# --------------------------------------------------------------------------- # +# check_budget # +# --------------------------------------------------------------------------- # +class TestCheckBudget: + def test_pass_within_all_limits(self): + e = estimate("short", model="gpt-4o-mini", output_tokens=10) + r = check_budget(e, max_cost_usd=1.0, max_tokens=1_000_000) + assert r.ok and r.violations == [] + + def test_no_limits_passes_when_in_window(self): + e = estimate("short", model="gpt-4o-mini") + r = check_budget(e) + assert r.ok + + def test_cost_violation(self): + e = estimate(model="claude-opus", input_tokens=100_000, output_tokens=10_000) + r = check_budget(e, max_cost_usd=0.01) + assert not r.ok + assert any("cost" in v for v in r.violations) + + def test_token_violation(self): + e = estimate(model="gpt-4o", input_tokens=5000, output_tokens=0) + r = check_budget(e, max_tokens=1000) + assert not r.ok + assert any("tokens" in v and "budget" in v for v in r.violations) + + def test_context_window_violation_without_ceilings(self): + e = estimate(model="gpt-4o", input_tokens=200_000, output_tokens=0) + r = check_budget(e) + assert not r.ok + assert any("context window" in v for v in r.violations) + + def test_multiple_violations_reported(self): + e = estimate(model="generic-1k", input_tokens=100_000, output_tokens=0) + r = check_budget(e, max_cost_usd=0.0, max_tokens=1) + # cost + token + context window all trip + assert len(r.violations) == 3 + + def test_boundary_exactly_at_limit_passes(self): + # tokens exactly equal to max_tokens is NOT a violation (strictly >) + e = estimate(model="gpt-4o", input_tokens=100, output_tokens=0) + r = check_budget(e, max_tokens=100) + assert r.ok + + def test_to_dict_shape(self): + e = estimate("hi", model="gpt-4o") + d = check_budget(e, max_cost_usd=1.0).to_dict() + assert set(d) >= {"ok", "max_cost_usd", "max_tokens", "violations", "estimate"} + assert isinstance(d["estimate"], dict) + + +# --------------------------------------------------------------------------- # +# compare_models # +# --------------------------------------------------------------------------- # +class TestCompareModels: + def test_ranked_cheapest_first(self): + ests = compare_models("hello world prompt", output_tokens=500) + costs = [e.total_cost for e in ests] + assert costs == sorted(costs) + + def test_covers_all_models_by_default(self): + ests = compare_models("x", output_tokens=0) + assert len(ests) == len(MODELS) + + def test_same_input_tokens_across_all(self): + ests = compare_models("the quick brown fox jumps", output_tokens=0) + assert len({e.input_tokens for e in ests}) == 1 + + def test_subset_only(self): + ests = compare_models("x y", output_tokens=10, models=["claude-opus", "gpt-4o-mini"]) + assert {e.model for e in ests} == {"claude-opus", "gpt-4o-mini"} + assert ests[0].model == "gpt-4o-mini" # cheaper first + + def test_empty_models_returns_empty(self): + assert compare_models("x", models=[]) == [] + + def test_unknown_model_in_subset_raises(self): + with pytest.raises(KeyError): + compare_models("x", models=["gpt-4o", "does-not-exist"]) + + def test_explicit_input_tokens(self): + ests = compare_models(input_tokens=1000, output_tokens=1000, models=["gpt-4o"]) + assert ests[0].input_tokens == 1000 + + def test_tie_break_by_name(self): + add_model("tie-a", 0.001, 0.001, 4096) + add_model("tie-b", 0.001, 0.001, 4096) + ests = compare_models(input_tokens=100, output_tokens=100, models=["tie-b", "tie-a"]) + # equal cost -> sorted by model name ascending + assert [e.model for e in ests] == ["tie-a", "tie-b"] + + +# --------------------------------------------------------------------------- # +# aggregate # +# --------------------------------------------------------------------------- # +class TestAggregate: + def test_empty_rollup(self): + r = aggregate([]) + assert r["files"] == 0 + assert r["total_tokens"] == 0 + assert r["total_cost_usd"] == 0 + + def test_single(self): + e = estimate(model="gpt-4o", input_tokens=1000, output_tokens=500) + r = aggregate([e]) + assert r["files"] == 1 + assert r["input_tokens"] == 1000 + assert r["output_tokens"] == 500 + assert r["total_tokens"] == 1500 + + def test_sum_of_many(self): + ests = [estimate("a b c", model="gpt-4o") for _ in range(5)] + r = aggregate(ests) + assert r["files"] == 5 + assert r["total_tokens"] == 5 * ests[0].input_tokens + + def test_accepts_generator(self): + gen = (estimate("x", model="gpt-4o") for _ in range(3)) + r = aggregate(gen) + assert r["files"] == 3 + + def test_cost_rounded_six_places(self): + ests = [estimate(model="claude-opus", input_tokens=333, output_tokens=333)] + r = aggregate(ests) + # value should be a float rounded to <= 6 decimals + assert round(r["total_cost_usd"], 6) == r["total_cost_usd"] diff --git a/tests/test_demos.py b/tests/test_demos.py index 0f2caf4..f6fce0b 100644 --- a/tests/test_demos.py +++ b/tests/test_demos.py @@ -21,6 +21,21 @@ "03_ci_budget_gate", "04_eng_manager_prompt_library", "05_rag_context_budget", + "06_chat_transcript_cost", + "07_context_window_guard", + "08_diff_review_cost", + "09_fewshot_tax", + "10_classifier_router", + "11_agent_step_budget", + "12_prompt_diet", + "13_rag_chunk_tuning", + "14_multi_model_portfolio", + "15_streaming_output_forecast", + "16_batch_library_audit", + "17_custom_model_pricing", + "18_json_output_pipeline", + "19_regression_guard", + "20_full_month_forecast", ] diff --git a/tests/test_demos_behavior.py b/tests/test_demos_behavior.py new file mode 100644 index 0000000..0aef8e1 --- /dev/null +++ b/tests/test_demos_behavior.py @@ -0,0 +1,81 @@ +"""Behavioral assertions for the demo scenarios beyond 'it prints something'. + +Each demo drives the real API; here we assert the *shape* of what a few of them +produce (idempotence, exit code, cost figure) so a regression in the estimator +or a broken fixture is caught at the demo layer too. +""" +from __future__ import annotations + +import contextlib +import importlib +import io +import os +import sys + +import pytest + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +DEMOS = os.path.join(REPO_ROOT, "demos") +for p in (REPO_ROOT, DEMOS): + if p not in sys.path: + sys.path.insert(0, p) + +ALL_SCENARIOS = [ + "01_ai_app_prompt_cost", "02_finops_model_selection", "03_ci_budget_gate", + "04_eng_manager_prompt_library", "05_rag_context_budget", + "06_chat_transcript_cost", "07_context_window_guard", "08_diff_review_cost", + "09_fewshot_tax", "10_classifier_router", "11_agent_step_budget", + "12_prompt_diet", "13_rag_chunk_tuning", "14_multi_model_portfolio", + "15_streaming_output_forecast", "16_batch_library_audit", + "17_custom_model_pricing", "18_json_output_pipeline", "19_regression_guard", + "20_full_month_forecast", +] + + +def _capture(name): + mod = importlib.import_module(name) + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + mod.main() + return buf.getvalue() + + +@pytest.mark.parametrize("name", ALL_SCENARIOS) +def test_every_demo_prints_cost(name): + out = _capture(name) + assert "$" in out, f"{name} printed no cost figure" + assert out.strip() + + +@pytest.mark.parametrize("name", ALL_SCENARIOS) +def test_every_demo_is_idempotent(name): + # Running twice must not raise (e.g. add_model overrides cleanly, no state + # leaks between runs). + a = _capture(name) + b = _capture(name) + assert a and b + + +def test_run_all_lists_twenty_scenarios(): + runner = importlib.import_module("run_all") + assert len(runner.SCENARIOS) == 20 + + +def test_run_all_completes(): + runner = importlib.import_module("run_all") + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + runner.main() + assert "All demo scenarios completed." in buf.getvalue() + + +def test_json_pipeline_demo_emits_valid_json(): + import json + out = _capture("18_json_output_pipeline") + # extract the JSON block (between first '{' and matching structure); the + # demo prints a full record, so find a parseable object. + start = out.index("{") + # find the end by scanning for the last closing brace before "Parsed" + end = out.rindex("}") + 1 + obj = json.loads(out[start:end]) + assert obj["estimate"]["model"] == "claude-sonnet" diff --git a/tests/test_identity.py b/tests/test_identity.py new file mode 100644 index 0000000..8b6606b --- /dev/null +++ b/tests/test_identity.py @@ -0,0 +1,46 @@ +"""Package identity, public-API surface, and metadata consistency tests.""" +from __future__ import annotations + +import importlib + +import pytest + + +def test_tool_name_and_version_from_core(): + core = importlib.import_module("tokenmeter.core") + assert core.TOOL_NAME == "tokenmeter" + assert isinstance(core.TOOL_VERSION, str) and core.TOOL_VERSION + + +def test_package_reexports_identity(): + m = importlib.import_module("tokenmeter") + assert m.TOOL_NAME == "tokenmeter" + assert m.__version__ == m.TOOL_VERSION + + +def test_public_api_symbols_present(): + m = importlib.import_module("tokenmeter") + for name in ( + "count_tokens", "estimate", "get_pricing", "list_models", + ): + assert hasattr(m, name), name + + +def test_core_public_api_symbols_present(): + core = importlib.import_module("tokenmeter.core") + for name in ( + "count_tokens", "estimate", "get_pricing", "list_models", + "add_model", "aggregate", "check_budget", "compare_models", + "MODELS", "ModelPricing", "Estimate", "BudgetResult", + ): + assert hasattr(core, name), name + + +def test_cli_has_main_and_build_parser(): + cli = importlib.import_module("tokenmeter.cli") + assert callable(cli.main) + assert callable(cli.build_parser) + + +def test_main_module_importable(): + assert importlib.import_module("tokenmeter.__main__") is not None diff --git a/tokenmeter/cli.py b/tokenmeter/cli.py index 90f1722..a8c9b49 100644 --- a/tokenmeter/cli.py +++ b/tokenmeter/cli.py @@ -175,9 +175,12 @@ def _cmd_compare(args: argparse.Namespace) -> int: ests = compare_models( text, output_tokens=args.output_tokens, models=models ) - except KeyError as exc: + except (KeyError, ValueError) as exc: print(f"error: {exc.args[0] if exc.args else exc}", file=sys.stderr) return 2 + if not ests: + print("error: no models to compare", file=sys.stderr) + return 2 payload = { "output_tokens": args.output_tokens, "input_tokens": ests[0].input_tokens if ests else 0, @@ -302,7 +305,7 @@ def main(argv: Optional[List[str]] = None) -> int: except KeyError as exc: print(f"error: {exc.args[0] if exc.args else exc}", file=sys.stderr) return 2 - except FileNotFoundError as exc: + except (ValueError, FileNotFoundError) as exc: print(f"error: {exc}", file=sys.stderr) return 2 except BrokenPipeError: diff --git a/tokenmeter/core.py b/tokenmeter/core.py index 2dd0d77..6ca76c1 100644 --- a/tokenmeter/core.py +++ b/tokenmeter/core.py @@ -17,6 +17,12 @@ from dataclasses import dataclass, field from typing import Dict, Iterable, List, Optional +# Tool identity. Kept here (not just in __init__) so `from tokenmeter.core import +# TOOL_NAME, TOOL_VERSION` — the path __init__ documents and prefers — actually +# resolves instead of silently falling back. +TOOL_NAME = "tokenmeter" +TOOL_VERSION = "0.1.1" + @dataclass(frozen=True) class ModelPricing: @@ -50,8 +56,25 @@ def cost(self, input_tokens: int, output_tokens: int) -> float: def add_model( name: str, input_per_1k: float, output_per_1k: float, context_window: int = 8192 ) -> ModelPricing: - """Register or override a model's pricing at runtime.""" - p = ModelPricing(name, float(input_per_1k), float(output_per_1k), int(context_window)) + """Register or override a model's pricing at runtime. + + Raises ``ValueError`` on a blank name, non-numeric prices, negative prices, + or a non-positive context window — a mispriced model would silently corrupt + every downstream estimate, so we fail loudly at registration time instead. + """ + if not name or not str(name).strip(): + raise ValueError("model name must be a non-empty string") + try: + in_price = float(input_per_1k) + out_price = float(output_per_1k) + ctx = int(context_window) + except (TypeError, ValueError) as exc: + raise ValueError(f"model pricing must be numeric: {exc}") from exc + if in_price < 0 or out_price < 0: + raise ValueError("model prices must be non-negative") + if ctx <= 0: + raise ValueError("context_window must be a positive integer") + p = ModelPricing(name, in_price, out_price, ctx) MODELS[name] = p return p @@ -84,6 +107,10 @@ def count_tokens(text: str) -> int: BPE tends to split numbers), min 1. * Single punctuation/symbol: 1 token. """ + if text is None: + return 0 + if not isinstance(text, str): + raise TypeError(f"count_tokens expects str, got {type(text).__name__}") if not text: return 0 @@ -144,8 +171,15 @@ def estimate( is the expected/observed completion length (defaults to 0). """ pricing = get_pricing(model) - in_tok = count_tokens(text) if input_tokens is None else int(input_tokens) + if input_tokens is not None: + in_tok = int(input_tokens) + if in_tok < 0: + raise ValueError(f"input_tokens must be >= 0, got {in_tok}") + else: + in_tok = count_tokens(text) out_tok = int(output_tokens) + if out_tok < 0: + raise ValueError(f"output_tokens must be >= 0, got {out_tok}") in_cost = in_tok / 1000.0 * pricing.input_per_1k out_cost = out_tok / 1000.0 * pricing.output_per_1k used = in_tok + out_tok