tokenjam 0.2.0__py3-none-any.whl

tokenjam/core/models.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any
+
+# Sessions with no spans for this long are considered stale (zombie).
+SESSION_STALE_THRESHOLD = timedelta(minutes=5)
+
+
+class Severity(str, Enum):
+    CRITICAL = "critical"
+    WARNING  = "warning"
+    INFO     = "info"
+
+
+class AlertType(str, Enum):
+    COST_BUDGET_DAILY        = "cost_budget_daily"
+    COST_BUDGET_SESSION      = "cost_budget_session"
+    SENSITIVE_ACTION         = "sensitive_action"
+    RETRY_LOOP               = "retry_loop"
+    TOKEN_ANOMALY            = "token_anomaly"
+    SESSION_DURATION         = "session_duration"
+    SCHEMA_VIOLATION         = "schema_violation"
+    DRIFT_DETECTED           = "drift_detected"
+    FAILURE_RATE             = "failure_rate"
+    NETWORK_EGRESS_BLOCKED   = "network_egress_blocked"
+    FILESYSTEM_ACCESS_DENIED = "filesystem_access_denied"
+    SYSCALL_DENIED           = "syscall_denied"
+    INFERENCE_REROUTED       = "inference_rerouted"
+
+
+class SpanStatus(str, Enum):
+    OK    = "ok"
+    ERROR = "error"
+    UNSET = "unset"
+
+
+class SpanKind(str, Enum):
+    INTERNAL  = "internal"
+    CLIENT    = "client"
+    SERVER    = "server"
+    PRODUCER  = "producer"
+    CONSUMER  = "consumer"
+
+
+@dataclass
+class NormalizedSpan:
+    span_id:        str
+    trace_id:       str
+    name:           str
+    kind:           SpanKind
+    status_code:    SpanStatus
+    start_time:     datetime
+    parent_span_id: str | None     = None
+    session_id:     str | None     = None
+    agent_id:       str | None     = None
+    end_time:       datetime | None = None
+    duration_ms:    float | None   = None
+    status_message: str | None     = None
+    attributes:     dict[str, Any] = field(default_factory=dict)
+    events:         list[dict]     = field(default_factory=list)
+    # Extracted indexed fields
+    provider:       str | None     = None
+    model:          str | None     = None
+    tool_name:      str | None     = None
+    input_tokens:   int | None     = None
+    output_tokens:  int | None     = None
+    cache_tokens:   int | None     = None
+    cost_usd:       float | None   = None
+    request_type:   str | None     = None
+    conversation_id: str | None    = None
+
+
+@dataclass
+class SessionRecord:
+    session_id:      str
+    agent_id:        str
+    started_at:      datetime
+    conversation_id: str | None   = None
+    ended_at:        datetime | None = None
+    status:          str          = "active"
+    total_cost_usd:  float | None = None
+    input_tokens:    int          = 0
+    output_tokens:   int          = 0
+    cache_tokens:    int          = 0
+    tool_call_count: int          = 0
+    error_count:     int          = 0
+
+    @property
+    def duration_seconds(self) -> float | None:
+        if self.started_at and self.ended_at:
+            return (self.ended_at - self.started_at).total_seconds()
+        return None
+
+    @property
+    def effective_status(self) -> str:
+        """Return 'stale' for zombie sessions whose process was killed."""
+        if self.status != "active":
+            return self.status
+        from tokenjam.utils.time_parse import utcnow
+        last_activity = self.ended_at or self.started_at
+        if last_activity and (utcnow() - last_activity) > SESSION_STALE_THRESHOLD:
+            return "stale"
+        return "active"
+
+
+@dataclass
+class AgentRecord:
+    agent_id:   str
+    first_seen: datetime
+    last_seen:  datetime
+    name:       str | None    = None
+    version:    str | None    = None
+    provider:   str | None    = None
+
+
+@dataclass
+class Alert:
+    alert_id:   str
+    fired_at:   datetime
+    type:       AlertType
+    severity:   Severity
+    title:      str
+    detail:     dict[str, Any]
+    agent_id:   str | None = None
+    session_id: str | None = None
+    span_id:    str | None = None
+    acknowledged: bool     = False
+    suppressed:   bool     = False
+
+
+@dataclass
+class DriftBaseline:
+    agent_id:               str
+    sessions_sampled:       int
+    computed_at:            datetime
+    avg_input_tokens:       float | None = None
+    stddev_input_tokens:    float | None = None
+    avg_output_tokens:      float | None = None
+    stddev_output_tokens:   float | None = None
+    avg_session_duration_s: float | None = None
+    stddev_session_duration: float | None = None
+    avg_tool_call_count:    float | None = None
+    stddev_tool_call_count: float | None = None
+    common_tool_sequences:  list | None  = None
+    output_schema_inferred: dict | None  = None
+
+
+@dataclass
+class DriftViolation:
+    dimension: str
+    z_score:   float | None   = None
+    expected:  str | None     = None
+    observed:  str | None     = None
+    detail:    str | None     = None
+
+
+@dataclass
+class DriftResult:
+    violations: list[DriftViolation]
+    drifted:    bool
+
+
+@dataclass
+class SchemaValidationResult:
+    validation_id: str
+    span_id:       str
+    validated_at:  datetime
+    passed:        bool
+    errors:        list[str]  = field(default_factory=list)
+    agent_id:      str | None = None
+
+
+@dataclass
+class TraceRecord:
+    trace_id:   str
+    agent_id:   str | None
+    name:       str
+    start_time: datetime
+    duration_ms: float | None = None
+    cost_usd:   float | None  = None
+    status_code: str          = "ok"
+    span_count:  int          = 0
+
+
+@dataclass
+class CostRow:
+    group:        str
+    agent_id:     str | None  = None
+    model:        str | None  = None
+    input_tokens: int         = 0
+    output_tokens: int        = 0
+    cost_usd:     float       = 0.0
+
+
+# -- Filter dataclasses used by StorageBackend --
+
+@dataclass
+class TraceFilters:
+    agent_id:   str | None   = None
+    since:      datetime | None = None
+    until:      datetime | None = None
+    span_name:  str | None   = None
+    status:     str | None   = None
+    limit:      int          = 50
+    offset:     int          = 0
+
+
+@dataclass
+class CostFilters:
+    agent_id:  str | None   = None
+    since:     datetime | None = None
+    until:     datetime | None = None
+    group_by:  str          = "day"   # agent | model | day | tool
+
+
+@dataclass
+class AlertFilters:
+    agent_id:  str | None   = None
+    since:     datetime | None = None
+    severity:  Severity | None = None
+    type:      AlertType | None = None
+    unread:    bool          = False
+    limit:     int           = 100

