tokenjam 0.2.0__py3-none-any.whl

tokenjam/core/api_backend.py

@@ -0,0 +1,235 @@
+"""
+API-based backend for CLI commands when DuckDB is locked by tj serve.
+Implements the subset of StorageBackend used by CLI commands by routing
+queries through the REST API.
+"""
+from __future__ import annotations
+
+from datetime import date, datetime
+
+import httpx
+
+from tokenjam.core.models import (
+    Alert,
+    AlertFilters,
+    AlertType,
+    CostFilters,
+    CostRow,
+    NormalizedSpan,
+    Severity,
+    SessionRecord,
+    SpanKind,
+    SpanStatus,
+    TraceFilters,
+    TraceRecord,
+)
+from tokenjam.utils.time_parse import utcnow
+
+
+class ApiBackend:
+    """Read-only backend that queries tj serve instead of DuckDB."""
+
+    def __init__(self, base_url: str, api_key: str | None = None) -> None:
+        self.base_url = base_url.rstrip("/")
+        headers: dict[str, str] = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        self.client = httpx.Client(base_url=self.base_url, headers=headers, timeout=10)
+
+    def _get(self, path: str, params: dict | None = None) -> dict:
+        resp = self.client.get(path, params=params)
+        resp.raise_for_status()
+        return resp.json()
+
+    def get_traces(self, filters: TraceFilters) -> list[TraceRecord]:
+        params: dict[str, str | int] = {"limit": filters.limit, "offset": filters.offset}
+        if filters.agent_id:
+            params["agent_id"] = filters.agent_id
+        if filters.since:
+            params["since"] = filters.since.isoformat()
+        if filters.until:
+            params["until"] = filters.until.isoformat()
+        if filters.status:
+            params["status"] = filters.status
+        if filters.span_name:
+            params["span_name"] = filters.span_name
+        data = self._get("/api/v1/traces", params)
+        return [
+            TraceRecord(
+                trace_id=t["trace_id"],
+                agent_id=t["agent_id"],
+                name=t["name"],
+                start_time=datetime.fromisoformat(t["start_time"]) if t.get("start_time") else None,
+                duration_ms=t.get("duration_ms"),
+                cost_usd=t.get("cost_usd"),
+                status_code=t.get("status_code", "ok"),
+                span_count=t.get("span_count", 0),
+            )
+            for t in data.get("traces", [])
+        ]
+
+    def get_trace_spans(self, trace_id: str) -> list[NormalizedSpan]:
+        data = self._get(f"/api/v1/traces/{trace_id}")
+        return [_dict_to_span(s) for s in data.get("spans", [])]
+
+    def get_cost_summary(self, filters: CostFilters) -> list[CostRow]:
+        params: dict[str, str] = {}
+        if filters.agent_id:
+            params["agent_id"] = filters.agent_id
+        if filters.since:
+            params["since"] = filters.since.isoformat()
+        if filters.until:
+            params["until"] = filters.until.isoformat()
+        if filters.group_by:
+            params["group_by"] = filters.group_by
+        data = self._get("/api/v1/cost", params)
+        return [
+            CostRow(
+                group=r["group"],
+                agent_id=r["agent_id"],
+                model=r["model"],
+                input_tokens=r.get("input_tokens", 0),
+                output_tokens=r.get("output_tokens", 0),
+                cost_usd=r.get("cost_usd", 0.0),
+            )
+            for r in data.get("rows", [])
+        ]
+
+    def get_alerts(self, filters: AlertFilters) -> list[Alert]:
+        params: dict[str, str | int | bool] = {}
+        if filters.agent_id:
+            params["agent_id"] = filters.agent_id
+        if filters.since:
+            params["since"] = filters.since.isoformat()
+        if filters.severity:
+            params["severity"] = filters.severity.value
+        if filters.type:
+            params["type"] = filters.type.value
+        if filters.unread:
+            params["unread"] = True
+        data = self._get("/api/v1/alerts", params)
+        return [
+            Alert(
+                alert_id=a["alert_id"],
+                fired_at=datetime.fromisoformat(a["fired_at"]),
+                type=AlertType(a["type"]),
+                severity=Severity(a["severity"]),
+                title=a["title"],
+                detail=a.get("detail", {}),
+                agent_id=a.get("agent_id"),
+                session_id=a.get("session_id"),
+                span_id=a.get("span_id"),
+                acknowledged=a.get("acknowledged", False),
+                suppressed=a.get("suppressed", False),
+            )
+            for a in data.get("alerts", [])
+        ]
+
+    def get_tool_calls(
+        self, agent_id: str | None, since: datetime | None, tool_name: str | None,
+    ) -> list[dict]:
+        params: dict[str, str] = {}
+        if agent_id:
+            params["agent_id"] = agent_id
+        if since:
+            params["since"] = since.isoformat()
+        if tool_name:
+            params["tool_name"] = tool_name
+        data = self._get("/api/v1/tools", params)
+        return data.get("tools", [])
+
+    def get_daily_cost(self, agent_id: str, day: date) -> float:
+        params: dict[str, str] = {
+            "agent_id": agent_id,
+            "since": datetime(day.year, day.month, day.day).isoformat(),
+            "group_by": "day",
+        }
+        data = self._get("/api/v1/cost", params)
+        return data.get("total_cost_usd", 0.0)
+
+    def get_completed_sessions(self, agent_id: str, limit: int) -> list[SessionRecord]:
+        # Use /api/v1/status which already returns the latest session per agent
+        # with token counts and tool_call_count populated. Without this,
+        # `tj status` over a running server shows every agent as "idle" with
+        # zeros even when sessions exist (U3).
+        try:
+            data = self._get("/api/v1/status", {"agent_id": agent_id})
+        except (httpx.HTTPError, ValueError):
+            return []
+        agents = data.get("agents", [])
+        if not agents:
+            return []
+        a = agents[0]
+        if not a.get("session_id"):
+            return []
+        started_at = a.get("started_at")
+        return [SessionRecord(
+            session_id=a["session_id"],
+            agent_id=a.get("agent_id") or agent_id,
+            started_at=datetime.fromisoformat(started_at) if started_at else utcnow(),
+            ended_at=None,
+            conversation_id=None,
+            status=a.get("status", "completed"),
+            total_cost_usd=a.get("total_cost_usd"),
+            input_tokens=a.get("input_tokens", 0) or 0,
+            output_tokens=a.get("output_tokens", 0) or 0,
+            cache_tokens=0,
+            tool_call_count=a.get("tool_call_count", 0) or 0,
+            error_count=a.get("error_count", 0) or 0,
+        )][:limit]
+
+    def get_completed_session_count(self, agent_id: str) -> int:
+        return 0
+
+    def get_session_cost(self, session_id: str) -> float:
+        return 0.0
+
+    def get_recent_spans(self, session_id: str, limit: int) -> list[NormalizedSpan]:
+        return []
+
+    def close(self) -> None:
+        self.client.close()
+
+
+def _dict_to_span(d: dict) -> NormalizedSpan:
+    return NormalizedSpan(
+        span_id=d["span_id"],
+        trace_id=d["trace_id"],
+        name=d["name"],
+        kind=SpanKind(d.get("kind", "internal")),
+        status_code=SpanStatus(d.get("status_code", "ok")),
+        start_time=datetime.fromisoformat(d["start_time"]) if d.get("start_time") else None,
+        parent_span_id=d.get("parent_span_id"),
+        session_id=d.get("session_id"),
+        agent_id=d.get("agent_id"),
+        end_time=datetime.fromisoformat(d["end_time"]) if d.get("end_time") else None,
+        duration_ms=d.get("duration_ms"),
+        status_message=d.get("status_message"),
+        attributes=d.get("attributes", {}),
+        events=d.get("events", []),
+        provider=d.get("provider"),
+        model=d.get("model"),
+        tool_name=d.get("tool_name"),
+        input_tokens=d.get("input_tokens"),
+        output_tokens=d.get("output_tokens"),
+        cache_tokens=d.get("cache_tokens"),
+        cost_usd=d.get("cost_usd"),
+        request_type=d.get("request_type"),
+        conversation_id=d.get("conversation_id"),
+    )
+
+
+def probe_api(host: str, port: int, api_key: str | None = None) -> ApiBackend | None:
+    """Check if tj serve is running and return an ApiBackend if so."""
+    base_url = f"http://{host}:{port}"
+    try:
+        headers: dict[str, str] = {}
+        if api_key:
+            headers["Authorization"] = f"Bearer {api_key}"
+        resp = httpx.get(f"{base_url}/api/v1/traces", params={"limit": 1},
+                         headers=headers, timeout=2)
+        if resp.status_code in (200, 401):
+            return ApiBackend(base_url, api_key)
+    except (httpx.ConnectError, httpx.TimeoutException):
+        pass
+    return None

