Skip to content

Commit ce86223

Browse files
yashmehrotrayashmehrotra
committed
refactor: ai fixture benchmarking
1 parent 3f5cbdd commit ce86223

16 files changed

Lines changed: 1869 additions & 1227 deletions

File tree

‎examples/ai-fixtures/cluster-health.yaml‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
name: mission-control-investigate
2+
description: |
3+
Get cluster health
4+
prompt: |
5+
Tell me about the general health of the kubernetes cluster, whats happening and what changed recently
6+
baseline: direct
7+
repeat: 1
8+
captureKubernetesProxy: true
9+
mcpProxy:
10+
capture: true
11+
headers:
12+
Accept: text/markdown
13+
14+
defaults:
15+
timeout: 3m
16+
permissionMode: bypassPermissions
17+
noSessionPersistence: true
18+
model: claude-sonnet-4-6
19+
20+
runs:
21+
- name: direct
22+
promptCaching: false
23+
tools: [Bash, Read]
24+
allowedTools:
25+
- Bash(kubectl *)
26+
- Bash(aws *)
27+
- Bash(curl *)
28+
- Read
29+
30+
- name: mission-control
31+
promptCaching: true
32+
tools: [default]
33+
mcpConfig:
34+
- .mcp.json
35+
allowedTools:
36+
- mcp__mission-control__*

‎pkg/ai/fixture/aggregate.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
// ABOUTME: Per-Run statistics accumulator — adds Summary records over N iterations and finalizes means.
2+
3+
package fixture
4+
5+
import "math"
6+
7+
type aggregate struct {
8+
N int
9+
Success bool
10+
Error string
11+
SessionID string
12+
DurationMean float64
13+
DurationStdev float64
14+
CostMean float64
15+
16+
Input int
17+
Output int
18+
CacheRead int
19+
CacheWrite int
20+
21+
ToolCalls int
22+
MCPCalls int
23+
BashCalls int
24+
KubectlCalls int
25+
KubectlAPICalls int
26+
MCPAPICalls int
27+
28+
KubectlAPILog []KubectlAPIEntry
29+
MCPAPILog []MCPAPIEntry
30+
ToolCallLog []ToolCallEntry
31+
ToolCounts map[string]int
32+
33+
durations []float64
34+
costs []float64
35+
}
36+
37+
func (a *aggregate) add(s Summary) {
38+
a.N++
39+
a.durations = append(a.durations, s.DurationMS)
40+
a.costs = append(a.costs, s.CostUSD)
41+
if s.SessionID != "" {
42+
a.SessionID = s.SessionID
43+
}
44+
if s.Error != "" && a.Error == "" {
45+
a.Error = s.Error
46+
}
47+
if s.Success {
48+
a.Success = true
49+
}
50+
a.Input += s.Input
51+
a.Output += s.Output
52+
a.CacheRead += s.CacheRead
53+
a.CacheWrite += s.CacheWrite
54+
a.ToolCalls += s.ToolCalls
55+
a.MCPCalls += s.MCPCalls
56+
a.BashCalls += s.BashCalls
57+
a.KubectlCalls += s.KubectlCalls
58+
a.KubectlAPICalls += s.KubectlAPICalls
59+
a.KubectlAPILog = append(a.KubectlAPILog, s.KubectlAPILog...)
60+
a.MCPAPICalls += s.MCPAPICalls
61+
a.MCPAPILog = append(a.MCPAPILog, s.MCPAPILog...)
62+
a.ToolCallLog = append(a.ToolCallLog, s.ToolCallLog...)
63+
if a.ToolCounts == nil {
64+
a.ToolCounts = map[string]int{}
65+
}
66+
for k, v := range s.ToolCounts {
67+
a.ToolCounts[k] += v
68+
}
69+
}
70+
71+
func (a *aggregate) finalize() {
72+
if a.N == 0 {
73+
return
74+
}
75+
a.DurationMean = mean(a.durations)
76+
a.DurationStdev = stdev(a.durations, a.DurationMean)
77+
a.CostMean = mean(a.costs)
78+
// Per-iteration averages for token/tool aggregate counts.
79+
a.Input /= a.N
80+
a.Output /= a.N
81+
a.CacheRead /= a.N
82+
a.CacheWrite /= a.N
83+
a.ToolCalls /= a.N
84+
a.MCPCalls /= a.N
85+
a.BashCalls /= a.N
86+
a.KubectlCalls /= a.N
87+
a.KubectlAPICalls /= a.N
88+
a.MCPAPICalls /= a.N
89+
for k, v := range a.ToolCounts {
90+
a.ToolCounts[k] = v / a.N
91+
}
92+
}
93+
94+
func mean(xs []float64) float64 {
95+
if len(xs) == 0 {
96+
return 0
97+
}
98+
var s float64
99+
for _, x := range xs {
100+
s += x
101+
}
102+
return s / float64(len(xs))
103+
}
104+
105+
func stdev(xs []float64, m float64) float64 {
106+
if len(xs) < 2 {
107+
return 0
108+
}
109+
var s float64
110+
for _, x := range xs {
111+
d := x - m
112+
s += d * d
113+
}
114+
return math.Sqrt(s / float64(len(xs)-1))
115+
}

