backend/observability.py

"""Pipeline observability helpers for step, API, LLM, and cost tracking."""
from __future__ import annotations

from contextlib import asynccontextmanager
from contextvars import ContextVar
from time import perf_counter
from typing import Any, AsyncIterator

from backend.config import settings

_current_step: ContextVar[dict[str, Any] | None] = ContextVar("current_pipeline_step", default=None)


def _elapsed_ms(start: float) -> int:
    return max(0, round((perf_counter() - start) * 1000))


def current_step_id() -> int | None:
    """pipeline_step id of the step currently running (None outside a step).

    Lets pipeline code link evidence_items / provider_calls to the live step
    without threading the id through every call.
    """
    context = _current_step.get()
    return context.get("step_id") if context else None


@asynccontextmanager
async def tracked_step(
    access_key: str,
    step_name: str,
    step_order: float | None = None,
    metadata: dict[str, Any] | None = None,
) -> AsyncIterator[dict[str, Any]]:
    from backend.storage import jobs as job_store

    started = perf_counter()
    step = await job_store.start_pipeline_step(
        access_key,
        step_name,
        step_order=step_order,
        metadata=metadata,
    )
    token = _current_step.set(
        {
            "access_key": access_key,
            "step_id": step.get("id"),
            "step_name": step_name,
        }
    )
    try:
        yield step
    except Exception as exc:
        await job_store.finish_pipeline_step(
            access_key,
            step["id"],
            "failed",
            latency_ms=_elapsed_ms(started),
            error=f"{type(exc).__name__}: {exc}",
        )
        raise
    else:
        # The body may stash its output on the step dict (see _run_tracked_step).
        await job_store.finish_pipeline_step(
            access_key,
            step["id"],
            "completed",
            latency_ms=_elapsed_ms(started),
            error=None,
            output_json=step.get("_output_json"),
            output_summary=step.get("_output_summary"),
        )
    finally:
        _current_step.reset(token)


async def record_llm_call(
    provider: str,
    model: str,
    response_format: str | None,
    prompt_tokens: int | None,
    completion_tokens: int | None,
    total_tokens: int | None,
    latency_ms: int | None,
    status: str,
    error: str | None,
) -> None:
    context = _current_step.get()
    if context is None:
        return
    from backend.storage import jobs as job_store

    await job_store.record_llm_call(
        access_key=context["access_key"],
        step_id=context.get("step_id"),
        step_name=context.get("step_name"),
        provider=provider,
        model=model,
        response_format=response_format,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        input_price_per_1m=_input_price_for_provider(provider),
        output_price_per_1m=_output_price_for_provider(provider),
        latency_ms=latency_ms,
        status=status,
        error=error,
    )


async def record_api_call(
    provider: str,
    operation: str,
    status_code: int | None = None,
    latency_ms: int | None = None,
    request_count: int = 1,
    status: str = "success",
    error: str | None = None,
    cost_usd: float | None = None,
    metadata: dict[str, Any] | None = None,
) -> None:
    context = _current_step.get()
    if context is None:
        return
    from backend.storage import jobs as job_store

    await job_store.record_api_call(
        access_key=context["access_key"],
        step_id=context.get("step_id"),
        step_name=context.get("step_name"),
        provider=provider,
        operation=operation,
        status_code=status_code,
        latency_ms=latency_ms,
        request_count=request_count,
        status=status,
        error=error,
        cost_usd=cost_usd,
        metadata=metadata,
    )


def _input_price_for_provider(provider: str) -> float | None:
    if provider == "openai":
        return settings.openai_input_price_per_1m
    if provider == "anthropic":
        return settings.anthropic_input_price_per_1m
    if provider == "gemini":
        return settings.gemini_input_price_per_1m
    if provider == "openrouter":
        return settings.openrouter_input_price_per_1m
    return None


def _output_price_for_provider(provider: str) -> float | None:
    if provider == "openai":
        return settings.openai_output_price_per_1m
    if provider == "anthropic":
        return settings.anthropic_output_price_per_1m
    if provider == "gemini":
        return settings.gemini_output_price_per_1m
    if provider == "openrouter":
        return settings.openrouter_output_price_per_1m
    return None