src/tracing.py

"""Optional LangSmith tracing and structured logging.

Everything here degrades to a no-op. If `langsmith` is not installed, or
LANGSMITH_TRACING is not "true", the decorator becomes an identity function
and the helpers do nothing -- so pipeline.py imports this unconditionally and
the Space runs fine with no API key set.

    export LANGSMITH_TRACING=true
    export LANGSMITH_API_KEY=<key>
    export LANGSMITH_PROJECT=diffusiondb-sd15-lora-inference

Image payloads are deliberately kept out of traces: LangSmith stores run I/O
as JSON, and a base64 512x512 PNG is ~400 KB per run, which gets truncated
and makes the trace view unusable. Set LANGSMITH_LOG_IMAGES=1 to attach a
small WEBP thumbnail (~10 KB) instead of the full image.
"""

from __future__ import annotations

import base64
import io
import logging
import os
from contextlib import contextmanager
from typing import Any, Callable

logger = logging.getLogger("diffusiondb.inference")


def _env_flag(name: str, default: str = "") -> bool:
    return os.environ.get(name, default).strip().lower() in {"1", "true", "yes", "on"}


TRACING_ENABLED = _env_flag("LANGSMITH_TRACING")
LOG_IMAGES = _env_flag("LANGSMITH_LOG_IMAGES")
THUMBNAIL_PX = int(os.environ.get("LANGSMITH_THUMBNAIL_PX", "256"))

# Resolved once. A missing package is not an error -- tracing is opt-in.
_langsmith: Any = None
if TRACING_ENABLED:
    try:
        import langsmith as _langsmith  # noqa: F401
    except ImportError:
        logger.warning(
            "LANGSMITH_TRACING is set but the langsmith package is not "
            "installed; running untraced. pip install langsmith"
        )
        _langsmith = None
    else:
        if not os.environ.get("LANGSMITH_API_KEY"):
            logger.warning("LANGSMITH_TRACING is set but LANGSMITH_API_KEY is empty.")

TRACING_ACTIVE = _langsmith is not None


def traceable(**decorator_kwargs: Any) -> Callable:
    """`langsmith.traceable` when tracing is on, identity decorator otherwise."""
    if not TRACING_ACTIVE:
        def passthrough(func: Callable) -> Callable:
            return func
        return passthrough
    return _langsmith.traceable(**decorator_kwargs)


def add_trace_metadata(**fields: Any) -> None:
    """Attach key/values to the run currently in scope. No-op when untraced.

    Used for facts only known at request time (resolved device, dtype, the
    checkpoint actually loaded) which cannot be passed at decoration time.
    """
    if not TRACING_ACTIVE:
        return
    try:
        run = _langsmith.get_current_run_tree()
        if run is not None:
            run.extra.setdefault("metadata", {}).update(fields)
    except Exception:  # never let telemetry break a generation
        logger.debug("could not attach trace metadata", exc_info=True)


def _incoming_headers(request: Any) -> dict:
    """The distributed-tracing headers on an inbound Gradio request."""
    try:
        headers = getattr(request, "headers", None) or {}
        lowered = {str(k).lower(): v for k, v in dict(headers).items()}
        found = {
            name: lowered[name] for name in ("langsmith-trace", "baggage")
            if lowered.get(name)
        }
        return found if "langsmith-trace" in found else {}
    except Exception:
        return {}


@contextmanager
def continue_trace(request: Any):
    """Attach this generation to the caller's trace, when there is one.

    The UI Space sends its run id in `langsmith-trace` and its metadata --
    including the visitor's location -- in `baggage`. Entering that context
    makes the run created inside a child of the UI's run rather than a
    separate root, so one trace spans UI and backend.

    Falls through untouched when tracing is off, when the caller sent no
    headers (a direct API call), or if anything goes wrong: a broken trace
    must never cost an image.
    """
    if not TRACING_ACTIVE:
        yield
        return
    headers = _incoming_headers(request)
    if not headers:
        yield
        return
    try:
        from langsmith.run_helpers import tracing_context
    except Exception:
        yield
        return
    # Only the context construction is guarded. Wrapping the yield itself
    # would swallow exceptions raised by the generation and then yield a
    # second time.
    try:
        context = tracing_context(parent=headers)
    except Exception:
        logger.debug("could not continue caller's trace", exc_info=True)
        yield
        return
    with context:
        yield


def thumbnail_data_uri(image: Any, max_px: int = THUMBNAIL_PX) -> str | None:
    """Downscale a PIL image to a small WEBP data URI, or None on failure."""
    try:
        thumb = image.copy()
        thumb.thumbnail((max_px, max_px))
        buffer = io.BytesIO()
        thumb.save(buffer, format="WEBP", quality=70)
        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
        return f"data:image/webp;base64,{encoded}"
    except Exception:
        logger.debug("thumbnail failed", exc_info=True)
        return None


def configure_logging(level: str | None = None) -> None:
    """Single-line structured logs. Safe to call more than once."""
    if logger.handlers:
        return
    handler = logging.StreamHandler()
    handler.setFormatter(
        logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
    )
    logger.addHandler(handler)
    logger.setLevel(level or os.environ.get("LOG_LEVEL", "INFO").upper())
    logger.propagate = False