gdmcode 0.1.0__py3-none-any.whl

src/enterprise/usage_analytics.py

@@ -0,0 +1,261 @@
+"""
+Usage analytics and cost attribution for enterprise deployments.
+Tracks token usage, tool calls, session costs, and team-level rollups.
+"""
+from __future__ import annotations
+
+import csv
+import sqlite3
+import threading
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+DB_PATH = Path.home() / ".config" / "gdm" / "analytics.db"
+
+
+@dataclass
+class UsageEvent:
+    session_id: str
+    actor_id: str
+    model: str
+    prompt_tokens: int
+    completion_tokens: int
+    tool_calls: int
+    task_id: Optional[str] = None
+    team: Optional[str] = None
+    timestamp: str = field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
+
+    @property
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+    def estimated_cost_usd(
+        self,
+        price_per_1k_input: float = 0.003,
+        price_per_1k_output: float = 0.015,
+    ) -> float:
+        return (
+            self.prompt_tokens / 1000 * price_per_1k_input
+            + self.completion_tokens / 1000 * price_per_1k_output
+        )
+
+
+@dataclass
+class SessionSummary:
+    session_id: str
+    actor_id: str
+    total_prompt_tokens: int
+    total_completion_tokens: int
+    total_tool_calls: int
+    total_cost_usd: float
+    model: str
+    started_at: str
+    last_event_at: str
+
+
+@dataclass
+class TeamRollup:
+    team: str
+    period: str  # "daily" | "weekly" | "monthly"
+    total_cost_usd: float
+    total_tokens: int
+    total_sessions: int
+    top_users: list[dict]  # [{actor_id, cost_usd, sessions}]
+
+
+class UsageAnalytics:
+    def __init__(self, db_path: Path = None):
+        self._path = db_path or DB_PATH
+        self._lock = threading.Lock()
+        self._conn: Optional[sqlite3.Connection] = None
+        self._init_db()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+            self._conn = sqlite3.connect(str(self._path), check_same_thread=False)
+            self._conn.row_factory = sqlite3.Row
+        return self._conn
+
+    def _init_db(self) -> None:
+        with self._lock:
+            conn = self._get_conn()
+            conn.executescript("""
+                CREATE TABLE IF NOT EXISTS usage_events (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    session_id TEXT NOT NULL,
+                    actor_id TEXT NOT NULL,
+                    model TEXT NOT NULL,
+                    prompt_tokens INTEGER NOT NULL DEFAULT 0,
+                    completion_tokens INTEGER NOT NULL DEFAULT 0,
+                    tool_calls INTEGER NOT NULL DEFAULT 0,
+                    task_id TEXT,
+                    team TEXT,
+                    timestamp TEXT NOT NULL,
+                    cost_usd REAL NOT NULL DEFAULT 0.0
+                );
+                CREATE INDEX IF NOT EXISTS idx_usage_session ON usage_events(session_id);
+                CREATE INDEX IF NOT EXISTS idx_usage_actor ON usage_events(actor_id);
+                CREATE INDEX IF NOT EXISTS idx_usage_team ON usage_events(team);
+                CREATE INDEX IF NOT EXISTS idx_usage_timestamp ON usage_events(timestamp);
+                CREATE TABLE IF NOT EXISTS cost_budgets (
+                    team TEXT PRIMARY KEY,
+                    monthly_limit_usd REAL NOT NULL,
+                    alert_threshold REAL NOT NULL DEFAULT 0.8
+                );
+            """)
+            conn.commit()
+
+    def record(self, event: "UsageEvent") -> None:
+        cost = event.estimated_cost_usd()
+        with self._lock:
+            self._get_conn().execute(
+                """INSERT INTO usage_events
+                   (session_id, actor_id, model, prompt_tokens, completion_tokens,
+                    tool_calls, task_id, team, timestamp, cost_usd)
+                   VALUES (?,?,?,?,?,?,?,?,?,?)""",
+                (
+                    event.session_id, event.actor_id, event.model,
+                    event.prompt_tokens, event.completion_tokens, event.tool_calls,
+                    event.task_id, event.team, event.timestamp, cost,
+                ),
+            )
+            self._get_conn().commit()
+
+    def get_session_summary(self, session_id: str) -> "Optional[SessionSummary]":
+        with self._lock:
+            row = self._get_conn().execute(
+                """SELECT session_id, actor_id, model,
+                          SUM(prompt_tokens) as total_prompt,
+                          SUM(completion_tokens) as total_completion,
+                          SUM(tool_calls) as total_tools,
+                          SUM(cost_usd) as total_cost,
+                          MIN(timestamp) as started_at,
+                          MAX(timestamp) as last_event_at
+                   FROM usage_events WHERE session_id=? GROUP BY session_id""",
+                (session_id,),
+            ).fetchone()
+        if row is None:
+            return None
+        return SessionSummary(
+            session_id=row["session_id"],
+            actor_id=row["actor_id"],
+            total_prompt_tokens=row["total_prompt"],
+            total_completion_tokens=row["total_completion"],
+            total_tool_calls=row["total_tools"],
+            total_cost_usd=row["total_cost"],
+            model=row["model"],
+            started_at=row["started_at"],
+            last_event_at=row["last_event_at"],
+        )
+
+    def get_team_rollup(self, team: str, period: str = "monthly") -> "TeamRollup":
+        since = _period_start(period)
+        with self._lock:
+            conn = self._get_conn()
+            agg = conn.execute(
+                """SELECT SUM(cost_usd) as total_cost,
+                          SUM(prompt_tokens + completion_tokens) as total_tokens,
+                          COUNT(DISTINCT session_id) as total_sessions
+                   FROM usage_events WHERE team=? AND timestamp>=?""",
+                (team, since),
+            ).fetchone()
+            top = conn.execute(
+                """SELECT actor_id, SUM(cost_usd) as cost_usd,
+                          COUNT(DISTINCT session_id) as sessions
+                   FROM usage_events WHERE team=? AND timestamp>=?
+                   GROUP BY actor_id ORDER BY cost_usd DESC LIMIT 10""",
+                (team, since),
+            ).fetchall()
+        return TeamRollup(
+            team=team,
+            period=period,
+            total_cost_usd=agg["total_cost"] or 0.0,
+            total_tokens=agg["total_tokens"] or 0,
+            total_sessions=agg["total_sessions"] or 0,
+            top_users=[
+                {"actor_id": r["actor_id"], "cost_usd": r["cost_usd"], "sessions": r["sessions"]}
+                for r in top
+            ],
+        )
+
+    def check_budget(self, team: str) -> dict:
+        """Returns {over_budget, alert, spent_usd, limit_usd}."""
+        since = _period_start("monthly")
+        with self._lock:
+            budget_row = self._get_conn().execute(
+                "SELECT monthly_limit_usd, alert_threshold FROM cost_budgets WHERE team=?",
+                (team,),
+            ).fetchone()
+            if budget_row is None:
+                return {"over_budget": False, "alert": False, "spent_usd": 0.0, "limit_usd": None}
+            spent = (
+                self._get_conn().execute(
+                    "SELECT SUM(cost_usd) as total FROM usage_events WHERE team=? AND timestamp>=?",
+                    (team, since),
+                ).fetchone()["total"]
+                or 0.0
+            )
+        limit = budget_row["monthly_limit_usd"]
+        threshold = budget_row["alert_threshold"]
+        return {
+            "over_budget": spent >= limit,
+            "alert": spent >= limit * threshold,
+            "spent_usd": spent,
+            "limit_usd": limit,
+        }
+
+    def set_budget(
+        self, team: str, monthly_limit_usd: float, alert_threshold: float = 0.8
+    ) -> None:
+        with self._lock:
+            self._get_conn().execute(
+                """INSERT OR REPLACE INTO cost_budgets (team, monthly_limit_usd, alert_threshold)
+                   VALUES (?,?,?)""",
+                (team, monthly_limit_usd, alert_threshold),
+            )
+            self._get_conn().commit()
+
+    def export_csv(
+        self, output_path: "Path", team: str = None, since: str = None
+    ) -> int:
+        clauses, params = [], []
+        if team:
+            clauses.append("team=?")
+            params.append(team)
+        if since:
+            clauses.append("timestamp>=?")
+            params.append(since)
+        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+        with self._lock:
+            rows = self._get_conn().execute(
+                f"SELECT * FROM usage_events {where} ORDER BY timestamp", params
+            ).fetchall()
+        if not rows:
+            return 0
+        with open(output_path, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=rows[0].keys())
+            writer.writeheader()
+            writer.writerows([dict(r) for r in rows])
+        return len(rows)
+
+    def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+
+def _period_start(period: str) -> str:
+    import datetime as _dt
+    now = datetime.now(timezone.utc)
+    if period == "daily":
+        return now.replace(hour=0, minute=0, second=0, microsecond=0).isoformat()
+    if period == "weekly":
+        monday = now - _dt.timedelta(days=now.weekday())
+        return monday.replace(hour=0, minute=0, second=0, microsecond=0).isoformat()
+    # monthly
+    return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0).isoformat()

