diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..c53b7c9 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,21 @@ +name: CI + +on: + push: + branches: [master] + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: "1.25" + - name: Build + run: go build ./... + - name: Vet + run: go vet ./... + - name: Test + run: go test ./... -count=1 diff --git a/.gitignore b/.gitignore index 9865439..365770d 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,11 @@ go.work.sum #IDE .idea + +# Augur build + generated artifacts (keep the repo clean of local runs) +/augur +/augur-report.md +/augur-report.json +/trace.jsonl +/report.md +/report.json diff --git a/README.md b/README.md index 4f12bee..984e825 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,90 @@ augur gate --traffic traffic.yaml --budget budget.yaml `augur gate` is the one you wire into CI. +### Record once, replay for free + +Running the agent against the real provider on every CI push spends real tokens. +Record the responses once, then replay them — `augur run --replay` re-executes +the agent against the recorded responses and regenerates the trace **without +contacting the provider**: + +```sh +# Once (locally or nightly), against the real provider: +augur run --scenarios scenarios.yaml --record cassette.jsonl + +# In CI, on every push — zero tokens spent: +augur run --scenarios scenarios.yaml --replay cassette.jsonl --trace trace.jsonl +augur gate --traffic traffic.yaml --budget budget.yaml +``` + +Because replay re-runs the *agent* (not just the trace), a cost regression from +an agent code change still surfaces — exercised against the old responses, for +free. Commit `cassette.jsonl` alongside your scenarios. (A call the agent makes +that wasn't recorded is a replay miss — reported loudly, so divergence from the +recording never passes silently.) + +### What-if sensitivity analysis + +`project` and `gate` take what-if multipliers that re-cost the *recorded* trace +under hypothetical agentic cost drivers — no agent re-run, no tokens: + +```sh +# "What if retries climb 30%, sub-agent fan-out adds 50% more calls, and +# context grows to 2× as conversations lengthen?" +augur project --retry-rate 0.3 --fanout 1.5 --context-growth 2 + +# Gate against a pessimistic scenario, not just today's happy path: +augur gate --context-growth 1.5 --budget budget.yaml +``` + +Retries and fan-out scale the whole call (more calls); context growth inflates +only the prompt side, not the completion — so the model reflects *which* driver +moved, not a flat fudge factor. + +### Self-hosted models (TCO) + +For a model you run yourself there's no per-token API price — you pay for an +instance by the hour. Describe the deployment in `tco.yaml` (instance $/hr, +serving throughput, utilization) and Augur derives the effective $/Mtok: + +```sh +augur tco --tco tco.yaml # show the derived effective $/Mtok +augur gate --tco tco.yaml ... # cost the trace against self-hosted pricing +``` + +`--tco` is accepted by `aggregate`, `project`, and `gate` as an alternative to +`--pricing`. Utilization is the honest part: you pay for the box 24/7 but it's +rarely saturated, and a half-idle instance doubles the effective token price. + +### GitHub Action + +Augur ships as a composite action that runs the gate on a pull request, posts +the report as a comment, and fails the check if the projection is over budget. +Combined with a committed cassette, the PR check spends no tokens: + +```yaml +# .github/workflows/cost-gate.yml +name: Cost gate +on: pull_request +jobs: + augur: + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write # to post the report comment + steps: + - uses: actions/checkout@v4 + - uses: your-org/augur@v1 + with: + cassette: cassette.jsonl # replay → zero tokens + traffic: traffic.yaml + budget: budget.yaml +``` + +See [`action.yml`](action.yml) for all inputs and +[`examples/github-workflow.yml`](examples/github-workflow.yml) for a fuller +example. + --- ## Design decisions @@ -182,10 +266,10 @@ dependency is `gopkg.in/yaml.v3`): | Hito 2 | scenario runner + per-scenario aggregation | | Hito 3 | projection engine with bootstrap confidence intervals | | Hito 4 | budget gate + Markdown/JSON report + CI exit codes | +| Hito 5 | record/replay cassette (`--record`/`--replay`), what-if knobs (`--retry-rate`/`--fanout`/`--context-growth`), self-hosted TCO mode (`augur tco`, `--tco`), GitHub Action (`action.yml`) | -**Roadmap (stretch, each independently shippable):** GitHub Action wrapper • -what-if multiplier knobs • record-once/replay to cut CI token cost • self-hosted -TCO mode • a Python output-length prediction sidecar. +**Roadmap (stretch):** a Python output-length prediction sidecar (the analytical +piece that would earn a second language). See [`SPEC.md`](SPEC.md) for the full design. diff --git a/action.yml b/action.yml new file mode 100644 index 0000000..6ea257b --- /dev/null +++ b/action.yml @@ -0,0 +1,144 @@ +name: "Augur cost gate" +description: "Project an AI agent's production cost from a recorded trace and fail the PR if it's over budget." +branding: + icon: "dollar-sign" + color: "purple" + +inputs: + trace: + description: "Path to the cost trace (JSONL). Ignored if `cassette` is set." + default: "trace.jsonl" + pricing: + description: "Path to the pricing snapshot." + default: "pricing.yaml" + tco: + description: "Path to a self-hosted TCO config. Overrides `pricing` when set." + default: "" + traffic: + description: "Path to the traffic profile." + default: "traffic.yaml" + budget: + description: "Path to the budget thresholds." + default: "budget.yaml" + scenarios: + description: "Path to scenarios.yaml. Only used when replaying a cassette." + default: "scenarios.yaml" + cassette: + description: "If set, replay this cassette to regenerate the trace (no tokens spent). Requires the agent command to be runnable in CI." + default: "" + extra-args: + description: "Extra flags passed to `augur gate` (e.g. what-if knobs: --context-growth 1.5)." + default: "" + comment: + description: "Post (or update) the report as a PR comment." + default: "true" + working-directory: + description: "Directory the config files live in." + default: "." + go-version: + description: "Go version used to build augur." + default: "1.25" + +outputs: + pass: + description: "'true' if within budget, 'false' if over." + value: ${{ steps.gate.outputs.pass }} + report-path: + description: "Absolute path to the generated Markdown report." + value: ${{ steps.gate.outputs.report-path }} + +runs: + using: "composite" + steps: + - name: Set up Go + uses: actions/setup-go@v5 + with: + go-version: ${{ inputs.go-version }} + + - name: Build augur + shell: bash + run: | + cd "${{ github.action_path }}" + go build -o "${RUNNER_TEMP}/augur" . + echo "AUGUR=${RUNNER_TEMP}/augur" >> "$GITHUB_ENV" + + - name: Replay cassette + if: ${{ inputs.cassette != '' }} + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + "$AUGUR" run \ + --scenarios "${{ inputs.scenarios }}" \ + --replay "${{ inputs.cassette }}" \ + --trace "${{ inputs.trace }}" + + - name: Augur gate + id: gate + shell: bash + working-directory: ${{ inputs.working-directory }} + run: | + pricing_args="--pricing ${{ inputs.pricing }}" + if [ -n "${{ inputs.tco }}" ]; then + pricing_args="--tco ${{ inputs.tco }}" + fi + + set +e + "$AUGUR" gate \ + --trace "${{ inputs.trace }}" \ + $pricing_args \ + --traffic "${{ inputs.traffic }}" \ + --budget "${{ inputs.budget }}" \ + --report-md augur-report.md \ + --report-json augur-report.json \ + ${{ inputs.extra-args }} + code=$? + set -e + + if [ "$code" -eq 0 ]; then + echo "pass=true" >> "$GITHUB_OUTPUT" + elif [ -f augur-report.md ]; then + # Report exists => a genuine over-budget verdict (not an internal error). + echo "pass=false" >> "$GITHUB_OUTPUT" + else + echo "::error::augur gate failed before producing a report (exit $code)" + exit "$code" + fi + echo "exit-code=$code" >> "$GITHUB_OUTPUT" + echo "report-path=$(pwd)/augur-report.md" >> "$GITHUB_OUTPUT" + + - name: Post report comment + if: ${{ always() && inputs.comment == 'true' && github.event_name == 'pull_request' && steps.gate.outputs.report-path != '' }} + uses: actions/github-script@v7 + env: + AUGUR_REPORT_PATH: ${{ steps.gate.outputs.report-path }} + with: + script: | + const fs = require('fs'); + const marker = ''; + let report; + try { + report = fs.readFileSync(process.env.AUGUR_REPORT_PATH, 'utf8'); + } catch (e) { + report = '_Augur produced no report._'; + } + const body = `${marker}\n${report}`; + const { owner, repo } = context.repo; + const issue_number = context.issue.number; + const comments = await github.paginate(github.rest.issues.listComments, { + owner, repo, issue_number, + }); + const existing = comments.find(c => c.body && c.body.includes(marker)); + if (existing) { + await github.rest.issues.updateComment({ owner, repo, comment_id: existing.id, body }); + } else { + await github.rest.issues.createComment({ owner, repo, issue_number, body }); + } + + - name: Enforce budget + shell: bash + run: | + if [ "${{ steps.gate.outputs.pass }}" != "true" ]; then + echo "::error::Augur: projected cost is over budget — see the report comment." + exit 1 + fi + echo "Augur: within budget." diff --git a/aggregate/aggregate.go b/aggregate/aggregate.go index 2f28669..3056863 100644 --- a/aggregate/aggregate.go +++ b/aggregate/aggregate.go @@ -61,6 +61,50 @@ type Result struct { Runs []Run `json:"runs"` } +// Knobs are what-if multipliers applied to every call's cost, for sensitivity +// analysis that reuses an existing trace instead of re-running the agent. The +// zero value is the identity (the observed cost). They model the agentic cost +// drivers: +// +// RetryRate extra fraction of calls retried (0.2 = 20% more calls) +// FanoutFactor multiplier on call count from sub-agent/tool fan-out +// ContextGrowth multiplier on prompt (input + cached) cost from history growth +// +// Retries and fan-out add whole calls, so they scale the entire call cost; +// context growth inflates the prompt only, not the completion. Per call: +// +// cost' = (1 + RetryRate) * FanoutFactor * (ContextGrowth*prompt + output) +type Knobs struct { + RetryRate float64 + FanoutFactor float64 + ContextGrowth float64 +} + +// normalized fills in identity defaults: FanoutFactor and ContextGrowth default +// to 1 (no change) when left at zero, RetryRate stays 0. +func (k Knobs) normalized() Knobs { + if k.FanoutFactor == 0 { + k.FanoutFactor = 1 + } + if k.ContextGrowth == 0 { + k.ContextGrowth = 1 + } + return k +} + +// IsIdentity reports whether the knobs leave cost unchanged. +func (k Knobs) IsIdentity() bool { + n := k.normalized() + return n.RetryRate == 0 && n.FanoutFactor == 1 && n.ContextGrowth == 1 +} + +// apply returns the what-if cost of a single call given its component breakdown. +func (k Knobs) apply(b cost.Breakdown) float64 { + n := k.normalized() + callMult := (1 + n.RetryRate) * n.FanoutFactor + return callMult * (n.ContextGrowth*b.PromptUSD() + b.OutputUSD) +} + // Aggregate prices every record against pricing and builds per-scenario cost // distributions. An un-priceable model is a hard error (wrapping // cost.ErrUnknownModel): an un-priced call means an unknowable bill, the exact @@ -68,6 +112,12 @@ type Result struct { // // An empty trace yields a zero Result and no error. func Aggregate(records []trace.Record, pricing cost.Pricing) (Result, error) { + return AggregateWithKnobs(records, pricing, Knobs{}) +} + +// AggregateWithKnobs is Aggregate with what-if multipliers applied to every +// call's cost (see Knobs). With the zero Knobs it is identical to Aggregate. +func AggregateWithKnobs(records []trace.Record, pricing cost.Pricing, knobs Knobs) (Result, error) { type runKey struct{ scenario, run string } runs := make(map[runKey]*Run) @@ -81,11 +131,12 @@ func Aggregate(records []trace.Record, pricing cost.Pricing) (Result, error) { OutputTokens: rec.OutputTokens, CachedTokens: rec.CachedTokens, } - c, err := pricing.Cost(rec.Model, u) + b, err := pricing.Breakdown(rec.Model, u) if err != nil { return Result{}, fmt.Errorf("aggregate: scenario %q run %q seq %d: %w", rec.ScenarioID, rec.RunID, rec.Seq, err) } + c := knobs.apply(b) k := runKey{rec.ScenarioID, rec.RunID} r, ok := runs[k] diff --git a/aggregate/aggregate_test.go b/aggregate/aggregate_test.go index cca404e..85a405c 100644 --- a/aggregate/aggregate_test.go +++ b/aggregate/aggregate_test.go @@ -135,6 +135,72 @@ func TestAggregateUnknownModelIsHardError(t *testing.T) { } } +func TestAggregateWithKnobs(t *testing.T) { + // One run, one call: gpt-4o, 1000 in (200 cached), 500 out. + // prompt = 800@2.50 + 200@1.25 = 0.0020 + 0.00025 = 0.00225 + // output = 500@10.00 = 0.0050 + records := []trace.Record{rec("s", "r", 0, "gpt-4o", 1000, 500, 200)} + + // Identity knobs reproduce the observed cost (0.00725). + base, err := AggregateWithKnobs(records, testPricing(), Knobs{}) + if err != nil { + t.Fatalf("Aggregate identity: %v", err) + } + if !approx(base.Runs[0].CostUSD, 0.00725) { + t.Errorf("identity run cost = %v, want 0.00725", base.Runs[0].CostUSD) + } + + // Knobs: retry +50%, fanout ×2, context ×3. + // callMult = 1.5 * 2 = 3 + // cost' = 3 * (3*0.00225 + 0.0050) = 3 * 0.01175 = 0.03525 + knobs := Knobs{RetryRate: 0.5, FanoutFactor: 2, ContextGrowth: 3} + res, err := AggregateWithKnobs(records, testPricing(), knobs) + if err != nil { + t.Fatalf("Aggregate with knobs: %v", err) + } + if !approx(res.Runs[0].CostUSD, 0.03525) { + t.Errorf("knob run cost = %v, want 0.03525", res.Runs[0].CostUSD) + } + // The scenario distribution and per-model totals scale identically. + if !approx(res.Scenarios[0].TotalCost, 0.03525) { + t.Errorf("scenario total = %v, want 0.03525", res.Scenarios[0].TotalCost) + } + if !approx(res.Scenarios[0].ByModel[0].CostUSD, 0.03525) { + t.Errorf("by-model cost = %v, want 0.03525", res.Scenarios[0].ByModel[0].CostUSD) + } +} + +func TestKnobsIdentity(t *testing.T) { + cases := []struct { + k Knobs + want bool + }{ + {Knobs{}, true}, + {Knobs{FanoutFactor: 1, ContextGrowth: 1}, true}, + {Knobs{RetryRate: 0.1}, false}, + {Knobs{FanoutFactor: 2}, false}, + {Knobs{ContextGrowth: 1.5}, false}, + } + for _, c := range cases { + if got := c.k.IsIdentity(); got != c.want { + t.Errorf("%+v IsIdentity = %v, want %v", c.k, got, c.want) + } + } +} + +// Context growth must inflate ONLY the prompt side, not the completion. +func TestKnobsContextGrowthPromptOnly(t *testing.T) { + // Pure-output call (0 input): context growth must leave it unchanged. + outOnly := []trace.Record{rec("s", "r", 0, "gpt-4o", 0, 1000, 0)} // 1000@10 = 0.01 + res, err := AggregateWithKnobs(outOnly, testPricing(), Knobs{ContextGrowth: 5}) + if err != nil { + t.Fatalf("Aggregate: %v", err) + } + if !approx(res.Runs[0].CostUSD, 0.01) { + t.Errorf("output-only cost under context growth = %v, want 0.01 (unchanged)", res.Runs[0].CostUSD) + } +} + func TestAggregateEmpty(t *testing.T) { res, err := Aggregate(nil, testPricing()) if err != nil { diff --git a/aggregate_cmd.go b/aggregate_cmd.go index 0fc8db2..cb6a288 100644 --- a/aggregate_cmd.go +++ b/aggregate_cmd.go @@ -7,7 +7,6 @@ import ( "os" "augur/aggregate" - "augur/cost" "augur/trace" ) @@ -18,6 +17,7 @@ func runAggregate(args []string) error { fs := flag.NewFlagSet("aggregate", flag.ContinueOnError) tracePath := fs.String("trace", "trace.jsonl", "path to the cost trace (JSONL) to aggregate") pricingPath := fs.String("pricing", "pricing.yaml", "path to the pricing snapshot") + tcoPath := fs.String("tco", "", "derive pricing from a self-hosted TCO config instead of -pricing") asJSON := fs.Bool("json", false, "emit the aggregation as JSON instead of a table") if err := fs.Parse(args); err != nil { return err @@ -34,7 +34,7 @@ func runAggregate(args []string) error { return err } - pricing, err := cost.LoadPricing(*pricingPath) + pricing, err := resolvePricing(*pricingPath, *tcoPath) if err != nil { return err } diff --git a/cassette/cassette.go b/cassette/cassette.go new file mode 100644 index 0000000..453bab2 --- /dev/null +++ b/cassette/cassette.go @@ -0,0 +1,132 @@ +// Package cassette is the record/replay store for the proxy: a recording of +// every LLM response keyed by the call's identity (scenario, run, seq). +// +// It exists to cut the token cost of re-running the gate (SPEC D3 / Hito 5). +// Record once against the real provider; thereafter replay the cassette so CI +// re-runs the agent against the recorded responses and regenerates the cost +// trace without spending a cent. Because replay re-executes the AGENT (not just +// the trace), a cost regression introduced by an agent code change still shows +// up — it is exercised against the old LLM responses, for free. +// +// The on-disk format is JSON Lines, one Entry per line, so a cassette is +// diffable and can be committed to the repo alongside the scenarios. +package cassette + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "os" + "sync" +) + +// Entry is one recorded response. Body holds the raw response bytes as a string +// (OpenAI-compatible responses are UTF-8 text — JSON or an SSE stream), so the +// cassette stays human-readable. +type Entry struct { + ScenarioID string `json:"scenario_id"` + RunID string `json:"run_id"` + Seq int `json:"seq"` + Status int `json:"status"` + ContentType string `json:"content_type,omitempty"` + LatencyMs int64 `json:"latency_ms"` + Body string `json:"body"` +} + +type key struct { + scenario, run string + seq int +} + +// Cassette records entries (in record mode) or serves them (in replay mode). +// A given Cassette is used for one or the other, not both. It is safe for +// concurrent use. +type Cassette struct { + mu sync.Mutex + entries map[key]Entry // populated for replay + w io.Writer // set for record + closer io.Closer +} + +// Create opens path for recording, truncating any existing cassette — a +// recording is a fresh capture, not an append. +func Create(path string) (*Cassette, error) { + f, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644) + if err != nil { + return nil, fmt.Errorf("cassette: creating cassette: %w", err) + } + return &Cassette{w: f, closer: f}, nil +} + +// Load reads a cassette from path into memory for replay. +func Load(path string) (*Cassette, error) { + f, err := os.Open(path) + if err != nil { + return nil, fmt.Errorf("cassette: opening cassette: %w", err) + } + defer f.Close() + + entries := make(map[key]Entry) + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + line := 0 + for sc.Scan() { + line++ + b := sc.Bytes() + if len(b) == 0 { + continue + } + var e Entry + if err := json.Unmarshal(b, &e); err != nil { + return nil, fmt.Errorf("cassette: parsing line %d: %w", line, err) + } + entries[key{e.ScenarioID, e.RunID, e.Seq}] = e + } + if err := sc.Err(); err != nil { + return nil, fmt.Errorf("cassette: reading cassette: %w", err) + } + return &Cassette{entries: entries}, nil +} + +// Record appends one entry to the cassette. Concurrent calls are serialized so +// lines never interleave. +func (c *Cassette) Record(e Entry) error { + b, err := json.Marshal(e) + if err != nil { + return fmt.Errorf("cassette: marshaling entry: %w", err) + } + b = append(b, '\n') + + c.mu.Lock() + defer c.mu.Unlock() + if _, err := c.w.Write(b); err != nil { + return fmt.Errorf("cassette: writing entry: %w", err) + } + return nil +} + +// Lookup returns the recorded entry for a call, or ok=false on a replay miss +// (the agent made a call that was not recorded — a sign its behavior diverged +// from the recording). +func (c *Cassette) Lookup(scenario, run string, seq int) (Entry, bool) { + c.mu.Lock() + defer c.mu.Unlock() + e, ok := c.entries[key{scenario, run, seq}] + return e, ok +} + +// Len reports how many entries a loaded cassette holds. +func (c *Cassette) Len() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.entries) +} + +// Close closes the underlying file when the cassette owns one (record mode). +func (c *Cassette) Close() error { + if c.closer == nil { + return nil + } + return c.closer.Close() +} diff --git a/cassette/cassette_test.go b/cassette/cassette_test.go new file mode 100644 index 0000000..54a5041 --- /dev/null +++ b/cassette/cassette_test.go @@ -0,0 +1,113 @@ +package cassette + +import ( + "os" + "path/filepath" + "testing" +) + +// appendLine appends a raw line to a file, for corrupting a cassette in tests. +func appendLine(path, line string) error { + f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o644) + if err != nil { + return err + } + defer f.Close() + _, err = f.WriteString(line + "\n") + return err +} + +func TestRecordThenLoad(t *testing.T) { + path := filepath.Join(t.TempDir(), "cassette.jsonl") + + c, err := Create(path) + if err != nil { + t.Fatalf("Create: %v", err) + } + entries := []Entry{ + {ScenarioID: "checkout", RunID: "checkout-000", Seq: 0, Status: 200, + ContentType: "application/json", LatencyMs: 120, Body: `{"usage":{"prompt_tokens":10}}`}, + {ScenarioID: "checkout", RunID: "checkout-000", Seq: 1, Status: 200, + ContentType: "text/event-stream", LatencyMs: 80, Body: "data: [DONE]\n\n"}, + } + for _, e := range entries { + if err := c.Record(e); err != nil { + t.Fatalf("Record: %v", err) + } + } + if err := c.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + + loaded, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if loaded.Len() != 2 { + t.Fatalf("Len = %d, want 2", loaded.Len()) + } + + got, ok := loaded.Lookup("checkout", "checkout-000", 0) + if !ok { + t.Fatal("lookup seq 0 missed") + } + if got != entries[0] { + t.Errorf("entry 0 = %+v, want %+v", got, entries[0]) + } + if _, ok := loaded.Lookup("checkout", "checkout-000", 1); !ok { + t.Error("lookup seq 1 missed") + } +} + +func TestLookupMiss(t *testing.T) { + path := filepath.Join(t.TempDir(), "c.jsonl") + c, _ := Create(path) + c.Record(Entry{ScenarioID: "a", RunID: "a-000", Seq: 0, Body: "{}"}) + c.Close() + + loaded, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if _, ok := loaded.Lookup("a", "a-000", 1); ok { + t.Error("expected miss for unrecorded seq") + } + if _, ok := loaded.Lookup("b", "b-000", 0); ok { + t.Error("expected miss for unrecorded scenario") + } +} + +func TestCreateTruncates(t *testing.T) { + path := filepath.Join(t.TempDir(), "c.jsonl") + c1, _ := Create(path) + c1.Record(Entry{ScenarioID: "a", RunID: "r", Seq: 0, Body: "{}"}) + c1.Close() + + // Re-creating must start fresh, not append. + c2, _ := Create(path) + c2.Record(Entry{ScenarioID: "b", RunID: "r", Seq: 0, Body: "{}"}) + c2.Close() + + loaded, _ := Load(path) + if loaded.Len() != 1 { + t.Errorf("Len = %d, want 1 (Create should truncate)", loaded.Len()) + } + if _, ok := loaded.Lookup("b", "r", 0); !ok { + t.Error("expected the second recording to be the only entry") + } +} + +func TestLoadRejectsMalformed(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "bad.jsonl") + c, _ := Create(path) + c.Record(Entry{ScenarioID: "a", RunID: "r", Seq: 0, Body: "{}"}) + c.Close() + // Append a malformed line. + if err := appendLine(path, "{not json"); err != nil { + t.Fatal(err) + } + if _, err := Load(path); err == nil { + t.Error("expected error loading malformed cassette") + } +} diff --git a/cost/cost.go b/cost/cost.go index 83b9ad9..d4a3e93 100644 --- a/cost/cost.go +++ b/cost/cost.go @@ -67,18 +67,48 @@ func (u Usage) Validate() error { return nil } -// Cost returns the USD cost of a single call priced at p. The cached portion of -// the prompt is billed at CachedInput, the remainder at Input, and completion -// tokens at Output. It returns an error if the usage is invalid (see Validate). -func (p ModelPrice) Cost(u Usage) (float64, error) { +// Breakdown is a single call's cost split into its three components. It exists +// so callers (the what-if knobs) can scale the prompt side independently of the +// completion side — context growth inflates input, not output. +type Breakdown struct { + // InputUSD is the cost of the non-cached prompt tokens. + InputUSD float64 + // CachedUSD is the cost of the cached prompt tokens. + CachedUSD float64 + // OutputUSD is the cost of the completion tokens. + OutputUSD float64 +} + +// Total is the full call cost: the three components summed. +func (b Breakdown) Total() float64 { return b.InputUSD + b.CachedUSD + b.OutputUSD } + +// PromptUSD is the cost attributable to the prompt (input + cached) — the part +// that scales with context growth. +func (b Breakdown) PromptUSD() float64 { return b.InputUSD + b.CachedUSD } + +// Breakdown returns the per-component cost of a single call priced at p. The +// cached portion of the prompt is billed at CachedInput, the remainder at +// Input, and completion tokens at Output. It errors if the usage is invalid. +func (p ModelPrice) Breakdown(u Usage) (Breakdown, error) { if err := u.Validate(); err != nil { - return 0, err + return Breakdown{}, err } fullInput := u.InputTokens - u.CachedTokens - dollars := float64(fullInput)/tokensPerMtok*p.Input + - float64(u.CachedTokens)/tokensPerMtok*p.CachedInput + - float64(u.OutputTokens)/tokensPerMtok*p.Output - return dollars, nil + return Breakdown{ + InputUSD: float64(fullInput) / tokensPerMtok * p.Input, + CachedUSD: float64(u.CachedTokens) / tokensPerMtok * p.CachedInput, + OutputUSD: float64(u.OutputTokens) / tokensPerMtok * p.Output, + }, nil +} + +// Cost returns the USD cost of a single call priced at p. It returns an error if +// the usage is invalid (see Validate). +func (p ModelPrice) Cost(u Usage) (float64, error) { + b, err := p.Breakdown(u) + if err != nil { + return 0, err + } + return b.Total(), nil } // Pricing is a loaded pricing snapshot: a set of per-model prices plus the date @@ -104,3 +134,13 @@ func (p Pricing) Cost(model string, u Usage) (float64, error) { } return mp.Cost(u) } + +// Breakdown computes the per-component cost of a single call for the named +// model, wrapping ErrUnknownModel when the model is absent. +func (p Pricing) Breakdown(model string, u Usage) (Breakdown, error) { + mp, ok := p.Models[model] + if !ok { + return Breakdown{}, fmt.Errorf("%w: %q", ErrUnknownModel, model) + } + return mp.Breakdown(u) +} diff --git a/cost/cost_test.go b/cost/cost_test.go index 6d26c7c..449256c 100644 --- a/cost/cost_test.go +++ b/cost/cost_test.go @@ -80,6 +80,32 @@ func TestModelPriceCost(t *testing.T) { } } +func TestBreakdown(t *testing.T) { + // 1000 input (200 cached), 500 output on gpt-4o: + // full input 800 @2.50 = 0.0020 ; cached 200 @1.25 = 0.00025 ; out 500 @10 = 0.0050 + b, err := gpt4o.Breakdown(Usage{InputTokens: 1000, OutputTokens: 500, CachedTokens: 200}) + if err != nil { + t.Fatalf("Breakdown: %v", err) + } + if !approxEqual(b.InputUSD, 0.0020) { + t.Errorf("InputUSD = %v, want 0.0020", b.InputUSD) + } + if !approxEqual(b.CachedUSD, 0.00025) { + t.Errorf("CachedUSD = %v, want 0.00025", b.CachedUSD) + } + if !approxEqual(b.OutputUSD, 0.0050) { + t.Errorf("OutputUSD = %v, want 0.0050", b.OutputUSD) + } + if !approxEqual(b.PromptUSD(), 0.00225) { + t.Errorf("PromptUSD = %v, want 0.00225", b.PromptUSD()) + } + // Total must equal Cost for the same usage. + c, _ := gpt4o.Cost(Usage{InputTokens: 1000, OutputTokens: 500, CachedTokens: 200}) + if !approxEqual(b.Total(), c) { + t.Errorf("Total %v != Cost %v", b.Total(), c) + } +} + func TestCostInvalidUsage(t *testing.T) { tests := []struct { name string diff --git a/examples/github-workflow.yml b/examples/github-workflow.yml new file mode 100644 index 0000000..471dc6b --- /dev/null +++ b/examples/github-workflow.yml @@ -0,0 +1,35 @@ +# Example: gate an AI agent's projected cost on every pull request. +# +# Copy this into your repo as .github/workflows/cost-gate.yml and adjust the +# paths. The lightest setup commits a recorded cassette (or a trace) so CI spends +# no tokens: Augur replays it, projects the cost, and fails the PR if it's over +# budget — posting the report as a comment. + +name: Cost gate + +on: pull_request + +jobs: + augur: + runs-on: ubuntu-latest + permissions: + contents: read + pull-requests: write # needed to post the report comment + steps: + - uses: actions/checkout@v4 + + # Replace with the published action ref, e.g.: + # uses: your-org/augur@v1 + # (Within the Augur repo itself, `uses: ./` runs the local action.) + - uses: your-org/augur@v1 + with: + # Replay a committed cassette so CI spends no tokens. Drop `cassette` + # and point `trace` at a committed trace.jsonl if you don't run the + # agent in CI. + cassette: cassette.jsonl + scenarios: scenarios.yaml + pricing: pricing.yaml + traffic: traffic.yaml + budget: budget.yaml + # Optional: gate against a pessimistic scenario. + # extra-args: "--context-growth 1.5 --retry-rate 0.2" diff --git a/examples/tco.yaml b/examples/tco.yaml new file mode 100644 index 0000000..870a1a1 --- /dev/null +++ b/examples/tco.yaml @@ -0,0 +1,26 @@ +# Augur — example TCO config (used by `augur tco`, and by project/gate via -tco) +# +# For self-hosted models there is no per-token API price; you pay for an +# instance by the hour and get some serving throughput. Augur derives the +# effective $/Mtok: +# +# $/Mtok = instance_cost_per_hour x 1e6 / (tokens_per_sec x 3600 x utilization) +# +# The deployment name must match the model id your agent sends, so the derived +# price lines up with the trace. A self-hosted price is token-type-agnostic: +# input, output, and cached tokens all carry the effective rate. + +version: 1 + +deployments: + # e.g. Llama-3.1-70B on an 8xH100 node. + llama-3.1-70b: + instance_cost_per_hour: 32.77 # on-demand 8xH100 instance + tokens_per_sec: 2500 # measured aggregate serving throughput + utilization: 0.6 # you pay 24/7 but the box isn't always busy + + # e.g. Mixtral-8x7B on a single A100. + mixtral-8x7b: + instance_cost_per_hour: 3.67 + tokens_per_sec: 1800 + # utilization omitted -> defaults to 1.0 (always busy) diff --git a/gate/report.go b/gate/report.go index 369b51b..b56c55e 100644 --- a/gate/report.go +++ b/gate/report.go @@ -46,7 +46,11 @@ func (r Report) WriteMarkdown(w io.Writer) error { bw.printf("- Pricing snapshot: `%s`\n", p.SnapshotDate) bw.printf("- Traffic: %d users × %g req/user/day, %d tenants, %d days/month\n", tr.Users, tr.RequestsPerUserPerDay, tr.Tenants, tr.DaysPerMonth) - bw.printf("- Bootstrap: %d samples, %.0f%% CI, seed %d\n\n", p.BootstrapSamples, p.CILevel*100, p.Seed) + bw.printf("- Bootstrap: %d samples, %.0f%% CI, seed %d\n", p.BootstrapSamples, p.CILevel*100, p.Seed) + if p.WhatIf != "" { + bw.printf("- **What-if:** %s\n", p.WhatIf) + } + bw.printf("\n") bw.printf("## Budget checks\n\n") bw.printf("| check | limit | projected | status |\n") diff --git a/gate_cmd.go b/gate_cmd.go index ca606d3..4035cae 100644 --- a/gate_cmd.go +++ b/gate_cmd.go @@ -7,7 +7,6 @@ import ( "os" "augur/aggregate" - "augur/cost" "augur/gate" "augur/project" "augur/trace" @@ -21,6 +20,7 @@ func runGate(args []string) error { fs := flag.NewFlagSet("gate", flag.ContinueOnError) tracePath := fs.String("trace", "trace.jsonl", "path to the cost trace (JSONL)") pricingPath := fs.String("pricing", "pricing.yaml", "path to the pricing snapshot") + tcoPath := fs.String("tco", "", "derive pricing from a self-hosted TCO config instead of -pricing") trafficPath := fs.String("traffic", "traffic.yaml", "path to the traffic profile") budgetPath := fs.String("budget", "budget.yaml", "path to the budget thresholds") reportMD := fs.String("report-md", "report.md", "path to write the Markdown report (empty to skip)") @@ -28,15 +28,17 @@ func runGate(args []string) error { ciLevel := fs.Float64("ci", 0.95, "confidence level for bootstrap intervals (0..1)") bootstrap := fs.Int("bootstrap", 2000, "number of bootstrap resamples") seed := fs.Uint64("seed", 1, "PRNG seed for reproducible bootstrap intervals") + knobFs := addKnobFlags(fs) if err := fs.Parse(args); err != nil { return err } + knobs := knobFs.knobs() records, err := readTrace(*tracePath) if err != nil { return err } - pricing, err := cost.LoadPricing(*pricingPath) + pricing, err := resolvePricing(*pricingPath, *tcoPath) if err != nil { return err } @@ -49,7 +51,7 @@ func runGate(args []string) error { return err } - res, err := aggregate.Aggregate(records, pricing) + res, err := aggregate.AggregateWithKnobs(records, pricing, knobs) if err != nil { return err } @@ -59,6 +61,7 @@ func runGate(args []string) error { if err != nil { return err } + proj.WhatIf = describeKnobs(knobs) report := gate.NewReport(proj, gate.Evaluate(proj, budget)) diff --git a/knobs.go b/knobs.go new file mode 100644 index 0000000..250bbc8 --- /dev/null +++ b/knobs.go @@ -0,0 +1,53 @@ +package main + +import ( + "flag" + "fmt" + "strings" + + "augur/aggregate" +) + +// knobFlags registers the shared what-if multiplier flags on fs and returns +// accessors. They let `project` and `gate` re-cost a recorded trace under +// hypothetical agentic cost drivers without re-running the agent. +type knobFlags struct { + retry *float64 + fanout *float64 + growth *float64 +} + +func addKnobFlags(fs *flag.FlagSet) knobFlags { + return knobFlags{ + retry: fs.Float64("retry-rate", 0, "what-if: extra fraction of calls retried (0.2 = 20% more calls)"), + fanout: fs.Float64("fanout", 1, "what-if: multiplier on call count from sub-agent/tool fan-out"), + growth: fs.Float64("context-growth", 1, "what-if: multiplier on prompt cost from history growth"), + } +} + +func (k knobFlags) knobs() aggregate.Knobs { + return aggregate.Knobs{ + RetryRate: *k.retry, + FanoutFactor: *k.fanout, + ContextGrowth: *k.growth, + } +} + +// describeKnobs renders a one-line summary of non-identity knobs (empty for the +// observed-as-is case), for the projection banner and the gate report. +func describeKnobs(k aggregate.Knobs) string { + if k.IsIdentity() { + return "" + } + var parts []string + if k.RetryRate != 0 { + parts = append(parts, fmt.Sprintf("retry +%.0f%%", k.RetryRate*100)) + } + if k.FanoutFactor != 0 && k.FanoutFactor != 1 { + parts = append(parts, fmt.Sprintf("fanout ×%g", k.FanoutFactor)) + } + if k.ContextGrowth != 0 && k.ContextGrowth != 1 { + parts = append(parts, fmt.Sprintf("context ×%g", k.ContextGrowth)) + } + return strings.Join(parts, ", ") +} diff --git a/main.go b/main.go index 0fece25..c05973c 100644 --- a/main.go +++ b/main.go @@ -40,10 +40,11 @@ var commands = map[string]command{ "aggregate": {runAggregate, "summarize a cost trace into per-scenario distributions"}, "project": {runProject, "project a trace to production unit economics with CIs"}, "gate": {runGate, "check a projection against budget.yaml (exit 1 if over)"}, + "tco": {runTCO, "show effective $/Mtok for self-hosted models (TCO)"}, } // order fixes the usage listing (maps don't iterate deterministically). -var order = []string{"proxy", "run", "aggregate", "project", "gate"} +var order = []string{"proxy", "run", "aggregate", "project", "gate", "tco"} func main() { if len(os.Args) < 2 { diff --git a/pricing_source.go b/pricing_source.go new file mode 100644 index 0000000..cec0004 --- /dev/null +++ b/pricing_source.go @@ -0,0 +1,23 @@ +package main + +import ( + "fmt" + + "augur/cost" + "augur/tco" +) + +// resolvePricing loads the pricing the cost pipeline should use. When tcoPath is +// set, prices are derived from a self-hosted TCO config (effective $/Mtok from +// instance cost + throughput); otherwise they come from the pricing snapshot. +// The two are alternatives — tco takes precedence when both are provided. +func resolvePricing(pricingPath, tcoPath string) (cost.Pricing, error) { + if tcoPath != "" { + tc, err := tco.LoadTCO(tcoPath) + if err != nil { + return cost.Pricing{}, err + } + return tc.Pricing(fmt.Sprintf("tco (%s)", tcoPath)), nil + } + return cost.LoadPricing(pricingPath) +} diff --git a/project/project.go b/project/project.go index 2e9448f..0f57ad9 100644 --- a/project/project.go +++ b/project/project.go @@ -51,6 +51,11 @@ type Projection struct { // ByModel decomposes the average request's cost by model, sorted by cost. ByModel []ModelShare `json:"by_model"` + // WhatIf, when non-empty, describes the what-if knobs applied to the + // underlying aggregation (set by the caller). Empty means the projection is + // of the observed cost as-is. + WhatIf string `json:"what_if,omitempty"` + Traffic Traffic `json:"-"` } diff --git a/project/table.go b/project/table.go index 0d5a2ce..5dd7c1b 100644 --- a/project/table.go +++ b/project/table.go @@ -16,7 +16,11 @@ func (p Projection) WriteTable(w io.Writer) error { } fmt.Fprintf(w, "traffic: %d users x %g req/user/day, %d tenants, %d days/month\n", tr.Users, tr.RequestsPerUserPerDay, tr.Tenants, tr.DaysPerMonth) - fmt.Fprintf(w, "bootstrap: %d samples, %.0f%% CI, seed %d\n\n", p.BootstrapSamples, p.CILevel*100, p.Seed) + fmt.Fprintf(w, "bootstrap: %d samples, %.0f%% CI, seed %d\n", p.BootstrapSamples, p.CILevel*100, p.Seed) + if p.WhatIf != "" { + fmt.Fprintf(w, "%s\n", p.WhatIf) + } + fmt.Fprintln(w) tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', 0) fmt.Fprintf(tw, " metric\tvalue\t%.0f%% CI\n", p.CILevel*100) diff --git a/project_cmd.go b/project_cmd.go index 8db7f89..3c7e14a 100644 --- a/project_cmd.go +++ b/project_cmd.go @@ -7,7 +7,6 @@ import ( "os" "augur/aggregate" - "augur/cost" "augur/project" "augur/trace" ) @@ -19,14 +18,17 @@ func runProject(args []string) error { fs := flag.NewFlagSet("project", flag.ContinueOnError) tracePath := fs.String("trace", "trace.jsonl", "path to the cost trace (JSONL)") pricingPath := fs.String("pricing", "pricing.yaml", "path to the pricing snapshot") + tcoPath := fs.String("tco", "", "derive pricing from a self-hosted TCO config instead of -pricing") trafficPath := fs.String("traffic", "traffic.yaml", "path to the traffic profile") asJSON := fs.Bool("json", false, "emit the projection as JSON instead of a table") ciLevel := fs.Float64("ci", 0.95, "confidence level for bootstrap intervals (0..1)") bootstrap := fs.Int("bootstrap", 2000, "number of bootstrap resamples") seed := fs.Uint64("seed", 1, "PRNG seed for reproducible bootstrap intervals") + knobFs := addKnobFlags(fs) if err := fs.Parse(args); err != nil { return err } + knobs := knobFs.knobs() f, err := os.Open(*tracePath) if err != nil { @@ -38,7 +40,7 @@ func runProject(args []string) error { return err } - pricing, err := cost.LoadPricing(*pricingPath) + pricing, err := resolvePricing(*pricingPath, *tcoPath) if err != nil { return err } @@ -47,7 +49,7 @@ func runProject(args []string) error { return err } - res, err := aggregate.Aggregate(records, pricing) + res, err := aggregate.AggregateWithKnobs(records, pricing, knobs) if err != nil { return err } @@ -59,6 +61,7 @@ func runProject(args []string) error { if err != nil { return err } + proj.WhatIf = describeKnobs(knobs) if *asJSON { enc := json.NewEncoder(os.Stdout) diff --git a/proxy/proxy.go b/proxy/proxy.go index dff968d..5425b3f 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -24,6 +24,7 @@ import ( "sync" "time" + "augur/cassette" "augur/trace" ) @@ -39,6 +40,18 @@ const ( // time.Now) so tests can stamp deterministic timestamps. type nowFunc func() time.Time +// Mode selects how the proxy obtains responses. +type Mode int + +const ( + // ModeLive forwards to the real provider (the default). + ModeLive Mode = iota + // ModeRecord forwards to the provider AND saves each response to a cassette. + ModeRecord + // ModeReplay serves responses from a cassette without calling the provider. + ModeReplay +) + // Server is the recording proxy. Construct it with New and mount it with // http.ListenAndServe; it implements http.Handler. type Server struct { @@ -47,12 +60,16 @@ type Server struct { client *http.Client now nowFunc + mode Mode + cassette *cassette.Cassette + mu sync.Mutex seq map[string]int // per (scenario|run) next call ordinal } // New returns a Server that forwards to upstream (e.g. https://api.openai.com) -// and appends trace rows via tracer. A nil client uses a sensible default. +// and appends trace rows via tracer. A nil client uses a sensible default. The +// server starts in ModeLive; call Record or Replay to change mode. func New(upstream *url.URL, tracer *trace.Writer, client *http.Client) *Server { if client == nil { client = &http.Client{Timeout: 10 * time.Minute} @@ -66,6 +83,21 @@ func New(upstream *url.URL, tracer *trace.Writer, client *http.Client) *Server { } } +// Record switches the server to ModeRecord: responses are still fetched from +// the provider and relayed, but each is also saved to c. +func (s *Server) Record(c *cassette.Cassette) { + s.mode = ModeRecord + s.cassette = c +} + +// Replay switches the server to ModeReplay: responses are served from c and the +// provider is never contacted (no tokens spent). The cost trace is regenerated +// from the recorded responses. +func (s *Server) Replay(c *cassette.Cassette) { + s.mode = ModeReplay + s.cassette = c +} + // nextSeq returns and increments the call ordinal for a (scenario, run) pair. func (s *Server) nextSeq(scenario, run string) int { key := scenario + "|" + run @@ -91,6 +123,15 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { _ = r.Body.Close() reqModel := modelFromRequest(reqBody) + // Assign the call's ordinal once, up front: it is both the trace's seq and + // the cassette key, so record and replay must compute it identically. + seq := s.nextSeq(scenario, run) + + if s.mode == ModeReplay { + s.replay(w, scenario, run, seq, reqModel, r.URL.Path) + return + } + outReq, err := s.buildUpstreamRequest(r, reqBody) if err != nil { http.Error(w, "augur proxy: building upstream request: "+err.Error(), http.StatusBadGateway) @@ -105,64 +146,92 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { } defer resp.Body.Close() - // Dispatch on the response shape. A streamed (SSE) body must be relayed - // chunk-by-chunk so the agent sees tokens in real time; a regular JSON body - // is read in full. Both extract usage the same way, so totals reconcile. - var ( - usage oaiUsage - respModel string - ok bool - ) - if isEventStream(resp.Header) { - usage, respModel = s.streamResponse(w, resp) - } else if usage, respModel, ok = s.bufferResponse(w, resp); !ok { - return // bufferResponse already wrote an error to the client + // Live mode streams an SSE body back chunk-by-chunk so the agent sees tokens + // in real time. Record mode (and every non-streaming response) goes through + // the buffered path so the full body can be captured for the cassette. + if s.mode == ModeLive && isEventStream(resp.Header) { + usage, respModel := s.streamResponse(w, resp) + latency := s.now().Sub(start) + s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode) + return + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + http.Error(w, "augur proxy: reading upstream response: "+err.Error(), http.StatusBadGateway) + return } latency := s.now().Sub(start) + contentType := resp.Header.Get("Content-Type") + + copyHeader(w.Header(), resp.Header) + w.WriteHeader(resp.StatusCode) + _, _ = w.Write(body) - // Record the call. Parse failures (a non-JSON body, an error response with - // no usage) yield zero tokens but the row is still written — a burned-but- - // failed call is precisely the surprise the trace exists to surface. - model := reqModel - if model == "" { - model = respModel + if s.mode == ModeRecord && s.cassette != nil { + if err := s.cassette.Record(cassette.Entry{ + ScenarioID: scenario, RunID: run, Seq: seq, + Status: resp.StatusCode, ContentType: contentType, + LatencyMs: latency.Milliseconds(), Body: string(body), + }); err != nil { + fmt.Printf("augur proxy: WARNING cassette write failed: %v\n", err) + } } + + usage, respModel := extractUsage(contentType, body) + s.recordTrace(start, scenario, run, seq, pickModel(reqModel, respModel), usage, latency.Milliseconds(), r.URL.Path, resp.StatusCode) +} + +// replay serves a previously recorded response from the cassette without +// contacting the provider, and regenerates the call's trace row from it. A miss +// means the agent made a call that was not recorded — surfaced as a 502 so the +// divergence is loud rather than silently mis-costed. +func (s *Server) replay(w http.ResponseWriter, scenario, run string, seq int, reqModel, path string) { + e, ok := s.cassette.Lookup(scenario, run, seq) + if !ok { + http.Error(w, fmt.Sprintf("augur proxy: replay miss for scenario %q run %q seq %d (agent diverged from the recording?)", + scenario, run, seq), http.StatusBadGateway) + return + } + body := []byte(e.Body) + if e.ContentType != "" { + w.Header().Set("Content-Type", e.ContentType) + } + w.WriteHeader(e.Status) + _, _ = w.Write(body) + + usage, respModel := extractUsage(e.ContentType, body) + s.recordTrace(s.now(), scenario, run, seq, pickModel(reqModel, respModel), usage, e.LatencyMs, path, e.Status) +} + +// recordTrace writes one trace row. A write failure must not corrupt the +// agent's response but must be loud: a dropped row means an under-counted bill. +func (s *Server) recordTrace(ts time.Time, scenario, run string, seq int, model string, u oaiUsage, latencyMs int64, path string, status int) { rec := trace.Record{ - Timestamp: start.UTC().Format(time.RFC3339Nano), + Timestamp: ts.UTC().Format(time.RFC3339Nano), ScenarioID: scenario, RunID: run, - Seq: s.nextSeq(scenario, run), + Seq: seq, Model: model, - InputTokens: usage.InputTokens, - OutputTokens: usage.OutputTokens, - CachedTokens: usage.CachedTokens, - LatencyMs: latency.Milliseconds(), - Endpoint: r.URL.Path, - Status: resp.StatusCode, + InputTokens: u.InputTokens, + OutputTokens: u.OutputTokens, + CachedTokens: u.CachedTokens, + LatencyMs: latencyMs, + Endpoint: path, + Status: status, } if err := s.tracer.Write(rec); err != nil { - // A trace write failure must not corrupt the agent's response, but it - // must be loud: a silently dropped row means an under-counted bill. fmt.Printf("augur proxy: WARNING trace write failed: %v\n", err) } } -// bufferResponse relays a non-streaming response: it reads the whole body, -// copies headers/status, writes the body through unchanged, and parses usage. -// It returns ok=false (after writing a 502) only if the upstream body cannot be -// read — in which case the caller must not record a row. -func (s *Server) bufferResponse(w http.ResponseWriter, resp *http.Response) (oaiUsage, string, bool) { - body, err := io.ReadAll(resp.Body) - if err != nil { - http.Error(w, "augur proxy: reading upstream response: "+err.Error(), http.StatusBadGateway) - return oaiUsage{}, "", false +// pickModel prefers the request's model (it matches pricing.yaml keys) and falls +// back to the model echoed in the response. +func pickModel(reqModel, respModel string) string { + if reqModel != "" { + return reqModel } - copyHeader(w.Header(), resp.Header) - w.WriteHeader(resp.StatusCode) - _, _ = w.Write(body) - - usage, _, model := parseUsageJSON(body) - return usage, model, true + return respModel } // buildUpstreamRequest clones the inbound request onto the upstream base URL, diff --git a/proxy/proxy_stream.go b/proxy/proxy_stream.go index 8acbfc7..d986492 100644 --- a/proxy/proxy_stream.go +++ b/proxy/proxy_stream.go @@ -12,10 +12,45 @@ import ( // response's Content-Type rather than the request's stream flag because the // response is the source of truth for how the body must be relayed. func isEventStream(h http.Header) bool { - ct := h.Get("Content-Type") + return isEventStreamCT(h.Get("Content-Type")) +} + +// isEventStreamCT reports whether a Content-Type value denotes SSE. +func isEventStreamCT(ct string) bool { return strings.HasPrefix(strings.TrimSpace(strings.ToLower(ct)), "text/event-stream") } +// extractUsage parses token usage and the resolved model from a COMPLETE +// response body, handling both a single JSON object and a full SSE stream. The +// buffered and replay paths use it so a recorded streaming response is costed +// the same as it was live. +func extractUsage(contentType string, body []byte) (oaiUsage, string) { + if isEventStreamCT(contentType) { + return usageFromSSE(body) + } + return usageFromResponse(body) +} + +// usageFromSSE scans a buffered SSE body line-by-line for the usage block and +// the resolved model, mirroring what streamResponse captures while relaying. +func usageFromSSE(body []byte) (oaiUsage, string) { + var ( + usage oaiUsage + model string + ) + for line := range bytes.SplitSeq(body, []byte("\n")) { + if u, has, m := parseSSEChunk(line); m != "" || has { + if m != "" { + model = m + } + if has { + usage = u + } + } + } + return usage, model +} + // streamResponse relays an SSE body to the client chunk-by-chunk, flushing after // each line so the agent receives tokens in real time, while scanning for the // usage block that providers emit in the final chunk when the request set diff --git a/proxy/record_test.go b/proxy/record_test.go new file mode 100644 index 0000000..b9319fa --- /dev/null +++ b/proxy/record_test.go @@ -0,0 +1,204 @@ +package proxy + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "strings" + "testing" + "time" + + "augur/cassette" + "augur/trace" +) + +// doTagged POSTs body through proxyURL with scenario/run headers and returns the +// response body. +func doTagged(t *testing.T, proxyURL, scenario, run, body string) []byte { + t.Helper() + req, _ := http.NewRequest(http.MethodPost, proxyURL+"/v1/chat/completions", strings.NewReader(body)) + req.Header.Set(HeaderScenarioID, scenario) + req.Header.Set(HeaderRunID, run) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("request: %v", err) + } + defer resp.Body.Close() + out, _ := io.ReadAll(resp.Body) + return out +} + +// TestRecordThenReplay records two calls against a fake provider, then replays +// the cassette with the provider shut down — proving replay spends nothing and +// regenerates an identical trace. +func TestRecordThenReplay(t *testing.T) { + dir := t.TempDir() + cassettePath := filepath.Join(dir, "cassette.jsonl") + fixed := time.Date(2026, 6, 21, 18, 0, 0, 0, time.UTC) + + // --- Record phase --- + var upstreamHits int + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamHits++ + io.Copy(io.Discard, r.Body) + io.WriteString(w, chatResponse) // from proxy_test.go: usage 1234/567/200 + })) + upURL, _ := url.Parse(upstream.URL) + + cRec, err := cassette.Create(cassettePath) + if err != nil { + t.Fatalf("Create cassette: %v", err) + } + var recTrace bytes.Buffer + recSrv := New(upURL, trace.NewWriter(&recTrace), upstream.Client()) + recSrv.now = func() time.Time { return fixed } + recSrv.Record(cRec) + recProxy := httptest.NewServer(recSrv) + + doTagged(t, recProxy.URL, "checkout", "checkout-000", `{"model":"gpt-4o"}`) + doTagged(t, recProxy.URL, "checkout", "checkout-001", `{"model":"gpt-4o"}`) + + recProxy.Close() + cRec.Close() + upstream.Close() // provider is GONE for replay + + if upstreamHits != 2 { + t.Fatalf("record phase hit upstream %d times, want 2", upstreamHits) + } + + recRows, err := trace.ReadAll(&recTrace) + if err != nil { + t.Fatalf("read record trace: %v", err) + } + if len(recRows) != 2 { + t.Fatalf("record trace has %d rows, want 2", len(recRows)) + } + + // --- Replay phase --- (no upstream; would panic on a real call) + cPlay, err := cassette.Load(cassettePath) + if err != nil { + t.Fatalf("Load cassette: %v", err) + } + var playTrace bytes.Buffer + // A client that fails any request, to prove replay never calls out. + failClient := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + t.Fatal("replay must not contact the provider") + return nil, nil + })} + playSrv := New(upURL, trace.NewWriter(&playTrace), failClient) + playSrv.now = func() time.Time { return fixed } + playSrv.Replay(cPlay) + playProxy := httptest.NewServer(playSrv) + defer playProxy.Close() + + gotBody := doTagged(t, playProxy.URL, "checkout", "checkout-000", `{"model":"gpt-4o"}`) + doTagged(t, playProxy.URL, "checkout", "checkout-001", `{"model":"gpt-4o"}`) + + // The agent gets the recorded body back. + if string(gotBody) != chatResponse { + t.Errorf("replay body mismatch:\n got %q", gotBody) + } + + playRows, err := trace.ReadAll(&playTrace) + if err != nil { + t.Fatalf("read replay trace: %v", err) + } + if len(playRows) != 2 { + t.Fatalf("replay trace has %d rows, want 2", len(playRows)) + } + // Replay regenerates an identical trace (same tokens, model, latency). + for i := range recRows { + if recRows[i] != playRows[i] { + t.Errorf("row %d diverged:\n record %+v\n replay %+v", i, recRows[i], playRows[i]) + } + } +} + +// A replay miss (a call that wasn't recorded) must fail loudly with 502. +func TestReplayMiss(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "c.jsonl") + c, _ := cassette.Create(path) + c.Record(cassette.Entry{ScenarioID: "a", RunID: "a-000", Seq: 0, Status: 200, + ContentType: "application/json", Body: chatResponse}) + c.Close() + + loaded, _ := cassette.Load(path) + var tr bytes.Buffer + srv := New(&url.URL{}, trace.NewWriter(&tr), nil) + srv.Replay(loaded) + proxy := httptest.NewServer(srv) + defer proxy.Close() + + // seq 0 hits; the agent's second call (seq 1) was never recorded -> 502. + req1, _ := http.NewRequest(http.MethodPost, proxy.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o"}`)) + req1.Header.Set(HeaderScenarioID, "a") + req1.Header.Set(HeaderRunID, "a-000") + resp1, _ := http.DefaultClient.Do(req1) + resp1.Body.Close() + if resp1.StatusCode != 200 { + t.Errorf("first replayed call status = %d, want 200", resp1.StatusCode) + } + + req2, _ := http.NewRequest(http.MethodPost, proxy.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o"}`)) + req2.Header.Set(HeaderScenarioID, "a") + req2.Header.Set(HeaderRunID, "a-000") + resp2, _ := http.DefaultClient.Do(req2) + resp2.Body.Close() + if resp2.StatusCode != http.StatusBadGateway { + t.Errorf("replay miss status = %d, want 502", resp2.StatusCode) + } +} + +// Recording a streaming response must capture the SSE body and replay it with +// the same usage accounting. +func TestRecordReplayStreaming(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "stream.jsonl") + fixed := time.Date(2026, 6, 21, 18, 0, 0, 0, time.UTC) + + upstream := newStreamingUpstream(streamChunks) // usage 1234/567/200 + upURL, _ := url.Parse(upstream.URL) + + cRec, _ := cassette.Create(path) + var recTrace bytes.Buffer + recSrv := New(upURL, trace.NewWriter(&recTrace), nil) + recSrv.now = func() time.Time { return fixed } + recSrv.Record(cRec) + recProxy := httptest.NewServer(recSrv) + + doTagged(t, recProxy.URL, "s", "s-000", `{"model":"gpt-4o","stream":true}`) + recProxy.Close() + cRec.Close() + upstream.Close() + + recRows, _ := trace.ReadAll(&recTrace) + if len(recRows) != 1 || recRows[0].InputTokens != 1234 || recRows[0].OutputTokens != 567 { + t.Fatalf("record trace wrong: %+v", recRows) + } + + cPlay, _ := cassette.Load(path) + var playTrace bytes.Buffer + playSrv := New(upURL, trace.NewWriter(&playTrace), nil) + playSrv.now = func() time.Time { return fixed } + playSrv.Replay(cPlay) + playProxy := httptest.NewServer(playSrv) + defer playProxy.Close() + + body := doTagged(t, playProxy.URL, "s", "s-000", `{"model":"gpt-4o","stream":true}`) + if string(body) != streamChunks { + t.Errorf("replayed SSE body mismatch:\n got %q", body) + } + playRows, _ := trace.ReadAll(&playTrace) + if len(playRows) != 1 || playRows[0] != recRows[0] { + t.Errorf("streaming replay row diverged:\n record %+v\n replay %+v", recRows[0], playRows) + } +} + +// roundTripFunc adapts a function to http.RoundTripper. +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } diff --git a/run_cmd.go b/run_cmd.go index b801c62..100fbe8 100644 --- a/run_cmd.go +++ b/run_cmd.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "augur/cassette" "augur/proxy" "augur/runner" "augur/trace" @@ -20,6 +21,10 @@ import ( // runRun is the all-in-one Hito 2 command: it starts the recording proxy // in-process, drives the agent against scenarios.yaml N times through it, and // leaves a complete cost trace ready for `augur aggregate`. +// +// With -record it also saves every response to a cassette; with -replay it +// serves responses from a cassette and never contacts the provider (Hito 5), +// so re-running the gate in CI costs no tokens. func runRun(args []string) error { fs := flag.NewFlagSet("run", flag.ContinueOnError) scenariosPath := fs.String("scenarios", "scenarios.yaml", "path to the scenarios file") @@ -27,20 +32,31 @@ func runRun(args []string) error { tracePath := fs.String("trace", "trace.jsonl", "path to append the cost trace to (JSONL)") listen := fs.String("listen", "127.0.0.1:0", "address for the in-process proxy (default: random local port)") runs := fs.Int("runs", 0, "override the repetitions per scenario (0 = use scenarios.yaml)") - session := fs.String("session", "", "run-id prefix to keep repeated invocations distinct (default: timestamp)") + session := fs.String("session", "", "run-id prefix (default: timestamp live, empty for record/replay so ids are stable)") continueOnError := fs.Bool("continue-on-error", false, "keep going after an agent invocation fails") + record := fs.String("record", "", "record every response to this cassette file (real provider calls)") + replay := fs.String("replay", "", "replay responses from this cassette file (no provider calls, no tokens)") if err := fs.Parse(args); err != nil { return err } + if *record != "" && *replay != "" { + return fmt.Errorf("-record and -replay are mutually exclusive") + } + replaying := *replay != "" cfg, err := runner.LoadConfig(*scenariosPath) if err != nil { return err } + // In replay mode the provider is never contacted, so the upstream URL is + // irrelevant; otherwise it must be valid. up, err := url.Parse(strings.TrimRight(*upstream, "/")) if err != nil || up.Scheme == "" || up.Host == "" { - return fmt.Errorf("invalid -upstream %q: need scheme and host (e.g. https://api.openai.com)", *upstream) + if !replaying { + return fmt.Errorf("invalid -upstream %q: need scheme and host (e.g. https://api.openai.com)", *upstream) + } + up = &url.URL{} } tracer, err := trace.OpenFile(*tracePath) @@ -49,6 +65,15 @@ func runRun(args []string) error { } defer tracer.Close() + pxy := proxy.New(up, tracer, nil) + cass, err := configureCassette(pxy, *record, *replay) + if err != nil { + return err + } + if cass != nil { + defer cass.Close() + } + // Bind the proxy listener up front so we know the real address (handles a // :0 random port) before pointing the agent at it. ln, err := net.Listen("tcp", *listen) @@ -57,7 +82,7 @@ func runRun(args []string) error { } baseURL := "http://" + ln.Addr().String() - srv := &http.Server{Handler: proxy.New(up, tracer, nil)} + srv := &http.Server{Handler: pxy} serveErr := make(chan error, 1) go func() { serveErr <- srv.Serve(ln) }() @@ -66,12 +91,15 @@ func runRun(args []string) error { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) defer stop() + // Stable run ids matter for record/replay (the cassette is keyed by them), so + // default the session to empty there; live runs default to a timestamp to + // keep repeated invocations distinct in a shared trace. sess := *session - if sess == "" { + if sess == "" && *record == "" && !replaying { sess = time.Now().Format("20060102-150405") } - fmt.Printf("augur run: proxy on %s → %s, tracing to %s\n", baseURL, up.String(), *tracePath) + fmt.Printf("augur run: proxy on %s (%s), tracing to %s\n", baseURL, modeLabel(*record, *replay, up), *tracePath) fmt.Printf("augur run: %d scenario(s), %d run(s) each, session %q\n", len(cfg.Scenarios), effectiveRuns(cfg.Runs, *runs), sess) @@ -93,10 +121,49 @@ func runRun(args []string) error { if runErr != nil { return runErr } + if *record != "" { + fmt.Printf("augur run: cassette written to %s — replay it for free with: augur run -replay %s\n", *record, *record) + } fmt.Printf("augur run: trace written to %s — summarize it with: augur aggregate -trace %s\n", *tracePath, *tracePath) return nil } +// configureCassette puts the proxy into record or replay mode, opening the +// cassette file. It returns the cassette (for the caller to Close) or nil in +// live mode. +func configureCassette(pxy *proxy.Server, record, replay string) (*cassette.Cassette, error) { + switch { + case record != "": + c, err := cassette.Create(record) + if err != nil { + return nil, err + } + pxy.Record(c) + return c, nil + case replay != "": + c, err := cassette.Load(replay) + if err != nil { + return nil, err + } + pxy.Replay(c) + return c, nil + default: + return nil, nil + } +} + +// modeLabel describes the proxy's mode for the startup banner. +func modeLabel(record, replay string, up *url.URL) string { + switch { + case record != "": + return "RECORD → " + up.String() + ", cassette " + record + case replay != "": + return "REPLAY from " + replay + " (no provider calls)" + default: + return "LIVE → " + up.String() + } +} + // effectiveRuns reports the run count that will actually be used, for the // startup banner. func effectiveRuns(configRuns, override int) int { diff --git a/tco/table.go b/tco/table.go new file mode 100644 index 0000000..dffcd92 --- /dev/null +++ b/tco/table.go @@ -0,0 +1,30 @@ +package tco + +import ( + "fmt" + "io" + "sort" + "text/tabwriter" +) + +// WriteTable renders the derived effective $/Mtok for each deployment, with the +// inputs so the number can be checked at a glance. +func (t TCO) WriteTable(w io.Writer) error { + names := make([]string, 0, len(t.Deployments)) + for name := range t.Deployments { + names = append(names, name) + } + sort.Strings(names) + + if _, err := fmt.Fprintln(w, "Augur — self-hosted effective pricing (TCO)"); err != nil { + return err + } + tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, " deployment\t$/hr\ttok/s\tutil\teffective $/Mtok") + for _, name := range names { + d := t.Deployments[name] + fmt.Fprintf(tw, " %s\t%.2f\t%.0f\t%.0f%%\t$%.4f\n", + name, d.InstanceCostPerHour, d.TokensPerSec, d.Utilization*100, d.EffectivePerMtok()) + } + return tw.Flush() +} diff --git a/tco/tco.go b/tco/tco.go new file mode 100644 index 0000000..0fac357 --- /dev/null +++ b/tco/tco.go @@ -0,0 +1,124 @@ +// Package tco derives an effective $/Mtok price for self-hosted models from +// their total cost of ownership, so Augur can gate self-hosted deployments with +// the same machinery it uses for API pricing (SPEC Hito 5). +// +// For a hosted API you are quoted $/Mtok directly. For a model you run +// yourself, you instead pay for an instance by the hour and get some serving +// throughput; the effective token price is the standard benchmark-then-size +// calculation: +// +// $/Mtok = instance_$/hr × 1e6 / (tokens_per_sec × 3600 × utilization) +// +// Utilization is the fraction of paid time the instance is actually serving +// tokens. It matters: you pay for the box 24/7, but it is rarely saturated, and +// a half-idle instance doubles the effective token price. Modelling it is the +// honest part — ignoring it is how self-hosting looks cheaper than it is. +// +// A self-hosted price is token-type-agnostic: it is the same compute whether the +// tokens are prompt, completion, or cached, so input/output/cached all carry the +// effective rate. +package tco + +import ( + "fmt" + "os" + + "augur/cost" + + "gopkg.in/yaml.v3" +) + +// secondsPerHour and tokensPerMtok convert between the units a deployment is +// described in and the $/Mtok the pricing pipeline expects. +const ( + secondsPerHour = 3600.0 + tokensPerMtok = 1_000_000.0 +) + +// Deployment is one self-hosted model: what its instance costs and how fast it +// serves. The deployment name (the map key in tco.yaml) must match the model id +// the agent sends, so the derived price lines up with the trace. +type Deployment struct { + // InstanceCostPerHour is the all-in hourly cost of the serving instance(s). + InstanceCostPerHour float64 + // TokensPerSec is the measured aggregate serving throughput (prompt + + // completion tokens processed per second) at your batch size. + TokensPerSec float64 + // Utilization is the fraction of paid time the instance is serving (0..1]. + // Defaults to 1.0 (always busy) when omitted. + Utilization float64 +} + +// EffectivePerMtok is the derived $/Mtok for the deployment. +func (d Deployment) EffectivePerMtok() float64 { + return d.InstanceCostPerHour * tokensPerMtok / (d.TokensPerSec * secondsPerHour * d.Utilization) +} + +// TCO is a parsed tco.yaml: a set of named self-hosted deployments. +type TCO struct { + Deployments map[string]Deployment +} + +type yamlTCO struct { + Version int `yaml:"version"` + Deployments map[string]struct { + InstanceCostPerHour float64 `yaml:"instance_cost_per_hour"` + TokensPerSec float64 `yaml:"tokens_per_sec"` + Utilization *float64 `yaml:"utilization"` + } `yaml:"deployments"` +} + +// LoadTCO reads and validates a tco.yaml from path. +func LoadTCO(path string) (TCO, error) { + raw, err := os.ReadFile(path) + if err != nil { + return TCO{}, fmt.Errorf("tco: reading tco file: %w", err) + } + return ParseTCO(raw) +} + +// ParseTCO parses and validates a TCO config from YAML bytes. +func ParseTCO(data []byte) (TCO, error) { + var yt yamlTCO + if err := yaml.Unmarshal(data, &yt); err != nil { + return TCO{}, fmt.Errorf("tco: parsing tco yaml: %w", err) + } + if len(yt.Deployments) == 0 { + return TCO{}, fmt.Errorf("tco: config has no deployments") + } + + deployments := make(map[string]Deployment, len(yt.Deployments)) + for name, d := range yt.Deployments { + util := 1.0 + if d.Utilization != nil { + util = *d.Utilization + } + switch { + case d.InstanceCostPerHour < 0: + return TCO{}, fmt.Errorf("tco: %q instance_cost_per_hour must be >= 0, got %g", name, d.InstanceCostPerHour) + case d.TokensPerSec <= 0: + return TCO{}, fmt.Errorf("tco: %q tokens_per_sec must be > 0, got %g", name, d.TokensPerSec) + case util <= 0 || util > 1: + return TCO{}, fmt.Errorf("tco: %q utilization must be in (0, 1], got %g", name, util) + } + deployments[name] = Deployment{ + InstanceCostPerHour: d.InstanceCostPerHour, + TokensPerSec: d.TokensPerSec, + Utilization: util, + } + } + return TCO{Deployments: deployments}, nil +} + +// Pricing converts the deployments into a cost.Pricing: each deployment becomes +// a model priced at its effective $/Mtok for input, output, and cached tokens +// alike (self-hosted compute does not distinguish token types). snapshotDate is +// stamped onto the pricing for reporting. +func (t TCO) Pricing(snapshotDate string) cost.Pricing { + models := make(map[string]cost.ModelPrice, len(t.Deployments)) + for name, d := range t.Deployments { + eff := d.EffectivePerMtok() + models[name] = cost.ModelPrice{Input: eff, Output: eff, CachedInput: eff} + } + return cost.Pricing{SnapshotDate: snapshotDate, Models: models} +} diff --git a/tco/tco_test.go b/tco/tco_test.go new file mode 100644 index 0000000..76ce236 --- /dev/null +++ b/tco/tco_test.go @@ -0,0 +1,109 @@ +package tco + +import ( + "math" + "testing" +) + +func approx(a, b float64) bool { return math.Abs(a-b) < 1e-6 } + +func TestEffectivePerMtok(t *testing.T) { + // $36/hr, 2500 tok/s, 100% util: + // tokens/hr = 2500 * 3600 = 9,000,000 = 9 Mtok/hr + // $/Mtok = 36 / 9 = 4.00 + d := Deployment{InstanceCostPerHour: 36, TokensPerSec: 2500, Utilization: 1} + if !approx(d.EffectivePerMtok(), 4.0) { + t.Errorf("effective = %v, want 4.00", d.EffectivePerMtok()) + } + + // At 50% utilization the effective price doubles: $8.00/Mtok. + d.Utilization = 0.5 + if !approx(d.EffectivePerMtok(), 8.0) { + t.Errorf("effective at 50%% util = %v, want 8.00", d.EffectivePerMtok()) + } +} + +func TestParseTCO(t *testing.T) { + data := []byte(` +version: 1 +deployments: + llama-70b: + instance_cost_per_hour: 36 + tokens_per_sec: 2500 + utilization: 0.5 + mixtral: + instance_cost_per_hour: 12 + tokens_per_sec: 3000 +`) + tc, err := ParseTCO(data) + if err != nil { + t.Fatalf("ParseTCO: %v", err) + } + if len(tc.Deployments) != 2 { + t.Fatalf("got %d deployments, want 2", len(tc.Deployments)) + } + // Omitted utilization defaults to 1.0. + if tc.Deployments["mixtral"].Utilization != 1.0 { + t.Errorf("mixtral util = %v, want 1.0 default", tc.Deployments["mixtral"].Utilization) + } + if !approx(tc.Deployments["llama-70b"].EffectivePerMtok(), 8.0) { + t.Errorf("llama-70b effective = %v, want 8.00", tc.Deployments["llama-70b"].EffectivePerMtok()) + } +} + +func TestParseTCORejects(t *testing.T) { + cases := map[string]string{ + "no deployments": `version: 1`, + "zero throughput": ` +deployments: + m: + instance_cost_per_hour: 10 + tokens_per_sec: 0 +`, + "negative cost": ` +deployments: + m: + instance_cost_per_hour: -1 + tokens_per_sec: 100 +`, + "util over 1": ` +deployments: + m: + instance_cost_per_hour: 10 + tokens_per_sec: 100 + utilization: 1.5 +`, + "util zero": ` +deployments: + m: + instance_cost_per_hour: 10 + tokens_per_sec: 100 + utilization: 0 +`, + } + for name, data := range cases { + t.Run(name, func(t *testing.T) { + if _, err := ParseTCO([]byte(data)); err == nil { + t.Errorf("%s: expected error, got nil", name) + } + }) + } +} + +func TestPricingFromTCO(t *testing.T) { + tc := TCO{Deployments: map[string]Deployment{ + "llama-70b": {InstanceCostPerHour: 36, TokensPerSec: 2500, Utilization: 1}, + }} + p := tc.Pricing("2026-06-21") + if p.SnapshotDate != "2026-06-21" { + t.Errorf("SnapshotDate = %q", p.SnapshotDate) + } + mp, ok := p.Price("llama-70b") + if !ok { + t.Fatal("llama-70b missing from derived pricing") + } + // Self-hosted prices input/output/cached identically at the effective rate. + if !approx(mp.Input, 4.0) || !approx(mp.Output, 4.0) || !approx(mp.CachedInput, 4.0) { + t.Errorf("derived price = %+v, want all 4.00", mp) + } +} diff --git a/tco_cmd.go b/tco_cmd.go new file mode 100644 index 0000000..d0f680d --- /dev/null +++ b/tco_cmd.go @@ -0,0 +1,33 @@ +package main + +import ( + "encoding/json" + "flag" + "os" + + "augur/tco" +) + +// runTCO shows the effective $/Mtok derived from a self-hosted TCO config. Feed +// the same config to project/gate via their -tco flag to cost a trace against +// it. +func runTCO(args []string) error { + fs := flag.NewFlagSet("tco", flag.ContinueOnError) + tcoPath := fs.String("tco", "tco.yaml", "path to the TCO config") + asJSON := fs.Bool("json", false, "emit the derived pricing as JSON instead of a table") + if err := fs.Parse(args); err != nil { + return err + } + + tc, err := tco.LoadTCO(*tcoPath) + if err != nil { + return err + } + + if *asJSON { + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + return enc.Encode(tc.Pricing("tco (" + *tcoPath + ")")) + } + return tc.WriteTable(os.Stdout) +}