tokenjam/core/pricing.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+import sys
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib  # type: ignore[no-redef]
+
+
+PRICING_FILE = Path(__file__).parent.parent / "pricing" / "models.toml"
+
+# Default rate used when a model is not in the pricing table.
+# 0.50 per MTok input, 2.00 per MTok output — conservative mid-range estimate.
+DEFAULT_INPUT_PER_MTOK = 0.50
+DEFAULT_OUTPUT_PER_MTOK = 2.00
+
+
+@dataclass(frozen=True)
+class ModelRates:
+    input_per_mtok: float
+    output_per_mtok: float
+    cache_read_per_mtok: float = 0.0
+    cache_write_per_mtok: float = 0.0
+
+
+@lru_cache(maxsize=1)
+def load_pricing_table() -> dict[str, dict[str, ModelRates]]:
+    """
+    Load pricing/models.toml and return a nested dict:
+      { provider: { model_name: ModelRates } }
+    Cached after first load — restart process to pick up changes.
+    """
+    with open(PRICING_FILE, "rb") as f:
+        raw = tomllib.load(f)
+    result: dict[str, dict[str, ModelRates]] = {}
+    for provider, models in raw.items():
+        result[provider] = {}
+        for model_name, rates in models.items():
+            result[provider][model_name] = ModelRates(
+                input_per_mtok=rates.get("input_per_mtok", DEFAULT_INPUT_PER_MTOK),
+                output_per_mtok=rates.get("output_per_mtok", DEFAULT_OUTPUT_PER_MTOK),
+                cache_read_per_mtok=rates.get("cache_read_per_mtok", 0.0),
+                cache_write_per_mtok=rates.get("cache_write_per_mtok", 0.0),
+            )
+    return result
+
+
+def get_rates(provider: str, model: str) -> ModelRates | None:
+    """Return ModelRates for the given provider/model, or None if not found."""
+    table = load_pricing_table()
+    return table.get(provider, {}).get(model)

tokenjam/core/retention.py

