ai-pipeline-core 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/observability/_tracking/_writer.py

@@ -0,0 +1,229 @@
+"""Background writer thread for ClickHouse inserts and summary jobs."""
+
+import asyncio
+import contextlib
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass, field
+from threading import Event, Thread
+
+from lmnr.opentelemetry_lib.tracing import context as laminar_context
+from opentelemetry import context as otel_context
+from opentelemetry.context import Context
+from pydantic import BaseModel
+
+from ai_pipeline_core.logging import get_pipeline_logger
+
+from ._client import ClickHouseClient
+from ._models import SummaryRowBuilder
+
+type SpanSummaryFn = Callable[[str, str, str], Coroutine[None, None, str]]
+
+logger = get_pipeline_logger(__name__)
+
+
+@dataclass(frozen=True)
+class InsertBatch:
+    """Batch of rows to insert into a single table."""
+
+    table: str
+    rows: list[BaseModel] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class SummaryJob:
+    """Job requesting LLM-generated summary for a span."""
+
+    span_id: str
+    label: str
+    output_hint: str
+    summary_model: str = "gemini-3-flash"
+    parent_otel_context: Context | None = field(default=None, hash=False, compare=False)
+    parent_laminar_context: Context | None = field(default=None, hash=False, compare=False)
+
+
+_SENTINEL = object()
+
+
+class ClickHouseWriter:
+    """Background writer that batches inserts and processes summary jobs.
+
+    Uses a dedicated thread with its own asyncio event loop. External callers
+    push work via ``write()`` and ``write_job()`` which use
+    ``loop.call_soon_threadsafe()``. The writer drains the queue in batches
+    for efficiency.
+    """
+
+    def __init__(
+        self,
+        client: ClickHouseClient,
+        *,
+        summary_row_builder: SummaryRowBuilder | None = None,
+        span_summary_fn: SpanSummaryFn | None = None,
+        batch_size: int = 100,
+        flush_interval_seconds: float = 2.0,
+    ) -> None:
+        """Store config. Does NOT start the writer thread."""
+        self._client = client
+        self._summary_row_builder = summary_row_builder
+        self._span_summary_fn = span_summary_fn
+        self._batch_size = batch_size
+        self._flush_interval = flush_interval_seconds
+
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._queue: asyncio.Queue[InsertBatch | SummaryJob | object] | None = None
+        self._thread: Thread | None = None
+        self._shutdown = False
+        self._disabled = False
+        self._ready = Event()
+
+    def start(self) -> None:
+        """Start the background writer thread."""
+        if self._thread is not None:
+            return
+        self._thread = Thread(target=self._thread_main, name="ch-writer", daemon=True)
+        self._thread.start()
+        if not self._ready.wait(timeout=10.0):
+            logger.warning("ClickHouse writer thread did not start within 10 seconds")
+
+    def _thread_main(self) -> None:
+        """Entry point for background thread — creates event loop and runs."""
+        self._loop = asyncio.new_event_loop()
+        self._queue = asyncio.Queue()
+        self._ready.set()
+        try:
+            self._loop.run_until_complete(self._run())
+        finally:
+            self._loop.close()
+            self._loop = None
+
+    async def _run(self) -> None:
+        """Main async loop: connect, create tables, then drain queue."""
+        assert self._queue is not None, "_run() must be called after _queue is initialized"
+
+        if not await self._connect_with_retry():
+            return
+
+        pending: dict[str, list[BaseModel]] = {}
+        pending_jobs: list[SummaryJob] = []
+
+        while True:
+            try:
+                item = await asyncio.wait_for(self._queue.get(), timeout=self._flush_interval)
+            except TimeoutError:
+                self._flush_batches(pending)
+                await self._process_summary_jobs(pending_jobs)
+                continue
+
+            if item is _SENTINEL:
+                self._flush_batches(pending)
+                await self._process_summary_jobs(pending_jobs)
+                break
+
+            if isinstance(item, InsertBatch):
+                pending.setdefault(item.table, []).extend(item.rows)
+                if sum(len(v) for v in pending.values()) >= self._batch_size:
+                    self._flush_batches(pending)
+            elif isinstance(item, SummaryJob):
+                pending_jobs.append(item)
+            elif isinstance(item, Event):
+                self._flush_batches(pending)
+                await self._process_summary_jobs(pending_jobs)
+                item.set()
+
+    async def _connect_with_retry(self, max_retries: int = 5, base_delay: float = 1.0) -> bool:
+        """Attempt to connect to ClickHouse with exponential backoff.
+
+        Returns True on success, False if all retries exhausted (writer is disabled).
+        """
+        for attempt in range(max_retries):
+            try:
+                self._client.connect()
+                self._client.ensure_tables()
+                return True
+            except Exception as e:
+                delay = base_delay * (2**attempt)
+                if attempt < max_retries - 1:
+                    logger.warning(f"ClickHouse connection attempt {attempt + 1}/{max_retries} failed: {e}. Retrying in {delay:.0f}s")
+                    await asyncio.sleep(delay)
+                else:
+                    logger.warning(f"ClickHouse connection failed after {max_retries} attempts, tracking disabled: {e}")
+                    self._disabled = True
+                    return False
+        return False  # unreachable
+
+    def _flush_batches(self, pending: dict[str, list[BaseModel]]) -> None:
+        """Flush all pending inserts to ClickHouse."""
+        flushed: list[str] = []
+        for table, rows in list(pending.items()):
+            if rows:
+                try:
+                    self._client._insert_rows(table, rows)  # pyright: ignore[reportPrivateUsage]
+                    flushed.append(table)
+                except Exception as e:
+                    logger.warning(f"Failed to insert {len(rows)} rows into {table}: {e}")
+        for table in flushed:
+            del pending[table]
+
+    async def _process_summary_jobs(self, jobs: list[SummaryJob]) -> None:
+        """Process all pending summary jobs in parallel."""
+        if not jobs:
+            return
+        await asyncio.gather(*[self._process_summary_job(job) for job in jobs])
+        jobs.clear()
+
+    async def _process_summary_job(self, job: SummaryJob) -> None:
+        """Process a span summary generation job."""
+        try:
+            if self._span_summary_fn:
+                otel_token = otel_context.attach(job.parent_otel_context) if job.parent_otel_context is not None else None
+                laminar_token = laminar_context.attach_context(job.parent_laminar_context) if job.parent_laminar_context is not None else None
+                try:
+                    summary = await self._span_summary_fn(job.label, job.output_hint, job.summary_model)
+                finally:
+                    if laminar_token is not None:
+                        laminar_context.detach_context(laminar_token)
+                    if otel_token is not None:
+                        otel_context.detach(otel_token)
+                if summary and self._summary_row_builder:
+                    row = self._summary_row_builder.build_span_summary_update(job.span_id, summary)
+                    if row:
+                        self._client.update_span(row)
+        except Exception as e:
+            logger.warning(f"Summary job failed: {e}")
+
+    def write(self, table: str, rows: list[BaseModel]) -> None:
+        """Enqueue rows for insertion. Thread-safe, non-blocking."""
+        if self._disabled or self._shutdown or self._loop is None or self._queue is None:
+            return
+        batch = InsertBatch(table=table, rows=rows)
+        with contextlib.suppress(RuntimeError):
+            self._loop.call_soon_threadsafe(self._queue.put_nowait, batch)
+
+    def write_job(self, job: SummaryJob) -> None:
+        """Enqueue a summary job. Thread-safe, non-blocking."""
+        if self._disabled or self._shutdown or self._loop is None or self._queue is None:
+            return
+        with contextlib.suppress(RuntimeError):
+            self._loop.call_soon_threadsafe(self._queue.put_nowait, job)
+
+    def flush(self, timeout: float = 30.0) -> None:
+        """Block until all queued items (including summary jobs) are processed."""
+        if self._disabled or self._shutdown or self._loop is None or self._queue is None:
+            return
+        barrier = Event()
+        try:
+            self._loop.call_soon_threadsafe(self._queue.put_nowait, barrier)
+        except RuntimeError:
+            return
+        barrier.wait(timeout=timeout)
+
+    def shutdown(self, timeout: float = 60.0) -> None:
+        """Signal shutdown and wait for the writer thread to finish."""
+        if self._shutdown:
+            return
+        self._shutdown = True
+        if self._loop is not None and self._queue is not None:
+            with contextlib.suppress(RuntimeError):
+                self._loop.call_soon_threadsafe(self._queue.put_nowait, _SENTINEL)
+        if self._thread is not None:
+            self._thread.join(timeout=timeout)