traceguard 0.2.0__py3-none-any.whl

traceguard/__init__.py

@@ -0,0 +1,11 @@
+"""TraceGuard SDK — Phase 0 MVP.
+
+See ../TRACEGUARD_SPEC.md for the binding interface contract.
+"""
+from __future__ import annotations
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "__version__",
+]

traceguard/registry/__init__.py

@@ -0,0 +1,11 @@
+"""Model + prompt registries (SPEC §4.2, §4.3)."""
+from traceguard.registry.models import NoEligibleModelError, register_model, select_model
+from traceguard.registry.prompts import PromptTemplate, load_prompt
+
+__all__ = [
+    "NoEligibleModelError",
+    "PromptTemplate",
+    "load_prompt",
+    "register_model",
+    "select_model",
+]

traceguard/registry/models.py

@@ -0,0 +1,145 @@
+"""Model registry queries (SPEC §4.2)."""
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Literal, overload
+
+from sqlalchemy import or_, select
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+
+from traceguard.store.models import ModelRegistryEntry, make_engine
+
+
+class NoEligibleModelError(LookupError):
+    """No model satisfies the look-ahead constraints (SPEC §4.2 strict mode)."""
+
+
+def register_model(
+    model_id: str,
+    *,
+    model_family: str,
+    capability_class: str,
+    released_at: datetime,
+    available_to_us_at: datetime,
+    deprecated_at: datetime | None = None,
+    engine: Engine | None = None,
+) -> None:
+    """Insert a new model registry entry.
+
+    Per SPEC §3.2, ``model_id`` is the stable primary key and entries are
+    insert-only — re-registering the same id with different metadata is
+    rejected. To "upgrade" a model, register a new ``model_id``.
+    """
+    if released_at > available_to_us_at:
+        raise ValueError(
+            "released_at MUST be <= available_to_us_at "
+            f"(got released_at={released_at!r}, available_to_us_at={available_to_us_at!r})"
+        )
+    eng = engine if engine is not None else make_engine()
+    with Session(eng) as sess:
+        existing = sess.get(ModelRegistryEntry, model_id)
+        if existing is not None:
+            raise ValueError(
+                f"model_id {model_id!r} already registered; "
+                "register a new model_id instead of modifying existing entries"
+            )
+        sess.add(
+            ModelRegistryEntry(
+                model_id=model_id,
+                model_family=model_family,
+                capability_class=capability_class,
+                released_at=released_at,
+                available_to_us_at=available_to_us_at,
+                deprecated_at=deprecated_at,
+            )
+        )
+        sess.commit()
+
+
+@overload
+def select_model(
+    capability_class: str,
+    *,
+    available_at: datetime,
+    strict: Literal[True],
+    engine: Engine | None = None,
+) -> str: ...
+
+
+@overload
+def select_model(
+    capability_class: str,
+    *,
+    available_at: datetime,
+    strict: Literal[False],
+    engine: Engine | None = None,
+) -> tuple[str, bool]: ...
+
+
+def select_model(
+    capability_class: str,
+    *,
+    available_at: datetime,
+    strict: bool,
+    engine: Engine | None = None,
+) -> str | tuple[str, bool]:
+    """Pick a model_id for ``capability_class`` at the ``available_at`` instant.
+
+    ``strict`` is keyword-only and has no default (SPEC §4.2). Callers MUST
+    explicitly state which mode they want — this is the friction that
+    prevents accidental look-ahead.
+
+    strict=True
+        Returns the most recently available model whose
+        ``available_to_us_at <= available_at`` and which is not deprecated as
+        of ``available_at``. Raises ``NoEligibleModelError`` if none.
+
+    strict=False
+        Returns ``(model_id, is_anachronistic)`` for the latest currently
+        active model in the capability class regardless of timing.
+        ``is_anachronistic`` is True iff that model was not yet available at
+        ``available_at``.
+    """
+    eng = engine if engine is not None else make_engine()
+    with Session(eng) as sess:
+        if strict:
+            stmt = (
+                select(ModelRegistryEntry)
+                .where(ModelRegistryEntry.capability_class == capability_class)
+                .where(ModelRegistryEntry.available_to_us_at <= available_at)
+                .where(
+                    or_(
+                        ModelRegistryEntry.deprecated_at.is_(None),
+                        ModelRegistryEntry.deprecated_at > available_at,
+                    )
+                )
+                .order_by(ModelRegistryEntry.available_to_us_at.desc())
+            )
+            entry = sess.scalars(stmt).first()
+            if entry is None:
+                raise NoEligibleModelError(
+                    f"no model registered for capability_class={capability_class!r} "
+                    f"available at {available_at.isoformat()} (strict mode)"
+                )
+            return entry.model_id
+
+        now = datetime.now(timezone.utc)
+        stmt = (
+            select(ModelRegistryEntry)
+            .where(ModelRegistryEntry.capability_class == capability_class)
+            .where(
+                or_(
+                    ModelRegistryEntry.deprecated_at.is_(None),
+                    ModelRegistryEntry.deprecated_at > now,
+                )
+            )
+            .order_by(ModelRegistryEntry.available_to_us_at.desc())
+        )
+        entry = sess.scalars(stmt).first()
+        if entry is None:
+            raise NoEligibleModelError(
+                f"no active model registered for capability_class={capability_class!r}"
+            )
+        is_anachronistic = entry.available_to_us_at > available_at
+        return entry.model_id, is_anachronistic

