AI_Agro_Support/pipeline/observability.py

"""Opt-in OpenTelemetry/OpenInference tracing to a self-hosted Phoenix
instance (see `phoenix` in docker-compose.yml).

Disabled by default (`config.yaml` `observability.enabled: false`). Every
public function in this module is fail-open:

- If tracing is disabled, every helper below is a cheap no-op.
- If `setup_tracing` fails for any reason (missing packages, a bad endpoint
  URL, an instrumentor import error), it logs a warning and leaves tracing
  disabled for the rest of the process -- callers never need to check
  whether setup succeeded.
- `phoenix.otel.register()` only builds an exporter; it does not contact the
  collector, so a stopped/unreachable Phoenix container cannot fail startup.
- Spans are exported on a background thread (`batch=True`), so a slow or
  down collector cannot block a worker thread mid-job.

`call_llm` must keep returning `str` (see pipeline/stages/llm.py); nothing
here changes that. `job_span` / `llm_stage_span` / `retrieval_span` wrap
existing calls without altering their signatures or return values.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
import sys
from datetime import date
from typing import TYPE_CHECKING, Any, Iterator, Sequence

if TYPE_CHECKING:
    from pipeline.config import Settings
    from pipeline.worklist import Job

logger = logging.getLogger(__name__)

_tracer: Any = None
_tracer_provider: Any = None


class _NullSpan:
    """Stand-in for a span when tracing is off or failed to start; every
    method is a no-op so callers never need an `if span is not None` check."""

    def set_attribute(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def set_attributes(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def set_input(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def set_output(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def set_status(self, *_args: Any, **_kwargs: Any) -> None:
        pass

    def record_exception(self, *_args: Any, **_kwargs: Any) -> None:
        pass


class _SafeSpan:
    """Proxy around a real OTel span: attribute writes must never fail a job.

    `run_job`'s contract is that it never raises. If a successful
    `_run_job_inner` were followed by a raising `span.set_attribute`, the
    batch would record that job as failed even though advice was already
    stored. Every mutating method here is therefore fail-open.
    """

    def __init__(self, inner: Any) -> None:
        self._inner = inner

    def set_attribute(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.set_attribute(*args, **kwargs)
        except Exception:
            logger.debug("Failed to set span attribute.", exc_info=True)

    def set_attributes(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.set_attributes(*args, **kwargs)
        except Exception:
            logger.debug("Failed to set span attributes.", exc_info=True)

    def set_input(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.set_input(*args, **kwargs)
        except Exception:
            logger.debug("Failed to set span input.", exc_info=True)

    def set_output(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.set_output(*args, **kwargs)
        except Exception:
            logger.debug("Failed to set span output.", exc_info=True)

    def set_status(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.set_status(*args, **kwargs)
        except Exception:
            logger.debug("Failed to set span status.", exc_info=True)

    def record_exception(self, *args: Any, **kwargs: Any) -> None:
        try:
            self._inner.record_exception(*args, **kwargs)
        except Exception:
            logger.debug("Failed to record span exception.", exc_info=True)


_NULL_SPAN = _SafeSpan(_NullSpan())


def setup_tracing(settings: "Settings") -> None:
    """
    Register the OTel tracer provider and instrument the Gemini/OpenAI/
    Anthropic clients.

    Call once at CLI startup, right after `load_settings()` and before any
    provider client is constructed (clients are created inside
    `pipeline.stages.llm`'s `_call_gemini` / `_call_openai` /
    `_call_anthropic`, so instrumentation must be wired up before the first
    call, not before the first client object).
    """
    global _tracer, _tracer_provider

    cfg = settings.observability
    if not cfg.enabled:
        return
    if _tracer is not None:
        return  # already set up in this process

    try:
        # Must be set before importing phoenix.otel / opentelemetry.sdk, which
        # read span-limit env vars at import/construction time.
        os.environ.setdefault("OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT", str(cfg.max_attribute_chars))
        # Caps how long a flush can block at process exit if Phoenix is down
        # (see shutdown_tracing); the OTLP/HTTP exporter default is 10s/try.
        os.environ.setdefault("OTEL_EXPORTER_OTLP_TRACES_TIMEOUT", "5")

        from openinference.instrumentation import TraceConfig
        from phoenix.otel import register

        trace_config = TraceConfig(
            hide_inputs=cfg.hide_prompts,
            hide_input_messages=cfg.hide_prompts,
            hide_prompts=cfg.hide_prompts,
        )

        tracer_provider = register(
            project_name=cfg.project_name,
            endpoint=cfg.endpoint,
            protocol="http/protobuf",
            batch=True,
            auto_instrument=False,
            set_global_tracer_provider=True,
            verbose=False,
        )

        _instrument_providers(tracer_provider, trace_config)

        _tracer_provider = tracer_provider
        _tracer = tracer_provider.get_tracer("pipeline")
        logger.info(
            "Tracing enabled: project=%s endpoint=%s", cfg.project_name, cfg.endpoint
        )
    except Exception:
        logger.warning(
            "Failed to set up tracing; continuing without it (batch/one are unaffected).",
            exc_info=True,
        )
        _tracer = None
        _tracer_provider = None


def _instrument_providers(tracer_provider: Any, trace_config: Any) -> None:
    """Instrument each provider independently: one provider's failure (or
    absence) must not prevent the others from being traced."""
    try:
        from openinference.instrumentation.google_genai import GoogleGenAIInstrumentor

        GoogleGenAIInstrumentor().instrument(tracer_provider=tracer_provider, config=trace_config)
    except Exception:
        logger.warning(
            "Gemini tracing instrumentation failed; Gemini calls will run untraced "
            "(token counts/cost will be absent from LLM stage spans).",
            exc_info=True,
        )

    try:
        from openinference.instrumentation.openai import OpenAIInstrumentor

        OpenAIInstrumentor().instrument(tracer_provider=tracer_provider, config=trace_config)
    except Exception:
        logger.debug("OpenAI tracing instrumentation unavailable or failed.", exc_info=True)

    try:
        from openinference.instrumentation.anthropic import AnthropicInstrumentor

        AnthropicInstrumentor().instrument(tracer_provider=tracer_provider, config=trace_config)
    except Exception:
        logger.debug("Anthropic tracing instrumentation unavailable or failed.", exc_info=True)


def shutdown_tracing() -> None:
    """
    Flush and close the tracer provider, if one was set up.

    Call once at process exit -- after operational artifacts
    (`_run_report.json`, per-job output files) are already written, so a
    slow or unreachable collector delays nothing that matters. Bounded by
    `OTEL_EXPORTER_OTLP_TRACES_TIMEOUT` (set in `setup_tracing`), not the
    SDK's 30s default.
    """
    global _tracer, _tracer_provider
    if _tracer_provider is None:
        return
    try:
        _tracer_provider.shutdown()
    except Exception:
        logger.warning("Error shutting down the tracer provider.", exc_info=True)
    finally:
        _tracer = None
        _tracer_provider = None


@contextlib.contextmanager
def _span(name: str, kind: str, metadata: dict[str, Any] | None = None) -> Iterator[Any]:
    """
    Shared, fail-open span opener used by every helper below.

    Only span *creation* is guarded by try/except; once a span is
    successfully open, exceptions raised by the caller's own code inside the
    `with` block propagate normally (recorded on the span as an error, then
    re-raised) -- tracing must never swallow a real pipeline failure.
    """
    if _tracer is None:
        yield _NULL_SPAN
        return

    try:
        cm = _tracer.start_as_current_span(name, openinference_span_kind=kind)
        span = cm.__enter__()
    except Exception:
        logger.debug("Failed to open span %r; continuing untraced.", name, exc_info=True)
        yield _NULL_SPAN
        return

    if metadata:
        try:
            span.set_attribute("metadata", json.dumps(metadata, ensure_ascii=False, default=str))
        except Exception:
            logger.debug("Failed to set metadata on span %r.", name, exc_info=True)

    try:
        yield _SafeSpan(span)
    finally:
        try:
            cm.__exit__(*sys.exc_info())
        except Exception:
            logger.debug("Failed to close span %r.", name, exc_info=True)


def job_span(job: "Job", as_of: date, dry_run: bool) -> contextlib.AbstractContextManager:
    """
    One root CHAIN span per `(field, crop, disease, as_of)` job.

    `run_job` sets the final `job.status` (and, on failure, `job.error_type`)
    on the yielded span itself before the `with` block exits; this function
    only opens the span and stamps the metadata known up front.
    """
    metadata = {
        "field_id": job.field_id,
        "crop": job.crop_english,
        "disease": job.disease_english,
        "as_of": as_of.isoformat(),
        "organic": job.organic,
        "station": job.station,
        "dry_run": dry_run,
    }
    return _span("run_job", "chain", metadata)


def llm_stage_span(stage: str, provider: str, model: str) -> contextlib.AbstractContextManager:
    """CHAIN span around one `Resources.call_llm` invocation (one of
    `query_synthesis` / `advice_generation`), wrapping the provider's own
    LLM span (added by the OpenInference auto-instrumentor, if any)."""
    metadata = {"stage": stage, "provider": provider, "model": model}
    return _span(stage, "chain", metadata)


def retrieval_span(queries: Sequence[str], candidate_count: int) -> contextlib.AbstractContextManager:
    """
    RETRIEVER span around one `Resources.search_products` call.

    The queries are short, LLM-synthesised search phrases (already
    anonymised relative to the full advice payload), so they are safe to
    record as the span input even when `hide_prompts` is set -- `hide_prompts`
    only affects the LLM instrumentor's own input-message capture.
    """
    metadata = {"query_count": len(queries), "candidate_count": candidate_count}
    return _retrieval_span(queries, metadata)


@contextlib.contextmanager
def _retrieval_span(queries: Sequence[str], metadata: dict[str, Any]) -> Iterator[Any]:
    with _span("search_products", "retriever", metadata) as span:
        span.set_input(list(queries), mime_type="application/json")
        yield span


def set_retrieval_results(span: Any, results: Sequence[Any]) -> None:
    """
    Record retrieved products on a retrieval span as `document.id` /
    `document.score` pairs -- product IDs and Weaviate distances only, never
    product names or label content.

    `results` items are expected to expose `.product_id` and `.distance`
    (see `pipeline.stages.vector.RecommendedProduct`); any other shape is
    ignored rather than raising.
    """
    try:
        span.set_attribute("retrieval.documents_count", len(results))
        for i, item in enumerate(results):
            product_id = getattr(item, "product_id", None)
            distance = getattr(item, "distance", None)
            if product_id is not None:
                span.set_attribute(f"retrieval.documents.{i}.document.id", str(product_id))
            if distance is not None:
                span.set_attribute(f"retrieval.documents.{i}.document.score", float(distance))
    except Exception:
        logger.debug("Failed to record retrieval results on span.", exc_info=True)


def record_llm_usage(span: Any, usage_metadata: Any) -> None:
    """
    Fallback B (see plan step 5/8): manually copy Gemini token counts onto
    the *current* span from a `google.genai` response's `usage_metadata`,
    for use only if the `openinference-instrumentation-google-genai`
    auto-instrumentation is verified to not capture them for the installed
    `google-genai` version. Not wired into `pipeline/stages/llm.py` unless
    that verification fails -- see README "Observability".

    Mirrors the folding behaviour of the upstream instrumentor: thinking
    tokens (`thoughts_token_count`) are counted once as
    `completion_details.reasoning` and again inside `completion` (Gemini
    bills them as output tokens), never double-added beyond that.
    """
    if usage_metadata is None:
        return
    try:
        prompt = int(getattr(usage_metadata, "prompt_token_count", None) or 0)
        candidates = int(getattr(usage_metadata, "candidates_token_count", None) or 0)
        thoughts = int(getattr(usage_metadata, "thoughts_token_count", None) or 0)
        total = int(getattr(usage_metadata, "total_token_count", None) or 0)
        completion = candidates + thoughts

        if prompt:
            span.set_attribute("llm.token_count.prompt", prompt)
        if thoughts:
            span.set_attribute("llm.token_count.completion_details.reasoning", thoughts)
        if completion:
            span.set_attribute("llm.token_count.completion", completion)
        if total or (prompt + completion):
            span.set_attribute("llm.token_count.total", max(total, prompt + completion))
    except Exception:
        logger.debug("Failed to record manual LLM usage attributes.", exc_info=True)