diff --git a/netra/__init__.py b/netra/__init__.py index 1f51d8b..83eb98e 100644 --- a/netra/__init__.py +++ b/netra/__init__.py @@ -257,10 +257,25 @@ def shutdown(cls) -> None: meter_provider = otel_metrics.get_meter_provider() if hasattr(meter_provider, "force_flush"): meter_provider.force_flush() - if hasattr(meter_provider, "shutdown"): + # _NetraOwnedMeterProvider.shutdown() is a no-op for third-party + # callers (see netra/meter.py); Netra's own teardown must use the + # owner entry point or metrics are never flushed at exit. + if hasattr(meter_provider, "shutdown_as_owner"): + meter_provider.shutdown_as_owner() + elif hasattr(meter_provider, "shutdown"): meter_provider.shutdown() except Exception: pass + # Backstop for LiveKit calls whose session never closed cleanly, so + # their captured audio is flushed rather than abandoned in a queue. + try: + from netra.instrumentation.livekit.audio_capture import close_all_audio_capture + + close_all_audio_capture() + except ImportError: + pass + except Exception: + logger.warning("Failed to shut down LiveKit audio capture", exc_info=True) # Close simulation HTTP client if hasattr(cls, "simulation") and cls.simulation is not None: try: diff --git a/netra/config.py b/netra/config.py index e6cf26b..3dada19 100644 --- a/netra/config.py +++ b/netra/config.py @@ -1,4 +1,5 @@ import json +import logging import os from typing import Any, Dict, List, Optional @@ -6,6 +7,8 @@ from netra.version import __version__ +logger = logging.getLogger(__name__) + # Fallback limits used when no Config has been activated yet (e.g. code paths that # run before ``Netra.init()``, or tests that never call it). Once ``init()`` runs, # the active Config instance's values (resolved from env at init time) take over. @@ -13,6 +16,25 @@ _DEFAULT_CONVERSATION_CONTENT_MAX_LEN = 50000 _DEFAULT_TRIAL_BLOCK_DURATION_SECONDS = 15 * 60 +# --- Voice-agent audio capture (env-only; no Netra.init() parameter) ------------ +# The path segment appended to the OTLP endpoint when no explicit audio endpoint +# is given. Deliberately not derived via UsageHttpClient._resolve_base_url, which +# strips a "/telemetry" suffix — correct for the REST APIs, wrong here. +_AUDIO_CHUNK_PATH = "/v1/audio/chunk" + +# Header names that count as an audio-ingest credential. An unauthenticated PCM +# POST is never attempted. +_AUDIO_AUTH_HEADERS = ("x-api-key", "Authorization") + +_DEFAULT_AUDIO_BATCH_BYTES = 32768 +_DEFAULT_AUDIO_BATCH_INTERVAL_MS = 1000 +_DEFAULT_AUDIO_BUFFER_BYTES = 2097152 +_DEFAULT_AUDIO_MAX_REQUEST_BYTES = 262144 + +_MIN_AUDIO_BATCH_BYTES = 1024 +_MIN_AUDIO_BATCH_INTERVAL_MS = 100 +_MAX_AUDIO_BATCH_INTERVAL_MS = 30000 + class Config: """ @@ -97,8 +119,148 @@ def __init__( None, "TRIAL_BLOCK_DURATION_SECONDS", default=_DEFAULT_TRIAL_BLOCK_DURATION_SECONDS ) + self._resolve_audio_settings() + self._set_trace_content_env() + def _resolve_audio_settings(self) -> None: + """Resolve and validate the voice-agent audio-capture settings. + + Env-only by design: there is no ``capture_audio`` parameter on + ``Netra.init()``. Whether audio is captured at all is decided by + :attr:`audio_capture_enabled`, not by a flag. + + Every validation failure logs a ``WARNING`` naming the setting and the + value actually used, then falls back to a safe value. A bad number MUST + NOT raise out of ``Netra.init()``. + """ + self.audio_endpoint_override = os.getenv("NETRA_AUDIO_ENDPOINT") + self.audio_batch_bytes = self._get_int_config( + None, "NETRA_AUDIO_BATCH_BYTES", default=_DEFAULT_AUDIO_BATCH_BYTES + ) + self.audio_batch_interval_ms = self._get_int_config( + None, "NETRA_AUDIO_BATCH_INTERVAL_MS", default=_DEFAULT_AUDIO_BATCH_INTERVAL_MS + ) + self.audio_buffer_bytes = self._get_int_config( + None, "NETRA_AUDIO_BUFFER_BYTES", default=_DEFAULT_AUDIO_BUFFER_BYTES + ) + self.audio_max_request_bytes = self._get_int_config( + None, "NETRA_AUDIO_MAX_REQUEST_BYTES", default=_DEFAULT_AUDIO_MAX_REQUEST_BYTES + ) + + # Order matters: audio_batch_bytes is clamped against the resolved + # max-request size first, then the two ceilings are raised to whatever + # batch size survived. Doing it the other way round lets a tiny + # max_request_bytes silently shrink the batch below its floor. + if self.audio_max_request_bytes < _MIN_AUDIO_BATCH_BYTES: + logger.warning( + "netra.audio: NETRA_AUDIO_MAX_REQUEST_BYTES=%d is below the minimum batch size; using %d", + self.audio_max_request_bytes, + _MIN_AUDIO_BATCH_BYTES, + ) + self.audio_max_request_bytes = _MIN_AUDIO_BATCH_BYTES + + clamped_batch = min(max(self.audio_batch_bytes, _MIN_AUDIO_BATCH_BYTES), self.audio_max_request_bytes) + if clamped_batch != self.audio_batch_bytes: + logger.warning( + "netra.audio: NETRA_AUDIO_BATCH_BYTES=%d out of range [%d, %d]; using %d", + self.audio_batch_bytes, + _MIN_AUDIO_BATCH_BYTES, + self.audio_max_request_bytes, + clamped_batch, + ) + self.audio_batch_bytes = clamped_batch + + clamped_interval = min( + max(self.audio_batch_interval_ms, _MIN_AUDIO_BATCH_INTERVAL_MS), + _MAX_AUDIO_BATCH_INTERVAL_MS, + ) + if clamped_interval != self.audio_batch_interval_ms: + logger.warning( + "netra.audio: NETRA_AUDIO_BATCH_INTERVAL_MS=%d out of range [%d, %d]; using %d", + self.audio_batch_interval_ms, + _MIN_AUDIO_BATCH_INTERVAL_MS, + _MAX_AUDIO_BATCH_INTERVAL_MS, + clamped_interval, + ) + self.audio_batch_interval_ms = clamped_interval + + if self.audio_buffer_bytes < self.audio_batch_bytes: + logger.warning( + "netra.audio: NETRA_AUDIO_BUFFER_BYTES=%d is below the batch size; using %d", + self.audio_buffer_bytes, + self.audio_batch_bytes, + ) + self.audio_buffer_bytes = self.audio_batch_bytes + + if self.audio_max_request_bytes < self.audio_batch_bytes: + logger.warning( + "netra.audio: NETRA_AUDIO_MAX_REQUEST_BYTES=%d is below the batch size; using %d", + self.audio_max_request_bytes, + self.audio_batch_bytes, + ) + self.audio_max_request_bytes = self.audio_batch_bytes + + # Last, so it sees the settled endpoint override, and once, so the + # missing-credential warning is not re-emitted per LiveKit session. + self._audio_endpoint: Optional[str] = self._resolve_audio_endpoint() + + def audio_endpoint(self) -> Optional[str]: + """Return the audio ingest URL, or None if audio must not be sent. + + This is the ONLY gate on audio capture: there is no ``capture_audio`` + flag. A non-None return means audio WILL be captured and streamed once a + LiveKit session starts. Returns None unless a concrete endpoint resolves + AND an auth header is present. + + Callers treat None as "disable capture entirely", not "retry later" — the + result is resolved from init-time state and does not change during the + process. Resolved once in ``_resolve_audio_settings`` for that reason: the + instrumentor, the per-session hook and the startup log line all ask, and a + misconfiguration should be reported once rather than once per call. + + Returns: + The absolute audio ingest URL, or ``None`` when audio must not be sent. + """ + return self._audio_endpoint + + def _resolve_audio_endpoint(self) -> Optional[str]: + """Work out the audio ingest URL, warning if a credential is missing. + + Returns: + The absolute audio ingest URL, or ``None`` when audio must not be sent. + """ + if self.audio_endpoint_override: + url = self.audio_endpoint_override + elif self.otlp_endpoint: + url = self.otlp_endpoint.rstrip("/") + _AUDIO_CHUNK_PATH + else: + return None + + if not any(header in self.headers for header in _AUDIO_AUTH_HEADERS): + logger.warning( + "netra.audio: an audio endpoint resolved but no credential is configured; " + "audio capture is disabled. Set NETRA_API_KEY or pass an auth header." + ) + return None + + return url + + @property + def audio_capture_enabled(self) -> bool: + """Whether call audio will be captured and streamed. + + The single derived predicate behind audio capture, so the instrumentor, + the session hooks and the startup log line cannot disagree about it. A + resolved endpoint is the whole of it: capture is all of the call's audio or + none of it, never one speaker. + + Returns: + True when an audio endpoint resolves; False otherwise, meaning no + audio is captured or streamed. + """ + return self.audio_endpoint() is not None + def _get_app_name(self, app_name: Optional[str]) -> str: """Get application name from param or environment variables.""" return app_name or os.getenv("NETRA_APP_NAME") or os.getenv("OTEL_SERVICE_NAME") or "llm_tracing_service" diff --git a/netra/instrumentation/__init__.py b/netra/instrumentation/__init__.py index 445bb41..b9f2293 100644 --- a/netra/instrumentation/__init__.py +++ b/netra/instrumentation/__init__.py @@ -196,6 +196,10 @@ def init_instrumentations( if CustomInstruments.DEEPGRAM in netra_custom_instruments: init_deepgram_instrumentation() + # Initialize LiveKit voice-agent instrumentation. + if CustomInstruments.LIVEKIT in netra_custom_instruments: + init_livekit_instrumentation() + # Initialize ADK instrumentation. if CustomInstruments.ADK in netra_custom_instruments: init_adk_instrumentation() @@ -430,6 +434,25 @@ def init_deepgram_instrumentation() -> bool: return False +def init_livekit_instrumentation() -> bool: + """Initialize LiveKit voice-agent instrumentation. + + Returns: + bool: True if initialization was successful, False otherwise. + """ + try: + if is_package_installed("livekit-agents"): + from netra.instrumentation.livekit import NetraLiveKitInstrumentor + + instrumentor = NetraLiveKitInstrumentor() + if not instrumentor.is_instrumented_by_opentelemetry: + instrumentor.instrument() + return True + except Exception: + logging.exception("Error initializing LiveKit instrumentor") + return False + + def init_adk_instrumentation() -> bool: """Initialize ADK instrumentation. diff --git a/netra/instrumentation/instruments.py b/netra/instrumentation/instruments.py index 79f5002..8fd380b 100644 --- a/netra/instrumentation/instruments.py +++ b/netra/instrumentation/instruments.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Optional, Type +from typing import Any, Dict, Optional, Type from traceloop.sdk import Instruments @@ -74,6 +74,7 @@ class CustomInstruments(Enum): ELEVENLABS = "elevenlabs" CLAUDE_AGENT_SDK = "claude_agent_sdk" HERMES_AGENT = "hermes_agent" + LIVEKIT = "livekit" class InstrumentSet(Enum): @@ -146,6 +147,7 @@ def __new__(cls, value: Any, origin: Optional[Type[Enum]] = None) -> "Instrument LANCEDB = ("lancedb", Instruments) LANGCHAIN = ("langchain", Instruments) LITELLM = ("litellm", CustomInstruments) + LIVEKIT = ("livekit", CustomInstruments) LLAMA_INDEX = ("llama_index", Instruments) LOGGING = ("logging", CustomInstruments) MARQO = ("marqo", Instruments) @@ -197,6 +199,23 @@ def __new__(cls, value: Any, origin: Optional[Type[Enum]] = None) -> "Instrument NetraInstruments = InstrumentSet +# Instrumentation scopes that Netra enables but does not author, so their scope +# name does not follow the ``netra.instrumentation.*`` / +# ``opentelemetry.instrumentation.*`` convention that the span processors key +# off. Mapping the scope to its ``InstrumentSet`` value is what puts these +# spans under the same instrument-name machinery as every other +# instrumentation — ``root_instruments`` filtering and the +# ``netra.instrumentation.name`` attribute. +# +# Matched exactly, never as a prefix: an alias claims one specific scope, and a +# prefix match here would start pulling in unrelated third-party tracers. +THIRD_PARTY_INSTRUMENTATION_SCOPES: Dict[str, str] = { + # livekit-agents emits its own span tree (agent_session -> agent_turn -> + # llm_node / tts_node / function_tool) under this scope. + "livekit-agents": InstrumentSet.LIVEKIT.value, +} + + # Default instrument sets # These sets are intentionally independent. Removing an @@ -220,6 +239,7 @@ def __new__(cls, value: Any, origin: Optional[Type[Enum]] = None) -> "Instrument InstrumentSet.GROQ, InstrumentSet.LANGCHAIN, InstrumentSet.LITELLM, + InstrumentSet.LIVEKIT, InstrumentSet.CEREBRAS, InstrumentSet.MISTRALAI, InstrumentSet.OPENAI, @@ -260,6 +280,11 @@ def __new__(cls, value: Any, origin: Optional[Type[Enum]] = None) -> "Instrument ) # Subset of DEFAULT_INSTRUMENTS allowed to produce root-level spans. +# +# InstrumentSet.LIVEKIT must stay listed here: ``agent_session`` is the root of +# every voice trace, so dropping LiveKit from the root allow-list peels that +# span, then recursively peels ``agent_turn`` / ``llm_node`` / ... — the whole +# voice tree — leaving only the provider spans underneath as orphaned roots. DEFAULT_INSTRUMENTS_FOR_ROOT: frozenset[InstrumentSet] = frozenset( { InstrumentSet.ANTHROPIC, @@ -274,6 +299,7 @@ def __new__(cls, value: Any, origin: Optional[Type[Enum]] = None) -> "Instrument InstrumentSet.GROQ, InstrumentSet.LANGCHAIN, InstrumentSet.LITELLM, + InstrumentSet.LIVEKIT, InstrumentSet.CEREBRAS, InstrumentSet.MISTRALAI, InstrumentSet.OPENAI, diff --git a/netra/instrumentation/livekit/__init__.py b/netra/instrumentation/livekit/__init__.py new file mode 100644 index 0000000..b26c849 --- /dev/null +++ b/netra/instrumentation/livekit/__init__.py @@ -0,0 +1,241 @@ +"""LiveKit voice-agent instrumentation for Netra.""" + +import logging +import threading +from typing import Any, Collection, Optional + +from opentelemetry import trace +from opentelemetry.instrumentation.instrumentor import BaseInstrumentor +from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.sdk import trace as trace_sdk +from wrapt import wrap_function_wrapper + +from netra.config import Config, get_active_config +from netra.instrumentation.livekit.audio_processor import AudioSpanProcessor +from netra.instrumentation.livekit.provider_binding import bind_livekit_tracer +from netra.instrumentation.livekit.trace_processor import SpanMappingProcessor +from netra.instrumentation.livekit.wrappers import wrap_aclose, wrap_start + +logger = logging.getLogger(__name__) + +_instruments = ("livekit-agents >= 1.6.0, < 2.0.0",) + +_AGENT_SESSION_MODULE = "livekit.agents.voice.agent_session" +_AGENT_SESSION_CLASS = "AgentSession" +_START_METHOD = "start" +# ``_aclose_impl`` rather than ``aclose``: see ``wrap_aclose`` for why the public +# method covers only one of the five close reasons. +_ACLOSE_METHOD = "_aclose_impl" +# wrapt resolves a dotted attribute path against the module; ``unwrap`` does not +# — see ``_uninstrument``. +_SESSION_START_METHOD = f"{_AGENT_SESSION_CLASS}.{_START_METHOD}" +_SESSION_ACLOSE_METHOD = f"{_AGENT_SESSION_CLASS}.{_ACLOSE_METHOD}" + +# Set on the provider once our processor is attached. OTel has no +# remove_span_processor, so a double registration would silently double every +# mapped attribute write; this flag is the only thing preventing that. Mirrors +# ``_netra_processors_installed`` in netra/tracer.py. +_PROCESSORS_FLAG = "_netra_livekit_processors_installed" + +# Guards against double-wrapping. BaseInstrumentor.is_instrumented_by_opentelemetry +# already prevents a repeat instrument(); this covers a direct _instrument() call. +_session_hook_lock = threading.Lock() +_session_hook_installed = False + + +class NetraLiveKitInstrumentor(BaseInstrumentor): # type: ignore[misc] + """Binds livekit-agents' OTel tracer to Netra's provider and installs session hooks. + + Unlike most Netra instrumentors this one creates no spans of its own on the + trace path — livekit-agents already emits a full span tree + (``agent_session`` → ``agent_turn`` → ``llm_node`` / ``tts_node`` / + ``function_tool``). Our job is to make that tree land in Netra's pipeline, + shield the providers from LiveKit's per-job telemetry teardown, and stamp the + Netra session id on the session root. + + Note on session-id scope: the id is attached for the duration of + ``AgentSession.start`` and inherited by every task LiveKit creates during it, + then detached. Code running in the entrypoint task *after* + ``await session.start(...)`` therefore carries no session id — call + ``Netra.set_session_id()`` for that, which is process-wide by design. + """ + + def instrumentation_dependencies(self) -> Collection[str]: + """Return the package requirement this instrumentor applies to. + + Returns: + The ``livekit-agents`` version range this instrumentation was written + against. + """ + return _instruments + + def _instrument(self, **kwargs: Any) -> None: + """Install the LiveKit integration. + + Each step is isolated so that a LiveKit signature change disables one + feature rather than the whole integration — and never ``Netra.init()``. + + Args: + **kwargs: Optional ``config`` and ``tracer_provider`` overrides. + """ + config: Optional[Config] = kwargs.get("config") or get_active_config() + if config is None: + logger.warning( + "netra.livekit: no active Netra config; LiveKit instrumentation is disabled. " + "Call Netra.init() before instrumenting" + ) + return + + provider = kwargs.get("tracer_provider") or trace.get_tracer_provider() + + try: + bind_livekit_tracer(provider) + except Exception: + logger.exception( + "netra.livekit: could not bind the LiveKit tracer to Netra's provider; " + "LiveKit spans will NOT reach Netra. Session hooks are unaffected" + ) + + try: + self._register_processors(provider, config) + except Exception: + logger.exception("netra.livekit: could not register span processors; lk.* mapping is disabled") + + try: + _install_session_hook() + except Exception: + logger.exception("netra.livekit: could not install the session hook; netra.session_id will be missing") + + self._log_audio_decision(config) + + def _uninstrument(self, **kwargs: Any) -> None: + """Remove the session hooks. + + Does not un-bind the tracer provider or unregister the processors: OTel + offers no ``remove_span_processor`` and LiveKit offers no way to restore a + previous provider. Both are documented limitations; the processors are + inert without LiveKit spans to act on, so leaving them registered is + harmless. + + Args: + **kwargs: Unused. + """ + global _session_hook_installed + + try: + # MUST pass the class, not ``(module, "AgentSession.start")``: unwrap() + # resolves its second argument with a single ``getattr``, which cannot + # walk a dotted path, and it defaults to None rather than raising — so + # the dotted form is a silent no-op that leaves the wrapper installed + # and reports success. + from livekit.agents.voice.agent_session import AgentSession + + # unwrap() is a no-op when the attribute is absent or unwrapped, so the + # aclose hook not having been installed is not an error here. + unwrap(AgentSession, _START_METHOD) + unwrap(AgentSession, _ACLOSE_METHOD) + except (AttributeError, ImportError): + logger.error("netra.livekit: failed to uninstrument %s", _AGENT_SESSION_CLASS) + + with _session_hook_lock: + _session_hook_installed = False + + @staticmethod + def _register_processors(provider: Any, config: Config) -> None: + """Append this integration's span processors to *provider*. + + Called from ``_instrument()``, so it only runs when livekit-agents is + installed and ``InstrumentSet.LIVEKIT`` is enabled — exactly the gate we + want, without ``netra/tracer.py`` having to reimplement it. + + These are appended *after* ``BatchSpanProcessor``; see the module + docstring in ``trace_processor.py`` for the invariant that makes it safe before + adding a third. + + Args: + provider: The tracer provider to register on. + config: The active Netra config, read for the audio-capture decision. + """ + if not isinstance(provider, trace_sdk.TracerProvider): + logger.warning("netra.livekit: provider is not an SDK TracerProvider; span mapping disabled") + return + if getattr(provider, _PROCESSORS_FLAG, False): + return + + provider.add_span_processor(SpanMappingProcessor()) + logger.debug("netra.livekit: registered SpanMappingProcessor") + + if config.audio_capture_enabled: + provider.add_span_processor(AudioSpanProcessor()) + logger.debug("netra.livekit: registered AudioSpanProcessor") + + setattr(provider, _PROCESSORS_FLAG, True) + + @staticmethod + def _log_audio_decision(config: Config) -> None: + """State whether call-audio capture resolved on or off, at INFO. + + An operator must be able to tell from the logs alone whether PCM is + leaving the process, without reading the source. Logs the endpoint *host* + only — never the full URL, never the credential. + + Args: + config: The active Netra config. + """ + try: + if not config.audio_capture_enabled: + logger.info( + "netra.livekit: call audio capture is OFF (no authenticated audio endpoint " + "resolved). Traces are unaffected" + ) + return + + endpoint = config.audio_endpoint() or "" + host = endpoint.split("://")[-1].split("/")[0] + logger.info( + "netra.livekit: call audio capture is ON for both speakers, streaming to host %s", + host, + ) + except Exception: + logger.debug("netra.livekit: could not log the audio capture decision", exc_info=True) + + +def _install_session_hook() -> None: + """Wrap ``AgentSession.start`` and ``AgentSession._aclose_impl``. + + ``start`` carries the session id onto the session root span; ``_aclose_impl`` + runs the per-session teardown that closes the audio sender. Each target is + wrapped in its own ``try``/``except`` so a LiveKit signature change to one + leaves the other working — and so a build without ``_aclose_impl`` still gets + the session id. + + Guarded by a module flag because ``wrapt`` would otherwise double-wrap on a + repeat ``_instrument()`` call. + """ + global _session_hook_installed + + with _session_hook_lock: + if _session_hook_installed: + return + + try: + wrap_function_wrapper(_AGENT_SESSION_MODULE, _SESSION_START_METHOD, wrap_start) + except Exception: + logger.exception( + "netra.livekit: could not wrap AgentSession.start; netra.session_id will be missing " + "from LiveKit spans" + ) + return + + try: + wrap_function_wrapper(_AGENT_SESSION_MODULE, _SESSION_ACLOSE_METHOD, wrap_aclose) + except Exception: + logger.exception( + "netra.livekit: could not wrap AgentSession._aclose_impl; per-session teardown " + "(audio sender close) is disabled" + ) + + _session_hook_installed = True + + +__all__ = ["NetraLiveKitInstrumentor"] diff --git a/netra/instrumentation/livekit/audio_capture.py b/netra/instrumentation/livekit/audio_capture.py new file mode 100644 index 0000000..d22f419 --- /dev/null +++ b/netra/instrumentation/livekit/audio_capture.py @@ -0,0 +1,808 @@ +"""Captures a LiveKit session's audio and attributes it to speaking spans. + +:class:`SessionAudioCoordinator` sits between livekit-agents' audio I/O and +:class:`~netra.instrumentation.livekit.audio_sender.AudioChunkSender`. It owns +two things: + +* **where a frame belongs** — the ``user_speaking``/``agent_speaking`` span open + at the moment of capture, pushed in by + :class:`~netra.instrumentation.livekit.audio_processor.AudioSpanProcessor`. + Frames captured between turns are still sent, attributed to the call but to no + span; +* **what the caller actually heard** — when a caller interrupts the agent, + LiveKit discards the un-played tail of the utterance, so the coordinator stops + forwarding agent frames and reports the playback position so the recorded + audio can be trimmed to match. + +Nothing here may change the behaviour of the user's agent: every patched method +forwards to the original whether or not our own work succeeded. +""" + +from __future__ import annotations + +import asyncio +import functools +import logging +import threading +import time +from concurrent.futures import TimeoutError as FuturesTimeoutError +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple + +from netra.instrumentation.livekit.audio_sender import AudioChunkSender +from netra.instrumentation.livekit.audio_types import ( + CREDENTIAL_HEADER_NAMES, + NETRA_AUDIO_CIRCUIT_TRIPPED, + NETRA_AUDIO_DROPPED_FRAMES, + NETRA_AUDIO_ERRORS, + NETRA_AUDIO_SENT_BYTES, + NETRA_AUDIO_SENT_CHUNKS, + SpeakerRole, +) + +if TYPE_CHECKING: + from livekit.agents import AgentSession + from livekit.rtc import AudioFrame + + from netra.config import Config + +logger = logging.getLogger(__name__) + +# Nominal PCM bytes in one captured frame: 20ms of 24kHz mono 16-bit audio, what +# livekit-agents delivers by default. Used only to turn the byte budget +# ``NETRA_AUDIO_BUFFER_BYTES`` into the frame count the queue is actually bounded +# by — a different frame size simply makes the queue hold proportionally more or +# less audio than the budget names. +_NOMINAL_FRAME_BYTES = 960 + +_MILLISECONDS_PER_SECOND = 1000 + +# Slack added to the wait in ``_close_from_outside`` on top of the drain budget it +# hands the coordinator, so the coordinator's own deadline is the one that fires. +_TEARDOWN_GRACE_SECONDS = 1.0 + + +@dataclass(frozen=True) +class _ActiveSpeech: + """The speaking span currently open for one speaker. + + Attributes: + span_id: Hex id of the open ``*_speaking`` span. + trace_id: Hex trace id of the call the span belongs to. + """ + + span_id: str + trace_id: str + + +class SessionAudioCoordinator: + """Routes one AgentSession's audio frames to the sender, tagged by span. + + Lifecycle: + + 1. :meth:`attach` patches the session's audio I/O, once ``start()`` has + returned and ``session.input``/``session.output`` exist; + 2. :class:`AudioSpanProcessor` calls :meth:`on_speaking_start` / + :meth:`on_speaking_end` as LiveKit opens and closes speaking spans; + 3. the patched I/O calls :meth:`on_frame` for every frame in either + direction; + 4. :meth:`aclose` closes any span still recording and shuts the sender down. + + Confined to the agent's event loop, like the sender it feeds. + """ + + def __init__(self, *, sender: Optional[AudioChunkSender] = None) -> None: + """Bind the coordinator to a sender. + + Args: + sender: Where frames are handed off. ``None`` makes the coordinator + inert, which is what the span processor's callbacks expect when + audio capture is off. + """ + self._sender = sender + + self._active_speech: Dict[SpeakerRole, Optional[_ActiveSpeech]] = {role: None for role in SpeakerRole} + self._session_trace_id = "" + + # The agent span most recently opened, kept after it closes: LiveKit + # routinely ends the ``agent_speaking`` span *before* it reports the + # interrupt that cut it short, so the id would otherwise be gone by the + # time there is something to report about it. + self._last_agent_span_id = "" + self._is_agent_interrupted = False + self._interrupted_agent_span_id = "" + + # -- attachment --------------------------------------------------------- + + def attach(self, session: "AgentSession") -> None: + """Patch *session*'s audio input and output to feed this coordinator. + + Must run after ``session.start()``: before that, ``session.input`` and + ``session.output`` are not yet populated. + + Args: + session: The started LiveKit ``AgentSession``. + """ + self._session_trace_id = _current_trace_id() + self._patch_audio_input(session) + self._patch_audio_output(session) + + # -- span callbacks ----------------------------------------------------- + + def on_speaking_start(self, role: SpeakerRole, *, trace_id: str, span_id: str) -> None: + """Attribute subsequent frames from *role* to a newly opened span. + + Args: + role: The speaker whose span opened. + trace_id: Hex trace id of the span. + span_id: Hex id of the span. + """ + self._active_speech[role] = _ActiveSpeech(span_id=span_id, trace_id=trace_id) + if role is SpeakerRole.AGENT: + self._last_agent_span_id = span_id + self._is_agent_interrupted = False + self._interrupted_agent_span_id = "" + logger.debug("netra.audio: %s speaking started — span_id=%s", role.value, span_id) + + def on_speaking_end(self, role: SpeakerRole) -> None: + """Close the recording for *role*'s open span. + + An interrupted agent span is left for :meth:`on_playback_finished` to + finalize: only the playback report says how much of the utterance was + heard, and finalizing here would fix the recording at its full length. + + Args: + role: The speaker whose span closed. + """ + active = self._active_speech[role] + self._active_speech[role] = None + if active is None: + return + if role is SpeakerRole.AGENT and self._is_agent_interrupted: + return + if self._sender is not None: + self._sender.mark_audio_end(role=role, span_id=active.span_id) + + # -- frame callbacks ---------------------------------------------------- + + def on_frame(self, role: SpeakerRole, frame: "AudioFrame") -> None: + """Hand one captured frame to the sender. + + Stamps the capture time here, the earliest point the frame is seen, so + the timeline is not skewed by time spent queued. + + Args: + role: The speaker the frame came from. + frame: The frame LiveKit just captured. + """ + if self._sender is None: + return + if role is SpeakerRole.AGENT and self._is_agent_interrupted: + # Produced after the caller cut in, so never played out. + return + + active = self._active_speech[role] + self._sender.enqueue( + frame, + role=role, + span_id=active.span_id if active is not None else "", + trace_id=(active.trace_id if active is not None else "") or self._session_trace_id, + timestamp_ns=time.time_ns(), + ) + + # -- interrupt callbacks ------------------------------------------------ + + def on_output_buffer_cleared(self) -> None: + """Note that LiveKit dropped the agent's queued audio — a caller interrupt. + + Stops further agent frames from being forwarded and remembers which span + was cut, for :meth:`on_playback_finished` to trim. + """ + active = self._active_speech[SpeakerRole.AGENT] + self._is_agent_interrupted = True + self._interrupted_agent_span_id = active.span_id if active is not None else self._last_agent_span_id + logger.debug( + "netra.audio: agent audio buffer cleared — utterance interrupted (span_id=%s)", + self._interrupted_agent_span_id, + ) + + def on_playback_finished(self, event: Any) -> None: + """Trim an interrupted utterance to the audio that was played out. + + Args: + event: LiveKit's ``playback_finished`` event. Only an event flagged + ``interrupted`` is acted on; a normal end of playback needs no + correction. + """ + if not getattr(event, "interrupted", False): + return + span_id = self._interrupted_agent_span_id + if not span_id or self._sender is None: + return + + playback_ms = int(getattr(event, "playback_position", 0.0) * _MILLISECONDS_PER_SECOND) + self._sender.interrupt_agent_span(span_id=span_id, playback_ms=playback_ms) + logger.debug( + "netra.audio: interrupted playback finished — span_id=%s heard=%dms", + span_id, + playback_ms, + ) + + # -- teardown ----------------------------------------------------------- + + def close(self) -> None: + """Close every span still recording, without touching the sender. + + Separate from :meth:`aclose` because the session span has to be stamped + with the sender's final statistics, which means the two teardown halves + run at different points. + """ + for role in SpeakerRole: + self.on_speaking_end(role) + + async def aclose(self, *, drain_timeout_seconds: Optional[float] = None) -> None: + """Close the open recordings and shut the sender down. + + Args: + drain_timeout_seconds: Total budget for the sender's drain. ``None`` + leaves the sender's own default in place, which is what the normal + per-session teardown wants; ``Netra.shutdown()`` passes the budget + it is willing to wait so the two cannot disagree. + """ + self.close() + if self._sender is None: + return + if drain_timeout_seconds is None: + await self._sender.end_session() + else: + await self._sender.end_session(drain_timeout_seconds=drain_timeout_seconds) + + @property + def sender(self) -> Optional[AudioChunkSender]: + """The sender this coordinator feeds, if audio capture is on.""" + return self._sender + + # -- audio input -------------------------------------------------------- + + def _patch_audio_input(self, session: "AgentSession") -> None: + """Intercept the caller's audio by proxying the session's input stream. + + Args: + session: The started LiveKit ``AgentSession``. + """ + session_input = getattr(session, "input", None) + audio_input = getattr(session_input, "audio", None) + if session_input is None or audio_input is None: + logger.warning("netra.audio: session.input.audio is unavailable — caller audio is not captured") + return + + leaf = _leaf_audio_source(audio_input) + proxy = _AudioInputProxy(leaf, self) + + for holder, attribute in _proxy_mount_points(session_input, audio_input, leaf): + if _try_set(holder, attribute, proxy): + logger.debug("netra.audio: caller audio proxied at %s.%s", type(holder).__name__, attribute) + return + + # Every mount point is read-only, so there is nowhere to insert a proxy; + # patch the iteration protocol on the leaf itself instead. + _patch_anext(leaf, self) + + # -- audio output ------------------------------------------------------- + + def _patch_audio_output(self, session: "AgentSession") -> None: + """Intercept the agent's audio and the events describing its playback. + + Args: + session: The started LiveKit ``AgentSession``. + """ + audio_output = getattr(getattr(session, "output", None), "audio", None) + if audio_output is None: + logger.warning("netra.audio: session.output.audio is unavailable — agent audio is not captured") + return + + self._patch_capture_frame(audio_output) + self._patch_clear_buffer(audio_output) + self._subscribe_to_playback_finished(audio_output) + + def _patch_capture_frame(self, audio_output: Any) -> None: + """Wrap ``capture_frame`` so every outgoing frame is seen. + + Args: + audio_output: LiveKit's agent audio output. + """ + original = audio_output.capture_frame + + @functools.wraps(original) + async def capture_frame(frame: "AudioFrame") -> Any: + _run_hook_safely(lambda: self.on_frame(SpeakerRole.AGENT, frame), "agent frame") + return await original(frame) + + audio_output.capture_frame = capture_frame + logger.debug("netra.audio: wrapped agent capture_frame") + + def _patch_clear_buffer(self, audio_output: Any) -> None: + """Wrap ``clear_buffer``, LiveKit's signal that the caller interrupted. + + Args: + audio_output: LiveKit's agent audio output. + """ + original = getattr(audio_output, "clear_buffer", None) + if not callable(original): + logger.debug("netra.audio: no clear_buffer on the audio output — interrupts are not detected") + return + + @functools.wraps(original) + def clear_buffer() -> Any: + _run_hook_safely(self.on_output_buffer_cleared, "clear_buffer") + return original() + + audio_output.clear_buffer = clear_buffer + logger.debug("netra.audio: wrapped clear_buffer for interrupt detection") + + def _subscribe_to_playback_finished(self, audio_output: Any) -> None: + """Listen for playback reports, which say how much audio was heard. + + Args: + audio_output: LiveKit's agent audio output. + """ + subscribe = getattr(audio_output, "on", None) + if not callable(subscribe): + logger.debug("netra.audio: audio output is not an event emitter — interrupts are not trimmed") + return + try: + subscribe("playback_finished", self.on_playback_finished) + except (TypeError, ValueError): + logger.debug("netra.audio: could not subscribe to playback_finished", exc_info=True) + return + logger.debug("netra.audio: subscribed to playback_finished") + + +# --------------------------------------------------------------------------- +# Audio input plumbing +# --------------------------------------------------------------------------- + + +def _run_hook_safely(action: Callable[[], None], description: str) -> None: + """Run one of our own hooks without letting it reach the user's agent. + + The one place this package swallows an exception, and deliberately: these + hooks run inline in the agent's audio path, where a raise would drop the + caller's audio or kill the playout task. The failure is logged, and losing + observability is always preferable to breaking the call. + + Args: + action: The hook to run. + description: What it was doing, for the log line. + """ + try: + action() + except Exception: + logger.debug("netra.audio: %s hook failed", description, exc_info=True) + + +class _AudioInputProxy: + """Transparent proxy over an async audio iterator, tapping each frame. + + ``__aiter__``/``__anext__`` are defined on the class rather than the + instance because ``async for`` resolves them on the *type*: an instance + attribute would simply be ignored. + """ + + def __init__(self, source: Any, coordinator: SessionAudioCoordinator) -> None: + """Wrap *source*, reporting each frame it yields to *coordinator*. + + Args: + source: The audio iterator being proxied. + coordinator: Where captured frames are reported. + """ + self._source = source + self._coordinator = coordinator + + def __aiter__(self) -> "_AudioInputProxy": + """Return self; the proxy is its own iterator.""" + return self + + async def __anext__(self) -> "AudioFrame": + """Yield the next frame from the wrapped source, tapping it on the way. + + Returns: + The frame, untouched. + """ + frame: "AudioFrame" = await self._source.__anext__() + _run_hook_safely(lambda: self._coordinator.on_frame(SpeakerRole.USER, frame), "caller frame") + return frame + + def __getattr__(self, name: str) -> Any: + """Forward every other attribute to the wrapped source. + + Args: + name: The attribute being looked up. + + Returns: + The wrapped source's attribute. + """ + return getattr(self._source, name) + + +def _leaf_audio_source(audio_input: Any) -> Any: + """Follow the ``.source`` chain to the object actually producing frames. + + LiveKit stacks audio streams (resamplers, buffers) each holding the next in + ``.source``. Tapping the innermost one captures the caller's audio before + any of that processing. + + Args: + audio_input: The outermost audio input. + + Returns: + The innermost source, which may be *audio_input* itself. + """ + current = audio_input + while getattr(current, "source", None) is not None: + current = current.source + return current + + +def _proxy_mount_points(session_input: Any, audio_input: Any, leaf: Any) -> List[Tuple[Any, str]]: + """Return the places a proxy over *leaf* could be installed, best first. + + Args: + session_input: The session's input container. + audio_input: The outermost audio input. + leaf: The innermost audio source. + + Returns: + ``(holder, attribute)`` pairs to try assigning the proxy to. + """ + if leaf is audio_input: + return [(session_input, "audio")] + + parent = _parent_of(audio_input, leaf) + return [(parent, "source")] if parent is not None else [] + + +def _parent_of(audio_input: Any, leaf: Any) -> Optional[Any]: + """Return the object whose ``.source`` is *leaf*. + + Args: + audio_input: The outermost audio input to search from. + leaf: The innermost audio source. + + Returns: + The holder of *leaf*, or ``None`` when *leaf* is not in the chain. + """ + current = audio_input + while current is not None: + if getattr(current, "source", None) is leaf: + return current + current = getattr(current, "source", None) + return None + + +def _try_set(holder: Any, attribute: str, value: Any) -> bool: + """Assign *attribute* on *holder*, reporting whether it took. + + Args: + holder: The object to assign on. + attribute: The attribute name. + value: The value to assign. + + Returns: + True on success; False when the attribute is read-only or slotted. + """ + try: + setattr(holder, attribute, value) + except (AttributeError, TypeError): + return False + return True + + +def _patch_anext(leaf: Any, coordinator: SessionAudioCoordinator) -> None: + """Tap frames by replacing ``__anext__`` on the leaf instance itself. + + Last resort: it only works for code that calls ``leaf.__anext__()`` + explicitly, since ``async for`` looks the method up on the type. + + Args: + leaf: The innermost audio source. + coordinator: Where captured frames are reported. + """ + original = leaf.__anext__ + + @functools.wraps(original) + async def traced_anext() -> "AudioFrame": + frame = await original() + _run_hook_safely(lambda: coordinator.on_frame(SpeakerRole.USER, frame), "caller frame") + return frame + + if not _try_set(leaf, "__anext__", traced_anext): + logger.warning("netra.audio: could not intercept the caller audio stream — caller audio is not captured") + return + logger.debug("netra.audio: fell back to patching __anext__ on the audio source") + + +def _current_trace_id() -> str: + """Return the active span's trace id as hex, or ``""`` when there is none. + + Frames captured between speaking spans still belong to the call, so they are + attributed to this trace rather than dropped. + """ + from opentelemetry import context, trace + + span_context = trace.get_current_span(context.get_current()).get_span_context() + if span_context is None or not span_context.is_valid: + return "" + return format(span_context.trace_id, "032x") + + +# --------------------------------------------------------------------------- +# Per-session registry +# --------------------------------------------------------------------------- + + +class AudioCoordinatorRegistry: + """Finds the coordinator for a call, given the trace its spans belong to. + + :class:`AudioSpanProcessor` is registered once for the process but speaking + spans arrive for every concurrent call, so the span's trace id is what says + which call's audio a span delimits. + + Locked rather than loop-confined. Most traffic is on the agent's event loop — + registration from the session wrapper, lookups from span callbacks — but + ``Netra.shutdown()`` reaches :meth:`pop_all` from whichever thread called it, + and that has to be atomic against a concurrent :meth:`register` or a call's + coordinator is dropped on the floor with its audio still queued. Contention is + a handful of operations per call, so a plain lock costs nothing measurable. + """ + + def __init__(self) -> None: + """Start with no calls registered.""" + self._by_trace_id: Dict[int, SessionAudioCoordinator] = {} + self._lock = threading.Lock() + + def register(self, trace_id: int, coordinator: SessionAudioCoordinator) -> None: + """Record the coordinator capturing audio for a call. + + Args: + trace_id: The ``agent_session`` span's trace id. + coordinator: The call's coordinator. + """ + with self._lock: + self._by_trace_id[trace_id] = coordinator + + def get(self, trace_id: int) -> Optional[SessionAudioCoordinator]: + """Return the coordinator for a call, or ``None`` if it is not capturing. + + Args: + trace_id: The trace id off a speaking span. + + Returns: + The call's coordinator, if one is registered. + """ + with self._lock: + return self._by_trace_id.get(trace_id) + + def unregister(self, trace_id: int) -> Optional[SessionAudioCoordinator]: + """Remove and return a call's coordinator. Idempotent. + + Args: + trace_id: The ``agent_session`` span's trace id. + + Returns: + The coordinator that was registered, if any. + """ + with self._lock: + return self._by_trace_id.pop(trace_id, None) + + def pop_all(self) -> List[SessionAudioCoordinator]: + """Remove and return every registered coordinator. + + Used by ``Netra.shutdown()`` as a backstop for calls whose session never + closed cleanly. Atomic, so a call registering concurrently is either + returned here or left registered — never lost between the read and the + clear. + + Returns: + The coordinators that were registered. + """ + with self._lock: + coordinators = list(self._by_trace_id.values()) + self._by_trace_id.clear() + return coordinators + + +audio_coordinators = AudioCoordinatorRegistry() + + +# --------------------------------------------------------------------------- +# Session wiring +# --------------------------------------------------------------------------- + + +def build_audio_sender(config: "Config", session_id: str) -> Optional[AudioChunkSender]: + """Construct the sender for one call from the active Netra config. + + Args: + config: The active Netra config. + session_id: The Netra session id for this call. + + Returns: + A configured, unstarted sender, or ``None`` when no audio endpoint + resolves — which is the single gate on audio capture. + """ + url = config.audio_endpoint() + if not url: + return None + + credential_headers = { + name: value for name, value in (config.headers or {}).items() if name.lower() in CREDENTIAL_HEADER_NAMES + } + return AudioChunkSender( + url=url, + session_id=session_id, + api_key=config.api_key or "", + auth_headers=credential_headers, + batch_interval_seconds=config.audio_batch_interval_ms / _MILLISECONDS_PER_SECOND, + flush_at_bytes=config.audio_batch_bytes, + max_request_bytes=config.audio_max_request_bytes, + max_queue_frames=max(1, config.audio_buffer_bytes // _NOMINAL_FRAME_BYTES), + ) + + +async def start_audio_capture(session: Any, *, config: "Config", session_id: str, trace_id: int) -> None: + """Begin capturing a started session's call audio. + + Isolated from the caller by design: audio capture failing must never make + ``AgentSession.start()`` fail, and traces are unaffected either way. + + Args: + session: The started LiveKit ``AgentSession``. + config: The active Netra config. + session_id: The Netra session id for this call. + trace_id: The ``agent_session`` span's trace id, under which the + coordinator is registered for the span processor to find. + """ + try: + sender = build_audio_sender(config, session_id) + if sender is None: + return + + coordinator = SessionAudioCoordinator(sender=sender) + await sender.start() + + # Registered before attaching, not after: from here on the sender owns a + # background task and an HTTP client, and the registry is the only handle + # anything has for closing them. ``attach`` patches third-party objects + # that may refuse assignment, so it is exactly the step that can raise — + # and a raise between start() and register() would strand both resources + # for the life of the process. ``attach`` does not need the registry. + audio_coordinators.register(trace_id, coordinator) + try: + coordinator.attach(session) + except Exception: + await stop_audio_capture(trace_id) + raise + logger.debug("netra.audio: capture attached for trace_id=%032x", trace_id) + except Exception: + logger.warning("netra.livekit: audio capture setup failed; the call is traced without audio", exc_info=True) + + +async def stop_audio_capture(trace_id: int, session_span: Optional[Any] = None) -> None: + """Stop capturing a call's audio and record what was delivered. + + Idempotent: a call whose coordinator has already been removed does nothing. + + Args: + trace_id: The ``agent_session`` span's trace id. + session_span: The still-recording ``agent_session`` span, stamped with + the delivery statistics when given. + """ + coordinator = audio_coordinators.unregister(trace_id) + if coordinator is None: + return + + try: + await coordinator.aclose() + except Exception: + logger.warning("netra.audio: audio capture teardown failed", exc_info=True) + + sender = coordinator.sender + if session_span is not None and sender is not None: + _stamp_audio_stats(session_span, sender) + + +def close_all_audio_capture(timeout_seconds: float = 5.0) -> None: + """Shut down every call still capturing audio. Backstop for ``Netra.shutdown()``. + + A sender's queue and task belong to the event loop its call was running on, + so it cannot simply be awaited from wherever shutdown happens to be called. + Each one is driven through its own loop instead — and a call whose loop is + already gone is reported rather than silently skipped, because its unsent + audio is genuinely lost. + + Args: + timeout_seconds: How long to wait for one call's audio to drain when + shutting it down from outside its event loop. Passed down as the + sender's own drain budget too, so the inner deadline expires first and + a timeout here means the audio really could not be delivered rather + than that the two limits were set inconsistently. + """ + coordinators = audio_coordinators.pop_all() + if not coordinators: + return + + logger.info("netra.audio: shutting down %d call(s) still capturing audio", len(coordinators)) + try: + current_loop: Optional[asyncio.AbstractEventLoop] = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + for coordinator in coordinators: + _close_from_outside(coordinator, current_loop, timeout_seconds) + + +def _close_from_outside( + coordinator: SessionAudioCoordinator, + current_loop: Optional["asyncio.AbstractEventLoop"], + timeout_seconds: float, +) -> None: + """Drive one coordinator's teardown from whichever loop is available. + + Args: + coordinator: The coordinator to shut down. + current_loop: The loop the caller is running on, if any. + timeout_seconds: How long to wait when driving another loop, and the drain + budget handed to the coordinator either way. + """ + sender = coordinator.sender + target_loop = sender.loop if sender is not None else None + + if target_loop is None or target_loop.is_closed(): + logger.warning("netra.audio: a call's event loop is gone; its unsent audio is lost") + return + + if target_loop is current_loop: + # Cannot block the loop we are on, so this is scheduled and not awaited: + # whether it finishes depends on the caller keeping the loop alive, which + # a synchronous shutdown() cannot promise. Said plainly rather than left + # looking like a completed teardown. + target_loop.create_task(coordinator.aclose(drain_timeout_seconds=timeout_seconds)) + logger.warning( + "netra.audio: shutdown was called from a call's own event loop; its drain is scheduled " + "but cannot be awaited. Await AgentSession.aclose() before Netra.shutdown() to be sure " + "the audio is delivered" + ) + return + + future = asyncio.run_coroutine_threadsafe(coordinator.aclose(drain_timeout_seconds=timeout_seconds), target_loop) + try: + # A shade past the inner budget, so the coordinator's own deadline is what + # gives up and it still gets to log its statistics. + future.result(timeout=timeout_seconds + _TEARDOWN_GRACE_SECONDS) + except FuturesTimeoutError: + logger.warning("netra.audio: a call did not finish sending within %.0fs", timeout_seconds) + except Exception: + logger.warning("netra.audio: a call failed to shut down cleanly", exc_info=True) + + +def _stamp_audio_stats(session_span: Any, sender: AudioChunkSender) -> None: + """Record the call's audio delivery counters on its session span. + + Args: + session_span: The still-recording ``agent_session`` span. + sender: The sender whose statistics to record. + """ + stats = sender.stats + try: + session_span.set_attributes( + { + NETRA_AUDIO_SENT_BYTES: stats.bytes_sent, + NETRA_AUDIO_SENT_CHUNKS: stats.chunks_sent, + NETRA_AUDIO_DROPPED_FRAMES: stats.frames_dropped, + NETRA_AUDIO_ERRORS: stats.errors, + NETRA_AUDIO_CIRCUIT_TRIPPED: stats.circuit_tripped, + } + ) + except Exception: + logger.debug("netra.audio: could not stamp audio stats on the session span", exc_info=True) diff --git a/netra/instrumentation/livekit/audio_processor.py b/netra/instrumentation/livekit/audio_processor.py new file mode 100644 index 0000000..d410689 --- /dev/null +++ b/netra/instrumentation/livekit/audio_processor.py @@ -0,0 +1,120 @@ +"""Tells the audio coordinator which turn is being spoken, as spans open and close. + +LiveKit brackets each run of speech in a ``user_speaking`` or ``agent_speaking`` +span. This processor is the only thing that sees those spans start and end, so +it is what lets a frame captured milliseconds later be filed under the turn it +belongs to. + +Registered once for the process, while coordinators are per call — hence the +lookup by the span's trace id in +:data:`~netra.instrumentation.livekit.audio_capture.audio_coordinators`. +""" + +from __future__ import annotations + +import logging +from typing import NamedTuple, Optional, Union + +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor +from opentelemetry.trace import SpanContext + +from netra.instrumentation.livekit.audio_capture import SessionAudioCoordinator, audio_coordinators +from netra.instrumentation.livekit.audio_types import SPEAKING_SPAN_ROLES, SpeakerRole + +logger = logging.getLogger(__name__) + +_TRACE_ID_HEX_DIGITS = "032x" +_SPAN_ID_HEX_DIGITS = "016x" + + +class _SpeakingSpan(NamedTuple): + """A span that delimits speech, resolved to the call it belongs to. + + Attributes: + role: The speaker the span delimits. + coordinator: The coordinator capturing that call's audio. + span_context: The span's own context, for its trace and span ids. + """ + + role: SpeakerRole + coordinator: SessionAudioCoordinator + span_context: SpanContext + + +class AudioSpanProcessor(SpanProcessor): # type: ignore[misc] + """Opens and closes an audio recording alongside each speaking span.""" + + def on_start(self, span: Span, parent_context: Optional[otel_context.Context] = None) -> None: + """Start attributing this speaker's audio to the span that just opened. + + Args: + span: The span that was started. + parent_context: The parent context (unused). + """ + speaking = _resolve_speaking_span(span) + if speaking is None: + return + + speaking.coordinator.on_speaking_start( + speaking.role, + trace_id=format(speaking.span_context.trace_id, _TRACE_ID_HEX_DIGITS), + span_id=format(speaking.span_context.span_id, _SPAN_ID_HEX_DIGITS), + ) + + def on_end(self, span: ReadableSpan) -> None: + """Close the recording for the speaking span that just ended. + + Args: + span: The span that has ended. + """ + speaking = _resolve_speaking_span(span) + if speaking is None: + return + + speaking.coordinator.on_speaking_end(speaking.role) + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """No-op flush; this processor holds nothing pending. + + Args: + timeout_millis: Maximum time to wait (unused). + + Returns: + Always True. + """ + return True + + def shutdown(self) -> None: + """No-op shutdown; coordinator teardown belongs to the session wrapper.""" + + +def _resolve_speaking_span(span: Union[Span, ReadableSpan]) -> Optional[_SpeakingSpan]: + """Identify a speaking span and the call whose audio it delimits. + + Never raises: this runs on every span the process produces, so a failure + here would be a failure of the user's tracing, not just of audio capture. + + Args: + span: The span that started or ended. + + Returns: + The resolved speaking span, or ``None`` when *span* does not delimit + speech or its call is not capturing audio — the common case by far. + """ + try: + role = SPEAKING_SPAN_ROLES.get(span.name or "") + if role is None: + return None + + span_context = span.get_span_context() + if span_context is None or not span_context.is_valid: + return None + + coordinator = audio_coordinators.get(span_context.trace_id) + if coordinator is None: + return None + return _SpeakingSpan(role=role, coordinator=coordinator, span_context=span_context) + except Exception: + logger.debug("netra.audio: could not resolve a speaking span", exc_info=True) + return None diff --git a/netra/instrumentation/livekit/audio_sender.py b/netra/instrumentation/livekit/audio_sender.py new file mode 100644 index 0000000..e4d4533 --- /dev/null +++ b/netra/instrumentation/livekit/audio_sender.py @@ -0,0 +1,1200 @@ +"""Streams captured call audio to the Netra audio-ingest endpoint. + +:class:`SessionAudioCoordinator` hands frames to :meth:`AudioChunkSender.enqueue` +from the agent's event loop; a background task batches them and POSTs raw PCM +with the metadata in ``x-audio-*`` headers. Enqueueing never blocks and never +raises into the agent: a full queue drops the frame and a failing endpoint trips +a circuit breaker for the rest of the call. + +Three request shapes reach the endpoint, all defined in ``audio_types``: + +**Span chunk** — audio captured while a ``user_speaking``/``agent_speaking`` span +was open. Body is raw PCM; carries ``x-audio-span-id`` and a per-span +``x-audio-seq``, and the final one carries ``x-audio-last`` (plus +``x-audio-heard-ms`` when the utterance was interrupted). + +**Noise chunk** — audio captured between speaking spans. Same shape without the +span headers, so it can be laid out on the call timeline but belongs to no turn. + +**Session end** — one bodyless request carrying ``x-audio-session-last``. +""" + +from __future__ import annotations + +import asyncio +import logging +import random +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +import httpx +from opentelemetry import context as otel_context + +from netra.instrumentation.livekit.audio_types import ( + CONTENT_TYPE_PCM, + DEFAULT_CHANNEL_COUNT, + DEFAULT_SAMPLE_RATE_HZ, + HEADER_API_KEY, + HEADER_BIT_DEPTH, + HEADER_CHANNELS, + HEADER_CONTENT_TYPE, + HEADER_HEARD_MS, + HEADER_LAST_CHUNK, + HEADER_ROLE, + HEADER_SAMPLE_RATE, + HEADER_SEQUENCE, + HEADER_SESSION_ID, + HEADER_SESSION_LAST, + HEADER_SPAN_ID, + HEADER_START_MS, + HEADER_TRACE_ID, + HEADER_VALUE_TRUE, + PCM_BIT_DEPTH, + SpeakerRole, + pcm_byte_offset_at, +) + +if TYPE_CHECKING: + from livekit.rtc import AudioFrame + +logger = logging.getLogger(__name__) + +# Defaults for the knobs ``Config`` does not resolve. Every other limit reaches +# the sender from ``Config`` — see ``audio_capture.start_audio_capture``. +DEFAULT_BATCH_INTERVAL_SECONDS = 0.5 +DEFAULT_MAX_BATCH_FRAMES = 200 +DEFAULT_FLUSH_AT_BYTES = 32768 +DEFAULT_MAX_REQUEST_BYTES = 262144 + +_HTTP_TIMEOUT_SECONDS = 5.0 + +# Attempts per chunk, total. A chunk POST is safe to repeat: the endpoint keys on +# (session, span, sequence) and the sequence only advances once a chunk has been +# accepted, so a retry re-sends identical bytes under an identical key. +_POST_ATTEMPTS = 2 +_RETRY_BASE_DELAY_SECONDS = 0.05 + +# Consecutive failed chunks after which the rest of the call is abandoned. Audio +# is best-effort: a backend that has been failing this long will not be fixed by +# the next frame, and retrying every 20ms frame for a 10-minute call is worse for +# the agent than sending nothing. +_MAX_CONSECUTIVE_FAILURES = 5 + +# How long ``end_session`` spends draining, in total, before giving up. It runs +# inline in ``AgentSession._aclose_impl``, so this delays the caller's own session +# teardown — a few seconds of best-effort audio is worth that, half a minute is +# not. A backend too slow to drain inside it has usually tripped the circuit +# already. +_DEFAULT_DRAIN_TIMEOUT_SECONDS = 5.0 + +_HTTP_STATUS_BAD_REQUEST = 400 +_UNAUTHENTICATED_STATUSES = frozenset({401, 403}) + + +# --------------------------------------------------------------------------- +# Queue messages +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _FrameMessage: + """One captured audio frame awaiting batching.""" + + pcm_bytes: bytes + role: SpeakerRole + span_id: str + trace_id: str + sample_rate_hz: int + channel_count: int + timestamp_ns: int + + +@dataclass(frozen=True) +class _SpanEndMarker: + """A speaking span closed normally; its recording is complete.""" + + role: SpeakerRole + span_id: str + + +@dataclass(frozen=True) +class _SpanInterruptMarker: + """An agent utterance was cut off after *playback_ms* of audible playback.""" + + span_id: str + playback_ms: int + + +@dataclass(frozen=True) +class _SessionEndMarker: + """The session is closing; drain everything and stop the loop.""" + + +_QueueMessage = Union[_FrameMessage, _SpanEndMarker, _SpanInterruptMarker, _SessionEndMarker] + + +# --------------------------------------------------------------------------- +# Sender state +# --------------------------------------------------------------------------- + + +@dataclass +class _PendingBatch: + """Frames of one speaker accumulating until a flush condition is met. + + Reused in place across flushes rather than reallocated, so the send loop can + hold one per :class:`SpeakerRole` in a plain dict with no rebinding. + """ + + role: SpeakerRole + span_id: str = "" + trace_id: str = "" + sample_rate_hz: int = 0 + channel_count: int = 0 + start_ms: int = 0 + frame_count: int = 0 + byte_count: int = 0 + _pcm_parts: List[bytes] = field(default_factory=list) + + @property + def is_empty(self) -> bool: + """Whether the batch holds no frames yet.""" + return self.frame_count == 0 + + @property + def pcm_bytes(self) -> bytes: + """The accumulated frames as one contiguous PCM buffer.""" + return b"".join(self._pcm_parts) + + def frames_within(self, byte_count: int) -> int: + """Estimate how many accumulated frames fit in the first *byte_count* bytes. + + Used only for the ``frames_sent`` statistic when an interrupt trims the + batch: pro-rating by the mean frame size is accurate whenever the frames + are uniformly sized, which is every case livekit-agents produces, and is + never worse than reporting the untrimmed count. + + Args: + byte_count: Length of the prefix actually being sent. + + Returns: + The frame count attributable to that prefix. + """ + if self.byte_count <= 0: + return 0 + capped = min(max(byte_count, 0), self.byte_count) + return round(self.frame_count * capped / self.byte_count) + + def add(self, frame: _FrameMessage) -> None: + """Append *frame*, adopting its span and format if this is the first one. + + Args: + frame: The frame to accumulate. + """ + if self.is_empty: + self.span_id = frame.span_id + self.trace_id = frame.trace_id + self.sample_rate_hz = frame.sample_rate_hz + self.channel_count = frame.channel_count + self.start_ms = frame.timestamp_ns // 1_000_000 + self._pcm_parts.append(frame.pcm_bytes) + self.frame_count += 1 + self.byte_count += len(frame.pcm_bytes) + + def clear(self) -> None: + """Discard the accumulated frames, keeping the batch's speaker role.""" + self.span_id = "" + self.trace_id = "" + self.sample_rate_hz = 0 + self.channel_count = 0 + self.start_ms = 0 + self.frame_count = 0 + self.byte_count = 0 + self._pcm_parts.clear() + + +@dataclass +class _SpanAudioState: + """Everything the sender tracks about one speaking span's audio stream. + + One record per span replaces the parallel per-span dictionaries this class + used to keep, so a span's sequence number, byte position and terminal state + cannot disagree about which spans exist. + + Attributes: + role: The speaker the span belongs to. + trace_id: Hex trace id, so a terminator posted after the batch holding the + span is gone can still be attributed. + next_sequence: The number the span's next chunk will carry. + bytes_consumed: How many PCM bytes of this span have already left the + pending batch — a *position* in the span's stream, so it counts a + chunk the sender gave up on as well as an accepted one. Trimming an + interrupted utterance measures against this; counting bytes actually + delivered here would make the trim offset drift by whatever was lost. + is_finalized: Whether the span's terminal chunk has been accepted. + is_interrupted: Whether the caller cut this utterance short. + """ + + role: SpeakerRole + trace_id: str = "" + next_sequence: int = 0 + bytes_consumed: int = 0 + is_finalized: bool = False + is_interrupted: bool = False + + +@dataclass +class AudioSenderStats: + """Delivery counters for one call, stamped onto the ``agent_session`` span. + + The ``sent`` counters record what the endpoint *accepted*: a chunk that + failed every attempt raises ``errors``, never ``chunks_sent``. + + Attributes: + chunks_sent: Accepted HTTP requests carrying audio or a terminal marker. + frames_sent: Captured frames inside those accepted requests. + bytes_sent: PCM bytes inside those accepted requests. + frames_dropped: Frames discarded because the queue was full. + errors: Failed POST attempts, including ones a retry then recovered. + circuit_tripped: Whether the call gave up on the endpoint entirely. + total_send_time_ms: Wall-clock spent inside POSTs, for the average below. + """ + + chunks_sent: int = 0 + frames_sent: int = 0 + bytes_sent: int = 0 + frames_dropped: int = 0 + errors: int = 0 + circuit_tripped: bool = False + total_send_time_ms: float = 0.0 + + def __str__(self) -> str: + """Render the counters as a single log-friendly line.""" + average_ms = self.total_send_time_ms / self.chunks_sent if self.chunks_sent else 0.0 + return ( + f"chunks={self.chunks_sent} frames={self.frames_sent} " + f"bytes={self.bytes_sent} dropped={self.frames_dropped} " + f"errors={self.errors} avg_latency={average_ms:.1f}ms" + ) + + +# --------------------------------------------------------------------------- +# Sender +# --------------------------------------------------------------------------- + + +class AudioChunkSender: + """Batches captured frames and POSTs them to the audio-ingest endpoint. + + Single-consumer by construction: :meth:`enqueue` and the marker methods are + called from the agent's event loop and only hand work to a bounded queue, and + exactly one background task drains it. Nothing here is safe to call from + another thread. + """ + + def __init__( + self, + *, + url: str, + session_id: str, + api_key: str = "", + auth_headers: Optional[Dict[str, str]] = None, + batch_interval_seconds: float = DEFAULT_BATCH_INTERVAL_SECONDS, + max_batch_frames: int = DEFAULT_MAX_BATCH_FRAMES, + flush_at_bytes: int = DEFAULT_FLUSH_AT_BYTES, + max_request_bytes: int = DEFAULT_MAX_REQUEST_BYTES, + max_queue_frames: int = 0, + ) -> None: + """Configure the sender without starting it. + + Args: + url: Absolute audio-ingest URL, from ``Config.audio_endpoint()``. + session_id: Identifies the call; sent as ``x-audio-session-id``. + api_key: Credential sent as ``x-api-key`` when non-empty. + auth_headers: Further credential headers from the Netra config. + Applied only where they do not already have a value. + batch_interval_seconds: Longest a frame waits before being flushed. + max_batch_frames: Flush once this many frames have accumulated. + flush_at_bytes: Target request size — flush once this many PCM bytes + have accumulated. + max_request_bytes: Hard ceiling on one request body. A frame that + would push the batch past it flushes the batch first, so the + ceiling holds even when it sits just above *flush_at_bytes*. + max_queue_frames: Bound on frames awaiting batching; further frames + are dropped rather than queued. 0 means unbounded. + """ + self._url = url.rstrip("/") + self._session_id = session_id + self._api_key = api_key + self._auth_headers = auth_headers or {} + self._batch_interval_seconds = batch_interval_seconds + self._max_batch_frames = max_batch_frames + self._flush_at_bytes = flush_at_bytes + self._max_request_bytes = max(flush_at_bytes, max_request_bytes) + + self._queue: asyncio.Queue[_QueueMessage] = asyncio.Queue(maxsize=max(0, max_queue_frames)) + self._span_states: Dict[str, _SpanAudioState] = {} + self._send_task: Optional[asyncio.Task[None]] = None + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._client: Optional[httpx.AsyncClient] = None + self._is_closed = False + + self._consecutive_failures = 0 + self._circuit_tripped = False + self._has_warned_about_drops = False + + self.stats = AudioSenderStats() + + # -- lifecycle ---------------------------------------------------------- + + @property + def loop(self) -> Optional[asyncio.AbstractEventLoop]: + """The event loop this sender's queue and task belong to, once started. + + Everything here is bound to that loop, so a shutdown path reaching the + sender from elsewhere has to drive it through this rather than awaiting + it directly. ``None`` before :meth:`start`. + """ + return self._loop + + async def start(self) -> None: + """Open the HTTP client and start the background send loop.""" + self._loop = asyncio.get_running_loop() + self._client = httpx.AsyncClient(timeout=_HTTP_TIMEOUT_SECONDS) + self._send_task = asyncio.create_task(self._run_send_loop(), name="netra-audio-chunk-sender") + logger.info( + "netra.audio: sender started -> %s (batch=%.1fs, max_frames=%d, flush_at=%dB, max_request=%dB)", + self._url, + self._batch_interval_seconds, + self._max_batch_frames, + self._flush_at_bytes, + self._max_request_bytes, + ) + + async def end_session(self, *, drain_timeout_seconds: float = _DEFAULT_DRAIN_TIMEOUT_SECONDS) -> None: + """Drain the queue, close every open span, and signal the session's end. + + Idempotent: a second call returns immediately. Once this has been called + no further frames are accepted, so a late frame from a task that has not + noticed the shutdown is dropped rather than queued behind the terminal + marker it would never get past. + + Args: + drain_timeout_seconds: Total budget for the whole teardown. The two + waits inside share one deadline rather than each taking the full + timeout, because a caller that allowed *n* seconds for the session + to close means *n* seconds, not 2*n*. + """ + if self._is_closed: + return + self._is_closed = True + + deadline = time.monotonic() + max(0.0, drain_timeout_seconds) + await self._enqueue_session_end(deadline) + if self._send_task is not None: + await self._await_send_task(deadline) + if self._client is not None: + await self._client.aclose() + logger.info("netra.audio: sender closed — %s", self.stats) + + async def _enqueue_session_end(self, deadline: float) -> None: + """Get the terminal marker onto the queue, waiting for room if need be. + + ``put_nowait`` is wrong here: on a bounded queue that is currently full + the marker would be dropped and the send loop would never learn to stop, + so the drain below would spend its whole timeout before cancelling. No + producer can refill the queue at this point — ``_is_closed`` is already + set — so waiting for the consumer to make room terminates. + + Args: + deadline: ``time.monotonic()`` value the whole teardown must finish by. + """ + try: + await asyncio.wait_for(self._queue.put(_SessionEndMarker()), timeout=_seconds_until(deadline)) + except asyncio.TimeoutError: + logger.warning("netra.audio: could not signal session end before the teardown deadline") + + async def _await_send_task(self, deadline: float) -> None: + """Wait for the send loop to drain, cancelling it if it overruns. + + The cancellation is awaited rather than merely requested: ``end_session`` + closes the HTTP client next, and a send loop still inside a POST would + otherwise find the client shut from under it. + + Args: + deadline: ``time.monotonic()`` value the whole teardown must finish by. + """ + task = self._send_task + if task is None: + return + try: + await asyncio.wait_for(task, timeout=_seconds_until(deadline)) + except asyncio.TimeoutError: + logger.warning("netra.audio: send loop did not drain before the teardown deadline; cancelling") + task.cancel() + await asyncio.gather(task, return_exceptions=True) + except asyncio.CancelledError: + raise + except Exception: + logger.warning("netra.audio: send loop ended with an error", exc_info=True) + + # -- producer side (agent event loop) ----------------------------------- + + def enqueue( + self, + frame: "AudioFrame", + *, + role: SpeakerRole, + trace_id: str, + span_id: str = "", + timestamp_ns: Optional[int] = None, + ) -> None: + """Queue one captured frame. Never blocks, never raises into the agent. + + Copies the PCM out of the frame via the public ``frame.data`` + memoryview: LiveKit reuses the underlying buffer for the next frame, so + holding a reference would corrupt the batch. + + Args: + frame: The LiveKit frame just captured. + role: Which speaker produced it. + trace_id: Hex trace id to attribute the audio to. + span_id: Hex id of the open speaking span, or ``""`` for audio + captured between turns. + timestamp_ns: Capture time, defaulting to now. Passed in by the + coordinator so the timestamp is taken at capture rather than + after any queuing delay. + """ + if self._is_closed or self._circuit_tripped: + return + try: + message = _FrameMessage( + pcm_bytes=bytes(frame.data), + role=role, + span_id=span_id, + trace_id=trace_id, + sample_rate_hz=frame.sample_rate, + channel_count=frame.num_channels, + timestamp_ns=timestamp_ns if timestamp_ns is not None else time.time_ns(), + ) + except (AttributeError, TypeError, ValueError): + # A frame shaped differently from what livekit-agents documents. Not + # recoverable and not the agent's problem — drop this one frame. + logger.debug("netra.audio: unreadable audio frame dropped", exc_info=True) + self.stats.frames_dropped += 1 + return + + if not self._offer(message): + self.stats.frames_dropped += 1 + self._warn_about_drops_once() + + def mark_audio_end(self, *, role: SpeakerRole, span_id: str) -> None: + """Signal that the recording for *span_id* is complete. + + Args: + role: The speaker whose span closed. + span_id: Hex id of the closed speaking span. + """ + if self._is_closed or not span_id: + return + state = self._span_states.get(span_id) + if state is not None and state.is_finalized: + return + if not self._offer(_SpanEndMarker(role=role, span_id=span_id)): + logger.debug("netra.audio: queue full; end marker for span=%s dropped", span_id) + + def interrupt_agent_span(self, *, span_id: str, playback_ms: int) -> None: + """Signal that an agent utterance was cut off *playback_ms* into playback. + + The send loop trims the pending audio for the span to what was heard and + finalizes it. This is still correct when the span was already finalized + through :meth:`mark_audio_end` — LiveKit routinely ends the + ``agent_speaking`` span before it reports the interrupt — in which case a + bodyless correction carrying only ``x-audio-heard-ms`` follows. + + Args: + span_id: Hex id of the interrupted ``agent_speaking`` span. + playback_ms: Milliseconds of the utterance the caller heard. + """ + if self._is_closed or not span_id: + return + if not self._offer(_SpanInterruptMarker(span_id=span_id, playback_ms=playback_ms)): + logger.debug("netra.audio: queue full; interrupt marker for span=%s dropped", span_id) + + def _offer(self, message: _QueueMessage) -> bool: + """Hand *message* to the send loop without ever blocking the caller. + + ``asyncio.Queue`` is not thread-safe, and the marker methods are reachable + from :class:`AudioSpanProcessor`, which OTel invokes on whichever thread + ends the span — normally the agent's loop thread, but nothing enforces + that. An off-loop caller is therefore bounced onto the sender's own loop + instead of corrupting the queue. + + Args: + message: The message to enqueue. + + Returns: + True when it was queued or handed to the loop, False when the queue is + at its bound. The caller decides how a drop is accounted for — a + dropped frame is a statistic, a dropped marker is not. + """ + loop = self._loop + if loop is not None and loop is not _running_loop(): + # Whether the queue had room is not knowable from here; the hop itself + # succeeding is all this can report. + loop.call_soon_threadsafe(self._offer_on_loop, message) + return True + return self._put_nowait(message) + + def _offer_on_loop(self, message: _QueueMessage) -> None: + """Enqueue a message that arrived from another thread. Runs on the loop. + + Args: + message: The message to enqueue. + """ + if not self._put_nowait(message): + logger.debug("netra.audio: queue full; cross-thread %s dropped", type(message).__name__) + + def _put_nowait(self, message: _QueueMessage) -> bool: + """Put *message* on the queue if it has room. + + Args: + message: The message to enqueue. + + Returns: + True when it was queued, False when the queue is at its bound. + """ + try: + self._queue.put_nowait(message) + except asyncio.QueueFull: + return False + return True + + def _warn_about_drops_once(self) -> None: + """Warn that frames are being dropped, at most once per session.""" + if self._has_warned_about_drops: + return + self._has_warned_about_drops = True + logger.warning( + "netra.audio: queue full, dropping frames (session=%s). This is logged once per session", + self._session_id, + ) + + # -- consumer side (background task) ------------------------------------ + + async def _run_send_loop(self) -> None: + """Drain the queue until the session ends, with instrumentation muted. + + The loop's own HTTP calls run under ``_SUPPRESS_INSTRUMENTATION_KEY`` so + Netra's httpx instrumentation does not trace them: every audio chunk + would otherwise produce a span, inside the very trace the audio belongs + to. + """ + from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY + + token = otel_context.attach(otel_context.set_value(_SUPPRESS_INSTRUMENTATION_KEY, True)) + try: + await self._consume_queue() + finally: + otel_context.detach(token) + + async def _consume_queue(self) -> None: + """Batch queued frames and post them until the session-end marker.""" + batches = {role: _PendingBatch(role=role) for role in SpeakerRole} + + while True: + try: + message = await asyncio.wait_for(self._queue.get(), timeout=self._batch_interval_seconds) + except asyncio.TimeoutError: + await self._flush_idle_batches(batches) + continue + + if isinstance(message, _SessionEndMarker): + await self._drain_batches(batches) + return + + await self._handle_message(message, batches) + + async def _handle_message(self, message: _QueueMessage, batches: Dict[SpeakerRole, _PendingBatch]) -> None: + """Dispatch one queued message to its handler. + + Args: + message: The message the loop dequeued. + batches: The pending batch for each speaker. + """ + if isinstance(message, _FrameMessage): + await self._handle_frame(message, batches[message.role]) + elif isinstance(message, _SpanEndMarker): + await self._handle_span_end(message, batches[message.role]) + elif isinstance(message, _SpanInterruptMarker): + await self._handle_span_interrupt(message, batches[SpeakerRole.AGENT]) + + async def _handle_frame(self, frame: _FrameMessage, batch: _PendingBatch) -> None: + """Accumulate one frame, flushing first or after if a boundary is hit. + + Args: + frame: The frame to accumulate. + batch: The pending batch for that frame's speaker. + """ + state = self._span_states.get(frame.span_id) if frame.span_id else None + if state is not None and state.is_interrupted: + # Queued before the interrupt was observed but captured after the + # caller cut in — this audio was never heard. + return + + # A batch holds one span's audio: the chunk's span id is a single header. + # It also has to stay under the request ceiling, so a frame that would + # burst it closes the batch instead of joining it. + spans_differ = not batch.is_empty and batch.span_id != frame.span_id + would_overflow = batch.byte_count + len(frame.pcm_bytes) > self._max_request_bytes + if spans_differ or would_overflow: + await self._flush(batch) + + batch.add(frame) + + if batch.frame_count >= self._max_batch_frames or batch.byte_count >= self._flush_at_bytes: + await self._flush(batch) + + async def _handle_span_end(self, marker: _SpanEndMarker, batch: _PendingBatch) -> None: + """Finalize a speaking span, flushing whatever audio is still pending. + + Args: + marker: The end marker for the span. + batch: The pending batch for that span's speaker. + """ + state = self._span_states.get(marker.span_id) + if state is not None and state.is_finalized: + return + + if batch.span_id == marker.span_id and not batch.is_empty: + await self._flush(batch, is_final=True) + return + + if batch.span_id == marker.span_id: + batch.clear() + await self._post_span_terminator(role=marker.role, span_id=marker.span_id) + + async def _handle_span_interrupt(self, marker: _SpanInterruptMarker, batch: _PendingBatch) -> None: + """Trim an interrupted agent span to the audio heard, then finalize it. + + Args: + marker: The interrupt marker, carrying the playback position. + batch: The pending agent batch. + """ + state = self._state_for(marker.span_id, SpeakerRole.AGENT) + state.is_interrupted = True + + if batch.span_id != marker.span_id or batch.is_empty: + # Nothing pending: the audio already went out, so all the endpoint + # needs is where to cut it. Forced, because the normal end marker has + # usually finalized the span by now. + await self._post_span_terminator( + role=state.role, + span_id=marker.span_id, + heard_ms=marker.playback_ms, + force=True, + ) + return + + await self._flush_heard_prefix(batch, marker.playback_ms) + + async def _flush_heard_prefix(self, batch: _PendingBatch, playback_ms: int) -> None: + """Post only the part of *batch* the caller heard, marked final. + + The heard prefix is measured from the start of the *span*, so whatever + earlier chunks already consumed of it has to come off the offset before + the pending batch can be trimmed. + + Args: + batch: The pending agent batch, known to hold audio for the span. + playback_ms: Milliseconds of the utterance the caller heard. + """ + # Read the batch's identity out before any flush: ``_PendingBatch.clear`` + # resets ``span_id``, so a terminator addressed from a cleared batch would + # carry ``""`` and be silently dropped by ``_post_span_terminator``. + span_id = batch.span_id + role = batch.role + heard_offset = pcm_byte_offset_at( + playback_ms=playback_ms, + sample_rate_hz=batch.sample_rate_hz or DEFAULT_SAMPLE_RATE_HZ, + channel_count=batch.channel_count or DEFAULT_CHANNEL_COUNT, + ) + already_consumed = self._state_for(span_id, role).bytes_consumed + remaining = heard_offset - already_consumed + + if remaining <= 0: + # Everything heard has already been sent; the endpoint only needs the + # cut point so it can discard the overshoot. Forced, because the normal + # end marker may already have finalized the span. + batch.clear() + await self._post_span_terminator( + role=role, + span_id=span_id, + heard_ms=playback_ms, + force=True, + ) + return + + heard_pcm = batch.pcm_bytes[:remaining] + logger.debug( + "netra.audio: trimmed interrupted span=%s to %d of %d pending bytes (heard=%dms, consumed=%d)", + span_id, + len(heard_pcm), + batch.byte_count, + playback_ms, + already_consumed, + ) + frame_count = batch.frames_within(len(heard_pcm)) + start_ms = batch.start_ms + sample_rate_hz = batch.sample_rate_hz + channel_count = batch.channel_count + trace_id = batch.trace_id + batch.clear() + await self._post_chunk( + role=role, + span_id=span_id, + trace_id=trace_id, + sample_rate_hz=sample_rate_hz, + channel_count=channel_count, + pcm=heard_pcm, + frame_count=frame_count, + start_ms=start_ms, + is_last=True, + heard_ms=playback_ms, + ) + + async def _flush_idle_batches(self, batches: Dict[SpeakerRole, _PendingBatch]) -> None: + """Flush both speakers' pending audio after an idle interval. + + Args: + batches: The pending batch for each speaker. + """ + for batch in batches.values(): + await self._flush(batch) + + async def _drain_batches(self, batches: Dict[SpeakerRole, _PendingBatch]) -> None: + """Send everything still held, then close the session on the wire. + + Args: + batches: The pending batch for each speaker. + """ + for batch in batches.values(): + await self._flush(batch, is_final=bool(batch.span_id)) + await self._finalize_open_spans() + await self._post_session_terminator() + + async def _flush(self, batch: _PendingBatch, *, is_final: bool = False) -> None: + """Post *batch*'s audio and clear it. + + A final flush is two requests, not one: the audio chunk, then an empty + chunk carrying ``x-audio-last``. Keeping the terminator separate means + the span closes the same way whether or not audio happened to be pending + when it ended. + + Args: + batch: The batch to send. + is_final: Whether this closes the batch's span. + """ + if batch.is_empty and not is_final: + return + + span_id = batch.span_id + role = batch.role + trace_id = batch.trace_id + + if not batch.is_empty: + await self._post_chunk( + role=role, + span_id=span_id, + trace_id=trace_id, + sample_rate_hz=batch.sample_rate_hz, + channel_count=batch.channel_count, + pcm=batch.pcm_bytes, + frame_count=batch.frame_count, + start_ms=batch.start_ms, + is_last=False, + ) + batch.clear() + + if is_final and span_id: + await self._post_span_terminator(role=role, span_id=span_id) + + async def _finalize_open_spans(self) -> None: + """Close any span that never received an end marker. + + A span left open would leave the endpoint waiting for audio that is + never coming, so this is a backstop rather than a normal path — hence + the warning. + + Skipped entirely once the circuit has tripped: every span is open in that + case, by definition, and ``_trip_circuit`` has already said why once. + Warning per span would bury it under hundreds of lines. + """ + if self._circuit_tripped: + return + + open_span_ids = sorted(span_id for span_id, state in self._span_states.items() if not state.is_finalized) + for span_id in open_span_ids: + state = self._span_states[span_id] + logger.warning( + "netra.audio: finalizing span left open at session end: span_id=%s role=%s", + span_id, + state.role.value, + ) + await self._post_span_terminator(role=state.role, span_id=span_id) + + # -- requests ----------------------------------------------------------- + + async def _post_span_terminator( + self, + *, + role: SpeakerRole, + span_id: str, + heard_ms: int = 0, + force: bool = False, + ) -> None: + """Post the empty chunk that closes a span. + + Args: + role: The speaker the span belongs to. + span_id: Hex id of the span to close. + heard_ms: Milliseconds heard, for an interrupted agent span only. + force: Send even though the span is already finalized. Used for an + interrupt correction arriving after the normal terminator. + """ + if not span_id: + return + state = self._span_states.get(span_id) + if state is not None and state.is_finalized and not force: + return + + await self._post_chunk( + role=role, + span_id=span_id, + trace_id=state.trace_id if state is not None else "", + sample_rate_hz=DEFAULT_SAMPLE_RATE_HZ, + channel_count=DEFAULT_CHANNEL_COUNT, + pcm=b"", + frame_count=0, + start_ms=0, + is_last=True, + heard_ms=heard_ms, + ) + + async def _post_session_terminator(self) -> None: + """Post the bodyless request that marks the whole session complete. + + Skipped once the circuit has tripped: "no further audio will be sent for + this session" has to include this request, or a session abandoned over a + rejected credential would still end with one more rejected POST. + """ + if self._circuit_tripped: + return + + headers = { + HEADER_SESSION_ID: self._session_id, + HEADER_SESSION_LAST: HEADER_VALUE_TRUE, + } + self._apply_credentials(headers) + await self._post(b"", headers) + + async def _post_chunk( + self, + *, + role: SpeakerRole, + span_id: str, + trace_id: str, + sample_rate_hz: int, + channel_count: int, + pcm: bytes, + frame_count: int, + start_ms: int, + is_last: bool, + heard_ms: int = 0, + ) -> None: + """Send one chunk and record what it did to the span's state. + + Args: + role: The speaker the audio came from. + span_id: Hex id of the speaking span, or ``""`` for between-turn audio. + trace_id: Hex trace id the audio belongs to. + sample_rate_hz: Samples per second, per channel. + channel_count: Interleaved channel count. + pcm: The body — signed 16-bit little-endian PCM. + frame_count: How many captured frames the body holds, for the stats. + start_ms: Epoch milliseconds of the body's first frame. + is_last: Whether this closes the span. + heard_ms: Milliseconds heard, for an interrupted agent span only. + """ + if self._circuit_tripped: + return + + state = self._state_for(span_id, role, trace_id) if span_id else None + headers = self._chunk_headers( + role=role, + span_id=span_id, + trace_id=trace_id, + sample_rate_hz=sample_rate_hz, + channel_count=channel_count, + start_ms=start_ms, + is_last=is_last, + heard_ms=heard_ms, + state=state, + ) + + accepted = await self._post(pcm, headers) + + logger.debug( + "netra.audio: chunk span_id=%s role=%s frames=%d bytes=%d last=%s accepted=%s", + span_id or "(between turns)", + role.value, + frame_count, + len(pcm), + is_last, + accepted, + ) + + if state is not None: + # Advanced whether or not the chunk landed. Both are positions in the + # span's stream, not delivery counts: a chunk the sender gave up on + # still occupied its slot, so reusing its number for the *next*, + # different audio would break the idempotency key the endpoint dedupes + # on. A gap is how the endpoint learns audio was lost. + state.next_sequence += 1 + state.bytes_consumed += len(pcm) + if accepted and is_last: + state.is_finalized = True + + if not accepted: + return + + self.stats.chunks_sent += 1 + self.stats.frames_sent += frame_count + self.stats.bytes_sent += len(pcm) + + def _chunk_headers( + self, + *, + role: SpeakerRole, + span_id: str, + trace_id: str, + sample_rate_hz: int, + channel_count: int, + start_ms: int, + is_last: bool, + heard_ms: int, + state: Optional[_SpanAudioState], + ) -> Dict[str, str]: + """Build the ``x-audio-*`` headers describing one chunk. + + Args: + role: The speaker the audio came from. + span_id: Hex span id, or ``""`` for between-turn audio. + trace_id: Hex trace id the audio belongs to. + sample_rate_hz: Samples per second, per channel. + channel_count: Interleaved channel count. + start_ms: Epoch milliseconds of the first frame. + is_last: Whether this closes the span. + heard_ms: Milliseconds heard, for an interrupted agent span only. + state: The span's state, or ``None`` for between-turn audio. + + Returns: + The complete header set for the request. + """ + headers = { + HEADER_CONTENT_TYPE: CONTENT_TYPE_PCM, + HEADER_SESSION_ID: self._session_id, + HEADER_TRACE_ID: trace_id, + HEADER_ROLE: role.value, + HEADER_START_MS: str(start_ms), + HEADER_SAMPLE_RATE: str(sample_rate_hz or DEFAULT_SAMPLE_RATE_HZ), + HEADER_CHANNELS: str(channel_count or DEFAULT_CHANNEL_COUNT), + HEADER_BIT_DEPTH: str(PCM_BIT_DEPTH), + } + self._apply_credentials(headers) + + if state is not None: + headers[HEADER_SPAN_ID] = span_id + headers[HEADER_SEQUENCE] = str(state.next_sequence) + if is_last: + headers[HEADER_LAST_CHUNK] = HEADER_VALUE_TRUE + if heard_ms > 0: + headers[HEADER_HEARD_MS] = str(heard_ms) + return headers + + def _apply_credentials(self, headers: Dict[str, str]) -> None: + """Add the configured credential headers, without overwriting any. + + Args: + headers: The header set being built, mutated in place. + """ + if self._api_key: + headers[HEADER_API_KEY] = self._api_key + for name, value in self._auth_headers.items(): + headers.setdefault(name, value) + + async def _post(self, pcm: bytes, headers: Dict[str, str]) -> bool: + """POST one request, retrying a transient failure. + + Args: + pcm: The request body. + headers: The request headers. + + Returns: + True when the endpoint accepted the request. + """ + client = self._client + if client is None: + logger.debug("netra.audio: post attempted before start(); dropping chunk") + return False + + for attempt in range(_POST_ATTEMPTS): + accepted, is_fatal = await self._post_once(client, pcm, headers, attempt) + if accepted or is_fatal: + return accepted + if attempt < _POST_ATTEMPTS - 1: + await asyncio.sleep(_retry_delay_seconds(attempt)) + + logger.warning("netra.audio: giving up on a chunk after %d attempts", _POST_ATTEMPTS) + return False + + async def _post_once( + self, + client: httpx.AsyncClient, + pcm: bytes, + headers: Dict[str, str], + attempt: int, + ) -> tuple[bool, bool]: + """Make one POST attempt and account for its outcome. + + Args: + client: The open HTTP client. + pcm: The request body. + headers: The request headers. + attempt: 0-based attempt number, for the log line. + + Returns: + ``(accepted, is_fatal)`` — ``is_fatal`` means retrying cannot help, + either because the credential was rejected or because the circuit + breaker has now tripped. + """ + started_at = time.monotonic() + try: + response = await client.post(self._url, content=pcm, headers=headers) + except httpx.HTTPError as exc: + self.stats.total_send_time_ms += (time.monotonic() - started_at) * 1000 + self.stats.errors += 1 + logger.warning("netra.audio: chunk POST error (attempt=%d): %s", attempt + 1, exc) + return False, self._record_failure() + + self.stats.total_send_time_ms += (time.monotonic() - started_at) * 1000 + + if response.status_code < _HTTP_STATUS_BAD_REQUEST: + self._consecutive_failures = 0 + return True, False + + self.stats.errors += 1 + if response.status_code in _UNAUTHENTICATED_STATUSES: + self._trip_circuit(f"HTTP {response.status_code} — a credential will not become valid mid-call") + return False, True + + logger.warning( + "netra.audio: chunk POST rejected (attempt=%d): %d %s", + attempt + 1, + response.status_code, + response.text[:200], + ) + return False, self._record_failure() + + # -- failure handling --------------------------------------------------- + + def _record_failure(self) -> bool: + """Count one failure and trip the circuit if the run is long enough. + + Returns: + True when the circuit is now open, meaning retrying is pointless. + """ + self._consecutive_failures += 1 + if self._consecutive_failures >= _MAX_CONSECUTIVE_FAILURES: + self._trip_circuit(f"{self._consecutive_failures} consecutive failures") + return self._circuit_tripped + + def _trip_circuit(self, reason: str) -> None: + """Abandon audio for the rest of the call. + + Args: + reason: What went wrong, for the operator-facing log line. + """ + if self._circuit_tripped: + return + self._circuit_tripped = True + self.stats.circuit_tripped = True + logger.warning( + "netra.audio: circuit breaker tripped (session=%s): %s. " + "No further audio will be sent for this session; traces are unaffected", + self._session_id, + reason, + ) + + # -- span state --------------------------------------------------------- + + def _state_for(self, span_id: str, role: SpeakerRole, trace_id: str = "") -> _SpanAudioState: + """Return the state record for *span_id*, creating it on first sight. + + Args: + span_id: Hex id of a speaking span. + role: The speaker it belongs to. + trace_id: Hex trace id, remembered so a later terminator for this + span can still be attributed once the batch holding it is gone. + + Returns: + The span's mutable state record. + """ + state = self._span_states.get(span_id) + if state is None: + state = _SpanAudioState(role=role, trace_id=trace_id) + self._span_states[span_id] = state + elif trace_id and not state.trace_id: + state.trace_id = trace_id + return state + + +def _running_loop() -> Optional[asyncio.AbstractEventLoop]: + """Return the loop running on this thread, or ``None`` on a plain thread. + + Returns: + The current event loop, if there is one. + """ + try: + return asyncio.get_running_loop() + except RuntimeError: + return None + + +def _seconds_until(deadline: float) -> float: + """Return the time left before *deadline*, never negative. + + Args: + deadline: A ``time.monotonic()`` value. + + Returns: + Seconds remaining. 0.0 once the deadline has passed, which makes the + ``wait_for`` it is handed to give up immediately rather than restart the + full budget. + """ + return max(0.0, deadline - time.monotonic()) + + +def _retry_delay_seconds(attempt: int) -> float: + """Return the backoff before retrying, exponential with full jitter. + + Args: + attempt: 0-based number of the attempt that just failed. + + Returns: + Seconds to wait. Jittered so that a backend recovering from an outage is + not hit by every concurrent call's sender at the same instant. + """ + ceiling = _RETRY_BASE_DELAY_SECONDS * (2**attempt) + return random.uniform(0.0, ceiling) diff --git a/netra/instrumentation/livekit/audio_types.py b/netra/instrumentation/livekit/audio_types.py new file mode 100644 index 0000000..0a1a454 --- /dev/null +++ b/netra/instrumentation/livekit/audio_types.py @@ -0,0 +1,156 @@ +"""Vocabulary shared by every part of the LiveKit call-audio pipeline. + +Three kinds of thing live here, and nothing else: + +* :class:`SpeakerRole` — the two speakers a frame can belong to, as an enum + rather than the string ``"user"``/``"agent"`` that used to be threaded through + the sender, the coordinator and the span processor independently; +* the PCM format constants and the one arithmetic helper that converts a + playback duration to a byte offset; +* the wire contract — the ``x-audio-*`` request headers the ingest endpoint + reads, and the ``netra.audio.*`` span attributes the session root is stamped + with. + +Free of OTel and LiveKit imports, so the wire contract can be asserted against +in tests without a tracer or a livekit-agents install. +""" + +from enum import Enum +from typing import Dict + +# --------------------------------------------------------------------------- +# Speakers +# --------------------------------------------------------------------------- + + +class SpeakerRole(str, Enum): + """Which side of the call a run of audio came from. + + A ``str`` enum because the value is also the wire value of the + ``x-audio-role`` header, so the two cannot drift apart. + """ + + USER = "user" + AGENT = "agent" + + +# LiveKit span name -> the speaker whose audio that span delimits. The audio for +# a call is addressed by these spans' ids, so a frame arriving while one is open +# is attributed to it and a frame arriving between them is attributed to nobody +# (see ``SessionAudioCoordinator``). +SPEAKING_SPAN_ROLES: Dict[str, SpeakerRole] = { + "user_speaking": SpeakerRole.USER, + "agent_speaking": SpeakerRole.AGENT, +} + + +# --------------------------------------------------------------------------- +# PCM format +# --------------------------------------------------------------------------- + +# What the ingest endpoint is told when a frame reported no format of its own — +# the terminal empty chunk of a span, which carries no frame to read it from. +DEFAULT_SAMPLE_RATE_HZ = 16000 +DEFAULT_CHANNEL_COUNT = 1 + +# The body is always signed 16-bit little-endian PCM. Not negotiable per chunk: +# the ingest endpoint reads it off the header only so a future format change can +# be rolled out without breaking stored audio. +PCM_BIT_DEPTH = 16 +PCM_BYTES_PER_SAMPLE = PCM_BIT_DEPTH // 8 + +_MILLISECONDS_PER_SECOND = 1000 + + +def pcm_byte_offset_at(*, playback_ms: int, sample_rate_hz: int, channel_count: int) -> int: + """Return the PCM byte offset *playback_ms* into a stream, on a frame boundary. + + Used to trim an interrupted agent utterance down to the audio the caller + actually heard. The result is rounded *down* to a whole sample frame: + cutting mid-sample would leave the stored audio one byte out of phase for + its whole remaining length. + + Args: + playback_ms: Milliseconds of audio played out. A non-positive value means + nothing was heard and yields 0. + sample_rate_hz: Samples per second, per channel. Must be positive. + channel_count: Number of interleaved channels. Must be positive. + + Returns: + The byte offset, never negative and always a multiple of the frame size. + + Raises: + ValueError: If the PCM format is not playable. Callers substitute + :data:`DEFAULT_SAMPLE_RATE_HZ` / :data:`DEFAULT_CHANNEL_COUNT` for a + frame that reported neither, so reaching this is a programming error + rather than bad input. + """ + if sample_rate_hz <= 0 or channel_count <= 0: + raise ValueError(f"unplayable PCM format: sample_rate_hz={sample_rate_hz} channel_count={channel_count}") + if playback_ms <= 0: + return 0 + + frame_size = channel_count * PCM_BYTES_PER_SAMPLE + bytes_per_ms = sample_rate_hz * frame_size / _MILLISECONDS_PER_SECOND + return int(playback_ms * bytes_per_ms) // frame_size * frame_size + + +# --------------------------------------------------------------------------- +# Wire contract: request headers +# --------------------------------------------------------------------------- + +HEADER_CONTENT_TYPE = "Content-Type" +CONTENT_TYPE_PCM = "application/octet-stream" + +HEADER_API_KEY = "x-api-key" + +HEADER_SESSION_ID = "x-audio-session-id" +HEADER_TRACE_ID = "x-audio-trace-id" +HEADER_SPAN_ID = "x-audio-span-id" +HEADER_ROLE = "x-audio-role" +HEADER_SAMPLE_RATE = "x-audio-sample-rate" +HEADER_CHANNELS = "x-audio-channels" +HEADER_BIT_DEPTH = "x-audio-bit-depth" + +# Epoch milliseconds at which the first frame of this chunk was captured. +HEADER_START_MS = "x-audio-start-ms" + +# 0-based and monotonic *per span* — a chunk's position in that span's stream, +# not a count of what arrived. Two properties follow, and the endpoint depends on +# both: +# +# * the retries of a single chunk all carry the same number and the same bytes, +# so the endpoint can treat them as idempotent; +# * a chunk the sender gave up on still consumes its number, so a gap in the +# sequence is the endpoint's signal that audio was lost — never a number +# reused for different bytes. +HEADER_SEQUENCE = "x-audio-seq" + +# Present on the final chunk of a span, and on that chunk only. +HEADER_LAST_CHUNK = "x-audio-last" + +# Only on the final chunk of an *interrupted* agent span: how many milliseconds +# of the utterance the caller heard before cutting in. +HEADER_HEARD_MS = "x-audio-heard-ms" + +# Present on the bodyless request that closes the session. +HEADER_SESSION_LAST = "x-audio-session-last" + +HEADER_VALUE_TRUE = "true" + +# The request headers a Netra config may contribute as an audio-ingest +# credential. Lower-cased for comparison against user-supplied header names. +CREDENTIAL_HEADER_NAMES = frozenset({"x-api-key", "authorization"}) + + +# --------------------------------------------------------------------------- +# Wire contract: span attributes +# --------------------------------------------------------------------------- + +# Stamped on the ``agent_session`` span as it closes, so a trace shows what the +# audio pipeline actually managed to deliver for that call. +NETRA_AUDIO_SENT_BYTES = "netra.audio.sent_bytes" +NETRA_AUDIO_SENT_CHUNKS = "netra.audio.sent_chunks" +NETRA_AUDIO_DROPPED_FRAMES = "netra.audio.dropped_frames" +NETRA_AUDIO_ERRORS = "netra.audio.errors" +NETRA_AUDIO_CIRCUIT_TRIPPED = "netra.audio.circuit_tripped" diff --git a/netra/instrumentation/livekit/provider_binding.py b/netra/instrumentation/livekit/provider_binding.py new file mode 100644 index 0000000..62b3d88 --- /dev/null +++ b/netra/instrumentation/livekit/provider_binding.py @@ -0,0 +1,160 @@ +"""Binds livekit-agents' OTel tracer to Netra's provider, behind a shield. + +``livekit-agents`` does two things to whatever ``TracerProvider`` it is handed, +both of which are wrong for us: + +* it calls ``shutdown()`` on every job cleanup, which would permanently disable + Netra's ``BatchSpanProcessor`` for every later job in the process; +* it calls ``add_span_processor()`` to install its LiveKit Cloud exporter and a + metadata processor, which are process-wide and would therefore export *every* + Netra span to a third party. + +``_ShieldedTracerProvider`` delegates the reads LiveKit needs and absorbs both. +""" + +import logging +from typing import Any + +from opentelemetry import trace as trace_api +from opentelemetry.sdk import trace as trace_sdk +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import SpanProcessor + +logger = logging.getLogger(__name__) + +# Set on the *delegate* once bound, mirroring ``_netra_processors_installed`` +# in netra/tracer.py, so repeat instrument() calls are idempotent. +_BOUND_FLAG = "_netra_livekit_tracer_bound" + + +class _ShieldedTracerProvider(trace_sdk.TracerProvider): # type: ignore[misc] + """Delegates to Netra's TracerProvider but absorbs everything LiveKit does to it. + + Holds no mutable state, so it needs no lock. + + MUST subclass ``trace_sdk.TracerProvider``: LiveKit's ``_setup_cloud_tracer`` + and ``_shutdown_telemetry`` both gate on + ``isinstance(..., trace_sdk.TracerProvider)``, and a duck-typed object would + take a different branch — in the cloud-tracer case, one that never reads our + resource. + """ + + def __init__(self, delegate: trace_api.TracerProvider) -> None: + """Wrap *delegate* without initialising a second provider. + + Deliberately does not call ``super().__init__()``: every method LiveKit + touches is overridden and delegated, and constructing real SDK provider + state here would create a second, useless span pipeline. The contact + surface was verified against livekit-agents 1.6.7 + (``telemetry/traces.py`` ``set_tracer_provider`` / + ``_setup_cloud_tracer`` / ``_shutdown_telemetry``). + + Args: + delegate: Netra's real SDK ``TracerProvider``. + """ + self._delegate = delegate + + def get_tracer(self, *args: Any, **kwargs: Any) -> Any: + """Return a tracer from Netra's provider — the whole point of the shield. + + Args: + *args: Positional arguments forwarded verbatim to the delegate + (``instrumenting_module_name`` and friends). + **kwargs: Keyword arguments forwarded verbatim to the delegate. + + Returns: + A tracer created by Netra's provider, so LiveKit's spans enter Netra's + pipeline. + """ + return self._delegate.get_tracer(*args, **kwargs) + + @property + def resource(self) -> Any: + """Expose Netra's resource; LiveKit reads it in ``_setup_cloud_tracer``. + + Falls back to an empty ``Resource`` when the delegate has none: LiveKit + reaches this behind an ``isinstance(..., trace_sdk.TracerProvider)`` check + that we satisfy by subclassing, so an API-only delegate would otherwise + raise ``AttributeError`` inside LiveKit's code. + + Returns: + The delegate's ``Resource``, or an empty one when it has none. + """ + return getattr(self._delegate, "resource", Resource.get_empty()) + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """Propagate flushes: we do want the tail of a session exported. + + Args: + timeout_millis: Maximum time to wait for the delegate's processors to + flush. + + Returns: + Whatever the delegate reports, or True when it exposes no + ``force_flush`` — there is nothing pending in that case. + """ + flush = getattr(self._delegate, "force_flush", None) + if flush is None: + return True + result: bool = flush(timeout_millis) + return result + + def shutdown(self) -> None: + """Absorb LiveKit's per-job teardown of Netra's tracing pipeline.""" + logger.debug( + "netra.livekit: absorbed a TracerProvider shutdown; Netra owns this provider's lifecycle", + ) + + def add_span_processor(self, span_processor: SpanProcessor) -> None: + """Refuse every processor LiveKit tries to install on Netra's provider. + + LiveKit registers a ``_MetadataSpanProcessor`` and a Cloud + ``BatchSpanProcessor`` whenever recording is enabled. Both are + process-wide, so accepting them would (a) export every Netra span — + ``openai.chat``, ``httpx``, ``@task`` — to LiveKit Cloud, and (b) stamp + ``room_id``/``job_id`` on spans from unrelated work, because + ``_MetadataSpanProcessor.on_start`` is unconditional. + + Netra spans are never exported to a third party. There is no flag to + change this. + + Args: + span_processor: The processor LiveKit asked us to install. Discarded. + """ + logger.info( + "netra.livekit: refused LiveKit-added span processor %s; Netra spans are never " + "exported to LiveKit Cloud. LiveKit Cloud trace recording is inactive in this " + "process (its logs and session reports are unaffected)", + type(span_processor).__name__, + ) + + +def bind_livekit_tracer(provider: trace_api.TracerProvider) -> None: + """Hand LiveKit a shielded view of Netra's TracerProvider. Idempotent. + + Takes no ``Config``: there is nothing left to configure about the binding. + + Accepts the API type rather than the SDK one because + ``trace.get_tracer_provider()`` may hand back a proxy — binding is still + correct in that case, since the shield only delegates. + + Args: + provider: The tracer provider LiveKit's spans should be created from. + + Raises: + ImportError: If ``livekit.agents.telemetry.set_tracer_provider`` cannot be + imported. The caller logs this and continues — losing trace binding + must not disable the session hooks. + """ + if getattr(provider, _BOUND_FLAG, False): + return + + from livekit.agents.telemetry import set_tracer_provider + + shield = _ShieldedTracerProvider(provider) + # No metadata= argument, ever: that path calls add_span_processor() on the + # object we hand over, so keeping the call single-argument means the + # guarantee does not depend on our gate holding in a future LiveKit version. + set_tracer_provider(shield) + setattr(provider, _BOUND_FLAG, True) + logger.info("netra.livekit: bound livekit-agents tracer to Netra's TracerProvider") diff --git a/netra/instrumentation/livekit/trace_processor.py b/netra/instrumentation/livekit/trace_processor.py new file mode 100644 index 0000000..ec882bf --- /dev/null +++ b/netra/instrumentation/livekit/trace_processor.py @@ -0,0 +1,554 @@ +"""Normalises the shape of livekit-agents' trace into Netra's conventions. + +The trace half of this package's two span processors: it rewrites what LiveKit +puts *on* a span — the ``lk.*`` attributes, the conversation events, the +classification markers a span's name implies. The audio half, +``audio_processor.py``, uses spans only as timing boundaries for captured PCM +and shares none of this module's machinery. + +INVARIANT for anything added here: ``on_end`` must never mutate the span that is +ending. By the time it runs, ``BatchSpanProcessor`` — registered earlier in the +chain — has already queued that span, and the exporter serialises it on another +thread. ``on_end`` may only mutate *other* spans that are still recording, which is +exactly what the parent-ward content propagation below does. +""" + +from __future__ import annotations + +import itertools +import logging +import threading +import weakref +from typing import Any, Callable, Dict, Iterator, Mapping, Optional, Tuple + +from opentelemetry import context as otel_context +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor +from opentelemetry.util.types import Attributes + +from netra.instrumentation.livekit.utils import ( + ATTRIBUTE_MAP, + AUDIO_TYPE_BY_SPAN_NAME, + CHAT_CTX_ATTRIBUTE, + CONVERSATION_MAP, + EVENT_CHOICE, + EVENT_ROLE, + GEN_AI_COMPLETION_CONTENT, + GEN_AI_COMPLETION_ROLE, + GEN_AI_PROMPT_CONTENT, + GEN_AI_PROMPT_ROLE, + GEN_AI_REQUEST_MODEL, + GEN_AI_USAGE_CHARACTER_COUNT, + IO_FROM_CHILD_SPAN_NAMES, + LIVEKIT_SCOPE_NAME, + MAX_CONVERSATION_MESSAGES_PER_SIDE, + NETRA_AUDIO_TYPE, + NETRA_CONVERSATION_TRUNCATED, + NETRA_ENTITY_TYPE, + NETRA_ENTITY_TYPE_BY_NAME, + NETRA_SPAN_TYPE, + NETRA_USAGE_SOURCE, + TTS_METRICS_ATTRIBUTE, + USAGE_SOURCE_FRAMEWORK, + ConversationSide, + as_attribute_text, + content_of_choice_event, + content_of_event, + conversation_from_attributes, + is_absent, + is_usage_attribute, + is_zero_usage, + messages_for_parent, + messages_from_chat_ctx, + netra_span_type_for, + role_of_choice_event, + tts_pricing_attributes_from, +) + +logger = logging.getLogger(__name__) + +SetAttributeFunc = Callable[[str, Any], None] + +# The indexed key pair to write for each side of the conversation convention. +_KEYS_BY_SIDE: Dict[ConversationSide, Tuple[str, str]] = { + ConversationSide.PROMPT: (GEN_AI_PROMPT_ROLE, GEN_AI_PROMPT_CONTENT), + ConversationSide.COMPLETION: (GEN_AI_COMPLETION_ROLE, GEN_AI_COMPLETION_CONTENT), +} + +# Instance attribute holding a span's ``_ConversationRecorder``. Stored on the span +# itself so the registry in ``SpanMappingProcessor`` can stay a +# ``WeakValueDictionary`` keyed on span id: the span's own lifetime then decides how +# long the entry lives, with no risk of the processor pinning finished spans in +# memory. +_RECORDER_FIELD = "_netra_livekit_recorder" + + +def _is_livekit_span(span: Any) -> bool: + """Whether *span* was produced by livekit-agents' own instrumentation. + + Args: + span: The span to test. + + Returns: + True only for spans whose instrumentation scope is ``livekit-agents``. + """ + scope = getattr(span, "instrumentation_scope", None) + return getattr(scope, "name", None) == LIVEKIT_SCOPE_NAME + + +def _class_level_writer(span: Span) -> SetAttributeFunc: + """Return a writer that bypasses every instance-level wrapper on *span*. + + Args: + span: The span to write to. + + Returns: + A single-attribute writer calling ``type(span).set_attributes`` directly. + """ + class_set_attributes = type(span).set_attributes + + def write(key: str, value: Any) -> None: + """Write one attribute straight to the class method. + + Args: + key: The attribute name. + value: The attribute value. + """ + class_set_attributes(span, {key: value}) + + return write + + +def _write_tts_pricing(span: Span, metrics_payload: Any) -> None: + """Lift the priceable fields out of LiveKit's TTS metrics blob into Netra keys. + + Writes through ``span.set_attribute`` — the outermost wrapper — so the model + reaches the rest of the processor chain and the character count takes the + usage branch, which stamps ``netra.usage.source`` on it like every other + framework-reported usage number. + + Args: + span: The LiveKit span the metrics were written on (``tts_request``). + metrics_payload: The value of ``lk.tts_metrics``. + """ + pricing = tts_pricing_attributes_from(metrics_payload) + if pricing.model is not None: + span.set_attribute(GEN_AI_REQUEST_MODEL, pricing.model) + if pricing.character_count is not None: + span.set_attribute(GEN_AI_USAGE_CHARACTER_COUNT, pricing.character_count) + + +class _ConversationRecorder: + """Appends messages to one span's indexed gen_ai prompt/completion sequences. + + The single place an indexed conversation attribute is written, so every source + that contributes to a span — mapped ``lk.*`` attributes, an expanded chat + context, conversation events, and a child span's content — advances the same + counters and cannot overwrite another source's entries. One instance per + LiveKit span, created in ``SpanMappingProcessor.on_start``. + """ + + __slots__ = ("_span", "_next_index", "_truncated") + + def __init__(self, span: Span) -> None: + """Start both index sequences at zero for *span*. + + Args: + span: The span whose conversation this records. + """ + self._span = span + self._next_index: Dict[ConversationSide, Iterator[int]] = { + ConversationSide.PROMPT: itertools.count(), + ConversationSide.COMPLETION: itertools.count(), + } + self._truncated = False + + def append(self, side: ConversationSide, role: str, content: Any) -> None: + """Append one message to the given side of the conversation. + + Writes through ``span.set_attribute`` — the outermost wrapper — so the + values reach ``SpanIOProcessor``, which assembles them into + ``input``/``output``. + + Silently stops at ``MAX_CONVERSATION_MESSAGES_PER_SIDE`` and marks the + span instead — see that constant for why an unbounded sequence is not + merely wasteful but destructive. The cap lives here, rather than at each + call site, so it covers every source that feeds a recorder: mapped + ``lk.*`` attributes, an expanded chat context, conversation events, and a + child span's propagated content. + + Args: + side: Which indexed sequence to append to. + role: The conversation role to stamp alongside the text. + content: The message text. + """ + # The budget is read off the counter itself rather than a separate + # decrement: ``next()`` on an ``itertools.count`` is atomic, and a span's + # attributes can be written from more than one thread (a child ending on + # another thread propagates content up through here). + index = next(self._next_index[side]) + if index >= MAX_CONVERSATION_MESSAGES_PER_SIDE: + self._mark_truncated() + return + + role_key, content_key = _KEYS_BY_SIDE[side] + self._span.set_attribute(role_key.format(index=index), role) + self._span.set_attribute(content_key.format(index=index), as_attribute_text(content)) + + def _mark_truncated(self) -> None: + """Record on the span that the conversation was cut short by the cap. + + Written at most once. The guard is not synchronised: two threads racing + here both write the same value, so the only cost is a duplicate write. + """ + if self._truncated: + return + self._truncated = True + self._span.set_attribute(NETRA_CONVERSATION_TRUNCATED, True) + + def append_attribute(self, key: str, value: Any) -> bool: + """Route an ``lk.*`` conversation-content attribute into the sequences. + + Args: + key: The LiveKit attribute name being written. + value: The value being written. + + Returns: + True when *key* belongs to the conversation convention — whether or not + it carried a value — so the caller knows not to fall through to + ``ATTRIBUTE_MAP``. + """ + if key == CHAT_CTX_ATTRIBUTE: + messages = messages_from_chat_ctx(value) + # Keep the newest turns. ``append`` caps the sequence either way, but it + # can only drop what arrives last, so feeding it the whole context + # oldest-first would preserve the opening of the call and discard the + # turns this span is actually about. The full context stays on the span + # verbatim as ``lk.chat_ctx``. + if len(messages) > MAX_CONVERSATION_MESSAGES_PER_SIDE: + self._mark_truncated() + messages = messages[-MAX_CONVERSATION_MESSAGES_PER_SIDE:] + for role, content in messages: + self.append(ConversationSide.PROMPT, role, content) + return True + + target = CONVERSATION_MAP.get(key) + if target is None: + return False + if not is_absent(value): + self.append(target.side, target.role, value) + return True + + def append_event(self, name: str, attributes: Attributes) -> None: + """Route a LiveKit conversation event into the sequences. + + Args: + name: The event name LiveKit passed to ``add_event``. + attributes: The event attributes. Events that are not conversation + content, or that carry no text, contribute nothing. + """ + if name == EVENT_CHOICE: + content = content_of_choice_event(attributes) + if content: + self.append(ConversationSide.COMPLETION, role_of_choice_event(attributes), content) + return + + role = EVENT_ROLE.get(name) + if role is None: + return + content = content_of_event(attributes) + if content: + self.append(ConversationSide.PROMPT, role, content) + + def append_child_conversation(self, child: ReadableSpan) -> None: + """Append a finished child span's conversation content. + + Args: + child: The span that has ended directly beneath this one. + """ + conversation = conversation_from_attributes(child.attributes) + # A child no LLM-aware instrumentation touched carries an ``input`` that is + # not a conversation at all — an HTTP envelope, a SQL statement — and must + # not be copied up as if it were one. + allow_raw_io = conversation.carries_gen_ai or _is_livekit_span(child) + for message in messages_for_parent(conversation, allow_raw_io=allow_raw_io): + self.append(message.side, message.role, message.content) + + +class SpanMappingProcessor(SpanProcessor): # type: ignore[misc] + """Mirrors LiveKit's ``lk.*`` attributes and conversation events into Netra keys. + + Additive throughout: an ``lk.*`` attribute is never deleted or rewritten, and a + conversation event is always still recorded on the span. The one exception is + a zero token count, which is dropped rather than mirrored — see + ``is_zero_usage``. + + Conversation content — LiveKit's own ``lk.*`` attributes, its serialised chat + contexts, and its conversation events — all land in the indexed + ``gen_ai.prompt.*``/``gen_ai.completion.*`` pair that ``SpanIOProcessor`` + assembles into ``input``/``output``. That is the convention every other Netra + instrumentation emits, so a voice turn renders like any other span. + + Two spans carry no conversation content of their own and inherit it from a + direct child when that child ends — see ``IO_FROM_CHILD_SPAN_NAMES`` and + ``on_end``. + + Two values are additionally *derived* rather than mirrored: the model and the + character count that price a TTS call, which LiveKit reports only inside the + opaque ``lk.tts_metrics`` JSON blob — see ``_write_tts_pricing``. + """ + + def __init__(self) -> None: + """Create the registry of spans awaiting content from a child.""" + # Weak values: an entry costs nothing once the span itself is collected, so + # a span that somehow never ends cannot leak. Guarded by a lock because + # spans can start and end on threads other than the agent's event loop. + self._io_parents: "weakref.WeakValueDictionary[int, Span]" = weakref.WeakValueDictionary() + self._io_parents_lock = threading.Lock() + + def on_start(self, span: Span, parent_context: Optional[otel_context.Context] = None) -> None: + """Stamp the Netra markers and install the mapping wrappers on a LiveKit span. + + Args: + span: The span that was started. + parent_context: The parent context (unused). + """ + try: + if not _is_livekit_span(span): + return + self._stamp_markers(span) + + recorder = _ConversationRecorder(span) + setattr(span, _RECORDER_FIELD, recorder) + self._wrap_set_attribute(span, recorder) + self._wrap_add_event(span, recorder) + + if span.name in IO_FROM_CHILD_SPAN_NAMES: + self._register_io_parent(span) + except Exception: + logger.warning("netra.livekit: span mapping could not be installed", exc_info=True) + + def on_end(self, span: ReadableSpan) -> None: + """Copy a finished span's conversation content up to its parent, if wanted. + + Deliberately *not* gated on ``_is_livekit_span``: the child holding the + content is usually the provider's own span (``openai.chat`` and friends), + which belongs to another instrumentation scope entirely. + + Never touches *this* span — see the module docstring. It only appends to a + still-recording parent, which the exporter has not seen yet. + + Args: + span: The span that has ended. + """ + try: + self._propagate_content_to_parent(span) + except Exception: + logger.debug("netra.livekit: content propagation to the parent span failed", exc_info=True) + try: + self._deregister_io_parent(span) + except Exception: + logger.debug("netra.livekit: span could not be deregistered", exc_info=True) + + def force_flush(self, timeout_millis: int = 30000) -> bool: + """No-op flush. + + Args: + timeout_millis: Maximum time to wait (unused). + + Returns: + Always True. + """ + return True + + def shutdown(self) -> None: + """No-op shutdown.""" + + @staticmethod + def _stamp_markers(span: Span) -> None: + """Write the Netra classification markers a LiveKit span's name implies. + + Args: + span: The LiveKit span to stamp. + """ + span.set_attribute(NETRA_SPAN_TYPE, netra_span_type_for(span.name)) + + entity_type = NETRA_ENTITY_TYPE_BY_NAME.get(span.name) + if entity_type is not None: + span.set_attribute(NETRA_ENTITY_TYPE, entity_type) + + audio_type = AUDIO_TYPE_BY_SPAN_NAME.get(span.name) + if audio_type is not None: + span.set_attribute(NETRA_AUDIO_TYPE, audio_type) + + def _register_io_parent(self, span: Span) -> None: + """Record *span* as one whose conversation content arrives from a child. + + Args: + span: The LiveKit span to register. + """ + context = span.get_span_context() + if context is None: + return + with self._io_parents_lock: + self._io_parents[context.span_id] = span + + def _propagate_content_to_parent(self, span: ReadableSpan) -> None: + """Append *span*'s conversation content to its parent's gen_ai sequences. + + Args: + span: The span that has ended. + """ + parent_context = span.parent + if parent_context is None or not self._io_parents: + return + with self._io_parents_lock: + parent = self._io_parents.get(parent_context.span_id) + if parent is None or not parent.is_recording(): + return + + recorder: Optional[_ConversationRecorder] = getattr(parent, _RECORDER_FIELD, None) + if recorder is None: + return + recorder.append_child_conversation(span) + + def _deregister_io_parent(self, span: ReadableSpan) -> None: + """Drop *span* from the registry if it was awaiting content from a child. + + Args: + span: The span that has ended. + """ + if span.name not in IO_FROM_CHILD_SPAN_NAMES: + return + context = span.get_span_context() + if context is None: + return + with self._io_parents_lock: + self._io_parents.pop(context.span_id, None) + + @staticmethod + def _wrap_set_attribute(span: Span, recorder: _ConversationRecorder) -> None: + """Wrap ``span.set_attribute`` so mapped ``lk.*`` writes also write Netra keys. + + Chains through the previously-installed wrapper rather than the class + method, so writes still pass down through ``SpanIOProcessor`` and + ``InstrumentationSpanProcessor``. ``set_attributes`` (plural) is wrapped + too because the OTel SDK writes it straight to ``_attributes`` without + going through ``set_attribute`` — LiveKit uses it, e.g. for the + ``gen_ai.*`` request attributes on ``llm_request`` and for + ``lk.user_transcript`` on ``user_turn``. + + Args: + span: The LiveKit span to wrap. + recorder: The span's conversation recorder. + """ + previous: SetAttributeFunc = span.set_attribute + if "set_attribute" not in vars(span): + # Nothing has wrapped this span, so ``previous`` is the raw SDK method + # — and from opentelemetry-sdk 1.41 that method is implemented as + # ``self.set_attributes({key: value})``, which resolves the plural + # wrapper installed below and recurses until RecursionError. In Netra's + # own pipeline this branch never runs (``InstrumentationSpanProcessor`` + # always wraps first, and terminates its own writes at the class + # method for exactly this reason); it keeps the processor correct when + # it is registered on a provider by itself. + previous = _class_level_writer(span) + + def map_attribute(key: str, value: Any) -> None: + """Write *key* through, then write whatever Netra key it implies. + + Args: + key: The attribute name LiveKit is writing. + value: The attribute value LiveKit is writing. + """ + if is_usage_attribute(key): + if is_zero_usage(value): + return + previous(key, value) + # Marks whose accounting this is, so the backend can prefer a + # provider span's tokens over the framework's for the same call. + previous(NETRA_USAGE_SOURCE, USAGE_SOURCE_FRAMEWORK) + return + + previous(key, value) + + if key == TTS_METRICS_ATTRIBUTE: + # Not conversation content and not in ATTRIBUTE_MAP: the blob + # is forwarded as-is and its priceable fields are lifted out. + _write_tts_pricing(span, value) + return + + if recorder.append_attribute(key, value): + return + + target = ATTRIBUTE_MAP.get(key) + if target is None or is_absent(value): + return + previous(target, value) + + def patched_set_attribute(key: str, value: Any) -> None: + """Map *key* onto its Netra keys, falling back to a plain write. + + Args: + key: The attribute name LiveKit is writing. + value: The attribute value LiveKit is writing. + """ + try: + map_attribute(key, value) + except Exception: + logger.debug("netra.livekit: attribute mapping failed for %s", key, exc_info=True) + try: + previous(key, value) + except Exception: + logger.debug("netra.livekit: set_attribute failed for %s", key, exc_info=True) + + def patched_set_attributes(attributes: Mapping[str, Any]) -> None: + """Route a bulk write through the single-attribute mapping. + + Args: + attributes: The attributes LiveKit is writing. + """ + for key, value in (attributes or {}).items(): + patched_set_attribute(key, value) + + setattr(span, "set_attribute", patched_set_attribute) + setattr(span, "set_attributes", patched_set_attributes) + + @staticmethod + def _wrap_add_event(span: Span, recorder: _ConversationRecorder) -> None: + """Wrap ``span.add_event`` so conversation events become attributes. + + LiveKit emits conversation content as span *events* + (``_chat_ctx_to_otel_events`` for the request, ``gen_ai.choice`` for the + reply), which a ``set_attribute`` wrapper structurally cannot see — so + LiveKit spans would otherwise export with empty ``input``/``output``. + + Assigning ``span.add_event`` shadows the class method, because + ``add_event`` is not a dunder and attribute lookup hits the instance dict. + + Args: + span: The LiveKit span to wrap. + recorder: The span's conversation recorder. + """ + original = span.add_event + + def patched_add_event( + name: str, + attributes: Attributes = None, + timestamp: Optional[int] = None, + ) -> None: + """Record the event as conversation content, then forward it verbatim. + + Args: + name: The event name LiveKit passed to ``add_event``. + attributes: The event attributes, if any. + timestamp: The event timestamp, if any. Forwarded untouched. + """ + try: + recorder.append_event(name, attributes) + except Exception: + logger.debug("netra.livekit: event -> attribute mapping failed for %s", name, exc_info=True) + # ALWAYS forward: the user's event must be recorded whatever happens + # on our side. + original(name, attributes, timestamp) + + setattr(span, "add_event", patched_add_event) diff --git a/netra/instrumentation/livekit/utils.py b/netra/instrumentation/livekit/utils.py new file mode 100644 index 0000000..e384974 --- /dev/null +++ b/netra/instrumentation/livekit/utils.py @@ -0,0 +1,796 @@ +"""Pure mapping tables and helpers for the LiveKit span processor. + +Free of OTel *tracing* imports so the mapping rules can be unit-tested as plain +data — the one exception is Netra's own ``SpanType`` enum, imported so the +``netra.span.type`` values here cannot drift from the vocabulary every other +Netra instrumentation stamps. Every table here was checked against +``livekit-agents`` 1.6.7 (``telemetry/trace_types.py`` for the attribute names, +``telemetry/traces.py`` ``_chat_ctx_to_otel_events`` for the event shape). + +Layout: constants, then the types the mapping tables are built from, then the +tables themselves, then one section of helpers per payload LiveKit produces +(span attributes, conversation events, chat contexts, TTS metrics). +""" + +import json +import re +from enum import Enum +from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Tuple + +from netra.span_wrapper import SpanType + +# --------------------------------------------------------------------------- +# LiveKit instrumentation scope +# --------------------------------------------------------------------------- + +# livekit-agents' OTel instrumentation scope. Everything this package does is +# gated on it: our processors are registered process-wide and must not touch a +# span from any other instrumentation. +LIVEKIT_SCOPE_NAME = "livekit-agents" + +# --------------------------------------------------------------------------- +# Netra target attribute keys +# --------------------------------------------------------------------------- + +NETRA_TOOL_NAME = "netra.tool.name" +NETRA_USAGE_SOURCE = "netra.usage.source" +USAGE_SOURCE_FRAMEWORK = "framework" + +# The ``netra.span.type`` every other Netra instrumentation stamps +# (``hermes_agent``, ``google_adk``, ``agno``, ``claude_agent_sdk``). The only +# span-type contract this package emits: LiveKit spans carry no package-local +# ``span_type`` attribute. +NETRA_SPAN_TYPE = "netra.span.type" + +# ``SpanType`` has no TTS or STT member, so the audio spans take this default +# rather than being given a value that means something else. +DEFAULT_NETRA_SPAN_TYPE = SpanType.SPAN + +# The entity marker the ``@workflow``/``@agent``/``@task`` decorators stamp +# (``netra/decorators.py:_add_span_attributes``) and that ``agno`` emits as +# ``ATTR_ENTITY``. Separate from ``netra.span.type``: ``SpanType`` has no +# ``WORKFLOW`` member, so the workflow marking rides on the entity contract while +# the span type stays at its ``SPAN`` default. +NETRA_ENTITY_TYPE = "netra.entity.type" +ENTITY_TYPE_WORKFLOW = "workflow" + +# The audio marker, written on the interaction-level spans named in +# ``AUDIO_TYPE_BY_SPAN_NAME`` rather than on every LiveKit span. Its value says at +# which granularity the call audio for that span is addressable: the whole call +# (``session``) versus a single turn (``span``). +NETRA_AUDIO_TYPE = "netra.audio.type" +AUDIO_TYPE_SESSION = "session" +AUDIO_TYPE_SPAN = "span" + +# --------------------------------------------------------------------------- +# The gen_ai conventions this package emits into +# --------------------------------------------------------------------------- + +# The conversation-attribute convention SpanIOProcessor already consumes +# (``_PROMPT_RE`` in netra/processors/span_io_processor.py). Emitting into this +# shape rather than inventing a third convention is what makes voice turns render +# like every other LLM span. +GEN_AI_PROMPT_ROLE = "gen_ai.prompt.{index}.role" +GEN_AI_PROMPT_CONTENT = "gen_ai.prompt.{index}.content" + +# The completion-side counterpart (``_COMPLETION_RE`` in the same processor), +# which fills ``output``. +GEN_AI_COMPLETION_ROLE = "gen_ai.completion.{index}.role" +GEN_AI_COMPLETION_CONTENT = "gen_ai.completion.{index}.content" + +# Marks a span as one an LLM-aware instrumentation wrote. Gates the verbatim +# ``input``/``output`` fallback in ``messages_for_parent``: without it, an +# ``HTTP POST`` span under ``llm_request_run`` (netra/instrumentation/httpx/utils.py:124 +# writes the URL, headers and body into ``input``) would be copied up as if it +# were a user message. +GEN_AI_ATTRIBUTE_PREFIX = "gen_ai." + +# Prefix identifying token-usage attributes, whoever wrote them. +GEN_AI_USAGE_PREFIX = "gen_ai.usage." + +# The pair of keys Netra's backend prices a TTS call from. Identical to what every +# Netra TTS provider instrumentation emits (``cartesia``, ``elevenlabs``, +# ``deepgram``), so a LiveKit-hosted synthesis prices through the same path as a +# directly-instrumented one. +GEN_AI_REQUEST_MODEL = "gen_ai.request.model" +GEN_AI_USAGE_CHARACTER_COUNT = "gen_ai.usage.prompt.character_count" + +# The assembled input/output ``SpanIOProcessor`` builds from the indexed pairs. Read +# off a child span as the fallback when it carries no indexed pairs of its own. +INPUT_ATTRIBUTE = "input" +OUTPUT_ATTRIBUTE = "output" + +# The fallback carries text with no role attached, so one has to be supplied. Named +# for what the side means to the parent: its request and its reply. +FALLBACK_PROMPT_ROLE = "user" +FALLBACK_COMPLETION_ROLE = "assistant" + +# The most conversation messages this package writes onto one span, per side. +# +# A span's attribute capacity is bounded — ``OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT``, +# default 128 — and OTel's ``BoundedAttributes`` evicts the *oldest* entry on +# overflow. So an unbounded conversation does not merely truncate itself: it +# silently deletes the attributes written earliest, which are exactly the markers +# ``SpanMappingProcessor.on_start`` and ``SessionSpanProcessor`` stamp +# (``netra.span.type``, ``netra.instrumentation.name``, ``netra.session_id``). +# +# Two LiveKit sources grow with the length of the call, both verified against +# livekit-agents 1.6.7: +# * ``lk.chat_ctx`` on ``llm_node`` — the whole serialised ChatContext +# (``voice/generation.py``); +# * one conversation event per context item on ``llm_request`` +# (``llm/llm.py`` -> ``_chat_ctx_to_otel_events``). +# Each message costs two attributes (role + content), so without a cap a ~30-turn +# call is enough to evict every marker. +# +# 20 per side is 80 attributes at worst, which leaves the markers, LiveKit's own +# attributes and the latencies inside the default budget. Nothing is lost that is +# not still on the span: the full context remains verbatim in ``lk.chat_ctx``. +# LiveKit bounds its own ``eou_detection`` context the same way, via +# ``_EOU_MAX_HISTORY_TURNS``. +MAX_CONVERSATION_MESSAGES_PER_SIDE = 20 + +# Marks a span whose conversation was cut short by the cap above, so the +# truncation is visible on the span rather than silent. Written once per span. +NETRA_CONVERSATION_TRUNCATED = "netra.conversation.truncated" + +# --------------------------------------------------------------------------- +# LiveKit attributes this package reads by name +# --------------------------------------------------------------------------- + +# The attribute holding a serialised ``ChatContext`` (on ``llm_node`` and +# ``eou_detection``). Expanded into indexed ``gen_ai.prompt.*`` attributes rather +# than mirrored verbatim into ``input``: the raw JSON also contains non-message +# items (``agent_config_update``, handoffs) and reads as an opaque blob. +CHAT_CTX_ATTRIBUTE = "lk.chat_ctx" + +# LiveKit's serialised ``TTSMetrics`` (``trace_types.ATTR_TTS_METRICS``, written on +# ``tts_request``). It is the only place on that span carrying the two values +# pricing needs — ``characters_count`` and the model name, nested under +# ``metadata`` — and as one opaque JSON blob the backend cannot read either. The +# sibling ``tts_node`` span does carry ``gen_ai.request.model``, but pricing needs +# the model and the character count on the *same* span. +TTS_METRICS_ATTRIBUTE = "lk.tts_metrics" + +# LiveKit's completion event (``trace_types.EVENT_GEN_AI_CHOICE``, emitted from +# ``llm/llm.py`` once the reply is complete). Handled separately from +# ``EVENT_ROLE`` because it carries the model's reply and so belongs in the +# completion convention — without it, ``llm_request`` exports with an empty +# ``output`` even though the reply is right there on the span. +EVENT_CHOICE = "gen_ai.choice" + +# Role LiveKit puts on the choice event; only used if the event omits it. +DEFAULT_CHOICE_ROLE = "assistant" + + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + + +class ConversationSide(Enum): + """Which half of the ``gen_ai`` conversation convention a value belongs to. + + ``PROMPT`` feeds ``input``, ``COMPLETION`` feeds ``output`` — the assembly + happens in ``SpanIOProcessor``, which this package only has to emit into. + """ + + PROMPT = "prompt" + COMPLETION = "completion" + + +class ConversationTarget(NamedTuple): + """The gen_ai slot an ``lk.*`` content attribute is mirrored into. + + Attributes: + side: Whether the text is a prompt message or a completion message. + role: The conversation role to stamp alongside the text. Named for the + *speaker*, not for the span's position in the pipeline, so a chat + preview assembled from these attributes reads as the actual dialogue. + """ + + side: ConversationSide + role: str + + +class ConversationMessage(NamedTuple): + """One message to append to a span's indexed gen_ai sequences. + + Attributes: + side: Which sequence the message belongs in. + role: The conversation role to stamp alongside the text. + content: The message text. + """ + + side: ConversationSide + role: str + content: str + + +class SpanConversation(NamedTuple): + """The conversation content read back off a finished span. + + Attributes: + prompts: The ``(role, content)`` pairs of the indexed prompt sequence, in + index order. + completions: The same for the completion sequence. + raw_input: The span's assembled ``input``, but only when it carries no + indexed prompt pairs — otherwise it is derived from them and copying + both would duplicate the conversation. + raw_output: The same for ``output``. + carries_gen_ai: Whether the span has any ``gen_ai.*`` attribute at all, + i.e. whether an LLM-aware instrumentation wrote it. + """ + + prompts: List[Tuple[str, str]] + completions: List[Tuple[str, str]] + raw_input: Optional[str] + raw_output: Optional[str] + carries_gen_ai: bool + + +class TtsPricingAttributes(NamedTuple): + """The billable facts of one TTS synthesis, as LiveKit reported them. + + Attributes: + model: The synthesis model, verbatim from LiveKit — including the + ``provider/model`` prefix it uses for its inference gateway + (``cartesia/sonic-3``). ``None`` when LiveKit reported none. + character_count: The number of characters synthesised, or ``None`` when + LiveKit reported none or a count of zero. + """ + + model: Optional[str] + character_count: Optional[int] + + +# --------------------------------------------------------------------------- +# Mapping tables +# --------------------------------------------------------------------------- + +# lk.* -> Netra key. Additive: the original lk.* attribute is always preserved. +# +# Conversation *content* is deliberately absent from this table — see +# ``CONVERSATION_MAP``, which routes it through the indexed ``gen_ai.*`` +# convention instead of writing ``input``/``output`` directly. +ATTRIBUTE_MAP: Dict[str, str] = { + # Function tools + "lk.function_tool.name": NETRA_TOOL_NAME, + "lk.function_tool.arguments": INPUT_ATTRIBUTE, + "lk.function_tool.output": OUTPUT_ATTRIBUTE, + # Latencies, all in seconds + "lk.response.ttft": "netra.latency.ttft", + "lk.response.ttfb": "netra.latency.ttfb", + "lk.e2e_latency": "netra.latency.e2e", + "lk.end_of_turn_delay": "netra.latency.end_of_turn_delay", + # Turn quality + "lk.transcript_confidence": "netra.stt.confidence", + "lk.interrupted": "netra.turn.interrupted", +} + +# lk.* content attribute -> gen_ai slot. One table for every LiveKit span: +# ``lk.response.text`` means the same thing on ``agent_turn`` as it does on +# ``llm_node``, so nothing here needs gating on the span name. +# +# Additive: the original lk.* attribute is always preserved, and these never write +# ``input``/``output`` directly — indexed ``gen_ai.prompt.*``/``gen_ai.completion.*`` +# attributes are emitted instead, the same convention Netra's own provider +# instrumentations use, which is what lets ``SpanIOProcessor`` assemble a +# multi-message ``input``. +CONVERSATION_MAP: Dict[str, ConversationTarget] = { + # agent_turn: the system instructions in force for the turn. LiveKit writes it + # before ``lk.user_input``, so it lands at prompt index 0 and the assembled + # ``input`` reads [system, user] like any other LLM span. + "lk.instructions": ConversationTarget(ConversationSide.PROMPT, "system"), + # agent_turn: the utterance that opened the turn (LiveKit sets this only when a + # new message did). + "lk.user_input": ConversationTarget(ConversationSide.PROMPT, "user"), + # agent_turn and llm_node: the generated reply. + "lk.response.text": ConversationTarget(ConversationSide.COMPLETION, "assistant"), + # tts_request: the text handed to the TTS provider. The words are the agent's. + "lk.input_text": ConversationTarget(ConversationSide.PROMPT, "assistant"), + # user_turn: the STT transcript — the output of the transcription. The words are + # the caller's. + "lk.user_transcript": ConversationTarget(ConversationSide.COMPLETION, "user"), +} + +# span name -> ``netra.span.type``. +NETRA_SPAN_TYPE_BY_NAME: Dict[str, SpanType] = { + "agent_turn": SpanType.AGENT, + "llm_node": SpanType.GENERATION, + "llm_request": SpanType.GENERATION, + "function_tool": SpanType.TOOL, + "tts_request": SpanType.GENERATION, +} + +# span name -> ``netra.entity.type``. ``job_entrypoint`` is livekit-agents' own +# root span for a job (``ipc/job_proc_lazy_main.py``: ``_traceable_entrypoint``), +# so it wraps everything the user's entrypoint does — the agent session, and any +# work before or after it — which is exactly a workflow. +NETRA_ENTITY_TYPE_BY_NAME: Dict[str, str] = { + "job_entrypoint": ENTITY_TYPE_WORKFLOW, +} + +# LiveKit span name -> the ``netra.audio.type`` value it carries. Matched against +# the LiveKit span name, so only spans this package already gates on (scope +# ``livekit-agents``) are eligible — a nested provider span such as ``openai.chat`` +# never reaches the lookup. +AUDIO_TYPE_BY_SPAN_NAME: Dict[str, str] = { + "agent_session": AUDIO_TYPE_SESSION, + "agent_turn": AUDIO_TYPE_SPAN, + "user_turn": AUDIO_TYPE_SPAN, +} + +# LiveKit spans that carry no conversation content of their own: the text exists +# only on a direct child. ``llm_request_run`` wraps the provider call, so the +# prompt and completion are on the provider's own span (``openai.chat`` and +# friends, a *non*-LiveKit scope); ``tts_node`` wraps the synthesis, so the text +# is on ``tts_request``. See ``SpanMappingProcessor.on_end``. +# +# Verified against livekit-agents 1.6.7 that in both cases the child ends while +# the parent is still recording: the provider span ends inside +# ``LLMStream._run()``, and ``tts_request`` is ended by the ``async with +# wrapped_tts.stream()`` exit inside the generator ``_tts_inference_task`` +# iterates. +IO_FROM_CHILD_SPAN_NAMES = frozenset({"llm_request_run", "tts_node"}) + +# LiveKit conversation event name -> gen_ai role. From ``trace_types.EVENT_*``; +# note LiveKit folds OpenAI's ``developer`` role into the system message event. +# These are all request-side messages, hence the prompt convention. +EVENT_ROLE: Dict[str, str] = { + "gen_ai.system.message": "system", + "gen_ai.user.message": "user", + "gen_ai.assistant.message": "assistant", + "gen_ai.tool.message": "tool", +} + +# --------------------------------------------------------------------------- +# Payload field names (private) +# --------------------------------------------------------------------------- + +# Reads an indexed conversation attribute back off a span, for the child-to-parent +# propagation in ``conversation_from_attributes``. The plural forms match what +# ``SpanIOProcessor`` accepts, so a child written by any instrumentation in the SDK +# is readable here. +_INDEXED_MESSAGE_RE = re.compile(r"^gen_ai\.(prompt|completion)s?\.(\d+)\.(role|content)$") + +_PROMPT_GROUP = "prompt" +_ROLE_FIELD = "role" +_CONTENT_FIELD = "content" + +# The attribute key LiveKit puts message text under in a conversation event +# (``_chat_ctx_to_otel_events``: ``{"content": item.raw_text_content or ""}``). +_EVENT_CONTENT_KEY = "content" +_EVENT_ROLE_KEY = "role" + +# On the choice event, a tool-only reply carries no ``content`` — the requested +# calls are the whole output. LiveKit sends them as a list of JSON strings. +_EVENT_TOOL_CALLS_KEY = "tool_calls" + +_CHAT_CTX_ITEMS_KEY = "items" +_CHAT_CTX_MESSAGE_TYPE = "message" +_CHAT_CTX_TYPE_KEY = "type" +_CHAT_CTX_ROLE_KEY = "role" +_CHAT_CTX_CONTENT_KEY = "content" + +_TTS_METRICS_CHARACTERS_KEY = "characters_count" +_TTS_METRICS_METADATA_KEY = "metadata" +_TTS_METRICS_MODEL_KEY = "model_name" + + +# --------------------------------------------------------------------------- +# Value helpers +# --------------------------------------------------------------------------- + + +def is_absent(value: Any) -> bool: + """Whether *value* should be treated as "not set". + + Empty and missing values are treated as absent so a mapped write can never + blank out a value another processor supplied. + + Args: + value: The candidate attribute value. + + Returns: + True if the value carries no information. + """ + return value is None or value == "" + + +def as_attribute_text(value: Any) -> str: + """Render a value as the string an OTel text attribute needs. + + Args: + value: The value LiveKit wrote, or one read back off a span. + + Returns: + The value unchanged if it is already a string, else its ``str()``. + """ + return value if isinstance(value, str) else str(value) + + +def is_usage_attribute(key: str) -> bool: + """Whether *key* is a token-usage attribute. + + Args: + key: An attribute name. + + Returns: + True for ``gen_ai.usage.*`` keys. + """ + return key.startswith(GEN_AI_USAGE_PREFIX) + + +def is_zero_usage(value: Any) -> bool: + """Whether a usage value is a zero that is worse than no value at all. + + A framework that cannot surface token counts reports 0 rather than omitting + them, and livekit-agents forwards ``metrics.prompt_tokens`` verbatim — so a + custom LLM node that does not report usage produces + ``gen_ai.usage.input_tokens = 0``. Writing that claims a measurement nobody + made, and the real counts are on the provider span underneath. Dropping the + attribute lets the provider's numbers stand unopposed. + + Args: + value: The candidate usage value. + + Returns: + True for a numeric zero; False for every other value, including ``None`` + and booleans. + """ + if isinstance(value, bool): + return False + if isinstance(value, (int, float)): + return value == 0 + return False + + +def netra_span_type_for(span_name: Optional[str]) -> str: + """Return the ``netra.span.type`` value for a LiveKit span name. + + Args: + span_name: The LiveKit span's name, or ``None``. + + Returns: + A ``SpanType`` value — ``AGENT``, ``GENERATION``, ``TOOL``, or ``SPAN``. + """ + return NETRA_SPAN_TYPE_BY_NAME.get(span_name or "", DEFAULT_NETRA_SPAN_TYPE).value + + +# --------------------------------------------------------------------------- +# Conversation content on spans +# --------------------------------------------------------------------------- + + +def conversation_from_attributes(attributes: Optional[Mapping[str, Any]]) -> SpanConversation: + """Read a finished span's conversation content back out of its attributes. + + The inverse of what this package (and every other Netra instrumentation) writes + when it emits indexed ``gen_ai.prompt.*``/``gen_ai.completion.*`` pairs, so a + child span's conversation can be re-emitted onto its parent. + + An entry carrying a role but no text is skipped — a role alone is not a message. + An entry carrying text but no role keeps the text; the caller supplies a role. + + Args: + attributes: The finished span's attributes, or ``None``. + + Returns: + The prompt and completion pairs in index order, the verbatim + ``input``/``output`` for the sides that have no pairs, and whether the span + carries any ``gen_ai.*`` attribute. A malformed or empty mapping yields an + empty result rather than an error — a mapping failure must never break the + user's trace. + """ + prompt_entries: Dict[int, Dict[str, str]] = {} + completion_entries: Dict[int, Dict[str, str]] = {} + raw_input: Optional[str] = None + raw_output: Optional[str] = None + carries_gen_ai = False + + for key, value in (attributes or {}).items(): + if key == INPUT_ATTRIBUTE: + raw_input = None if is_absent(value) else as_attribute_text(value) + continue + if key == OUTPUT_ATTRIBUTE: + raw_output = None if is_absent(value) else as_attribute_text(value) + continue + if not key.startswith(GEN_AI_ATTRIBUTE_PREFIX): + continue + carries_gen_ai = True + match = _INDEXED_MESSAGE_RE.match(key) + if match is None: + continue + entries = prompt_entries if match.group(1) == _PROMPT_GROUP else completion_entries + entries.setdefault(int(match.group(2)), {})[match.group(3)] = as_attribute_text(value) + + prompts = _ordered_messages(prompt_entries) + completions = _ordered_messages(completion_entries) + return SpanConversation( + prompts=prompts, + completions=completions, + raw_input=raw_input if not prompts else None, + raw_output=raw_output if not completions else None, + carries_gen_ai=carries_gen_ai, + ) + + +def messages_for_parent(conversation: SpanConversation, *, allow_raw_io: bool) -> List[ConversationMessage]: + """Turn a child's conversation into the messages to append to its parent. + + Args: + conversation: What ``conversation_from_attributes`` read off the child. + allow_raw_io: Whether the verbatim ``input``/``output`` fallback may be + used. False for a child no LLM-aware instrumentation touched, whose + ``input`` is something else entirely — an HTTP request envelope, a SQL + statement — and would read as a fabricated user message on the parent. + + Returns: + The messages to append, prompts first. Empty when the child carried no + conversation. + """ + messages = [ + ConversationMessage(ConversationSide.PROMPT, role or FALLBACK_PROMPT_ROLE, content) + for role, content in conversation.prompts + ] + messages.extend( + ConversationMessage(ConversationSide.COMPLETION, role or FALLBACK_COMPLETION_ROLE, content) + for role, content in conversation.completions + ) + if not allow_raw_io: + return messages + + if conversation.raw_input is not None: + messages.append(ConversationMessage(ConversationSide.PROMPT, FALLBACK_PROMPT_ROLE, conversation.raw_input)) + if conversation.raw_output is not None: + messages.append( + ConversationMessage(ConversationSide.COMPLETION, FALLBACK_COMPLETION_ROLE, conversation.raw_output) + ) + return messages + + +def _ordered_messages(entries: Mapping[int, Mapping[str, str]]) -> List[Tuple[str, str]]: + """Flatten indexed role/content entries into ``(role, content)`` pairs. + + Args: + entries: Index -> the fields collected for that index. + + Returns: + The pairs in index order, skipping any index that carried no text. The + indices themselves are discarded: the caller re-numbers against the + parent's own counters. + """ + pairs: List[Tuple[str, str]] = [] + for index in sorted(entries): + fields = entries[index] + content = fields.get(_CONTENT_FIELD) + if content is None or content == "": + continue + pairs.append((fields.get(_ROLE_FIELD) or "", content)) + return pairs + + +# --------------------------------------------------------------------------- +# LiveKit conversation events +# --------------------------------------------------------------------------- + + +def content_of_event(attributes: Optional[Mapping[str, Any]]) -> Optional[str]: + """Extract the message text from a LiveKit conversation event's attributes. + + Args: + attributes: The event attributes LiveKit passed to ``add_event``. + + Returns: + The message text, or ``None`` when the event carries none. A + ``function_call`` event legitimately has no ``content`` — its payload is + already on the ``function_tool`` span as ``lk.function_tool.*``, so + returning ``None`` here drops nothing from the trace. + """ + if not attributes: + return None + content = attributes.get(_EVENT_CONTENT_KEY) + if is_absent(content): + return None + return as_attribute_text(content) + + +def content_of_choice_event(attributes: Optional[Mapping[str, Any]]) -> Optional[str]: + """Extract the reply text from LiveKit's ``gen_ai.choice`` event. + + Args: + attributes: The event attributes LiveKit passed to ``add_event``. + + Returns: + The reply text; the serialised tool calls when the reply was tool-only; + or ``None`` when the event carries neither. + """ + if not attributes: + return None + + content = attributes.get(_EVENT_CONTENT_KEY) + if not is_absent(content): + return as_attribute_text(content) + + return _joined_tool_calls(attributes.get(_EVENT_TOOL_CALLS_KEY)) + + +def role_of_choice_event(attributes: Optional[Mapping[str, Any]]) -> str: + """Return the role LiveKit put on a ``gen_ai.choice`` event. + + Args: + attributes: The event attributes LiveKit passed to ``add_event``. + + Returns: + The event's ``role``, or ``assistant`` when it carries none. + """ + role = (attributes or {}).get(_EVENT_ROLE_KEY) + if isinstance(role, str) and role: + return role + return DEFAULT_CHOICE_ROLE + + +def _joined_tool_calls(tool_calls: Any) -> Optional[str]: + """Render LiveKit's list of JSON tool-call strings as one JSON array. + + Args: + tool_calls: The event's ``tool_calls`` value. + + Returns: + A JSON array string, or ``None`` when there is nothing to render. + """ + if isinstance(tool_calls, str): + return tool_calls or None + if not isinstance(tool_calls, (list, tuple)) or not tool_calls: + return None + # Each element is already a JSON object string, so concatenating them into an + # array yields valid JSON. + return "[" + ", ".join(str(call) for call in tool_calls) + "]" + + +# --------------------------------------------------------------------------- +# LiveKit chat contexts +# --------------------------------------------------------------------------- + + +def messages_from_chat_ctx(payload: Any) -> List[Tuple[str, str]]: + """Extract ``(role, text)`` pairs from a serialised LiveKit ``ChatContext``. + + Accepts either the JSON string LiveKit puts in ``lk.chat_ctx`` or the dict + ``ChatContext.to_dict()`` returns, so the same rules apply whether the + context is read off a span or off a live session. + + Non-message items (``agent_config_update``, ``function_call``, + ``agent_handoff``, ...) are skipped: they are not conversation turns, and + their payloads are already on the spans that produced them. + + Args: + payload: A ``ChatContext`` JSON string or dict. + + Returns: + The conversation turns in order. Empty when *payload* is malformed, + which is treated as "no messages" rather than an error — a mapping + failure must never break the user's trace. + """ + items = _chat_ctx_items(payload) + + messages: List[Tuple[str, str]] = [] + for item in items: + if not isinstance(item, Mapping): + continue + if item.get(_CHAT_CTX_TYPE_KEY) != _CHAT_CTX_MESSAGE_TYPE: + continue + role = item.get(_CHAT_CTX_ROLE_KEY) + if not isinstance(role, str) or not role: + continue + text = _text_of_chat_content(item.get(_CHAT_CTX_CONTENT_KEY)) + if text is None: + continue + messages.append((role, text)) + + return messages + + +def _chat_ctx_items(payload: Any) -> List[Any]: + """Decode a ``ChatContext`` payload down to its item list. + + Args: + payload: A ``ChatContext`` JSON string or mapping. + + Returns: + The raw items, or an empty list for any payload that is not a decodable + ``ChatContext``. + """ + if isinstance(payload, str): + try: + payload = json.loads(payload) + except ValueError: + return [] + + if not isinstance(payload, Mapping): + return [] + + items = payload.get(_CHAT_CTX_ITEMS_KEY) + return items if isinstance(items, list) else [] + + +def _text_of_chat_content(content: Any) -> Optional[str]: + """Join the text parts of a ``ChatContext`` message's ``content``. + + Mirrors ``ChatMessage.raw_text_content`` in livekit-agents + (``llm/chat_context.py``): string parts joined by newline, non-text parts + (image/audio content objects) skipped. + + Args: + content: The item's ``content`` value. + + Returns: + The message text, or ``None`` when the item carries none. + """ + if isinstance(content, str): + return content or None + if not isinstance(content, list): + return None + + parts = [part for part in content if isinstance(part, str) and part] + if not parts: + return None + return "\n".join(parts) + + +# --------------------------------------------------------------------------- +# LiveKit TTS metrics +# --------------------------------------------------------------------------- + + +def tts_pricing_attributes_from(payload: Any) -> TtsPricingAttributes: + """Extract the priceable fields from a serialised LiveKit ``TTSMetrics``. + + Accepts either the JSON string LiveKit puts in ``lk.tts_metrics`` or the + equivalent dict, so the same rules apply whether the metrics are read off a + span or off a live ``TTSMetrics.model_dump()``. + + Args: + payload: A ``TTSMetrics`` JSON string or mapping. + + Returns: + The model and character count, each ``None`` when absent. Malformed input + yields both ``None`` rather than an error — a mapping failure must never + break the user's trace. + """ + if isinstance(payload, str): + try: + payload = json.loads(payload) + except ValueError: + return TtsPricingAttributes(None, None) + + if not isinstance(payload, Mapping): + return TtsPricingAttributes(None, None) + + metadata = payload.get(_TTS_METRICS_METADATA_KEY) + model = metadata.get(_TTS_METRICS_MODEL_KEY) if isinstance(metadata, Mapping) else None + + return TtsPricingAttributes( + model=model if isinstance(model, str) and model else None, + character_count=_positive_count(payload.get(_TTS_METRICS_CHARACTERS_KEY)), + ) + + +def _positive_count(value: Any) -> Optional[int]: + """Coerce a reported count to a positive int, or None if it is not one. + + A zero or negative count is treated as absent: it prices to nothing and would + only claim a measurement that says less than no attribute at all. + + Args: + value: The candidate count. + + Returns: + The count as an int, or ``None``. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + if value <= 0: + return None + return int(value) diff --git a/netra/instrumentation/livekit/version.py b/netra/instrumentation/livekit/version.py new file mode 100644 index 0000000..e4adfb8 --- /dev/null +++ b/netra/instrumentation/livekit/version.py @@ -0,0 +1 @@ +__version__ = "1.6.0" diff --git a/netra/instrumentation/livekit/wrappers.py b/netra/instrumentation/livekit/wrappers.py new file mode 100644 index 0000000..358710f --- /dev/null +++ b/netra/instrumentation/livekit/wrappers.py @@ -0,0 +1,318 @@ +"""wrapt wrappers for LiveKit's ``AgentSession`` lifecycle. + +Two things hang off the session's lifecycle, and this module is where both are +bolted on: + +* **the Netra session id** — the LiveKit room SID, falling back to the room name + — attached as OTel baggage *around* ``AgentSession.start`` so the + ``agent_session`` root span created inside it carries the id, then detached so + the caller's context is restored. See :func:`wrap_start` and + :func:`_resolve_session_id`; +* **call-audio capture** — started once ``start()`` has returned and torn down + before the session closes. The capture itself lives in ``audio_capture.py``; + this module only decides when it begins and ends. + +Nothing in here may change the behaviour of the user's application: every hook +runs the wrapped function whether or not our own logic succeeded, and exceptions +raised by the user's code propagate untouched. +""" + +from __future__ import annotations + +import logging +from contextlib import ExitStack +from typing import Any, Awaitable, Callable, Dict, Optional, Tuple + +from netra.config import get_active_config +from netra.instrumentation.livekit.audio_capture import start_audio_capture, stop_audio_capture +from netra.session_manager import SessionManager + +logger = logging.getLogger(__name__) + +# The wrapt quadruple is (wrapped, instance, args, kwargs). ``instance`` is a +# livekit AgentSession, which cannot be imported at module scope — the SDK must +# stay importable with livekit-agents absent. +WrappedAsync = Callable[..., Awaitable[Any]] + + +# --------------------------------------------------------------------------- +# Session-id resolution +# --------------------------------------------------------------------------- + + +def _resolve_session_id(kwargs: Dict[str, Any]) -> Optional[str]: + """Derive the Netra session id for an ``AgentSession.start`` call. + + Prefers the LiveKit room SID — the id LiveKit itself identifies the session by — + and falls back to the room name when there is no job context to read it from. + + Args: + kwargs: The keyword arguments ``start()`` was called with. + + Returns: + The session id, or ``None`` when neither source yields one — in which case + the session simply carries no Netra session id. + """ + return _room_sid_from_job_context() or _room_name(kwargs) + + +def _room_sid_from_job_context() -> Optional[str]: + """Read the room SID off the job assignment, or None if it is unavailable. + + Taken from ``JobContext.job.room.sid`` rather than ``rtc.Room.sid``: the latter + is an *async* property that only resolves once the room is connected, and in the + usual entrypoint ``session.start()`` runs before ``ctx.connect()`` — awaiting it + here would stall the user's agent, and in console mode (no real room) it would + never resolve. The job assignment carries the same server-issued SID + synchronously, before connect, which is what lets the ``agent_session`` root + span be stamped with it. + + Returns: + The room SID, or ``None`` outside a job (eval mode, direct library use) or + when livekit-agents does not expose one. + """ + try: + from livekit.agents import get_job_context + + job_context = get_job_context(required=False) + except Exception: + logger.debug("netra.livekit: could not read the job context", exc_info=True) + return None + + if job_context is None: + return None + + try: + sid = getattr(getattr(job_context.job, "room", None), "sid", None) + except Exception: + logger.debug("netra.livekit: could not read the room sid off the job", exc_info=True) + return None + + if isinstance(sid, str) and sid: + return sid + return None + + +def _room_name(kwargs: Dict[str, Any]) -> Optional[str]: + """Read the room name from ``AgentSession.start``'s ``room`` kwarg. + + ``room`` is keyword-only and defaults to ``NOT_GIVEN``, so it MUST NOT be read + positionally and MUST be checked with LiveKit's ``is_given`` before touching + ``room.name``. + + Args: + kwargs: The keyword arguments ``start()`` was called with. + + Returns: + The room name, or ``None`` when it cannot be determined. + """ + room = kwargs.get("room") + if room is None: + return None + + try: + from livekit.agents.utils import is_given + + if not is_given(room): + return None + except Exception: + logger.debug("netra.livekit: could not check room kwarg with is_given", exc_info=True) + return None + + name = getattr(room, "name", None) + if isinstance(name, str) and name: + return name + return None + + +# --------------------------------------------------------------------------- +# Session-span helpers +# --------------------------------------------------------------------------- + + +def _session_span(instance: Any) -> Optional[Any]: + """Return the live ``agent_session`` span, or ``None`` once it is gone. + + Args: + instance: The ``AgentSession``. + + Returns: + LiveKit's own session span while the session is open. + """ + return getattr(instance, "_session_span", None) + + +def _trace_id_of(session_span: Optional[Any]) -> Optional[int]: + """Read the trace id off the ``agent_session`` span. + + Args: + session_span: The session span, or ``None``. + + Returns: + The trace id, or ``None`` when there is no usable span context. This is + the key every per-session resource is filed under, so ``None`` means the + session gets no session-scoped wiring at all. + """ + if session_span is None: + return None + try: + span_context = session_span.get_span_context() + except Exception: + logger.debug("netra.livekit: could not read the session span context", exc_info=True) + return None + if span_context is None or not span_context.trace_id: + return None + return int(span_context.trace_id) + + +# --------------------------------------------------------------------------- +# Session lifecycle hooks +# --------------------------------------------------------------------------- + + +async def _after_start(instance: Any, session_id: Optional[str]) -> None: + """Run the per-session wiring, now that ``start()`` has returned. + + Args: + instance: The ``AgentSession`` that has started. + session_id: The Netra session id resolved for it, if any. + """ + trace_id = _trace_id_of(_session_span(instance)) + if trace_id is None: + logger.debug( + "netra.livekit: no agent_session span after start(); session-scoped wiring skipped " + "(session_id=%s). Spans still flow normally", + session_id, + ) + return + + logger.debug("netra.livekit: agent session started session_id=%s trace_id=%032x", session_id, trace_id) + + config = get_active_config() + if config is None or not config.audio_capture_enabled: + return + + await start_audio_capture(instance, config=config, session_id=session_id or "", trace_id=trace_id) + + +async def _before_close(instance: Any) -> None: + """Run the per-session teardown, *before* LiveKit closes the session. + + Ordering is load-bearing: ``_aclose_impl`` ends ``_session_span`` before it + emits ``close``, after which the span is gone and its trace id — the key + every per-session resource is filed under — is unreachable. + + Idempotent, in two layers: a second call finds no ``_session_span``, and the + coordinator registry only hands out a coordinator once. + + Args: + instance: The ``AgentSession`` that is closing. + """ + session_span = _session_span(instance) + trace_id = _trace_id_of(session_span) + if trace_id is None: + logger.debug("netra.livekit: session close with no live agent_session span; nothing to tear down") + return + + logger.debug("netra.livekit: agent session closing trace_id=%032x", trace_id) + await stop_audio_capture(trace_id, session_span=session_span) + + +# --------------------------------------------------------------------------- +# wrapt wrapper functions (public — referenced from __init__.py) +# --------------------------------------------------------------------------- + + +async def wrap_start( + wrapped: WrappedAsync, + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + """Attach the session id around ``AgentSession.start``. + + The attach happens **before** the await, because the ``agent_session`` span is + created inside ``start()`` and ``SessionSpanProcessor.on_start`` reads baggage + at that moment. Attaching afterwards would leave the trace's root span as the + one span missing ``netra.session_id``. + + The detach happens in ``finally``, in the same task, as OTel requires. Every + LiveKit task that produces spans for this session is created *during* + ``start()`` and snapshots the context at creation, so those tasks keep the + baggage for their whole lifetime while the caller's context is restored. + + Documented consequence: the session id is scoped to the LiveKit session's task + tree, not the whole job. Code running in the entrypoint task *after* + ``await session.start(...)`` carries no session id; a user who wants that + calls ``Netra.set_session_id()``, which is process-wide by design. + + Args: + wrapped: LiveKit's ``AgentSession.start``. + instance: The ``AgentSession``, needed by ``_after_start`` to reach the + session span and the session's audio I/O. + args: Positional arguments (``agent``). + kwargs: Keyword arguments, including the keyword-only ``room``. + + Returns: + Whatever ``start()`` returns, untouched. + """ + session_id = _resolve_session_id(kwargs) + + # ExitStack rather than a bare token so the detach is ordinary context-manager + # unwinding: it runs in this same coroutine, on both the success and error + # paths, and an attach failure degrades to "no session id" instead of + # propagating into the user's start() call. + scope = ExitStack() + if session_id is not None: + try: + scope.enter_context(SessionManager.session_scope(session_id=session_id)) + except Exception: + logger.warning("netra.livekit: could not attach session context", exc_info=True) + + try: + result = await wrapped(*args, **kwargs) + finally: + try: + scope.close() + except Exception: + logger.debug("netra.livekit: session context detach failed", exc_info=True) + + # Awaited after the detach so its own failures cannot leak session context, + # and isolated so they can never surface in the user's start() call. + try: + await _after_start(instance, session_id) + except Exception: + logger.warning("netra.livekit: post-start wiring failed", exc_info=True) + + return result + + +async def wrap_aclose( + wrapped: WrappedAsync, + instance: Any, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + """Run per-session teardown before LiveKit closes the session. + + Wraps ``_aclose_impl`` rather than ``aclose``: ``aclose()`` covers only the + ``USER_INITIATED`` close reason, while the other four — including + ``PARTICIPANT_DISCONNECTED``, i.e. the caller hanging up — reach + ``_aclose_impl`` directly. Wrapping ``aclose`` would mean the teardown never + runs on a normal phone call. + + Args: + wrapped: LiveKit's ``AgentSession._aclose_impl``. + instance: The ``AgentSession``. + args: Positional arguments. + kwargs: Keyword arguments, including the ``reason``. + + Returns: + Whatever ``_aclose_impl`` returns, untouched. + """ + try: + await _before_close(instance) + except Exception: + logger.warning("netra.livekit: pre-close teardown failed", exc_info=True) + + return await wrapped(*args, **kwargs) diff --git a/netra/meter.py b/netra/meter.py index 50bf1e5..538386a 100644 --- a/netra/meter.py +++ b/netra/meter.py @@ -112,6 +112,59 @@ def export( return MetricExportResult.FAILURE +class _NetraOwnedMeterProvider(MeterProvider): # type: ignore[misc] + """A MeterProvider only Netra may shut down. + + Third-party teardown paths call ``shutdown()`` on the global meter provider. + ``livekit-agents`` does it on every job cleanup + (``telemetry/traces.py:_shutdown_telemetry``, reached unconditionally from + ``ipc/job.py``'s ``_on_cleanup``), which would permanently stop Netra's + metrics pipeline for the rest of the process — including every later job in + a multi-job worker. + + This cannot be retrofitted by wrapping the provider after the fact: OTel's + ``set_meter_provider`` is set-once and merely logs *"Overriding of current + MeterProvider is not allowed"* on a second call. The guard therefore has to + live where the provider is constructed. + + ``Netra.shutdown()`` calls :meth:`shutdown_as_owner`, which is the only way + to actually tear this provider down. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + """Build the provider with shutdown locked to Netra's own teardown. + + Args: + *args: Positional arguments forwarded to ``MeterProvider``. + **kwargs: Keyword arguments forwarded to ``MeterProvider``. + """ + super().__init__(*args, **kwargs) + self._netra_shutdown_allowed = False + + def shutdown(self, timeout_millis: float = 30_000, **kwargs: Any) -> None: + """Ignore shutdown requests that do not come from Netra's own teardown. + + Args: + timeout_millis: Maximum time to wait for the shutdown, forwarded to + ``MeterProvider.shutdown`` only when the caller is Netra itself. + **kwargs: Additional arguments forwarded to ``MeterProvider.shutdown``. + """ + if not self._netra_shutdown_allowed: + logger.debug("Ignoring third-party MeterProvider shutdown; Netra owns this provider's lifecycle") + return + super().shutdown(timeout_millis=timeout_millis, **kwargs) + + def shutdown_as_owner(self, timeout_millis: float = 30_000) -> None: + """Shut the provider down for real. Called only by ``Netra.shutdown()``. + + Args: + timeout_millis: Maximum time to wait for the metric readers to flush + and shut down. + """ + self._netra_shutdown_allowed = True + super().shutdown(timeout_millis=timeout_millis) + + class MetricsSetup: """ Configures Netra's OpenTelemetry metrics pipeline. @@ -194,7 +247,7 @@ def _setup_meter(self) -> None: views = self._build_views() - provider = MeterProvider( + provider = _NetraOwnedMeterProvider( resource=resource, metric_readers=[reader], views=views, diff --git a/netra/processors/instrumentation_span_processor.py b/netra/processors/instrumentation_span_processor.py index ff400d5..c917ac2 100644 --- a/netra/processors/instrumentation_span_processor.py +++ b/netra/processors/instrumentation_span_processor.py @@ -8,7 +8,7 @@ from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor from netra.config import Config, get_attribute_max_len -from netra.instrumentation.instruments import InstrumentSet +from netra.instrumentation.instruments import THIRD_PARTY_INSTRUMENTATION_SCOPES, InstrumentSet logger = logging.getLogger(__name__) @@ -241,8 +241,12 @@ def _extract_instrumentation_name(span: Span) -> Optional[str]: """Extracts the instrumentation name from the span's scope. For scopes with known prefixes (opentelemetry.instrumentation.* or - netra.instrumentation.*), returns just the final component. - Otherwise, returns the full scope name. + netra.instrumentation.*), returns just the final component. A + third-party scope registered in ``THIRD_PARTY_INSTRUMENTATION_SCOPES`` + resolves to its ``InstrumentSet`` value — ``livekit-agents`` to + ``livekit`` — so those spans get stamped with the same instrumentation + name as every other instrumentation instead of being skipped for + carrying a non-conforming scope. Otherwise, returns the full scope name. Args: span: The span to extract the instrumentation name from. @@ -258,6 +262,10 @@ def _extract_instrumentation_name(span: Span) -> Optional[str]: if not isinstance(name, str) or not name: return None + alias = THIRD_PARTY_INSTRUMENTATION_SCOPES.get(name) + if alias is not None: + return alias + if name.startswith(_OTEL_INSTRUMENTATION_PREFIX) or name.startswith(_NETRA_INSTRUMENTATION_PREFIX): base_name = name.rsplit(".", 1)[-1].strip() return base_name if base_name else name diff --git a/netra/processors/root_instrument_filter_processor.py b/netra/processors/root_instrument_filter_processor.py index ee1ba94..6cc63ae 100644 --- a/netra/processors/root_instrument_filter_processor.py +++ b/netra/processors/root_instrument_filter_processor.py @@ -8,6 +8,8 @@ from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor from opentelemetry.trace import INVALID_SPAN_ID, SpanContext +from netra.instrumentation.instruments import THIRD_PARTY_INSTRUMENTATION_SCOPES + logger = logging.getLogger(__name__) _INSTRUMENTATION_PREFIXES = ("opentelemetry.instrumentation.", "netra.instrumentation.") @@ -100,8 +102,11 @@ class RootInstrumentFilterProcessor(SpanProcessor): # type: ignore[misc] Spans created directly through netra decorators or ``Netra.start_span`` are never candidates — only spans from recognised auto-instrumentation - libraries (scope prefix ``opentelemetry.instrumentation.*`` or - ``netra.instrumentation.*``) are subject to the allow-list. + libraries are subject to the allow-list: those whose scope carries the + ``opentelemetry.instrumentation.*`` / ``netra.instrumentation.*`` prefix, plus + the third-party scopes named in ``THIRD_PARTY_INSTRUMENTATION_SCOPES`` (e.g. + ``livekit-agents``), which Netra enables but does not author and which + therefore do not follow that naming convention. Args: allowed_root_instrument_names: Instrumentation-name strings @@ -185,10 +190,7 @@ def _process_span_start(self, span: Span) -> None: Args: span: The span that is being started. """ - if not self._is_from_instrumentation_library(span): - return - - instr_name = self._extract_instrumentation_name(span) + instr_name = self._resolve_instrument_name(span) if instr_name is None or instr_name in self._allowed: return @@ -292,40 +294,28 @@ def _get_parent_span_context(span: Span) -> Optional[SpanContext]: return cast(Optional[SpanContext], parent) @staticmethod - def _is_from_instrumentation_library(span: Span) -> bool: - """Return ``True`` when *span* originates from a known - auto-instrumentation library. - - Spans created by netra decorators or ``Netra.start_span`` use - arbitrary tracer names that do not match the instrumentation - naming convention and will return ``False``. - - Args: - span: The span to check. - - Returns: - Whether the span's scope starts with a recognised prefix. - """ - scope = getattr(span, "instrumentation_scope", None) - if scope is None: - return False - name = getattr(scope, "name", None) - if not isinstance(name, str) or not name: - return False - return name.startswith(_INSTRUMENTATION_PREFIXES) + def _resolve_instrument_name(span: Span) -> Optional[str]: + """Return the instrument name *span* is subject to, or ``None``. - @staticmethod - def _extract_instrumentation_name(span: Span) -> Optional[str]: - """Extract the short instrumentation name from *span*'s scope. + Answers both "did an auto-instrumentation library produce this span?" and + "which instrument is it?" from the single scope string, deliberately in + one function. As two separate predicates they could disagree about a + scope, and a scope that the first accepted but the second could not name + would slip past the allow-list unchecked. - For a scope named ``netra.instrumentation.fastapi`` this returns - ``"fastapi"``. + A scope named ``netra.instrumentation.fastapi`` resolves to ``fastapi``. + A third-party scope registered in ``THIRD_PARTY_INSTRUMENTATION_SCOPES`` + resolves to its ``InstrumentSet`` value — ``livekit-agents`` to + ``livekit`` — which is what brings those spans under ``root_instruments`` + control despite the non-conforming scope name. Args: span: The span to inspect. Returns: - The short name, or ``None`` if extraction fails. + The short instrumentation name, or ``None`` when the scope belongs to + no recognised instrumentation — a netra decorator, ``Netra.start_span`` + or any user tracer — in which case the span is never a candidate. """ scope = getattr(span, "instrumentation_scope", None) if scope is None: @@ -333,11 +323,18 @@ def _extract_instrumentation_name(span: Span) -> Optional[str]: name = getattr(scope, "name", None) if not isinstance(name, str) or not name: return None - for prefix in _INSTRUMENTATION_PREFIXES: - if name.startswith(prefix): - base = name.rsplit(".", 1)[-1].strip() - return base if base else name - return name + + # Exact-match aliases first: a third-party scope carries no prefix, so + # the two branches cannot both claim the same name. + alias = THIRD_PARTY_INSTRUMENTATION_SCOPES.get(name) + if alias is not None: + return alias + + if name.startswith(_INSTRUMENTATION_PREFIXES): + base = name.rsplit(".", 1)[-1].strip() + return base if base else name + + return None def _evict_stale_candidates(self) -> None: """Evict entries whose span ended more than ``_ROOT_CANDIDATE_TTL_SECONDS`` ago. diff --git a/netra/session_manager.py b/netra/session_manager.py index 51c32f6..95cb7c9 100644 --- a/netra/session_manager.py +++ b/netra/session_manager.py @@ -1,8 +1,9 @@ import contextvars import logging +from contextlib import contextmanager from datetime import datetime from enum import Enum -from typing import Any, Dict, List, Optional, Tuple, Union, cast +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast from opentelemetry import baggage from opentelemetry import context as otel_context @@ -94,6 +95,42 @@ def _current_entity_name(entity_type: str) -> Optional[str]: return frames[-1][1] if frames else None +# The baggage keys that carry session identity. ``SessionSpanProcessor.on_start`` +# reads exactly these names off the ambient context, so every writer must go +# through ``_build_session_context`` rather than calling ``set_baggage`` inline — +# otherwise the global setter and the scoped attach can drift apart. +_SESSION_BAGGAGE_KEYS: Tuple[str, ...] = ("session_id", "user_id", "tenant_id") + + +def _build_session_context( + ctx: otel_context.Context, + *, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + tenant_id: Optional[str] = None, +) -> Optional[otel_context.Context]: + """Return *ctx* with the supplied session fields set as baggage. + + Args: + ctx: The context to derive from. Not mutated. + session_id: Session identifier, or ``None`` to leave unset. + user_id: User identifier, or ``None`` to leave unset. + tenant_id: Tenant identifier, or ``None`` to leave unset. + + Returns: + A new ``Context`` carrying the supplied fields, or ``None`` when every + field was ``None`` or empty — signalling that there is nothing to attach. + """ + values = {"session_id": session_id, "user_id": user_id, "tenant_id": tenant_id} + changed = False + for key in _SESSION_BAGGAGE_KEYS: + value = values[key] + if isinstance(value, str) and value: + ctx = baggage.set_baggage(key, value, ctx) + changed = True + return ctx if changed else None + + class ConversationType(str, Enum): INPUT = "input" OUTPUT = "output" @@ -343,23 +380,98 @@ def set_session_context( """ Set session context attributes in OpenTelemetry baggage. + The attach is deliberately never detached: ``Netra.set_session_id()`` and + friends are documented as process-sticky, and existing users rely on the + session id outliving the call that set it. For a scoped session id that + is restored on exit — what instrumentation wants — use + :meth:`attach_session_context` or :meth:`session_scope` instead. + Args: session_key: Key to set in baggage (session_id, user_id, tenant_id, or custom_attributes) value: Value to set for the key """ try: - ctx = otel_context.get_current() - if isinstance(value, str) and value: - if session_key == "session_id": - ctx = baggage.set_baggage("session_id", value, ctx) - elif session_key == "user_id": - ctx = baggage.set_baggage("user_id", value, ctx) - elif session_key == "tenant_id": - ctx = baggage.set_baggage("tenant_id", value, ctx) - otel_context.attach(ctx) + if isinstance(value, str) and value and session_key in _SESSION_BAGGAGE_KEYS: + ctx = _build_session_context(otel_context.get_current(), **{session_key: value}) + if ctx is not None: + otel_context.attach(ctx) except Exception as e: logger.exception(f"Failed to set session context for key={session_key}: {e}") + @staticmethod + def attach_session_context( + *, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + tenant_id: Optional[str] = None, + ) -> Optional[object]: + """Attach session baggage to the current OTel context and return its token. + + Unlike :meth:`set_session_context`, the caller owns the returned token and + MUST detach it — in the same context it was attached in, as OTel requires. + Prefer :meth:`session_scope` where the scope is lexical. + + Args: + session_id: Session identifier to put in baggage, if any. + user_id: User identifier to put in baggage, if any. + tenant_id: Tenant identifier to put in baggage, if any. + + Returns: + The token to pass to ``opentelemetry.context.detach``, or ``None`` + when no field was supplied — nothing was attached, so there is + nothing to detach and callers need no emptiness branch. + """ + ctx = _build_session_context( + otel_context.get_current(), + session_id=session_id, + user_id=user_id, + tenant_id=tenant_id, + ) + if ctx is None: + return None + # Declared as ``object`` so the public signature does not leak OTel's + # Token generic; callers only ever hand it back to ``otel_context.detach``. + token: object = otel_context.attach(ctx) + return token + + @staticmethod + @contextmanager + def session_scope( + *, + session_id: Optional[str] = None, + user_id: Optional[str] = None, + tenant_id: Optional[str] = None, + ) -> Iterator[None]: + """Scoped form of :meth:`attach_session_context`. + + Detaches on exit, including when the body raises. + + Args: + session_id: Session identifier to put in baggage, if any. + user_id: User identifier to put in baggage, if any. + tenant_id: Tenant identifier to put in baggage, if any. + + Yields: + None. The session baggage is active for the duration of the block. + """ + # Attaches directly rather than via attach_session_context() so the token + # keeps its concrete OTel type here; both paths share _build_session_context, + # which is what keeps the baggage keys from drifting. + ctx = _build_session_context( + otel_context.get_current(), + session_id=session_id, + user_id=user_id, + tenant_id=tenant_id, + ) + if ctx is None: + yield + return + token = otel_context.attach(ctx) + try: + yield + finally: + otel_context.detach(token) + @staticmethod def set_custom_event(name: str, attributes: Dict[str, Any]) -> None: """ diff --git a/poetry.lock b/poetry.lock index 27d09b9..55e76b2 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.4.1 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -209,7 +209,7 @@ description = "Timeout context manager for asyncio programs" optional = false python-versions = ">=3.8" groups = ["main"] -markers = "python_version < \"3.11\"" +markers = "python_version == \"3.10\"" files = [ {file = "async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c"}, {file = "async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3"}, @@ -487,7 +487,7 @@ decli = ">=0.6.0,<1.0" importlib-metadata = {version = ">=8.0.0,<9.0.0", markers = "python_version != \"3.9\""} jinja2 = ">=2.10.3" packaging = ">=19" -pyyaml = ">=3.08" +pyyaml = ">=3.8" questionary = ">=2.0,<3.0" termcolor = ">=1.1.0,<4.0.0" tomlkit = ">=0.5.3,<1.0.0" @@ -553,7 +553,7 @@ description = "Backport of PEP 654 (exception groups)" optional = false python-versions = ">=3.7" groups = ["main", "dev"] -markers = "python_version < \"3.11\"" +markers = "python_version == \"3.10\"" files = [ {file = "exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10"}, {file = "exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88"}, @@ -3668,7 +3668,7 @@ description = "A lil' TOML parser" optional = false python-versions = ">=3.8" groups = ["dev"] -markers = "python_version < \"3.11\"" +markers = "python_version == \"3.10\"" files = [ {file = "tomli-2.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:678e4fa69e4575eb77d103de3df8a895e1591b48e740211bd1067378c69e8249"}, {file = "tomli-2.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:023aa114dd824ade0100497eb2318602af309e5a55595f76b626d6d9f3b7b0a6"}, @@ -4091,4 +4091,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.14" -content-hash = "ec711d22faad500af30d50cfdf754ffdf038509c9bd94a404bd4135b0197b046" +content-hash = "193f2cbec44769c5e7e9c434e3780b2e8387ccf5b35ca6de1f88f88dd98ebd2a" diff --git a/pyproject.toml b/pyproject.toml index 16b728a..359582d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ dependencies = [ "opentelemetry-instrumentation-tortoiseorm>=0.55b1,<=0.62b1", "opentelemetry-instrumentation-urllib>=0.55b1,<=0.62b1", "opentelemetry-instrumentation-urllib3>=0.55b1,<=0.62b1", - "json-repair==0.44.1", + "json-repair>=0.44.1,<1.0.0", "httpx>=0.27.0,<1.0.0", ] diff --git a/tests/test_audio_integration.py b/tests/test_audio_integration.py new file mode 100644 index 0000000..84d3220 --- /dev/null +++ b/tests/test_audio_integration.py @@ -0,0 +1,967 @@ +"""Tests for the LiveKit call-audio capture pipeline. + +The sender is exercised end to end against a recording HTTP server defined in +this module, so the ``x-audio-*`` wire contract the Netra backend depends on is +asserted on real requests rather than on a mock's call args. The coordinator and +the span processor are tested directly, with a stub sender. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +import time +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Awaitable, Callable, Dict, List, Optional +from unittest.mock import MagicMock + +import pytest + +from netra.config import Config +from netra.instrumentation.livekit.audio_capture import ( + AudioCoordinatorRegistry, + SessionAudioCoordinator, + audio_coordinators, + build_audio_sender, + start_audio_capture, + stop_audio_capture, +) +from netra.instrumentation.livekit.audio_processor import AudioSpanProcessor +from netra.instrumentation.livekit.audio_sender import AudioChunkSender +from netra.instrumentation.livekit.audio_types import ( + HEADER_HEARD_MS, + HEADER_LAST_CHUNK, + HEADER_ROLE, + HEADER_SEQUENCE, + HEADER_SESSION_ID, + HEADER_SESSION_LAST, + HEADER_SPAN_ID, + HEADER_TRACE_ID, + NETRA_AUDIO_DROPPED_FRAMES, + NETRA_AUDIO_SENT_BYTES, + NETRA_AUDIO_SENT_CHUNKS, + SpeakerRole, + pcm_byte_offset_at, +) + +# 24kHz mono 16-bit — what livekit-agents delivers by default. +SAMPLE_RATE_HZ = 24000 +BYTES_PER_MS = SAMPLE_RATE_HZ * 2 // 1000 +SAMPLES_PER_FRAME = 480 +FRAME_BYTES = SAMPLES_PER_FRAME * 2 +FRAME_MS = FRAME_BYTES // BYTES_PER_MS + +USER_SPAN_ID = "aaaabbbbccccdddd" +AGENT_SPAN_ID = "1111222233334444" +TRACE_ID = "0123456789abcdef0123456789abcdef" + +# Large enough that no test hits a batch boundary it did not ask for. +UNBOUNDED_BYTES = 10_000_000 +UNBOUNDED_FRAMES = 10_000 +LONG_INTERVAL_SECONDS = 30.0 + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +@dataclass +class FakeAudioFrame: + """Minimal stand-in for ``livekit.rtc.AudioFrame``.""" + + pcm: bytes + sample_rate: int = SAMPLE_RATE_HZ + num_channels: int = 1 + + @property + def data(self) -> memoryview: + return memoryview(self.pcm) + + +def make_frame(sample_count: int = SAMPLES_PER_FRAME, value: int = 1000) -> FakeAudioFrame: + """Build one frame of constant-amplitude PCM.""" + return FakeAudioFrame(pcm=value.to_bytes(2, "little", signed=True) * sample_count) + + +async def _async_noop(*args: Any, **kwargs: Any) -> None: + """Stand in for an awaitable the test does not care about.""" + + +@dataclass +class RecordedRequest: + """One request the ingest server received.""" + + headers: Dict[str, str] + body: bytes + + @property + def span_id(self) -> Optional[str]: + return self.headers.get(HEADER_SPAN_ID) + + @property + def is_last(self) -> bool: + return self.headers.get(HEADER_LAST_CHUNK) == "true" + + +@dataclass +class IngestRecorder: + """Thread-safe record of what the ingest server received.""" + + status_code: int = 200 + # Held before replying, so a test can stall the send loop mid-POST the way a + # degraded backend would. Only the teardown-budget tests set it. + delay_seconds: float = 0.0 + requests: List[RecordedRequest] = field(default_factory=list) + _lock: threading.Lock = field(default_factory=threading.Lock) + + def record(self, request: RecordedRequest) -> None: + with self._lock: + self.requests.append(request) + + def snapshot(self) -> List[RecordedRequest]: + with self._lock: + return list(self.requests) + + def chunks_for(self, span_id: str) -> List[RecordedRequest]: + return [request for request in self.snapshot() if request.span_id == span_id] + + def bytes_for(self, span_id: str) -> int: + return sum(len(request.body) for request in self.chunks_for(span_id)) + + +class _IngestHandler(BaseHTTPRequestHandler): + """Records every POST into the server's recorder and replies with its status.""" + + recorder: IngestRecorder + + def do_POST(self) -> None: # noqa: N802 - name fixed by BaseHTTPRequestHandler + length = int(self.headers.get("Content-Length") or 0) + body = self.rfile.read(length) if length else b"" + self.recorder.record( + RecordedRequest( + headers={name.lower(): value for name, value in self.headers.items()}, + body=body, + ) + ) + if self.recorder.delay_seconds: + time.sleep(self.recorder.delay_seconds) + self.send_response(self.recorder.status_code) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, format: str, *args: Any) -> None: + """Silence the default stderr access log.""" + + +@pytest.fixture() +def ingest_server(): + """Serve the audio-ingest endpoint on a random port; yield (url, recorder).""" + recorder = IngestRecorder() + handler = type("_BoundIngestHandler", (_IngestHandler,), {"recorder": recorder}) + + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/telemetry/v1/audio/chunk", recorder + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def build_sender(url: str, **overrides: Any) -> AudioChunkSender: + """Build a sender that only flushes when a test tells it to.""" + settings: Dict[str, Any] = { + "url": url, + "session_id": "session-under-test", + "api_key": "test-key", + "batch_interval_seconds": LONG_INTERVAL_SECONDS, + "max_batch_frames": UNBOUNDED_FRAMES, + "flush_at_bytes": UNBOUNDED_BYTES, + "max_request_bytes": UNBOUNDED_BYTES, + } + settings.update(overrides) + return AudioChunkSender(**settings) + + +def enqueue_frames(sender: AudioChunkSender, count: int, *, role: SpeakerRole, span_id: str) -> None: + """Enqueue *count* identical frames for one span.""" + for _ in range(count): + sender.enqueue(make_frame(), role=role, trace_id=TRACE_ID, span_id=span_id) + + +def run_call(url: str, scenario: Callable[[AudioChunkSender], Awaitable[None]], **overrides: Any) -> AudioChunkSender: + """Drive one whole call against the ingest server and return its sender. + + The repo has no pytest-asyncio, so each async scenario gets its own loop. + + Args: + url: The ingest URL to send to. + scenario: What the call does between start and close. + **overrides: Sender settings to override for this call. + + Returns: + The closed sender, for its statistics. + """ + + async def drive() -> AudioChunkSender: + sender = build_sender(url, **overrides) + await sender.start() + await scenario(sender) + await sender.end_session() + return sender + + return asyncio.run(drive()) + + +# --------------------------------------------------------------------------- +# audio_types +# --------------------------------------------------------------------------- + + +class TestAudioTypes: + @pytest.mark.parametrize( + "playback_ms,expected_bytes", + [ + (0, 0), + (-100, 0), + (1, 48), + (400, 19200), + (1000, 48000), + ], + ) + def test_pcm_byte_offset_converts_playback_time_to_bytes(self, playback_ms: int, expected_bytes: int) -> None: + offset = pcm_byte_offset_at(playback_ms=playback_ms, sample_rate_hz=SAMPLE_RATE_HZ, channel_count=1) + assert offset == expected_bytes + + def test_pcm_byte_offset_rounds_down_to_a_whole_sample_frame(self) -> None: + # 11025Hz stereo: 44.1 bytes/ms, so 7ms is 308.7 bytes — not a frame boundary. + offset = pcm_byte_offset_at(playback_ms=7, sample_rate_hz=11025, channel_count=2) + + frame_size = 2 * 2 + assert offset % frame_size == 0 + assert offset == 308 + + @pytest.mark.parametrize("sample_rate_hz,channel_count", [(0, 1), (24000, 0), (-1, 1)]) + def test_pcm_byte_offset_rejects_an_unplayable_format(self, sample_rate_hz: int, channel_count: int) -> None: + with pytest.raises(ValueError, match="unplayable PCM format"): + pcm_byte_offset_at(playback_ms=100, sample_rate_hz=sample_rate_hz, channel_count=channel_count) + + +# --------------------------------------------------------------------------- +# Sender: wire contract +# --------------------------------------------------------------------------- + + +class TestAudioChunkSenderWireContract: + def test_frames_reach_the_endpoint_and_the_session_is_closed(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 10, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + sender = run_call(url, scenario) + + requests = recorder.snapshot() + assert recorder.bytes_for(USER_SPAN_ID) == 10 * FRAME_BYTES + assert sender.stats.frames_sent == 10 + assert sender.stats.errors == 0 + + session_end = [r for r in requests if r.headers.get(HEADER_SESSION_LAST) == "true"] + assert len(session_end) == 1 + assert session_end[0].body == b"" + assert session_end[0].headers[HEADER_SESSION_ID] == "session-under-test" + + def test_every_request_carries_the_session_and_credential(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 3, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + run_call(url, scenario, auth_headers={"Authorization": "Bearer token"}) + + assert recorder.snapshot() + for request in recorder.snapshot(): + assert request.headers[HEADER_SESSION_ID] == "session-under-test" + assert request.headers["x-api-key"] == "test-key" + assert request.headers["authorization"] == "Bearer token" + + def test_sequence_is_zero_based_and_monotonic_per_span(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 6, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + enqueue_frames(sender, 4, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + + # Two frames per request, so each span spans several sequence numbers. + run_call(url, scenario, max_batch_frames=2) + + for span_id, expected_chunks in ((USER_SPAN_ID, 3), (AGENT_SPAN_ID, 2)): + sequences = [int(r.headers[HEADER_SEQUENCE]) for r in recorder.chunks_for(span_id)] + # Each span sends its data chunks plus one terminator. + assert sequences == list(range(expected_chunks + 1)), span_id + + def test_a_span_is_terminated_exactly_once(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 3, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + sender.mark_audio_end(role=SpeakerRole.USER, span_id=USER_SPAN_ID) + # A duplicate end marker must not produce a second terminator. + sender.mark_audio_end(role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + run_call(url, scenario) + + terminators = [r for r in recorder.chunks_for(USER_SPAN_ID) if r.is_last] + assert len(terminators) == 1 + assert terminators[0].body == b"" + assert terminators[0].headers[HEADER_ROLE] == "user" + + def test_a_span_left_open_is_terminated_at_session_end(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + # No mark_audio_end: the session closes with the span still recording. + enqueue_frames(sender, 2, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + + run_call(url, scenario) + + terminators = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if r.is_last] + assert len(terminators) == 1 + + def test_audio_outside_a_span_carries_no_span_headers(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 4, role=SpeakerRole.USER, span_id="") + + run_call(url, scenario) + + chunks = [r for r in recorder.snapshot() if r.body] + assert chunks + for chunk in chunks: + assert HEADER_SPAN_ID not in chunk.headers + assert HEADER_SEQUENCE not in chunk.headers + assert HEADER_LAST_CHUNK not in chunk.headers + assert chunk.headers[HEADER_TRACE_ID] == TRACE_ID + + def test_a_request_body_never_exceeds_the_configured_ceiling(self, ingest_server) -> None: + url, recorder = ingest_server + ceiling = FRAME_BYTES * 3 + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 10, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + run_call(url, scenario, flush_at_bytes=ceiling, max_request_bytes=ceiling) + + bodies = [len(r.body) for r in recorder.chunks_for(USER_SPAN_ID)] + assert bodies, "expected at least one chunk" + assert max(bodies) <= ceiling + assert sum(bodies) == 10 * FRAME_BYTES + + +# --------------------------------------------------------------------------- +# Sender: interrupts +# --------------------------------------------------------------------------- + + +class TestAudioChunkSenderInterrupts: + def test_interrupt_trims_pending_audio_to_what_was_heard(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + # 50 frames = 1000ms of agent speech, of which 400ms was heard. + enqueue_frames(sender, 50, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=400) + + run_call(url, scenario) + + heard_bytes = 400 * BYTES_PER_MS + assert recorder.bytes_for(AGENT_SPAN_ID) == heard_bytes + + terminators = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if r.is_last] + assert len(terminators) == 1 + assert terminators[0].headers[HEADER_HEARD_MS] == "400" + + def test_frames_queued_after_an_interrupt_are_discarded(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 5, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=50) + # Frames already in flight when the caller cut in; never played out. + enqueue_frames(sender, 10, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + + run_call(url, scenario) + + assert recorder.bytes_for(AGENT_SPAN_ID) == 50 * BYTES_PER_MS + + def test_interrupt_after_the_span_closed_sends_a_correction(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 5, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + sender.mark_audio_end(role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + # LiveKit routinely reports the interrupt only after the span ended. + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=40) + + run_call(url, scenario) + + corrections = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if HEADER_HEARD_MS in r.headers] + assert len(corrections) == 1 + assert corrections[0].headers[HEADER_HEARD_MS] == "40" + assert corrections[0].headers[HEADER_LAST_CHUNK] == "true" + + def test_interrupt_after_the_heard_audio_was_already_sent_only_marks_the_cut(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 5, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + await asyncio.sleep(0.1) + # Only the first frame's worth was heard, but 5 already went out. + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=FRAME_MS) + + # Flush every frame immediately, so all the audio is delivered up front. + run_call(url, scenario, max_batch_frames=1) + + terminators = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if r.is_last] + assert len(terminators) == 1 + assert terminators[0].body == b"", "nothing more to send; the endpoint trims" + assert terminators[0].headers[HEADER_HEARD_MS] == str(FRAME_MS) + + def test_interrupt_marks_the_cut_when_a_pending_batch_is_past_the_heard_point(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + # Two frames flush on the batch boundary, so they are already + # delivered; a third stays pending. + enqueue_frames(sender, 3, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + await asyncio.sleep(0.1) + # Only the first frame was heard, which is behind what already went + # out — so the pending batch trims to nothing but the endpoint still + # has to be told where to cut. + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=FRAME_MS) + + run_call(url, scenario, max_batch_frames=2) + + terminators = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if r.is_last] + assert len(terminators) == 1, "the span must be terminated, not left open until session end" + assert terminators[0].headers[HEADER_HEARD_MS] == str(FRAME_MS) + assert terminators[0].headers[HEADER_SPAN_ID] == AGENT_SPAN_ID + + def test_a_trimmed_chunk_counts_only_the_frames_it_actually_sent(self, ingest_server) -> None: + url, _ = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + # 5 frames pending, of which 2 frames' worth was heard. + enqueue_frames(sender, 5, role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + sender.interrupt_agent_span(span_id=AGENT_SPAN_ID, playback_ms=FRAME_MS * 2) + + sender = run_call(url, scenario) + + assert sender.stats.bytes_sent == FRAME_BYTES * 2 + assert sender.stats.frames_sent == 2, "the untrimmed frame count would report 5" + + +# --------------------------------------------------------------------------- +# Sender: failure handling +# --------------------------------------------------------------------------- + + +class TestAudioChunkSenderFailures: + def test_a_full_queue_drops_frames_instead_of_blocking(self, ingest_server) -> None: + url, _ = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + # Enqueued without awaiting, so the loop cannot drain any of them. + enqueue_frames(sender, 50, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + assert sender.stats.frames_dropped == 45 + + run_call(url, scenario, max_queue_frames=5) + + def test_a_rejected_credential_stops_the_call_from_sending_more(self, ingest_server) -> None: + url, recorder = ingest_server + recorder.status_code = 401 + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 10, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + sender = run_call(url, scenario, max_batch_frames=1) + + assert sender.stats.circuit_tripped is True + assert sender.stats.chunks_sent == 0 + # One rejected attempt, then nothing further — not one per frame, and no + # retry of a credential that cannot become valid mid-call. + assert len(recorder.snapshot()) == 1 + + def test_a_server_error_is_retried_then_given_up_on(self, ingest_server) -> None: + url, recorder = ingest_server + recorder.status_code = 500 + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 1, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + sender = run_call(url, scenario) + + assert sender.stats.chunks_sent == 0 + assert sender.stats.errors > 1, "a 5xx is worth retrying" + assert sender.stats.bytes_sent == 0, "nothing was accepted" + + def test_end_session_is_idempotent(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + enqueue_frames(sender, 2, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + # run_call closes it again once this returns. + await sender.end_session() + + run_call(url, scenario) + + session_ends = [r for r in recorder.snapshot() if r.headers.get(HEADER_SESSION_LAST) == "true"] + assert len(session_ends) == 1 + + def test_a_given_up_chunk_does_not_leave_its_sequence_number_to_the_next_chunk(self, ingest_server) -> None: + url, recorder = ingest_server + + async def scenario(sender: AudioChunkSender) -> None: + recorder.status_code = 500 + sender.enqueue(make_frame(value=1111), role=SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + await asyncio.sleep(0.3) # both attempts fail; the chunk is given up on + recorder.status_code = 200 + sender.enqueue(make_frame(value=2222), role=SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + await asyncio.sleep(0.3) + + run_call(url, scenario, max_batch_frames=1) + + # A number may repeat only across retries of the *same* bytes: that is what + # makes it an idempotency key. Reusing it for different audio would have the + # endpoint either drop the new chunk as a duplicate or overwrite the old. + audio_by_sequence: Dict[str, set] = {} + for request in recorder.chunks_for(AGENT_SPAN_ID): + if request.body: + audio_by_sequence.setdefault(request.headers[HEADER_SEQUENCE], set()).add(request.body) + reused = {seq: len(bodies) for seq, bodies in audio_by_sequence.items() if len(bodies) > 1} + assert not reused, f"sequence reused across distinct audio: {reused}" + # The lost chunk still consumed its slot, so the gap is visible. + assert sorted(audio_by_sequence) == ["0", "1"] + + def test_end_session_spends_one_total_budget_not_one_per_wait(self, ingest_server) -> None: + url, recorder = ingest_server + recorder.delay_seconds = 2.0 + + async def drive() -> float: + sender = build_sender(url, max_batch_frames=1) + await sender.start() + enqueue_frames(sender, 4, role=SpeakerRole.USER, span_id=USER_SPAN_ID) + started_at = time.monotonic() + await sender.end_session(drain_timeout_seconds=0.5) + return time.monotonic() - started_at + + elapsed = asyncio.run(drive()) + + # The two internal waits share the 0.5s deadline. Taking it each would put + # this at 1s+, and the pre-fix 30s-per-wait default at a minute. + assert elapsed < 1.5, f"teardown took {elapsed:.2f}s for a 0.5s budget" + + def test_a_tripped_circuit_does_not_warn_once_per_open_span(self, ingest_server, caplog) -> None: + url, recorder = ingest_server + recorder.status_code = 500 + + async def scenario(sender: AudioChunkSender) -> None: + for index in range(8): + span_id = f"{index:016x}" + enqueue_frames(sender, 1, role=SpeakerRole.USER, span_id=span_id) + await asyncio.sleep(0.05) + + with caplog.at_level("WARNING"): + sender = run_call(url, scenario, max_batch_frames=1) + + assert sender.stats.circuit_tripped is True + left_open = [r for r in caplog.records if "finalizing span left open" in r.getMessage()] + assert left_open == [], "the circuit breaker already said why once; per-span warnings bury it" + + def test_a_marker_from_another_thread_is_enqueued_on_the_loop_thread(self, ingest_server) -> None: + url, recorder = ingest_server + threads: Dict[str, Any] = {} + + async def scenario(sender: AudioChunkSender) -> None: + # OTel invokes span callbacks on whichever thread ended the span, and + # asyncio.Queue is not thread-safe. The marker must therefore reach the + # queue from the loop's own thread, never from the foreign one. + threads["loop"] = threading.get_ident() + enqueued_from: List[int] = [] + original_put = sender._queue.put_nowait + + def recording_put(message: Any) -> None: + enqueued_from.append(threading.get_ident()) + original_put(message) + + sender._queue.put_nowait = recording_put # type: ignore[method-assign] + worker = threading.Thread( + target=lambda: sender.mark_audio_end(role=SpeakerRole.AGENT, span_id=AGENT_SPAN_ID) + ) + worker.start() + await asyncio.to_thread(worker.join) + threads["worker"] = worker.ident + await asyncio.sleep(0.1) + threads["enqueued_from"] = enqueued_from + + run_call(url, scenario) + + assert threads["enqueued_from"], "the marker never reached the queue" + off_loop = [ident for ident in threads["enqueued_from"] if ident != threads["loop"]] + assert off_loop == [], ( + f"queue touched from thread(s) {off_loop} instead of the loop thread " + f"{threads['loop']} (the worker was {threads['worker']})" + ) + terminators = [r for r in recorder.chunks_for(AGENT_SPAN_ID) if r.is_last] + assert len(terminators) == 1 + + +# --------------------------------------------------------------------------- +# Coordinator +# --------------------------------------------------------------------------- + + +class TestSessionAudioCoordinator: + def test_a_frame_inside_a_speaking_span_carries_that_span_id(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.USER, trace_id=TRACE_ID, span_id=USER_SPAN_ID) + + coordinator.on_frame(SpeakerRole.USER, make_frame()) + + kwargs = sender.enqueue.call_args.kwargs + assert kwargs["role"] is SpeakerRole.USER + assert kwargs["span_id"] == USER_SPAN_ID + assert kwargs["trace_id"] == TRACE_ID + + def test_a_frame_between_turns_is_sent_with_no_span_id(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator._session_trace_id = TRACE_ID + + coordinator.on_frame(SpeakerRole.USER, make_frame()) + + kwargs = sender.enqueue.call_args.kwargs + assert kwargs["span_id"] == "" + assert kwargs["trace_id"] == TRACE_ID + + def test_both_speakers_are_streamed(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + coordinator.on_speaking_start(SpeakerRole.USER, trace_id=TRACE_ID, span_id=USER_SPAN_ID) + + coordinator.on_frame(SpeakerRole.AGENT, make_frame()) + coordinator.on_frame(SpeakerRole.USER, make_frame()) + + # Capture is all of the call's audio or none of it — there is no per-role + # gate to leave one side out. + streamed = [call.kwargs["role"] for call in sender.enqueue.call_args_list] + assert streamed == [SpeakerRole.AGENT, SpeakerRole.USER] + + def test_closing_a_span_finalizes_its_recording(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.USER, trace_id=TRACE_ID, span_id=USER_SPAN_ID) + + coordinator.on_speaking_end(SpeakerRole.USER) + + sender.mark_audio_end.assert_called_once_with(role=SpeakerRole.USER, span_id=USER_SPAN_ID) + + def test_close_finalizes_every_span_still_recording(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.USER, trace_id=TRACE_ID, span_id=USER_SPAN_ID) + coordinator.on_speaking_start(SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + + coordinator.close() + + finalized = {call.kwargs["role"] for call in sender.mark_audio_end.call_args_list} + assert finalized == {SpeakerRole.USER, SpeakerRole.AGENT} + + def test_close_is_idempotent(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.USER, trace_id=TRACE_ID, span_id=USER_SPAN_ID) + + coordinator.close() + coordinator.close() + + assert sender.mark_audio_end.call_count == 1 + + +class TestSessionAudioCoordinatorInterrupts: + @staticmethod + def _interrupted_coordinator(sender: MagicMock) -> SessionAudioCoordinator: + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + coordinator.on_output_buffer_cleared() + return coordinator + + def test_agent_frames_stop_once_the_caller_cuts_in(self) -> None: + sender = MagicMock() + coordinator = self._interrupted_coordinator(sender) + + coordinator.on_frame(SpeakerRole.AGENT, make_frame()) + + sender.enqueue.assert_not_called() + + def test_the_playback_position_is_reported_as_the_audio_heard(self) -> None: + sender = MagicMock() + coordinator = self._interrupted_coordinator(sender) + + event = MagicMock(interrupted=True, playback_position=0.75) + coordinator.on_playback_finished(event) + + sender.interrupt_agent_span.assert_called_once_with(span_id=AGENT_SPAN_ID, playback_ms=750) + + def test_an_interrupted_span_is_not_finalized_at_its_full_length(self) -> None: + sender = MagicMock() + coordinator = self._interrupted_coordinator(sender) + + coordinator.on_speaking_end(SpeakerRole.AGENT) + + sender.mark_audio_end.assert_not_called() + + def test_the_span_id_survives_the_span_closing_before_the_interrupt_is_reported(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + # LiveKit's ordering: the span ends, and only then does clear_buffer fire. + coordinator.on_speaking_end(SpeakerRole.AGENT) + coordinator.on_output_buffer_cleared() + + coordinator.on_playback_finished(MagicMock(interrupted=True, playback_position=0.2)) + + sender.interrupt_agent_span.assert_called_once_with(span_id=AGENT_SPAN_ID, playback_ms=200) + + def test_playback_that_was_not_interrupted_needs_no_correction(self) -> None: + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + coordinator.on_speaking_start(SpeakerRole.AGENT, trace_id=TRACE_ID, span_id=AGENT_SPAN_ID) + + coordinator.on_playback_finished(MagicMock(interrupted=False, playback_position=2.0)) + + sender.interrupt_agent_span.assert_not_called() + + +# --------------------------------------------------------------------------- +# Registry and span processor +# --------------------------------------------------------------------------- + + +def make_span(name: str, *, trace_id: int, span_id: int = 0xABCD) -> MagicMock: + """Build a span whose context reports the given ids.""" + span = MagicMock() + span.name = name + span.get_span_context.return_value = MagicMock(is_valid=True, trace_id=trace_id, span_id=span_id) + return span + + +class TestAudioCoordinatorRegistry: + def test_a_coordinator_is_handed_out_only_once(self) -> None: + registry = AudioCoordinatorRegistry() + coordinator = SessionAudioCoordinator() + registry.register(1234, coordinator) + + assert registry.unregister(1234) is coordinator + assert registry.unregister(1234) is None + assert registry.get(1234) is None + + def test_pop_all_drains_the_registry(self) -> None: + registry = AudioCoordinatorRegistry() + registry.register(1, SessionAudioCoordinator()) + registry.register(2, SessionAudioCoordinator()) + + assert len(registry.pop_all()) == 2 + assert registry.pop_all() == [] + + +class TestAudioSpanProcessor: + @pytest.fixture(autouse=True) + def _clear_registry(self): + audio_coordinators.pop_all() + yield + audio_coordinators.pop_all() + + @pytest.mark.parametrize( + "span_name,role", + [("user_speaking", SpeakerRole.USER), ("agent_speaking", SpeakerRole.AGENT)], + ) + def test_a_speaking_span_opens_and_closes_a_recording(self, span_name: str, role: SpeakerRole) -> None: + trace_id = 0xAAAABBBBCCCCDDDD + sender = MagicMock() + coordinator = SessionAudioCoordinator(sender=sender) + audio_coordinators.register(trace_id, coordinator) + processor = AudioSpanProcessor() + + span = make_span(span_name, trace_id=trace_id, span_id=0x1234567890ABCDEF) + processor.on_start(span) + + assert sender.enqueue.call_count == 0 + coordinator.on_frame(role, make_frame()) + assert sender.enqueue.call_args.kwargs["span_id"] == format(0x1234567890ABCDEF, "016x") + + processor.on_end(span) + sender.mark_audio_end.assert_called_once_with(role=role, span_id=format(0x1234567890ABCDEF, "016x")) + + def test_a_span_from_another_call_is_ignored(self) -> None: + sender = MagicMock() + audio_coordinators.register(0x1111, SessionAudioCoordinator(sender=sender)) + processor = AudioSpanProcessor() + + processor.on_start(make_span("user_speaking", trace_id=0x2222)) + + sender.enqueue.assert_not_called() + + def test_a_span_that_is_not_speech_is_ignored(self) -> None: + processor = AudioSpanProcessor() + span = make_span("llm_request", trace_id=0x1111) + + processor.on_start(span) + + span.get_span_context.assert_not_called() + + +# --------------------------------------------------------------------------- +# Session wiring +# --------------------------------------------------------------------------- + + +class TestSessionWiring: + @pytest.fixture(autouse=True) + def _clear_registry(self): + audio_coordinators.pop_all() + yield + audio_coordinators.pop_all() + + def test_the_sender_is_built_from_the_configured_limits(self) -> None: + config = MagicMock( + api_key="key", + headers={"x-api-key": "key", "x-tenant": "acme"}, + audio_batch_interval_ms=250, + audio_batch_bytes=4096, + audio_max_request_bytes=65536, + audio_buffer_bytes=960_000, + ) + config.audio_endpoint.return_value = "https://ingest.example/v1/audio/chunk" + + sender = build_audio_sender(config, "session-1") + + assert sender is not None + assert sender._batch_interval_seconds == 0.25 + assert sender._flush_at_bytes == 4096 + assert sender._max_request_bytes == 65536 + assert sender._queue.maxsize == 1000 + # Only credential headers are forwarded, never arbitrary config headers. + assert sender._auth_headers == {"x-api-key": "key"} + + def test_no_endpoint_means_no_sender(self) -> None: + config = MagicMock() + config.audio_endpoint.return_value = None + + assert build_audio_sender(config, "session-1") is None + + def test_stopping_a_call_unregisters_it_and_records_what_was_sent(self) -> None: + sender = MagicMock() + sender.stats = MagicMock(bytes_sent=4096, chunks_sent=3, frames_dropped=1, errors=0, circuit_tripped=False) + sender.end_session = _async_noop + coordinator = SessionAudioCoordinator(sender=sender) + audio_coordinators.register(0x99, coordinator) + session_span = MagicMock() + + asyncio.run(stop_audio_capture(0x99, session_span=session_span)) + + assert audio_coordinators.get(0x99) is None + stamped = session_span.set_attributes.call_args.args[0] + assert stamped[NETRA_AUDIO_SENT_BYTES] == 4096 + assert stamped[NETRA_AUDIO_SENT_CHUNKS] == 3 + assert stamped[NETRA_AUDIO_DROPPED_FRAMES] == 1 + + def test_stopping_a_call_twice_is_harmless(self) -> None: + sender = MagicMock() + sender.end_session = _async_noop + audio_coordinators.register(0x99, SessionAudioCoordinator(sender=sender)) + session_span = MagicMock() + + asyncio.run(stop_audio_capture(0x99, session_span=session_span)) + asyncio.run(stop_audio_capture(0x99, session_span=session_span)) + + assert session_span.set_attributes.call_count == 1 + + def test_a_failed_attach_leaves_no_sender_running(self, ingest_server) -> None: + url, _ = ingest_server + config = MagicMock( + api_key="key", + headers={"x-api-key": "key"}, + audio_batch_interval_ms=1000, + audio_batch_bytes=32768, + audio_max_request_bytes=262144, + audio_buffer_bytes=2097152, + ) + config.audio_endpoint.return_value = url + + # A custom AudioOutput whose capture_frame cannot be reassigned, which is + # what attach() trips over. + session = MagicMock() + session.input.audio = None + unpatchable = MagicMock() + type(unpatchable).capture_frame = property(lambda self: _async_noop) + session.output.audio = unpatchable + + async def drive() -> List[asyncio.Task]: + await start_audio_capture(session, config=config, session_id="s", trace_id=0xABC) + await asyncio.sleep(0.05) + return [task for task in asyncio.all_tasks() if task.get_name() == "netra-audio-chunk-sender"] + + leaked = asyncio.run(drive()) + + # The sender owns a background task and an HTTP client from start() + # onwards; if attach() fails after that, nothing else can ever close them. + assert leaked == [], "a started sender was stranded with no coordinator registered to close it" + assert audio_coordinators.get(0xABC) is None + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +class TestAudioConfigResolution: + def test_a_missing_credential_is_reported_once_not_once_per_session(self, monkeypatch, caplog) -> None: + for name in list(os.environ): + if name.startswith(("NETRA_", "OTEL_")): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("NETRA_AUDIO_ENDPOINT", "https://ingest.example/v1/audio/chunk") + + with caplog.at_level("WARNING"): + config = Config() + # Every per-session hook asks; the answer is fixed at init time. + for _ in range(5): + assert config.audio_endpoint() is None + assert config.audio_capture_enabled is False + + missing_credential = [r for r in caplog.records if "no credential is configured" in r.getMessage()] + assert len(missing_credential) == 1 + + def test_a_resolved_endpoint_is_the_whole_gate(self, monkeypatch) -> None: + for name in list(os.environ): + if name.startswith(("NETRA_", "OTEL_")): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://collector.getnetra.com") + monkeypatch.setenv("NETRA_AUDIO_ENDPOINT", "https://ingest.example/v1/audio/chunk") + monkeypatch.setenv("NETRA_API_KEY", "key") + # A leftover role list from before capture became all-or-nothing must not + # still gate anything. + monkeypatch.setenv("NETRA_AUDIO_ROLES", "") + + config = Config() + + assert config.audio_capture_enabled is True + assert not hasattr(config, "audio_roles") diff --git a/tests/test_livekit_instrumentation.py b/tests/test_livekit_instrumentation.py new file mode 100644 index 0000000..9c1729a --- /dev/null +++ b/tests/test_livekit_instrumentation.py @@ -0,0 +1,835 @@ +"""Tests for Netra's LiveKit voice-agent instrumentation. + +The suite exercises the package through the real OpenTelemetry SDK rather than +mocks: spans are created from a tracer whose instrumentation scope is +``livekit-agents`` — the only thing the processor gates on — so no +``livekit-agents`` install is required. +""" + +import asyncio +import json +from typing import Any, Dict, List, Optional, Tuple + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from wrapt import ObjectProxy + +from netra.instrumentation import livekit as livekit_instrumentation +from netra.instrumentation.livekit import NetraLiveKitInstrumentor +from netra.instrumentation.livekit.provider_binding import _ShieldedTracerProvider +from netra.instrumentation.livekit.trace_processor import SpanMappingProcessor +from netra.instrumentation.livekit.utils import ( + LIVEKIT_SCOPE_NAME, + MAX_CONVERSATION_MESSAGES_PER_SIDE, + NETRA_CONVERSATION_TRUNCATED, + NETRA_SPAN_TYPE, + ConversationSide, + content_of_choice_event, + content_of_event, + conversation_from_attributes, + is_zero_usage, + messages_for_parent, + messages_from_chat_ctx, + netra_span_type_for, + role_of_choice_event, + tts_pricing_attributes_from, +) +from netra.span_wrapper import SpanType + +pytestmark = pytest.mark.unit + + +class _Harness: + """A tracer provider carrying the LiveKit processor and an in-memory exporter.""" + + def __init__(self) -> None: + self.exporter = InMemorySpanExporter() + self.provider = TracerProvider() + # Registration order mirrors production: the exporting processor is + # installed by netra/tracer.py first, the LiveKit one by _instrument(). + self.provider.add_span_processor(SimpleSpanProcessor(self.exporter)) + self.provider.add_span_processor(SpanMappingProcessor()) + self.livekit_tracer = self.provider.get_tracer(LIVEKIT_SCOPE_NAME) + + def tracer(self, scope_name: str) -> Any: + """Return a tracer for some other instrumentation scope.""" + return self.provider.get_tracer(scope_name) + + def finished(self, name: str) -> ReadableSpan: + """Return the single finished span called *name*.""" + matches = [span for span in self.exporter.get_finished_spans() if span.name == name] + assert len(matches) == 1, f"expected exactly one {name!r} span, got {len(matches)}" + return matches[0] + + def attributes(self, name: str) -> Dict[str, Any]: + """Return the exported attributes of the finished span called *name*.""" + return dict(self.finished(name).attributes or {}) + + +@pytest.fixture +def harness() -> _Harness: + return _Harness() + + +def _record(harness: _Harness, span_name: str, attributes: Dict[str, Any]) -> Dict[str, Any]: + """Run one LiveKit span through the processor and return its exported attributes.""" + span = harness.livekit_tracer.start_span(span_name) + for key, value in attributes.items(): + span.set_attribute(key, value) + span.end() + return harness.attributes(span_name) + + +def _messages(attributes: Dict[str, Any], side: str) -> List[Tuple[str, str]]: + """Read the indexed ``gen_ai..N.role/content`` pairs back, in index order.""" + indices = sorted( + int(key.split(".")[2]) for key in attributes if key.startswith(f"gen_ai.{side}.") and key.endswith(".content") + ) + return [ + (attributes[f"gen_ai.{side}.{index}.role"], attributes[f"gen_ai.{side}.{index}.content"]) for index in indices + ] + + +class TestSpanTypeMapping: + @pytest.mark.parametrize( + "span_name,expected", + [ + ("agent_turn", SpanType.AGENT.value), + ("llm_node", SpanType.GENERATION.value), + ("llm_request", SpanType.GENERATION.value), + ("function_tool", SpanType.TOOL.value), + ("tts_request", SpanType.GENERATION.value), + ("agent_session", SpanType.SPAN.value), + ("something_unmapped", SpanType.SPAN.value), + (None, SpanType.SPAN.value), + ("", SpanType.SPAN.value), + ], + ) + def test_returns_span_type_for_name(self, span_name: Optional[str], expected: str) -> None: + assert netra_span_type_for(span_name) == expected + + def test_span_type_is_stamped_on_livekit_spans(self, harness: _Harness) -> None: + assert _record(harness, "agent_turn", {})["netra.span.type"] == SpanType.AGENT.value + + def test_job_entrypoint_is_marked_as_a_workflow(self, harness: _Harness) -> None: + assert _record(harness, "job_entrypoint", {})["netra.entity.type"] == "workflow" + + def test_non_entity_spans_carry_no_entity_marker(self, harness: _Harness) -> None: + assert "netra.entity.type" not in _record(harness, "agent_turn", {}) + + @pytest.mark.parametrize( + "span_name,expected", + [("agent_session", "session"), ("agent_turn", "span"), ("user_turn", "span")], + ) + def test_audio_type_is_stamped_on_interaction_spans(self, harness: _Harness, span_name: str, expected: str) -> None: + assert _record(harness, span_name, {})["netra.audio.type"] == expected + + def test_audio_type_is_absent_from_other_spans(self, harness: _Harness) -> None: + assert "netra.audio.type" not in _record(harness, "llm_node", {}) + + def test_spans_from_other_scopes_are_left_untouched(self, harness: _Harness) -> None: + span = harness.tracer("openai").start_span("agent_turn") + span.set_attribute("lk.function_tool.name", "lookup") + span.end() + + attributes = harness.attributes("agent_turn") + assert "netra.span.type" not in attributes + assert "netra.tool.name" not in attributes + assert attributes["lk.function_tool.name"] == "lookup" + + +class TestAttributeMirroring: + @pytest.mark.parametrize( + "lk_key,netra_key", + [ + ("lk.function_tool.name", "netra.tool.name"), + ("lk.function_tool.arguments", "input"), + ("lk.function_tool.output", "output"), + ("lk.response.ttft", "netra.latency.ttft"), + ("lk.response.ttfb", "netra.latency.ttfb"), + ("lk.e2e_latency", "netra.latency.e2e"), + ("lk.end_of_turn_delay", "netra.latency.end_of_turn_delay"), + ("lk.transcript_confidence", "netra.stt.confidence"), + ("lk.interrupted", "netra.turn.interrupted"), + ], + ) + def test_mapped_attribute_is_mirrored_and_original_preserved( + self, harness: _Harness, lk_key: str, netra_key: str + ) -> None: + attributes = _record(harness, "function_tool", {lk_key: 0.25}) + + assert attributes[netra_key] == 0.25 + assert attributes[lk_key] == 0.25, "the original lk.* attribute must be preserved" + + def test_unmapped_attribute_is_written_through_unchanged(self, harness: _Harness) -> None: + attributes = _record(harness, "agent_turn", {"lk.speech_id": "sp_1"}) + + assert attributes["lk.speech_id"] == "sp_1" + + @pytest.mark.parametrize("empty", ["", None]) + def test_empty_value_is_not_mirrored(self, harness: _Harness, empty: Any) -> None: + span = harness.livekit_tracer.start_span("function_tool") + span.set_attribute("lk.function_tool.name", empty) + span.end() + + assert "netra.tool.name" not in harness.attributes("function_tool") + + def test_set_attributes_plural_goes_through_the_mapping(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("user_turn") + span.set_attributes({"lk.user_transcript": "book me a table", "lk.e2e_latency": 1.5}) + span.end() + + attributes = harness.attributes("user_turn") + assert attributes["netra.latency.e2e"] == 1.5 + assert _messages(attributes, "completion") == [("user", "book me a table")] + + def test_empty_set_attributes_is_a_no_op(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("agent_turn") + span.set_attributes({}) + span.end() + + assert "gen_ai.prompt.0.content" not in harness.attributes("agent_turn") + + +class TestUsageAttributes: + def test_reported_usage_is_marked_as_framework_sourced(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"gen_ai.usage.input_tokens": 120}) + + assert attributes["gen_ai.usage.input_tokens"] == 120 + assert attributes["netra.usage.source"] == "framework" + + def test_zero_usage_is_dropped_rather_than_claimed(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"gen_ai.usage.input_tokens": 0}) + + assert "gen_ai.usage.input_tokens" not in attributes + assert "netra.usage.source" not in attributes + + @pytest.mark.parametrize( + "value,expected", + [(0, True), (0.0, True), (1, False), (-1, False), (True, False), (False, False), (None, False), ("0", False)], + ) + def test_is_zero_usage_only_matches_numeric_zero(self, value: Any, expected: bool) -> None: + assert is_zero_usage(value) is expected + + +class TestConversationMapping: + def test_agent_turn_reads_as_system_then_user_then_reply(self, harness: _Harness) -> None: + attributes = _record( + harness, + "agent_turn", + { + "lk.instructions": "You are a helpful agent.", + "lk.user_input": "what is the weather?", + "lk.response.text": "It is sunny.", + }, + ) + + assert _messages(attributes, "prompt") == [ + ("system", "You are a helpful agent."), + ("user", "what is the weather?"), + ] + assert _messages(attributes, "completion") == [("assistant", "It is sunny.")] + + def test_tts_input_text_is_attributed_to_the_assistant(self, harness: _Harness) -> None: + attributes = _record(harness, "tts_request", {"lk.input_text": "It is sunny."}) + + assert _messages(attributes, "prompt") == [("assistant", "It is sunny.")] + + def test_user_transcript_is_the_callers_words(self, harness: _Harness) -> None: + attributes = _record(harness, "user_turn", {"lk.user_transcript": "hello there"}) + + assert _messages(attributes, "completion") == [("user", "hello there")] + + def test_chat_ctx_is_expanded_into_indexed_prompts(self, harness: _Harness) -> None: + chat_ctx = json.dumps( + { + "items": [ + {"type": "message", "role": "system", "content": ["Be brief."]}, + {"type": "agent_config_update", "role": "system", "content": ["ignored"]}, + {"type": "message", "role": "user", "content": ["hi"]}, + ] + } + ) + attributes = _record(harness, "llm_node", {"lk.chat_ctx": chat_ctx}) + + assert _messages(attributes, "prompt") == [("system", "Be brief."), ("user", "hi")] + assert attributes["lk.chat_ctx"] == chat_ctx, "the original blob must be preserved" + + def test_sources_on_one_span_share_the_index_counters(self, harness: _Harness) -> None: + chat_ctx = json.dumps({"items": [{"type": "message", "role": "system", "content": ["Be brief."]}]}) + attributes = _record( + harness, + "agent_turn", + {"lk.chat_ctx": chat_ctx, "lk.user_input": "hi", "lk.response.text": "hello"}, + ) + + assert _messages(attributes, "prompt") == [("system", "Be brief."), ("user", "hi")] + assert _messages(attributes, "completion") == [("assistant", "hello")] + + def test_conversation_content_does_not_write_input_directly(self, harness: _Harness) -> None: + attributes = _record(harness, "agent_turn", {"lk.user_input": "hi"}) + + assert "input" not in attributes, "input is assembled downstream by SpanIOProcessor" + + +# A conversation long enough that expanding all of it would overflow the OTel +# default attribute budget (2 attributes per message vs OTEL_SPAN_ATTRIBUTE_COUNT_LIMIT +# = 128), so the tests below fail on an unbounded recorder rather than merely +# asserting the cap's arithmetic. +_OVERFLOWING_MESSAGE_COUNT = MAX_CONVERSATION_MESSAGES_PER_SIDE * 4 +_DEFAULT_SPAN_ATTRIBUTE_LIMIT = 128 + + +def _chat_ctx(num_messages: int) -> str: + """A serialised ChatContext of *num_messages* alternating turns.""" + items = [ + { + "type": "message", + "role": "user" if index % 2 == 0 else "assistant", + "content": [f"turn {index}"], + } + for index in range(num_messages) + ] + return json.dumps({"items": items}) + + +class TestConversationIsBounded: + """The conversation must not grow past the span's bounded attribute capacity. + + OTel's ``BoundedAttributes`` evicts the *oldest* attribute on overflow, so an + unbounded sequence deletes the markers stamped in ``on_start`` rather than + dropping its own tail. See ``MAX_CONVERSATION_MESSAGES_PER_SIDE``. + """ + + def test_long_chat_ctx_does_not_evict_the_netra_markers(self, harness: _Harness) -> None: + # The regression: the unbounded expansion pushed the span past the + # 128-attribute default, and the markers — written first, in on_start — + # were the first things evicted. + attributes = _record(harness, "llm_node", {"lk.chat_ctx": _chat_ctx(_OVERFLOWING_MESSAGE_COUNT)}) + + assert attributes[NETRA_SPAN_TYPE] == SpanType.GENERATION.value + assert attributes["lk.chat_ctx"], "the original blob must still be preserved" + assert len(attributes) < _DEFAULT_SPAN_ATTRIBUTE_LIMIT + + def test_long_chat_ctx_keeps_the_newest_turns(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"lk.chat_ctx": _chat_ctx(_OVERFLOWING_MESSAGE_COUNT)}) + prompts = _messages(attributes, "prompt") + + assert len(prompts) == MAX_CONVERSATION_MESSAGES_PER_SIDE + # The tail of the conversation survives; its opening is dropped. + first_kept = _OVERFLOWING_MESSAGE_COUNT - MAX_CONVERSATION_MESSAGES_PER_SIDE + assert prompts[0][1] == f"turn {first_kept}" + assert prompts[-1][1] == f"turn {_OVERFLOWING_MESSAGE_COUNT - 1}" + + def test_truncation_is_marked_on_the_span(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"lk.chat_ctx": _chat_ctx(_OVERFLOWING_MESSAGE_COUNT)}) + + assert attributes[NETRA_CONVERSATION_TRUNCATED] is True + + def test_a_short_conversation_is_not_marked_truncated(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"lk.chat_ctx": _chat_ctx(4)}) + + assert len(_messages(attributes, "prompt")) == 4 + assert NETRA_CONVERSATION_TRUNCATED not in attributes + + def test_a_conversation_exactly_at_the_cap_is_not_marked(self, harness: _Harness) -> None: + attributes = _record(harness, "llm_node", {"lk.chat_ctx": _chat_ctx(MAX_CONVERSATION_MESSAGES_PER_SIDE)}) + + assert len(_messages(attributes, "prompt")) == MAX_CONVERSATION_MESSAGES_PER_SIDE + assert NETRA_CONVERSATION_TRUNCATED not in attributes + + def test_unbounded_events_are_capped_too(self, harness: _Harness) -> None: + # livekit-agents emits one conversation event per context item on + # llm_request (llm/llm.py -> _chat_ctx_to_otel_events), so the event path + # grows with the call exactly like lk.chat_ctx does. + span = harness.livekit_tracer.start_span("llm_request") + for index in range(_OVERFLOWING_MESSAGE_COUNT): + span.add_event("gen_ai.user.message", {"content": f"turn {index}"}) + span.end() + attributes = harness.attributes("llm_request") + + assert len(_messages(attributes, "prompt")) == MAX_CONVERSATION_MESSAGES_PER_SIDE + assert attributes[NETRA_CONVERSATION_TRUNCATED] is True + assert attributes[NETRA_SPAN_TYPE] == SpanType.GENERATION.value + + def test_the_reply_survives_a_saturated_prompt_side(self, harness: _Harness) -> None: + # The two sides count independently, so a long context can never crowd + # out the completion — the single most important value on the span. + attributes = _record( + harness, + "llm_node", + {"lk.chat_ctx": _chat_ctx(_OVERFLOWING_MESSAGE_COUNT), "lk.response.text": "THE REPLY"}, + ) + + assert _messages(attributes, "completion") == [("assistant", "THE REPLY")] + + def test_propagated_child_content_is_capped(self, harness: _Harness) -> None: + # A provider span's own indexed prompts also grow with the conversation, + # and llm_request_run inherits them wholesale. + parent = harness.livekit_tracer.start_span("llm_request_run") + child_attributes: Dict[str, Any] = {} + for index in range(_OVERFLOWING_MESSAGE_COUNT): + child_attributes[f"gen_ai.prompt.{index}.role"] = "user" + child_attributes[f"gen_ai.prompt.{index}.content"] = f"turn {index}" + with trace.use_span(parent, end_on_exit=False): + child = harness.tracer("openai").start_span("openai.chat", attributes=child_attributes) + child.end() + parent.end() + attributes = harness.attributes("llm_request_run") + + assert len(_messages(attributes, "prompt")) == MAX_CONVERSATION_MESSAGES_PER_SIDE + assert attributes[NETRA_CONVERSATION_TRUNCATED] is True + + +class TestConversationEvents: + @pytest.mark.parametrize( + "event_name,role", + [ + ("gen_ai.system.message", "system"), + ("gen_ai.user.message", "user"), + ("gen_ai.assistant.message", "assistant"), + ("gen_ai.tool.message", "tool"), + ], + ) + def test_conversation_event_becomes_an_indexed_prompt(self, harness: _Harness, event_name: str, role: str) -> None: + span = harness.livekit_tracer.start_span("llm_request") + span.add_event(event_name, {"content": "some text"}) + span.end() + + assert _messages(harness.attributes("llm_request"), "prompt") == [(role, "some text")] + + def test_choice_event_becomes_an_indexed_completion(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("llm_request") + span.add_event("gen_ai.choice", {"role": "assistant", "content": "the reply"}) + span.end() + + assert _messages(harness.attributes("llm_request"), "completion") == [("assistant", "the reply")] + + def test_event_is_still_recorded_on_the_span(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("llm_request") + span.add_event("gen_ai.user.message", {"content": "hi"}) + span.end() + + events = harness.finished("llm_request").events + assert [event.name for event in events] == ["gen_ai.user.message"] + + def test_unrelated_event_is_recorded_but_not_mapped(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("llm_request") + span.add_event("some.other.event", {"content": "hi"}) + span.end() + + attributes = harness.attributes("llm_request") + assert "gen_ai.prompt.0.content" not in attributes + assert [event.name for event in harness.finished("llm_request").events] == ["some.other.event"] + + def test_event_without_content_contributes_no_message(self, harness: _Harness) -> None: + span = harness.livekit_tracer.start_span("llm_request") + span.add_event("gen_ai.user.message", {"role": "user"}) + span.end() + + assert "gen_ai.prompt.0.content" not in harness.attributes("llm_request") + + +class TestTtsPricing: + def test_priceable_fields_are_lifted_out_of_the_metrics_blob(self, harness: _Harness) -> None: + metrics = json.dumps({"characters_count": 42, "metadata": {"model_name": "cartesia/sonic-3"}}) + attributes = _record(harness, "tts_request", {"lk.tts_metrics": metrics}) + + assert attributes["gen_ai.request.model"] == "cartesia/sonic-3" + assert attributes["gen_ai.usage.prompt.character_count"] == 42 + assert attributes["lk.tts_metrics"] == metrics, "the original blob must be preserved" + + def test_lifted_character_count_is_marked_framework_sourced(self, harness: _Harness) -> None: + metrics = json.dumps({"characters_count": 42, "metadata": {"model_name": "cartesia/sonic-3"}}) + attributes = _record(harness, "tts_request", {"lk.tts_metrics": metrics}) + + assert attributes["netra.usage.source"] == "framework" + + @pytest.mark.parametrize( + "payload,expected_model,expected_count", + [ + ({"characters_count": 7, "metadata": {"model_name": "sonic"}}, "sonic", 7), + ({"characters_count": 7.9, "metadata": {"model_name": "sonic"}}, "sonic", 7), + ({"characters_count": 0, "metadata": {"model_name": "sonic"}}, "sonic", None), + ({"characters_count": -3, "metadata": {"model_name": "sonic"}}, "sonic", None), + ({"characters_count": True, "metadata": {"model_name": "sonic"}}, "sonic", None), + ({"characters_count": 7}, None, 7), + ({"characters_count": 7, "metadata": {"model_name": ""}}, None, 7), + ({"metadata": {"model_name": "sonic"}}, "sonic", None), + ("not json", None, None), + (None, None, None), + ([], None, None), + ], + ) + def test_extraction_tolerates_every_shape( + self, payload: Any, expected_model: Optional[str], expected_count: Optional[int] + ) -> None: + pricing = tts_pricing_attributes_from(payload) + + assert pricing.model == expected_model + assert pricing.character_count == expected_count + + def test_accepts_a_json_string_and_a_dict_identically(self) -> None: + payload = {"characters_count": 9, "metadata": {"model_name": "sonic"}} + + assert tts_pricing_attributes_from(payload) == tts_pricing_attributes_from(json.dumps(payload)) + + +class TestChildToParentPropagation: + def _child_under(self, harness: _Harness, parent_name: str, scope: str, attributes: Dict[str, Any]) -> None: + parent = harness.livekit_tracer.start_span(parent_name) + with trace.use_span(parent, end_on_exit=False): + child = harness.tracer(scope).start_span("provider.call") + for key, value in attributes.items(): + child.set_attribute(key, value) + child.end() + parent.end() + + def test_provider_span_content_is_lifted_onto_llm_request_run(self, harness: _Harness) -> None: + self._child_under( + harness, + "llm_request_run", + "openai", + { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "hi", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "hello", + }, + ) + + attributes = harness.attributes("llm_request_run") + assert _messages(attributes, "prompt") == [("user", "hi")] + assert _messages(attributes, "completion") == [("assistant", "hello")] + + def test_non_llm_child_input_is_not_copied_up_as_a_message(self, harness: _Harness) -> None: + self._child_under(harness, "llm_request_run", "httpx", {"input": "POST https://api.example.com/v1/chat"}) + + attributes = harness.attributes("llm_request_run") + assert "gen_ai.prompt.0.content" not in attributes + + def test_llm_child_raw_io_is_copied_when_it_has_no_indexed_pairs(self, harness: _Harness) -> None: + self._child_under( + harness, + "llm_request_run", + "openai", + {"gen_ai.request.model": "gpt-4o", "input": "hi", "output": "hello"}, + ) + + attributes = harness.attributes("llm_request_run") + assert _messages(attributes, "prompt") == [("user", "hi")] + assert _messages(attributes, "completion") == [("assistant", "hello")] + + def test_tts_node_inherits_from_its_tts_request_child(self, harness: _Harness) -> None: + parent = harness.livekit_tracer.start_span("tts_node") + with trace.use_span(parent, end_on_exit=False): + child = harness.livekit_tracer.start_span("tts_request") + child.set_attribute("lk.input_text", "It is sunny.") + child.end() + parent.end() + + assert _messages(harness.attributes("tts_node"), "prompt") == [("assistant", "It is sunny.")] + + def test_spans_not_awaiting_child_content_are_unaffected(self, harness: _Harness) -> None: + self._child_under( + harness, + "agent_turn", + "openai", + {"gen_ai.prompt.0.role": "user", "gen_ai.prompt.0.content": "hi"}, + ) + + assert "gen_ai.prompt.0.content" not in harness.attributes("agent_turn") + + def test_content_arriving_after_the_parent_ended_is_dropped(self, harness: _Harness) -> None: + parent = harness.livekit_tracer.start_span("llm_request_run") + with trace.use_span(parent, end_on_exit=False): + child = harness.tracer("openai").start_span("provider.call") + child.set_attribute("gen_ai.prompt.0.role", "user") + child.set_attribute("gen_ai.prompt.0.content", "hi") + parent.end() + child.end() + + assert "gen_ai.prompt.0.content" not in harness.attributes("llm_request_run") + + +class TestConversationReading: + def test_indexed_pairs_are_ordered_numerically_not_lexically(self) -> None: + conversation = conversation_from_attributes( + { + "gen_ai.prompt.10.role": "user", + "gen_ai.prompt.10.content": "eleventh", + "gen_ai.prompt.2.role": "user", + "gen_ai.prompt.2.content": "third", + } + ) + + assert conversation.prompts == [("user", "third"), ("user", "eleventh")] + + def test_plural_prompts_form_is_accepted(self) -> None: + conversation = conversation_from_attributes({"gen_ai.prompts.0.role": "user", "gen_ai.prompts.0.content": "hi"}) + + assert conversation.prompts == [("user", "hi")] + + def test_role_without_content_is_not_a_message(self) -> None: + conversation = conversation_from_attributes({"gen_ai.prompt.0.role": "user"}) + + assert conversation.prompts == [] + + def test_content_without_role_keeps_the_text(self) -> None: + conversation = conversation_from_attributes({"gen_ai.prompt.0.content": "hi"}) + + assert conversation.prompts == [("", "hi")] + + def test_raw_input_is_suppressed_when_indexed_pairs_exist(self) -> None: + conversation = conversation_from_attributes( + {"gen_ai.prompt.0.role": "user", "gen_ai.prompt.0.content": "hi", "input": "hi"} + ) + + assert conversation.raw_input is None + + def test_raw_io_is_kept_when_there_are_no_indexed_pairs(self) -> None: + conversation = conversation_from_attributes({"input": "hi", "output": "hello"}) + + assert (conversation.raw_input, conversation.raw_output) == ("hi", "hello") + + @pytest.mark.parametrize("attributes,expected", [({"gen_ai.request.model": "x"}, True), ({"input": "hi"}, False)]) + def test_gen_ai_authorship_is_detected(self, attributes: Dict[str, Any], expected: bool) -> None: + assert conversation_from_attributes(attributes).carries_gen_ai is expected + + def test_empty_attributes_yield_an_empty_conversation(self) -> None: + conversation = conversation_from_attributes(None) + + assert conversation.prompts == [] + assert conversation.completions == [] + assert conversation.carries_gen_ai is False + + def test_raw_io_is_dropped_when_not_allowed(self) -> None: + conversation = conversation_from_attributes({"input": "hi", "output": "hello"}) + + assert messages_for_parent(conversation, allow_raw_io=False) == [] + + def test_raw_io_takes_fallback_roles_when_allowed(self) -> None: + conversation = conversation_from_attributes({"input": "hi", "output": "hello"}) + + messages = messages_for_parent(conversation, allow_raw_io=True) + assert [(message.side, message.role, message.content) for message in messages] == [ + (ConversationSide.PROMPT, "user", "hi"), + (ConversationSide.COMPLETION, "assistant", "hello"), + ] + + +class TestChatContextParsing: + def test_string_and_dict_payloads_agree(self) -> None: + payload = {"items": [{"type": "message", "role": "user", "content": ["hi"]}]} + + assert messages_from_chat_ctx(payload) == messages_from_chat_ctx(json.dumps(payload)) == [("user", "hi")] + + def test_multiple_text_parts_are_joined_by_newline(self) -> None: + payload = {"items": [{"type": "message", "role": "user", "content": ["one", "two"]}]} + + assert messages_from_chat_ctx(payload) == [("user", "one\ntwo")] + + def test_non_text_content_parts_are_skipped(self) -> None: + payload = {"items": [{"type": "message", "role": "user", "content": [{"type": "image"}, "caption"]}]} + + assert messages_from_chat_ctx(payload) == [("user", "caption")] + + @pytest.mark.parametrize( + "payload", + [ + "not json", + None, + [], + {"items": "not a list"}, + {}, + {"items": [{"type": "message", "role": "", "content": ["hi"]}]}, + {"items": [{"type": "message", "role": "user", "content": []}]}, + {"items": [{"type": "message", "role": "user"}]}, + {"items": ["not a mapping"]}, + ], + ) + def test_malformed_payloads_yield_no_messages(self, payload: Any) -> None: + assert messages_from_chat_ctx(payload) == [] + + +class TestEventPayloadParsing: + def test_choice_content_is_preferred_over_tool_calls(self) -> None: + assert content_of_choice_event({"content": "reply", "tool_calls": ['{"name": "x"}']}) == "reply" + + def test_tool_only_reply_falls_back_to_the_tool_calls(self) -> None: + attributes = {"tool_calls": ['{"name": "lookup"}', '{"name": "book"}']} + + assert content_of_choice_event(attributes) == '[{"name": "lookup"}, {"name": "book"}]' + + def test_tool_call_fallback_is_valid_json(self) -> None: + rendered = content_of_choice_event({"tool_calls": ['{"name": "lookup"}']}) + + assert json.loads(str(rendered)) == [{"name": "lookup"}] + + @pytest.mark.parametrize("attributes", [None, {}, {"content": ""}, {"tool_calls": []}, {"tool_calls": ""}]) + def test_empty_choice_events_carry_no_content(self, attributes: Any) -> None: + assert content_of_choice_event(attributes) is None + + @pytest.mark.parametrize( + "attributes,expected", + [({"role": "tool"}, "tool"), ({}, "assistant"), ({"role": ""}, "assistant"), (None, "assistant")], + ) + def test_choice_role_defaults_to_assistant(self, attributes: Any, expected: str) -> None: + assert role_of_choice_event(attributes) == expected + + @pytest.mark.parametrize("attributes", [None, {}, {"content": ""}, {"content": None}]) + def test_event_without_content_returns_none(self, attributes: Any) -> None: + assert content_of_event(attributes) is None + + def test_non_string_event_content_is_stringified(self) -> None: + assert content_of_event({"content": 42}) == "42" + + +class TestShieldedTracerProvider: + def test_livekit_added_processors_are_refused(self) -> None: + delegate = TracerProvider() + shield = _ShieldedTracerProvider(delegate) + exporter = InMemorySpanExporter() + + shield.add_span_processor(SimpleSpanProcessor(exporter)) + delegate.get_tracer("x").start_span("s").end() + + assert exporter.get_finished_spans() == () + + def test_shutdown_is_absorbed(self) -> None: + delegate = TracerProvider() + exporter = InMemorySpanExporter() + delegate.add_span_processor(SimpleSpanProcessor(exporter)) + + _ShieldedTracerProvider(delegate).shutdown() + delegate.get_tracer("x").start_span("s").end() + + assert len(exporter.get_finished_spans()) == 1, "the delegate's pipeline must survive LiveKit's teardown" + + def test_get_tracer_delegates(self) -> None: + delegate = TracerProvider() + + assert _ShieldedTracerProvider(delegate).get_tracer("x") is delegate.get_tracer("x") + + def test_resource_is_the_delegates(self) -> None: + delegate = TracerProvider(resource=Resource.create({"service.name": "voice-agent"})) + + assert _ShieldedTracerProvider(delegate).resource is delegate.resource + + def test_resource_falls_back_when_the_delegate_has_none(self) -> None: + class _ApiOnlyProvider: + pass + + shield = _ShieldedTracerProvider(_ApiOnlyProvider()) # type: ignore[arg-type] + + assert shield.resource == Resource.get_empty() + + def test_force_flush_propagates(self) -> None: + calls: List[int] = [] + + class _RecordingProvider: + def force_flush(self, timeout_millis: int = 30000) -> bool: + calls.append(timeout_millis) + return True + + assert _ShieldedTracerProvider(_RecordingProvider()).force_flush(500) is True # type: ignore[arg-type] + assert calls == [500] + + def test_force_flush_tolerates_a_provider_without_one(self) -> None: + class _ApiOnlyProvider: + pass + + assert _ShieldedTracerProvider(_ApiOnlyProvider()).force_flush() is True # type: ignore[arg-type] + + +class TestSessionStartHook: + def test_start_result_is_returned_untouched(self) -> None: + from netra.instrumentation.livekit.wrappers import wrap_start + + async def fake_start(**kwargs: Any) -> str: + return "started" + + result = asyncio.run(wrap_start(fake_start, object(), (), {"room": None})) + + assert result == "started" + + def test_start_exceptions_propagate_unchanged(self) -> None: + from netra.instrumentation.livekit.wrappers import wrap_start + + async def failing_start(**kwargs: Any) -> None: + raise RuntimeError("livekit blew up") + + with pytest.raises(RuntimeError, match="livekit blew up"): + asyncio.run(wrap_start(failing_start, object(), (), {})) + + +@pytest.fixture +def fake_agent_session(monkeypatch: pytest.MonkeyPatch) -> Any: + """Install a stand-in ``livekit.agents.voice.agent_session`` module. + + The hook is installed and removed by module path, so the wrap/unwrap round + trip can be exercised without a livekit-agents install — the shape of the + module tree is the only thing that matters to it. + """ + import sys + from types import ModuleType + + class AgentSession: + async def start(self, agent: Any = None, **kwargs: Any) -> str: + return "started" + + modules: Dict[str, ModuleType] = {} + for path in ("livekit", "livekit.agents", "livekit.agents.voice", "livekit.agents.voice.agent_session"): + modules[path] = ModuleType(path) + modules["livekit.agents.voice.agent_session"].AgentSession = AgentSession # type: ignore[attr-defined] + modules["livekit.agents.voice"].agent_session = modules["livekit.agents.voice.agent_session"] # type: ignore[attr-defined] + modules["livekit.agents"].voice = modules["livekit.agents.voice"] # type: ignore[attr-defined] + modules["livekit"].agents = modules["livekit.agents"] # type: ignore[attr-defined] + for path, module in modules.items(): + monkeypatch.setitem(sys.modules, path, module) + + # The hook guard is a module global; leaving it set would leak into later tests. + monkeypatch.setattr(livekit_instrumentation, "_session_hook_installed", False) + yield AgentSession + monkeypatch.setattr(livekit_instrumentation, "_session_hook_installed", False) + + +class TestSessionHookLifecycle: + @staticmethod + def _is_wrapped(agent_session: Any) -> bool: + return isinstance(agent_session.start, ObjectProxy) + + def test_hook_wraps_agent_session_start(self, fake_agent_session: Any) -> None: + livekit_instrumentation._install_session_hook() + + assert self._is_wrapped(fake_agent_session) + + def test_uninstrument_actually_removes_the_wrapper(self, fake_agent_session: Any) -> None: + livekit_instrumentation._install_session_hook() + assert self._is_wrapped(fake_agent_session), "precondition: the hook is installed" + + NetraLiveKitInstrumentor()._uninstrument() + + assert not self._is_wrapped( + fake_agent_session + ), "unwrap() cannot walk a dotted attribute path and returns None instead of raising" + + def test_reinstalling_after_uninstrument_wraps_exactly_once(self, fake_agent_session: Any) -> None: + livekit_instrumentation._install_session_hook() + NetraLiveKitInstrumentor()._uninstrument() + livekit_instrumentation._install_session_hook() + + # A stale wrapper left behind by uninstrument would nest here, running the + # session-id hook twice per start(). + assert self._is_wrapped(fake_agent_session) + assert not isinstance(fake_agent_session.start.__wrapped__, ObjectProxy) + + def test_install_is_idempotent(self, fake_agent_session: Any) -> None: + livekit_instrumentation._install_session_hook() + livekit_instrumentation._install_session_hook() + + assert not isinstance(fake_agent_session.start.__wrapped__, ObjectProxy)