diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 218f59f8..34a1e53d 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -32,6 +32,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config from haiku.rag.skills.analysis import AnalysisState from haiku.rag.skills.rag import RAGState, get_agent_preamble +from haiku.rag.telemetry import configure as configure_telemetry from haiku.skills.agent import ( SkillToolset, run_agui_stream, @@ -39,13 +40,7 @@ from haiku.skills.agent import ( from haiku.skills.models import Skill from haiku.skills.prompts import build_system_prompt -try: - import logfire - - logfire.configure(send_to_logfire="if-token-present", console=False) - logfire.instrument_pydantic_ai() -except ImportError: - pass +configure_telemetry() if TYPE_CHECKING: from textual.app import ComposeResult diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 8e0a2c66..26b3dd1c 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -134,18 +134,10 @@ def main( # Configure logging for CLI context configure_cli_logging() - # Configure logfire (only sends data if token is present) - try: - import logfire + from haiku.rag.telemetry import configure as configure_telemetry - is_production = get_config().environment != "development" - logfire.configure( - send_to_logfire="if-token-present", - console=False if is_production else None, - ) - logfire.instrument_pydantic_ai() - except Exception: # pragma: no cover - pass + is_production = get_config().environment != "development" + configure_telemetry(console=False if is_production else None) if get_config().environment != "development": # Suppress warnings in production diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index ec5c2d60..efc48165 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -3,8 +3,6 @@ from pathlib import Path from typing import TYPE_CHECKING from urllib.parse import unquote, urlparse -import logfire - from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.client.processing import ( ensure_chunks_embedded, @@ -16,6 +14,7 @@ from haiku.rag.ingester.sources import FetchResult, resolve_fetcher from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import extract_items +from haiku.rag.telemetry import logfire if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument diff --git a/haiku_rag_slim/haiku/rag/converters/pdf_split.py b/haiku_rag_slim/haiku/rag/converters/pdf_split.py index 6139bb14..83dfb147 100644 --- a/haiku_rag_slim/haiku/rag/converters/pdf_split.py +++ b/haiku_rag_slim/haiku/rag/converters/pdf_split.py @@ -6,10 +6,10 @@ from collections.abc import Generator from pathlib import Path from typing import TYPE_CHECKING -import logfire import pypdfium2 as pdfium from haiku.rag.client.exceptions import UnsupportedSourceError +from haiku.rag.telemetry import logfire # pypdfium2 wraps libpdfium, which has global C state and is not thread-safe. # Multiple workers calling iter_pdf_slices concurrently race on that state diff --git a/haiku_rag_slim/haiku/rag/ingester/cli.py b/haiku_rag_slim/haiku/rag/ingester/cli.py index 3cb0b81e..ab112015 100644 --- a/haiku_rag_slim/haiku/rag/ingester/cli.py +++ b/haiku_rag_slim/haiku/rag/ingester/cli.py @@ -50,28 +50,12 @@ def main( _load_config_with_override(config) -def _configure_logfire() -> None: - """Logfire emits spans only when LOGFIRE_TOKEN is set; otherwise it - stays silent. Console output is disabled in either case so span lines - don't interleave with the ingester's own RichHandler logs — telemetry - lives in the Logfire UI.""" - try: - import logfire - - logfire.configure( - service_name="haiku-ingester", - send_to_logfire="if-token-present", - console=False, - ) - logfire.instrument_pydantic_ai() - except Exception: # pragma: no cover - pass - - def cli() -> None: """Entry point that translates store-state errors into a clean exit.""" + from haiku.rag.telemetry import configure as configure_telemetry + configure_cli_logging() - _configure_logfire() + configure_telemetry(service_name="haiku-ingester") try: _cli() except (MigrationRequiredError, ReadOnlyError) as e: diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py index eb5a0068..f9348d4d 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/base.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/base.py @@ -2,8 +2,6 @@ import asyncio import logging from datetime import UTC, datetime -import logfire - from haiku.rag.config import SourceConfig from haiku.rag.ingester.pollers.circuit_breaker import CircuitBreaker from haiku.rag.ingester.queue.models import JobOp @@ -13,6 +11,7 @@ from haiku.rag.ingester.sources.base import ( SourceEvent, SourceEventKind, ) +from haiku.rag.telemetry import get_context, logfire logger = logging.getLogger(__name__) @@ -29,7 +28,7 @@ def _enqueue_extra(cfg: SourceConfig) -> dict | None: headers = getattr(cfg, "headers", None) if headers: extra["headers"] = dict(headers) - carrier = logfire.get_context() + carrier = get_context() if carrier: extra["_otel"] = dict(carrier) return extra or None diff --git a/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py b/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py index 22d8634d..b01a50df 100644 --- a/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/pollers/fs.py @@ -3,12 +3,12 @@ import logging from pathlib import Path from typing import TYPE_CHECKING -import logfire from watchfiles import Change, awatch from haiku.rag.ingester.pollers.base import BasePoller, _enqueue_extra from haiku.rag.ingester.queue.models import JobOp from haiku.rag.ingester.sources.filter import FileFilter +from haiku.rag.telemetry import logfire if TYPE_CHECKING: from haiku.rag.config import FSSourceConfig diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py index 05aaa3ae..fade3726 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pipeline.py @@ -3,12 +3,12 @@ from contextlib import nullcontext from typing import TYPE_CHECKING import httpx -import logfire from pydantic import BaseModel from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.ingester.exceptions import PermanentError, TransientError from haiku.rag.ingester.queue.models import Job, JobOp +from haiku.rag.telemetry import attach_context, logfire if TYPE_CHECKING: from haiku.rag.client import HaikuRAG @@ -75,7 +75,7 @@ async def run_job(client: "HaikuRAG", job: Job) -> JobResult: # Restore the poller's trace context (if any) so the job span nests # under the `ingester.poller.sweep` that enqueued it. parent_ctx = extra.get("_otel") - attach = logfire.attach_context(parent_ctx) if parent_ctx else nullcontext() + attach = attach_context(parent_ctx) if parent_ctx else nullcontext() with ( attach, diff --git a/haiku_rag_slim/haiku/rag/telemetry.py b/haiku_rag_slim/haiku/rag/telemetry.py new file mode 100644 index 00000000..0a34fcbb --- /dev/null +++ b/haiku_rag_slim/haiku/rag/telemetry.py @@ -0,0 +1,47 @@ +from typing import Literal + +from logfire import Logfire, attach_context, get_context + +# Scoped Logfire instance — every span emitted through `logfire.span(...)` +# on this object carries `instrumentation_scope.name = "haiku.rag"` instead +# of the default "logfire". The scope is OTel's identifier for *which +# library* produced the span, separate from `service.name` which is the +# running process. Downstream consumers (Logfire UI saved views, OTel +# collectors, alert rules) can then filter on `scope.name = haiku.rag` +# rather than catching every span the SDK ever exports. +# +# Cross-library instrumentations (pydantic-ai, FastAPI, OpenAI) keep their +# own scopes — this only retags the spans WE write. +logfire = Logfire(otel_scope="haiku.rag") + + +def configure( + *, + service_name: str | None = None, + console: Literal[False] | None = False, +) -> None: + """Configure Logfire and enable pydantic-ai instrumentation for the + running process. Each CLI entry point calls this once at startup. + Silently no-ops on failure so a missing/misconfigured LOGFIRE_TOKEN + never crashes the app. + + - service_name: identifies the process in the Logfire UI (e.g. + "haiku-ingester"). Falls back to logfire's default when None. + - console: False (default) suppresses span lines on stderr so they + don't interleave with RichHandler logs. Pass None to let logfire + decide (its own default applies). + """ + try: + import logfire as _lf + + _lf.configure( + service_name=service_name, + send_to_logfire="if-token-present", + console=console, + ) + _lf.instrument_pydantic_ai() + except Exception: # pragma: no cover + pass + + +__all__ = ["attach_context", "configure", "get_context", "logfire"]