refactor(agent/insights): shared JSON/tool-call iteration helpers, dict-driven model accumulation, unified cost-bucket lines (SQL byte-identical)

This commit is contained in:
Teknium
2026-09-02 18:27:21 -07:00
parent 3b17aa46cc
commit b73837cff4
+102 -198
View File
@@ -1,9 +1,7 @@
"""
Session Insights Engine: aggregates the SQLite state DB into usage insights
"""Session Insights Engine: aggregates the SQLite state DB into usage insights
(tokens, cost estimates, tool/skill usage, activity, model/platform breakdowns).
engine = InsightsEngine(db)
report = engine.generate(days=30)
engine = InsightsEngine(db); report = engine.generate(days=30)
print(engine.format_terminal(report))
"""
@@ -28,7 +26,7 @@ _SKILL_TOOLS = {"skill_view", "skill_manage"}
def _fmt_est_cost(est_cost: float) -> str:
"""Aggregate cost via the shared label helper so sub-cent totals render at 4dp, not "~$0.00"."""
"""Shared label helper so sub-cent totals render at 4dp, not "~$0.00"."""
return format_cost_label(Decimal(str(est_cost)))
@@ -50,18 +48,12 @@ def _estimate_cost(
provider, base_url = s.get("billing_provider"), s.get("billing_base_url")
else:
model = session_or_model or ""
usage = CanonicalUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
cache_read_tokens=cache_read_tokens,
cache_write_tokens=cache_write_tokens,
)
usage = CanonicalUsage(input_tokens, output_tokens, cache_read_tokens, cache_write_tokens)
result = estimate_usage_cost(model, usage, provider=provider, base_url=base_url)
return float(result.amount_usd or 0.0), result.status
def _bar_chart(values: List[int], max_width: int = 20) -> List[str]:
"""Create simple horizontal bar chart strings from values."""
peak = max(values) if values else 1
if peak == 0:
return ["" for _ in values]
@@ -73,14 +65,21 @@ def _short_model(model: Optional[str]) -> str:
return (model or "unknown").split("/")[-1]
def _parse_calls(raw: Any) -> Optional[list]:
"""tool_calls column → list, or None when not decodable as a JSON list."""
def _parse_json(raw: Any, kind: type) -> Any:
"""JSON-decode *raw* when it is a string; the value if it is a *kind*, else None."""
try:
if isinstance(raw, str):
raw = json.loads(raw)
except (json.JSONDecodeError, TypeError):
return None
return raw if isinstance(raw, list) else None
return raw if isinstance(raw, kind) else None
def _iter_functions(raw_calls: Any):
"""Yield the ``function`` dict of every well-formed entry in a tool_calls column."""
for call in _parse_json(raw_calls, list) or []:
if isinstance(call, dict):
yield call.get("function", {})
def _hour12(hr: int) -> str:
@@ -173,12 +172,10 @@ class InsightsEngine:
self.db = db
self._conn = db._conn
try:
self._has_assistant_calls_index = bool(
self._conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='index' AND name=?",
(self._MESSAGES_ASSISTANT_CALLS_INDEX,),
).fetchone()
)
self._has_assistant_calls_index = bool(self._conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='index' AND name=?",
(self._MESSAGES_ASSISTANT_CALLS_INDEX,),
).fetchone())
except sqlite3.Error:
self._has_assistant_calls_index = False
if not self._has_assistant_calls_index:
@@ -187,18 +184,17 @@ class InsightsEngine:
for suffix in ("_ALL", "_WITH_SOURCE"):
setattr(self, base + suffix, getattr(self, base + suffix).replace(strip, ""))
def _query(self, base: str, cutoff: float, source: Optional[str]):
"""Run ``<base>_WITH_SOURCE`` or ``<base>_ALL`` (instance attrs, so the
unpinned fallback applies) and return the cursor."""
def _query(self, base: str, cutoff: float, source: Optional[str]) -> list:
"""Rows of ``<base>_WITH_SOURCE`` or ``<base>_ALL`` (instance attrs, so the
unpinned fallback applies)."""
if source:
return self._conn.execute(getattr(self, base + "_WITH_SOURCE"), (cutoff, source))
return self._conn.execute(getattr(self, base + "_ALL"), (cutoff,))
return self._conn.execute(getattr(self, base + "_WITH_SOURCE"), (cutoff, source)).fetchall()
return self._conn.execute(getattr(self, base + "_ALL"), (cutoff,)).fetchall()
def generate(self, days: int = 30, source: str = None) -> Dict[str, Any]:
"""Generate a complete insights report for the last ``days`` days,
optionally filtered by source platform."""
cutoff = time.time() - (days * 86400)
# Drain the SessionDB's async accounting queue so counters are exact
# (self.db may be a raw sqlite3 connection in tests — guard).
flush = getattr(self.db, "flush_token_counts", None)
@@ -209,27 +205,16 @@ class InsightsEngine:
tool_usage = self._get_tool_usage(cutoff, source)
skill_usage = self._get_skill_usage(cutoff, source)
message_stats = self._get_message_stats(cutoff, source)
if not sessions:
return {
"days": days,
"source_filter": source,
"empty": True,
"overview": {},
"models": [],
"platforms": [],
"tools": [],
"skills": self._compute_skill_breakdown([]),
"activity": {},
"top_sessions": [],
"days": days, "source_filter": source, "empty": True, "overview": {}, "models": [],
"platforms": [], "tools": [], "skills": self._compute_skill_breakdown([]),
"activity": {}, "top_sessions": [],
}
models = self._compute_model_breakdown(sessions, cutoff, source)
return {
"days": days,
"source_filter": source,
"empty": False,
"generated_at": time.time(),
"days": days, "source_filter": source, "empty": False, "generated_at": time.time(),
"overview": self._compute_overview(sessions, message_stats, models),
"models": models,
"platforms": self._compute_platform_breakdown(sessions),
@@ -240,11 +225,8 @@ class InsightsEngine:
}
def get_usage_breakdown(self, days: int = 30, source: str = None) -> Dict[str, Any]:
"""Analytics-usage payload (tools + skills) without a full generate().
Uses the instr()-prefiltered skill query so only skill_view/skill_manage
messages are loaded, while keeping the per-tool breakdown the dashboard uses.
"""
"""Analytics-usage payload (tools + skills) without a full generate(); the
instr()-prefiltered skill query loads only skill_view/skill_manage messages."""
cutoff = time.time() - (days * 86400)
return {
"tools": self._compute_tool_breakdown(self._get_tool_usage(cutoff, source)),
@@ -254,97 +236,69 @@ class InsightsEngine:
# ------------------------------------------------------------------ SQL
def _get_sessions(self, cutoff: float, source: str = None) -> List[Dict]:
return [dict(row) for row in self._query("_GET_SESSIONS", cutoff, source).fetchall()]
return [dict(row) for row in self._query("_GET_SESSIONS", cutoff, source)]
def _get_tool_usage(self, cutoff: float, source: str = None) -> List[Dict]:
"""Tool call counts from two sources: ``tool_name`` on 'tool' rows (set
by the gateway) and ``tool_calls`` JSON on assistant rows (covers CLI,
where tool_name is not populated). Overlapping tools take the max."""
tool_counts = Counter()
for row in self._query("_GET_TOOL_NAMES", cutoff, source).fetchall():
for row in self._query("_GET_TOOL_NAMES", cutoff, source):
tool_counts[row["tool_name"]] += row["count"]
tool_calls_counts = Counter()
for row in self._query("_GET_TOOL_CALLS", cutoff, source).fetchall():
for row in self._query("_GET_TOOL_CALLS", cutoff, source):
try:
for call in _parse_calls(row["tool_calls"]) or []:
name = (call.get("function", {}) if isinstance(call, dict) else {}).get("name")
if name:
tool_calls_counts[name] += 1
tool_calls_counts.update(filter(None, (fn.get("name") for fn in _iter_functions(row["tool_calls"]))))
except (TypeError, AttributeError):
continue
if tool_calls_counts:
if tool_counts:
tool_counts = Counter({
tool: max(tool_counts.get(tool, 0), tool_calls_counts.get(tool, 0))
for tool in set(tool_counts) | set(tool_calls_counts)
})
else:
tool_counts = tool_calls_counts
if tool_calls_counts and tool_counts:
tool_counts = Counter({
tool: max(tool_counts.get(tool, 0), tool_calls_counts.get(tool, 0))
for tool in set(tool_counts) | set(tool_calls_counts)
})
elif tool_calls_counts:
tool_counts = tool_calls_counts
return [{"tool_name": name, "count": count} for name, count in tool_counts.most_common()]
def _get_skill_usage(self, cutoff: float, source: str = None) -> List[Dict]:
"""Extract per-skill usage from assistant tool calls."""
skill_counts: Dict[str, Dict[str, Any]] = {}
for row in self._query("_GET_SKILL_CALLS", cutoff, source).fetchall():
calls = _parse_calls(row["tool_calls"])
if calls is None:
continue
for row in self._query("_GET_SKILL_CALLS", cutoff, source):
timestamp = row["timestamp"]
for call in calls:
if not isinstance(call, dict):
continue
func = call.get("function", {})
for func in _iter_functions(row["tool_calls"]):
tool_name = func.get("name")
if tool_name not in _SKILL_TOOLS:
continue
args = func.get("arguments")
if isinstance(args, str):
try:
args = json.loads(args)
except (json.JSONDecodeError, TypeError):
continue
if not isinstance(args, dict):
continue
skill_name = args.get("name")
skill_name = (_parse_json(func.get("arguments"), dict) or {}).get("name")
if not isinstance(skill_name, str) or not skill_name.strip():
continue
entry = skill_counts.setdefault(
skill_name,
{"skill": skill_name, "view_count": 0, "manage_count": 0, "last_used_at": None},
skill_name, {"skill": skill_name, "view_count": 0, "manage_count": 0, "last_used_at": None},
)
entry["view_count" if tool_name == "skill_view" else "manage_count"] += 1
if timestamp is not None and (
entry["last_used_at"] is None or timestamp > entry["last_used_at"]
):
if timestamp is not None and (entry["last_used_at"] is None or timestamp > entry["last_used_at"]):
entry["last_used_at"] = timestamp
return list(skill_counts.values())
def _get_message_stats(self, cutoff: float, source: str = None) -> Dict:
row = self._query("_GET_MESSAGE_STATS", cutoff, source).fetchone()
return dict(row) if row else {
"total_messages": 0, "user_messages": 0,
"assistant_messages": 0, "tool_messages": 0,
rows = self._query("_GET_MESSAGE_STATS", cutoff, source)
return dict(rows[0]) if rows else {
"total_messages": 0, "user_messages": 0, "assistant_messages": 0, "tool_messages": 0,
}
def _get_model_usage(self, cutoff: float, source: str = None) -> List[Dict]:
"""Per-model usage rows; [] when the table is missing (older DB) so the
caller falls back to the per-session aggregate."""
try:
return [dict(row) for row in self._query("_GET_MODEL_USAGE", cutoff, source).fetchall()]
return [dict(row) for row in self._query("_GET_MODEL_USAGE", cutoff, source)]
except sqlite3.OperationalError:
return []
# -------------------------------------------------------------- Compute
def _compute_overview(
self,
sessions: List[Dict],
message_stats: Dict,
models: Optional[List[Dict]] = None,
) -> Dict:
def _compute_overview(self, sessions: List[Dict], message_stats: Dict, models: Optional[List[Dict]] = None) -> Dict:
# Per-model breakdown includes auxiliary usage rows (vision/compression/
# titles) plus reconciled residuals, while session counters carry
# main-loop usage only — sum the breakdown when available so overview
@@ -413,30 +367,25 @@ class InsightsEngine:
Sessions without per-model rows (pre-table data) fall back to their
single recorded aggregate. Tool calls aren't tied to an API call, so
they stay attributed to the session's recorded model."""
count_keys = _TOKEN_KEYS + ("reasoning_tokens", "api_call_count")
model_data = defaultdict(lambda: {
"sessions": set(), "input_tokens": 0, "output_tokens": 0,
"cache_read_tokens": 0, "cache_write_tokens": 0,
"reasoning_tokens": 0, "total_tokens": 0, "api_calls": 0,
"tool_calls": 0, "cost": 0.0, "actual_cost": 0.0,
"sessions": set(), **dict.fromkeys(_TOKEN_KEYS, 0), "reasoning_tokens": 0, "total_tokens": 0,
"api_calls": 0, "tool_calls": 0, "cost": 0.0, "actual_cost": 0.0,
})
def _accumulate(model, provider, base_url, session_id, inp, out,
cache_read, cache_write, reasoning, *,
def _accumulate(model, provider, base_url, session_id, counts: Dict[str, int], *,
stored_cost=None, actual_cost=None, cost_status=None):
model = model or "unknown"
display_model = _short_model(model)
d: Dict[str, Any] = model_data[display_model]
d: Dict[str, Any] = model_data[_short_model(model)]
d["sessions"].add(session_id)
d["input_tokens"] += inp
d["output_tokens"] += out
d["cache_read_tokens"] += cache_read
d["cache_write_tokens"] += cache_write
d["reasoning_tokens"] += reasoning
d["total_tokens"] += inp + out + cache_read + cache_write
for key in _TOKEN_KEYS + ("reasoning_tokens",):
d[key] += counts[key]
d["total_tokens"] += sum(counts[k] for k in _TOKEN_KEYS)
d["api_calls"] += counts["api_call_count"]
if stored_cost is None:
estimate, status = _estimate_cost(
model, inp, out,
cache_read_tokens=cache_read, cache_write_tokens=cache_write,
model, counts["input_tokens"], counts["output_tokens"],
cache_read_tokens=counts["cache_read_tokens"], cache_write_tokens=counts["cache_write_tokens"],
provider=provider or None, base_url=base_url,
)
else:
@@ -448,61 +397,37 @@ class InsightsEngine:
d["has_pricing"] = True
else:
d.setdefault("has_pricing", False)
return display_model
count_keys = _TOKEN_KEYS + ("reasoning_tokens", "api_call_count")
usage_totals = defaultdict(lambda: dict.fromkeys(count_keys, 0) | {
"estimated_cost_usd": 0.0, "actual_cost_usd": 0.0,
})
for r in self._get_model_usage(cutoff, source):
totals: Dict[str, Any] = usage_totals[r["session_id"]]
counts = {key: r[key] or 0 for key in count_keys}
for key in count_keys:
totals[key] += r[key] or 0
totals[key] += counts[key]
totals["estimated_cost_usd"] += r["estimated_cost_usd"] or 0.0
totals["actual_cost_usd"] += r["actual_cost_usd"] or 0.0
d = _accumulate(
r["model"], r["billing_provider"], r.get("billing_base_url"),
r["session_id"], r["input_tokens"] or 0, r["output_tokens"] or 0,
r["cache_read_tokens"] or 0, r["cache_write_tokens"] or 0,
r["reasoning_tokens"] or 0,
stored_cost=(
r["estimated_cost_usd"]
if r.get("cost_status") or r.get("cost_source")
else None
),
actual_cost=r["actual_cost_usd"],
cost_status=r.get("cost_status"),
_accumulate(
r["model"], r["billing_provider"], r.get("billing_base_url"), r["session_id"], counts,
stored_cost=r["estimated_cost_usd"] if r.get("cost_status") or r.get("cost_source") else None,
actual_cost=r["actual_cost_usd"], cost_status=r.get("cost_status"),
)
model_data[d]["api_calls"] += r["api_call_count"] or 0
# Reconcile against the aggregate row: covers legacy sessions,
# interrupted migrations, and absolute cumulative updates without
# double-counting already-attributed route deltas.
for s in sessions:
totals = usage_totals[s["id"]]
inp, out, cache_read, cache_write, residual_calls = (
max(0, (s.get(k) or 0) - totals[k]) for k in _TOKEN_KEYS + ("api_call_count",)
)
residual_cost = max(
0.0, float(s.get("estimated_cost_usd") or 0.0) - totals["estimated_cost_usd"],
)
residual_actual = max(
0.0, float(s.get("actual_cost_usd") or 0.0) - totals["actual_cost_usd"],
)
if not (
inp or out or cache_read or cache_write or residual_cost
or residual_actual or residual_calls
):
continue
d = _accumulate(
s.get("model"), s.get("billing_provider"),
s.get("billing_base_url"), s["id"],
inp, out, cache_read, cache_write, 0,
stored_cost=residual_cost,
actual_cost=residual_actual,
cost_status=s.get("cost_status"),
)
model_data[d]["api_calls"] += residual_calls
residual = {k: max(0, (s.get(k) or 0) - totals[k]) for k in _TOKEN_KEYS + ("api_call_count",)}
residual["reasoning_tokens"] = 0
residual_cost = max(0.0, float(s.get("estimated_cost_usd") or 0.0) - totals["estimated_cost_usd"])
residual_actual = max(0.0, float(s.get("actual_cost_usd") or 0.0) - totals["actual_cost_usd"])
if any(residual.values()) or residual_cost or residual_actual:
_accumulate(
s.get("model"), s.get("billing_provider"), s.get("billing_base_url"), s["id"], residual,
stored_cost=residual_cost, actual_cost=residual_actual, cost_status=s.get("cost_status"),
)
for s in sessions:
tool_calls = s.get("tool_call_count") or 0
@@ -522,9 +447,7 @@ class InsightsEngine:
def _compute_platform_breakdown(self, sessions: List[Dict]) -> List[Dict]:
platform_data = defaultdict(lambda: {
"sessions": 0, "messages": 0, "input_tokens": 0,
"output_tokens": 0, "cache_read_tokens": 0,
"cache_write_tokens": 0, "total_tokens": 0, "tool_calls": 0,
"sessions": 0, "messages": 0, **dict.fromkeys(_TOKEN_KEYS, 0), "total_tokens": 0, "tool_calls": 0,
})
for s in sessions:
d = platform_data[s.get("source") or "unknown"]
@@ -541,37 +464,29 @@ class InsightsEngine:
def _compute_tool_breakdown(self, tool_usage: List[Dict]) -> List[Dict]:
"""Ranked tool list with percentages."""
total_calls = sum(t["count"] for t in tool_usage) if tool_usage else 0
total_calls = sum(t["count"] for t in tool_usage)
return [
{
"tool": t["tool_name"],
"count": t["count"],
"percentage": (t["count"] / total_calls * 100) if total_calls else 0,
}
{"tool": t["tool_name"], "count": t["count"], "percentage": (t["count"] / total_calls * 100) if total_calls else 0}
for t in tool_usage
]
def _compute_skill_breakdown(self, skill_usage: List[Dict]) -> Dict[str, Any]:
"""Per-skill usage → summary + ranked list."""
total_skill_loads = sum(s["view_count"] for s in skill_usage) if skill_usage else 0
total_skill_edits = sum(s["manage_count"] for s in skill_usage) if skill_usage else 0
total_skill_loads = sum(s["view_count"] for s in skill_usage)
total_skill_edits = sum(s["manage_count"] for s in skill_usage)
total_skill_actions = total_skill_loads + total_skill_edits
top_skills = []
for skill in skill_usage:
total_count = skill["view_count"] + skill["manage_count"]
top_skills.append({
"skill": skill["skill"],
"view_count": skill["view_count"],
"manage_count": skill["manage_count"],
"skill": skill["skill"], "view_count": skill["view_count"], "manage_count": skill["manage_count"],
"total_count": total_count,
"percentage": (total_count / total_skill_actions * 100) if total_skill_actions else 0,
"last_used_at": skill.get("last_used_at"),
})
top_skills.sort(
key=lambda s: (
s["total_count"], s["view_count"], s["manage_count"], s["last_used_at"] or 0, s["skill"],
),
key=lambda s: (s["total_count"], s["view_count"], s["manage_count"], s["last_used_at"] or 0, s["skill"]),
reverse=True,
)
return {
@@ -632,8 +547,7 @@ class InsightsEngine:
if timed:
longest = max(timed, key=lambda s: s["ended_at"] - s["started_at"])
top.append({
"label": "Longest session",
"session_id": longest["id"][:16],
"label": "Longest session", "session_id": longest["id"][:16],
"value": format_duration_compact(longest["ended_at"] - longest["started_at"]),
"date": _day(longest["started_at"]),
})
@@ -641,12 +555,7 @@ class InsightsEngine:
best = max(sessions, key=metric)
value = metric(best)
if value > 0:
top.append({
"label": label,
"session_id": best["id"][:16],
"value": fmt.format(value),
"date": _day(best.get("started_at")),
})
top.append({"label": label, "session_id": best["id"][:16], "value": fmt.format(value), "date": _day(best.get("started_at"))})
return top
# ------------------------------------------------------------- Formatting
@@ -655,6 +564,16 @@ class InsightsEngine:
def _section(title: str) -> List[str]:
return [f" {title}", " " + "─" * 56]
@staticmethod
def _cost_lines(o: Dict, templates: tuple) -> List[str]:
"""One formatted line per non-zero cost bucket (estimated, included, unknown)."""
est_cost = o.get("estimated_cost", 0.0)
values = (
_fmt_est_cost(est_cost) if est_cost > 0 else "",
o.get("included_cost_sessions", 0), o.get("unknown_cost_sessions", 0),
)
return [tpl.format(v) for tpl, v in zip(templates, values) if v]
def format_terminal(self, report: Dict) -> str:
"""Format the insights report for terminal display (CLI)."""
if report.get("empty"):
@@ -692,18 +611,12 @@ class InsightsEngine:
lines += [f" Avg msgs/session: {o['avg_messages_per_session']:.1f}", ""]
# Cost buckets: show included/unknown sessions instead of collapsing to $0.
est_cost = o.get("estimated_cost", 0.0)
included_sessions = o.get("included_cost_sessions", 0)
unknown_sessions = o.get("unknown_cost_sessions", 0)
if est_cost > 0 or included_sessions > 0 or unknown_sessions > 0:
lines += self._section("💰 Cost")
if est_cost > 0:
lines.append(f" Estimated: {_fmt_est_cost(est_cost)}")
if included_sessions > 0:
lines.append(f" Included: {included_sessions} session(s) (subscription — no provider invoice)")
if unknown_sessions > 0:
lines.append(f" Unknown: {unknown_sessions} session(s) (no pricing data)")
lines.append("")
cost_lines = self._cost_lines(o, (
" Estimated: {}", " Included: {} session(s) (subscription — no provider invoice)",
" Unknown: {} session(s) (no pricing data)",
))
if cost_lines:
lines += self._section("💰 Cost") + cost_lines + [""]
if report["models"]:
lines += self._section("🤖 Models Used")
@@ -789,16 +702,7 @@ class InsightsEngine:
lines.append(f"**Active time:** ~{format_duration_compact(o['total_hours'] * 3600)} | **Avg session:** ~{format_duration_compact(o['avg_session_duration'])}")
lines.append("")
est_cost = o.get("estimated_cost", 0.0)
included = o.get("included_cost_sessions", 0)
unknown = o.get("unknown_cost_sessions", 0)
cost_parts: list[str] = []
if est_cost > 0:
cost_parts.append(f"{_fmt_est_cost(est_cost)} estimated")
if included > 0:
cost_parts.append(f"{included} included (subscription)")
if unknown > 0:
cost_parts.append(f"{unknown} unknown")
cost_parts = self._cost_lines(o, ("{} estimated", "{} included (subscription)", "{} unknown"))
if cost_parts:
lines += [f"**Cost:** {' | '.join(cost_parts)}", ""]