traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/sinks/otel_exporter.py

@@ -0,0 +1,212 @@
+"""OTel exporter sink — emit governance results as OpenTelemetry spans."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+
+class OtelExporterSink(StorageSink):
+    """Exports enriched events as OTel spans via OTLP/HTTP JSON.
+
+    Uses a simplified OTLP JSON payload (not protobuf) to avoid heavy
+    dependencies. Compatible with any OTLP/HTTP collector.
+    """
+
+    def __init__(
+        self,
+        endpoint: str = "http://localhost:4318/v1/traces",
+        service_name: str = "traceforge",
+        headers: dict[str, str] | None = None,
+        max_backlog: int = 1024,
+    ) -> None:
+        if not endpoint.startswith(("http://", "https://")):
+            raise ValueError(f"OTel endpoint must use http:// or https:// scheme, got: {endpoint}")
+        self._endpoint = endpoint
+        self._service_name = service_name
+        self._headers = headers or {}
+        self._batch: list[dict] = []
+        self._batch_size = 32
+        self._max_backlog = max_backlog
+
+    async def on_event(self, event: SessionEvent) -> None:
+        await self._enqueue(self._event_to_span(event))
+
+    async def _enqueue(self, span: dict) -> None:
+        self._batch.append(span)
+
+        if len(self._batch) >= self._max_backlog:
+            dropped = len(self._batch) - self._batch_size
+            self._batch = self._batch[-self._batch_size :]
+            logger.warning(
+                "OtelExporterSink: backlog exceeded %d, dropped %d oldest spans",
+                self._max_backlog,
+                dropped,
+            )
+
+        if len(self._batch) >= self._batch_size:
+            await self.flush()
+
+    async def flush(self) -> None:
+        if not self._batch:
+            return
+
+        payload = {
+            "resourceSpans": [
+                {
+                    "resource": {
+                        "attributes": [
+                            {"key": "service.name", "value": {"stringValue": self._service_name}},
+                        ]
+                    },
+                    "scopeSpans": [
+                        {
+                            "scope": {"name": "traceforge.governance"},
+                            "spans": list(self._batch),
+                        }
+                    ],
+                }
+            ]
+        }
+
+        body = json.dumps(payload, default=str).encode("utf-8")
+        headers = {
+            "Content-Type": "application/json",
+            **self._headers,
+        }
+
+        try:
+            req = Request(self._endpoint, data=body, headers=headers, method="POST")
+            status = await asyncio.to_thread(self._do_request, req)
+            if status >= 300:
+                logger.warning("OtelExporterSink: OTLP endpoint returned %d", status)
+                return  # keep batch for retry on next flush
+            self._batch.clear()  # only clear after successful send
+        except (URLError, OSError, TimeoutError) as exc:
+            logger.error("OtelExporterSink: failed to export %d spans: %s", len(self._batch), exc)
+
+    def _do_request(self, req: Request) -> int:
+        """Synchronous HTTP request — returns status code. Ensures response body is consumed."""
+        with urlopen(req, timeout=10) as resp:
+            resp.read()
+            return resp.status
+
+    async def close(self) -> None:
+        await self.flush()
+
+    def _event_to_span(self, event: SessionEvent) -> dict:
+        """Convert a SessionEvent to an OTLP span dict."""
+        import uuid
+
+        ts_ns = int(event.timestamp.timestamp() * 1_000_000_000) if event.timestamp else 0
+        span_id = uuid.uuid4().hex[:16]
+        trace_id = uuid.uuid4().hex
+
+        attributes = [
+            {"key": "traceforge.event.kind", "value": {"stringValue": event.kind}},
+            {"key": "traceforge.session.id", "value": {"stringValue": event.session_id}},
+        ]
+
+        if event.payload:
+            tool_name = event.payload.get("tool_name")
+            if tool_name:
+                attributes.append(
+                    {"key": "gen_ai.tool.name", "value": {"stringValue": str(tool_name)}}
+                )
+            tool_args = event.payload.get("arguments") or event.payload.get("tool_input")
+            if tool_args:
+                attributes.append(
+                    {
+                        "key": "gen_ai.tool.call.arguments",
+                        "value": {
+                            "stringValue": str(tool_args)
+                            if not isinstance(tool_args, str)
+                            else tool_args
+                        },
+                    }
+                )
+            tool_result = event.payload.get("result") or event.payload.get("tool_result")
+            if tool_result:
+                attributes.append(
+                    {"key": "gen_ai.tool.call.result", "value": {"stringValue": str(tool_result)}}
+                )
+            tool_call_id = event.payload.get("tool_call_id")
+            if tool_call_id:
+                attributes.append(
+                    {"key": "gen_ai.tool.call.id", "value": {"stringValue": str(tool_call_id)}}
+                )
+
+        gov = event.metadata.governance if event.metadata else None
+        if gov is not None:
+            if gov.risk_assessment is not None:
+                attributes.append(
+                    {
+                        "key": "traceforge.risk.score",
+                        "value": {"intValue": gov.risk_assessment.score},
+                    }
+                )
+                attributes.append(
+                    {
+                        "key": "traceforge.risk.level",
+                        "value": {"stringValue": gov.risk_assessment.level},
+                    }
+                )
+            if gov.recommendation is not None:
+                attributes.append(
+                    {
+                        "key": "traceforge.action",
+                        "value": {"stringValue": gov.recommendation.recommended_action.value},
+                    }
+                )
+
+        return {
+            "traceId": trace_id,
+            "spanId": span_id,
+            "name": f"traceforge.{event.kind}",
+            "kind": 1,  # SPAN_KIND_INTERNAL
+            "startTimeUnixNano": str(ts_ns),
+            "endTimeUnixNano": str(ts_ns),
+            "attributes": attributes,
+        }
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        pass
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        pass
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        await self._enqueue(self._title_update_to_span(update))
+
+    def _title_update_to_span(self, update: TitleUpdate) -> dict:
+        """Convert a TitleUpdate to an OTLP span dict."""
+        import uuid
+
+        attributes = [
+            {"key": "traceforge.session.id", "value": {"stringValue": update.session_id}},
+            {"key": "traceforge.segment.id", "value": {"stringValue": update.segment_id}},
+            {"key": "traceforge.segment.kind", "value": {"stringValue": update.kind}},
+            {"key": "traceforge.segment.title", "value": {"stringValue": update.title}},
+            {"key": "traceforge.segment.title_version", "value": {"intValue": update.version}},
+        ]
+        if update.parent_id is not None:
+            attributes.append(
+                {"key": "traceforge.segment.parent_id", "value": {"stringValue": update.parent_id}}
+            )
+        return {
+            "traceId": uuid.uuid4().hex,
+            "spanId": uuid.uuid4().hex[:16],
+            "name": f"traceforge.title.{update.kind}",
+            "kind": 1,  # SPAN_KIND_INTERNAL
+            "startTimeUnixNano": "0",
+            "endTimeUnixNano": "0",
+            "attributes": attributes,
+        }

traceforge/sinks/parquet.py

@@ -0,0 +1,260 @@
+"""Parquet output sink — per-session columnar files for analytics consumers.
+
+Buffers SessionEvents in memory by session_id; flushes a parquet file when
+SESSION_ENDED arrives, when the buffer exceeds ``max_buffered_events``, or
+on ``close()``. One parquet file per session.
+
+``pyarrow`` is a required core dependency — the canonical analytics format
+is parquet, so we don't ship a parquet sink that pretends pyarrow is
+optional.
+
+See ``research/docs/06-pipeline-architecture.md`` for the design rationale,
+schema, and consumer story.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import re
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from traceforge.phase.event_rows import event_to_feature_row as _row_from_event
+from traceforge.sinks.base import StorageSink
+from traceforge.types import EventKind, SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MAX_BUFFERED_EVENTS = 5_000
+_DEFAULT_COMPRESSION = "zstd"
+_DEFAULT_ROW_GROUP_SIZE = 10_000
+_SAFE_SESSION_RE = re.compile(r"[^a-zA-Z0-9_\-]")
+_SESSION_END_KINDS = frozenset({EventKind.SESSION_ENDED, EventKind.SESSION_PAUSED})
+
+
+def _build_title_schema() -> pa.Schema:
+    """Schema for the per-session title sidecar (``<session>.titles.parquet``).
+
+    Titles are append-only :class:`~traceforge.types.TitleUpdate` records keyed to
+    a segment by ``segment_id``; consumers join them onto the events that carry
+    the matching ``activity_id``/``step_id``. Kept as a sidecar so the stable
+    event schema is never widened for late-arriving titles.
+    """
+    return pa.schema(
+        [
+            pa.field("session_id", pa.string()),
+            pa.field("segment_id", pa.string()),
+            pa.field("kind", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("title", pa.string()),
+            pa.field("version", pa.int64()),
+            pa.field("parent_id", pa.string()),
+        ]
+    )
+
+
+def _build_schema() -> pa.Schema:
+    """Stable column schema mirroring the canonical ``Classification`` and
+    ``EventMetadata`` shapes.
+
+    Multi-valued classification dimensions (``scope``, ``role``, ``action``,
+    ``capability``, ``structure``, ``source_labels``) are ``list<string>``
+    because the underlying type is ``frozenset[str]`` — not a scalar.
+
+    New fields land in ``payload_json`` / ``metadata_json`` until promoted to
+    typed columns.
+    """
+    return pa.schema(
+        [
+            pa.field("event_id", pa.string()),
+            pa.field("session_id", pa.string()),
+            pa.field("kind", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("timestamp_ns", pa.timestamp("ns", tz="UTC")),
+            pa.field("seq", pa.int64()),
+            pa.field("tool_name", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("mechanism", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("effect", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("scope", pa.list_(pa.string())),
+            pa.field("role", pa.list_(pa.string())),
+            pa.field("action", pa.list_(pa.string())),
+            pa.field("capability", pa.list_(pa.string())),
+            pa.field("structure", pa.list_(pa.string())),
+            pa.field("source_labels", pa.list_(pa.string())),
+            pa.field("shell_dialect", pa.dictionary(pa.int32(), pa.string())),
+            pa.field("binaries", pa.list_(pa.string())),
+            pa.field("phase_signals", pa.list_(pa.string())),
+            pa.field("activity", pa.string()),
+            pa.field("motivation", pa.string()),
+            pa.field("payload_json", pa.string()),
+            pa.field("metadata_json", pa.string()),
+            pa.field("duration_ms", pa.int64()),
+        ]
+    )
+
+
+class ParquetSink(StorageSink):
+    """Per-session parquet sink.
+
+    Buffers events in memory keyed by ``session_id``; emits one parquet file
+    per session. Flushes when:
+
+    - the session ends (``EventKind.SESSION_ENDED``)
+    - the per-session buffer exceeds ``max_buffered_events``
+    - ``flush()`` or ``close()`` is called
+
+    Output path supports ``{session_id}`` as a template variable. If the
+    template has no ``{session_id}``, ``"<session_id>.parquet"`` is appended.
+
+    The schema is intentionally stable — see ``_build_schema`` and
+    ``research/docs/06-pipeline-architecture.md``. New fields land in
+    ``payload_json`` / ``metadata_json`` until promoted.
+
+    Resumed sessions: if a session emits more events after a flush
+    (``SESSION_RESUMED`` then more events), this sink writes to
+    ``<session_id>.<n>.parquet`` for n=1, 2, ... so existing files are not
+    overwritten.
+    """
+
+    def __init__(
+        self,
+        path: str | Path,
+        max_buffered_events: int = _DEFAULT_MAX_BUFFERED_EVENTS,
+        compression: str = _DEFAULT_COMPRESSION,
+        row_group_size: int = _DEFAULT_ROW_GROUP_SIZE,
+    ) -> None:
+        self._path_template = str(path)
+        self._max_buffered_events = max_buffered_events
+        self._compression = compression
+        self._row_group_size = row_group_size
+
+        self._schema = _build_schema()
+        self._title_schema = _build_title_schema()
+
+        self._buffers: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        self._title_buffers: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        self._seq_counters: dict[str, int] = defaultdict(int)
+        self._flush_index: dict[str, int] = defaultdict(int)
+        self._lock = asyncio.Lock()
+
+    def _resolve_path(self, session_id: str, flush_idx: int) -> Path:
+        sanitized = _SAFE_SESSION_RE.sub("_", session_id)[:128]
+        suffix = f".{flush_idx}" if flush_idx > 0 else ""
+
+        if "{session_id}" in self._path_template:
+            resolved = self._path_template.replace("{session_id}", sanitized)
+            base = Path(resolved)
+            if suffix:
+                base = base.with_suffix(f"{suffix}{base.suffix}")
+        else:
+            base = Path(self._path_template) / f"{sanitized}{suffix}.parquet"
+
+        path = base.expanduser().resolve()
+
+        # Containment check: resolved path must stay under the template's
+        # parent directory.
+        if "{session_id}" in self._path_template:
+            base_dir = Path(self._path_template.split("{session_id}")[0]).expanduser().resolve()
+        else:
+            base_dir = Path(self._path_template).expanduser().resolve()
+        if not str(path).startswith(str(base_dir)):
+            raise ValueError(f"ParquetSink: resolved path escapes base directory: {path}")
+        return path
+
+    async def on_event(self, event: SessionEvent) -> None:
+        async with self._lock:
+            sid = event.session_id
+            seq = self._seq_counters[sid]
+            self._seq_counters[sid] = seq + 1
+
+            row = _row_from_event(event, seq)
+            self._buffers[sid].append(row)
+
+            should_flush = (
+                event.kind in _SESSION_END_KINDS
+                or len(self._buffers[sid]) >= self._max_buffered_events
+            )
+
+        if should_flush:
+            await self._flush_session(sid)
+
+    async def _flush_session(self, session_id: str) -> None:
+        """Write the buffered rows for one session and clear them."""
+        async with self._lock:
+            rows = self._buffers.pop(session_id, [])
+            titles = self._title_buffers.pop(session_id, [])
+            if not rows and not titles:
+                return
+            flush_idx = self._flush_index[session_id]
+            self._flush_index[session_id] = flush_idx + 1
+
+        if rows:
+            await asyncio.to_thread(self._write_rows, session_id, flush_idx, rows)
+        if titles:
+            await asyncio.to_thread(self._write_titles, session_id, flush_idx, titles)
+
+    def _write_rows(self, session_id: str, flush_idx: int, rows: list[dict[str, Any]]) -> None:
+        path = self._resolve_path(session_id, flush_idx)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Conversion errors are programming bugs (schema vs. row shape) and
+        # must surface — silently dropping a session's data is worse than a
+        # loud failure. OSError is the only "expected" runtime failure
+        # (disk full, permission denied) and gets a structured log.
+        table = pa.Table.from_pylist(rows, schema=self._schema)
+        try:
+            pq.write_table(
+                table,
+                path,
+                compression=self._compression,
+                row_group_size=self._row_group_size,
+            )
+        except OSError as exc:
+            logger.error("ParquetSink: failed to write %s: %s", path, exc)
+            raise
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        # Spans are not part of the per-event schema; sidecar parquet for
+        # spans is a future addition.
+        pass
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        pass
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        async with self._lock:
+            self._title_buffers[update.session_id].append(
+                {
+                    "session_id": update.session_id,
+                    "segment_id": update.segment_id,
+                    "kind": update.kind,
+                    "title": update.title,
+                    "version": update.version,
+                    "parent_id": update.parent_id,
+                }
+            )
+
+    def _write_titles(self, session_id: str, flush_idx: int, rows: list[dict[str, Any]]) -> None:
+        evt_path = self._resolve_path(session_id, flush_idx)
+        path = evt_path.with_name(f"{evt_path.stem}.titles{evt_path.suffix}")
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        table = pa.Table.from_pylist(rows, schema=self._title_schema)
+        try:
+            pq.write_table(table, path, compression=self._compression)
+        except OSError as exc:
+            logger.error("ParquetSink: failed to write titles %s: %s", path, exc)
+            raise
+
+    async def flush(self) -> None:
+        """Flush every buffered session (events and/or titles)."""
+        async with self._lock:
+            sids = set(self._buffers) | set(self._title_buffers)
+        for sid in sids:
+            await self._flush_session(sid)
+
+    async def close(self) -> None:
+        await self.flush()

traceforge/sinks/s3.py

@@ -0,0 +1,206 @@
+"""S3 sink — buffer events and flush as JSONL objects to S3-compatible storage."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import re
+import time
+import uuid
+from datetime import datetime, timezone
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_BUFFER_SIZE = 100
+_DEFAULT_FLUSH_INTERVAL_SECONDS = 60.0
+_SAFE_SESSION_RE = re.compile(r"[^a-zA-Z0-9_\-]")
+
+
+def _require_boto3():
+    """Import and return boto3, raising a helpful error if missing."""
+    try:
+        import boto3
+
+        return boto3
+    except ImportError:
+        raise ImportError(
+            "boto3 is required for S3Sink. Install it with: pip install traceforge-toolkit[s3]"
+        ) from None
+
+
+class S3Sink(StorageSink):
+    """Buffers events in memory and flushes to S3 as JSONL objects.
+
+    Object key format: {prefix}{session_id}/{date}/{timestamp}-{uuid_short}.jsonl
+
+    Requires boto3 (optional dependency). Install via ``pip install traceforge-toolkit[s3]``.
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: str = "",
+        region: str | None = None,
+        endpoint_url: str | None = None,
+        buffer_size: int = _DEFAULT_BUFFER_SIZE,
+        flush_interval: float = _DEFAULT_FLUSH_INTERVAL_SECONDS,
+    ) -> None:
+        self._bucket = bucket
+        self._prefix = prefix
+        self._region = region
+        self._endpoint_url = endpoint_url
+        self._buffer_size = buffer_size
+        self._flush_interval = flush_interval
+
+        self._buffer: list[dict] = []
+        self._last_flush_time: float = time.monotonic()
+        self._session_id: str | None = None
+        self._lock = asyncio.Lock()
+
+        # Validate boto3 is available at construction time
+        self._boto3 = _require_boto3()
+        self._client = None
+
+    def _get_client(self):
+        """Lazily create the S3 client."""
+        if self._client is None:
+            kwargs: dict = {}
+            if self._region:
+                kwargs["region_name"] = self._region
+            if self._endpoint_url:
+                kwargs["endpoint_url"] = self._endpoint_url
+            self._client = self._boto3.client("s3", **kwargs)
+        return self._client
+
+    def _make_object_key(self, session_id: str) -> str:
+        """Generate the S3 object key for the current flush.
+
+        Sanitizes session_id to prevent unexpected characters in S3 keys.
+        """
+        sanitized = _SAFE_SESSION_RE.sub("_", session_id)[:128]
+        sanitized = sanitized or "unknown"
+        now = datetime.now(timezone.utc)
+        date_str = now.strftime("%Y-%m-%d")
+        timestamp_str = now.strftime("%Y%m%dT%H%M%S")
+        short_uuid = uuid.uuid4().hex[:8]
+        return f"{self._prefix}{sanitized}/{date_str}/{timestamp_str}-{short_uuid}.jsonl"
+
+    def _serialize_event(self, event: SessionEvent) -> dict:
+        return {
+            "id": event.id,
+            "kind": event.kind,
+            "session_id": event.session_id,
+            "timestamp": event.timestamp.isoformat(),
+            "payload": event.payload,
+            "metadata": event.metadata.model_dump(exclude_none=True) if event.metadata else None,
+        }
+
+    def _serialize_span(self, span: TelemetrySpan) -> dict:
+        return {
+            "type": "span",
+            "name": span.name,
+            "session_id": span.session_id,
+            "start_time": span.start_time.isoformat(),
+            "end_time": span.end_time.isoformat(),
+            "attributes": span.attributes,
+        }
+
+    def _serialize_usage(self, usage: UsageRecord) -> dict:
+        return {
+            "type": "usage",
+            "session_id": usage.session_id,
+            "timestamp": usage.timestamp.isoformat(),
+            "model": usage.model,
+            "input_tokens": usage.input_tokens,
+            "output_tokens": usage.output_tokens,
+            "cost_usd": usage.cost_usd,
+        }
+
+    def _serialize_title_update(self, update: TitleUpdate) -> dict:
+        return {
+            "type": "title_update",
+            "session_id": update.session_id,
+            "segment_id": update.segment_id,
+            "kind": update.kind,
+            "title": update.title,
+            "version": update.version,
+            "parent_id": update.parent_id,
+        }
+
+    async def on_event(self, event: SessionEvent) -> None:
+        async with self._lock:
+            if self._session_id is None:
+                self._session_id = event.session_id
+            self._buffer.append(self._serialize_event(event))
+            if self._should_flush():
+                await self._flush_buffer()
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        async with self._lock:
+            if self._session_id is None:
+                self._session_id = span.session_id
+            self._buffer.append(self._serialize_span(span))
+            if self._should_flush():
+                await self._flush_buffer()
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        async with self._lock:
+            if self._session_id is None:
+                self._session_id = usage.session_id
+            self._buffer.append(self._serialize_usage(usage))
+            if self._should_flush():
+                await self._flush_buffer()
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        async with self._lock:
+            if self._session_id is None:
+                self._session_id = update.session_id
+            self._buffer.append(self._serialize_title_update(update))
+            if self._should_flush():
+                await self._flush_buffer()
+
+    def _should_flush(self) -> bool:
+        if len(self._buffer) >= self._buffer_size:
+            return True
+        elapsed = time.monotonic() - self._last_flush_time
+        if elapsed >= self._flush_interval:
+            return True
+        return False
+
+    async def _flush_buffer(self) -> None:
+        """Upload buffered items to S3 and reset the buffer."""
+        if not self._buffer:
+            return
+
+        session_id = self._session_id or "unknown"
+        key = self._make_object_key(session_id)
+        body = "\n".join(json.dumps(item, default=str) for item in self._buffer) + "\n"
+
+        try:
+            client = self._get_client()
+            await asyncio.to_thread(
+                client.put_object,
+                Bucket=self._bucket,
+                Key=key,
+                Body=body.encode("utf-8"),
+                ContentType="application/x-ndjson",
+            )
+            logger.debug(
+                "S3Sink: flushed %d items to s3://%s/%s", len(self._buffer), self._bucket, key
+            )
+        except Exception as exc:
+            logger.error("S3Sink: failed to upload to s3://%s/%s: %s", self._bucket, key, exc)
+
+        self._buffer.clear()
+        self._last_flush_time = time.monotonic()
+
+    async def flush(self) -> None:
+        async with self._lock:
+            await self._flush_buffer()
+
+    async def close(self) -> None:
+        await self.flush()