‎pkg/ai/fixture/api_log.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
// ABOUTME: Parsers for the JSONL files produced by kubeproxy and mcpproxy.
2+
// ABOUTME: Tool calls themselves come from the live stream-json output, not from these.
3+
4+
package fixture
5+
6+
import (
7+
"bufio"
8+
"bytes"
9+
"encoding/json"
10+
"os"
11+
"time"
12+
)
13+
14+
// scanJSONL streams a JSONL file and feeds each line to decode. decode returns
15+
// (parsed, true) when the line should be appended to the result, or (_, false)
16+
// when it should be skipped. Returns an empty slice when the file is missing
17+
// or unreadable — callers treat that as "no log available."
18+
func scanJSONL[T any](path string, decode func([]byte) (T, bool)) []T {
19+
data, err := os.ReadFile(path)
20+
if err != nil {
21+
return nil
22+
}
23+
scanner := bufio.NewScanner(bytes.NewReader(data))
24+
scanner.Buffer(make([]byte, jsonlScannerInitialBuf), jsonlScannerMaxBuf)
25+
var out []T
26+
for scanner.Scan() {
27+
v, ok := decode(scanner.Bytes())
28+
if !ok {
29+
continue
30+
}
31+
out = append(out, v)
32+
}
33+
return out
34+
}
35+
36+
// readKubectlAPILog parses a kubeproxy JSONL file into chronological API entries.
37+
func readKubectlAPILog(path string) []KubectlAPIEntry {
38+
return scanJSONL(path, decodeKubectlEntry)
39+
}
40+
41+
func decodeKubectlEntry(line []byte) (KubectlAPIEntry, bool) {
42+
var ev struct {
43+
Type string `json:"type"`
44+
Time time.Time `json:"time"`
45+
Method string `json:"method"`
46+
Path string `json:"path"`
47+
Query string `json:"query"`
48+
Status int `json:"status"`
49+
Duration string `json:"duration"`
50+
Bytes int64 `json:"bytes"`
51+
}
52+
if err := json.Unmarshal(line, &ev); err != nil || ev.Type != "request" {
53+
return KubectlAPIEntry{}, false
54+
}
55+
return KubectlAPIEntry{
56+
Time: ev.Time,
57+
Method: ev.Method,
58+
URL: joinURL(ev.Path, ev.Query),
59+
Status: ev.Status,
60+
Duration: ev.Duration,
61+
Bytes: ev.Bytes,
62+
}, true
63+
}
64+
65+
// readMCPAPILog parses an mcpproxy JSONL file into chronological API entries.
66+
func readMCPAPILog(path string) []MCPAPIEntry {
67+
return scanJSONL(path, decodeMCPEntry)
68+
}
69+
70+
func decodeMCPEntry(line []byte) (MCPAPIEntry, bool) {
71+
var ev struct {
72+
Type string `json:"type"`
73+
Time time.Time `json:"time"`
74+
Server string `json:"server"`
75+
Method string `json:"method"`
76+
Path string `json:"path"`
77+
Query string `json:"query"`
78+
Status int `json:"status"`
79+
Duration string `json:"duration"`
80+
Bytes int64 `json:"bytes"`
81+
RPCMethod string `json:"rpcMethod"`
82+
Tool string `json:"tool"`
83+
}
84+
if err := json.Unmarshal(line, &ev); err != nil || ev.Type != "request" {
85+
return MCPAPIEntry{}, false
86+
}
87+
return MCPAPIEntry{
88+
Time: ev.Time,
89+
Server: ev.Server,
90+
Method: ev.Method,
91+
URL: joinURL(ev.Path, ev.Query),
92+
RPCMethod: ev.RPCMethod,
93+
Tool: ev.Tool,
94+
Status: ev.Status,
95+
Duration: ev.Duration,
96+
Bytes: ev.Bytes,
97+
}, true
98+
}
99+
100+
func joinURL(path, query string) string {
101+
if query == "" {
102+
return path
103+
}
104+
return path + "?" + query
105+
}

