diff --git a/netra/__init__.py b/netra/__init__.py index 1f51d8b..63b4cb1 100644 --- a/netra/__init__.py +++ b/netra/__init__.py @@ -257,7 +257,12 @@ 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 diff --git a/netra/config.py b/netra/config.py index e6cf26b..21c6c03 100644 --- a/netra/config.py +++ b/netra/config.py @@ -1,11 +1,14 @@ import json +import logging import os -from typing import Any, Dict, List, Optional +from typing import Any, Dict, FrozenSet, List, Optional from opentelemetry.util.re import parse_env_headers 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,28 @@ _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") + +# The only recognised speaker roles. +AUDIO_ROLES: FrozenSet[str] = frozenset({"user", "agent"}) + +_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 +122,170 @@ 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 + ) + self.audio_roles = self._get_role_set("NETRA_AUDIO_ROLES") + self.audio_save_local = self._get_bool_config(None, "NETRA_AUDIO_SAVE_LOCAL", default=False) + + # 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 + + if not self.audio_roles: + logger.warning( + "netra.audio: NETRA_AUDIO_ROLES resolved empty; no call audio will be captured. " + "Traces are unaffected." + ) + + if self.audio_save_local: + logger.warning( + "netra.audio: NETRA_AUDIO_SAVE_LOCAL is enabled. Local WAV capture is a " + "development-only aid and retains full-session PCM in memory." + ) + + def _get_role_set(self, env_var: str) -> FrozenSet[str]: + """Parse a comma-separated speaker-role list, dropping unknown roles. + + An explicitly empty value (``NETRA_AUDIO_ROLES=``) is the documented way + to disable audio capture without affecting traces, so it resolves to an + empty set rather than the default. + + Args: + env_var: Name of the environment variable holding the role list. + + Returns: + The recognised roles, or the full default set when *env_var* is unset. + """ + raw = os.getenv(env_var) + if raw is None: + return AUDIO_ROLES + + requested = {part.strip().lower() for part in raw.split(",") if part.strip()} + unknown = requested - AUDIO_ROLES + if unknown: + logger.warning( + "netra.audio: %s contains unknown role(s) %s; recognised roles are %s", + env_var, + sorted(unknown), + sorted(AUDIO_ROLES), + ) + return frozenset(requested & AUDIO_ROLES) + + def audio_endpoint(self) -> Optional[str]: + """Resolve 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. + + 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. + """ + return self.audio_endpoint() is not None and bool(self.audio_roles) + 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..d7fa73f --- /dev/null +++ b/netra/instrumentation/livekit/__init__.py @@ -0,0 +1,208 @@ +"""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.processors import LiveKitSpanProcessor +from netra.instrumentation.livekit.provider_binding import bind_livekit_tracer +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" + +# Set on the provider once our processors are 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 +_wrappers_lock = threading.Lock() +_wrappers_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.""" + 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. + """ + cfg: Optional[Config] = kwargs.get("config") or get_active_config() + if cfg 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, cfg) + except Exception: + logger.exception("netra.livekit: could not register span processors; lk.* mapping is disabled") + + try: + _install_wrappers() + except Exception: + logger.exception("netra.livekit: could not install session hooks; netra.session_id will be missing") + + self._log_audio_decision(cfg) + + 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 _wrappers_installed + + try: + unwrap(_AGENT_SESSION_MODULE, "AgentSession.start") + unwrap(_AGENT_SESSION_MODULE, "AgentSession._aclose_impl") + except (AttributeError, ModuleNotFoundError): + logger.error("netra.livekit: failed to uninstrument AgentSession") + + with _wrappers_lock: + _wrappers_installed = False + + @staticmethod + def _register_processors(provider: Any, cfg: 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 ``processors.py`` for the invariant that makes it safe before + adding a third. + + Args: + provider: The tracer provider to register on. + cfg: The active Netra config. + """ + 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(LiveKitSpanProcessor()) + setattr(provider, _PROCESSORS_FLAG, True) + logger.debug("netra.livekit: registered LiveKitSpanProcessor") + + @staticmethod + def _log_audio_decision(cfg: 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: + cfg: The active Netra config. + """ + try: + if not cfg.audio_capture_enabled: + logger.info( + "netra.livekit: call audio capture is OFF (no authenticated audio endpoint " + "resolved, or NETRA_AUDIO_ROLES is empty). Traces are unaffected" + ) + return + + endpoint = cfg.audio_endpoint() or "" + host = endpoint.split("://")[-1].split("/")[0] + logger.info( + "netra.livekit: call audio capture is ON for role(s) %s, streaming to host %s", + ",".join(sorted(cfg.audio_roles)), + host, + ) + except Exception: + logger.debug("netra.livekit: could not log the audio capture decision", exc_info=True) + + +def _install_wrappers() -> None: + """Wrap ``AgentSession.start`` and ``AgentSession._aclose_impl``. + + Each target is wrapped in its own ``try``/``except`` so a signature change to + one leaves the other working. Guarded by a module flag because ``wrapt`` would + otherwise double-wrap on a repeat ``_instrument()`` call. + """ + global _wrappers_installed + + with _wrappers_lock: + if _wrappers_installed: + return + + wrapped_any = False + + try: + wrap_function_wrapper(_AGENT_SESSION_MODULE, "AgentSession.start", wrap_start) + wrapped_any = True + except Exception: + logger.exception( + "netra.livekit: could not wrap AgentSession.start; netra.session_id will be missing " + "from LiveKit spans" + ) + + try: + wrap_function_wrapper(_AGENT_SESSION_MODULE, "AgentSession._aclose_impl", wrap_aclose) + wrapped_any = True + except Exception: + logger.exception("netra.livekit: could not wrap AgentSession._aclose_impl; session teardown is disabled") + + _wrappers_installed = wrapped_any + + +__all__ = ["NetraLiveKitInstrumentor"] diff --git a/netra/instrumentation/livekit/processors.py b/netra/instrumentation/livekit/processors.py new file mode 100644 index 0000000..0d6d19c --- /dev/null +++ b/netra/instrumentation/livekit/processors.py @@ -0,0 +1,375 @@ +"""Span processors that normalise livekit-agents spans into Netra's conventions.""" + +import itertools +import logging +from typing import Any, Callable, Iterator, Mapping, Optional + +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 ( + AUDIO_TYPE_BY_SPAN_NAME, + CHAT_CTX_ATTRIBUTE, + EVENT_CHOICE, + EVENT_ROLE, + GEN_AI_COMPLETION_CONTENT, + GEN_AI_COMPLETION_ROLE, + GEN_AI_INSTRUCTIONS, + GEN_AI_PROMPT_CONTENT, + GEN_AI_PROMPT_ROLE, + GEN_AI_TRANSCRIPT_RESPONSE, + GEN_AI_TRANSCRIPT_ROLE, + INSTRUCTIONS_ATTRIBUTE, + LIVEKIT_SCOPE_NAME, + NETRA_AUDIO_TYPE, + NETRA_USAGE_SOURCE, + TRANSCRIPT_SPAN_NAMES, + USAGE_SOURCE_FRAMEWORK, + ConversationSide, + ConversationTarget, + content_of_choice_event, + content_of_event, + conversation_target_for, + is_absent, + is_usage_attribute, + is_zero_usage, + mapped_key_for, + messages_from_chat_ctx, + role_of_choice_event, + span_type_for, + transcript_role_for, +) + +logger = logging.getLogger(__name__) + +SetAttributeFunc = Callable[[str, Any], None] + +# Routes one ``lk.*`` conversation-content attribute into a span's conversation +# convention. Returns True when the key belongs to that convention — whether or not +# it carried a value — so the caller knows not to fall through to ``ATTRIBUTE_MAP``. +ContentWriter = Callable[[str, Any], bool] + +# The indexed key pair to write for each side of the conversation convention. +_KEYS_BY_SIDE = { + ConversationSide.PROMPT: (GEN_AI_PROMPT_ROLE, GEN_AI_PROMPT_CONTENT), + ConversationSide.COMPLETION: (GEN_AI_COMPLETION_ROLE, GEN_AI_COMPLETION_CONTENT), +} + + +def _as_text(value: Any) -> str: + """Render an attribute value as the string an OTel text attribute needs. + + Args: + value: The value LiveKit wrote. + + Returns: + The value unchanged if it is already a string, else its ``str()``. + """ + return value if isinstance(value, str) else str(value) + + +def _transcript_content_writer(span: Span) -> ContentWriter: + """Build the content writer for a span in ``TRANSCRIPT_SPAN_NAMES``. + + Emits ``gen_ai.instructions`` plus one ``gen_ai.transcript.{index}`` entry per + utterance, in LiveKit's write order. The index counter is per span and is + advanced only on a write that carries text, so the sequence has no holes. + + Args: + span: The LiveKit span being wrapped. + + Returns: + A writer that consumes the transcript-carrying ``lk.*`` keys. + """ + transcript_index = itertools.count() + + def _append(role: str, value: Any) -> None: + """Append one utterance to the span's transcript sequence. + + Writes through ``span.set_attribute`` — the outermost wrapper — so the + values still pass down through the rest of the processor chain. + """ + index = next(transcript_index) + span.set_attribute(GEN_AI_TRANSCRIPT_ROLE.format(index=index), role) + span.set_attribute(GEN_AI_TRANSCRIPT_RESPONSE.format(index=index), _as_text(value)) + + def write(key: str, value: Any) -> bool: + if key == CHAT_CTX_ATTRIBUTE: + for role, content in messages_from_chat_ctx(value): + _append(role, content) + return True + + if key == INSTRUCTIONS_ATTRIBUTE: + if not is_absent(value): + span.set_attribute(GEN_AI_INSTRUCTIONS, _as_text(value)) + return True + + utterance_role = transcript_role_for(key) + if utterance_role is None: + return False + if not is_absent(value): + _append(utterance_role, value) + return True + + return write + + +def _conversation_content_writer( + span: Span, + prompt_index: Iterator[int], + completion_index: Iterator[int], +) -> ContentWriter: + """Build the content writer for a span on the prompt/completion convention. + + Args: + span: The LiveKit span being wrapped. + prompt_index: Shared counter for the indexed prompt sequence. + completion_index: Shared counter for the indexed completion sequence. + + Returns: + A writer that consumes ``lk.chat_ctx`` and the ``CONVERSATION_MAP`` keys. + """ + + def _write_conversation(target: ConversationTarget, value: Any) -> None: + """Append one message to the span's indexed gen_ai sequence. + + Writes through ``span.set_attribute`` — the outermost wrapper — so the + values reach ``SpanIOProcessor``, which assembles them into + ``input``/``output``. + """ + role_key, content_key = _KEYS_BY_SIDE[target.side] + counter = prompt_index if target.side is ConversationSide.PROMPT else completion_index + index = next(counter) + span.set_attribute(role_key.format(index=index), target.role) + span.set_attribute(content_key.format(index=index), _as_text(value)) + + def _expand_chat_ctx(value: Any) -> None: + """Turn a serialised ChatContext into indexed prompt attributes.""" + for role, content in messages_from_chat_ctx(value): + index = next(prompt_index) + span.set_attribute(GEN_AI_PROMPT_ROLE.format(index=index), role) + span.set_attribute(GEN_AI_PROMPT_CONTENT.format(index=index), content) + + def write(key: str, value: Any) -> bool: + if key == CHAT_CTX_ATTRIBUTE: + _expand_chat_ctx(value) + return True + + conversation = conversation_target_for(key) + if conversation is None: + return False + if not is_absent(value): + _write_conversation(conversation, value) + return True + + return write + + +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 + + +class LiveKitSpanProcessor(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 takes one of two shapes depending on the span, chosen in + ``on_start``: the interaction-level turn spans in ``TRANSCRIPT_SPAN_NAMES`` emit + ``gen_ai.instructions`` plus an indexed ``gen_ai.transcript.*`` sequence, while + every other span — the LLM spans, which are real generations — keeps the + indexed ``gen_ai.prompt.*``/``gen_ai.completion.*`` pair that + ``SpanIOProcessor`` assembles into ``input``/``output``. + """ + + def on_start(self, span: Span, parent_context: Optional[otel_context.Context] = None) -> None: + """Install the attribute and event 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 + span.set_attribute("span_type", span_type_for(span.name)) + audio_type = AUDIO_TYPE_BY_SPAN_NAME.get(span.name) + if audio_type is not None: + span.set_attribute(NETRA_AUDIO_TYPE, audio_type) + if span.name in TRANSCRIPT_SPAN_NAMES: + self._install_transcript_mapping(span) + else: + self._install_conversation_mapping(span) + except Exception: + logger.warning("netra.livekit: span mapping could not be installed", exc_info=True) + + @classmethod + def _install_transcript_mapping(cls, span: Span) -> None: + """Map an interaction-level turn span onto the transcript convention. + + No event wrapper is installed: LiveKit emits its conversation events + (``gen_ai.system.message`` and friends, ``gen_ai.choice``) only on + ``llm_request``, never on these spans, and mapping them would reintroduce + the prompt/completion keys the transcript convention replaces here. + + Args: + span: The LiveKit span to wrap. + """ + cls._wrap_set_attribute(span, _transcript_content_writer(span)) + + @classmethod + def _install_conversation_mapping(cls, span: Span) -> None: + """Map every other LiveKit span onto the prompt/completion convention. + + Args: + span: The LiveKit span to wrap. + """ + # One counter per span, shared by both wrappers: chat-context expansion, + # conversation-content attributes (``lk.response.text``) and conversation + # events all write into the same indexed sequences, and separate counters + # would overwrite each other's entries on a span that has more than one of + # them. Advanced only on a *mapped* write, so a stray attribute or event + # cannot leave a hole in the sequence. + prompt_index = itertools.count() + completion_index = itertools.count() + cls._wrap_set_attribute(span, _conversation_content_writer(span, prompt_index, completion_index)) + cls._wrap_add_event(span, prompt_index, completion_index) + + def on_end(self, span: ReadableSpan) -> None: + """No-op. See the module docstring: end-time mutation would race the exporter.""" + + 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 _wrap_set_attribute(span: Span, write_content: ContentWriter) -> 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. + write_content: The conversation convention this span emits into. + """ + previous: SetAttributeFunc = span.set_attribute + + def patched_set_attribute(key: str, value: Any) -> None: + try: + 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 write_content(key, value): + return + + target = mapped_key_for(key) + if target is None or is_absent(value): + return + previous(target, 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: + if not attributes: + return + for key, value in attributes.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, prompt_index: Iterator[int], completion_index: Iterator[int]) -> 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. + prompt_index: Shared counter for the indexed prompt sequence. + completion_index: Counter for the indexed completion sequence. + """ + original = span.add_event + + def _map_choice(attributes: Attributes) -> None: + content = content_of_choice_event(attributes) + if not content: + return + index = next(completion_index) + span.set_attribute(GEN_AI_COMPLETION_ROLE.format(index=index), role_of_choice_event(attributes)) + span.set_attribute(GEN_AI_COMPLETION_CONTENT.format(index=index), content) + + def _map_prompt(role: str, attributes: Attributes) -> None: + content = content_of_event(attributes) + if not content: + return + index = next(prompt_index) + span.set_attribute(GEN_AI_PROMPT_ROLE.format(index=index), role) + span.set_attribute(GEN_AI_PROMPT_CONTENT.format(index=index), content) + + def patched_add_event( + name: str, + attributes: Attributes = None, + timestamp: Optional[int] = None, + ) -> None: + try: + if name == EVENT_CHOICE: + _map_choice(attributes) + else: + role = EVENT_ROLE.get(name) + if role is not None: + _map_prompt(role, 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/provider_binding.py b/netra/instrumentation/livekit/provider_binding.py new file mode 100644 index 0000000..002a073 --- /dev/null +++ b/netra/instrumentation/livekit/provider_binding.py @@ -0,0 +1,138 @@ +"""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.""" + 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. + """ + 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.""" + 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/utils.py b/netra/instrumentation/livekit/utils.py new file mode 100644 index 0000000..d1ab9a5 --- /dev/null +++ b/netra/instrumentation/livekit/utils.py @@ -0,0 +1,472 @@ +"""Pure mapping tables and helpers for the LiveKit span processor. + +Deliberately free of OTel imports so the mapping rules can be unit-tested as +plain data. 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). +""" + +import json +from enum import Enum +from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Tuple + +# 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 keys. +NETRA_TOOL_NAME = "netra.tool.name" +NETRA_USAGE_SOURCE = "netra.usage.source" +USAGE_SOURCE_FRAMEWORK = "framework" + +# 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" + +# 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. +# +# ``job_entrypoint`` is included by request. It is absent from the span tree +# verified against livekit-agents 1.6.7, so it currently matches nothing; the +# entry is inert rather than wrong, and becomes live if a version emits that span. +AUDIO_TYPE_BY_SPAN_NAME: Dict[str, str] = { + "agent_session": AUDIO_TYPE_SESSION, + "agent_turn": AUDIO_TYPE_SPAN, + "user_turn": AUDIO_TYPE_SPAN, +} + +# The interaction-level spans whose conversation content is emitted as the +# transcript convention (``GEN_AI_INSTRUCTIONS`` plus the indexed +# ``gen_ai.transcript.*`` sequence) instead of the indexed +# ``gen_ai.prompt``/``gen_ai.completion`` pair. +# +# Gated on the span *name* rather than on the attribute key, because two content +# attributes appear on both families: ``lk.response.text`` is set on ``agent_turn`` +# and on ``llm_node``, and ``lk.chat_ctx`` on ``eou_detection`` and on +# ``llm_node``. The LLM spans (``llm_node``, ``llm_request``) deliberately stay on +# the prompt/completion convention — they are real generations, and that is the +# convention every other Netra provider instrumentation emits for one. +TRANSCRIPT_SPAN_NAMES = frozenset({"agent_turn", "tts_request", "user_turn", "eou_detection"}) + +# 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" + +# The transcript convention, used on the ``TRANSCRIPT_SPAN_NAMES`` spans in place +# of the prompt/completion pair. One ordered sequence per span — who spoke and +# what they said, in LiveKit's write order — rather than two sequences split by +# request/response side, because on a voice turn the split is arbitrary: the +# caller's utterance and the agent's reply are both just turns of one dialogue. +# +# ``SpanIOProcessor`` recognises neither key, so it builds no ``input``/``output`` +# from them and both stay empty on these spans. That is deliberate: the transcript +# is read from these attributes directly. +GEN_AI_INSTRUCTIONS = "gen_ai.instructions" +GEN_AI_TRANSCRIPT_ROLE = "gen_ai.transcript.{index}.role" +GEN_AI_TRANSCRIPT_RESPONSE = "gen_ai.transcript.{index}.response" + +# Prefix identifying token-usage attributes, whoever wrote them. +GEN_AI_USAGE_PREFIX = "gen_ai.usage." + +# 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", + "lk.function_tool.output": "output", + # 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", +} + + +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 + + +# lk.* content attribute -> gen_ai slot, for spans *not* in +# ``TRANSCRIPT_SPAN_NAMES`` — i.e. ``llm_node``. 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] = { + # llm_node: the generated reply. (``lk.response.text`` is also set on + # ``agent_turn``, which takes the transcript route below instead.) + "lk.response.text": ConversationTarget(ConversationSide.COMPLETION, "assistant"), +} + +# lk.* content attribute -> transcript role, for the spans in +# ``TRANSCRIPT_SPAN_NAMES``. Each entry appends one entry to the span's +# ``gen_ai.transcript.*`` sequence; the role names the *speaker*, so the sequence +# reads as the actual dialogue. +# +# ``lk.instructions`` is deliberately absent: it is not a turn of the dialogue, and +# goes to the flat ``GEN_AI_INSTRUCTIONS`` key instead. +TRANSCRIPT_ROLE_MAP: Dict[str, str] = { + # agent_turn: the utterance that opened the turn (LiveKit sets this only when a + # new message did), then the agent's reply. + "lk.user_input": "user", + "lk.response.text": "assistant", + # tts_request: the text handed to the TTS provider. The words are the agent's. + "lk.input_text": "assistant", + # user_turn: the STT transcript. The words are the caller's. + "lk.user_transcript": "user", +} + +# agent_turn: the system instructions in force for the turn. +INSTRUCTIONS_ATTRIBUTE = "lk.instructions" + +# span name -> span_type. Anything not listed here is a plain "span"; that +# includes user_turn, eou_detection and the *_request_run retry spans, which +# have no distinct type in Netra's vocabulary. +SPAN_TYPE_BY_NAME: Dict[str, str] = { + "agent_turn": "agent", + "llm_node": "llm", + "llm_request": "llm", + "tts_node": "tts", + "tts_request": "tts", + "function_tool": "tool", +} + +DEFAULT_SPAN_TYPE = "span" + +# 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", +} + +# 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" + +# 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" + +# 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" +_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" + + +def span_type_for(span_name: Optional[str]) -> str: + """Return the Netra ``span_type`` for a LiveKit span name. + + Derived from the span *name* rather than an attribute, because LiveKit sets + no attribute that identifies the pipeline stage. + + Args: + span_name: The LiveKit span's name, or ``None``. + + Returns: + One of ``agent``, ``llm``, ``tts``, ``tool``, or ``span``. + """ + if not span_name: + return DEFAULT_SPAN_TYPE + return SPAN_TYPE_BY_NAME.get(span_name, DEFAULT_SPAN_TYPE) + + +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 mapped_key_for(lk_key: str) -> Optional[str]: + """Return the Netra target key for *lk_key*, or None if it is not mapped. + + Args: + lk_key: A LiveKit ``lk.*`` attribute name. + + Returns: + The Netra attribute name to mirror the value into, or ``None``. + """ + return ATTRIBUTE_MAP.get(lk_key) + + +def conversation_target_for(lk_key: str) -> Optional[ConversationTarget]: + """Return the gen_ai slot *lk_key*'s text belongs in, or None if it carries none. + + Args: + lk_key: A LiveKit ``lk.*`` attribute name. + + Returns: + The ``ConversationTarget`` for a conversation-content attribute, or + ``None`` for every other attribute. + """ + return CONVERSATION_MAP.get(lk_key) + + +def transcript_role_for(lk_key: str) -> Optional[str]: + """Return the transcript role for *lk_key*, or None if it carries no utterance. + + Only meaningful for a span in ``TRANSCRIPT_SPAN_NAMES``; the caller applies + that gate. + + Args: + lk_key: A LiveKit ``lk.*`` attribute name. + + Returns: + The role to stamp on the transcript entry, or ``None`` for every attribute + that is not an utterance. + """ + return TRANSCRIPT_ROLE_MAP.get(lk_key) + + +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 content if isinstance(content, str) else str(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 content if isinstance(content, str) else str(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) + "]" + + +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. + """ + 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) + if not isinstance(items, list): + return [] + + 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 _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) + + +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 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..7b5950f --- /dev/null +++ b/netra/instrumentation/livekit/wrappers.py @@ -0,0 +1,288 @@ +"""wrapt wrappers for LiveKit's ``AgentSession`` lifecycle. + +The session id — the LiveKit room SID, falling back to the room name — is attached +as OTel baggage *around* ``AgentSession.start`` so that the ``agent_session`` root +span, created inside ``start()``, carries it, then detached so the caller's context +is restored. See ``wrap_start`` and ``_resolve_session_id``. + +Nothing in here may change the behaviour of the user's application: every hook +calls the wrapped function even if our own logic raises, and exceptions raised by +the user's code propagate untouched. +""" + +import logging +from contextlib import ExitStack +from typing import Any, Awaitable, Callable, Dict, Optional, Tuple + +from netra.session_manager import SessionManager + +logger = logging.getLogger(__name__) + +# The wrapt quadruple: (wrapped, instance, args, kwargs). ``instance`` is a +# livekit AgentSession, which is not importable at module scope — the SDK must +# stay importable with livekit-agents absent. +WrappedAsync = Callable[..., Awaitable[Any]] + + +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 + + +def _session_span(instance: Any) -> Optional[Any]: + """Return the live ``agent_session`` span, or ``None`` once it is gone. + + Args: + instance: The LiveKit ``AgentSession``. + + Returns: + The span object LiveKit holds for the session, if it still has one. + """ + 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. + + This is the per-session state key used across the integration: it is + reachable both here and from any span processor via ``span.context.trace_id``. + + Args: + session_span: The session span, or ``None``. + + Returns: + The trace id as an int, or ``None`` if the session span is unavailable. + """ + 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) + + +def _after_start(instance: Any, session_id: Optional[str]) -> None: + """Per-session wiring that runs once ``start()`` has returned. + + Phase 1 only records the session at DEBUG. From Phase 4 this is where the + audio sender is registered, which is why it must tolerate being called for a + session it did not start: ``start()`` returns early on an already-started + session, and a restarted session arrives here with a *new* trace id. + + Args: + instance: The LiveKit ``AgentSession``. + session_id: The resolved Netra session id, 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, + ) + + +def _before_close(instance: Any) -> None: + """Per-session teardown that runs *before* LiveKit closes the session. + + Must run before the wrapped call: ``_aclose_impl`` ends ``_session_span`` and + only then emits ``close``, after which the span is ``None``, its trace id — + our state key — is unreachable, and it can no longer be written to. + + ``_aclose_impl`` can be reached more than once for one session, so this must + be idempotent. Reading the trace id off ``_session_span`` gives that for + free: the second call sees ``None`` and does nothing. + + Args: + instance: The LiveKit ``AgentSession``. + """ + 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) + + +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``. + 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) + + try: + _after_start(instance, session_id) + except Exception as exc: + logger.warning("netra.livekit: post-start wiring failed: %s", exc) + + 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: + _before_close(instance) + except Exception as exc: + logger.warning("netra.livekit: pre-close teardown failed: %s", exc) + + return await wrapped(*args, **kwargs) diff --git a/netra/meter.py b/netra/meter.py index 50bf1e5..b2fad58 100644 --- a/netra/meter.py +++ b/netra/meter.py @@ -112,6 +112,42 @@ 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: + 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.""" + 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()``.""" + self._netra_shutdown_allowed = True + super().shutdown(timeout_millis=timeout_millis) + + class MetricsSetup: """ Configures Netra's OpenTelemetry metrics pipeline. @@ -194,7 +230,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 af628bd..824df40 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", ]