@@ -0,0 +1,21 @@
+"""Storage retention cleanup. Deletes spans older than config.retention_days."""
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from tokenjam.utils.time_parse import utcnow
+
+if TYPE_CHECKING:
+    from tokenjam.core.config import StorageConfig
+    from tokenjam.core.db import StorageBackend
+
+
+def run_retention_cleanup(db: StorageBackend, config: StorageConfig) -> int:
+    """
+    Delete spans older than config.retention_days.
+    Returns the number of spans deleted.
+    Called by the apscheduler background job in tj serve.
+    """
+    cutoff = utcnow() - timedelta(days=config.retention_days)
+    return db.delete_spans_before(cutoff)

tokenjam/core/schema_validator.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import jsonschema
+from genson import SchemaBuilder
+
+from tokenjam.core.models import SchemaValidationResult, AlertType, Severity
+from tokenjam.otel.semconv import GenAIAttributes
+from tokenjam.utils.ids import new_uuid
+from tokenjam.utils.time_parse import utcnow
+
+if TYPE_CHECKING:
+    from tokenjam.core.db import StorageBackend
+    from tokenjam.core.models import NormalizedSpan
+    from tokenjam.core.alerts import AlertEngine
+    from tokenjam.core.config import TjConfig
+
+logger = logging.getLogger("tokenjam.schema")
+
+
+class SchemaValidator:
+    """
+    Post-ingest hook. Called by IngestPipeline after each span.
+
+    Only runs when ALL of these are true:
+    1. span.name == "gen_ai.tool.call"
+    2. span has gen_ai.tool.output in its attributes (capture.tool_outputs = true)
+    3. The agent has an output_schema configured, OR an inferred schema exists in baseline
+
+    If output_schema is a file path in config, load the JSON Schema from that file.
+    If no schema is declared, use inferred schema from drift baseline.
+    If neither exists yet, silently skip.
+    """
+
+    def __init__(self, db: StorageBackend, alert_engine: AlertEngine, config: TjConfig):
+        self.db = db
+        self.alert_engine = alert_engine
+        self.config = config
+        self._schema_cache: dict[str, dict] = {}
+
+    def validate(self, span: NormalizedSpan) -> None:
+        """
+        Validate tool output against schema if applicable.
+        Persists a SchemaValidationResult to DB regardless of pass/fail.
+        Fires SCHEMA_VIOLATION alert on failure.
+        """
+        if span.name != GenAIAttributes.SPAN_TOOL_CALL:
+            return
+
+        tool_output = span.attributes.get(GenAIAttributes.TOOL_OUTPUT)
+        if tool_output is None:
+            return
+
+        schema = self._get_schema(span.agent_id)
+        if schema is None:
+            return
+
+        # Parse tool_output if it's a string
+        parsed_output = tool_output
+        if isinstance(parsed_output, str):
+            try:
+                parsed_output = json.loads(parsed_output)
+            except (json.JSONDecodeError, TypeError):
+                pass  # Validate the raw string against the schema
+
+        errors = list(jsonschema.Draft7Validator(schema).iter_errors(parsed_output))
+
+        result = SchemaValidationResult(
+            validation_id=new_uuid(),
+            span_id=span.span_id,
+            agent_id=span.agent_id,
+            validated_at=utcnow(),
+            passed=len(errors) == 0,
+            errors=[e.message for e in errors],
+        )
+        self.db.insert_validation(result)
+
+        if not result.passed:
+            self.alert_engine.fire(
+                alert_type=AlertType.SCHEMA_VIOLATION,
+                span_or_session=span,
+                detail={"errors": result.errors, "tool_name": span.tool_name},
+                severity=Severity.WARNING,
+            )
+
+    def _get_schema(self, agent_id: str | None) -> dict | None:
+        """
+        Return the JSON Schema for this agent, or None if unavailable.
+        Priority: 1) declared schema file in config, 2) inferred schema in baseline.
+        Caches loaded schemas in-memory.
+        """
+        if agent_id is None:
+            return None
+
+        if agent_id in self._schema_cache:
+            return self._schema_cache[agent_id]
+
+        # 1. Check declared schema in agent config
+        agent_cfg = self.config.agents.get(agent_id)
+        if agent_cfg and agent_cfg.output_schema:
+            schema = self._load_schema_file(agent_cfg.output_schema)
+            if schema is not None:
+                self._schema_cache[agent_id] = schema
+                return schema
+
+        # 2. Fall back to inferred schema from drift baseline
+        baseline = self.db.get_baseline(agent_id)
+        if baseline and baseline.output_schema_inferred:
+            self._schema_cache[agent_id] = baseline.output_schema_inferred
+            return baseline.output_schema_inferred
+
+        return None
+
+    def _load_schema_file(self, path_str: str) -> dict | None:
+        """Load a JSON Schema from a file path.
+
+        Relative paths are resolved relative to the config file's parent directory
+        (config.config_path). Falls back to CWD-relative resolution when no config
+        file path is available (e.g. in tests with a synthetic TjConfig).
+        """
+        path = Path(path_str)
+        if not path.is_absolute():
+            config_file_path = getattr(self.config, "config_path", None)
+            if config_file_path is not None:
+                path = Path(config_file_path).parent / path_str
+        if not path.exists():
+            logger.warning("Schema file not found: %s", path)
+            return None
+        try:
+            with open(path, "r") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, OSError) as exc:
+            logger.warning("Failed to load schema file %s: %s", path, exc)
+            return None
+
+    def infer_schema_from_outputs(self, tool_outputs: list[Any]) -> dict | None:
+        """
+        Use genson to infer a JSON Schema from a list of tool outputs.
+        Returns None if the list is empty.
+        """
+        if not tool_outputs:
+            return None
+
+        builder = SchemaBuilder()
+        for output in tool_outputs:
+            if isinstance(output, str):
+                try:
+                    output = json.loads(output)
+                except (json.JSONDecodeError, TypeError):
+                    pass
+            builder.add_object(output)
+        return builder.to_schema()

