agentpool-cli 0.1.0__py3-none-any.whl

agentpool/stats/queries.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from typing import Any
+
+from agentpool.models import AgentSession, CapacitySnapshot
+from agentpool.stats.window import Window
+from agentpool.store import Store
+
+
+def list_sessions_in_window(
+    store: Store,
+    window: Window,
+    provider_id: str | None = None,
+    coordinator_id: str | None = None,
+    scope_mine: bool = False,
+) -> list[AgentSession]:
+    clauses = ["created_at >= ?", "created_at < ?"]
+    args: list[Any] = [window.start.isoformat(), window.end.isoformat()]
+    if provider_id:
+        clauses.append("provider_id = ?")
+        args.append(provider_id)
+    where = " WHERE " + " AND ".join(clauses)
+    with store.connect() as conn:
+        rows = conn.execute(f"SELECT * FROM sessions{where} ORDER BY created_at ASC", args).fetchall()
+    sessions = [store._row_to_session(row) for row in rows]
+    if scope_mine and coordinator_id:
+        sessions = [
+            session
+            for session in sessions
+            if (session.metadata or {}).get("coordinator_id") == coordinator_id
+        ]
+    return sessions
+
+
+def list_events_in_window(
+    store: Store,
+    window: Window,
+    event_types: list[str] | None = None,
+    session_ids: set[str] | None = None,
+) -> list[dict[str, Any]]:
+    if session_ids is not None and not session_ids:
+        return []
+    clauses = ["e.ts >= ?", "e.ts < ?"]
+    args: list[Any] = [window.start.isoformat(), window.end.isoformat()]
+    if event_types:
+        clauses.append(f"e.event_type IN ({','.join('?' for _ in event_types)})")
+        args.extend(event_types)
+    if session_ids:
+        clauses.append(f"e.session_id IN ({','.join('?' for _ in session_ids)})")
+        args.extend(sorted(session_ids))
+    where = " WHERE " + " AND ".join(clauses)
+    with store.connect() as conn:
+        rows = conn.execute(
+            f"""
+            SELECT e.*, s.provider_id AS session_provider_id, s.metadata_json AS session_metadata_json
+            FROM events e
+            JOIN sessions s ON s.id = e.session_id
+            {where}
+            ORDER BY e.ts ASC, e.id ASC
+            """,
+            args,
+        ).fetchall()
+    return [
+        {
+            "id": row["id"],
+            "session_id": row["session_id"],
+            "ts": row["ts"],
+            "event_type": row["event_type"],
+            "state": row["state"],
+            "metadata": json.loads(row["metadata_json"]),
+            "provider_id": row["session_provider_id"],
+            "session_metadata": json.loads(row["session_metadata_json"]),
+        }
+        for row in rows
+    ]
+
+
+def usage_snapshot_at_or_before(
+    store: Store,
+    provider_id: str,
+    ts: datetime,
+    max_age_seconds: float,
+) -> CapacitySnapshot | None:
+    with store.connect() as conn:
+        row = conn.execute(
+            """
+            SELECT raw_json, ts
+            FROM usage_snapshots
+            WHERE provider_id = ? AND ts <= ?
+            ORDER BY ts DESC, id DESC
+            LIMIT 1
+            """,
+            (provider_id, ts.isoformat()),
+        ).fetchone()
+    if row is None:
+        return None
+    snapshot_ts = datetime.fromisoformat(row["ts"].replace("Z", "+00:00"))
+    if snapshot_ts.tzinfo is None:
+        snapshot_ts = snapshot_ts.replace(tzinfo=timezone.utc)
+    age = (ts - snapshot_ts).total_seconds()
+    if age > max_age_seconds:
+        return None
+    return CapacitySnapshot.model_validate_json(row["raw_json"])
+
+
+def usage_snapshots_in_window(store: Store, window: Window, provider_id: str | None = None) -> list[CapacitySnapshot]:
+    clauses = ["ts >= ?", "ts < ?"]
+    args: list[Any] = [window.start.isoformat(), window.end.isoformat()]
+    if provider_id:
+        clauses.append("provider_id = ?")
+        args.append(provider_id)
+    where = " WHERE " + " AND ".join(clauses)
+    with store.connect() as conn:
+        rows = conn.execute(
+            f"SELECT raw_json FROM usage_snapshots{where} ORDER BY ts ASC, id ASC",
+            args,
+        ).fetchall()
+    return [CapacitySnapshot.model_validate_json(row["raw_json"]) for row in rows]
+
+
+def has_any_usage_snapshots(store: Store) -> bool:
+    with store.connect() as conn:
+        row = conn.execute("SELECT 1 FROM usage_snapshots LIMIT 1").fetchone()
+    return row is not None
+
+
+def count_usage_snapshots_in_window(store: Store, window: Window, provider_id: str | None = None) -> int:
+    clauses = ["ts >= ?", "ts < ?"]
+    args: list[Any] = [window.start.isoformat(), window.end.isoformat()]
+    if provider_id:
+        clauses.append("provider_id = ?")
+        args.append(provider_id)
+    where = " WHERE " + " AND ".join(clauses)
+    with store.connect() as conn:
+        row = conn.execute(f"SELECT COUNT(*) AS count FROM usage_snapshots{where}", args).fetchone()
+    return int(row["count"]) if row else 0

