diff --git a/agent/core/session.py b/agent/core/session.py index 3a7f50c3..85e7dd56 100644 --- a/agent/core/session.py +++ b/agent/core/session.py @@ -560,6 +560,7 @@ def get_trajectory(self) -> dict: "session_id": self.session_id, "user_id": self.user_id, "hf_username": self.hf_username, + "user_plan": self.user_plan, "session_start_time": self.session_start_time, "session_end_time": datetime.now().isoformat(), "model_name": self.config.model_name, diff --git a/agent/core/session_uploader.py b/agent/core/session_uploader.py index 268c8459..3c2c09a7 100644 --- a/agent/core/session_uploader.py +++ b/agent/core/session_uploader.py @@ -290,6 +290,7 @@ def _write_row_payload(data: dict, tmp_path: str) -> None: session_row = { "session_id": data["session_id"], "user_id": data.get("user_id"), + "user_plan": data.get("user_plan") or "unknown", "session_start_time": data["session_start_time"], "session_end_time": data["session_end_time"], "model_name": data["model_name"], diff --git a/agent/core/telemetry.py b/agent/core/telemetry.py index ef6623dd..1c170f2f 100644 --- a/agent/core/telemetry.py +++ b/agent/core/telemetry.py @@ -17,7 +17,10 @@ from __future__ import annotations import asyncio +import hashlib +import hmac import logging +import os import time from typing import Any @@ -150,6 +153,92 @@ def _infer_push_to_hub(script_or_cmd: Any) -> bool: ) +def _kpi_hash_identifier(value: Any) -> str | None: + salt = os.environ.get("KPI_USER_HASH_SALT") + if not salt or value is None: + return None + raw = str(value) + if not raw: + return None + return hmac.new( + salt.encode("utf-8"), raw.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + +def _sanitize_expected_hub_artifacts(artifacts: Any) -> list[dict[str, Any]]: + sanitized: list[dict[str, Any]] = [] + if not isinstance(artifacts, list): + return sanitized + for artifact in artifacts: + if not isinstance(artifact, dict): + continue + repo_type = str(artifact.get("repo_type") or "model").strip().lower() + repo_id = artifact.get("repo_id") + if repo_type not in {"model", "dataset", "space"} or not repo_id: + continue + artifact_hash = _kpi_hash_identifier(f"{repo_type}:{repo_id}") + if artifact_hash is None: + continue + sanitized.append( + { + "repo_type": repo_type, + "artifact_hash": artifact_hash, + "private": artifact.get("private"), + "is_sandbox": bool(artifact.get("is_sandbox")), + } + ) + return sanitized + + +async def record_hub_artifact( + session: Any, + *, + repo_type: str, + repo_id: str, + source: str, + is_sandbox: bool | None = None, + private: bool | None = None, + success: bool = True, +) -> dict[str, Any]: + """Emit a sanitized Hub artifact event. + + The raw repo id never leaves this function. Consumers dedupe artifacts by + ``artifact_hash``. + """ + from agent.core.session import Event + + try: + normalized_type = str(repo_type or "").strip().lower() + if normalized_type not in {"model", "dataset", "space"}: + return {} + if is_sandbox is None: + try: + from agent.core.hub_artifacts import is_sandbox_hub_repo + + is_sandbox = is_sandbox_hub_repo(repo_id, normalized_type) + except Exception: + is_sandbox = False + artifact_hash = _kpi_hash_identifier(f"{normalized_type}:{repo_id}") + if artifact_hash is None: + logger.debug( + "record_hub_artifact skipped because KPI_USER_HASH_SALT is unset" + ) + return {} + payload = { + "repo_type": normalized_type, + "artifact_hash": artifact_hash, + "source": str(source or "unknown"), + "is_sandbox": bool(is_sandbox), + "private": private, + "success": bool(success), + } + await session.send_event(Event(event_type="hub_artifact", data=payload)) + return payload + except Exception as e: + logger.debug("record_hub_artifact failed (non-fatal): %s", e) + return {} + + async def record_hf_job_submit( session: Any, job: Any, @@ -165,6 +254,9 @@ async def record_hf_job_submit( t_start = time.monotonic() try: script_text = args.get("script") or args.get("command") or "" + expected_artifacts = _sanitize_expected_hub_artifacts( + args.get("expected_hub_artifacts") + ) await session.send_event( Event( event_type="hf_job_submit", @@ -177,6 +269,8 @@ async def record_hf_job_submit( "image": image, "namespace": args.get("namespace"), "push_to_hub": _infer_push_to_hub(script_text), + "expected_hub_artifacts": expected_artifacts, + "expected_hub_artifacts_count": len(expected_artifacts), }, ) ) @@ -231,6 +325,34 @@ async def record_hf_job_complete( return {} +async def record_hf_job_cancel( + session: Any, + *, + job_id: str, + namespace: str | None = None, + flavor: str | None = None, +) -> dict: + from agent.core.session import Event + + try: + payload = { + "job_id": job_id, + "namespace": namespace, + "flavor": flavor, + "final_status": "cancelled", + "wall_time_s": 0, + "billable_seconds_estimate": 0, + "price_usd_per_hour": None, + "estimated_cost_usd": None, + "cost_estimate_source": "manual_cancel", + } + await session.send_event(Event(event_type="hf_job_complete", data=payload)) + return payload + except Exception as e: + logger.debug("record_hf_job_cancel failed (non-fatal): %s", e) + return {} + + # ── sandbox ───────────────────────────────────────────────────────────────── diff --git a/agent/tools/hf_repo_git_tool.py b/agent/tools/hf_repo_git_tool.py index 672186c6..0468ecce 100644 --- a/agent/tools/hf_repo_git_tool.py +++ b/agent/tools/hf_repo_git_tool.py @@ -548,6 +548,17 @@ async def _create_repo(self, args: Dict[str, Any]) -> ToolResult: session=self.session, extra_metadata=extra_metadata, ) + if self.session: + from agent.core import telemetry + + await telemetry.record_hub_artifact( + self.session, + repo_type=repo_type, + repo_id=repo_id, + source="hf_repo_git", + private=private, + success=True, + ) return { "formatted": f"**Repository created:** {repo_id}\n**Private:** {private}\n{result}", diff --git a/agent/tools/jobs_tool.py b/agent/tools/jobs_tool.py index f9afe782..4eae8dbe 100644 --- a/agent/tools/jobs_tool.py +++ b/agent/tools/jobs_tool.py @@ -148,6 +148,43 @@ def _add_environment_variables( return result +def _normalize_expected_hub_artifacts(raw: Any) -> list[dict[str, Any]]: + artifacts: list[dict[str, Any]] = [] + if not isinstance(raw, list): + return artifacts + for item in raw: + if not isinstance(item, dict): + continue + repo_type = str(item.get("repo_type") or "model").strip().lower() + repo_id = ( + item.get("repo_id") + or item.get("hub_model_id") + or item.get("hub_dataset_id") + or item.get("space_id") + ) + if repo_type not in {"model", "dataset", "space"} or not repo_id: + continue + artifacts.append( + { + "repo_type": repo_type, + "repo_id": str(repo_id), + "private": item.get("private"), + "is_sandbox": bool(item.get("is_sandbox")), + } + ) + return artifacts + + +def _job_status_completed(status: Any) -> bool: + return str(status or "").strip().upper() in { + "COMPLETED", + "COMPLETE", + "SUCCEEDED", + "SUCCESS", + "DONE", + } + + def _build_uv_command( script: str, with_deps: list[str] | None = None, @@ -389,7 +426,7 @@ async def execute(self, params: Dict[str, Any]) -> ToolResult: "isError": True, } - async def _seed_trackio_dashboard(self, space_id: str) -> None: + async def _seed_trackio_dashboard(self, space_id: str) -> bool: """Idempotently install trackio dashboard files into *space_id* before the job runs. Surfaces seed progress as tool_log events but never raises — a seed failure should not block job submission, since trackio @@ -410,9 +447,41 @@ def _log(msg: str) -> None: await asyncio.to_thread( ensure_trackio_dashboard, space_id, self.hf_token, _log ) + if self.session: + from agent.core import telemetry + + await telemetry.record_hub_artifact( + self.session, + repo_type="space", + repo_id=space_id, + source="trackio", + is_sandbox=False, + private=True, + success=True, + ) + return True except Exception as e: logger.warning(f"trackio dashboard seed failed for {space_id}: {e}") _log(f"trackio dashboard seed failed: {e}") + return False + + async def _record_successful_job_artifacts( + self, artifacts: list[dict[str, Any]] + ) -> None: + if not self.session: + return + from agent.core import telemetry + + for artifact in artifacts: + await telemetry.record_hub_artifact( + self.session, + repo_type=artifact["repo_type"], + repo_id=artifact["repo_id"], + source="hf_job", + is_sandbox=artifact["is_sandbox"], + private=artifact.get("private"), + success=True, + ) async def _wait_for_job_completion( self, job_id: str, namespace: Optional[str] = None @@ -569,6 +638,9 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: # Run the job flavor = args.get("hardware_flavor", "cpu-basic") timeout_str = args.get("timeout", "30m") + expected_hub_artifacts = _normalize_expected_hub_artifacts( + args.get("expected_hub_artifacts") + ) # Trackio: agent-declared space + project become env vars on the job # so trackio.init() picks them up automatically. We also surface them @@ -658,6 +730,7 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: "hardware_flavor": flavor, "timeout": timeout_str, "namespace": self.namespace, + "expected_hub_artifacts": expected_hub_artifacts, }, image=image, job_type=job_type, @@ -711,6 +784,8 @@ async def _run_job(self, args: Dict[str, Any]) -> ToolResult: else None, allow_zero_actual=True, ) + if _job_status_completed(final_status): + await self._record_successful_job_artifacts(expected_hub_artifacts) # Untrack job ID (completed or failed, no longer needs cancellation) if self.session: @@ -880,6 +955,14 @@ async def _cancel_job(self, args: Dict[str, Any]) -> ToolResult: job_id=job_id, namespace=self.namespace, ) + if self.session: + from agent.core import telemetry + + await telemetry.record_hf_job_cancel( + self.session, + job_id=job_id, + namespace=self.namespace, + ) response = f"""✓ Job {job_id} has been cancelled. @@ -1249,6 +1332,26 @@ async def _resume_scheduled_job(self, args: Dict[str, Any]) -> ToolResult: "the embedded dashboard to this project." ), }, + "expected_hub_artifacts": { + "type": "array", + "items": { + "type": "object", + "properties": { + "repo_type": { + "type": "string", + "enum": ["model", "dataset", "space"], + }, + "repo_id": {"type": "string"}, + "private": {"type": "boolean"}, + "is_sandbox": {"type": "boolean"}, + }, + "required": ["repo_type", "repo_id"], + }, + "description": ( + "Optional. Hub artifacts the job is expected to create. " + "They are only counted after the job completes successfully." + ), + }, "namespace": { "type": "string", "description": ( diff --git a/agent/tools/sandbox_tool.py b/agent/tools/sandbox_tool.py index a550c018..14e0e39d 100644 --- a/agent/tools/sandbox_tool.py +++ b/agent/tools/sandbox_tool.py @@ -289,6 +289,17 @@ def _log(msg: str) -> None: await asyncio.to_thread( ensure_trackio_dashboard, space_id, session.hf_token, _log ) + from agent.core import telemetry + + await telemetry.record_hub_artifact( + session, + repo_type="space", + repo_id=space_id, + source="trackio", + is_sandbox=False, + private=True, + success=True, + ) except Exception as e: _log(f"trackio dashboard seed failed: {e}") diff --git a/backend/kpis_scheduler.py b/backend/kpis_scheduler.py index 9b2199c6..fcfb0b2b 100644 --- a/backend/kpis_scheduler.py +++ b/backend/kpis_scheduler.py @@ -16,6 +16,7 @@ HF_KPI_WRITE_TOKEN | HF_SESSION_UPLOAD_TOKEN | HF_TOKEN | HF_ADMIN_TOKEN First one found is used. Least-privilege first. + KPI_USER_HASH_SALT required by the v2 builder to hash user/session/artifact ids KPI_SOURCE_REPO default smolagents/ml-intern-sessions KPI_TARGET_REPO default smolagents/ml-intern-kpis ML_INTERN_KPIS_DISABLED if truthy, the scheduler is not started @@ -80,7 +81,14 @@ async def _run_hour(hour_dt: datetime) -> None: api = HfApi() source = os.environ.get("KPI_SOURCE_REPO", "smolagents/ml-intern-sessions") target = os.environ.get("KPI_TARGET_REPO", "smolagents/ml-intern-kpis") - await asyncio.to_thread(mod.run_for_hour, api, source, target, hour_dt, token) + await asyncio.to_thread( + mod.run_for_hour, + api, + source_repo=source, + target_repo=target, + hour_dt=hour_dt, + token=token, + ) except Exception as e: logger.warning("kpis_scheduler: rollup for %s failed: %s", hour_dt, e) @@ -104,6 +112,9 @@ def start(backfill_hours: int = 6) -> None: if os.environ.get("ML_INTERN_KPIS_DISABLED"): logger.info("kpis_scheduler: disabled via ML_INTERN_KPIS_DISABLED") return + if not os.environ.get("KPI_USER_HASH_SALT"): + logger.error("kpis_scheduler: KPI_USER_HASH_SALT is required, skipping") + return if _scheduler is not None: return diff --git a/backend/routes/agent.py b/backend/routes/agent.py index 06cf1c38..4aa0ae83 100644 --- a/backend/routes/agent.py +++ b/backend/routes/agent.py @@ -659,6 +659,17 @@ async def upload_session_dataset( hf_username=hf_username, hf_token=hf_token, ) + from agent.core import telemetry + + await telemetry.record_hub_artifact( + agent_session.session, + repo_type=uploaded.repo_type, + repo_id=uploaded.repo_id, + source="dataset_upload", + is_sandbox=False, + private=uploaded.private, + success=True, + ) agent_session.session.context_manager.add_message( Message(role="user", content=dataset_context_note(uploaded)) ) diff --git a/scripts/build_kpis.py b/scripts/build_kpis.py index a4b2f499..fc002fc6 100644 --- a/scripts/build_kpis.py +++ b/scripts/build_kpis.py @@ -1,767 +1,1108 @@ #!/usr/bin/env python3 -"""Hourly KPI rollup for the session-trajectory dataset. - -================================================================================ - Data flow -================================================================================ - - ┌────────────────────┐ heartbeat ┌────────────────────────────────┐ - │ agent (CLI/web) │ ───────────────▶ │ hf-agent-sessions (dataset) │ - │ Session.send_event│ │ sessions/YYYY-MM-DD/.jsonl│ - └────────────────────┘ └───────────────┬────────────────┘ - │ cron @:05 each hour - ▼ - ┌──────────────────────────────────┐ - │ scripts/build_kpis.py │ - │ (GitHub Actions) │ - └───────────────┬──────────────────┘ - │ upload CSV - ▼ - ┌──────────────────────────────────┐ - │ hf-agent-kpis (dataset) │ - │ hourly/YYYY-MM-DD/HH.csv │ - └──────────────────────────────────┘ - -Each hourly run reads today's + yesterday's session folders (to cover sessions -that crossed midnight), filters events into the target hour window -``[hour, hour+1h)``, computes aggregates, and writes one CSV at -``hourly//.csv`` in the target dataset. Uploads are idempotent — -re-running the same hour overwrites. - -================================================================================ - Metrics (one row per hour) -================================================================================ - - sessions — distinct session_ids with ≥1 event in window - users — distinct user ids (when present on session rows) - turns — sum of user-message counts across active sessions - llm_calls — count of llm_call events - tokens_prompt / _completion / _cache_read / _cache_creation - cost_usd — sum of llm_call.cost_usd - cost_per_session_mean / _p50 / _p95 — per-session cost distribution - cache_hit_ratio — cache_read / (cache_read + prompt) - tool_calls_total / _succeeded / _failed — per-tool_output reliability counts - tool_success_rate — succeeded / total (kept for back-compat) - successful_sessions / errored_sessions / regenerated_sessions — outcome counts - failure_rate / regenerate_rate — kept for back-compat - time_to_first_action_s_p50 / _p95 — from session_start to first tool_call - thumbs_up / thumbs_down - hf_jobs_submitted / _succeeded / _blocked - sandboxes_created / _cpu / _gpu — sandbox_create events bucketed by hardware - pro_cta_clicks - gpu_hours_by_flavor_json — JSON-serialised {flavor: gpu-hours} - research_calls — total `research` tool_call events - sessions_with_research — sessions that called `research` ≥1 - research_calls_per_session_p50 / _p95 — among sessions that did any (zero-only sessions excluded) - distinct_tools_per_session_p50 / _p95 — among sessions with ≥1 named tool_call - tool_calls_per_session_p50 / _p95 — among sessions with ≥1 named tool_call - tool_calls_per_turn_p50 / _p95 — calls / turns, among sessions with turns>0 - tool_calls_by_name_json — JSON {tool: total_calls} (all tools seen) - sessions_using_tool_json — JSON {tool: distinct_sessions_using} - sessions_by_model_json — JSON {model_name: count} (router/local split) - -================================================================================ - Usage -================================================================================ - - # Run for the most recently completed hour (default — the cron path): - python scripts/build_kpis.py - - # Backfill last 24 hours: - python scripts/build_kpis.py --hours 24 - - # Explicit hour (UTC): - python scripts/build_kpis.py --datetime 2026-04-24T14 - -Env: - HF_TOKEN (or HF_KPI_WRITE_TOKEN) — write access to the target dataset. - -================================================================================ - Deploy -================================================================================ - -See ``.github/workflows/build-kpis.yml`` — runs every hour at :05. To provision: - - 1. Create the target dataset (once): - huggingface-cli repo create hf-agent-kpis --type dataset - 2. Put ``HF_KPI_WRITE_TOKEN`` (or ``HF_TOKEN``) into repo Actions secrets. - 3. Merge this file; the first scheduled run fires within the hour. -""" +"""Build KPI dataset v2 facts and rollups. -from __future__ import annotations +The v2 pipeline deliberately stops producing the legacy hourly/daily schema. +It emits exact per-session facts plus hourly, daily, and monthly rollups under +``v2/``. Raw user ids, session ids, and Hub artifact repo ids are never written; +all identifiers are salted HMAC-SHA256 hashes. +""" import argparse +import calendar +import csv +import hashlib +import hmac import io import json import logging import os import sys import tempfile -from collections import defaultdict +from collections import Counter, defaultdict from datetime import date, datetime, timedelta, timezone from typing import Any, Iterable logger = logging.getLogger("build_kpis") -# Rough gpu-hour pricing for hf_jobs flavor strings. Keep conservative; used -# only to compute gpu-hours (not dollars) — wall_time_s * flavor_gpu_count. -_FLAVOR_GPU_COUNT = { - "cpu-basic": 0, - "cpu-upgrade": 0, - "t4-small": 1, - "t4-medium": 1, - "l4x1": 1, - "l4x4": 4, - "l40sx1": 1, - "l40sx4": 4, - "l40sx8": 8, - "a10g-small": 1, - "a10g-large": 1, - "a10g-largex2": 2, - "a10g-largex4": 4, - "a100-large": 1, - "a100x2": 2, - "a100x4": 4, - "a100x8": 8, - "h100": 1, - "h100x8": 8, -} - - -def _percentile(values: list[float], p: float) -> float: - if not values: +SCHEMA_VERSION = 2 +V2_PREFIX = "v2" +VALID_PLANS = {"free", "pro", "unknown"} +VALID_REPO_TYPES = {"model", "dataset", "space"} + + +def _json_dumps(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _require_hash_salt(salt: str | None = None) -> str: + resolved = salt if salt is not None else os.environ.get("KPI_USER_HASH_SALT") + if not resolved: + raise RuntimeError("KPI_USER_HASH_SALT is required for KPI v2 exports") + return resolved + + +def _hash_identifier(value: Any, salt: str) -> str: + raw = "" if value is None else str(value) + return hmac.new( + salt.encode("utf-8"), raw.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, json.JSONDecodeError): + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _as_list(value: Any) -> list[Any]: + if isinstance(value, list): + return value + if isinstance(value, str): + try: + parsed = json.loads(value) + except (TypeError, json.JSONDecodeError): + return [] + return parsed if isinstance(parsed, list) else [] + return [] + + +def _number(value: Any) -> float: + if isinstance(value, bool) or value is None: + return 0.0 + try: + return float(value) + except (TypeError, ValueError): return 0.0 - values = sorted(values) - k = (len(values) - 1) * p - f = int(k) - c = min(f + 1, len(values) - 1) - if f == c: - return float(values[f]) - return float(values[f] + (values[c] - values[f]) * (k - f)) -def _parse_ts(s: Any) -> datetime | None: - if not s or not isinstance(s, str): - return None +def _integer(value: Any) -> int: + if isinstance(value, bool) or value is None: + return 0 try: - dt = datetime.fromisoformat(s) - except Exception: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _parse_ts(value: Any) -> datetime | None: + if isinstance(value, datetime): + dt = value + elif isinstance(value, str) and value: + try: + dt = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + else: return None - # Normalise to aware UTC so comparisons work against window bounds. if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) - return dt + return dt.astimezone(timezone.utc) -def _iter_session_files(api, repo_id: str, day: date, token: str) -> Iterable[str]: - """Yield repo-relative paths for all sessions under ``sessions/YYYY-MM-DD/``.""" - prefix = f"sessions/{day.isoformat()}/" - try: - files = api.list_repo_files(repo_id=repo_id, repo_type="dataset", token=token) - except Exception as e: - logger.warning("list_repo_files(%s) failed: %s", repo_id, e) - return [] - return [f for f in files if f.startswith(prefix) and f.endswith(".jsonl")] +def _date_key(dt: datetime) -> str: + return dt.astimezone(timezone.utc).strftime("%Y-%m-%d") -def _download_session(repo_id: str, path: str, token: str) -> dict | None: - """Fetch one session JSONL and decode its single row. +def _hour_key(dt: datetime) -> str: + return dt.astimezone(timezone.utc).strftime("%Y-%m-%dT%H") - ``hf_hub_download`` caches; second run within the same process / runner - directory is near-free. - """ - from huggingface_hub import hf_hub_download - try: - local = hf_hub_download( - repo_id=repo_id, - filename=path, - repo_type="dataset", - token=token, - ) - except Exception as e: - logger.warning("hf_hub_download(%s) failed: %s", path, e) - return None - try: - with open(local, "r") as f: - line = f.readline().strip() - if not line: - return None - row = json.loads(line) - # Session uploader stores messages/events as JSON strings — unpack. - for key in ("messages", "events", "tools"): - v = row.get(key) - if isinstance(v, str): - try: - row[key] = json.loads(v) - except Exception: - row[key] = [] - return row - except Exception as e: - logger.warning("parse(%s) failed: %s", path, e) - return None +def _month_key_from_date_key(day: str) -> str: + return day[:7] -def _filter_session_to_window( - session: dict, - start: datetime, - end: datetime, -) -> dict | None: - """Return a copy of ``session`` whose events are only those in ``[start, end)``. - - ``None`` if no event falls in the window — the caller drops the session - from this hour's aggregate. - """ - events = session.get("events") or [] - in_window = [] - for ev in events: - ts = _parse_ts(ev.get("timestamp")) - if ts is None: +def _normalize_plan(value: Any) -> str: + plan = str(value or "unknown").strip().lower() + return plan if plan in VALID_PLANS else "unknown" + + +def _normalize_job_status(value: Any) -> str: + status = str(value or "unknown").strip().lower() + if status in {"completed", "complete", "succeeded", "success", "done"}: + return "completed" + if status in {"cancelled", "canceled", "cancelled_by_user", "canceled_by_user"}: + return "cancelled" + if "cancel" in status or status in {"stopped", "killed"}: + return "cancelled" + if status in {"failed", "failure", "error", "errored", "timeout", "timed_out"}: + return "failed" + if "fail" in status or "error" in status or "timeout" in status: + return "failed" + return "unknown" + + +def _session_events(session: dict[str, Any]) -> list[dict[str, Any]]: + return [ + event for event in _as_list(session.get("events")) if isinstance(event, dict) + ] + + +def _session_messages(session: dict[str, Any]) -> list[dict[str, Any]]: + return [msg for msg in _as_list(session.get("messages")) if isinstance(msg, dict)] + + +def _event_data(event: dict[str, Any]) -> dict[str, Any]: + return _as_dict(event.get("data")) + + +def _session_usage_metrics(session: dict[str, Any]) -> dict[str, Any]: + metrics = _as_dict(session.get("usage_metrics")) + if metrics: + return metrics + # Historical rows can lack usage_metrics. Keep the fallback minimal and do + # not use llm_call.cost_usd for the billing fields. + prompt = completion = cache_read = cache_creation = total = calls = 0 + calls_by_model: Counter[str] = Counter() + for event in _session_events(session): + if event.get("event_type") != "llm_call": continue - if start <= ts < end: - in_window.append(ev) - if not in_window: - return None - return {**session, "events": in_window} - - -def _session_metrics(session: dict) -> dict: - """Reduce a single session trajectory to its KPI contributions. - - Assumes ``events`` are already filtered to the target window by the caller. - """ - # Pre-seed every numeric key so downstream aggregation can sum without - # having to special-case empty sessions. - out: dict = { - "sessions": 0, - "turns": 0, - "llm_calls": 0, - "tokens_prompt": 0, - "tokens_completion": 0, - "tokens_cache_read": 0, - "tokens_cache_creation": 0, - "cost_usd": 0.0, - "tool_calls_total": 0, - "tool_calls_success": 0, - "failures": 0, - "regenerate_sessions": 0, - "thumbs_up": 0, - "thumbs_down": 0, - "hf_jobs_submitted": 0, - "hf_jobs_succeeded": 0, - "hf_jobs_blocked": 0, - "pro_cta_clicks": 0, - "pro_conversions": 0, - "credits_topped_up": 0, - "sandboxes_created": 0, - "sandboxes_cpu": 0, - "sandboxes_gpu": 0, - "first_tool_s": -1, + data = _event_data(event) + calls += 1 + prompt += _integer(data.get("prompt_tokens")) + completion += _integer(data.get("completion_tokens")) + cache_read += _integer(data.get("cache_read_tokens")) + cache_creation += _integer(data.get("cache_creation_tokens")) + total += _integer(data.get("total_tokens")) or ( + _integer(data.get("prompt_tokens")) + + _integer(data.get("completion_tokens")) + + _integer(data.get("cache_read_tokens")) + + _integer(data.get("cache_creation_tokens")) + ) + calls_by_model[ + str(data.get("model") or session.get("model_name") or "unknown") + ] += 1 + return { + "total_usd": 0.0, + "total_usd_source": "missing_usage_metrics", + "app_total_usd": 0.0, + "hf_billing_total_usd": None, + "app_telemetry": { + "inference_usd": 0.0, + "hf_jobs_estimated_usd": 0.0, + "sandbox_estimated_usd": 0.0, + }, + "hf_billing": {"available": False, "current_session": None}, + "llm": { + "calls": calls, + "calls_by_model": dict(calls_by_model), + "prompt_tokens": prompt, + "completion_tokens": completion, + "cache_read_tokens": cache_read, + "cache_creation_tokens": cache_creation, + "total_tokens": total, + }, + "hf_jobs": {"submits": 0, "estimated_usd": 0.0}, + "sandboxes": {"creates": 0, "estimated_usd": 0.0}, } - events = session.get("events") or [] - messages = session.get("messages") or [] - - turn_count = sum(1 for m in messages if m.get("role") == "user") - out["turns"] = turn_count - out["sessions"] = 1 - - tool_success = 0 - tool_total = 0 - had_error = False - had_undo = False - first_tool_ts = None - session_start = session.get("session_start_time") - gpu_hours_by_flavor: dict[str, float] = defaultdict(float) - jobs_submitted = 0 - jobs_succeeded = 0 - thumbs_up = 0 - thumbs_down = 0 - sandboxes_created = 0 - sandboxes_cpu = 0 - sandboxes_gpu = 0 - jobs_blocked = 0 - pro_cta_clicks = 0 - pro_conversions = 0 - credits_topped_up = 0 - pro_cta_by_source: dict[str, int] = defaultdict(int) - # Per-tool counters from tool_call events. Counted off tool_call (which - # carries data["tool"]) rather than tool_output (which only carries - # success/output) so we can attribute calls to specific tools. - tool_calls_by_name: dict[str, int] = defaultdict(int) - total_named_tool_calls = 0 - - start_dt = _parse_ts(session_start) - - for ev in events: - et = ev.get("event_type") - data = ev.get("data") or {} - ts = _parse_ts(ev.get("timestamp")) - - if et == "llm_call": - out["llm_calls"] += 1 - out["tokens_prompt"] += int(data.get("prompt_tokens") or 0) - out["tokens_completion"] += int(data.get("completion_tokens") or 0) - out["tokens_cache_read"] += int(data.get("cache_read_tokens") or 0) - out["tokens_cache_creation"] += int(data.get("cache_creation_tokens") or 0) - out["cost_usd"] += float(data.get("cost_usd") or 0.0) - - elif et == "tool_output": - tool_total += 1 - if data.get("success"): - tool_success += 1 - if first_tool_ts is None and ts is not None and start_dt is not None: - first_tool_ts = (ts - start_dt).total_seconds() - - elif et == "tool_call": - name = data.get("tool") - if name: - tool_calls_by_name[name] += 1 - total_named_tool_calls += 1 - if first_tool_ts is None and ts is not None and start_dt is not None: - first_tool_ts = (ts - start_dt).total_seconds() - - elif et == "error": - had_error = True - - elif et == "undo_complete": - had_undo = True - - elif et == "feedback": - rating = data.get("rating") - if rating == "up": - thumbs_up += 1 - elif rating == "down": - thumbs_down += 1 - - elif et == "hf_job_submit": - jobs_submitted += 1 - - elif et == "hf_job_complete": - flavor = data.get("flavor") or "unknown" - status = (data.get("final_status") or "").lower() - wall = float(data.get("wall_time_s") or 0.0) - gpus = _FLAVOR_GPU_COUNT.get(flavor, 0) - gpu_hours_by_flavor[flavor] += wall * gpus / 3600.0 - if status in ("completed", "succeeded", "success"): - jobs_succeeded += 1 - - elif et == "jobs_access_blocked": - jobs_blocked += 1 - - elif et == "pro_cta_click": - pro_cta_clicks += 1 - source = str(data.get("source") or "unknown") - pro_cta_by_source[source] += 1 - - elif et == "pro_conversion": - pro_conversions += 1 - - elif et == "credits_topped_up": - credits_topped_up += 1 - - elif et == "sandbox_create": - sandboxes_created += 1 - hardware = (data.get("hardware") or "").lower() - # CPU flavors are explicitly named "cpu-*". Everything else - # (including unknown/missing hardware strings) lands in the GPU - # bucket, since the auto-create default is "cpu-basic" which is - # matched here — anything that isn't is almost always an explicit - # GPU choice. - if hardware.startswith("cpu-"): - sandboxes_cpu += 1 - else: - sandboxes_gpu += 1 - - out["tool_calls_total"] = tool_total - out["tool_calls_success"] = tool_success - out["failures"] = 1 if had_error else 0 - out["regenerate_sessions"] = 1 if had_undo else 0 - out["thumbs_up"] = thumbs_up - out["thumbs_down"] = thumbs_down - out["hf_jobs_submitted"] = jobs_submitted - out["hf_jobs_succeeded"] = jobs_succeeded - out["sandboxes_created"] = sandboxes_created - out["sandboxes_cpu"] = sandboxes_cpu - out["sandboxes_gpu"] = sandboxes_gpu - out["hf_jobs_blocked"] = jobs_blocked - out["pro_cta_clicks"] = pro_cta_clicks - out["pro_conversions"] = pro_conversions - out["credits_topped_up"] = credits_topped_up - out["first_tool_s"] = first_tool_ts if first_tool_ts is not None else -1 - out["_gpu_hours_by_flavor"] = dict(gpu_hours_by_flavor) - out["_pro_cta_by_source"] = dict(pro_cta_by_source) - out["_user"] = session.get("user_id") or session.get("session_id") - # Intra-session tool fields. Underscore-prefixed = consumed by _aggregate - # only, never written to CSV directly. - out["_tool_calls_by_name"] = dict(tool_calls_by_name) - out["_research_calls"] = tool_calls_by_name.get("research", 0) - out["_distinct_tools_used"] = len(tool_calls_by_name) - out["_total_named_tool_calls"] = total_named_tool_calls - out["_model_name"] = session.get("model_name") or "unknown" - return dict(out) - - -def _aggregate(per_session: list[dict]) -> dict: - """Collapse a bucket's worth of session rollups into the final KPI row.""" - ttfa_values = [ - s["first_tool_s"] for s in per_session if s.get("first_tool_s", -1) >= 0 - ] - gpu_hours: dict[str, float] = defaultdict(float) - for s in per_session: - for f, h in (s.get("_gpu_hours_by_flavor") or {}).items(): - gpu_hours[f] += h - - # Per-tool aggregates. ``sessions_using_tool`` counts each session at most - # once per tool, so the dashboard can show "how many sessions reached for - # research" alongside "how many research calls overall". - tool_calls_by_name: dict[str, int] = defaultdict(int) - sessions_using_tool: dict[str, int] = defaultdict(int) - sessions_by_model: dict[str, int] = defaultdict(int) - for s in per_session: - for name, count in (s.get("_tool_calls_by_name") or {}).items(): - tool_calls_by_name[name] += int(count) - sessions_using_tool[name] += 1 - sessions_by_model[s.get("_model_name") or "unknown"] += 1 - - # Percentile inputs. All "per session" percentiles exclude sessions that - # never reached for the relevant signal — otherwise quiet hours - # (status-check sessions, abandoned new conversations) drag every median - # to 0 and the chart tells you nothing. - research_calls_nz = [ - s.get("_research_calls", 0) - for s in per_session - if s.get("_research_calls", 0) > 0 - ] - distinct_tools_values = [ - s.get("_distinct_tools_used", 0) - for s in per_session - if s.get("_distinct_tools_used", 0) > 0 - ] - total_calls_values = [ - s.get("_total_named_tool_calls", 0) - for s in per_session - if s.get("_total_named_tool_calls", 0) > 0 - ] - # Per-turn intensity: turns>0 is the natural filter here (a session with - # 5 turns and 0 tools is a meaningful 0). Don't strip those. - calls_per_turn_values = [ - s.get("_total_named_tool_calls", 0) / s["turns"] - for s in per_session - if s.get("turns", 0) > 0 + + +def _active_times( + session: dict[str, Any], + events: list[dict[str, Any]], +) -> tuple[list[str], list[str]]: + event_timestamps = [ + ts + for ts in ( + _parse_ts(event.get("created_at") or event.get("timestamp")) + for event in events + ) + if ts is not None ] + timestamps = list(event_timestamps) + for key in ("session_start_time", "session_end_time"): + ts = _parse_ts(session.get(key)) + if ts is not None: + timestamps.append(ts) + if not timestamps: + now = datetime.now(timezone.utc) + timestamps.append(now) + + active_hours = {_hour_key(ts) for ts in timestamps} + start_ts = _parse_ts(session.get("session_start_time")) + end_ts = _parse_ts(session.get("session_end_time")) + if start_ts is not None and end_ts is not None and end_ts >= start_ts: + current = start_ts.replace(minute=0, second=0, microsecond=0) + final = end_ts.replace(minute=0, second=0, microsecond=0) + while current <= final: + active_hours.add(_hour_key(current)) + current += timedelta(hours=1) + + active_dates = {_date_key(ts) for ts in timestamps} + active_dates.update(hour[:10] for hour in active_hours) + return (sorted(active_dates), sorted(active_hours)) + + +def _session_status(events: list[dict[str, Any]]) -> str: + event_types = {event.get("event_type") for event in events} + if "error" in event_types: + return "failed" + if ( + "interrupted" in event_types + or "cancelled" in event_types + or "canceled" in event_types + ): + return "cancelled" + if "turn_complete" in event_types: + return "completed" + return "unknown" + + +def _usage_components(metrics: dict[str, Any]) -> dict[str, float | str]: + app = _as_dict(metrics.get("app_telemetry")) + hf_billing = _as_dict(metrics.get("hf_billing")) + current_session = _as_dict(hf_billing.get("current_session")) + source = str(metrics.get("total_usd_source") or "usage_metrics") + + if hf_billing.get("available") and current_session: + inference = _number(current_session.get("inference_providers_usd")) + jobs = _number(current_session.get("hf_jobs_usd")) + else: + inference = _number(app.get("inference_usd")) + jobs = _number(app.get("hf_jobs_estimated_usd")) + + sandboxes = _number(app.get("sandbox_estimated_usd")) + return { + "usage_total_usd": round(_number(metrics.get("total_usd")), 6), + "usage_inference_providers_usd": round(inference, 6), + "usage_hf_jobs_usd": round(jobs, 6), + "usage_sandboxes_usd": round(sandboxes, 6), + "usage_cost_source": source, + } + - total_sessions = sum(s["sessions"] for s in per_session) - total_turns = sum(s["turns"] for s in per_session) - tokens_prompt = sum(s["tokens_prompt"] for s in per_session) - tokens_cache_read = sum(s["tokens_cache_read"] for s in per_session) - tool_total = sum(s["tool_calls_total"] for s in per_session) - tool_success = sum(s["tool_calls_success"] for s in per_session) - failures = int(sum(s["failures"] for s in per_session)) - regenerates = int(sum(s["regenerate_sessions"] for s in per_session)) - research_calls_total = int(sum(s.get("_research_calls", 0) for s in per_session)) - sessions_with_research = sum( - 1 for s in per_session if s.get("_research_calls", 0) > 0 +def _model_usage_from_events( + session: dict[str, Any], + events: list[dict[str, Any]], + metrics: dict[str, Any], +) -> dict[str, dict[str, int]]: + usage: dict[str, dict[str, int]] = defaultdict( + lambda: { + "sessions": 0, + "calls": 0, + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + } ) + for event in events: + if event.get("event_type") != "llm_call": + continue + data = _event_data(event) + model = str(data.get("model") or session.get("model_name") or "unknown") + input_tokens = ( + _integer(data.get("prompt_tokens")) + + _integer(data.get("cache_read_tokens")) + + _integer(data.get("cache_creation_tokens")) + ) + output_tokens = _integer(data.get("completion_tokens")) + total_tokens = _integer(data.get("total_tokens")) or ( + input_tokens + output_tokens + ) + usage[model]["calls"] += 1 + usage[model]["input_tokens"] += input_tokens + usage[model]["output_tokens"] += output_tokens + usage[model]["total_tokens"] += total_tokens + + if not usage: + llm = _as_dict(metrics.get("llm")) + calls_by_model = _as_dict(llm.get("calls_by_model")) + if calls_by_model: + for model, calls in calls_by_model.items(): + usage[str(model)]["calls"] = _integer(calls) + if len(calls_by_model) == 1: + model = next(iter(calls_by_model)) + usage[str(model)]["input_tokens"] = ( + _integer(llm.get("prompt_tokens")) + + _integer(llm.get("cache_read_tokens")) + + _integer(llm.get("cache_creation_tokens")) + ) + usage[str(model)]["output_tokens"] = _integer( + llm.get("completion_tokens") + ) + usage[str(model)]["total_tokens"] = _integer(llm.get("total_tokens")) + else: + model = str(session.get("model_name") or "unknown") + usage[model]["calls"] = _integer(llm.get("calls")) + usage[model]["input_tokens"] = ( + _integer(llm.get("prompt_tokens")) + + _integer(llm.get("cache_read_tokens")) + + _integer(llm.get("cache_creation_tokens")) + ) + usage[model]["output_tokens"] = _integer(llm.get("completion_tokens")) + usage[model]["total_tokens"] = _integer(llm.get("total_tokens")) + + for model_usage in usage.values(): + if model_usage["calls"] or model_usage["total_tokens"]: + model_usage["sessions"] = 1 + return dict(sorted((model, dict(values)) for model, values in usage.items())) + + +def _job_counts(events: list[dict[str, Any]]) -> dict[str, int]: + submitted = 0 + terminal_by_id: dict[str, str] = {} + anonymous_terminal_statuses: list[str] = [] + for idx, event in enumerate(events): + event_type = event.get("event_type") + data = _event_data(event) + if event_type == "hf_job_submit": + submitted += 1 + elif event_type in {"hf_job_complete", "hf_job_cancel"}: + status = _normalize_job_status( + data.get("final_status") or data.get("status") + ) + job_id = data.get("job_id") + if job_id: + terminal_by_id[str(job_id)] = status + else: + anonymous_terminal_statuses.append(status or f"unknown-{idx}") - # Per-session cost percentiles — chart "median session cost" alongside the - # mean so a few $700 outliers don't make you think every session is pricey. - session_costs = [float(s.get("cost_usd") or 0.0) for s in per_session] - cost_p50 = _percentile(session_costs, 0.5) - cost_p95 = _percentile(session_costs, 0.95) + statuses = Counter(terminal_by_id.values()) + statuses.update(anonymous_terminal_statuses) + return { + "hf_jobs_submitted": submitted, + "hf_jobs_completed": int(statuses.get("completed", 0)), + "hf_jobs_failed": int(statuses.get("failed", 0)), + "hf_jobs_cancelled": int(statuses.get("cancelled", 0)), + } - unique_users = {s.get("_user") for s in per_session if s.get("_user")} +def _sandbox_counts( + events: list[dict[str, Any]], metrics: dict[str, Any] +) -> dict[str, int]: + created = cpu = gpu = 0 + for event in events: + if event.get("event_type") != "sandbox_create": + continue + data = _event_data(event) + created += 1 + hardware = str(data.get("hardware") or "cpu-basic").lower() + if hardware.startswith("cpu-"): + cpu += 1 + else: + gpu += 1 + if created: + return { + "sandboxes_created": created, + "sandboxes_cpu": cpu, + "sandboxes_gpu": gpu, + } + + sandboxes = _as_dict(metrics.get("sandboxes")) + hardware = _as_dict(sandboxes.get("hardware")) + created = _integer(sandboxes.get("creates")) + cpu = sum( + _integer(count) + for flavor, count in hardware.items() + if str(flavor).lower().startswith("cpu-") + ) + gpu = max(0, created - cpu) return { - "sessions": total_sessions, - "users": len(unique_users), - "turns": total_turns, - "llm_calls": int(sum(s["llm_calls"] for s in per_session)), - "tokens_prompt": int(tokens_prompt), - "tokens_completion": int(sum(s["tokens_completion"] for s in per_session)), - "tokens_cache_read": int(tokens_cache_read), - "tokens_cache_creation": int( - sum(s["tokens_cache_creation"] for s in per_session) - ), - "cost_usd": round(sum(s["cost_usd"] for s in per_session), 4), - # Per-session cost summaries. - "cost_per_session_mean": round( - sum(s["cost_usd"] for s in per_session) / total_sessions, 6 - ) - if total_sessions > 0 - else 0.0, - "cost_per_session_p50": round(cost_p50, 6), - "cost_per_session_p95": round(cost_p95, 6), - "cache_hit_ratio": round( - tokens_cache_read / (tokens_cache_read + tokens_prompt), 4 + "sandboxes_created": created, + "sandboxes_cpu": int(cpu), + "sandboxes_gpu": int(gpu), + } + + +def _artifact_counts(events: list[dict[str, Any]]) -> dict[str, Any]: + by_hash: dict[str, dict[str, Any]] = {} + for event in events: + if event.get("event_type") != "hub_artifact": + continue + data = _event_data(event) + if data.get("success") is False: + continue + artifact_hash = str(data.get("artifact_hash") or "").strip() + if not artifact_hash: + continue + repo_type = str(data.get("repo_type") or "").strip().lower() + if repo_type not in VALID_REPO_TYPES: + continue + by_hash[artifact_hash] = data + + type_counts: Counter[str] = Counter() + non_sandbox_spaces = 0 + for data in by_hash.values(): + repo_type = str(data.get("repo_type") or "").lower() + type_counts[repo_type] += 1 + if repo_type == "space" and not bool(data.get("is_sandbox")): + non_sandbox_spaces += 1 + + by_type = { + "model": int(type_counts.get("model", 0)), + "dataset": int(type_counts.get("dataset", 0)), + "space": int(type_counts.get("space", 0)), + "non_sandbox_space": int(non_sandbox_spaces), + } + return { + "hub_models_created": by_type["model"], + "hub_datasets_created": by_type["dataset"], + "hub_spaces_created": by_type["space"], + "hub_non_sandbox_spaces_created": by_type["non_sandbox_space"], + "hub_artifacts_by_type_json": _json_dumps(by_type), + } + + +def _session_fact(session: dict[str, Any], salt: str | None = None) -> dict[str, Any]: + salt = _require_hash_salt(salt) + events = _session_events(session) + messages = _session_messages(session) + metrics = _session_usage_metrics(session) + llm = _as_dict(metrics.get("llm")) + active_dates, active_hours = _active_times(session, events) + usage = _usage_components(metrics) + jobs = _job_counts(events) + sandboxes = _sandbox_counts(events, metrics) + artifacts = _artifact_counts(events) + model_usage = _model_usage_from_events(session, events, metrics) + + session_id = str(session.get("session_id") or "") + user_id = session.get("user_id") or f"session:{session_id}" + prompt_tokens = _integer(llm.get("prompt_tokens")) + completion_tokens = _integer(llm.get("completion_tokens")) + cache_read_tokens = _integer(llm.get("cache_read_tokens")) + cache_creation_tokens = _integer(llm.get("cache_creation_tokens")) + input_tokens = prompt_tokens + cache_read_tokens + cache_creation_tokens + output_tokens = completion_tokens + total_tokens = _integer(llm.get("total_tokens")) or (input_tokens + output_tokens) + + models_used = sorted(model_usage.keys()) + return { + "schema_version": SCHEMA_VERSION, + "session_id_hash": _hash_identifier(session_id, salt), + "user_id_hash": _hash_identifier(user_id, salt), + "user_plan": _normalize_plan(session.get("user_plan")), + "session_start_time": session.get("session_start_time"), + "session_end_time": session.get("session_end_time"), + "active_dates": active_dates, + "active_hours": active_hours, + "turns": sum(1 for msg in messages if msg.get("role") == "user"), + "status": _session_status(events), + "models_used_json": _json_dumps(models_used), + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cache_read_tokens": cache_read_tokens, + "cache_creation_tokens": cache_creation_tokens, + **jobs, + **sandboxes, + **artifacts, + **usage, + "model_usage_json": _json_dumps(model_usage), + } + + +def _fact_active_in_hour(fact: dict[str, Any], hour_key: str) -> bool: + return hour_key in (fact.get("active_hours") or []) + + +def _fact_active_in_day(fact: dict[str, Any], day_key: str) -> bool: + return day_key in (fact.get("active_dates") or []) + + +def _fact_active_in_month(fact: dict[str, Any], month_key: str) -> bool: + return any( + str(day).startswith(f"{month_key}-") for day in fact.get("active_dates") or [] + ) + + +def _avg(values: list[float]) -> float: + return round(sum(values) / len(values), 6) if values else 0.0 + + +def _stats(prefix: str, values: list[float]) -> dict[str, float]: + if not values: + return { + f"{prefix}_min": 0.0, + f"{prefix}_max": 0.0, + f"{prefix}_avg": 0.0, + } + return { + f"{prefix}_min": round(min(values), 6), + f"{prefix}_max": round(max(values), 6), + f"{prefix}_avg": _avg(values), + } + + +def _latest_plan_counts(facts: list[dict[str, Any]]) -> dict[str, int]: + latest: dict[str, tuple[datetime, str]] = {} + for fact in facts: + user_hash = fact.get("user_id_hash") + if not user_hash: + continue + ts = _parse_ts(fact.get("session_end_time")) or _parse_ts( + fact.get("session_start_time") ) - if (tokens_cache_read + tokens_prompt) > 0 - else 0.0, - # Raw reliability COUNTS (these are what the dashboard shows directly). - "tool_calls_total": int(tool_total), - "tool_calls_succeeded": int(tool_success), - "tool_calls_failed": int(tool_total - tool_success), - "errored_sessions": failures, - # Successful = "did not raise an error event". Mutually exclusive - # with errored_sessions; sums with errored_sessions to total sessions. - "successful_sessions": int(total_sessions - failures), - # Regenerated is an orthogonal dimension (the user retried) — a - # session can be both successful and regenerated, or both errored - # and regenerated. - "regenerated_sessions": regenerates, - # Rates kept for backwards compatibility with anything reading the - # KPI dataset directly. - "tool_success_rate": round(tool_success / tool_total, 4) - if tool_total > 0 - else 0.0, - "failure_rate": round(failures / total_sessions, 4) - if total_sessions > 0 - else 0.0, - "regenerate_rate": round(regenerates / total_sessions, 4) - if total_sessions > 0 - else 0.0, - "time_to_first_action_s_p50": round(_percentile(ttfa_values, 0.5), 2), - "time_to_first_action_s_p95": round(_percentile(ttfa_values, 0.95), 2), - "thumbs_up": int(sum(s["thumbs_up"] for s in per_session)), - "thumbs_down": int(sum(s["thumbs_down"] for s in per_session)), - "hf_jobs_submitted": int(sum(s["hf_jobs_submitted"] for s in per_session)), - "hf_jobs_succeeded": int(sum(s["hf_jobs_succeeded"] for s in per_session)), - "sandboxes_created": int( - sum(s.get("sandboxes_created", 0) for s in per_session) - ), - "sandboxes_cpu": int(sum(s.get("sandboxes_cpu", 0) for s in per_session)), - "sandboxes_gpu": int(sum(s.get("sandboxes_gpu", 0) for s in per_session)), - "hf_jobs_blocked": int(sum(s.get("hf_jobs_blocked", 0) for s in per_session)), - "pro_cta_clicks": int(sum(s.get("pro_cta_clicks", 0) for s in per_session)), - "pro_conversions": int(sum(s.get("pro_conversions", 0) for s in per_session)), - "credits_topped_up": int( - sum(s.get("credits_topped_up", 0) for s in per_session) - ), - "gpu_hours_by_flavor_json": json.dumps(dict(gpu_hours), sort_keys=True), - # Research KPIs — answer "is the agent reaching for research?". - "research_calls": research_calls_total, - "sessions_with_research": int(sessions_with_research), - "research_calls_per_session_p50": round(_percentile(research_calls_nz, 0.5), 2), - "research_calls_per_session_p95": round( - _percentile(research_calls_nz, 0.95), 2 - ), - # Intra-session breadth + intensity. p50 + p95 over per-session values. - "distinct_tools_per_session_p50": round( - _percentile(distinct_tools_values, 0.5), 2 - ), - "distinct_tools_per_session_p95": round( - _percentile(distinct_tools_values, 0.95), 2 - ), - "tool_calls_per_session_p50": round(_percentile(total_calls_values, 0.5), 2), - "tool_calls_per_session_p95": round(_percentile(total_calls_values, 0.95), 2), - "tool_calls_per_turn_p50": round(_percentile(calls_per_turn_values, 0.5), 2), - "tool_calls_per_turn_p95": round(_percentile(calls_per_turn_values, 0.95), 2), - # JSON columns let the dashboard add/remove tools without schema churn. - "tool_calls_by_name_json": json.dumps(dict(tool_calls_by_name), sort_keys=True), - "sessions_using_tool_json": json.dumps( - dict(sessions_using_tool), sort_keys=True - ), - # Surface split by selected model for dashboard drilldowns. - "sessions_by_model_json": json.dumps(dict(sessions_by_model), sort_keys=True), + if ts is None: + ts = datetime.min.replace(tzinfo=timezone.utc) + plan = _normalize_plan(fact.get("user_plan")) + previous = latest.get(user_hash) + if previous is None or ts >= previous[0]: + latest[user_hash] = (ts, plan) + counter = Counter(plan for _, plan in latest.values()) + return { + "free_users": int(counter.get("free", 0)), + "pro_users": int(counter.get("pro", 0)), + "unknown_plan_users": int(counter.get("unknown", 0)), } -# Back-compat alias: older tests call _aggregate_day. -_aggregate_day = _aggregate +def _merge_model_usage(facts: list[dict[str, Any]]) -> dict[str, dict[str, int]]: + merged: dict[str, dict[str, int]] = defaultdict( + lambda: { + "sessions": 0, + "calls": 0, + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + } + ) + sessions_by_model: dict[str, set[str]] = defaultdict(set) + for fact in facts: + session_hash = str(fact.get("session_id_hash") or "") + usage = _as_dict(fact.get("model_usage_json")) + for model, raw_values in usage.items(): + values = _as_dict(raw_values) + if session_hash: + sessions_by_model[str(model)].add(session_hash) + merged[str(model)]["calls"] += _integer(values.get("calls")) + merged[str(model)]["input_tokens"] += _integer(values.get("input_tokens")) + merged[str(model)]["output_tokens"] += _integer(values.get("output_tokens")) + merged[str(model)]["total_tokens"] += _integer(values.get("total_tokens")) + for model, sessions in sessions_by_model.items(): + merged[model]["sessions"] = len(sessions) + return dict(sorted((model, dict(values)) for model, values in merged.items())) + + +def _rollup(facts: list[dict[str, Any]], bucket: str) -> dict[str, Any]: + facts = list(facts) + unique_users = { + fact.get("user_id_hash") for fact in facts if fact.get("user_id_hash") + } + unique_sessions = { + fact.get("session_id_hash") for fact in facts if fact.get("session_id_hash") + } + status_counts = Counter(str(fact.get("status") or "unknown") for fact in facts) + plan_counts = _latest_plan_counts(facts) + model_usage = _merge_model_usage(facts) + + row: dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "bucket": bucket, + "active_users": len(unique_users), + "active_sessions": len(unique_sessions), + **plan_counts, + "completed_sessions": int(status_counts.get("completed", 0)), + "failed_sessions": int(status_counts.get("failed", 0)), + "cancelled_sessions": int(status_counts.get("cancelled", 0)), + "unknown_status_sessions": int(status_counts.get("unknown", 0)), + } + stat_fields = [ + "turns", + "input_tokens", + "output_tokens", + "total_tokens", + "usage_total_usd", + "usage_inference_providers_usd", + "usage_hf_jobs_usd", + "usage_sandboxes_usd", + ] + for field in stat_fields: + row.update(_stats(field, [_number(fact.get(field)) for fact in facts])) + + sum_fields = [ + "input_tokens", + "output_tokens", + "total_tokens", + "prompt_tokens", + "completion_tokens", + "cache_read_tokens", + "cache_creation_tokens", + "hf_jobs_submitted", + "hf_jobs_completed", + "hf_jobs_failed", + "hf_jobs_cancelled", + "sandboxes_created", + "sandboxes_cpu", + "sandboxes_gpu", + "hub_models_created", + "hub_datasets_created", + "hub_spaces_created", + "hub_non_sandbox_spaces_created", + "usage_total_usd", + "usage_inference_providers_usd", + "usage_hf_jobs_usd", + "usage_sandboxes_usd", + ] + for field in sum_fields: + total = sum(_number(fact.get(field)) for fact in facts) + if field.endswith("_usd"): + row[f"total_{field}"] = round(total, 6) + else: + row[f"total_{field}"] = int(total) + + row["model_usage_json"] = _json_dumps(model_usage) + return row -def _csv_cell(v: Any) -> str: - s = str(v) - if "," in s or '"' in s or "\n" in s: - return '"' + s.replace('"', '""') + '"' - return s +def _hourly_rollup(facts: list[dict[str, Any]], hour_key: str) -> dict[str, Any]: + return _rollup( + [fact for fact in facts if _fact_active_in_hour(fact, hour_key)], hour_key + ) -def _write_csv( - api, - row: dict, - bucket_key: str, - path_in_repo: str, - target_repo: str, - token: str, -) -> None: - """Render ``row`` to CSV with a leading ``bucket`` column and upload. - - ``bucket_key`` is the hour string (ISO ``YYYY-MM-DDTHH``) or date string; - written as the ``bucket`` column so downstream consumers can union all - CSVs without date-parsing paths. ``api`` is the caller's ``HfApi`` - instance — reused so we don't spin up a fresh one per CSV. - """ - columns = list(row.keys()) + +def _daily_rollup(facts: list[dict[str, Any]], day_key: str) -> dict[str, Any]: + return _rollup( + [fact for fact in facts if _fact_active_in_day(fact, day_key)], day_key + ) + + +def _monthly_rollup(facts: list[dict[str, Any]], month_key: str) -> dict[str, Any]: + return _rollup( + [fact for fact in facts if _fact_active_in_month(fact, month_key)], + month_key, + ) + + +def _csv_bytes(rows: list[dict[str, Any]]) -> bytes: + if not rows: + return b"" buf = io.StringIO() - buf.write(",".join(["bucket", *columns]) + "\n") - buf.write(",".join([bucket_key, *[_csv_cell(row[c]) for c in columns]]) + "\n") + writer = csv.DictWriter(buf, fieldnames=list(rows[0].keys())) + writer.writeheader() + for row in rows: + writer.writerow(row) + return buf.getvalue().encode("utf-8") - with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as tmp: - tmp.write(buf.getvalue()) - tmp_path = tmp.name +def _jsonl_bytes(rows: list[dict[str, Any]]) -> bytes: + return b"".join( + json.dumps(row, sort_keys=True).encode("utf-8") + b"\n" + for row in sorted(rows, key=lambda item: str(item.get("session_id_hash") or "")) + ) + + +def _upload_bytes( + api: Any, + *, + repo_id: str, + token: str, + path_in_repo: str, + content: bytes, + commit_message: str, +) -> None: try: api.create_repo( - repo_id=target_repo, + repo_id=repo_id, repo_type="dataset", - exist_ok=True, token=token, + private=False, + exist_ok=True, ) + except Exception as e: + logger.debug("create_repo(%s) skipped: %s", repo_id, e) + + with tempfile.NamedTemporaryFile(suffix=".tmp", delete=False) as tmp: + tmp.write(content) + tmp_path = tmp.name + try: api.upload_file( path_or_fileobj=tmp_path, path_in_repo=path_in_repo, - repo_id=target_repo, + repo_id=repo_id, repo_type="dataset", token=token, - commit_message=f"KPIs for {bucket_key}", + commit_message=commit_message, ) finally: try: os.unlink(tmp_path) - except Exception: + except OSError: pass +def _write_csv( + api: Any, + *, + repo_id: str, + token: str, + path_in_repo: str, + row: dict[str, Any], +) -> None: + _upload_bytes( + api, + repo_id=repo_id, + token=token, + path_in_repo=path_in_repo, + content=_csv_bytes([row]), + commit_message=f"Update KPI v2 {path_in_repo}", + ) + + +def _fact_start_day(fact: dict[str, Any]) -> str: + start_ts = _parse_ts(fact.get("session_start_time")) + if start_ts: + return _date_key(start_ts) + return str((fact.get("active_dates") or ["unknown"])[0]) + + +def _write_session_facts( + api: Any, + *, + repo_id: str, + token: str, + facts: list[dict[str, Any]], +) -> None: + by_start_day: dict[str, list[dict[str, Any]]] = defaultdict(list) + for fact in facts: + by_start_day[_fact_start_day(fact)].append(fact) + + for day_key, day_facts in by_start_day.items(): + _upload_bytes( + api, + repo_id=repo_id, + token=token, + path_in_repo=f"{V2_PREFIX}/session_facts/{day_key}.jsonl", + content=_jsonl_bytes(day_facts), + commit_message=f"Update KPI v2 session facts {day_key}", + ) + + +def _iter_session_files(api: Any, repo_id: str, day: date, token: str) -> Iterable[str]: + prefix = f"sessions/{day.isoformat()}/" + try: + files = api.list_repo_files(repo_id=repo_id, repo_type="dataset", token=token) + except Exception as e: + logger.warning("list_repo_files(%s) failed: %s", repo_id, e) + return [] + return [ + path for path in files if path.startswith(prefix) and path.endswith(".jsonl") + ] + + +def _download_session(repo_id: str, path: str, token: str) -> dict[str, Any] | None: + from huggingface_hub import hf_hub_download + + try: + local = hf_hub_download( + repo_id=repo_id, + filename=path, + repo_type="dataset", + token=token, + ) + except Exception as e: + logger.warning("hf_hub_download(%s) failed: %s", path, e) + return None + + try: + with open(local, "r") as f: + line = f.readline().strip() + if not line: + return None + row = json.loads(line) + for key in ("messages", "events", "tools", "usage_metrics"): + value = row.get(key) + if isinstance(value, str): + try: + row[key] = json.loads(value) + except (TypeError, json.JSONDecodeError): + row[key] = [] if key in {"messages", "events", "tools"} else {} + return row + except Exception as e: + logger.warning("parse(%s) failed: %s", path, e) + return None + + +def _sessions_for_start_dates( + api: Any, + source_repo: str, + dates: Iterable[date], + token: str, +) -> list[dict[str, Any]]: + sessions: list[dict[str, Any]] = [] + seen_paths: set[str] = set() + for day in sorted(set(dates)): + for path in _iter_session_files(api, source_repo, day, token): + if path in seen_paths: + continue + seen_paths.add(path) + session = _download_session(source_repo, path, token) + if session: + sessions.append(session) + return sessions + + +def _facts_for_start_dates( + api: Any, + source_repo: str, + dates: Iterable[date], + token: str, + salt: str, +) -> list[dict[str, Any]]: + return [ + _session_fact(session, salt=salt) + for session in _sessions_for_start_dates(api, source_repo, dates, token) + ] + + +def _daily_start_dates(day: date) -> list[date]: + # Rollups read the target day and the previous start-date partition to + # include normal midnight-spanning sessions. Sessions active for more than + # one day before this window are not included; ML Intern sessions are + # expected to be short-lived. + return [day - timedelta(days=1), day] + + +def _month_start_dates(month_key: str, through_day: date | None = None) -> list[date]: + year, month = (int(part) for part in month_key.split("-", 1)) + last_day = calendar.monthrange(year, month)[1] + start = date(year, month, 1) + end = date(year, month, last_day) + if through_day is not None and through_day.strftime("%Y-%m") == month_key: + end = min(end, through_day) + dates = [] + # Include the previous day to catch ordinary sessions that started before + # midnight on the first of the month and remained active after it. + current = start - timedelta(days=1) + while current <= end: + dates.append(current) + current += timedelta(days=1) + return dates + + +def _write_daily_rollup( + api: Any, + *, + source_repo: str, + target_repo: str, + day: date, + token: str, + salt: str, +) -> dict[str, Any]: + facts = _facts_for_start_dates( + api, source_repo, _daily_start_dates(day), token, salt + ) + row = _daily_rollup(facts, day.isoformat()) + _write_csv( + api, + repo_id=target_repo, + token=token, + path_in_repo=f"{V2_PREFIX}/daily/{day.isoformat()}.csv", + row=row, + ) + return row + + +def _write_monthly_rollup( + api: Any, + *, + source_repo: str, + target_repo: str, + month_key: str, + token: str, + salt: str, + through_day: date | None = None, +) -> dict[str, Any]: + facts = _facts_for_start_dates( + api, + source_repo, + _month_start_dates(month_key, through_day=through_day), + token, + salt, + ) + row = _monthly_rollup(facts, month_key) + _write_csv( + api, + repo_id=target_repo, + token=token, + path_in_repo=f"{V2_PREFIX}/monthly/{month_key}.csv", + row=row, + ) + return row + + def run_for_hour( - api, + api: Any, + *, source_repo: str, target_repo: str, hour_dt: datetime, token: str, -) -> dict: - """Roll up one UTC hour [hour_dt, hour_dt+1h). - - Reads today's + yesterday's session folders so sessions that crossed - midnight land in the right hourly bucket. - """ - if hour_dt.tzinfo is None: - hour_dt = hour_dt.replace(tzinfo=timezone.utc) - window_start = hour_dt.replace(minute=0, second=0, microsecond=0) - window_end = window_start + timedelta(hours=1) - - # Sessions partition by session_start_time date. A session that started - # at 23:50 yesterday can still emit events in today's first hours, so we - # look at both folders. - candidate_dates = {window_start.date(), (window_start - timedelta(days=1)).date()} - - per_session: list[dict] = [] - for d in sorted(candidate_dates): - for path in _iter_session_files(api, source_repo, d, token): - sess = _download_session(source_repo, path, token) - if not sess: - continue - windowed = _filter_session_to_window(sess, window_start, window_end) - if windowed is None: - continue - per_session.append(_session_metrics(windowed)) + salt: str | None = None, +) -> dict[str, Any]: + salt = _require_hash_salt(salt) + hour_dt = hour_dt.astimezone(timezone.utc).replace( + minute=0, second=0, microsecond=0 + ) + day = hour_dt.date() + facts = _facts_for_start_dates( + api, + source_repo, + _daily_start_dates(day), + token, + salt, + ) + current_day = day.isoformat() + _write_session_facts( + api, + repo_id=target_repo, + token=token, + facts=[fact for fact in facts if _fact_start_day(fact) == current_day], + ) - if not per_session: - logger.info("No sessions in window %s — skipping", window_start.isoformat()) - return {} + hour = _hour_key(hour_dt) + row = _hourly_rollup(facts, hour) + _write_csv( + api, + repo_id=target_repo, + token=token, + path_in_repo=f"{V2_PREFIX}/hourly/{day.isoformat()}/{hour_dt:%H}.csv", + row=row, + ) - row = _aggregate(per_session) - bucket_key = window_start.strftime("%Y-%m-%dT%H") - path_in_repo = ( - f"hourly/{window_start.strftime('%Y-%m-%d')}/{window_start.strftime('%H')}.csv" + _write_daily_rollup( + api, + source_repo=source_repo, + target_repo=target_repo, + day=day, + token=token, + salt=salt, ) - _write_csv(api, row, bucket_key, path_in_repo, target_repo, token) - logger.info( - "Wrote KPIs for %s (%d sessions): %s", - bucket_key, - per_session and len(per_session), - row, + _write_monthly_rollup( + api, + source_repo=source_repo, + target_repo=target_repo, + month_key=day.strftime("%Y-%m"), + token=token, + salt=salt, + through_day=day, ) return row -# Back-compat for daily backfills — unchanged behaviour. -def run_for_day(api, source_repo: str, target_repo: str, day: date, token: str) -> dict: - paths = _iter_session_files(api, source_repo, day, token) - per_session: list[dict] = [] - for path in paths: - sess = _download_session(source_repo, path, token) - if not sess: - continue - per_session.append(_session_metrics(sess)) - if not per_session: - logger.info("No sessions found for %s — skipping", day) - return {} - row = _aggregate(per_session) - path_in_repo = f"daily/{day.isoformat()}.csv" - _write_csv(api, row, day.isoformat(), path_in_repo, target_repo, token) +def run_for_day( + api: Any, + *, + source_repo: str, + target_repo: str, + day: date, + token: str, + salt: str | None = None, +) -> dict[str, Any]: + salt = _require_hash_salt(salt) + row = _write_daily_rollup( + api, + source_repo=source_repo, + target_repo=target_repo, + day=day, + token=token, + salt=salt, + ) + _write_monthly_rollup( + api, + source_repo=source_repo, + target_repo=target_repo, + month_key=day.strftime("%Y-%m"), + token=token, + salt=salt, + through_day=day, + ) return row -def _parse_hour_arg(s: str) -> datetime: - """Accept ``YYYY-MM-DDTHH`` or full ISO — always pinned to the start of the hour, UTC.""" - dt = datetime.fromisoformat(s) +def _parse_hour(value: str) -> datetime: + if len(value) == 13: + value = value + ":00:00" + dt = datetime.fromisoformat(value.replace("Z", "+00:00")) if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc) - return dt.replace(minute=0, second=0, microsecond=0) + return dt.astimezone(timezone.utc).replace(minute=0, second=0, microsecond=0) def main(argv: list[str] | None = None) -> int: - logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") - ap = argparse.ArgumentParser() - ap.add_argument("--source", default="smolagents/ml-intern-sessions") - ap.add_argument("--target", default="smolagents/ml-intern-kpis") - ap.add_argument( - "--hours", - type=int, - default=1, - help="Number of trailing hours to roll up (default: 1 = last completed hour).", - ) - ap.add_argument( - "--datetime", - type=str, - default=None, - help="Single hour, ISO ``YYYY-MM-DDTHH`` (UTC); overrides --hours.", + parser = argparse.ArgumentParser(description="Build KPI dataset v2 rollups") + parser.add_argument("--source-repo", default="smolagents/ml-intern-sessions") + parser.add_argument("--target-repo", default="smolagents/ml-intern-kpis") + parser.add_argument( + "--hours", type=int, default=1, help="Number of completed hours to build" ) - ap.add_argument( + parser.add_argument("--datetime", help="Explicit UTC hour, e.g. 2026-04-24T14") + parser.add_argument( "--daily-backfill", - type=str, - default=None, - help="Escape hatch: aggregate a whole day at once (YYYY-MM-DD). " - "Writes to daily/.csv. Use for historical backfill only.", + type=int, + default=0, + help="Backfill N UTC days ending yesterday", ) - args = ap.parse_args(argv) + parser.add_argument("--month", help="Build one monthly rollup, e.g. 2026-06") + args = parser.parse_args(argv) - token = ( - os.environ.get("HF_KPI_WRITE_TOKEN") - or os.environ.get("HF_SESSION_UPLOAD_TOKEN") - or os.environ.get("HF_TOKEN") - or os.environ.get("HF_ADMIN_TOKEN") - ) + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + + token = os.environ.get("HF_KPI_WRITE_TOKEN") or os.environ.get("HF_TOKEN") if not token: - logger.error( - "No HF token found. Set one of: HF_KPI_WRITE_TOKEN, " - "HF_SESSION_UPLOAD_TOKEN, HF_TOKEN, HF_ADMIN_TOKEN." - ) - return 1 + logger.error("HF_KPI_WRITE_TOKEN or HF_TOKEN is required") + return 2 + try: + salt = _require_hash_salt() + except RuntimeError as e: + logger.error("%s", e) + return 2 from huggingface_hub import HfApi api = HfApi() - - if args.daily_backfill: - run_for_day( + if args.month: + _write_monthly_rollup( api, - args.source, - args.target, - date.fromisoformat(args.daily_backfill), - token, + source_repo=args.source_repo, + target_repo=args.target_repo, + month_key=args.month, + token=token, + salt=salt, ) return 0 + if args.daily_backfill: + yesterday = datetime.now(timezone.utc).date() - timedelta(days=1) + for offset in range(args.daily_backfill): + day = yesterday - timedelta(days=offset) + run_for_day( + api, + source_repo=args.source_repo, + target_repo=args.target_repo, + day=day, + token=token, + salt=salt, + ) + return 0 + if args.datetime: - target_hours = [_parse_hour_arg(args.datetime)] + hours = [_parse_hour(args.datetime)] else: - now = datetime.now(timezone.utc).replace(minute=0, second=0, microsecond=0) - # Roll up *completed* hours: start from the hour before ``now``. - target_hours = [now - timedelta(hours=i) for i in range(1, args.hours + 1)] - - for hour in target_hours: - run_for_hour(api, args.source, args.target, hour, token) + last_completed = datetime.now(timezone.utc).replace( + minute=0, second=0, microsecond=0 + ) - timedelta(hours=1) + hours = [ + last_completed - timedelta(hours=offset) for offset in range(args.hours) + ] + + for hour_dt in hours: + run_for_hour( + api, + source_repo=args.source_repo, + target_repo=args.target_repo, + hour_dt=hour_dt, + token=token, + salt=salt, + ) return 0 diff --git a/tests/unit/test_build_kpis.py b/tests/unit/test_build_kpis.py index 2f4ce80f..cca01ef8 100644 --- a/tests/unit/test_build_kpis.py +++ b/tests/unit/test_build_kpis.py @@ -1,16 +1,18 @@ -"""Unit tests for the KPI rollup math. - -We exercise the pure functions (``_session_metrics`` and ``_aggregate_day``) -on hand-crafted session trajectories — no network, no HF Hub. -""" +"""Unit tests for KPI dataset v2 facts and rollups.""" import importlib.util +import json import sys +from datetime import datetime, timezone from pathlib import Path +import pytest + + +SALT = "stable-test-salt" + def _load(): - """Load ``scripts/build_kpis.py`` without treating ``scripts`` as a package.""" path = Path(__file__).parent.parent.parent / "scripts" / "build_kpis.py" spec = importlib.util.spec_from_file_location("build_kpis", path) mod = importlib.util.module_from_spec(spec) @@ -19,413 +21,391 @@ def _load(): return mod -def _ev(event_type, data=None, ts="2026-04-24T10:00:00"): +def _ev(event_type, data=None, ts="2026-06-01T10:00:00+00:00"): return {"timestamp": ts, "event_type": event_type, "data": data or {}} -def _session(events, user_id="u1", start="2026-04-24T09:59:00"): +def _usage_metrics( + *, + total_usd=0.0, + inference_usd=0.0, + jobs_usd=0.0, + sandboxes_usd=0.0, + prompt=0, + completion=0, + cache_read=0, + cache_creation=0, + total_tokens=0, +): return { - "session_id": "sess-" + user_id, - "session_start_time": start, - "session_end_time": "2026-04-24T10:05:00", - "model_name": "anthropic/claude-opus-4.8:fal-ai", - "messages": [{"role": "user", "content": "hi"}], - "events": events, + "total_usd": total_usd, + "total_usd_source": "hf_billing_plus_sandbox_estimate", + "app_total_usd": total_usd, + "app_telemetry": { + "inference_usd": inference_usd, + "hf_jobs_estimated_usd": jobs_usd, + "sandbox_estimated_usd": sandboxes_usd, + }, + "hf_billing": { + "available": True, + "current_session": { + "inference_providers_usd": inference_usd, + "hf_jobs_usd": jobs_usd, + }, + }, + "llm": { + "calls": 1, + "calls_by_model": {"model-a": 1}, + "prompt_tokens": prompt, + "completion_tokens": completion, + "cache_read_tokens": cache_read, + "cache_creation_tokens": cache_creation, + "total_tokens": total_tokens, + }, + "hf_jobs": {"submits": 0, "estimated_usd": jobs_usd}, + "sandboxes": {"creates": 0, "estimated_usd": sandboxes_usd}, + } + + +def _session( + *, + session_id="session-1", + user_id="user-1", + user_plan="free", + start="2026-06-01T10:00:00+00:00", + end="2026-06-01T10:05:00+00:00", + events=None, + messages=None, + usage_metrics=None, +): + return { + "session_id": session_id, "user_id": user_id, + "user_plan": user_plan, + "session_start_time": start, + "session_end_time": end, + "model_name": "model-a", + "messages": messages or [{"role": "user", "content": "hi"}], + "events": events or [], + "usage_metrics": usage_metrics or _usage_metrics(), } -def test_llm_call_accumulates_tokens_and_cost(): - mod = _load() - events = [ - _ev( - "llm_call", - { - "prompt_tokens": 100, - "completion_tokens": 50, - "cache_read_tokens": 40, - "cache_creation_tokens": 10, - "cost_usd": 0.01, - }, - ), - _ev( - "llm_call", - { - "prompt_tokens": 200, - "completion_tokens": 100, - "cache_read_tokens": 80, - "cost_usd": 0.02, - }, - ), - ] - m = mod._session_metrics(_session(events)) - assert m["llm_calls"] == 2 - assert m["tokens_prompt"] == 300 - assert m["tokens_completion"] == 150 - assert m["tokens_cache_read"] == 120 - assert m["tokens_cache_creation"] == 10 - assert abs(m["cost_usd"] - 0.03) < 1e-9 +class FakeHfApi: + def __init__(self, files): + self.files = files + self.uploads = [] + def list_repo_files(self, *, repo_id, repo_type, token): + return list(self.files) -def test_tool_success_rate_and_first_action(): - mod = _load() - events = [ - _ev("tool_call", {"tool": "bash"}, ts="2026-04-24T10:00:05"), - _ev("tool_output", {"success": True}), - _ev("tool_output", {"success": False}), - ] - m = mod._session_metrics(_session(events)) - assert m["tool_calls_total"] == 2 - assert m["tool_calls_success"] == 1 - # 65s from start to first action - assert m["first_tool_s"] == 65 + def create_repo(self, **kwargs): + return None + + def upload_file(self, *, path_or_fileobj, path_in_repo, **kwargs): + with open(path_or_fileobj, "rb") as f: + content = f.read() + self.uploads.append({"path_in_repo": path_in_repo, "content": content}) -def test_hf_job_gpu_hours(): +def test_session_fact_hashes_ids_preserves_plan_and_uses_usage_metrics_for_billing(): mod = _load() - events = [ - _ev("hf_job_submit", {"flavor": "a100-large", "job_id": "j1"}), - _ev( - "hf_job_complete", - { - "flavor": "a100-large", - "final_status": "COMPLETED", - "wall_time_s": 3600, - }, + session = _session( + session_id="raw-session", + user_id="raw-user", + user_plan="pro", + start="2026-06-01T23:50:00+00:00", + end="2026-06-02T00:10:00+00:00", + events=[ + _ev( + "llm_call", + { + "model": "model-a", + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_read_tokens": 2, + "cache_creation_tokens": 3, + "total_tokens": 20, + "cost_usd": 999.0, + }, + ts="2026-06-01T23:59:00+00:00", + ), + _ev("turn_complete", ts="2026-06-02T00:01:00+00:00"), + ], + usage_metrics=_usage_metrics( + total_usd=3.0, + inference_usd=1.0, + jobs_usd=0.5, + sandboxes_usd=0.25, + prompt=10, + completion=5, + cache_read=2, + cache_creation=3, + total_tokens=20, ), - ] - m = mod._session_metrics(_session(events)) - assert m["hf_jobs_submitted"] == 1 - assert m["hf_jobs_succeeded"] == 1 - # a100-large = 1 gpu * 1 hour = 1 gpu-hour - assert abs(m["_gpu_hours_by_flavor"]["a100-large"] - 1.0) < 1e-6 - + ) -def test_hf_job_blocked_and_pro_clicks_are_counted(): - mod = _load() - events = [ - _ev("jobs_access_blocked", {"tool_call_ids": ["tc1"], "plan": "free"}), - _ev("pro_cta_click", {"source": "hf_jobs_upgrade_dialog"}), - _ev("pro_cta_click", {"source": "claude_cap_dialog"}), - ] - m = mod._session_metrics(_session(events)) - assert m["hf_jobs_blocked"] == 1 - assert m["pro_cta_clicks"] == 2 - assert m["_pro_cta_by_source"] == { - "hf_jobs_upgrade_dialog": 1, - "claude_cap_dialog": 1, + fact = mod._session_fact(session, salt=SALT) + second = mod._session_fact(session, salt=SALT) + + assert fact["session_id_hash"] == second["session_id_hash"] + assert fact["user_id_hash"] == second["user_id_hash"] + assert fact["session_id_hash"] != "raw-session" + assert fact["user_id_hash"] != "raw-user" + assert fact["user_plan"] == "pro" + assert fact["active_dates"] == ["2026-06-01", "2026-06-02"] + assert fact["active_hours"] == ["2026-06-01T23", "2026-06-02T00"] + assert fact["input_tokens"] == 15 + assert fact["output_tokens"] == 5 + assert fact["total_tokens"] == 20 + assert fact["usage_total_usd"] == 3.0 + assert fact["usage_inference_providers_usd"] == 1.0 + assert fact["usage_hf_jobs_usd"] == 0.5 + assert fact["usage_sandboxes_usd"] == 0.25 + + model_usage = json.loads(fact["model_usage_json"]) + assert model_usage == { + "model-a": { + "sessions": 1, + "calls": 1, + "input_tokens": 15, + "output_tokens": 5, + "total_tokens": 20, + } } -def test_pro_conversions_and_credits_topped_up_per_session(): +def test_session_fact_requires_hash_salt(): mod = _load() - events = [ - _ev("pro_conversion", {"first_seen_at": "2026-04-20T10:00:00"}), - _ev("credits_topped_up", {"namespace": "smolagents"}), - _ev("credits_topped_up", {"namespace": "smolagents"}), - ] - m = mod._session_metrics(_session(events)) - assert m["pro_conversions"] == 1 - assert m["credits_topped_up"] == 2 + with pytest.raises(RuntimeError): + mod._session_fact(_session(), salt="") -def test_aggregate_sums_pro_conversions_and_credits_topped_up(): - mod = _load() - s1 = mod._session_metrics( - _session( - [ - _ev("pro_conversion", {}), - ], - user_id="u1", - ) - ) - s2 = mod._session_metrics( - _session( - [ - _ev("credits_topped_up", {"namespace": "ns"}), - ], - user_id="u2", - ) - ) - s3 = mod._session_metrics(_session([], user_id="u3")) - row = mod._aggregate([s1, s2, s3]) - assert row["pro_conversions"] == 1 - assert row["credits_topped_up"] == 1 - -def test_feedback_counts(): +def test_daily_rollup_dedupes_users_and_uses_latest_plan(): mod = _load() - events = [ - _ev("feedback", {"rating": "up"}), - _ev("feedback", {"rating": "up"}), - _ev("feedback", {"rating": "down"}), + facts = [ + mod._session_fact( + _session( + session_id="s1", + user_id="u1", + user_plan="free", + end="2026-06-01T10:00:00+00:00", + messages=[{"role": "user"}], + usage_metrics=_usage_metrics( + total_usd=1.0, + inference_usd=0.5, + prompt=10, + completion=5, + total_tokens=15, + ), + ), + salt=SALT, + ), + mod._session_fact( + _session( + session_id="s2", + user_id="u1", + user_plan="pro", + end="2026-06-01T11:00:00+00:00", + messages=[{"role": "user"}, {"role": "user"}, {"role": "user"}], + usage_metrics=_usage_metrics( + total_usd=3.0, + inference_usd=1.5, + prompt=30, + completion=10, + total_tokens=40, + ), + ), + salt=SALT, + ), + mod._session_fact( + _session( + session_id="s3", + user_id="u2", + user_plan=None, + messages=[{"role": "user"}, {"role": "user"}], + usage_metrics=_usage_metrics( + total_usd=2.0, + inference_usd=1.0, + prompt=20, + completion=5, + total_tokens=25, + ), + ), + salt=SALT, + ), ] - m = mod._session_metrics(_session(events)) - assert m["thumbs_up"] == 2 - assert m["thumbs_down"] == 1 + row = mod._daily_rollup(facts, "2026-06-01") -def test_aggregate_day_cache_hit_and_users(): - mod = _load() - s1 = mod._session_metrics( - _session( - [ - _ev( - "llm_call", - {"prompt_tokens": 100, "cache_read_tokens": 400, "cost_usd": 0.5}, - ) - ], - user_id="u1", - ) - ) - s2 = mod._session_metrics( - _session( - [ - _ev( - "llm_call", - {"prompt_tokens": 200, "cache_read_tokens": 100, "cost_usd": 1.0}, - ) - ], - user_id="u2", - ) - ) - row = mod._aggregate_day([s1, s2]) - assert row["sessions"] == 2 - assert row["users"] == 2 - assert row["tokens_prompt"] == 300 - assert row["tokens_cache_read"] == 500 - # 500 / (500 + 300) = 0.625 - assert abs(row["cache_hit_ratio"] - 0.625) < 1e-9 - assert abs(row["cost_usd"] - 1.5) < 1e-9 + assert row["active_users"] == 2 + assert row["active_sessions"] == 3 + assert row["free_users"] == 0 + assert row["pro_users"] == 1 + assert row["unknown_plan_users"] == 1 + assert row["turns_min"] == 1 + assert row["turns_max"] == 3 + assert row["turns_avg"] == 2.0 + assert row["input_tokens_min"] == 10 + assert row["input_tokens_max"] == 30 + assert row["total_usage_total_usd"] == 6.0 + assert row["usage_total_usd_avg"] == 2.0 -def test_per_tool_counts_in_session_metrics(): +def test_job_artifact_and_model_rollups_normalize_and_dedupe(): mod = _load() + model_hash = "hash-model" + dataset_hash = "hash-dataset" + sandbox_space_hash = "hash-sandbox-space" + trackio_space_hash = "hash-trackio-space" events = [ - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "research"}), - _ev("tool_call", {"tool": "read"}), - _ev("tool_call", {}), # nameless tool_call must be ignored + _ev("hf_job_submit", {"job_id": "j1"}), + _ev("hf_job_submit", {"job_id": "j2"}), + _ev("hf_job_complete", {"job_id": "j1", "final_status": "COMPLETED"}), + _ev("hf_job_complete", {"job_id": "j2", "final_status": "FAILED"}), + _ev("hf_job_complete", {"job_id": "j3", "final_status": "CANCELED"}), + _ev("sandbox_create", {"hardware": "cpu-basic"}), + _ev("sandbox_create", {"hardware": "a10g-large"}), + _ev( + "hub_artifact", + {"repo_type": "model", "artifact_hash": model_hash, "success": True}, + ), + _ev( + "hub_artifact", + {"repo_type": "model", "artifact_hash": model_hash, "success": True}, + ), + _ev( + "hub_artifact", + {"repo_type": "dataset", "artifact_hash": dataset_hash, "success": True}, + ), + _ev( + "hub_artifact", + { + "repo_type": "space", + "artifact_hash": sandbox_space_hash, + "is_sandbox": True, + "success": True, + }, + ), + _ev( + "hub_artifact", + { + "repo_type": "space", + "artifact_hash": trackio_space_hash, + "is_sandbox": False, + "success": True, + }, + ), + _ev( + "llm_call", + { + "model": "model-b", + "prompt_tokens": 4, + "completion_tokens": 6, + "total_tokens": 10, + }, + ), ] - m = mod._session_metrics(_session(events, user_id="u1")) - assert m["_tool_calls_by_name"] == {"bash": 2, "research": 1, "read": 1} - assert m["_research_calls"] == 1 - assert m["_distinct_tools_used"] == 3 - assert m["_total_named_tool_calls"] == 4 - assert m["_model_name"] == "anthropic/claude-opus-4.8:fal-ai" - - -def test_aggregate_research_kpis_only_count_doer_sessions(): - mod = _load() - s1 = mod._session_metrics( - _session( - [ - _ev("tool_call", {"tool": "research"}), - _ev("tool_call", {"tool": "research"}), - _ev("tool_call", {"tool": "research"}), - ], - user_id="u1", - ) - ) - s2 = mod._session_metrics( + fact = mod._session_fact( _session( - [ - _ev("tool_call", {"tool": "research"}), - ], - user_id="u2", - ) - ) - s3 = mod._session_metrics( - _session( - [ - _ev("tool_call", {"tool": "bash"}), - ], - user_id="u3", - ) + events=events, + usage_metrics=_usage_metrics(prompt=4, completion=6, total_tokens=10), + ), + salt=SALT, ) - row = mod._aggregate([s1, s2, s3]) - assert row["sessions"] == 3 - assert row["sessions_with_research"] == 2 - assert row["research_calls"] == 4 - # Median among sessions that did any research = (1, 3) -> 2.0 - assert row["research_calls_per_session_p50"] == 2.0 - -def test_aggregate_tool_breadth_and_intensity(): - import json as _json - - mod = _load() - s1 = mod._session_metrics( - _session( - [ - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "research"}), - ], - user_id="u1", - ) - ) - # Two user turns so calls/turn = 4/2 = 2 - s2 = _session( - [ - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "edit"}), - _ev("tool_call", {"tool": "edit"}), - ], - user_id="u2", - ) - s2["messages"] = [{"role": "user"}, {"role": "user"}] - s2_metrics = mod._session_metrics(s2) - row = mod._aggregate([s1, s2_metrics]) - assert _json.loads(row["tool_calls_by_name_json"]) == { - "bash": 3, - "research": 1, - "edit": 2, - } - assert _json.loads(row["sessions_using_tool_json"]) == { - "bash": 2, - "research": 1, - "edit": 1, + assert fact["hf_jobs_submitted"] == 2 + assert fact["hf_jobs_completed"] == 1 + assert fact["hf_jobs_failed"] == 1 + assert fact["hf_jobs_cancelled"] == 1 + assert fact["sandboxes_created"] == 2 + assert fact["sandboxes_cpu"] == 1 + assert fact["sandboxes_gpu"] == 1 + assert fact["hub_models_created"] == 1 + assert fact["hub_datasets_created"] == 1 + assert fact["hub_spaces_created"] == 2 + assert fact["hub_non_sandbox_spaces_created"] == 1 + + row = mod._daily_rollup([fact], "2026-06-01") + assert row["total_hf_jobs_completed"] == 1 + assert row["total_hub_models_created"] == 1 + assert row["total_hub_non_sandbox_spaces_created"] == 1 + assert json.loads(row["model_usage_json"]) == { + "model-b": { + "sessions": 1, + "calls": 1, + "input_tokens": 4, + "output_tokens": 6, + "total_tokens": 10, + } } - # u1: 2 distinct, u2: 2 distinct -> p50 = 2 - assert row["distinct_tools_per_session_p50"] == 2.0 - # tool_calls_per_session: u1=2, u2=4 -> p50=3 - assert row["tool_calls_per_session_p50"] == 3.0 - # u1: 2 turns(?) — _session() default has one user message, so calls/turn=2/1=2; u2=4/2=2 - assert row["tool_calls_per_turn_p50"] == 2.0 -def test_breadth_intensity_percentiles_exclude_zero_tool_sessions(): - """Sessions that never called a tool would otherwise crush the median.""" +def test_monthly_rollup_filters_by_active_month(): mod = _load() - # Two productive sessions and three idle ones (no tool calls). Without - # the doer-only filter, median of [0,0,0,2,4] = 0, which is useless. - productive_a = mod._session_metrics( + may_fact = mod._session_fact( _session( - [ - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "research"}), - ], - user_id="prod_a", - ) - ) - productive_b = _session( - [ - _ev("tool_call", {"tool": "bash"}), - _ev("tool_call", {"tool": "edit"}), - _ev("tool_call", {"tool": "edit"}), - _ev("tool_call", {"tool": "edit"}), - ], - user_id="prod_b", - ) - productive_b["messages"] = [{"role": "user"}, {"role": "user"}] - productive_b_metrics = mod._session_metrics(productive_b) - idle = [ - mod._session_metrics(_session([], user_id="idle_a")), - mod._session_metrics(_session([], user_id="idle_b")), - mod._session_metrics(_session([], user_id="idle_c")), - ] - row = mod._aggregate([productive_a, productive_b_metrics, *idle]) - # Median of [2 distinct, 2 distinct] = 2 (idle sessions filtered). - assert row["distinct_tools_per_session_p50"] == 2.0 - # Median of [2 calls, 4 calls] = 3 (idle sessions filtered). - assert row["tool_calls_per_session_p50"] == 3.0 - - -def test_pro_clicks_and_blocked_jobs_in_aggregate(): - """The aggregate row keeps pro_cta_clicks + hf_jobs_blocked columns - even if the dashboard doesn't currently chart them — they're cheap to - keep and downstream consumers may still depend on the schema.""" - mod = _load() - s1 = mod._session_metrics( - _session( - [ - _ev("pro_cta_click", {"source": "hf_jobs_upgrade_dialog"}), - _ev("pro_cta_click", {"source": "claude_cap_dialog"}), - _ev("jobs_access_blocked", {}), - ], - user_id="u1", - ) + session_id="may", + start="2026-05-31T23:55:00+00:00", + end="2026-06-01T00:05:00+00:00", + events=[_ev("turn_complete", ts="2026-06-01T00:01:00+00:00")], + ), + salt=SALT, ) - s2 = mod._session_metrics( + july_fact = mod._session_fact( _session( - [ - _ev("jobs_access_blocked", {}), - _ev("jobs_access_blocked", {}), - ], - user_id="u2", - ) + session_id="july", + start="2026-07-01T00:00:00+00:00", + end="2026-07-01T00:05:00+00:00", + ), + salt=SALT, ) - row = mod._aggregate([s1, s2]) - assert row["pro_cta_clicks"] == 2 - assert row["hf_jobs_blocked"] == 3 + row = mod._monthly_rollup([may_fact, july_fact], "2026-06") -def test_aggregate_sessions_by_model_split(): - import json as _json + assert row["active_sessions"] == 1 + assert row["completed_sessions"] == 1 - mod = _load() - s_claude = _session([], user_id="a") - s_claude["model_name"] = "anthropic/claude-opus-4.8:fal-ai" - s_kimi = _session([], user_id="b") - s_kimi["model_name"] = "moonshotai/Kimi-K2.7-Code" - s_kimi2 = _session([], user_id="c") - s_kimi2["model_name"] = "moonshotai/Kimi-K2.7-Code" - row = mod._aggregate( - [ - mod._session_metrics(s_claude), - mod._session_metrics(s_kimi), - mod._session_metrics(s_kimi2), - ] - ) - assert _json.loads(row["sessions_by_model_json"]) == { - "anthropic/claude-opus-4.8:fal-ai": 1, - "moonshotai/Kimi-K2.7-Code": 2, - } - -def test_failure_and_regenerate_rates(): +def test_run_for_hour_uploads_session_facts_only_once_for_current_day(monkeypatch): mod = _load() - s1 = mod._session_metrics(_session([_ev("error", {"error": "boom"})], user_id="a")) - s2 = mod._session_metrics(_session([_ev("undo_complete")], user_id="b")) - s3 = mod._session_metrics(_session([], user_id="c")) - row = mod._aggregate_day([s1, s2, s3]) - assert row["failure_rate"] == round(1 / 3, 4) - assert row["regenerate_rate"] == round(1 / 3, 4) - + files = [ + "sessions/2026-06-01/previous.jsonl", + "sessions/2026-06-02/current.jsonl", + ] + sessions = { + "sessions/2026-06-01/previous.jsonl": _session( + session_id="previous", + start="2026-06-01T23:55:00+00:00", + end="2026-06-02T00:05:00+00:00", + ), + "sessions/2026-06-02/current.jsonl": _session( + session_id="current", + start="2026-06-02T10:00:00+00:00", + end="2026-06-02T10:10:00+00:00", + ), + } + api = FakeHfApi(files) + monkeypatch.setattr( + mod, + "_download_session", + lambda repo_id, path, token: sessions[path], + ) -def test_window_filter_keeps_only_events_in_range(): - from datetime import datetime, timezone + row = mod.run_for_hour( + api, + source_repo="source", + target_repo="target", + hour_dt=datetime(2026, 6, 2, 10, tzinfo=timezone.utc), + token="token", + salt=SALT, + ) - mod = _load() - events = [ - _ev("llm_call", {"prompt_tokens": 100}, ts="2026-04-24T09:45:00"), - _ev("llm_call", {"prompt_tokens": 200}, ts="2026-04-24T10:05:00"), - _ev("tool_call", {"tool": "bash"}, ts="2026-04-24T10:30:00"), - _ev("llm_call", {"prompt_tokens": 400}, ts="2026-04-24T11:10:00"), + assert row["active_sessions"] == 1 + assert [upload["path_in_repo"] for upload in api.uploads] == [ + "v2/session_facts/2026-06-02.jsonl", + "v2/hourly/2026-06-02/10.csv", + "v2/daily/2026-06-02.csv", + "v2/monthly/2026-06.csv", ] - session = _session(events, start="2026-04-24T09:44:00") - # Only events in [10:00, 11:00) should remain. - window_start = datetime(2026, 4, 24, 10, 0, 0, tzinfo=timezone.utc) - window_end = datetime(2026, 4, 24, 11, 0, 0, tzinfo=timezone.utc) - windowed = mod._filter_session_to_window(session, window_start, window_end) - assert windowed is not None - types = [e["event_type"] for e in windowed["events"]] - assert types == ["llm_call", "tool_call"] - # Metrics only reflect in-window events. - m = mod._session_metrics(windowed) - assert m["tokens_prompt"] == 200 - assert m["llm_calls"] == 1 - assert m["tool_calls_total"] == 0 # tool_call not tool_output - - -def test_window_filter_returns_none_when_nothing_in_range(): - from datetime import datetime, timezone - - mod = _load() - events = [_ev("llm_call", {"prompt_tokens": 100}, ts="2026-04-24T09:45:00")] - session = _session(events) - window_start = datetime(2026, 4, 24, 10, 0, 0, tzinfo=timezone.utc) - window_end = datetime(2026, 4, 24, 11, 0, 0, tzinfo=timezone.utc) - assert mod._filter_session_to_window(session, window_start, window_end) is None diff --git a/tests/unit/test_kpis_scheduler.py b/tests/unit/test_kpis_scheduler.py index cba24d7f..91cf1214 100644 --- a/tests/unit/test_kpis_scheduler.py +++ b/tests/unit/test_kpis_scheduler.py @@ -85,10 +85,23 @@ def test_start_is_no_op_when_disabled(monkeypatch): assert mod._scheduler is None # never instantiated +def test_start_is_no_op_without_hash_salt(monkeypatch, caplog): + mod = _load() + mod._scheduler = None + monkeypatch.delenv("ML_INTERN_KPIS_DISABLED", raising=False) + monkeypatch.delenv("KPI_USER_HASH_SALT", raising=False) + + mod.start() + + assert mod._scheduler is None + assert "KPI_USER_HASH_SALT is required" in caplog.text + + def test_start_skips_cleanly_without_apscheduler(monkeypatch): mod = _load() mod._scheduler = None monkeypatch.delenv("ML_INTERN_KPIS_DISABLED", raising=False) + monkeypatch.setenv("KPI_USER_HASH_SALT", "stable-test-salt") # Force the apscheduler import to fail — start() should log and return. real_import = ( diff --git a/tests/unit/test_session_uploader.py b/tests/unit/test_session_uploader.py index b156e804..9164b2e7 100644 --- a/tests/unit/test_session_uploader.py +++ b/tests/unit/test_session_uploader.py @@ -138,6 +138,7 @@ def test_row_payload_scrubs_messages_events_and_tools(tmp_path): data = { "session_id": "session-123", "user_id": "lewtun", + "user_plan": "pro", "session_start_time": "2026-01-01T00:00:00", "session_end_time": "2026-01-01T00:00:03", "model_name": "anthropic/claude-opus-4.8:fal-ai", @@ -163,6 +164,7 @@ def test_row_payload_includes_usage_scalars_and_parseable_metrics(tmp_path): data = { "session_id": "session-123", "user_id": "lewtun", + "user_plan": "pro", "session_start_time": "2026-01-01T00:00:00", "session_end_time": "2026-01-01T00:30:00", "model_name": "anthropic/claude-opus-4.8:fal-ai", @@ -193,6 +195,7 @@ def test_row_payload_includes_usage_scalars_and_parseable_metrics(tmp_path): row = json.loads(tmp_file.read_text()) assert row["session_id"] == "session-123" assert row["user_id"] == "lewtun" + assert row["user_plan"] == "pro" assert row["session_start_time"] == "2026-01-01T00:00:00" assert row["session_end_time"] == "2026-01-01T00:30:00" assert row["model_name"] == "anthropic/claude-opus-4.8:fal-ai" diff --git a/tests/unit/test_telemetry_usage.py b/tests/unit/test_telemetry_usage.py index dfb27bbb..0d01e98e 100644 --- a/tests/unit/test_telemetry_usage.py +++ b/tests/unit/test_telemetry_usage.py @@ -32,6 +32,93 @@ def test_extract_usage_reads_hf_router_cache_write_tokens(): assert usage["cache_creation_tokens"] == 20 +@pytest.mark.asyncio +async def test_record_hub_artifact_hashes_repo_id(monkeypatch): + monkeypatch.setenv("KPI_USER_HASH_SALT", "stable-test-salt") + session = FakeSession() + + await telemetry.record_hub_artifact( + session, + repo_type="model", + repo_id="alice/private-model", + source="hf_repo_git", + private=True, + success=True, + ) + + event = session.events[0] + assert event.event_type == "hub_artifact" + assert event.data["repo_type"] == "model" + assert event.data["artifact_hash"] != "alice/private-model" + assert "repo_id" not in event.data + assert event.data["source"] == "hf_repo_git" + assert event.data["private"] is True + assert event.data["success"] is True + + +@pytest.mark.asyncio +async def test_record_hub_artifact_keeps_sandbox_space_separate(monkeypatch): + monkeypatch.setenv("KPI_USER_HASH_SALT", "stable-test-salt") + session = FakeSession() + + await telemetry.record_hub_artifact( + session, + repo_type="space", + repo_id="alice/ml-intern-sandbox", + source="hf_repo_git", + is_sandbox=True, + ) + + event = session.events[0] + assert event.event_type == "hub_artifact" + assert event.data["repo_type"] == "space" + assert event.data["is_sandbox"] is True + assert "repo_id" not in event.data + + +@pytest.mark.asyncio +async def test_record_hf_job_submit_sanitizes_expected_artifacts(monkeypatch): + monkeypatch.setenv("KPI_USER_HASH_SALT", "stable-test-salt") + session = FakeSession() + + await telemetry.record_hf_job_submit( + session, + SimpleNamespace(id="job-1", url="https://hf.co/jobs/job-1"), + { + "hardware_flavor": "a100-large", + "namespace": "alice", + "expected_hub_artifacts": [ + {"repo_type": "model", "repo_id": "alice/model", "private": True} + ], + }, + image="python:3.12", + job_type="Python", + ) + + data = session.events[0].data + assert data["expected_hub_artifacts_count"] == 1 + assert data["expected_hub_artifacts"][0]["repo_type"] == "model" + assert data["expected_hub_artifacts"][0]["artifact_hash"] != "alice/model" + assert "repo_id" not in data["expected_hub_artifacts"][0] + + +@pytest.mark.asyncio +async def test_record_hf_job_cancel_emits_terminal_cancelled(): + session = FakeSession() + + await telemetry.record_hf_job_cancel( + session, + job_id="job-cancel", + namespace="alice", + ) + + data = session.events[0].data + assert session.events[0].event_type == "hf_job_complete" + assert data["job_id"] == "job-cancel" + assert data["final_status"] == "cancelled" + assert data["cost_estimate_source"] == "manual_cancel" + + @pytest.mark.asyncio async def test_record_hf_job_complete_emits_runtime_cost(monkeypatch): async def fake_catalog():