tokenjam/demo/env.py

@@ -0,0 +1,96 @@
+"""
+Shared demo environment: wires IngestPipeline + InMemoryBackend directly,
+bypassing the SDK/OTel TracerProvider so scenarios have zero setup friction.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from tokenjam.core.config import AgentConfig
+    from tokenjam.core.models import NormalizedSpan
+
+
+@dataclass
+class DemoResult:
+    agent_id: str
+    span_count: int
+    alert_count: int
+    alert_types: list[str]
+    total_cost_usd: float
+    trace_count: int
+
+
+class DemoEnvironment:
+    """
+    Self-contained observability stack for demo scenarios.
+    No API keys, no OTel global state, no config files needed.
+    Alerts are silenced (no stdout/file channels) so scenarios run cleanly.
+    """
+
+    def __init__(
+        self,
+        agent_configs: "dict[str, AgentConfig] | None" = None,
+    ) -> None:
+        from tokenjam.core.alerts import AlertEngine
+        from tokenjam.core.config import AlertsConfig, TjConfig, SecurityConfig
+        from tokenjam.core.cost import CostEngine
+        from tokenjam.core.db import InMemoryBackend
+        from tokenjam.core.drift import DriftDetector
+        from tokenjam.core.ingest import IngestPipeline
+
+        self.db = InMemoryBackend()
+
+        self.config = TjConfig(
+            version="1",
+            security=SecurityConfig(ingest_secret="demo"),
+            alerts=AlertsConfig(channels=[], cooldown_seconds=0),
+            agents=agent_configs or {},
+        )
+
+        cost_engine = CostEngine(db=self.db)
+        alert_engine = AlertEngine(db=self.db, config=self.config)
+        drift_detector = DriftDetector(
+            db=self.db, alert_engine=alert_engine, config=self.config
+        )
+
+        self.pipeline = IngestPipeline(
+            db=self.db,
+            config=self.config,
+            cost_engine=cost_engine,
+            alert_engine=alert_engine,
+            drift_detector=drift_detector,
+        )
+
+    def process(self, span: "NormalizedSpan") -> None:
+        self.pipeline.process(span)
+
+    def get_alerts(self) -> list:
+        from tokenjam.core.models import AlertFilters
+        return self.db.get_alerts(AlertFilters(limit=1000))
+
+    def total_cost_usd(self) -> float:
+        row = self.db.conn.execute(
+            "SELECT COALESCE(SUM(cost_usd), 0.0) FROM spans"
+        ).fetchone()
+        return float(row[0]) if row else 0.0
+
+    def span_count(self) -> int:
+        row = self.db.conn.execute("SELECT COUNT(*) FROM spans").fetchone()
+        return int(row[0]) if row else 0
+
+    def trace_count(self) -> int:
+        from tokenjam.core.models import TraceFilters
+        return len(self.db.get_traces(TraceFilters()))
+
+    def build_result(self, agent_id: str) -> DemoResult:
+        alerts = self.get_alerts()
+        return DemoResult(
+            agent_id=agent_id,
+            span_count=self.span_count(),
+            alert_count=len(alerts),
+            alert_types=[a.type.value for a in alerts],
+            total_cost_usd=self.total_cost_usd(),
+            trace_count=self.trace_count(),
+        )