stack/src/llm/metrics.py

"""OpenTelemetry metric instruments for the llm service (#579).

Follows the P26 pipeline-telemetry convention: ``perf.meter`` hands back
a real OTel meter when telemetry is enabled and a zero-overhead
``MockMeter`` when it is not, so every call site here is free in tests
and in a container started without ``STACK_TELEMETRY``.

Instruments (all exported as ``stack_llm_…``)::

    stack_llm_dispatch_total{host,kind}        host slots handed out
    stack_llm_inflight{host}                   slots held right now
    stack_llm_dispatch_seconds{host,kind}      how long a slot was held
    stack_llm_embedded_texts_total{host}       texts embedded
    stack_llm_embed_batch_seconds{host}        one /api/embed round trip
    stack_llm_chat_seconds{stage}              retrieve / generate
    stack_llm_chat_total{outcome}              ok / error
    stack_llm_chat_tokens_total{direction}     prompt / completion
    stack_llm_chat_prompt_overflow_total{resolution}  rebuilt / host_truncated
    stack_llm_indexed_chunks_total{collection} chunks written to pgvector
    stack_llm_indexed_items_total{collection,outcome}

Attribute values are deliberately closed sets — host, kind, stage,
outcome, collection — never item keys, dockets or question text, so the
time series count stays bounded.

The tagging chain (#575) does not exist yet, so tag throughput and
abstain rate are not instrumented here; they belong with that module.
"""

from __future__ import annotations

import threading
import time
from contextlib import contextmanager
from typing import Any, Iterator

#: Instrumentation scope, mirroring ``perf.collector``'s "stack.pipeline".
METER_NAME = "stack.llm"

_instruments: "_Instruments | None" = None
#: ``embed_texts`` fans batches across a thread pool, so first use can be
#: concurrent — build the instrument set exactly once all the same.
_build_lock = threading.Lock()


class _Instruments:
    """The service's instrument set, built once on first use."""

    def __init__(self) -> None:
        from perf import meter

        m: Any = meter(METER_NAME)
        self.dispatch_total = m.create_counter(
            "stack_llm_dispatch_total",
            unit="1",
            description="Ollama host slots handed out by the pool",
        )
        self.inflight = m.create_up_down_counter(
            "stack_llm_inflight",
            unit="1",
            description="Requests currently holding an Ollama host slot",
        )
        self.dispatch_seconds = m.create_histogram(
            "stack_llm_dispatch_seconds",
            unit="s",
            description="Time a request held an Ollama host slot",
        )
        self.embedded_texts_total = m.create_counter(
            "stack_llm_embedded_texts_total",
            unit="1",
            description="Texts embedded through the host pool",
        )
        self.embed_batch_seconds = m.create_histogram(
            "stack_llm_embed_batch_seconds",
            unit="s",
            description="Latency of one /api/embed batch",
        )
        self.chat_seconds = m.create_histogram(
            "stack_llm_chat_seconds",
            unit="s",
            description="Chat answer latency by stage",
        )
        self.chat_total = m.create_counter(
            "stack_llm_chat_total",
            unit="1",
            description="Chat answers by outcome",
        )
        self.chat_tokens_total = m.create_counter(
            "stack_llm_chat_tokens_total",
            unit="1",
            description="Chat tokens reported by Ollama, by direction",
        )
        self.chat_prompt_overflow_total = m.create_counter(
            "stack_llm_chat_prompt_overflow_total",
            unit="1",
            description="Chat prompts the host refused as longer than the context, by how they were resolved",
        )
        self.indexed_chunks_total = m.create_counter(
            "stack_llm_indexed_chunks_total",
            unit="1",
            description="Chunks written to pgvector by the indexer",
        )
        self.indexed_items_total = m.create_counter(
            "stack_llm_indexed_items_total",
            unit="1",
            description="Items the indexer embedded or skipped",
        )


def instruments() -> _Instruments:
    """The cached instrument set (created on first call)."""
    global _instruments
    if _instruments is None:
        with _build_lock:
            if _instruments is None:
                _instruments = _Instruments()
    return _instruments


def reset() -> None:
    """Drop the cached instruments so the next call rebuilds them.

    Only tests need this — they patch ``perf.meter`` and then want the
    instruments built against the patched meter.
    """
    global _instruments
    _instruments = None


@contextmanager
def dispatch(host: str, kind: str) -> Iterator[None]:
    """Count and time one held pool slot on *host* (``embed``/``generate``)."""
    i = instruments()
    slot = {"host": host, "kind": kind}
    i.dispatch_total.add(1, slot)
    i.inflight.add(1, {"host": host})
    started = time.monotonic()
    try:
        yield
    finally:
        i.inflight.add(-1, {"host": host})
        i.dispatch_seconds.record(time.monotonic() - started, slot)


def embedded(host: str, texts: int, seconds: float) -> None:
    """Record one finished embed batch of *texts* items on *host*."""
    i = instruments()
    i.embedded_texts_total.add(texts, {"host": host})
    i.embed_batch_seconds.record(seconds, {"host": host})


def chat_stage(stage: str, seconds: float) -> None:
    """Record the duration of a chat *stage* (``retrieve``/``generate``)."""
    instruments().chat_seconds.record(seconds, {"stage": stage})


def chat_outcome(outcome: str) -> None:
    """Count one finished chat answer (``ok``/``error``)."""
    instruments().chat_total.add(1, {"outcome": outcome})


def chat_tokens(prompt: int, completion: int) -> None:
    """Count the tokens Ollama reported for one answer.

    Either count is absent on a host/model that does not report it; a
    zero is not recorded rather than being counted as a real zero.
    """
    i = instruments()
    if prompt:
        i.chat_tokens_total.add(prompt, {"direction": "prompt"})
    if completion:
        i.chat_tokens_total.add(completion, {"direction": "completion"})


def chat_prompt_overflow(resolution: str) -> None:
    """Count one prompt the host refused as longer than its context —
    ``rebuilt`` (trimmed tighter and resent) or ``host_truncated`` (the
    last resort: resent with host-side truncation allowed)."""
    instruments().chat_prompt_overflow_total.add(1, {"resolution": resolution})


def indexed_item(collection: str, outcome: str, chunks: int = 0) -> None:
    """Count one item the indexer finished (``indexed``/``skipped``)."""
    i = instruments()
    i.indexed_items_total.add(1, {"collection": collection, "outcome": outcome})
    if chunks:
        i.indexed_chunks_total.add(chunks, {"collection": collection})