tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 31 Packages Projects Releases Wiki Activity

fix(evals): report cache-aware Harbor costs (#10430)

Browse Source
This commit is contained in:
filip
2026-07-17 22:56:59 -07:00
committed by GitHub
parent 445da5dcdf
commit 8e78960e53
3 changed files with 182 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+44 -22
View File
@@ -79,6 +79,12 @@ struct JsonMetadata {
input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
output_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cache_read_input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cache_write_input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cost_usd: Option<f64>,
status: String,
}
@@ -102,6 +108,12 @@ enum StreamEvent {
input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
output_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cache_read_input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cache_write_input_tokens: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
cost_usd: Option<f64>,
},
}
@@ -1408,28 +1420,25 @@ impl CliSession {
.agent
.config
.session_manager
.get_session(&self.session_id, false)
.get_session_usage_totals(&self.session_id)
.await
{
Ok(session) => JsonMetadata {
total_tokens: session
.accumulated_usage
.total_tokens
.or(session.usage.total_tokens),
input_tokens: session
.accumulated_usage
.input_tokens
.or(session.usage.input_tokens),
output_tokens: session
.accumulated_usage
.output_tokens
.or(session.usage.output_tokens),
Ok(totals) => JsonMetadata {
total_tokens: totals.accumulated_usage.total_tokens,
input_tokens: totals.accumulated_usage.input_tokens,
output_tokens: totals.accumulated_usage.output_tokens,
cache_read_input_tokens: totals.accumulated_usage.cache_read_input_tokens,
cache_write_input_tokens: totals.accumulated_usage.cache_write_input_tokens,
cost_usd: totals.accumulated_cost,
status: "completed".to_string(),
},
Err(_) => JsonMetadata {
total_tokens: None,
input_tokens: None,
output_tokens: None,
cache_read_input_tokens: None,
cache_write_input_tokens: None,
cost_usd: None,
status: "completed".to_string(),
},
};
@@ -1439,25 +1448,38 @@ impl CliSession {
};
println!("{}", serde_json::to_string_pretty(&json_output)?);
} else if is_stream_json_mode {
let session = self
let totals = self
.agent
.config
.session_manager
.get_session(&self.session_id, false)
.get_session_usage_totals(&self.session_id)
.await
.ok();
let (total_tokens, input_tokens, output_tokens) = match session {
Some(s) => (
s.accumulated_usage.total_tokens.or(s.usage.total_tokens),
s.accumulated_usage.input_tokens.or(s.usage.input_tokens),
s.accumulated_usage.output_tokens.or(s.usage.output_tokens),
let (
total_tokens,
input_tokens,
output_tokens,
cache_read_input_tokens,
cache_write_input_tokens,
cost_usd,
) = match totals {
Some(totals) => (
totals.accumulated_usage.total_tokens,
totals.accumulated_usage.input_tokens,
totals.accumulated_usage.output_tokens,
totals.accumulated_usage.cache_read_input_tokens,
totals.accumulated_usage.cache_write_input_tokens,
totals.accumulated_cost,
),
None => (None, None, None),
None => (None, None, None, None, None, None),
};
emit_stream_event(&StreamEvent::Complete {
total_tokens,
input_tokens,
output_tokens,
cache_read_input_tokens,
cache_write_input_tokens,
cost_usd,
});
} else {
println!();
evals/harbor/agent.py
+45 -16
View File
@@ -247,24 +247,23 @@ class GooseBinaryAgent(Goose):
)
@staticmethod
def _extract_complete_event_tokens(
log_text: str,
) -> tuple[int | None, int | None, int | None]:
total = inp = out = None
def _extract_complete_event(log_text: str) -> dict | None:
complete_event = None
for line in log_text.strip().split("\n"):
line = line.strip()
if not line or '"complete"' not in line:
continue
event = json.loads(line)
if event.get("type") != "complete":
continue
total = event.get("total_tokens")
inp = event.get("input_tokens")
out = event.get("output_tokens")
return total, inp, out
if event.get("type") == "complete":
complete_event = event
return complete_event
def _compute_cost_from_pricing(
self, prompt_tokens: int | None, completion_tokens: int | None
self,
prompt_tokens: int | None,
completion_tokens: int | None,
cache_read_tokens: int | None,
cache_write_tokens: int | None,
) -> float | None:
if not self.model_name or not (prompt_tokens or completion_tokens):
return None
@@ -280,9 +279,17 @@ class GooseBinaryAgent(Goose):
break
if pricing is None:
return None
return (prompt_tokens or 0) * (pricing.get("input_cost_per_token") or 0.0) + (
completion_tokens or 0
) * (pricing.get("output_cost_per_token") or 0.0)
input_cost = pricing.get("input_cost_per_token") or 0.0
cache_read = cache_read_tokens or 0
cache_write = cache_write_tokens or 0
uncached = max((prompt_tokens or 0) - cache_read - cache_write, 0)
return (
uncached * input_cost
+ cache_read * (pricing.get("cache_read_input_token_cost") or input_cost)
+ cache_write
* (pricing.get("cache_creation_input_token_cost") or input_cost)
+ (completion_tokens or 0) * (pricing.get("output_cost_per_token") or 0.0)
)
def populate_context_post_run(self, context: AgentContext) -> None:
super().populate_context_post_run(context)
@@ -290,11 +297,33 @@ class GooseBinaryAgent(Goose):
if not txt_path.exists():
return
log_text = txt_path.read_text()
_total, inp, out = self._extract_complete_event_tokens(log_text)
complete_event = self._extract_complete_event(log_text)
if complete_event is None:
return
inp = complete_event.get("input_tokens")
out = complete_event.get("output_tokens")
cache_read = complete_event.get("cache_read_input_tokens")
cache_write = complete_event.get("cache_write_input_tokens")
if inp is not None:
context.n_input_tokens = inp
if out is not None:
context.n_output_tokens = out
cost = self._compute_cost_from_pricing(inp, out)
if cache_read is not None:
context.n_cache_tokens = cache_read
if cache_read is not None or cache_write is not None:
context.metadata = {
**(context.metadata or {}),
"cache_read_input_tokens": cache_read or 0,
"cache_write_input_tokens": cache_write or 0,
}
cost = complete_event.get("cost_usd")
if cost is None:
cost = self._compute_cost_from_pricing(inp, out, cache_read, cache_write)
if cost is not None:
context.metadata = {
**(context.metadata or {}),
"cost_source": "litellm_estimate",
}
if cost is not None:
context.cost_usd = cost
evals/harbor/reporter.py
+93 -32
View File
@@ -32,6 +32,15 @@ class LoadedJob:
return self.summary.started_at
@dataclass(frozen=True)
class TokenTotals:
input_tokens: int | None
cache_read_tokens: int | None
cache_write_tokens: int | None
output_tokens: int | None
cost_usd: float | None
def load_job(job_dir: Path) -> LoadedJob:
summary = JobResult.model_validate_json((job_dir / "result.json").read_text())
results: list[TrialResult] = []
@@ -65,9 +74,25 @@ def trial_duration(trial: TrialResult) -> float | None:
return (trial.finished_at - trial.started_at).total_seconds()
def trial_token_totals(trial: TrialResult) -> tuple[int | None, int | None, float | None]:
n_in, _n_cache, n_out, cost = trial.compute_token_cost_totals()
return n_in, n_out, cost
def trial_token_totals(trial: TrialResult) -> TokenTotals:
n_in, n_cache_read, n_out, cost = trial.compute_token_cost_totals()
if trial.agent_result is not None:
contexts = [trial.agent_result]
else:
contexts = [
step.agent_result
for step in trial.step_results or []
if step.agent_result is not None
]
cache_write_values = [
context.metadata.get("cache_write_input_tokens")
for context in contexts
if context.metadata is not None
and context.metadata.get("cache_write_input_tokens") is not None
]
n_cache_write = sum(cache_write_values) if cache_write_values else None
return TokenTotals(n_in, n_cache_read, n_cache_write, n_out, cost)
def _trial_dir(trial: TrialResult, job_dir: Path) -> Path:
@@ -145,12 +170,14 @@ def job_turn_totals(job: LoadedJob) -> int:
return sum((trial_turns(t, job.job_dir) or 0) for t in job.results)
def job_token_totals(job: LoadedJob) -> tuple[int, int, float]:
def job_token_totals(job: LoadedJob) -> TokenTotals:
totals = [trial_token_totals(t) for t in job.results]
return (
sum((n_in or 0) for n_in, _, _ in totals),
sum((n_out or 0) for _, n_out, _ in totals),
sum((c or 0.0) for _, _, c in totals),
return TokenTotals(
input_tokens=sum((total.input_tokens or 0) for total in totals),
cache_read_tokens=sum((total.cache_read_tokens or 0) for total in totals),
cache_write_tokens=sum((total.cache_write_tokens or 0) for total in totals),
output_tokens=sum((total.output_tokens or 0) for total in totals),
cost_usd=sum((total.cost_usd or 0.0) for total in totals),
)
@@ -250,7 +277,7 @@ def cmd_list(args: argparse.Namespace) -> int:
counts = status_counts(job.results)
total = len(job.results)
rate = f"{100 * counts['pass'] / total:.1f}%" if total else "-"
tok_in, tok_out, cost = job_token_totals(job)
tokens = job_token_totals(job)
breakdown = f"{counts['pass']}/{counts['fail']}/{counts['error']}/{counts['timeout']}"
rows.append(
(
@@ -258,10 +285,12 @@ def cmd_list(args: argparse.Namespace) -> int:
job_model(job),
rate,
fmt_duration(job_duration(job)),
fmt_tokens(tok_in),
fmt_tokens(tok_out),
fmt_tokens(tokens.input_tokens),
fmt_tokens(tokens.cache_read_tokens),
fmt_tokens(tokens.cache_write_tokens),
fmt_tokens(tokens.output_tokens),
fmt_tokens(job_turn_totals(job)),
fmt_cost(cost),
fmt_cost(tokens.cost_usd),
breakdown,
)
)
@@ -271,13 +300,15 @@ def cmd_list(args: argparse.Namespace) -> int:
return 0
print(
f"{'job_name':<40} {'model':<25} {'rate':>7} {'compute':>8} "
f"{'in':>7} {'out':>7} {'turns':>6} {'cost':>8} {'pass/fail/err/tout':>18}"
f"{'in':>7} {'read':>7} {'write':>7} {'out':>7} "
f"{'turns':>6} {'cost':>8} {'pass/fail/err/tout':>18}"
)
print("-" * 131)
print("-" * 147)
for row in rows:
print(
f"{row[0]:<40} {row[1]:<25} {row[2]:>7} {row[3]:>8} "
f"{row[4]:>7} {row[5]:>7} {row[6]:>6} {row[7]:>8} {row[8]:>18}"
f"{row[4]:>7} {row[5]:>7} {row[6]:>7} {row[7]:>7} "
f"{row[8]:>6} {row[9]:>8} {row[10]:>18}"
)
return 0
@@ -298,16 +329,22 @@ def cmd_show(args: argparse.Namespace) -> int:
)
if total:
print(f"Pass rate: {100 * counts['pass'] / total:.1f}%")
total_in, total_out, total_cost = job_token_totals(job)
print(f"Tokens: in={fmt_tokens(total_in)} out={fmt_tokens(total_out)}")
tokens = job_token_totals(job)
print(
f"Tokens: in={fmt_tokens(tokens.input_tokens)} "
f"cache-read={fmt_tokens(tokens.cache_read_tokens)} "
f"cache-write={fmt_tokens(tokens.cache_write_tokens)} "
f"out={fmt_tokens(tokens.output_tokens)}"
)
print(f"Turns: {fmt_tokens(job_turn_totals(job))}")
print(f"Cost: {fmt_cost(total_cost)}")
print(f"Cost: {fmt_cost(tokens.cost_usd)}")
print()
print(
f"{'task':<45} {'status':<10} {'reward':>7} {'dur':>7} "
f"{'in':>7} {'out':>7} {'turns':>6} {'cost':>7} error"
f"{'in':>7} {'read':>7} {'write':>7} {'out':>7} "
f"{'turns':>6} {'cost':>7} error"
)
print("-" * 137)
print("-" * 153)
for trial in sorted(job.results, key=task_name):
status = trial_status(trial)
if args.status and status != args.status:
@@ -323,16 +360,18 @@ def cmd_show(args: argparse.Namespace) -> int:
err_str = ""
if len(err_str) > 50:
err_str = err_str[:47] + "..."
n_in, n_out, cost = trial_token_totals(trial)
tokens = trial_token_totals(trial)
turns = trial_turns(trial, job.job_dir)
turns_str = str(turns) if turns is not None else "-"
print(
f"{task_name(trial):<45} {status:<10} {reward_str:>7} "
f"{fmt_duration(trial_duration(trial)):>7} "
f"{fmt_tokens(n_in):>7} "
f"{fmt_tokens(n_out):>7} "
f"{fmt_tokens(tokens.input_tokens):>7} "
f"{fmt_tokens(tokens.cache_read_tokens):>7} "
f"{fmt_tokens(tokens.cache_write_tokens):>7} "
f"{fmt_tokens(tokens.output_tokens):>7} "
f"{turns_str:>6} "
f"{fmt_cost(cost):>7} {err_str}"
f"{fmt_cost(tokens.cost_usd):>7} {err_str}"
)
return 0
@@ -354,11 +393,16 @@ def cmd_task(args: argparse.Namespace) -> int:
print(f"Duration: {fmt_duration(trial_duration(trial))}")
print(f"Started: {trial.started_at}")
print(f"Ended: {trial.finished_at}")
n_in, n_out, cost = trial_token_totals(trial)
print(f"Tokens: in={fmt_tokens(n_in)} out={fmt_tokens(n_out)}")
tokens = trial_token_totals(trial)
print(
f"Tokens: in={fmt_tokens(tokens.input_tokens)} "
f"cache-read={fmt_tokens(tokens.cache_read_tokens)} "
f"cache-write={fmt_tokens(tokens.cache_write_tokens)} "
f"out={fmt_tokens(tokens.output_tokens)}"
)
turns = trial_turns(trial, job_dir)
print(f"Turns: {turns if turns is not None else '-'}")
print(f"Cost: {fmt_cost(cost)}")
print(f"Cost: {fmt_cost(tokens.cost_usd)}")
error = trial_error(trial)
if error is not None:
exception_class, message = error
@@ -480,13 +524,30 @@ def cmd_compare(args: argparse.Namespace) -> int:
if na and nb:
row("pass rate %", 100 * ca["pass"] / na, 100 * cb["pass"] / nb, "{:.1f}")
a_in, a_out, a_cost = job_token_totals(job_a)
b_in, b_out, b_cost = job_token_totals(job_b)
print(f"{'tokens in':<18} {fmt_tokens(a_in):>10} {fmt_tokens(b_in):>10}")
print(f"{'tokens out':<18} {fmt_tokens(a_out):>10} {fmt_tokens(b_out):>10}")
a_tokens = job_token_totals(job_a)
b_tokens = job_token_totals(job_b)
print(
f"{'tokens in':<18} {fmt_tokens(a_tokens.input_tokens):>10} "
f"{fmt_tokens(b_tokens.input_tokens):>10}"
)
print(
f"{'cache reads':<18} {fmt_tokens(a_tokens.cache_read_tokens):>10} "
f"{fmt_tokens(b_tokens.cache_read_tokens):>10}"
)
print(
f"{'cache writes':<18} {fmt_tokens(a_tokens.cache_write_tokens):>10} "
f"{fmt_tokens(b_tokens.cache_write_tokens):>10}"
)
print(
f"{'tokens out':<18} {fmt_tokens(a_tokens.output_tokens):>10} "
f"{fmt_tokens(b_tokens.output_tokens):>10}"
)
print(f"{'turns':<18} {fmt_tokens(job_turn_totals(job_a)):>10} "
f"{fmt_tokens(job_turn_totals(job_b)):>10}")
print(f"{'cost':<18} {fmt_cost(a_cost):>10} {fmt_cost(b_cost):>10}")
print(
f"{'cost':<18} {fmt_cost(a_tokens.cost_usd):>10} "
f"{fmt_cost(b_tokens.cost_usd):>10}"
)
print(f"{'compute time':<18} {fmt_duration(job_duration(job_a)):>10} "
f"{fmt_duration(job_duration(job_b)):>10}")