traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/sources/sqlite.py

@@ -0,0 +1,187 @@
+"""SQLite polling source for reading agent session databases."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import sqlite3
+from collections.abc import AsyncIterator
+from datetime import datetime, timezone
+from pathlib import Path
+from types import TracebackType
+from typing import Literal
+
+from traceforge.sources.base import RawRecord, Source
+
+logger = logging.getLogger(__name__)
+
+
+class SqliteSource(Source):
+    """Poll a SQLite database for new rows in a target table.
+
+    Tracks position via a monotonically increasing column (e.g. rowid, id,
+    or timestamp). Each new row becomes a RawRecord whose payload is the
+    JSON-serialized row dict.
+
+    Designed for:
+    - Copilot CLI session-store.db (turns, forge_trajectory_events)
+    - Any agent that stores conversation data in SQLite
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        name: str,
+        table: str = "turns",
+        order_column: str = "id",
+        columns: list[str] | None = None,
+        where: str | None = None,
+        interval: float = 2.0,
+        start_at: Literal["beginning", "end"] = "end",
+        session_filter: str | None = None,
+    ) -> None:
+        """Initialize SqliteSource.
+
+        Args:
+            path: Path to the SQLite database file.
+            name: Source name for record attribution.
+            table: Table to poll for new rows.
+            order_column: Column used for ordering and cursor tracking (must be monotonic).
+            columns: Specific columns to select (None = all).
+            where: Additional WHERE clause (without 'WHERE' keyword).
+            interval: Seconds between polls.
+            start_at: "beginning" reads all existing rows; "end" starts from latest.
+            session_filter: If set, adds WHERE session_id = ? filter.
+        """
+        if interval < 0:
+            raise ValueError("interval must be non-negative")
+        self.path = Path(path).resolve()
+        self.name = name
+        self.table = table
+        self.order_column = order_column
+        self.columns = columns
+        self.where = where
+        self.interval = interval
+        self.start_at = start_at
+        self.session_filter = session_filter
+        self._cursor: int | str | None = None
+        self._sequence = 0
+        self._conn: sqlite3.Connection | None = None
+        self._iterating = False
+
+    async def __aenter__(self) -> "SqliteSource":
+        if not self.path.exists():
+            raise FileNotFoundError(f"SqliteSource database not found: {self.path}")
+        self._conn = await asyncio.to_thread(self._connect)
+        if self.start_at == "end":
+            self._cursor = await asyncio.to_thread(self._get_max_cursor)
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+        self._iterating = False
+
+    async def _iter_records(self) -> AsyncIterator[RawRecord]:
+        if self._conn is None:
+            raise RuntimeError("SqliteSource must be entered before iteration")
+        if self._iterating:
+            raise RuntimeError("SqliteSource does not support concurrent iteration")
+        self._iterating = True
+        try:
+            while True:
+                rows = await asyncio.to_thread(self._poll_new_rows)
+                for row in rows:
+                    yield self._make_record(row)
+                await asyncio.sleep(self.interval)
+        finally:
+            self._iterating = False
+
+    def __aiter__(self) -> AsyncIterator[RawRecord]:
+        return self._iter_records()
+
+    def _connect(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(str(self.path), timeout=5, check_same_thread=False)
+        conn.row_factory = sqlite3.Row
+        # WAL mode for concurrent reads without blocking writers
+        try:
+            conn.execute("PRAGMA journal_mode=WAL")
+        except sqlite3.OperationalError:
+            pass
+        return conn
+
+    def _get_max_cursor(self) -> int | str | None:
+        """Get the current maximum value of the order column."""
+        assert self._conn is not None
+        query = f"SELECT MAX({self.order_column}) FROM {self.table}"  # noqa: S608
+        filters, params = self._build_filters()
+        if filters:
+            query += f" WHERE {' AND '.join(filters)}"
+        row = self._conn.execute(query, params).fetchone()
+        return row[0] if row else None
+
+    def _poll_new_rows(self) -> list[dict[str, object]]:
+        """Query for rows newer than the current cursor."""
+        assert self._conn is not None
+
+        col_spec = ", ".join(self.columns) if self.columns else "*"
+        query = f"SELECT {col_spec} FROM {self.table}"  # noqa: S608
+
+        filters, params = self._build_filters()
+        if self._cursor is not None:
+            filters.append(f"{self.order_column} > ?")
+            params.append(self._cursor)
+
+        if filters:
+            query += f" WHERE {' AND '.join(filters)}"
+        query += f" ORDER BY {self.order_column} ASC"
+
+        try:
+            rows = self._conn.execute(query, params).fetchall()
+        except sqlite3.OperationalError as exc:
+            logger.warning("SqliteSource %s: query failed: %s", self.name, exc)
+            return []
+
+        results: list[dict[str, object]] = []
+        for row in rows:
+            row_dict = dict(row)
+            cursor_val = row_dict.get(self.order_column)
+            if cursor_val is not None:
+                self._cursor = cursor_val
+            # Serialize to JSON for the RawRecord payload
+            results.append({k: v for k, v in row_dict.items() if v is not None})
+
+        if results:
+            logger.debug("SqliteSource %s: polled %d new rows", self.name, len(results))
+
+        return results
+
+    def _build_filters(self) -> tuple[list[str], list[object]]:
+        """Build WHERE clause filters and parameters."""
+        filters: list[str] = []
+        params: list[object] = []
+        if self.session_filter:
+            filters.append("session_id = ?")
+            params.append(self.session_filter)
+        if self.where:
+            filters.append(f"({self.where})")
+        return filters, params
+
+    def _make_record(self, row_dict: dict[str, object]) -> RawRecord:
+
+        record = RawRecord(
+            payload=json.dumps(row_dict, default=str),
+            source_name=self.name,
+            mode="sqlite",
+            sequence=self._sequence,
+            received_at=datetime.now(timezone.utc),
+        )
+        self._sequence += 1
+        return record

traceforge/sources/sse.py

@@ -0,0 +1,198 @@
+"""Server-sent events source transport (WHATWG spec compliant)."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from types import TracebackType
+
+import httpx
+
+from traceforge.sources.base import RawRecord, Source
+
+
+@dataclass(slots=True)
+class SSEEvent:
+    """A parsed SSE event per the WHATWG spec."""
+
+    data: str
+    event_type: str = "message"
+    last_event_id: str = ""
+
+
+class SSESource(Source):
+    """Read SSE events from an endpoint with spec-compliant parsing and reconnect.
+
+    Implements the WHATWG Server-Sent Events specification:
+    - Parses data, event, id, retry fields
+    - Sends Last-Event-ID on reconnect
+    - Validates Content-Type: text/event-stream
+    - Exponential backoff with server-controlled retry
+    """
+
+    def __init__(
+        self,
+        url: str,
+        name: str,
+        headers: dict[str, str] | None = None,
+        reconnect_delay: float = 1.0,
+        max_reconnects: int | None = None,
+    ) -> None:
+        self.url = url
+        self.name = name
+        self.headers = dict(headers or {})
+        self.reconnect_delay = reconnect_delay
+        self.max_reconnects = max_reconnects
+        self._client: httpx.AsyncClient | None = None
+        self._sequence = 0
+        self._last_event_id: str = ""
+        self._iterating = False
+
+    async def __aenter__(self) -> "SSESource":
+        self._client = httpx.AsyncClient(timeout=None)
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+            self._client = None
+        self._iterating = False
+
+    async def _iter_records(self) -> AsyncIterator[RawRecord]:
+        if self._client is None:
+            raise RuntimeError("SSESource must be entered before iteration")
+        if self._iterating:
+            raise RuntimeError("SSESource does not support concurrent iteration")
+        self._iterating = True
+        try:
+            reconnects = 0
+            while True:
+                if self.max_reconnects is not None and reconnects > self.max_reconnects:
+                    raise RuntimeError(f"Exceeded maximum reconnects for SSE source {self.name}")
+                try:
+                    async for sse_event in self._stream_once():
+                        reconnects = 0
+                        yield self._make_record(sse_event)
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    reconnects += 1
+                    if self.max_reconnects is not None and reconnects > self.max_reconnects:
+                        raise
+                    await asyncio.sleep(self._backoff_delay(reconnects))
+                    continue
+
+                reconnects += 1
+                if self.max_reconnects is not None and reconnects > self.max_reconnects:
+                    return
+                await asyncio.sleep(self._backoff_delay(reconnects))
+        finally:
+            self._iterating = False
+
+    def __aiter__(self) -> AsyncIterator[RawRecord]:
+        return self._iter_records()
+
+    async def _stream_once(self) -> AsyncIterator[SSEEvent]:
+        """Connect and yield parsed SSE events until the stream ends."""
+        if self._client is None:
+            raise RuntimeError("SSESource must be entered before iteration")
+
+        request_headers = {"Accept": "text/event-stream", **self.headers}
+        if self._last_event_id:
+            request_headers["Last-Event-ID"] = self._last_event_id
+
+        async with self._client.stream("GET", self.url, headers=request_headers) as response:
+            response.raise_for_status()
+            content_type = response.headers.get("content-type", "")
+            if "text/event-stream" not in content_type:
+                raise ValueError(
+                    f"SSE endpoint returned Content-Type '{content_type}', "
+                    f"expected 'text/event-stream'"
+                )
+
+            data_buf: list[str] = []
+            event_type = ""
+            last_id: str | None = None
+
+            async for line in response.aiter_lines():
+                if line == "":
+                    # Dispatch event
+                    if data_buf:
+                        data = "\n".join(data_buf)
+                        # SSE is at-least-once: after a reconnect the server may
+                        # redeliver an already-seen id (e.g. resumed via Last-Event-ID).
+                        # Drop such redeliveries; only emit ids newer than the last one
+                        # we emitted, advancing _last_event_id solely for emitted events.
+                        if last_id is None or not self._is_duplicate(last_id):
+                            if last_id is not None:
+                                self._last_event_id = last_id
+                            yield SSEEvent(
+                                data=data,
+                                event_type=event_type or "message",
+                                last_event_id=self._last_event_id,
+                            )
+                    # Reset per-event state
+                    data_buf = []
+                    event_type = ""
+                    last_id = None
+                    continue
+
+                if line.startswith(":"):
+                    continue
+
+                field, _, value = line.partition(":")
+                if value.startswith(" "):
+                    value = value[1:]
+
+                if field == "data":
+                    data_buf.append(value)
+                elif field == "event":
+                    event_type = value
+                elif field == "id":
+                    # Per spec: empty id resets; id with NUL is ignored
+                    if "\x00" not in value:
+                        last_id = value
+                elif field == "retry":
+                    if value.isdigit():
+                        self.reconnect_delay = max(int(value) / 1000.0, 0.0)
+
+    def _is_duplicate(self, event_id: str) -> bool:
+        """Return True if ``event_id`` was already emitted (at-least-once redelivery).
+
+        After a reconnect an SSE server may redeliver events the client has already
+        seen. When both the last-emitted id and the candidate are integers, treat any
+        id not strictly greater than the last-seen id as a duplicate (monotonic ids,
+        the common case). For ids that cannot be ordered numerically, drop only an id
+        that exactly equals the last-emitted one (the immediately-preceding event).
+        """
+        last = self._last_event_id
+        if not last:
+            return False
+        if event_id == last:
+            return True
+        try:
+            return int(event_id) <= int(last)
+        except ValueError:
+            return False
+
+    def _backoff_delay(self, reconnects: int) -> float:
+        exponent = max(reconnects - 1, 0)
+        return self.reconnect_delay * min(2**exponent, 30)
+
+    def _make_record(self, sse_event: SSEEvent) -> RawRecord:
+        record = RawRecord(
+            payload=sse_event.data,
+            source_name=self.name,
+            mode="stream",
+            sequence=self._sequence,
+            received_at=datetime.now(timezone.utc),
+        )
+        self._sequence += 1
+        return record

traceforge/telemetry/__init__.py

@@ -0,0 +1,192 @@
+"""Opt-in, near-zero-footprint self-metrics for :class:`~traceforge.pipeline.EventPipeline`.
+
+This package is traceforge's *self*-observability seam (SPEC §14). It is distinct
+from OTLP *export* (``sinks/otel_exporter.OtelExporterSink``, which ships events /
+spans / usage to an external collector): this measures the pipeline's **own**
+operation — throughput, enrichment latency, per-sink write time, and dropped /
+failed-sink counts.
+
+Design contract (why this stays honest to the "must not add per-event cost"
+constraint of #48):
+
+- **Opt-in.** A pipeline created without a :class:`PipelineMetrics` instance does
+  no timing and makes no metrics allocations on the hot path — the pipeline
+  guards every instrumentation site on ``metrics is not None`` and never calls a
+  clock when disabled. Attaching an instance is the *only* way to turn it on.
+- **In-process, no dependencies.** No ``opentelemetry-sdk``, no ``prometheus``,
+  no background threads, no parallel transport. It is a plain accumulator read
+  back off the same object you passed in (surfaced on ``flush()`` / ``close()``).
+- **Bounded.** Every field is a scalar counter or timing sum, plus exactly one
+  entry per sink. Nothing grows per event, so a long-lived daemon never
+  accumulates unbounded state.
+
+Usage::
+
+    from traceforge import EventPipeline
+    from traceforge.telemetry import PipelineMetrics
+
+    metrics = PipelineMetrics()
+    pipeline = EventPipeline(sinks=[...], metrics=metrics)
+    ...
+    await pipeline.close()
+    snap = metrics.snapshot()
+    print(snap.events_per_second, snap.mean_enrichment_ms)
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+
+__all__ = ["MetricsSnapshot", "PipelineMetrics", "SinkMetrics"]
+
+
+@dataclass(frozen=True)
+class SinkMetrics:
+    """Immutable per-sink write totals within a :class:`MetricsSnapshot`.
+
+    ``label`` identifies the sink by class name and position in the pipeline's
+    sink list (e.g. ``"JsonlSink#1"``), so two sinks of the same class stay
+    distinguishable.
+    """
+
+    label: str
+    writes: int
+    failures: int
+    write_seconds: float
+
+    @property
+    def mean_write_ms(self) -> float:
+        """Mean wall time this sink spent per write, in milliseconds."""
+        return (self.write_seconds / self.writes * 1000.0) if self.writes else 0.0
+
+
+@dataclass(frozen=True)
+class MetricsSnapshot:
+    """An immutable point-in-time read of a :class:`PipelineMetrics`.
+
+    Returned by :meth:`PipelineMetrics.snapshot`. Raw counters are captured
+    verbatim; rates and means are exposed as derived properties so no divide-by-
+    zero can leak to callers.
+    """
+
+    events: int
+    dropped_events: int
+    sink_failures: int
+    enrichment_calls: int
+    enrichment_seconds: float
+    active_seconds: float
+    sinks: tuple[SinkMetrics, ...]
+
+    @property
+    def events_per_second(self) -> float:
+        """Throughput over the span from the first to the most recent event.
+
+        ``0.0`` until at least two events have been recorded far enough apart for
+        the monotonic clock to advance (a single event has no measurable span).
+        """
+        return (self.events / self.active_seconds) if self.active_seconds > 0 else 0.0
+
+    @property
+    def mean_enrichment_ms(self) -> float:
+        """Mean wall time spent in the enricher per event, in milliseconds."""
+        return (
+            (self.enrichment_seconds / self.enrichment_calls * 1000.0)
+            if self.enrichment_calls
+            else 0.0
+        )
+
+
+class PipelineMetrics:
+    """Mutable, in-process accumulator of :class:`EventPipeline` self-metrics.
+
+    Attach one instance per pipeline via ``EventPipeline(..., metrics=...)`` to
+    turn metrics on; leave it unset (the default) and the pipeline stays a true
+    no-op on the hot path. The pipeline calls the ``record_*`` methods; consumers
+    call :meth:`snapshot` (any time — the counters update live) to read a stable,
+    immutable summary.
+
+    Not thread-safe by design: the pipeline is asyncio single-threaded, so every
+    ``record_*`` call runs on the one event loop; there is nothing to lock.
+    """
+
+    __slots__ = (
+        "_dropped_events",
+        "_enrichment_calls",
+        "_enrichment_seconds",
+        "_events",
+        "_first_perf",
+        "_last_perf",
+        "_sink_failures",
+        "_sink_seconds",
+        "_sink_writes",
+    )
+
+    def __init__(self) -> None:
+        self._events = 0
+        self._dropped_events = 0
+        self._enrichment_calls = 0
+        self._enrichment_seconds = 0.0
+        self._first_perf: float | None = None
+        self._last_perf = 0.0
+        # Per-sink accumulators, keyed by the pipeline-assigned sink label. Bounded
+        # by the number of sinks; never grows per event.
+        self._sink_writes: dict[str, int] = {}
+        self._sink_failures: dict[str, int] = {}
+        self._sink_seconds: dict[str, float] = {}
+
+    def record_event(self) -> None:
+        """Count one event reaching the sink fan-out and stamp its arrival time."""
+        now = time.perf_counter()
+        if self._first_perf is None:
+            self._first_perf = now
+        self._last_perf = now
+        self._events += 1
+
+    def record_drop(self) -> None:
+        """Count one event dropped by the enricher (``process`` returned ``None``)."""
+        self._dropped_events += 1
+
+    def record_enrichment(self, seconds: float) -> None:
+        """Record one enrichment call and the wall time it took."""
+        self._enrichment_calls += 1
+        self._enrichment_seconds += seconds
+
+    def record_sink_write(self, label: str, seconds: float, *, failed: bool) -> None:
+        """Record one sink write: its wall time and whether it raised.
+
+        A failed write still counts as a write (the sink was invoked); ``failed``
+        additionally bumps that sink's failure tally.
+        """
+        self._sink_writes[label] = self._sink_writes.get(label, 0) + 1
+        self._sink_seconds[label] = self._sink_seconds.get(label, 0.0) + seconds
+        if failed:
+            self._sink_failures[label] = self._sink_failures.get(label, 0) + 1
+
+    @property
+    def active_seconds(self) -> float:
+        """Span from the first to the most recent recorded event, in seconds."""
+        if self._first_perf is None:
+            return 0.0
+        return max(0.0, self._last_perf - self._first_perf)
+
+    def snapshot(self) -> MetricsSnapshot:
+        """Return an immutable summary of the metrics recorded so far."""
+        sinks = tuple(
+            SinkMetrics(
+                label=label,
+                writes=writes,
+                failures=self._sink_failures.get(label, 0),
+                write_seconds=self._sink_seconds.get(label, 0.0),
+            )
+            for label, writes in self._sink_writes.items()
+        )
+        return MetricsSnapshot(
+            events=self._events,
+            dropped_events=self._dropped_events,
+            sink_failures=sum(self._sink_failures.values()),
+            enrichment_calls=self._enrichment_calls,
+            enrichment_seconds=self._enrichment_seconds,
+            active_seconds=self.active_seconds,
+            sinks=sinks,
+        )

traceforge/title/__init__.py

@@ -0,0 +1,23 @@
+"""Torch-free, CPU-only title generation for activity/step spans.
+
+See :mod:`traceforge.title.inference` for the served titler. Title hygiene
+(decode cleanup, parent/child de-dup) lives in :mod:`traceforge.title.hygiene`.
+"""
+
+from __future__ import annotations
+
+from .context import distilled_context
+from .hygiene import best_of, clean_title, norm_key, pick_distinct
+from .inference import TitleModel
+from .inferencer import SessionTitleStream, TitleInferencer
+
+__all__ = [
+    "SessionTitleStream",
+    "TitleInferencer",
+    "TitleModel",
+    "best_of",
+    "clean_title",
+    "distilled_context",
+    "norm_key",
+    "pick_distinct",
+]

traceforge/title/_resolve.py

@@ -0,0 +1,79 @@
+"""Locate the packaged titler weights (the activity/step *span* head).
+
+The ONNX weights ship in the separate :mod:`traceforge_title_model` package, which
+is a hard dependency of traceforge, so a normal install always has them. This
+module resolves the span head's directory, in order:
+
+1. the installed :mod:`traceforge_title_model` package (the shipped path);
+2. an in-tree dev fallback (``src/traceforge/title/data/``) so a source checkout
+   with the ONNX dropped in place still serves without the installed package;
+3. ``None`` -> the caller raises :data:`INSTALL_HINT`.
+
+Session naming no longer uses a packaged head (the distilled request head was
+dropped as proven-weak); it is served by :mod:`traceforge.title.naming` instead.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+#: The three files :meth:`traceforge.title.TitleModel.load` reads for one head.
+TRIAD = ("encoder.onnx", "decoder.onnx", "tokenizer.json")
+
+_LFS_MAGIC = b"version https://git-lfs.github.com/spec/v1"
+#: Real ONNX heads are tens of MB; a Git LFS pointer is ~133 bytes. This floor
+#: cleanly separates a smudged binary from a pointer (or a truncated artifact).
+_MIN_ONNX_BYTES = 1_000_000
+
+_HERE = Path(__file__).resolve().parent
+#: In-tree dev fallback. Empty in a normal install (weights live in the data
+#: package); populated only if a developer drops the ONNX here by hand.
+_DEV_SPAN = _HERE / "data"
+
+INSTALL_HINT = (
+    "traceforge titler weights are not installed. They ship with the "
+    "'traceforge-title-model' package (a dependency of traceforge), so reinstalling "
+    "should restore them:\n"
+    "    pip install --force-reinstall traceforge-title-model\n"
+    "or, if PyPI is unavailable, pull the GitHub-release mirror:\n"
+    "    traceforge download-model --source gh"
+)
+
+
+def _usable(f: Path) -> bool:
+    """A weight file that is present *and* a real binary (not an LFS pointer).
+
+    Guards the pointer-in-wheel failure mode: if a package was built from a
+    checkout that never smudged LFS, the files exist but are ~133-byte pointer
+    stubs. Existence alone would then resolve as "installed" and fail only later,
+    deep inside onnxruntime. Reject pointers (and implausibly small ONNX) here so
+    the caller falls back / surfaces :data:`INSTALL_HINT` instead.
+    """
+    try:
+        if not f.is_file():
+            return False
+        if f.suffix == ".onnx" and f.stat().st_size < _MIN_ONNX_BYTES:
+            return False
+        with f.open("rb") as fh:
+            return not fh.read(len(_LFS_MAGIC)).startswith(_LFS_MAGIC)
+    except OSError:
+        return False
+
+
+def _complete(d: Path | None) -> bool:
+    return d is not None and all(_usable(d / f) for f in TRIAD)
+
+
+def _pkg_dir() -> Path | None:
+    """The span head dir from the installed data package, if complete."""
+    try:
+        import traceforge_title_model as m
+    except ImportError:
+        return None
+    d = m.span_dir()
+    return d if _complete(d) else None
+
+
+def span_dir() -> Path | None:
+    """Resolved span (activity/step) head dir, or ``None`` if unavailable."""
+    return _pkg_dir() or (_DEV_SPAN if _complete(_DEV_SPAN) else None)