Skip to content

feat(llm): add shared model telemetry and instrumented decorator #2214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Pouyanpi wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

feat(llm): add shared model telemetry and instrumented decorator #2214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
71ca58d
feat(llm): add shared model telemetry and instrumented decorator
Pouyanpi Jul 23, 2026
5cd4a2b
refactor(guardrails): route IORails model calls through the shared de…
Pouyanpi Jul 23, 2026
8de6ee0
docs(llm): document model telemetry surfaces
Pouyanpi Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions nemoguardrails/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -82,6 +83,7 @@
"ChatMessage",
"FinishReason",
"Guardrails",
"InstrumentedLLMModel",
"LLMFramework",
"LLMModel",
"LLMRails",
Expand Down
234 changes: 74 additions & 160 deletions nemoguardrails/guardrails/engine_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
Loading