src/exceptions.py

@@ -0,0 +1,207 @@
+"""Custom exception hierarchy for gdm code.
+
+All gdm exceptions inherit from GdmError so callers can catch the entire
+family with a single `except GdmError` when needed.
+"""
+from __future__ import annotations
+
+__all__ = [
+    "GdmError",
+    "ConfigError",
+    "ApiError",
+    "ApiRateLimitError",
+    "ApiBudgetError",
+    "FatalApiError",
+    "ToolError",
+    "ToolPermissionError",
+    "SecurityError",
+    "InjectionDetectedError",
+    "GitError",
+    "BudgetError",
+    "DatabaseError",
+    "SchemaError",
+    "IndexError",
+    "AgentLoopError",
+    "PermissionDeniedError",
+    "BudgetExceededError",
+    "DirtyWorkingTreeError",
+    "MergeConflictError",
+    "DaemonUnavailableError",
+]
+
+
+class GdmError(Exception):
+    """Base class for all gdm errors."""
+
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+class ConfigError(GdmError):
+    """Invalid or missing configuration (e.g. XAI_API_KEY not set)."""
+
+
+# ---------------------------------------------------------------------------
+# API / model
+# ---------------------------------------------------------------------------
+
+class ApiError(GdmError):
+    """xAI API call failed."""
+
+    def __init__(self, message: str, status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class FatalApiError(ApiError):
+    """Both primary and fallback providers failed — no recovery possible."""
+
+
+class ApiRateLimitError(ApiError):
+    """429 rate limit hit — caller should back off and retry."""
+
+
+class ApiBudgetError(ApiError):
+    """Session cost cap reached — stop calling the API."""
+
+
+# ---------------------------------------------------------------------------
+# Tools
+# ---------------------------------------------------------------------------
+
+class ToolError(GdmError):
+    """A tool execution failed."""
+
+    def __init__(self, tool_name: str, message: str) -> None:
+        super().__init__(f"[{tool_name}] {message}")
+        self.tool_name = tool_name
+
+
+class ToolPermissionError(ToolError):
+    """Tool was denied by ToolPermissionContext."""
+
+    def __init__(self, tool: str, reason: str) -> None:
+        super().__init__(tool_name=tool, message=reason)
+        self.tool = tool
+        self.reason = reason
+
+
+# ---------------------------------------------------------------------------
+# Security
+# ---------------------------------------------------------------------------
+
+class SecurityError(GdmError):
+    """A security invariant was violated."""
+
+
+class InjectionDetectedError(SecurityError):
+    """Prompt injection pattern found in file content."""
+
+    def __init__(self, filename: str, pattern: str) -> None:
+        super().__init__(
+            f"Injection pattern '{pattern}' detected in '{filename}'. "
+            "Content blocked from prompt injection."
+        )
+        self.filename = filename
+        self.pattern = pattern
+
+
+# ---------------------------------------------------------------------------
+# Git
+# ---------------------------------------------------------------------------
+
+class GitError(GdmError):
+    """Git operation failed."""
+
+    def __init__(self, command: str, message: str) -> None:
+        super().__init__(f"git {command}: {message}")
+        self.command = command
+
+
+# ---------------------------------------------------------------------------
+# Cost / budget
+# ---------------------------------------------------------------------------
+
+class BudgetError(GdmError):
+    """Session cost limit exceeded."""
+
+    def __init__(self, limit_usd: float, current_usd: float) -> None:
+        super().__init__(
+            f"Cost limit ${limit_usd:.2f} exceeded (current: ${current_usd:.2f})"
+        )
+        self.limit_usd = limit_usd
+        self.current_usd = current_usd
+
+
+# ---------------------------------------------------------------------------
+# Storage
+# ---------------------------------------------------------------------------
+
+class DatabaseError(GdmError):
+    """SQLite operation failed."""
+
+
+class SchemaError(GdmError):
+    """Pydantic model parse/validation failed."""
+
+
+class IndexError(GdmError):
+    """Code index (Tree-sitter/FAISS) operation failed."""
+
+
+# ---------------------------------------------------------------------------
+# Agent loop
+# ---------------------------------------------------------------------------
+
+class AgentLoopError(GdmError):
+    """Agent loop hit an unrecoverable state (max turns, bad API response, etc.)."""
+
+
+# ---------------------------------------------------------------------------
+# Aliases (backwards-compatible names)
+# ---------------------------------------------------------------------------
+
+class PermissionDeniedError(ToolPermissionError):
+    """Alias for ToolPermissionError — backwards-compatible name."""
+
+
+class BudgetExceededError(ApiBudgetError):
+    """Alias for ApiBudgetError — raised by cost_tracker when limit exceeded."""
+
+
+# ---------------------------------------------------------------------------
+# Branch Farm
+# ---------------------------------------------------------------------------
+
+class DirtyWorkingTreeError(GitError):
+    """Branch farm requires a clean working tree but uncommitted changes exist."""
+
+    def __init__(self) -> None:
+        super().__init__(
+            command="status --porcelain",
+            message=(
+                "Branch farm requires a clean working tree. "
+                "Commit or stash your changes before running a branch farm."
+            ),
+        )
+
+
+class MergeConflictError(GitError):
+    """Dry-run merge detected conflicts; cannot auto-merge the farm branch."""
+
+    def __init__(self, branch: str, details: str = "") -> None:
+        super().__init__(
+            command=f"merge-tree HEAD {branch}",
+            message=f"Conflicts detected when merging {branch!r}.{' ' + details if details else ''}",
+        )
+        self.branch = branch
+
+
+# ---------------------------------------------------------------------------
+# Daemon
+# ---------------------------------------------------------------------------
+
+class DaemonUnavailableError(GdmError):
+    """Daemon crashed too frequently or is otherwise unavailable."""
+