Observability Feature 1: Event Model, Observer pub/sub class, and Writer protocol

packages/railtracks/src/railtracks/__init__.py

@@ -41,6 +41,7 @@
     "ToolManifest",
     "session_id",
     "evaluations",
+    "observability",
     "retrieval",
     "Flow",
     "enable_logging",
@@ -57,6 +58,7 @@
     guardrails,
     integrations,
     llm,
+    observability,
     prebuilt,
 )
 from ._session import ExecutionInfo, Session, session

packages/railtracks/src/railtracks/observability/__init__.py

@@ -0,0 +1,28 @@
+"""Observability submodule: streaming Event pipeline with per-writer queues.
+
+Feature 1 of the observability HLD. Independent of the rest of the framework:
+accepts fully-formed `Event` objects; `publish_event`, contextvar reads, and
+emission sites are Feature 2 follow-ups.
+"""
+
+from .models import (
+    SCOPE_EVALUATION,
+    SCOPE_RETRIEVAL,
+    SCOPE_SESSION,
+    Event,
+    Timestamp,
+)
+from .observer import Observer, QueuePolicy
+from .writers import JsonlWriter, Writer
+
+__all__ = [
+    "Event",
+    "Timestamp",
+    "Observer",
+    "QueuePolicy",
+    "Writer",
+    "JsonlWriter",
+    "SCOPE_SESSION",
+    "SCOPE_RETRIEVAL",
+    "SCOPE_EVALUATION",
+]

packages/railtracks/src/railtracks/observability/models.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+SCOPE_SESSION = "session"
+SCOPE_RETRIEVAL = "retrieval"
+SCOPE_EVALUATION = "evaluation"
+
+
+class Timestamp:
+    """Namespace helper for constructing `Event.stamp`. The field itself is a plain tz-aware UTC datetime."""
+
+    # Doing it this way to make potential changes easier
+    @staticmethod
+    def now() -> datetime:
+        return datetime.now(timezone.utc)
+
+
+@dataclass
+class Event:
+    event_type: str
+    scope_type: str
+    scope_id: str
+    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    stamp: datetime = field(default_factory=Timestamp.now)
+    parent_scope_id: str | None = None
+    payload: dict[str, Any] = field(default_factory=dict)

packages/railtracks/src/railtracks/observability/observer.py