traceguard/registry/prompts.py

@@ -0,0 +1,148 @@
+"""Prompt template registry — Phase 0 backend is YAML files in a directory.
+
+Per ROADMAP §3.1, prompt_registry in Phase 0 is filesystem-backed (git-tracked)
+rather than DB-backed. The contract is the same set of MUST fields from
+SPEC §3.3; only the backend differs.
+
+Directory layout (relative to prompts_root, default ``prompts/`` at CWD):
+
+    prompts/
+    └── <project>/
+        └── <component>/
+            └── v1.yaml
+            └── v2.yaml
+
+Each YAML file:
+
+    template_body: |
+      ... prompt content with {variables} ...
+    template_format: jinja2 | fstring | raw
+    introduced_at: 2026-05-18T00:00:00+00:00
+    expected_output_schema: null     # optional, JSON-Schema-style dict
+    superseded_at: null
+    superseded_by: null
+    notes: null
+
+``prompt_template_id`` is derived from the path: ``<project>/<component>/v<N>``.
+``prompt_template_hash`` is SHA-256 of ``template_body`` as bytes.
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+_TEMPLATE_ID_RE = re.compile(r"^(?P<project>[a-z0-9_-]+)/(?P<component>[a-z0-9_-]+)/v\d+$")
+
+
+@dataclass(frozen=True)
+class PromptTemplate:
+    prompt_template_id: str
+    prompt_template_hash: str
+    template_body: str
+    template_format: str
+    introduced_at: datetime
+    expected_output_schema: dict[str, Any] | None = None
+    superseded_at: datetime | None = None
+    superseded_by: str | None = None
+    notes: str | None = None
+
+    def render(self, **variables: Any) -> str:
+        """Render the template with ``variables`` injected."""
+        if self.template_format == "raw":
+            if variables:
+                raise ValueError("raw template does not accept variables")
+            return self.template_body
+        if self.template_format == "fstring":
+            return self.template_body.format(**variables)
+        if self.template_format == "jinja2":
+            try:
+                import jinja2
+            except ImportError as exc:  # pragma: no cover - optional dep
+                raise ImportError(
+                    "template_format=jinja2 requires the 'jinja2' package; "
+                    "install with `pip install jinja2`"
+                ) from exc
+            return jinja2.Template(self.template_body).render(**variables)
+        raise ValueError(
+            f"unknown template_format={self.template_format!r}; "
+            "expected one of jinja2 | fstring | raw"
+        )
+
+
+def _prompts_root(explicit: str | os.PathLike[str] | None) -> Path:
+    if explicit is not None:
+        return Path(explicit).expanduser().resolve()
+    env = os.environ.get("TRACEGUARD_PROMPTS_DIR")
+    if env:
+        return Path(env).expanduser().resolve()
+    return Path.cwd() / "prompts"
+
+
+def load_prompt(
+    template_id: str,
+    *,
+    prompts_root: str | os.PathLike[str] | None = None,
+) -> PromptTemplate:
+    """Load a prompt template by its id (``<project>/<component>/v<N>``)."""
+    if not _TEMPLATE_ID_RE.match(template_id):
+        raise ValueError(
+            f"invalid prompt_template_id {template_id!r}; "
+            "expected '<project>/<component>/v<N>' with snake-case names"
+        )
+    root = _prompts_root(prompts_root)
+    path = root / f"{template_id}.yaml"
+    if not path.is_file():
+        raise FileNotFoundError(f"prompt template not found: {path}")
+    with path.open("r", encoding="utf-8") as fp:
+        data = yaml.safe_load(fp) or {}
+
+    body = data.get("template_body")
+    if not isinstance(body, str):
+        raise ValueError(f"prompt template {template_id!r} missing string 'template_body'")
+    fmt = data.get("template_format", "raw")
+    if fmt not in {"jinja2", "fstring", "raw"}:
+        raise ValueError(
+            f"prompt template {template_id!r}: invalid template_format={fmt!r}; "
+            "expected jinja2 | fstring | raw"
+        )
+    introduced_raw = data.get("introduced_at")
+    if introduced_raw is None:
+        raise ValueError(f"prompt template {template_id!r} missing required 'introduced_at'")
+    introduced_at = (
+        introduced_raw if isinstance(introduced_raw, datetime) else datetime.fromisoformat(str(introduced_raw))
+    )
+    if introduced_at.tzinfo is None:
+        raise ValueError(
+            f"prompt template {template_id!r}: introduced_at must be timezone-aware"
+        )
+
+    superseded_raw = data.get("superseded_at")
+    superseded_at: datetime | None
+    if superseded_raw is None:
+        superseded_at = None
+    else:
+        superseded_at = (
+            superseded_raw
+            if isinstance(superseded_raw, datetime)
+            else datetime.fromisoformat(str(superseded_raw))
+        )
+
+    return PromptTemplate(
+        prompt_template_id=template_id,
+        prompt_template_hash=hashlib.sha256(body.encode("utf-8")).hexdigest(),
+        template_body=body,
+        template_format=fmt,
+        introduced_at=introduced_at,
+        expected_output_schema=data.get("expected_output_schema"),
+        superseded_at=superseded_at,
+        superseded_by=data.get("superseded_by"),
+        notes=data.get("notes"),
+    )

traceguard/sdk/__init__.py

@@ -0,0 +1,4 @@
+"""Instrumentation entry points (SPEC §4.1, §4.4)."""
+from traceguard.sdk import normalizer, tracer
+
+__all__ = ["normalizer", "tracer"]

traceguard/sdk/normalizer.py

@@ -0,0 +1,101 @@
+"""Canonical input normalization for SHA-256 hashing (SPEC §4.4).
+
+The normalization function is the single authoritative source for input_hash.
+Business code MUST NOT compute hashes by other means; doing so breaks dedup
+and replay consistency.
+
+Rules (SPEC §4.4):
+- dict: keys sorted, JSON serialized with separators=(",", ":"),
+  ensure_ascii=False
+- str: leading/trailing whitespace stripped, all line endings → "\n"
+- float: serialized via Python's json (shortest repr, stable across versions);
+  NaN / Inf raise (allow_nan=False) — business code must convert first
+- bytes: base64-encoded into a marker dict {"__bytes_b64__": "..."}
+- datetime: ISO 8601 with timezone (naive datetime → raises)
+- Decimal: serialized via str() to preserve precision
+- pydantic BaseModel: dumped via model_dump() then normalized
+- list / tuple: order preserved, elements normalized recursively
+- unknown types: str() fallback with a UserWarning
+
+Changing this algorithm is a SPEC-major bump (SPEC §6.1) because it
+invalidates historical input_hash comparisons.
+"""
+from __future__ import annotations
+
+import base64
+import hashlib
+import json
+import warnings
+from datetime import datetime
+from decimal import Decimal
+from typing import Any
+
+try:
+    from pydantic import BaseModel as _PydanticBaseModel
+except ImportError:  # pragma: no cover - pydantic is a required dep, but keep guard
+    _PydanticBaseModel = None  # type: ignore[assignment, misc]
+
+
+_BYTES_MARKER = "__bytes_b64__"
+
+
+def _normalize(value: Any) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        # json.dumps with allow_nan=False will raise on NaN/Inf below.
+        return value
+    if isinstance(value, str):
+        # Unify line endings then strip outer whitespace.
+        cleaned = value.replace("\r\n", "\n").replace("\r", "\n")
+        return cleaned.strip()
+    if isinstance(value, bytes):
+        return {_BYTES_MARKER: base64.b64encode(value).decode("ascii")}
+    if isinstance(value, datetime):
+        if value.tzinfo is None:
+            raise ValueError(
+                "normalize_input: naive datetime not supported "
+                "(attach a timezone, e.g. datetime.now(timezone.utc))"
+            )
+        return value.isoformat()
+    if isinstance(value, Decimal):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(k): _normalize(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_normalize(v) for v in value]
+    if _PydanticBaseModel is not None and isinstance(value, _PydanticBaseModel):
+        return _normalize(value.model_dump())
+    warnings.warn(
+        f"normalize_input: unsupported type {type(value).__name__!r}, "
+        "falling back to str(); hash stability is not guaranteed across versions",
+        UserWarning,
+        stacklevel=3,
+    )
+    return str(value)
+
+
+def normalize_input(data: Any) -> bytes:
+    """Return canonical UTF-8 bytes for ``data``.
+
+    Identical inputs always yield identical bytes; semantically equivalent
+    inputs (e.g. dicts with different key orders) also collapse to the same
+    bytes.
+    """
+    normalized = _normalize(data)
+    return json.dumps(
+        normalized,
+        sort_keys=True,
+        ensure_ascii=False,
+        separators=(",", ":"),
+        allow_nan=False,
+    ).encode("utf-8")
+
+
+def input_hash(data: Any) -> str:
+    """Return SHA-256 hex digest of ``normalize_input(data)``."""
+    return hashlib.sha256(normalize_input(data)).hexdigest()

traceguard/sdk/tracer.py

@@ -0,0 +1,269 @@
+"""Tracer SDK — decorator + context manager (SPEC §4.1).
+
+Phase 0 ships sync-only instrumentation. The two entry points share a single
+``Span`` object that accumulates state and flushes one row to ``traces`` on
+exit (success or failure).
+"""
+from __future__ import annotations
+
+import json
+import time
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from decimal import Decimal
+from functools import wraps
+from typing import Any, Callable, Iterator
+
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+
+from traceguard.sdk.normalizer import input_hash
+from traceguard.store.models import Trace, make_engine
+
+
+_INPUT_SUMMARY_MAX = 500
+
+
+def _summarize(data: Any) -> str | None:
+    if data is None:
+        return None
+    text = data if isinstance(data, str) else repr(data)
+    return text[:_INPUT_SUMMARY_MAX]
+
+
+def _to_jsonable(value: Any) -> Any:
+    """Best-effort coerce to a JSON-compatible structure for output_parsed."""
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, dict):
+        return {str(k): _to_jsonable(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_to_jsonable(v) for v in value]
+    dump = getattr(value, "model_dump", None)  # pydantic v2
+    if callable(dump):
+        try:
+            return _to_jsonable(dump())
+        except Exception:  # noqa: BLE001 - best effort
+            pass
+    try:
+        json.dumps(value)
+        return value
+    except (TypeError, ValueError):
+        return repr(value)[:_INPUT_SUMMARY_MAX]
+
+
+class Span:
+    """Accumulator for one trace row. Created by ``Tracer.span`` / ``Tracer.trace``."""
+
+    def __init__(
+        self,
+        *,
+        project: str,
+        component: str,
+        operation: str,
+        engine: Engine,
+        correlation_id: str | None = None,
+        feature_as_of: datetime | None = None,
+        parent_trace_id: int | None = None,
+    ) -> None:
+        self.project = project
+        self.component = component
+        self.operation = operation
+        self.correlation_id = correlation_id
+        self.feature_as_of = feature_as_of
+        self.parent_trace_id = parent_trace_id
+
+        self._engine = engine
+        self._start_perf: float = time.perf_counter()
+        self._committed = False
+
+        self._input_hash: str | None = None
+        self._input_summary: str | None = None
+        self._model_id: str | None = None
+        self._prompt_template_id: str | None = None
+        self._prompt_template_hash: str | None = None
+        self._output_parsed: Any = None
+        self._parse_status: str | None = None
+        self._latency_ms: int | None = None
+        self._tokens_in: int | None = None
+        self._tokens_out: int | None = None
+        self._cost_usd: Decimal | None = None
+        self._error_class: str | None = None
+        self._error_message: str | None = None
+
+        self.trace_id: int | None = None
+
+    def record_input(self, data: Any) -> None:
+        self._input_hash = input_hash(data)
+        self._input_summary = _summarize(data)
+
+    def record_model_prompt(
+        self,
+        *,
+        model_id: str | None = None,
+        prompt_template_id: str | None = None,
+        prompt_template_hash: str | None = None,
+    ) -> None:
+        if model_id is not None:
+            self._model_id = model_id
+        if prompt_template_id is not None:
+            self._prompt_template_id = prompt_template_id
+        if prompt_template_hash is not None:
+            self._prompt_template_hash = prompt_template_hash
+
+    def record_output(
+        self,
+        *,
+        parsed: Any = None,
+        parse_status: str = "success",
+    ) -> None:
+        if parse_status not in {"success", "partial", "failed"}:
+            raise ValueError(
+                f"parse_status must be one of success | partial | failed, got {parse_status!r}"
+            )
+        self._output_parsed = _to_jsonable(parsed) if parsed is not None else None
+        self._parse_status = parse_status
+
+    def record_perf(
+        self,
+        *,
+        latency_ms: int | None = None,
+        tokens_in: int | None = None,
+        tokens_out: int | None = None,
+        cost_usd: float | Decimal | None = None,
+    ) -> None:
+        if latency_ms is not None:
+            self._latency_ms = int(latency_ms)
+        if tokens_in is not None:
+            self._tokens_in = int(tokens_in)
+        if tokens_out is not None:
+            self._tokens_out = int(tokens_out)
+        if cost_usd is not None:
+            self._cost_usd = Decimal(str(cost_usd))
+
+    def record_error(self, exc: BaseException) -> None:
+        self._error_class = type(exc).__name__
+        self._error_message = str(exc)
+        if self._parse_status is None:
+            self._parse_status = "failed"
+
+    def _flush(self) -> None:
+        if self._committed:
+            return
+        if self._input_hash is None:
+            self._input_hash = input_hash(None)
+        if self._parse_status is None:
+            self._parse_status = "success" if self._error_class is None else "failed"
+        if self._latency_ms is None:
+            self._latency_ms = int((time.perf_counter() - self._start_perf) * 1000)
+        row = Trace(
+            project=self.project,
+            component=self.component,
+            operation=self.operation,
+            parent_trace_id=self.parent_trace_id,
+            correlation_id=self.correlation_id,
+            input_hash=self._input_hash,
+            input_summary=self._input_summary,
+            model_id=self._model_id,
+            prompt_template_id=self._prompt_template_id,
+            prompt_template_hash=self._prompt_template_hash,
+            output_parsed=self._output_parsed,
+            parse_status=self._parse_status,
+            latency_ms=self._latency_ms,
+            tokens_in=self._tokens_in,
+            tokens_out=self._tokens_out,
+            cost_usd=self._cost_usd,
+            feature_as_of=self.feature_as_of,
+            invoked_at=datetime.now(timezone.utc),
+            error_class=self._error_class,
+            error_message=self._error_message,
+        )
+        with Session(self._engine) as sess:
+            sess.add(row)
+            sess.commit()
+            self.trace_id = row.trace_id
+        self._committed = True
+
+
+class Tracer:
+    """Holds the persistence engine and emits ``Span`` objects."""
+
+    def __init__(self, engine: Engine | None = None) -> None:
+        self._engine = engine
+
+    @property
+    def engine(self) -> Engine:
+        if self._engine is None:
+            self._engine = make_engine()
+        return self._engine
+
+    def configure(self, engine: Engine) -> None:
+        """Override the engine — useful for tests."""
+        self._engine = engine
+
+    @contextmanager
+    def span(
+        self,
+        project: str,
+        component: str,
+        operation: str,
+        *,
+        correlation_id: str | None = None,
+        feature_as_of: datetime | None = None,
+        parent_trace_id: int | None = None,
+    ) -> Iterator[Span]:
+        span = Span(
+            project=project,
+            component=component,
+            operation=operation,
+            engine=self.engine,
+            correlation_id=correlation_id,
+            feature_as_of=feature_as_of,
+            parent_trace_id=parent_trace_id,
+        )
+        try:
+            yield span
+        except BaseException as exc:
+            span.record_error(exc)
+            span._flush()
+            raise
+        else:
+            span._flush()
+
+    def trace(
+        self,
+        project: str,
+        component: str,
+        operation: str,
+        *,
+        correlation_from: Callable[..., str] | None = None,
+        feature_as_of_from: Callable[..., datetime] | None = None,
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Decorator form. Best-effort auto-records (args, kwargs) as input and
+        the return value as ``output_parsed``. For finer control use ``span``.
+        """
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            @wraps(fn)
+            def wrapper(*args: Any, **kwargs: Any) -> Any:
+                corr = correlation_from(*args, **kwargs) if correlation_from else None
+                feat = feature_as_of_from(*args, **kwargs) if feature_as_of_from else None
+                with self.span(
+                    project,
+                    component,
+                    operation,
+                    correlation_id=corr,
+                    feature_as_of=feat,
+                ) as sp:
+                    sp.record_input({"args": list(args), "kwargs": dict(kwargs)})
+                    result = fn(*args, **kwargs)
+                    sp.record_output(parsed=result)
+                    return result
+
+            return wrapper
+
+        return decorator
+
+
+tracer = Tracer()
+"""Module-level default tracer. Configure via ``TRACEGUARD_DB_URL`` env var,
+or call ``tracer.configure(engine)`` to inject a custom engine."""

traceguard/sdk/wrappers/__init__.py

@@ -0,0 +1 @@
+"""Client wrappers — auto-instrumentation for common LLM SDKs."""