diff --git a/nemoguardrails/__init__.py b/nemoguardrails/__init__.py index 5869f28c0a..a99b86dbab 100644 --- a/nemoguardrails/__init__.py +++ b/nemoguardrails/__init__.py @@ -63,6 +63,7 @@ register_framework, set_default_framework, ) +from nemoguardrails.llm.models.instrumented import InstrumentedLLMModel # noqa: E402 from nemoguardrails.llm.providers import register_provider # noqa: E402 from nemoguardrails.types import ( # noqa: E402 ChatMessage, @@ -82,6 +83,7 @@ "ChatMessage", "FinishReason", "Guardrails", + "InstrumentedLLMModel", "LLMFramework", "LLMModel", "LLMRails", diff --git a/nemoguardrails/guardrails/engine_registry.py b/nemoguardrails/guardrails/engine_registry.py index 1b7ff6f242..4e6186bedf 100644 --- a/nemoguardrails/guardrails/engine_registry.py +++ b/nemoguardrails/guardrails/engine_registry.py @@ -20,31 +20,19 @@ """ import logging -import time from collections.abc import AsyncGenerator -from contextlib import nullcontext -from typing import TYPE_CHECKING, Any, Optional, TypeVar +from contextlib import aclosing +from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast from nemoguardrails.guardrails.api_engine import APIEngine from nemoguardrails.guardrails.base_engine import BaseEngine from nemoguardrails.guardrails.guardrails_types import get_request_id, truncate from nemoguardrails.guardrails.model_engine import ModelEngine -from nemoguardrails.guardrails.telemetry import ( - api_call_span, - llm_call_span, - set_llm_call_content, - set_llm_request_attributes, - set_llm_response_attributes, -) +from nemoguardrails.guardrails.telemetry import api_call_span from nemoguardrails.guardrails.tool_schema import ToolExchange, ToolResult, Toolset +from nemoguardrails.llm.models.instrumented import instrument_llm_model from nemoguardrails.rails.llm.config import Model, RailsConfigData -from nemoguardrails.tracing.constants import ( - llm_operation_duration, - record_time_per_output_chunk, - record_time_to_first_chunk, - record_token_usage, -) -from nemoguardrails.types import LLMResponse, LLMResponseChunk, UsageInfo +from nemoguardrails.types import ChatMessage, LLMModel, LLMResponse, LLMResponseChunk if TYPE_CHECKING: from opentelemetry.trace import Tracer @@ -54,6 +42,58 @@ _EngineT = TypeVar("_EngineT", bound=BaseEngine) +class _ModelEngineAdapter: + """Adapt an IORails ``ModelEngine`` to the canonical ``LLMModel`` protocol.""" + + def __init__(self, engine: ModelEngine) -> None: + self._engine = engine + + @property + def model_name(self) -> str: + return self._engine.model_name + + @property + def provider_name(self) -> Optional[str]: + return self._engine.model_config.engine + + @property + def provider_url(self) -> Optional[str]: + return self._engine.base_url + + @staticmethod + def _messages(prompt: str | list[ChatMessage] | list[dict]) -> list[dict]: + if isinstance(prompt, str): + return [{"role": "user", "content": prompt}] + return [message.to_dict() if isinstance(message, ChatMessage) else message for message in prompt] + + def _params(self, stop: Optional[list[str]], kwargs: dict[str, Any]) -> dict[str, Any]: + params = {**self._engine.body_param_defaults, **kwargs} + if stop is not None: + params["stop"] = stop + return params + + async def generate_async( + self, + prompt: str | list[ChatMessage], + *, + stop: Optional[list[str]] = None, + **kwargs: Any, + ) -> LLMResponse: + return await self._engine.chat_completion(self._messages(prompt), **self._params(stop, kwargs)) + + async def stream_async( + self, + prompt: str | list[ChatMessage], + *, + stop: Optional[list[str]] = None, + **kwargs: Any, + ) -> AsyncGenerator[LLMResponseChunk, None]: + stream = self._engine.stream_chat_completion(self._messages(prompt), **self._params(stop, kwargs)) + async with aclosing(stream): + async for chunk in stream: + yield chunk + + class EngineRegistry: """Registry of ModelEngine and APIEngine instances for IORails. @@ -88,14 +128,20 @@ def __init__( work. """ self._engines: dict[str, BaseEngine] = {} + self._models: dict[str, LLMModel] = {} self._running = False self._tracer = tracer - self._metrics_enabled = metrics_enabled - self._content_capture_enabled = content_capture_enabled for model_config in models: engine = ModelEngine(model_config) self._engines[model_config.type] = engine + self._models[model_config.type] = instrument_llm_model( + _ModelEngineAdapter(engine), + tracer=tracer, + metrics_enabled=metrics_enabled, + content_capture_enabled=content_capture_enabled, + default_request_params=engine.body_param_defaults, + ) log.info( "Registered model engine: type=%s, model=%s, base_url=%s", model_config.type, @@ -195,47 +241,9 @@ async def model_call(self, model_type: str, messages: list[dict], **kwargs: Any) req_id = get_request_id() log.debug("[%s] Model engine '%s' messages: %s", req_id, model_type, truncate(messages)) - engine = self._get_engine(model_type, ModelEngine) - # TODO: Replace with LLMModel.provider_name after refactoring - provider_name = engine.model_config.engine or "unknown" - operation_name = "chat" - - # Merge the model's config parameters with per-call kwargs (GenerationOptions.llm_params) - # Per-call kwargs have priority. - merged_params = {**engine.body_param_defaults, **kwargs} - - # Compose: span (always created — no-op when tracer is None) and - # duration metric (only when metrics enabled). Token usage is - # emitted after the call returns since it depends on - # ``result.usage`` — exception path skips it because control - # never reaches the line below. - duration_ctx = ( - llm_operation_duration(engine.model_name, provider_name, operation_name) - if self._metrics_enabled - else nullcontext() - ) - with llm_call_span(self._tracer, engine.model_name, provider_name, operation_name) as span: - # Request params are known before the call, so set them first — - # they land on the span even if the call raises. - set_llm_request_attributes(span, merged_params) - with duration_ctx: - result = await engine.chat_completion(messages, **merged_params) - # Set response/usage and content attrs inside the span context so - # the helpers see the live LLM CLIENT span and the attributes land - # before it closes. Both are skipped on exception, which never - # reaches here. - set_llm_response_attributes( - span, - model=result.model, - response_id=result.request_id, - finish_reason=result.finish_reason, - usage=result.usage, - ) - if self._content_capture_enabled: - set_llm_call_content(span, messages, result.content) - - if self._metrics_enabled: - record_token_usage(engine.model_name, provider_name, operation_name, result.usage) + self._get_engine(model_type, ModelEngine) + prompt = cast(list[ChatMessage], messages) + result = await self._models[model_type].generate_async(prompt, **kwargs) log.debug("[%s] Model engine '%s' response: %s", req_id, model_type, truncate(result)) return result @@ -274,106 +282,12 @@ async def stream_model_call( req_id = get_request_id() log.debug("[%s] Model engine '%s' stream messages: %s", req_id, model_type, truncate(messages)) - engine = self._get_engine(model_type, ModelEngine) - # TODO: Change to LLMModel.provider_name after refactor - provider_name = engine.model_config.engine or "unknown" - operation_name = "chat" - - # Merge the model's configured parameter defaults with the per-call - # kwargs (per-call wins), above set_llm_request_attributes, so the span - # reflects the request body. Excluding "stream"/"stream_options" from - # body_param_defaults also prevents a duplicate-keyword TypeError when - # stream_call() passes its own stream=True into _prepare_request(). - merged_params = {**engine.body_param_defaults, **kwargs} - - # Capture the latest non-None response fields from the stream so we - # can set the LLM span's response/usage attrs and emit the token - # metric after the stream completes. Providers spread these across - # the SSE chunks — OpenAI-compatible engines only populate ``usage`` - # on the terminal chunk (when ``stream_options.include_usage=true``) - # and finish_reason likewise arrives last — so each field keeps its - # latest non-None value. - captured_usage: Optional["UsageInfo"] = None - captured_model: Optional[str] = None - captured_response_id: Optional[str] = None - captured_finish_reason: Optional[str] = None - # Accumulate streamed delta_content here when content capture is on; - # joined and recorded onto the LLM span at stream end. The list is - # allocated unconditionally (cost: one empty list per stream); the - # per-chunk appends are gated on the flag so the disabled path - # doesn't carry chunk strings in memory. - content_parts: list[str] = [] - duration_ctx = ( - llm_operation_duration(engine.model_name, provider_name, operation_name) - if self._metrics_enabled - else nullcontext() - ) - with llm_call_span(self._tracer, engine.model_name, provider_name, operation_name) as span: - # Set request params + stream=True before the first chunk so they - # land on the span even if the stream errors mid-flight. - set_llm_request_attributes(span, merged_params, stream=True) - with duration_ctx: - # Gate timing-state setup on ``_metrics_enabled`` so the - # cold path skips ``time.monotonic()`` and the per-chunk - # bookkeeping entirely. ``t0`` defaults to ``0.0`` in - # the disabled path so the type stays a plain ``float`` - # — it's never read in that branch. - t0 = time.monotonic() if self._metrics_enabled else 0.0 - last_chunk_time: Optional[float] = None - async for chunk in engine.stream_chat_completion(messages, **merged_params): - if self._metrics_enabled: - # Per OTEL semconv, "first chunk" / "output chunk" - # mean content-bearing chunks — gate on - # ``delta_content`` / ``delta_reasoning`` to skip - # the terminal usage frame and any other cosmetic - # SSE events that the parser leaves in place. - if chunk.delta_content or chunk.delta_reasoning: - now = time.monotonic() - if last_chunk_time is None: - record_time_to_first_chunk(engine.model_name, provider_name, operation_name, now - t0) - else: - record_time_per_output_chunk( - engine.model_name, provider_name, operation_name, now - last_chunk_time - ) - last_chunk_time = now - # Keep the latest non-None response field from each chunk. - # Captured regardless of ``_metrics_enabled`` because they - # feed the span's response/usage attrs (set after the loop) - # as well as the token-usage metric. - if chunk.model is not None: - captured_model = chunk.model - if chunk.request_id is not None: - captured_response_id = chunk.request_id - if chunk.finish_reason is not None: - captured_finish_reason = chunk.finish_reason - if chunk.usage is not None: - captured_usage = chunk.usage - if self._content_capture_enabled and chunk.delta_content: - content_parts.append(chunk.delta_content) - yield chunk - # Set response/usage attrs and (when enabled) content inside the - # span context so the helpers see the live LLM CLIENT span before - # it closes. Reached only on natural exhaustion — consumer - # cancellation or provider error raises out of the ``with`` blocks - # above, so partial response data is intentionally not recorded. - set_llm_response_attributes( - span, - model=captured_model, - response_id=captured_response_id, - finish_reason=captured_finish_reason, - usage=captured_usage, - ) - # Empty ``content_parts`` -> output_text=None so we don't claim - # an empty assistant response (matches iorails.py's request-span - # streaming path). - if self._content_capture_enabled: - output_text = "".join(content_parts) if content_parts else None - set_llm_call_content(span, messages, output_text) - - # Reached only on natural exhaustion (not on consumer cancellation - # or provider error — those raise out of the ``with`` blocks above). - if self._metrics_enabled: - record_token_usage(engine.model_name, provider_name, operation_name, captured_usage) + self._get_engine(model_type, ModelEngine) + prompt = cast(list[ChatMessage], messages) + stream = cast(AsyncGenerator[LLMResponseChunk, None], self._models[model_type].stream_async(prompt, **kwargs)) + async with aclosing(stream): + async for chunk in stream: + yield chunk def parse_tools(self, model_type: str, llm_params: Optional[dict]) -> Toolset: """Parse the tool block in ``llm_params`` for the named model engine. diff --git a/nemoguardrails/guardrails/telemetry.py b/nemoguardrails/guardrails/telemetry.py index 16b22b422d..ad3a3c68b1 100644 --- a/nemoguardrails/guardrails/telemetry.py +++ b/nemoguardrails/guardrails/telemetry.py @@ -13,16 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Inline OpenTelemetry instrumentation for the IORails engine. - -All OpenTelemetry API imports are isolated in this module so the rest of the -guardrails package never imports ``opentelemetry`` directly. When the -``opentelemetry-api`` package is not installed, the public entry points -``is_tracing_enabled``, ``get_tracer``, ``get_meter``, and ``traced_request`` -degrade gracefully (returning ``False``, ``None``, or a no-span / no-metric -passthrough respectively). Lower-level helpers like ``request_span`` and -``trace_id_to_request_id`` require OTEL to be available and are only -reachable through ``traced_request`` when a non-``None`` tracer is provided. +"""OpenTelemetry instrumentation for Guardrails operations. + +Reusable model-call telemetry lives in ``nemoguardrails.llm.telemetry``. +When the ``opentelemetry-api`` package is not installed, the public entry +points in this module degrade gracefully to disabled telemetry. """ import json @@ -54,8 +49,8 @@ reset_request_id, set_new_request_id, ) +from nemoguardrails.llm.telemetry import record_span_error from nemoguardrails.tracing.constants import ( - EventNames, GenAIAttributes, GuardrailsAttributes, MetricNames, @@ -79,11 +74,10 @@ Observation, UpDownCounter, ) - from opentelemetry.trace import Span, SpanKind, StatusCode, Tracer, format_trace_id + from opentelemetry.trace import Span, SpanKind, Tracer, format_trace_id from nemoguardrails.guardrails.async_work_queue import AsyncWorkQueue from nemoguardrails.rails.llm.config import MetricsConfig, TracingConfig - from nemoguardrails.types import UsageInfo _OTEL_AVAILABLE = True else: @@ -91,12 +85,13 @@ from opentelemetry import metrics as otel_metrics from opentelemetry import trace from opentelemetry.metrics import CallbackOptions, Observation - from opentelemetry.trace import SpanKind, StatusCode, format_trace_id + from opentelemetry.trace import SpanKind, format_trace_id _OTEL_AVAILABLE = True except ImportError: # pragma: no cover _OTEL_AVAILABLE = False + # Module-level tracer singleton. Thread-safe: the OTEL spec requires that # ``Tracer`` methods are safe for concurrent use, and ``get_tracer()`` is # designed to be called once and cached (see "Get a Tracer" at @@ -342,29 +337,6 @@ def trace_id_to_request_id(span: "Span") -> str: return format_trace_id(ctx.trace_id)[-REQUEST_ID_HEX_CHARS:] -def record_span_error(span: Optional["Span"], exc: BaseException) -> None: - """Record an exception on an OTEL span and set its status to ERROR. - - Also sets the ``error.type`` attribute to the exception's class name - (per OTEL GenAI conditional-required convention). Safe to call with - ``None`` (no-op). Use from every span helper's ``except`` block and - from callers that swallow exceptions before they can propagate. - - Best-effort: any failure while annotating the span (e.g. a broken - exporter or SDK) is swallowed so it can never mask the original - exception the caller is about to re-raise — notably ``CancelledError`` - / ``GeneratorExit`` on a cancelled stream. Only ``Exception`` is - suppressed, so a ``BaseException`` raised *inside* the SDK still - propagates. - """ - if span is None: - return - with suppress(Exception): - span.set_attribute("error.type", type(exc).__name__) - span.record_exception(exc) - span.set_status(StatusCode.ERROR, str(exc)) - - def mark_rail_stop(span: Optional["Span"], is_safe: bool) -> None: """Set ``rail.stop=True`` on a rail span when the rail blocked the request. @@ -401,301 +373,6 @@ def set_speculative_span_attrs( # span.set_attribute(GuardrailsAttributes.SPECULATIVE_TIME_SAVED_MS, time_saved_ms) -# Maps an OpenAI-style ``role`` to the OTEL GenAI legacy event name used -# when content capture emits per-message span events (i.e. when the -# stability opt-in does NOT select the new structured attribute form). -# ``function`` role (OpenAI legacy function-call format) is deliberately -# excluded — it will need its own decision when function-call support lands. -_LEGACY_EVENT_BY_ROLE = { - "system": EventNames.GEN_AI_SYSTEM_MESSAGE, - "user": EventNames.GEN_AI_USER_MESSAGE, - "assistant": EventNames.GEN_AI_ASSISTANT_MESSAGE, - "tool": EventNames.GEN_AI_TOOL_MESSAGE, -} - - -def _use_json_span_format() -> bool: - """Return True iff OTEL_SEMCONV_STABILITY_OPT_IN selects JSON span attrs. - - The env var holds a comma-separated list of opt-in tokens. When - ``gen_ai_latest_experimental`` is present, content is emitted as - JSON-encoded span attributes, otherwise as legacy per-message span events. - Read fresh each call so runtime changes to the env var take effect - immediately. - """ - raw_env_value = os.environ.get(OtelContentCapture.STABILITY_OPT_IN_ENV, "") - tokens = {tok.strip() for tok in raw_env_value.split(",")} - return OtelContentCapture.STABILITY_OPT_IN_LATEST in tokens - - -def _system_parts_from_messages(messages: LLMMessages) -> list[dict]: - """Return the bare OTEL GenAI ``parts`` for system messages only. - - Feeds ``gen_ai.system_instructions``, which the spec defines as a flat - list of parts with no role wrapper (every entry is implicitly system). - Asymmetric with :func:`_non_system_input_messages`, which keeps the role - wrapper — the two attributes have different shapes by spec. Entries - missing ``role`` or ``content`` are skipped silently. - - Example:: - - >>> _system_parts_from_messages([ - ... {"role": "system", "content": "be helpful"}, - ... {"role": "user", "content": "hi"}, - ... ]) - [{"type": "text", "content": "be helpful"}] - """ - out: list[dict] = [] - for msg in messages: - role = msg.get("role") - content = msg.get("content") - if role is None or content is None: - continue - if role == "system": - out.append({"type": "text", "content": content}) - return out - - -def _non_system_input_messages(messages: LLMMessages) -> list[dict]: - """Return the OTEL GenAI ``gen_ai.input.messages`` form for non-system messages. - - Each non-system message is role-wrapped as ``{"role": role, "parts": - [{"type": "text", "content": content}]}``. Named for the attribute it - populates rather than "parts" because — unlike - :func:`_system_parts_from_messages` — it keeps the role wrapper. - - Example:: - - >>> _non_system_input_messages([ - ... {"role": "system", "content": "be helpful"}, - ... {"role": "user", "content": "hi"}, - ... ]) - [{"role": "user", "parts": [{"type": "text", "content": "hi"}]}] - """ - out: list[dict] = [] - for msg in messages: - role = msg.get("role") - content = msg.get("content") - if role is None or content is None: - continue - if role != "system": - out.append({"role": role, "parts": [{"type": "text", "content": content}]}) - return out - - -def _set_llm_call_content_json( - span: "Span", - input_messages: LLMMessages, - output_text: Optional[str], -) -> None: - """JSON-attribute branch of :func:`set_llm_call_content`. - - Sets ``gen_ai.input.messages``, ``gen_ai.output.messages``, and - ``gen_ai.system_instructions`` as JSON-encoded span attributes per - the latest experimental OTEL GenAI semantic conventions. Attributes - are only set when non-empty so backends can distinguish "no system - instructions" from "system instructions == ''". - """ - system_parts = _system_parts_from_messages(input_messages) - if system_parts: - span.set_attribute(GenAIAttributes.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps(system_parts)) - - non_system = _non_system_input_messages(input_messages) - if non_system: - span.set_attribute(GenAIAttributes.GEN_AI_INPUT_MESSAGES, json.dumps(non_system)) - - if output_text is not None: - output_messages = [{"role": "assistant", "parts": [{"type": "text", "content": output_text}]}] - span.set_attribute(GenAIAttributes.GEN_AI_OUTPUT_MESSAGES, json.dumps(output_messages)) - - -def _set_llm_call_content_events( - span: "Span", - input_messages: LLMMessages, - output_text: Optional[str], -) -> None: - """Legacy-event branch of :func:`set_llm_call_content`. - - Adds one span event per input message (``gen_ai.system.message`` / - ``gen_ai.user.message`` / ``gen_ai.assistant.message`` / - ``gen_ai.tool.message``) plus a ``gen_ai.choice`` event for the - assistant output. Roles not in :data:`_LEGACY_EVENT_BY_ROLE` - (e.g. ``function``) are skipped silently. - """ - for msg in input_messages: - role = msg.get("role") - content = msg.get("content") - if role is None or content is None: - continue - event_name = _LEGACY_EVENT_BY_ROLE.get(role) - if event_name is None: - continue - span.add_event(event_name, attributes={"role": role, "content": content}) - - if output_text is not None: - span.add_event( - EventNames.GEN_AI_CHOICE, - attributes={"index": 0, "message.role": "assistant", "message.content": output_text}, - ) - - -def set_llm_call_content( - span: Optional["Span"], - input_messages: LLMMessages, - output_text: Optional[str] = None, -) -> None: - """Capture input/output messages on a span representing a model interaction. - - Used for both ``gen_ai.*`` CLIENT spans (LLM calls) and the - ``guardrails.request`` SERVER span — the OTEL GenAI semconv - attribute names apply to any span that represents a model - interaction, so reusing the names lets backends correlate the outer - guardrails request with the inner LLM call by attribute name alone. - - Dispatches on :func:`_use_json_span_format`: - - * **JSON attrs** (``OTEL_SEMCONV_STABILITY_OPT_IN`` includes - ``gen_ai_latest_experimental``): :func:`_set_llm_call_content_json` - sets the JSON-encoded ``gen_ai.input.messages``, - ``gen_ai.output.messages``, and ``gen_ai.system_instructions`` - span attributes per the latest experimental OTEL GenAI semantic - conventions. - * **Legacy events** (default): :func:`_set_llm_call_content_events` - adds one span event per input message plus a ``gen_ai.choice`` - event for the assistant output. - - Safe to call with ``span=None`` (no-op) so callers don't have to - branch on whether tracing is enabled. Caller is responsible for - checking the content-capture flag — this helper does NOT re-check - :func:`is_content_capture_enabled` so it stays cheap on hot paths. - """ - if span is None: - return - if _use_json_span_format(): - _set_llm_call_content_json(span, input_messages, output_text) - else: - _set_llm_call_content_events(span, input_messages, output_text) - - -# Maps an LLM request kwarg (as forwarded into the provider request body) -# to the OTEL GenAI span attribute that records it. Both ``max_tokens`` -# and the OpenAI ``max_completion_tokens`` alias map to the same attribute. -# ``stop`` / ``stop_sequences`` is handled separately by -# :func:`_stop_sequences` because it needs list normalization. -_GENAI_REQUEST_PARAMS = { - "temperature": GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, - "max_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS, - "max_completion_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS, - "top_p": GenAIAttributes.GEN_AI_REQUEST_TOP_P, - "top_k": GenAIAttributes.GEN_AI_REQUEST_TOP_K, - "frequency_penalty": GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY, - "presence_penalty": GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY, -} - - -def _stop_sequences(params: dict) -> Optional[list]: - """Return the request's stop sequences as a list, or ``None`` if unset. - - Reads the provider-specific request field that carries them: - - * OpenAI: ``stop`` - * Anthropic: ``stop_sequences`` - - A bare string is wrapped into a single-element list - (``gen_ai.request.stop_sequences`` is a string[]); a non-empty list is - returned unchanged. An empty or missing value, or any other type, - yields ``None`` — an empty ``stop`` is skipped rather than recorded as - a misleading empty span attribute. - """ - raw = params.get("stop") - if raw is None: - raw = params.get("stop_sequences") - if not raw: - return None - if isinstance(raw, str): - return [raw] - if isinstance(raw, list): - return raw - return None - - -def set_llm_request_attributes( - span: Optional["Span"], - params: dict, - *, - stream: bool = False, -) -> None: - """Set ``gen_ai.request.*`` attributes on an LLM CLIENT span. - - *params* is the kwargs dict forwarded to the model engine - (``GenerationOptions.llm_params``); only the known request-parameter - keys are mapped — any other kwargs are ignored. ``stop`` / - ``stop_sequences`` is normalized to a list via :func:`_stop_sequences`. - ``gen_ai.request.stream`` is set only when *stream* is True (omitted - otherwise, per the spec's conditionally-required-iff-streaming rule). - - These are non-sensitive sampling parameters (Recommended by spec), so - unlike message content they are recorded whenever the span exists — - there is no content-capture gate. Safe to call with ``span=None`` - (no-op) so callers don't have to branch on whether tracing is enabled. - """ - if span is None: - return - for key, attr in _GENAI_REQUEST_PARAMS.items(): - value = params.get(key) - if value is not None: - span.set_attribute(attr, value) - stop_sequences = _stop_sequences(params) - if stop_sequences is not None: - span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES, stop_sequences) - if stream: - span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_STREAM, True) - - -def set_llm_response_attributes( - span: Optional["Span"], - *, - model: Optional[str] = None, - response_id: Optional[str] = None, - finish_reason: Optional[str] = None, - usage: Optional["UsageInfo"] = None, -) -> None: - """Set ``gen_ai.response.*`` and ``gen_ai.usage.*`` attrs on an LLM CLIENT span. - - Each attribute is set only when its source value is non-``None`` so - backends can distinguish an absent value from a real zero. - *finish_reason* is a single value wrapped into a one-element list to - match the spec's ``gen_ai.response.finish_reasons`` string[] shape. - Reasoning tokens are recorded only when the provider returned them. - ``gen_ai.usage.total_tokens`` is intentionally never emitted — it was - removed from the current spec. - - Callers feed the values from whatever source they have: the - non-streaming path reads them off the returned ``LLMResponse``; the - streaming path passes the fields accumulated across chunks (model and - id arrive early, finish_reason and usage on the terminal chunk). Like - :func:`set_llm_request_attributes`, these are non-sensitive telemetry - recorded whenever the span exists — no content-capture gate. Safe to - call with ``span=None`` (no-op). - """ - if span is None: - return - if model is not None: - span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, model) - if response_id is not None: - span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_ID, response_id) - if finish_reason is not None: - span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason]) - if usage is not None: - span.set_attribute(GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) - span.set_attribute(GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) - if usage.reasoning_tokens is not None: - span.set_attribute( - GenAIAttributes.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, - usage.reasoning_tokens, - ) - - def set_request_content( span: Optional["Span"], input_messages: LLMMessages, @@ -828,44 +505,6 @@ def action_span(tracer: Optional["Tracer"], action_name: str) -> Generator[Optio raise -@contextmanager -def llm_call_span( - tracer: Optional["Tracer"], - model_name: str, - provider_name: str, - operation_name: str = "chat", -) -> Generator[Optional["Span"], None, None]: - """Create a CLIENT span for an LLM call following GenAI semantic conventions. - - Span name follows the OTEL pattern: ``"{operation_name} {model_name}"``. - - ``operation_name`` defaults to ``"chat"`` because IORails only issues - chat completions. In the future if any other non-chat LLM operations are - supported, callers should pass an explicit ``operation_name`` from the - OTEL GenAI semantic conventions. - - Yields the span (or ``None`` when *tracer* is ``None``). - """ - if tracer is None: - yield None - return - span_name = f"{operation_name} {model_name}" - with tracer.start_as_current_span( - span_name, - kind=SpanKind.CLIENT, - record_exception=False, - set_status_on_exception=False, - ) as span: - span.set_attribute(GenAIAttributes.GEN_AI_OPERATION_NAME, operation_name) - span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model_name) - span.set_attribute(GenAIAttributes.GEN_AI_PROVIDER_NAME, provider_name) - try: - yield span - except BaseException as exc: - record_span_error(span, exc) - raise - - @contextmanager def api_call_span(tracer: Optional["Tracer"], api_name: str) -> Generator[Optional["Span"], None, None]: """Create a CLIENT span for a non-LLM API call (e.g., jailbreak detection). diff --git a/nemoguardrails/llm/models/instrumented.py b/nemoguardrails/llm/models/instrumented.py new file mode 100644 index 0000000000..617d6671a9 --- /dev/null +++ b/nemoguardrails/llm/models/instrumented.py @@ -0,0 +1,256 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time +import warnings +from collections.abc import AsyncIterator, Mapping +from contextlib import nullcontext +from typing import Any, Optional, Union + +from nemoguardrails.llm.telemetry import ( + llm_call_span, + set_llm_call_content, + set_llm_request_attributes, + set_llm_response_attributes, +) +from nemoguardrails.tracing.constants import ( + OperationNames, + SystemConstants, + llm_operation_duration, + record_time_per_output_chunk, + record_time_to_first_chunk, + record_token_usage, +) +from nemoguardrails.types import ChatMessage, LLMModel, LLMResponse, LLMResponseChunk, UsageInfo + + +class InstrumentedLLMModel: + """Decorate an ``LLMModel`` with tracing, metrics, and content capture. + + Requests and responses are delegated to the wrapped model without changing + their runtime contract. Wrapping an existing ``InstrumentedLLMModel`` is + idempotent and returns the existing decorator. + + The tracer controls client-span creation, while metrics and content capture + are independently opt-in. ``default_request_params`` contributes provider + defaults to telemetry without changing the parameters sent to the model. + """ + + def __new__(cls, model: LLMModel, *args: Any, **kwargs: Any): + if isinstance(model, cls): + return model + return super().__new__(cls) + + def __init__( + self, + model: LLMModel, + *, + tracer: Optional[Any] = None, + metrics_enabled: bool = False, + content_capture_enabled: bool = False, + default_request_params: Optional[Mapping[str, Any]] = None, + ) -> None: + if model is self: + if ( + tracer is not self._tracer + or metrics_enabled != self._metrics_enabled + or content_capture_enabled != self._content_capture_enabled + or dict(default_request_params or {}) != self._default_request_params + ): + warnings.warn( + "InstrumentedLLMModel is already instrumented; new instrumentation " + "settings are ignored. Re-instrument the underlying wrapped_model instead.", + stacklevel=2, + ) + return + self._model = model + self._tracer = tracer + self._metrics_enabled = metrics_enabled + self._content_capture_enabled = content_capture_enabled + self._default_request_params = dict(default_request_params or {}) + + @property + def model_name(self) -> str: + return self._model.model_name + + @property + def provider_name(self) -> Optional[str]: + return self._model.provider_name + + @property + def provider_url(self) -> Optional[str]: + return self._model.provider_url + + @property + def wrapped_model(self) -> LLMModel: + """Return the underlying model for direct access or re-instrumentation.""" + return self._model + + def _request_params(self, stop: Optional[list[str]], kwargs: dict[str, Any]) -> dict[str, Any]: + params = {**self._default_request_params, **kwargs} + if stop is not None: + params["stop"] = stop + return params + + @staticmethod + def _input_messages(prompt: Union[str, list[ChatMessage]]) -> list[dict[str, Any]]: + if isinstance(prompt, str): + return [{"role": "user", "content": prompt}] + return [message.to_dict() if isinstance(message, ChatMessage) else message for message in prompt] + + async def generate_async( + self, + prompt: Union[str, list[ChatMessage]], + *, + stop: Optional[list[str]] = None, + **kwargs: Any, + ) -> LLMResponse: + """Generate one response while recording the LLM call telemetry. + + The prompt, stop sequences, and provider-specific keyword arguments are + forwarded unchanged to the wrapped model. + """ + operation_name = OperationNames.CHAT + provider_name = self.provider_name or SystemConstants.UNKNOWN + params = self._request_params(stop, kwargs) + with llm_call_span(self._tracer, self.model_name, provider_name, operation_name) as span: + set_llm_request_attributes(span, params) + duration = ( + llm_operation_duration(self.model_name, provider_name, operation_name) + if self._metrics_enabled + else nullcontext() + ) + with duration: + response = await self._model.generate_async(prompt, stop=stop, **kwargs) + set_llm_response_attributes( + span, + model=response.model, + response_id=response.request_id, + finish_reason=response.finish_reason, + usage=response.usage, + ) + if self._content_capture_enabled: + set_llm_call_content(span, self._input_messages(prompt), response.content) + if self._metrics_enabled: + record_token_usage(self.model_name, provider_name, operation_name, response.usage) + return response + + async def stream_async( + self, + prompt: Union[str, list[ChatMessage]], + *, + stop: Optional[list[str]] = None, + **kwargs: Any, + ) -> AsyncIterator[LLMResponseChunk]: + """Stream response chunks while recording one LLM call lifecycle. + + Metrics and response attributes are finalized when the stream completes + or closes, and the wrapped stream is always closed when it supports + ``aclose``. + """ + operation_name = OperationNames.CHAT + provider_name = self.provider_name or SystemConstants.UNKNOWN + params = self._request_params(stop, kwargs) + captured_usage: Optional[UsageInfo] = None + captured_model: Optional[str] = None + captured_response_id: Optional[str] = None + captured_finish_reason: Optional[str] = None + content_parts: list[str] = [] + stream = self._model.stream_async(prompt, stop=stop, **kwargs) + with llm_call_span(self._tracer, self.model_name, provider_name, operation_name) as span: + set_llm_request_attributes(span, params, stream=True) + duration = ( + llm_operation_duration(self.model_name, provider_name, operation_name) + if self._metrics_enabled + else nullcontext() + ) + try: + with duration: + started_at = time.monotonic() if self._metrics_enabled else 0.0 + last_chunk_at: Optional[float] = None + async for chunk in stream: + if self._metrics_enabled and (chunk.delta_content or chunk.delta_reasoning): + now = time.monotonic() + if last_chunk_at is None: + record_time_to_first_chunk( + self.model_name, + provider_name, + operation_name, + now - started_at, + ) + else: + record_time_per_output_chunk( + self.model_name, + provider_name, + operation_name, + now - last_chunk_at, + ) + last_chunk_at = now + if chunk.model is not None: + captured_model = chunk.model + if chunk.request_id is not None: + captured_response_id = chunk.request_id + if chunk.finish_reason is not None: + captured_finish_reason = chunk.finish_reason + if chunk.usage is not None: + captured_usage = chunk.usage + if self._content_capture_enabled and chunk.delta_content: + content_parts.append(chunk.delta_content) + yield chunk + finally: + close = getattr(stream, "aclose", None) + if close is not None: + await close() + set_llm_response_attributes( + span, + model=captured_model, + response_id=captured_response_id, + finish_reason=captured_finish_reason, + usage=captured_usage, + ) + if self._content_capture_enabled: + output_text = "".join(content_parts) if content_parts else None + set_llm_call_content(span, self._input_messages(prompt), output_text) + if self._metrics_enabled: + record_token_usage(self.model_name, provider_name, operation_name, captured_usage) + + +def instrument_llm_model( + model: LLMModel, + *, + tracer: Optional[Any] = None, + metrics_enabled: bool = False, + content_capture_enabled: bool = False, + default_request_params: Optional[Mapping[str, Any]] = None, +) -> LLMModel: + """Return ``model`` decorated with the requested telemetry. + + The original model is returned when neither tracing nor metrics are enabled. + Existing ``InstrumentedLLMModel`` instances are returned unchanged. + """ + if isinstance(model, InstrumentedLLMModel): + return model + if tracer is None and not metrics_enabled: + return model + return InstrumentedLLMModel( + model, + tracer=tracer, + metrics_enabled=metrics_enabled, + content_capture_enabled=content_capture_enabled, + default_request_params=default_request_params, + ) + + +__all__ = ["InstrumentedLLMModel", "instrument_llm_model"] diff --git a/nemoguardrails/llm/telemetry.py b/nemoguardrails/llm/telemetry.py new file mode 100644 index 0000000000..7398a8e4ce --- /dev/null +++ b/nemoguardrails/llm/telemetry.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from contextlib import contextmanager, suppress +from typing import TYPE_CHECKING, Any, Generator, Optional + +from nemoguardrails.tracing.constants import EventNames, GenAIAttributes, OtelContentCapture + +if TYPE_CHECKING: + from opentelemetry.trace import Span, Tracer + + from nemoguardrails.types import UsageInfo +else: + try: + from opentelemetry.trace import SpanKind, StatusCode + except ImportError: + SpanKind = None + StatusCode = None + + +_LEGACY_EVENT_BY_ROLE = { + "system": EventNames.GEN_AI_SYSTEM_MESSAGE, + "user": EventNames.GEN_AI_USER_MESSAGE, + "assistant": EventNames.GEN_AI_ASSISTANT_MESSAGE, + "tool": EventNames.GEN_AI_TOOL_MESSAGE, +} + +_GENAI_REQUEST_PARAMS = { + "temperature": GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, + "max_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS, + "max_completion_tokens": GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS, + "top_p": GenAIAttributes.GEN_AI_REQUEST_TOP_P, + "top_k": GenAIAttributes.GEN_AI_REQUEST_TOP_K, + "frequency_penalty": GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "presence_penalty": GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY, +} + + +def _use_json_span_format() -> bool: + tokens = {token.strip() for token in os.environ.get(OtelContentCapture.STABILITY_OPT_IN_ENV, "").split(",")} + return OtelContentCapture.STABILITY_OPT_IN_LATEST in tokens + + +def _system_parts_from_messages(messages: list[dict[str, Any]]) -> list[dict]: + return [ + {"type": "text", "content": message["content"]} + for message in messages + if message.get("role") == "system" and message.get("content") is not None + ] + + +def _non_system_input_messages(messages: list[dict[str, Any]]) -> list[dict]: + return [ + { + "role": message["role"], + "parts": [{"type": "text", "content": message["content"]}], + } + for message in messages + if message.get("role") is not None and message.get("role") != "system" and message.get("content") is not None + ] + + +def _set_llm_call_content_json(span: "Span", input_messages: list[dict[str, Any]], output_text: Optional[str]) -> None: + system_parts = _system_parts_from_messages(input_messages) + if system_parts: + span.set_attribute(GenAIAttributes.GEN_AI_SYSTEM_INSTRUCTIONS, json.dumps(system_parts)) + non_system = _non_system_input_messages(input_messages) + if non_system: + span.set_attribute(GenAIAttributes.GEN_AI_INPUT_MESSAGES, json.dumps(non_system)) + if output_text is not None: + output_messages = [{"role": "assistant", "parts": [{"type": "text", "content": output_text}]}] + span.set_attribute(GenAIAttributes.GEN_AI_OUTPUT_MESSAGES, json.dumps(output_messages)) + + +def _set_llm_call_content_events( + span: "Span", input_messages: list[dict[str, Any]], output_text: Optional[str] +) -> None: + for message in input_messages: + role = message.get("role") + content = message.get("content") + event_name = _LEGACY_EVENT_BY_ROLE.get(role) + if event_name is not None and isinstance(role, str) and isinstance(content, str): + span.add_event(event_name, attributes={"role": role, "content": content}) + if output_text is not None: + span.add_event( + EventNames.GEN_AI_CHOICE, + attributes={"index": 0, "message.role": "assistant", "message.content": output_text}, + ) + + +def set_llm_call_content( + span: Optional["Span"], + input_messages: list[dict[str, Any]], + output_text: Optional[str] = None, +) -> None: + """Record LLM input and output content on a span when available. + + The configured OpenTelemetry content format determines whether content is + stored as JSON attributes or legacy events. Telemetry failures are ignored + so they cannot affect the model call. + """ + if span is None: + return + with suppress(Exception): + if _use_json_span_format(): + _set_llm_call_content_json(span, input_messages, output_text) + else: + _set_llm_call_content_events(span, input_messages, output_text) + + +def _stop_sequences(params: dict) -> Optional[list]: + value = params.get("stop") + if value is None: + value = params.get("stop_sequences") + if isinstance(value, str) and value: + return [value] + if isinstance(value, list) and value: + return value + return None + + +def set_llm_request_attributes(span: Optional["Span"], params: dict, *, stream: bool = False) -> None: + """Record supported GenAI request parameters on a span. + + Unknown parameters are ignored, as are telemetry failures. + """ + if span is None: + return + with suppress(Exception): + for key, attribute in _GENAI_REQUEST_PARAMS.items(): + value = params.get(key) + if value is not None: + span.set_attribute(attribute, value) + stop_sequences = _stop_sequences(params) + if stop_sequences is not None: + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_STOP_SEQUENCES, stop_sequences) + if stream: + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_STREAM, True) + + +def set_llm_response_attributes( + span: Optional["Span"], + *, + model: Optional[str] = None, + response_id: Optional[str] = None, + finish_reason: Optional[str] = None, + usage: Optional["UsageInfo"] = None, +) -> None: + """Record available model response metadata and token usage on a span.""" + if span is None: + return + with suppress(Exception): + if model is not None: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, model) + if response_id is not None: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_ID, response_id) + if finish_reason is not None: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason]) + if usage is not None: + span.set_attribute(GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) + span.set_attribute(GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) + if usage.reasoning_tokens is not None: + span.set_attribute( + GenAIAttributes.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, + usage.reasoning_tokens, + ) + + +def record_span_error(span: Optional["Span"], exc: BaseException) -> None: + """Record an exception and error status without changing call behavior.""" + if span is None: + return + with suppress(Exception): + span.set_attribute("error.type", type(exc).__name__) + span.record_exception(exc) + span.set_status(StatusCode.ERROR, str(exc)) + + +@contextmanager +def llm_call_span( + tracer: Optional["Tracer"], + model_name: str, + provider_name: str, + operation_name: str = "chat", +) -> Generator[Optional["Span"], None, None]: + """Create a GenAI client span for one LLM call. + + Yields ``None`` when no tracer is configured. Exceptions from the wrapped + call are recorded on the span and re-raised unchanged. + """ + if tracer is None: + yield None + return + with tracer.start_as_current_span( + f"{operation_name} {model_name}", + kind=SpanKind.CLIENT, + record_exception=False, + set_status_on_exception=False, + ) as span: + with suppress(Exception): + span.set_attribute(GenAIAttributes.GEN_AI_OPERATION_NAME, operation_name) + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model_name) + span.set_attribute(GenAIAttributes.GEN_AI_PROVIDER_NAME, provider_name) + try: + yield span + except BaseException as exc: + record_span_error(span, exc) + raise + + +__all__ = [ + "llm_call_span", + "record_span_error", + "set_llm_call_content", + "set_llm_request_attributes", + "set_llm_response_attributes", +] diff --git a/nemoguardrails/tracing/constants.py b/nemoguardrails/tracing/constants.py index bc44f66c91..f24747af3d 100644 --- a/nemoguardrails/tracing/constants.py +++ b/nemoguardrails/tracing/constants.py @@ -242,7 +242,7 @@ class SpanNames: class MetricNames: - """OTEL metric names emitted by the IORails engine. + """OTEL metric names emitted by Guardrails engines. These names are part of the library's public API — customers point dashboards and alerts at them. Tests deliberately assert on the raw @@ -266,9 +266,9 @@ class MetricNames: STREAM_ACTIVE = "guardrails.stream.active" STREAM_REJECTIONS = "guardrails.stream.rejections" - # OTEL GenAI semantic-convention metric names emitted by IORails for - # downstream LLM calls. These names are mandated by OTEL hence ``gen_ai`` - # prefix separate to ``guardrails`` metrics above. + # OTEL GenAI semantic-convention metric names emitted for downstream LLM + # calls. These names are mandated by OTEL hence ``gen_ai`` prefix separate + # to ``guardrails`` metrics above. GEN_AI_CLIENT_TOKEN_USAGE = "gen_ai.client.token.usage" GEN_AI_CLIENT_OPERATION_DURATION = "gen_ai.client.operation.duration" GEN_AI_CLIENT_OPERATION_TIME_TO_FIRST_CHUNK = "gen_ai.client.operation.time_to_first_chunk" diff --git a/tests/guardrails/test_telemetry_content_capture.py b/tests/guardrails/test_telemetry_content_capture.py index 2d817f32a1..6da069d305 100644 --- a/tests/guardrails/test_telemetry_content_capture.py +++ b/tests/guardrails/test_telemetry_content_capture.py @@ -31,15 +31,17 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from nemoguardrails.guardrails.telemetry import ( + is_content_capture_enabled, + set_rail_content, + set_request_content, +) +from nemoguardrails.llm.telemetry import ( _non_system_input_messages, _set_llm_call_content_events, _set_llm_call_content_json, _system_parts_from_messages, _use_json_span_format, - is_content_capture_enabled, set_llm_call_content, - set_rail_content, - set_request_content, ) from nemoguardrails.rails.llm.config import TracingConfig from nemoguardrails.tracing.constants import ( diff --git a/tests/guardrails/test_telemetry_spans.py b/tests/guardrails/test_telemetry_spans.py index 6fcd28fae1..853ee8b416 100644 --- a/tests/guardrails/test_telemetry_spans.py +++ b/tests/guardrails/test_telemetry_spans.py @@ -27,8 +27,10 @@ from nemoguardrails.guardrails.telemetry import ( action_span, api_call_span, - llm_call_span, rail_span, +) +from nemoguardrails.llm.telemetry import ( + llm_call_span, set_llm_request_attributes, set_llm_response_attributes, ) diff --git a/tests/llm/models/test_instrumented.py b/tests/llm/models/test_instrumented.py new file mode 100644 index 0000000000..6dfdbbf536 --- /dev/null +++ b/tests/llm/models/test_instrumented.py @@ -0,0 +1,383 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json +import warnings +from unittest.mock import patch + +import pytest +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from nemoguardrails.actions.llm.utils import llm_call +from nemoguardrails.guardrails import telemetry +from nemoguardrails.llm.models.instrumented import InstrumentedLLMModel, instrument_llm_model +from nemoguardrails.tracing import constants as tracing_constants +from nemoguardrails.tracing.constants import SystemConstants +from nemoguardrails.types import ChatMessage, LLMModel, LLMResponse, LLMResponseChunk, UsageInfo +from tests.guardrails.metric_helpers import collect_histogram_sum, collect_metric_points + + +class RecordingModel: + model_name = "test-model" + provider_name = "test-provider" + provider_url = "https://example.test/v1" + + def __init__(self, response=None, chunks=None, error=None): + self.response = response or LLMResponse(content="response") + self.chunks = chunks or [] + self.error = error + self.calls = [] + self.stream_closed = False + + async def generate_async(self, prompt, *, stop=None, **kwargs): + self.calls.append((prompt, stop, kwargs)) + if self.error is not None: + raise self.error + return self.response + + async def stream_async(self, prompt, *, stop=None, **kwargs): + self.calls.append((prompt, stop, kwargs)) + try: + for chunk in self.chunks: + yield chunk + if self.error is not None: + raise self.error + finally: + self.stream_closed = True + + +@pytest.fixture(autouse=True) +def reset_telemetry(): + telemetry._meter = None + tracing_constants._llm_instruments = None + yield + telemetry._meter = None + tracing_constants._llm_instruments = None + + +@pytest.fixture +def span_exporter(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + return provider.get_tracer("test"), exporter + + +@pytest.fixture +def metric_reader(): + reader = InMemoryMetricReader() + provider = MeterProvider(metric_readers=[reader]) + telemetry._meter = provider.get_meter(SystemConstants.SYSTEM_NAME) + return reader + + +@pytest.mark.asyncio +async def test_generate_delegates_and_returns_exact_response(span_exporter): + tracer, exporter = span_exporter + response = LLMResponse( + content="answer", + model="response-model", + request_id="request-id", + finish_reason="stop", + usage=UsageInfo(input_tokens=3, output_tokens=2), + ) + model = RecordingModel(response=response) + instrumented = InstrumentedLLMModel(model, tracer=tracer) + prompt = [ChatMessage(role="user", content="question")] + + result = await instrumented.generate_async(prompt, stop=["END"], temperature=0.2) + + assert isinstance(instrumented, LLMModel) + assert result is response + assert model.calls == [(prompt, ["END"], {"temperature": 0.2})] + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["gen_ai.request.model"] == "test-model" + assert attrs["gen_ai.provider.name"] == "test-provider" + assert attrs["gen_ai.request.temperature"] == 0.2 + assert list(attrs["gen_ai.request.stop_sequences"]) == ["END"] + assert attrs["gen_ai.response.model"] == "response-model" + assert attrs["gen_ai.response.id"] == "request-id" + assert list(attrs["gen_ai.response.finish_reasons"]) == ["stop"] + assert attrs["gen_ai.usage.input_tokens"] == 3 + assert attrs["gen_ai.usage.output_tokens"] == 2 + + +@pytest.mark.asyncio +async def test_generate_preserves_exception_identity(span_exporter): + tracer, exporter = span_exporter + error = RuntimeError("provider failed") + model = InstrumentedLLMModel(RecordingModel(error=error), tracer=tracer) + + with pytest.raises(RuntimeError) as exc_info: + await model.generate_async("question") + + assert exc_info.value is error + traceback = exc_info.value.__traceback__ + while traceback and traceback.tb_next: + traceback = traceback.tb_next + assert traceback is not None + assert traceback.tb_frame.f_code is RecordingModel.generate_async.__code__ + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["error.type"] == "RuntimeError" + + +@pytest.mark.asyncio +async def test_llm_call_uses_instrumented_custom_model(span_exporter): + tracer, exporter = span_exporter + model = InstrumentedLLMModel(RecordingModel(), tracer=tracer) + + result = await llm_call(model, "question") + + assert result.content == "response" + assert len(exporter.get_finished_spans()) == 1 + + +@pytest.mark.asyncio +async def test_stream_preserves_chunks_and_latest_non_null_metadata(span_exporter): + tracer, exporter = span_exporter + chunks = [ + LLMResponseChunk(delta_content="a", model="first", request_id="id-1"), + LLMResponseChunk(delta_reasoning="thinking", model="second"), + LLMResponseChunk(delta_content="b", finish_reason="stop"), + LLMResponseChunk(usage=UsageInfo(input_tokens=4, output_tokens=2)), + ] + underlying = RecordingModel(chunks=chunks) + model = InstrumentedLLMModel(underlying, tracer=tracer) + + received = [chunk async for chunk in model.stream_async("question")] + + assert received == chunks + assert all(received_chunk is source_chunk for received_chunk, source_chunk in zip(received, chunks)) + assert underlying.stream_closed + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["gen_ai.request.stream"] is True + assert attrs["gen_ai.response.model"] == "second" + assert attrs["gen_ai.response.id"] == "id-1" + assert list(attrs["gen_ai.response.finish_reasons"]) == ["stop"] + assert attrs["gen_ai.usage.input_tokens"] == 4 + assert attrs["gen_ai.usage.output_tokens"] == 2 + + +@pytest.mark.asyncio +async def test_stream_error_has_no_final_response_telemetry(span_exporter): + tracer, exporter = span_exporter + error = RuntimeError("stream failed") + source = RecordingModel( + chunks=[LLMResponseChunk(delta_content="partial", model="partial-model")], + error=error, + ) + model = InstrumentedLLMModel(source, tracer=tracer) + + with pytest.raises(RuntimeError) as exc_info: + async for _ in model.stream_async("question"): + pass + + assert exc_info.value is error + traceback = exc_info.value.__traceback__ + while traceback and traceback.tb_next: + traceback = traceback.tb_next + assert traceback is not None + assert traceback.tb_frame.f_code is RecordingModel.stream_async.__code__ + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["error.type"] == "RuntimeError" + assert "gen_ai.response.model" not in attrs + assert "gen_ai.usage.input_tokens" not in attrs + + +@pytest.mark.asyncio +async def test_stream_consumer_close_has_no_final_response_telemetry(span_exporter): + tracer, exporter = span_exporter + source = RecordingModel( + chunks=[ + LLMResponseChunk(delta_content="partial", model="partial-model"), + LLMResponseChunk(usage=UsageInfo(input_tokens=1, output_tokens=1)), + ] + ) + model = InstrumentedLLMModel(source, tracer=tracer) + stream = model.stream_async("question") + + assert await anext(stream) is source.chunks[0] + await stream.aclose() + + assert source.stream_closed + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["error.type"] == "GeneratorExit" + assert "gen_ai.response.model" not in attrs + assert "gen_ai.usage.input_tokens" not in attrs + + +@pytest.mark.asyncio +async def test_stream_task_cancellation_preserves_cancelled_error(span_exporter): + tracer, exporter = span_exporter + started = asyncio.Event() + + class BlockingModel(RecordingModel): + async def stream_async(self, prompt, *, stop=None, **kwargs): + try: + started.set() + await asyncio.Event().wait() + yield LLMResponseChunk(delta_content="unreachable") + finally: + self.stream_closed = True + + source = BlockingModel() + model = InstrumentedLLMModel(source, tracer=tracer) + + async def consume(): + async for _ in model.stream_async("question"): + pass + + task = asyncio.create_task(consume()) + await started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert source.stream_closed + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert attrs["error.type"] == "CancelledError" + assert "gen_ai.response.model" not in attrs + + +@pytest.mark.asyncio +async def test_content_capture_is_disabled_by_default(span_exporter): + tracer, exporter = span_exporter + model = InstrumentedLLMModel(RecordingModel(response=LLMResponse(content="secret output")), tracer=tracer) + + await model.generate_async("secret input") + + span = exporter.get_finished_spans()[0] + assert all("secret" not in str(value) for value in span.attributes.values()) + assert all("secret" not in str(event.attributes) for event in span.events) + + +@pytest.mark.asyncio +async def test_content_capture_uses_existing_format_switch(span_exporter): + tracer, exporter = span_exporter + model = InstrumentedLLMModel( + RecordingModel(response=LLMResponse(content="captured output")), + tracer=tracer, + content_capture_enabled=True, + ) + + with patch.dict("os.environ", {"OTEL_SEMCONV_STABILITY_OPT_IN": "gen_ai_latest_experimental"}): + await model.generate_async("captured input") + + attrs = dict(exporter.get_finished_spans()[0].attributes) + assert json.loads(attrs["gen_ai.input.messages"])[0]["parts"][0]["content"] == "captured input" + assert json.loads(attrs["gen_ai.output.messages"])[0]["parts"][0]["content"] == "captured output" + + +@pytest.mark.asyncio +async def test_metrics_and_chunk_timing(metric_reader): + usage = UsageInfo(input_tokens=5, output_tokens=2) + chunks = [ + LLMResponseChunk(delta_reasoning="think"), + LLMResponseChunk(delta_content="answer"), + LLMResponseChunk(usage=usage), + ] + model = InstrumentedLLMModel(RecordingModel(chunks=chunks), metrics_enabled=True) + + async for _ in model.stream_async("question"): + pass + + points = collect_metric_points(metric_reader) + assert len(points["gen_ai.client.token.usage"]) == 2 + assert points["gen_ai.client.operation.time_to_first_chunk"][0].value == 1 + assert points["gen_ai.client.operation.time_per_output_chunk"][0].value == 1 + assert points["gen_ai.client.operation.duration"][0].value == 1 + + +@pytest.mark.asyncio +async def test_disabled_telemetry_returns_original_model(metric_reader): + model = RecordingModel() + + result = instrument_llm_model(model) + response = await result.generate_async("question") + + assert result is model + assert response is model.response + assert collect_metric_points(metric_reader) == {} + + +@pytest.mark.asyncio +async def test_instrumentation_is_idempotent_and_does_not_own_model(span_exporter): + tracer, exporter = span_exporter + model = RecordingModel() + first = InstrumentedLLMModel(model, tracer=tracer) + second = InstrumentedLLMModel(first, metrics_enabled=True) + third = instrument_llm_model(second, tracer=tracer) + + await third.generate_async("question") + + assert second is first + assert third is first + assert len(exporter.get_finished_spans()) == 1 + assert first.wrapped_model is model + assert not hasattr(first, "close") + assert not hasattr(first, "aclose") + + +@pytest.mark.asyncio +async def test_stream_cleanup_runs_outside_duration_metric(metric_reader): + close_delay = 0.2 + + class SlowClosingModel(RecordingModel): + async def stream_async(self, prompt, *, stop=None, **kwargs): + try: + yield LLMResponseChunk(delta_content="a") + yield LLMResponseChunk(delta_content="b") + finally: + await asyncio.sleep(close_delay) + self.stream_closed = True + + source = SlowClosingModel() + model = InstrumentedLLMModel(source, metrics_enabled=True) + stream = model.stream_async("question") + + assert await anext(stream) is not None + await stream.aclose() + + assert source.stream_closed + points = collect_metric_points(metric_reader) + assert len(points["gen_ai.client.operation.duration"]) == 1 + recorded_duration = collect_histogram_sum(metric_reader, "gen_ai.client.operation.duration") + assert recorded_duration < close_delay / 2 + + +def test_reinstrument_with_changed_settings_warns_and_keeps_original(): + instrumented = InstrumentedLLMModel(RecordingModel(), metrics_enabled=True) + + with pytest.warns(UserWarning, match="already instrumented"): + again = InstrumentedLLMModel(instrumented, content_capture_enabled=True) + + assert again is instrumented + assert again._content_capture_enabled is False + + +def test_reinstrument_with_identical_settings_does_not_warn(): + instrumented = InstrumentedLLMModel(RecordingModel(), metrics_enabled=True) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + again = InstrumentedLLMModel(instrumented, metrics_enabled=True) + + assert again is instrumented