tokenjam/core/config.py

@@ -0,0 +1,360 @@
+from __future__ import annotations
+import sys
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib  # type: ignore[no-redef]
+
+import tomli_w
+
+
+# -- Nested config dataclasses --
+
+@dataclass
+class SensitiveAction:
+    name:     str
+    severity: str = "warning"   # critical | warning | info
+
+
+@dataclass
+class BudgetConfig:
+    daily_usd:   float | None = None
+    session_usd: float | None = None
+
+
+@dataclass
+class DriftConfig:
+    enabled:            bool  = True
+    baseline_sessions:  int   = 10
+    token_threshold:    float = 2.0
+    tool_sequence_diff: float = 0.4
+
+
+@dataclass
+class AgentConfig:
+    description:      str                  = ""
+    budget:           BudgetConfig         = field(default_factory=BudgetConfig)
+    sensitive_actions: list[SensitiveAction] = field(default_factory=list)
+    output_schema:    str | None           = None
+    drift:            DriftConfig          = field(default_factory=DriftConfig)
+
+
+@dataclass
+class DefaultsConfig:
+    budget: BudgetConfig = field(default_factory=BudgetConfig)
+
+
+@dataclass
+class StorageConfig:
+    path:           str = "~/.tj/telemetry.duckdb"
+    retention_days: int = 90
+
+
+@dataclass
+class OtlpConfig:
+    enabled:  bool        = False
+    endpoint: str         = "http://localhost:4318"
+    protocol: str         = "http"   # http | grpc
+    headers:  dict        = field(default_factory=dict)
+    insecure: bool        = True
+
+
+@dataclass
+class PrometheusConfig:
+    enabled: bool = True
+    port:    int  = 9464
+    path:    str  = "/metrics"
+
+
+@dataclass
+class ExportConfig:
+    otlp:       OtlpConfig       = field(default_factory=OtlpConfig)
+    prometheus: PrometheusConfig = field(default_factory=PrometheusConfig)
+
+
+@dataclass
+class AlertChannelConfig:
+    type: str
+    # stdout / file
+    path: str | None = None
+    # ntfy
+    topic:        str | None = None
+    server:       str        = "https://ntfy.sh"
+    token:        str        = ""
+    # webhook
+    url:     str | None = None
+    method:  str        = "POST"
+    headers: dict       = field(default_factory=dict)
+    # discord
+    webhook_url: str | None = None
+    # telegram
+    bot_token: str | None = None
+    chat_id:   str | None = None
+    # shared
+    min_severity: str = "info"
+
+
+@dataclass
+class AlertsConfig:
+    cooldown_seconds:        int  = 60
+    include_captured_content: bool = False
+    channels: list[AlertChannelConfig] = field(default_factory=lambda: [
+        AlertChannelConfig(type="stdout"),
+    ])
+
+
+@dataclass
+class SecurityConfig:
+    ingest_secret:          str = ""
+    max_attribute_bytes:    int = 65536
+    max_attributes_per_span: int = 256
+    max_attribute_depth:    int = 10
+    webhook_allowed_domains: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ApiAuthConfig:
+    enabled: bool = False
+    api_key: str  = ""
+
+
+@dataclass
+class ApiConfig:
+    enabled: bool         = True
+    host:    str          = "127.0.0.1"
+    port:    int          = 7391
+    auth:    ApiAuthConfig = field(default_factory=ApiAuthConfig)
+
+
+@dataclass
+class CaptureConfig:
+    prompts:      bool = False
+    completions:  bool = False
+    tool_inputs:  bool = False
+    tool_outputs: bool = False
+
+
+@dataclass
+class TjConfig:
+    version:  str
+    defaults: DefaultsConfig          = field(default_factory=DefaultsConfig)
+    agents:   dict[str, AgentConfig]  = field(default_factory=dict)
+    storage:  StorageConfig           = field(default_factory=StorageConfig)
+    export:   ExportConfig            = field(default_factory=ExportConfig)
+    alerts:   AlertsConfig            = field(default_factory=AlertsConfig)
+    security: SecurityConfig          = field(default_factory=SecurityConfig)
+    api:      ApiConfig               = field(default_factory=ApiConfig)
+    capture:  CaptureConfig           = field(default_factory=CaptureConfig)
+    # Path to the config file on disk; set by load_config() so that relative
+    # paths in the config (e.g. output_schema) can be resolved correctly.
+    config_path: Path | None          = field(default=None, repr=False, compare=False)
+
+
+# -- File discovery --
+
+SEARCH_PATHS = [
+    Path("tokenjam.toml"),
+    Path(".tj/config.toml"),
+    Path.home() / ".config" / "tj" / "config.toml",
+]
+
+
+def find_config_file(override: str | None = None) -> Path | None:
+    if override:
+        p = Path(override)
+        if p.exists():
+            return p
+        raise FileNotFoundError(f"Config file not found: {override}")
+    for path in SEARCH_PATHS:
+        if path.exists():
+            return path
+    return None
+
+
+def load_config(path: str | None = None) -> TjConfig:
+    """
+    Load config from file, merge with defaults, return TjConfig.
+
+    IMPORTANT: tomllib requires binary mode "rb" -- not text mode "r".
+    Using "r" raises TypeError at runtime.
+    """
+    config_path = find_config_file(path)
+    if config_path is None:
+        return TjConfig(version="1")
+
+    with open(config_path, "rb") as f:   # "rb" is REQUIRED
+        raw = tomllib.load(f)
+
+    cfg = _parse(raw)
+    cfg.config_path = config_path.resolve()
+    return cfg
+
+
+def write_config(config: TjConfig, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "wb") as f:
+        tomli_w.dump(_serialise(config), f)
+
+
+def _parse(raw: dict) -> TjConfig:
+    """Convert raw TOML dict to TjConfig, applying defaults for missing keys."""
+    agents = {}
+    for agent_id, agent_raw in raw.get("agents", {}).items():
+        budget = BudgetConfig(**agent_raw.get("budget", {}))
+        sensitive_actions = [
+            SensitiveAction(**sa) for sa in agent_raw.get("sensitive_actions", [])
+        ]
+        drift = DriftConfig(**agent_raw.get("drift", {}))
+        agents[agent_id] = AgentConfig(
+            description=agent_raw.get("description", ""),
+            budget=budget,
+            sensitive_actions=sensitive_actions,
+            output_schema=agent_raw.get("output_schema"),
+            drift=drift,
+        )
+
+    storage_raw = raw.get("storage", {})
+    storage = StorageConfig(
+        path=storage_raw.get("path", StorageConfig.path),
+        retention_days=storage_raw.get("retention_days", StorageConfig.retention_days),
+    )
+
+    export_raw = raw.get("export", {})
+    otlp_raw = export_raw.get("otlp", {})
+    otlp = OtlpConfig(
+        enabled=otlp_raw.get("enabled", False),
+        endpoint=otlp_raw.get("endpoint", OtlpConfig.endpoint),
+        protocol=otlp_raw.get("protocol", OtlpConfig.protocol),
+        headers=otlp_raw.get("headers", {}),
+        insecure=otlp_raw.get("insecure", True),
+    )
+    prom_raw = export_raw.get("prometheus", {})
+    prometheus = PrometheusConfig(
+        enabled=prom_raw.get("enabled", True),
+        port=prom_raw.get("port", PrometheusConfig.port),
+        path=prom_raw.get("path", PrometheusConfig.path),
+    )
+    export = ExportConfig(otlp=otlp, prometheus=prometheus)
+
+    alerts_raw = raw.get("alerts", {})
+    channels = []
+    for ch_raw in alerts_raw.get("channels", []):
+        channels.append(AlertChannelConfig(**ch_raw))
+    alerts = AlertsConfig(
+        cooldown_seconds=alerts_raw.get("cooldown_seconds", AlertsConfig.cooldown_seconds),
+        include_captured_content=alerts_raw.get("include_captured_content", False),
+        channels=channels if channels else [AlertChannelConfig(type="stdout")],
+    )
+
+    security_raw = raw.get("security", {})
+    security = SecurityConfig(
+        ingest_secret=security_raw.get("ingest_secret", ""),
+        max_attribute_bytes=security_raw.get("max_attribute_bytes", 65536),
+        max_attributes_per_span=security_raw.get("max_attributes_per_span", 256),
+        max_attribute_depth=security_raw.get("max_attribute_depth", 10),
+        webhook_allowed_domains=security_raw.get("webhook_allowed_domains", []),
+    )
+
+    api_raw = raw.get("api", {})
+    api_auth_raw = api_raw.get("auth", {})
+    api_auth = ApiAuthConfig(
+        enabled=api_auth_raw.get("enabled", False),
+        api_key=api_auth_raw.get("api_key", ""),
+    )
+    api = ApiConfig(
+        enabled=api_raw.get("enabled", True),
+        host=api_raw.get("host", ApiConfig.host),
+        port=api_raw.get("port", ApiConfig.port),
+        auth=api_auth,
+    )
+
+    capture_raw = raw.get("capture", {})
+    capture = CaptureConfig(
+        prompts=capture_raw.get("prompts", False),
+        completions=capture_raw.get("completions", False),
+        tool_inputs=capture_raw.get("tool_inputs", False),
+        tool_outputs=capture_raw.get("tool_outputs", False),
+    )
+
+    defaults_raw = raw.get("defaults", {})
+    defaults_budget_raw = defaults_raw.get("budget", {})
+    defaults = DefaultsConfig(budget=BudgetConfig(**defaults_budget_raw))
+
+    return TjConfig(
+        version=raw.get("version", "1"),
+        defaults=defaults,
+        agents=agents,
+        storage=storage,
+        export=export,
+        alerts=alerts,
+        security=security,
+        api=api,
+        capture=capture,
+    )
+
+
+def _serialise(config: TjConfig) -> dict:
+    """Convert TjConfig back to a plain dict suitable for tomli_w."""
+    def _dc_to_dict(obj: object) -> dict:
+        result = {}
+        for f in fields(obj):  # type: ignore[arg-type]
+            val = getattr(obj, f.name)
+            if isinstance(val, dict):
+                result[f.name] = val
+            elif isinstance(val, list):
+                result[f.name] = [
+                    _dc_to_dict(item) if hasattr(item, "__dataclass_fields__") else item
+                    for item in val
+                ]
+            elif hasattr(val, "__dataclass_fields__"):
+                result[f.name] = _dc_to_dict(val)
+            elif val is not None and not isinstance(val, Path):
+                result[f.name] = val
+        return result
+
+    d = _dc_to_dict(config)
+
+    # agents is a dict of str -> AgentConfig, handle specially
+    agents_out = {}
+    for agent_id, agent_cfg in config.agents.items():
+        agents_out[agent_id] = _dc_to_dict(agent_cfg)
+    d["agents"] = agents_out
+
+    return d
+
+
+def resolve_effective_budget(agent_id: str, config: TjConfig) -> BudgetConfig:
+    """
+    Return the effective budget for an agent, merging per-agent overrides
+    with global defaults on a per-field basis.
+
+    Each field (daily_usd, session_usd) independently uses the agent value
+    if set, otherwise falls back to the defaults value.
+    """
+    defaults = config.defaults.budget
+    agent_cfg = config.agents.get(agent_id)
+    if agent_cfg is None:
+        return BudgetConfig(
+            daily_usd=defaults.daily_usd,
+            session_usd=defaults.session_usd,
+        )
+    ab = agent_cfg.budget
+    return BudgetConfig(
+        daily_usd=ab.daily_usd if ab.daily_usd is not None else defaults.daily_usd,
+        session_usd=ab.session_usd if ab.session_usd is not None else defaults.session_usd,
+    )
+
+
+def validate_budget_value(value: float, field_name: str) -> float | None:
+    """
+    Validate and normalise a budget value from user input.
+
+    Positive values are returned as-is. Zero means 'remove limit' (returns None).
+    Negative values raise ValueError.
+    """
+    if value < 0:
+        raise ValueError(f"Budget {field_name} must be non-negative, got {value}")
+    return value if value > 0 else None