‎pkg/ai/fixture/cost.go‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// ABOUTME: Cost resolution and money/byte formatters.
2+
// ABOUTME: Falls back to OpenRouter pricing when the claude CLI doesn't report a cost.
3+
4+
package fixture
5+
6+
import (
7+
"fmt"
8+
9+
"github.com/flanksource/captain/pkg/ai/pricing"
10+
)
11+
12+
// resolveCost returns the per-iteration cost for a row. If the underlying
13+
// claude CLI didn't report a cost (CostMean == 0) we fall back to the
14+
// OpenRouter-backed pricing registry to estimate from the token counts.
15+
// Returns (cost, estimated) — estimated is true when the value is from the
16+
// fallback path.
17+
func resolveCost(model string, a *aggregate) (float64, bool) {
18+
if a.CostMean > 0 {
19+
return a.CostMean, false
20+
}
21+
if a.Input == 0 && a.Output == 0 && a.CacheRead == 0 && a.CacheWrite == 0 {
22+
return 0, false
23+
}
24+
res, err := pricing.CalculateCost(model, a.Input, a.Output, 0, a.CacheRead, a.CacheWrite)
25+
if err != nil {
26+
return 0, false
27+
}
28+
return res.TotalCost, true
29+
}
30+
31+
func formatCostWithEstimate(v float64, estimated bool) string {
32+
s := formatUSD(v)
33+
if estimated && v > 0 {
34+
return s + " (est)"
35+
}
36+
return s
37+
}
38+
39+
func formatUSD(v float64) string {
40+
if v == 0 {
41+
return "$0"
42+
}
43+
return fmt.Sprintf("$%.4f", v)
44+
}
45+
46+
// humanBytes renders a byte count with KB/MB/GB/... suffixes.
47+
func humanBytes(n int64) string {
48+
const unit = 1024
49+
if n < unit {
50+
return fmt.Sprintf("%d B", n)
51+
}
52+
div, exp := int64(unit), 0
53+
for x := n / unit; x >= unit; x /= unit {
54+
div *= unit
55+
exp++
56+
}
57+
suffix := "KMGTPE"[exp]
58+
return fmt.Sprintf("%.1f %cB", float64(n)/float64(div), suffix)
59+
}

‎pkg/ai/fixture/cost_test.go‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package fixture
2+
3+
import "testing"
4+
5+
func TestResolveCost_UsesReportedWhenPresent(t *testing.T) {
6+
a := &aggregate{CostMean: 0.42, Input: 100, Output: 200}
7+
got, est := resolveCost("claude-sonnet-4-6", a)
8+
if got != 0.42 || est {
9+
t.Errorf("got %v, est=%v, want 0.42, false", got, est)
10+
}
11+
}
12+
13+
func TestResolveCost_FallbackZeroWhenNoTokens(t *testing.T) {
14+
a := &aggregate{}
15+
got, est := resolveCost("claude-sonnet-4-6", a)
16+
if got != 0 || est {
17+
t.Errorf("got %v, est=%v, want 0, false", got, est)
18+
}
19+
}
20+
21+
func TestResolveCost_FallbackComputesFromTokens(t *testing.T) {
22+
// claude-sonnet-4-6 pricing: $3/M input, $15/M output. Cache R/W: $0.30 / $3.75 per M.
23+
a := &aggregate{Input: 1_000_000, Output: 1_000_000}
24+
got, est := resolveCost("claude-sonnet-4-6", a)
25+
if !est {
26+
t.Fatal("estimated should be true on fallback")
27+
}
28+
want := 3.0 + 15.0
29+
if got != want {
30+
t.Errorf("got %v, want %v", got, want)
31+
}
32+
}
33+
34+
func TestResolveCost_FallbackUnknownModelReturnsZero(t *testing.T) {
35+
a := &aggregate{Input: 100, Output: 100}
36+
got, est := resolveCost("definitely-not-a-real-model-zzz", a)
37+
if got != 0 || est {
38+
t.Errorf("got %v, est=%v, want 0, false (model unknown → no estimate)", got, est)
39+
}
40+
}
41+
42+
func TestFormatCostWithEstimate(t *testing.T) {
43+
if got := formatCostWithEstimate(0, false); got != "$0" {
44+
t.Errorf("zero cost: got %q", got)
45+
}
46+
if got := formatCostWithEstimate(0.0234, false); got != "$0.0234" {
47+
t.Errorf("real cost: got %q", got)
48+
}
49+
if got := formatCostWithEstimate(0.0234, true); got != "$0.0234 (est)" {
50+
t.Errorf("estimated cost: got %q", got)
51+
}
52+
if got := formatCostWithEstimate(0, true); got != "$0" {
53+
t.Errorf("zero estimated should not show (est) suffix: got %q", got)
54+
}
55+
}
56+
57+
func TestHumanBytes(t *testing.T) {
58+
cases := map[int64]string{
59+
0: "0 B",
60+
512: "512 B",
61+
1024: "1.0 KB",
62+
1024 * 1024: "1.0 MB",
63+
1536: "1.5 KB",
64+
1024 * 1024 * 5: "5.0 MB",
65+
}
66+
for in, want := range cases {
67+
if got := humanBytes(in); got != want {
68+
t.Errorf("humanBytes(%d) = %q, want %q", in, got, want)
69+
}
70+
}
71+
}

0 commit comments

Comments
 (0)