agentpool/stats/render.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from typing import Any
+
+from rich.console import Group, RenderableType
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+
+
+def render_stats_panel(stats: dict[str, Any], *, plain: bool = False) -> RenderableType | str:
+    if plain:
+        return render_stats_plain(stats)
+
+    sections: list[RenderableType] = []
+    window = stats.get("window", {})
+    header = Text.assemble(
+        ("AgentPool stats ", "bold"),
+        (window.get("label", ""), "cyan"),
+        (f"  scope={stats.get('scope', 'all')}", "dim"),
+    )
+    sections.append(header)
+
+    sessions = stats.get("sessions", {})
+    session_line = (
+        f"sessions: {sessions.get('total', 0)} total | "
+        f"spawned {sessions.get('spawned', 0)} | "
+        f"terminated {sessions.get('terminated', 0)}"
+    )
+    sections.append(Text(session_line))
+
+    parallelism = stats.get("parallelism", {})
+    ratio = parallelism.get("ratio")
+    ratio_text = f"{ratio:.2f}" if ratio is not None else "n/a"
+    sections.append(
+        Text(
+            f"parallelism: ratio {ratio_text} | peak {parallelism.get('peak_concurrent', 0)} "
+            f"| worker-hours {parallelism.get('sum_worker_hours', 0)}"
+        )
+    )
+
+    walls = stats.get("walls", {})
+    walls_line = f"walls: avoided {walls.get('avoided')} | hit {walls.get('hit')} | confidence {walls.get('confidence')}"
+    if walls.get("confidence") == "low":
+        walls_line = f"[yellow]{walls_line}[/yellow]"
+    sections.append(Text.from_markup(walls_line))
+
+    tokens = stats.get("tokens", {})
+    if tokens.get("by_provider"):
+        totals = tokens.get("totals", {})
+        token_prefix = "tokens (partial: claude-code only): "
+        sections.append(
+            Text(
+                f"{token_prefix}input {totals.get('input')} | output {totals.get('output')}"
+            )
+        )
+
+    quota = stats.get("quota", {})
+    if quota:
+        table = Table("Provider", "Current %", "Min", "Max", "Samples", show_header=True, header_style="bold")
+        for provider_id, row in sorted(quota.items()):
+            table.add_row(
+                provider_id,
+                _fmt(row.get("current_remaining_percent")),
+                _fmt(row.get("min_in_window")),
+                _fmt(row.get("max_in_window")),
+                str(row.get("samples", 0)),
+            )
+        sections.append(table)
+
+    data_quality = stats.get("data_quality") or []
+    if data_quality:
+        dq_lines = [f"- {entry.get('code')}: {entry.get('note') or entry.get('impact')}" for entry in data_quality]
+        sections.append(Panel("\n".join(dq_lines), title="data quality", border_style="yellow"))
+
+    return Panel(Group(*sections), title="agentpool stats", border_style="blue")
+
+
+def render_stats_plain(stats: dict[str, Any]) -> str:
+    lines = [
+        f"schema_version={stats.get('schema_version')}",
+        f"scope={stats.get('scope')}",
+        f"window={stats.get('window', {}).get('spec')}",
+    ]
+    sessions = stats.get("sessions", {})
+    lines.append(f"sessions.total={sessions.get('total', 0)}")
+    lines.append(f"sessions.spawned={sessions.get('spawned', 0)}")
+    parallelism = stats.get("parallelism", {})
+    lines.append(f"parallelism.ratio={parallelism.get('ratio')}")
+    lines.append(f"parallelism.peak_concurrent={parallelism.get('peak_concurrent', 0)}")
+    walls = stats.get("walls", {})
+    lines.append(f"walls.avoided={walls.get('avoided')}")
+    lines.append(f"walls.hit={walls.get('hit')}")
+    lines.append(f"walls.confidence={walls.get('confidence')}")
+    return "\n".join(lines)
+
+
+def _fmt(value: Any) -> str:
+    if value is None:
+        return "n/a"
+    if isinstance(value, float):
+        return f"{value:g}"
+    return str(value)