tokenjam/core/cost.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+import logging
+from tokenjam.core.models import NormalizedSpan
+from tokenjam.core.pricing import get_rates, ModelRates, DEFAULT_INPUT_PER_MTOK, DEFAULT_OUTPUT_PER_MTOK
+
+logger = logging.getLogger(__name__)
+
+
+def calculate_cost(
+    provider: str,
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+    cache_read_tokens: int = 0,
+    cache_write_tokens: int = 0,
+) -> float:
+    """
+    Calculate USD cost for a single LLM call.
+
+    Returns cost rounded to 8 decimal places.
+    Falls back to default rates if the provider/model is not in the pricing table.
+    Logs a warning on fallback so developers know to add the model.
+    Zero tokens -> zero cost (no warning).
+    """
+    if input_tokens == 0 and output_tokens == 0:
+        return 0.0
+
+    rates = get_rates(provider, model)
+    if rates is None:
+        logger.warning(
+            "No pricing data for %s/%s — using default rates. "
+            "Add to pricing/models.toml to get accurate costs.",
+            provider, model,
+        )
+        rates = ModelRates(
+            input_per_mtok=DEFAULT_INPUT_PER_MTOK,
+            output_per_mtok=DEFAULT_OUTPUT_PER_MTOK,
+        )
+
+    cost = (
+        (input_tokens / 1_000_000) * rates.input_per_mtok
+        + (output_tokens / 1_000_000) * rates.output_per_mtok
+        + (cache_read_tokens / 1_000_000) * rates.cache_read_per_mtok
+        + (cache_write_tokens / 1_000_000) * rates.cache_write_per_mtok
+    )
+    return round(cost, 8)
+
+
+class CostEngine:
+    """
+    Post-ingest hook. Called by IngestPipeline after each span is written.
+    Calculates cost and updates span.cost_usd + session.total_cost_usd in DB.
+    """
+
+    def __init__(self, db: object) -> None:
+        self.db = db
+
+    def process_span(self, span: NormalizedSpan) -> None:
+        """
+        If the span has token counts and a provider/model, calculate cost,
+        update span.cost_usd in DB, update session.total_cost_usd in DB.
+        No-op if tokens are missing or zero.
+        """
+        if not span.provider or not span.model:
+            return
+        input_tokens = span.input_tokens or 0
+        output_tokens = span.output_tokens or 0
+        if input_tokens == 0 and output_tokens == 0:
+            return
+
+        cache_read_tokens = span.cache_tokens or 0
+
+        # Record whether the span was already pre-priced before we compute.
+        # Pre-priced spans have their session cost handled by _build_or_update_session
+        # in ingest.py; updating the session again here would double-count.
+        was_pre_priced = span.cost_usd is not None
+
+        cost = calculate_cost(
+            provider=span.provider,
+            model=span.model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+        )
+
+        span.cost_usd = cost
+
+        # Update span cost in DB
+        if hasattr(self.db, 'conn'):
+            self.db.conn.execute(
+                "UPDATE spans SET cost_usd = $1 WHERE span_id = $2",
+                [cost, span.span_id],
+            )
+
+            # Only accumulate into session total when we computed the cost here.
+            # Skip the session update for pre-priced spans to avoid double-counting.
+            if span.session_id and not was_pre_priced:
+                self.db.conn.execute(
+                    "UPDATE sessions SET total_cost_usd = COALESCE(total_cost_usd, 0) + $1 "
+                    "WHERE session_id = $2",
+                    [cost, span.session_id],
+                )