@@ -0,0 +1,157 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from enum import Enum
+
+from ..utils.logging.create import get_rt_logger
+from .models import Event
+from .writers.base import Writer
+
+logger = get_rt_logger(__name__)  # for now, we can make another logger later
+
+
+class QueuePolicy(Enum):
+    """How a writer's queue behaves when it's full at publish time."""
+
+    DROP_OLDEST = "drop_oldest"
+
+
+class _EndOfStream:
+    """Marker pushed onto a writer's queue to tell its consumer task to drain and exit."""
+
+
+_END = _EndOfStream()
+
+_QueueItem = Event | _EndOfStream
+
+
+@dataclass
+class _Entry:
+    writer: Writer
+    queue: asyncio.Queue[_QueueItem]
+    task: asyncio.Task[None]
+    policy: QueuePolicy
+
+
+class Observer:
+    def __init__(self) -> None:
+        self._writers: dict[str, _Entry] = {}
+        self._drops: dict[str, int] = {}  # dropped events per writer
+        self._running = False
+
+    # async context manager support added for now, this will become more clear
+    # once we move to integrating with the other modules
+    async def __aenter__(self) -> Observer:
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.shutdown()
+
+    async def start(self) -> None:
+        if self._running:
+            return
+        self._running = True
+
+    async def shutdown(self) -> None:
+        if not self._running:
+            return
+        self._running = False
+        for name in list(self._writers.keys()):
+            await self._teardown(name)
+
+    async def register(
+        self,
+        writer: Writer,
+        name: str,
+        maxsize: int = 10_000,
+        policy: QueuePolicy = QueuePolicy.DROP_OLDEST,
+    ) -> None:
+        if not self._running:
+            raise RuntimeError(
+                "Observer is not running; call start() or use as an async context manager."
+            )
+        if name in self._writers:
+            raise ValueError(f"Writer {name!r} is already registered.")
+        await writer.start()
+        queue: asyncio.Queue[_QueueItem] = asyncio.Queue(
+            maxsize=maxsize
+        )  # Each writer has its own queue
+        task = asyncio.create_task(
+            self._consumer_loop(name, writer, queue),
+            name=f"observer-consumer:{name}",
+        )
+        self._writers[name] = _Entry(
+            writer=writer, queue=queue, task=task, policy=policy
+        )
+        self._drops[name] = 0
+
+    async def unregister(self, name: str) -> None:
+        if name not in self._writers:
+            raise KeyError(f"No writer registered as {name!r}.")
+        await self._teardown(name)
+
+    async def publish(self, event: Event) -> None:
+        if not self._running:
+            raise RuntimeError("Observer is not running.")
+        for name, entry in self._writers.items():
+            try:
+                entry.queue.put_nowait(event)
+            except asyncio.QueueFull:
+                self._handle_full_queue(name, entry, event)
+
+    def _handle_full_queue(self, name: str, entry: _Entry, event: Event) -> None:
+        match entry.policy:
+            case QueuePolicy.DROP_OLDEST:
+                self._drop_oldest(name, entry, event)
+
+    def _drop_oldest(self, name: str, entry: _Entry, event: Event) -> None:
+        try:
+            entry.queue.get_nowait()
+        except asyncio.QueueEmpty:
+            pass
+        entry.queue.put_nowait(event)
+        self._drops[name] += 1
+        logger.warning(
+            "observability writer %r queue full; dropped oldest event "
+            "(policy=%s, total drops for this writer: %d)",
+            name,
+            entry.policy.value,
+            self._drops[name],
+        )
+
+    async def _teardown(self, name: str) -> None:
+        entry = self._writers.pop(name)
+        self._drops.pop(name, None)
+        _enqueue_end(entry.queue)
+        await entry.task
+        await entry.writer.shutdown()
+
+    async def _consumer_loop(
+        self, name: str, writer: Writer, queue: asyncio.Queue[_QueueItem]
+    ) -> None:
+        while True:
+            item = await queue.get()
+            if isinstance(item, _EndOfStream):
+                return
+            try:
+                await writer.write(item)
+            except Exception as exc:
+                logger.warning(
+                    "observability writer %r failed on event %s: %s",
+                    name,
+                    item.event_id,
+                    exc,
+                )
+
+
+def _enqueue_end(queue: asyncio.Queue[_QueueItem]) -> None:
+    try:
+        queue.put_nowait(_END)
+    except asyncio.QueueFull:
+        try:
+            queue.get_nowait()
+        except asyncio.QueueEmpty:
+            pass
+        queue.put_nowait(_END)

packages/railtracks/src/railtracks/observability/writers/__init__.py

@@ -0,0 +1,4 @@
+from .base import Writer
+from .jsonl import JsonlWriter
+
+__all__ = ["Writer", "JsonlWriter"]

packages/railtracks/src/railtracks/observability/writers/base.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from ..models import Event
+
+
+class Writer(Protocol):
+    async def start(self) -> None: ...
+    async def write(self, event: Event) -> None: ...
+    async def shutdown(self) -> None: ...

packages/railtracks/src/railtracks/observability/writers/jsonl.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import dataclasses
+import json
+from pathlib import Path
+from typing import TextIO
+
+from ..models import Event
+
+
+class JsonlWriter:
+    def __init__(self, directory: Path):
+        self._directory = directory
+        self._files: dict[str, TextIO] = {}
+
+    async def start(self) -> None:
+        self._directory.mkdir(parents=True, exist_ok=True)
+
+    async def write(self, event: Event) -> None:
+        handle = self._files.get(event.scope_type)
+        if handle is None:
+            handle = (self._directory / f"{event.scope_type}.jsonl").open(
+                "a", encoding="utf-8"
+            )
+            self._files[event.scope_type] = handle
+        handle.write(_serialize(event) + "\n")
+        handle.flush()
+
+    async def shutdown(self) -> None:
+        for handle in self._files.values():
+            handle.flush()
+            handle.close()
+        self._files.clear()
+
+
+def _serialize(event: Event) -> str:
+    data = dataclasses.asdict(event)
+    data["stamp"] = event.stamp.isoformat()
+    return json.dumps(data)