agentpool/stats/window.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+
+from agentpool.models import ToolError
+
+_DURATION_RE = re.compile(r"^(\d+)(h|d|w)$", re.I)
+
+
+@dataclass(frozen=True)
+class Window:
+    start: datetime
+    end: datetime
+    label: str
+    spec: str
+
+
+def parse_window(spec: str, now: datetime | None = None) -> Window:
+    normalized = (spec or "").strip()
+    if not normalized:
+        raise ToolError("INVALID_WINDOW", "Window spec must not be empty.", {"spec": spec})
+
+    current = now or datetime.now(timezone.utc)
+    if current.tzinfo is None:
+        current = current.replace(tzinfo=timezone.utc)
+
+    lowered = normalized.lower()
+    if lowered == "all":
+        return Window(
+            start=datetime.fromtimestamp(0, tz=timezone.utc),
+            end=current,
+            label="all time",
+            spec="all",
+        )
+
+    if "/" in normalized:
+        start_text, end_text = normalized.split("/", 1)
+        start = _parse_iso(start_text.strip())
+        end = _parse_iso(end_text.strip())
+        if end <= start:
+            raise ToolError(
+                "INVALID_WINDOW",
+                "Window end must be after start.",
+                {"spec": spec, "start": start.isoformat(), "end": end.isoformat()},
+            )
+        return Window(start=start, end=end, label=f"{start.date()} to {end.date()}", spec=normalized)
+
+    duration = _DURATION_RE.match(lowered)
+    if duration:
+        amount = int(duration.group(1))
+        unit = duration.group(2).lower()
+        if unit == "h":
+            delta = timedelta(hours=amount)
+            label = f"last {amount}h"
+        elif unit == "d":
+            delta = timedelta(days=amount)
+            label = f"last {amount}d"
+        else:
+            delta = timedelta(weeks=amount)
+            label = f"last {amount}w"
+        return Window(start=current - delta, end=current, label=label, spec=lowered)
+
+    try:
+        point = _parse_iso(normalized)
+    except ToolError:
+        raise ToolError("INVALID_WINDOW", f"Unrecognized window spec: {spec!r}.", {"spec": spec}) from None
+    end = point + timedelta(days=1)
+    return Window(start=point, end=min(end, current), label=str(point.date()), spec=normalized)
+
+
+def _parse_iso(value: str) -> datetime:
+    text = value.strip()
+    if not text:
+        raise ToolError("INVALID_WINDOW", "ISO timestamp must not be empty.")
+    if text.endswith("Z"):
+        text = f"{text[:-1]}+00:00"
+    try:
+        parsed = datetime.fromisoformat(text)
+    except ValueError as exc:
+        raise ToolError("INVALID_WINDOW", f"Invalid ISO timestamp: {value!r}.", {"value": value}) from exc
+    if parsed.tzinfo is None:
+        parsed = parsed.replace(tzinfo=timezone.utc)
+    return parsed.astimezone(timezone.utc)