nemo_relay/api/
llm.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use bitflags::bitflags;
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use serde_json::json;
10use typed_builder::TypedBuilder;
11use uuid::Uuid;
12
13use crate::api::runtime::NemoRelayContextState;
14use crate::api::runtime::current_scope_stack;
15use crate::api::runtime::global_context;
16use crate::api::runtime::{
17    LlmCollectorFn, LlmExecutionNextFn, LlmFinalizerFn, LlmJsonStream, LlmStreamExecutionNextFn,
18};
19use crate::api::scope::event;
20use crate::api::scope::{EmitMarkEventParams, ScopeHandle};
21use crate::api::shared::{
22    ensure_runtime_owner, resolve_parent_uuid, run_request_intercepts_with_codec,
23    snapshot_event_subscribers,
24};
25use crate::codec::request::AnnotatedLlmRequest;
26use crate::codec::response::AnnotatedLlmResponse;
27use crate::codec::traits::{LlmCodec, LlmResponseCodec};
28use crate::error::{FlowError, Result};
29use crate::json::Json;
30use crate::stream::LlmStreamWrapper;
31
32bitflags! {
33    /// Bitflags that modify LLM-call behavior and observability.
34    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
35    pub struct LlmAttributes: u32 {
36        /// Marks the request as stateful from the runtime's perspective.
37        const STATEFUL = 0b01;
38        /// Marks the request as streaming.
39        const STREAMING = 0b10;
40    }
41}
42
43/// Runtime-owned handle identifying an active or completed LLM call.
44#[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)]
45#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
46pub struct LlmHandle {
47    /// Unique LLM-call identifier.
48    #[builder(default = Uuid::now_v7())]
49    pub uuid: Uuid,
50    /// Timestamp captured when the LLM handle was created.
51    #[builder(default = Utc::now())]
52    pub started_at: DateTime<Utc>,
53    /// Provider or logical call name recorded on lifecycle events.
54    #[builder(setter(into))]
55    pub name: String,
56    /// Optional application payload stored on the handle.
57    #[builder(default)]
58    pub data: Option<Json>,
59    /// Optional metadata attached to the LLM span.
60    #[builder(default)]
61    pub metadata: Option<Json>,
62    /// LLM behavior flags.
63    #[builder(default = LlmAttributes::empty())]
64    pub attributes: LlmAttributes,
65    /// UUID of the parent scope, if any.
66    #[builder(default)]
67    pub parent_uuid: Option<Uuid>,
68    /// Optional normalized model name for observability.
69    #[builder(default, setter(into))]
70    pub model_name: Option<String>,
71}
72
73/// JSON-shaped LLM request payload passed through the runtime.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct LlmRequest {
76    /// Provider-specific request headers.
77    pub headers: serde_json::Map<String, Json>,
78    /// Provider-specific request body.
79    pub content: Json,
80}
81
82/// Builder parameters for [`NemoRelayContextState::create_llm_handle`].
83#[derive(Debug, Clone, TypedBuilder)]
84#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
85pub struct CreateLlmHandleParams<'a> {
86    /// Logical provider or model family name.
87    pub name: &'a str,
88    /// Optional parent scope UUID.
89    #[builder(default)]
90    pub parent_uuid: Option<uuid::Uuid>,
91    /// LLM attribute bitflags.
92    #[builder(default = LlmAttributes::empty())]
93    pub attributes: LlmAttributes,
94    /// Optional application payload stored on the handle.
95    #[builder(default)]
96    pub data: Option<Json>,
97    /// Optional metadata stored on the handle.
98    #[builder(default)]
99    pub metadata: Option<Json>,
100    /// Optional normalized model name stored on the handle.
101    #[builder(default, setter(into))]
102    pub model_name: Option<String>,
103    /// Optional timestamp captured as the handle start time and reused by the
104    /// emitted start event. When omitted, the current UTC time is used.
105    #[builder(default)]
106    pub timestamp: Option<DateTime<Utc>>,
107}
108
109/// Builder parameters for [`NemoRelayContextState::build_llm_end_event`].
110#[derive(Clone, TypedBuilder)]
111#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
112pub struct EndLlmHandleParams<'a> {
113    /// LLM handle to serialize into the emitted end event.
114    pub handle: &'a LlmHandle,
115    /// Optional data payload merged over the handle data.
116    #[builder(default)]
117    pub data: Option<Json>,
118    /// Optional metadata payload merged over the handle metadata.
119    #[builder(default)]
120    pub metadata: Option<Json>,
121    /// Optional normalized response annotation produced by a response codec.
122    #[builder(default)]
123    pub annotated_response: Option<Arc<AnnotatedLlmResponse>>,
124    /// Optional timestamp recorded on the emitted end event. When omitted, the
125    /// runtime records the current UTC time, or one microsecond after the
126    /// handle start time if the current time is not later.
127    #[builder(default)]
128    pub timestamp: Option<DateTime<Utc>>,
129}
130
131/// Builder parameters for [`llm_call`].
132#[derive(TypedBuilder)]
133#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
134pub struct LlmCallParams<'a> {
135    /// Logical provider or model family name recorded on the span.
136    pub name: &'a str,
137    /// Raw request associated with the span.
138    pub request: &'a LlmRequest,
139    /// Optional explicit parent scope.
140    #[builder(default)]
141    pub parent: Option<&'a ScopeHandle>,
142    /// LLM attribute bitflags applied to the span.
143    #[builder(default = LlmAttributes::empty())]
144    pub attributes: LlmAttributes,
145    /// Optional application payload stored on the handle but not emitted as
146    /// Agent Trajectory Observability Format (ATOF) data.
147    #[builder(default)]
148    pub data: Option<Json>,
149    /// Optional JSON metadata recorded on the start event.
150    #[builder(default)]
151    pub metadata: Option<Json>,
152    /// Optional normalized model name recorded separately from the request payload.
153    #[builder(default, setter(into))]
154    pub model_name: Option<String>,
155    /// Optional normalized request annotation produced by a codec.
156    #[builder(default)]
157    pub annotated_request: Option<Arc<AnnotatedLlmRequest>>,
158    /// Optional timestamp captured as the handle start time and reused by the
159    /// emitted start event. When omitted, the current UTC time is used.
160    #[builder(default)]
161    pub timestamp: Option<DateTime<Utc>>,
162}
163
164/// Builder parameters for [`llm_call_execute`].
165#[derive(TypedBuilder)]
166#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
167pub struct LlmCallExecuteParams {
168    /// Logical provider or model family name recorded on emitted events.
169    #[builder(setter(into))]
170    pub name: String,
171    /// Raw request passed into the managed pipeline.
172    pub request: LlmRequest,
173    /// Provider callback or execution continuation.
174    pub func: LlmExecutionNextFn,
175    /// Optional explicit parent scope for the emitted LLM span.
176    #[builder(default)]
177    pub parent: Option<ScopeHandle>,
178    /// LLM attribute bitflags applied to the managed span.
179    #[builder(default = LlmAttributes::empty())]
180    pub attributes: LlmAttributes,
181    /// Optional application payload stored on the handle but not emitted as
182    /// Agent Trajectory Observability Format (ATOF) data.
183    #[builder(default)]
184    pub data: Option<Json>,
185    /// Optional JSON metadata recorded on emitted events.
186    #[builder(default)]
187    pub metadata: Option<Json>,
188    /// Optional normalized model name for observability output.
189    #[builder(default, setter(into))]
190    pub model_name: Option<String>,
191    /// Optional request codec used to produce annotated request data.
192    #[builder(default)]
193    pub codec: Option<Arc<dyn LlmCodec>>,
194    /// Optional response codec used to attach annotated response data.
195    #[builder(default)]
196    pub response_codec: Option<Arc<dyn LlmResponseCodec>>,
197}
198
199/// Builder parameters for [`llm_stream_call_execute`].
200#[derive(TypedBuilder)]
201#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
202pub struct LlmStreamCallExecuteParams {
203    /// Logical provider or model family name recorded on emitted events.
204    #[builder(setter(into))]
205    pub name: String,
206    /// Raw request passed into the managed pipeline.
207    pub request: LlmRequest,
208    /// Streaming provider callback or execution continuation.
209    pub func: LlmStreamExecutionNextFn,
210    /// Per-chunk collector callback used to accumulate stream state.
211    pub collector: LlmCollectorFn,
212    /// Finalizer callback used to construct the completed response.
213    pub finalizer: LlmFinalizerFn,
214    /// Optional explicit parent scope for the emitted LLM span.
215    #[builder(default)]
216    pub parent: Option<ScopeHandle>,
217    /// LLM attribute bitflags applied to the managed span.
218    #[builder(default = LlmAttributes::empty())]
219    pub attributes: LlmAttributes,
220    /// Optional application payload stored on the handle but not emitted as
221    /// Agent Trajectory Observability Format (ATOF) data.
222    #[builder(default)]
223    pub data: Option<Json>,
224    /// Optional JSON metadata recorded on emitted events.
225    #[builder(default)]
226    pub metadata: Option<Json>,
227    /// Optional normalized model name for observability output.
228    #[builder(default, setter(into))]
229    pub model_name: Option<String>,
230    /// Optional request codec used to produce annotated request data.
231    #[builder(default)]
232    pub codec: Option<Arc<dyn LlmCodec>>,
233    /// Optional response codec used to attach annotated response data.
234    #[builder(default)]
235    pub response_codec: Option<Arc<dyn LlmResponseCodec>>,
236}
237
238/// Builder parameters for [`llm_call_end`].
239#[derive(TypedBuilder)]
240#[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))]
241pub struct LlmCallEndParams<'a> {
242    /// LLM handle to close.
243    pub handle: &'a LlmHandle,
244    /// Raw provider response associated with the end event.
245    pub response: Json,
246    /// Optional application payload retained for compatibility; Agent
247    /// Trajectory Observability Format (ATOF) data is the response.
248    #[builder(default)]
249    pub data: Option<Json>,
250    /// Optional JSON metadata recorded on the end event.
251    #[builder(default)]
252    pub metadata: Option<Json>,
253    /// Optional normalized response annotation produced by a response codec.
254    #[builder(default)]
255    pub annotated_response: Option<Arc<AnnotatedLlmResponse>>,
256    /// Optional response codec used to produce an annotation from sanitized event data.
257    #[builder(default)]
258    pub response_codec: Option<Arc<dyn LlmResponseCodec>>,
259    /// Optional timestamp recorded on the emitted end event. When omitted, the
260    /// runtime records the current UTC time, or one microsecond after the
261    /// handle start time if the current time is not later.
262    #[builder(default)]
263    pub timestamp: Option<DateTime<Utc>>,
264}
265
266fn create_llm_handle(params: CreateLlmHandleParams<'_>) -> Result<LlmHandle> {
267    ensure_runtime_owner()?;
268    let context = global_context();
269    let state = context
270        .read()
271        .map_err(|error| FlowError::Internal(error.to_string()))?;
272    Ok(state.create_llm_handle(params))
273}
274
275fn emit_llm_start(
276    handle: &LlmHandle,
277    request: &LlmRequest,
278    annotated_request: Option<Arc<AnnotatedLlmRequest>>,
279) -> Result<()> {
280    ensure_runtime_owner()?;
281    let (event, subscribers) = {
282        let scope_stack = current_scope_stack();
283        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
284        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
285            &registries.llm_sanitize_request_guardrails
286        });
287        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
288        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
289        let context = global_context();
290        let state = context
291            .read()
292            .map_err(|error| FlowError::Internal(error.to_string()))?;
293
294        let sanitized_request = state.llm_sanitize_request_chain(request.clone(), &scope_locals);
295        let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null);
296        let event = state.build_llm_start_event(handle, Some(input), annotated_request);
297        (event, subscribers)
298    };
299    NemoRelayContextState::emit_event(&event, &subscribers);
300    Ok(())
301}
302
303/// Start a manual LLM lifecycle span.
304///
305/// This emits an LLM-start event after applying sanitize-request guardrails to
306/// the payload recorded for observability.
307///
308/// # Parameters
309/// - `name`: Logical provider or model family name recorded on the span.
310/// - `request`: Raw [`LlmRequest`] associated with the span.
311/// - `parent`: Optional explicit parent scope.
312/// - `attributes`: LLM attribute bitflags applied to the span.
313/// - `data`: Optional application payload stored on the returned handle. The
314///   emitted start event data is the sanitized `request` payload.
315/// - `metadata`: Optional JSON metadata recorded on the start event.
316/// - `model_name`: Optional normalized model name recorded separately from the
317///   request payload.
318/// - `annotated_request`: Optional normalized request annotation produced by a
319///   codec.
320/// - `timestamp`: Optional timestamp recorded as the handle start time and on
321///   the emitted start event. When `None`, the current UTC time is used.
322///
323/// # Returns
324/// A [`Result`] containing the created [`LlmHandle`].
325///
326/// # Errors
327/// Returns an error when the runtime owner check fails or when internal state
328/// cannot be read safely.
329///
330/// # Notes
331/// Sanitize-request guardrails affect only the emitted start-event payload, not
332/// the caller-owned [`LlmRequest`].
333pub fn llm_call(params: LlmCallParams<'_>) -> Result<LlmHandle> {
334    let handle_params = CreateLlmHandleParams::builder()
335        .name(params.name)
336        .parent_uuid_opt(resolve_parent_uuid(params.parent))
337        .attributes(params.attributes)
338        .data_opt(params.data)
339        .metadata_opt(params.metadata)
340        .model_name_opt(params.model_name)
341        .timestamp_opt(params.timestamp)
342        .build();
343    let handle = create_llm_handle(handle_params)?;
344    emit_llm_start(&handle, params.request, params.annotated_request)?;
345    Ok(handle)
346}
347
348/// Finish a manual LLM lifecycle span.
349///
350/// This emits an LLM-end event for a handle previously returned by
351/// [`llm_call`].
352///
353/// # Parameters
354/// - `handle`: LLM handle to close.
355/// - `response`: Raw provider response associated with the end event.
356/// - `data`: Optional application payload retained for compatibility. The
357///   emitted end event data is the sanitized `response` unless it sanitizes to
358///   JSON null, in which case this payload is used.
359/// - `metadata`: Optional JSON metadata recorded on the end event.
360/// - `annotated_response`: Optional normalized response annotation produced by
361///   a response codec. When omitted and `response_codec` is supplied, the
362///   annotation is decoded from the sanitized end-event payload.
363/// - `response_codec`: Optional response codec used to produce a normalized
364///   response annotation from the sanitized end-event payload.
365/// - `timestamp`: Optional timestamp recorded on the emitted end event. When
366///   `None`, the runtime uses the current UTC time, or one microsecond after
367///   the handle start time if the current time is not later.
368///
369/// # Returns
370/// A [`Result`] that is `Ok(())` when the end event has been emitted.
371///
372/// # Errors
373/// Returns an error when the runtime owner check fails, internal state cannot be
374/// read safely, or response codec decoding fails.
375///
376/// # Notes
377/// Sanitize-response guardrails affect only the emitted end-event payload, not
378/// the caller-owned `response` value.
379pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> {
380    let LlmCallEndParams {
381        handle,
382        response,
383        data,
384        metadata,
385        annotated_response,
386        response_codec,
387        timestamp,
388    } = params;
389    ensure_runtime_owner()?;
390    let mut decode_error = None;
391    let (event, subscribers) = {
392        let scope_stack = current_scope_stack();
393        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
394        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
395            &registries.llm_sanitize_response_guardrails
396        });
397        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
398        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
399        let context = global_context();
400        let state = context
401            .read()
402            .map_err(|error| FlowError::Internal(error.to_string()))?;
403
404        let sanitized_response = state.llm_sanitize_response_chain(response, &scope_locals);
405        let data = if sanitized_response.is_null() {
406            data
407        } else {
408            Some(sanitized_response)
409        };
410        let annotated_response = match annotated_response {
411            Some(annotated_response) => Some(annotated_response),
412            None => match (response_codec.as_ref(), data.as_ref()) {
413                (Some(codec), Some(response)) => match codec.decode_response(response) {
414                    Ok(decoded) => Some(Arc::new(decoded)),
415                    Err(error) => {
416                        decode_error = Some(error);
417                        None
418                    }
419                },
420                _ => None,
421            },
422        };
423        let event = state.build_llm_end_event(
424            EndLlmHandleParams::builder()
425                .handle(handle)
426                .data_opt(data)
427                .metadata_opt(metadata)
428                .annotated_response_opt(annotated_response)
429                .timestamp_opt(timestamp)
430                .build(),
431        );
432        (event, subscribers)
433    };
434    NemoRelayContextState::emit_event(&event, &subscribers);
435    if let Some(error) = decode_error {
436        Err(error)
437    } else {
438        Ok(())
439    }
440}
441
442fn emit_llm_end_without_output(handle: &LlmHandle, metadata: Option<Json>) -> Result<()> {
443    ensure_runtime_owner()?;
444    let (event, subscribers) = {
445        let scope_stack = current_scope_stack();
446        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
447        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
448        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
449        let context = global_context();
450        let state = context
451            .read()
452            .map_err(|error| FlowError::Internal(error.to_string()))?;
453        let event = state.end_llm_handle(handle, handle.data.clone(), metadata, None);
454        (event, subscribers)
455    };
456    NemoRelayContextState::emit_event(&event, &subscribers);
457    Ok(())
458}
459
460/// Execute an LLM call through the managed middleware pipeline.
461///
462/// This runs conditional-execution guardrails, request intercepts, and
463/// sanitize-request guardrails, emits the LLM-start event, then runs execution
464/// intercepts, the provider callback when it is not replaced, and
465/// sanitize-response guardrails in the runtime-defined order.
466///
467/// # Parameters
468/// - `name`: Logical provider or model family name recorded on emitted events.
469/// - `request`: Raw [`LlmRequest`] passed into the managed pipeline.
470/// - `func`: Provider callback or execution continuation.
471/// - `parent`: Optional explicit parent scope for the emitted LLM span.
472/// - `attributes`: LLM attribute bitflags applied to the managed span.
473/// - `data`: Optional application payload stored on the managed LLM handle. It
474///   may be used on failure end events that have no output payload.
475/// - `metadata`: Optional JSON metadata recorded on emitted events.
476/// - `model_name`: Optional normalized model name for observability output.
477/// - `codec`: Optional request codec used to produce annotated request data for
478///   intercepts and events.
479/// - `response_codec`: Optional response codec used to attach annotated
480///   response data to the end event.
481///
482/// # Returns
483/// A [`Result`] containing the raw JSON response returned by the callback or
484/// an execution intercept.
485///
486/// # Errors
487/// Returns [`FlowError::GuardrailRejected`] when conditional-execution
488/// guardrails block the call, or any error raised by request intercepts,
489/// execution intercepts, codecs, or the callback itself.
490///
491/// # Notes
492/// The LLM-start event is emitted before execution intercepts run. When
493/// execution fails after that point, the runtime still emits an LLM-end event
494/// without an output payload.
495///
496/// Response codecs enrich observability output only and do not change the
497/// value returned to the caller.
498pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
499    let LlmCallExecuteParams {
500        name,
501        request,
502        func,
503        parent,
504        attributes,
505        data,
506        metadata,
507        model_name,
508        codec,
509        response_codec,
510    } = params;
511    ensure_runtime_owner()?;
512    {
513        let (entries, subscribers, parent_uuid, guardrail_metadata) = {
514            let scope_stack = current_scope_stack();
515            let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
516            let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
517                &registries.llm_conditional_execution_guardrails
518            });
519            let scope_subscribers = scope_guard.collect_scope_local_subscribers();
520            let context = global_context();
521            let state = context
522                .read()
523                .map_err(|error| FlowError::Internal(error.to_string()))?;
524            let entries = state.llm_conditional_execution_entries(&scope_locals);
525            let subscribers = state.collect_event_subscribers(&scope_subscribers);
526            (
527                entries,
528                subscribers,
529                resolve_parent_uuid(parent.as_ref()),
530                metadata.clone(),
531            )
532        };
533        if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
534            &request,
535            &entries,
536            &subscribers,
537            parent_uuid,
538            guardrail_metadata,
539        )? {
540            let mut rejection_data = json!({});
541            if let Some(object) = rejection_data.as_object_mut() {
542                object.insert("rejected".into(), json!(true));
543                object.insert("rejection_reason".into(), json!(&error));
544            }
545            let _ = event(
546                EmitMarkEventParams::builder()
547                    .name(&name)
548                    .parent_opt(parent.as_ref())
549                    .data(rejection_data)
550                    .metadata_opt(metadata.clone())
551                    .build(),
552            );
553            return Err(FlowError::GuardrailRejected(error));
554        }
555    }
556
557    let (intercepted_request, annotated_request) =
558        run_request_intercepts_with_codec(&name, request, codec)?;
559
560    let handle = create_llm_handle(
561        CreateLlmHandleParams::builder()
562            .name(name.as_str())
563            .parent_uuid_opt(resolve_parent_uuid(parent.as_ref()))
564            .attributes(attributes)
565            .data_opt(data.clone())
566            .metadata_opt(metadata.clone())
567            .model_name_opt(model_name)
568            .build(),
569    )?;
570    emit_llm_start(&handle, &intercepted_request, annotated_request.clone())?;
571
572    let execution = {
573        let scope_stack = current_scope_stack();
574        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
575        let scope_locals = scope_guard
576            .collect_scope_local_registries(|registries| &registries.llm_execution_intercepts);
577        let context = global_context();
578        let state = context
579            .read()
580            .map_err(|error| FlowError::Internal(error.to_string()))?;
581        state.llm_build_execution_chain(&name, func, &scope_locals)
582    };
583
584    match execution(intercepted_request).await {
585        Ok(response) => {
586            let annotated_response = response_codec
587                .as_ref()
588                .and_then(|codec| codec.decode_response(&response).ok())
589                .map(Arc::new);
590            llm_call_end(
591                LlmCallEndParams::builder()
592                    .handle(&handle)
593                    .response(response.clone())
594                    .data_opt(data)
595                    .metadata_opt(metadata)
596                    .annotated_response_opt(annotated_response)
597                    .build(),
598            )?;
599            Ok(response)
600        }
601        Err(error) => {
602            let _ = emit_llm_end_without_output(&handle, metadata);
603            Err(error)
604        }
605    }
606}
607
608/// Execute a streaming LLM call through the managed middleware pipeline.
609///
610/// This runs the same pre-execution middleware as [`llm_call_execute`], emits
611/// the LLM-start event, and then wraps the provider stream so chunk callbacks
612/// and finalization can emit a single LLM-end event when streaming completes.
613///
614/// # Parameters
615/// - `name`: Logical provider or model family name recorded on emitted events.
616/// - `request`: Raw [`LlmRequest`] passed into the managed pipeline.
617/// - `func`: Streaming provider callback or execution continuation.
618/// - `collector`: Per-chunk collector callback used to accumulate stream state.
619/// - `finalizer`: Finalizer callback used to construct the completed response.
620/// - `parent`: Optional explicit parent scope for the emitted LLM span.
621/// - `attributes`: LLM attribute bitflags applied to the managed span.
622/// - `data`: Optional application payload stored on the managed LLM handle. It
623///   may be used on failure end events that have no output payload.
624/// - `metadata`: Optional JSON metadata recorded on emitted events.
625/// - `model_name`: Optional normalized model name for observability output.
626/// - `codec`: Optional request codec used to produce annotated request data for
627///   intercepts and events.
628/// - `response_codec`: Optional response codec used to attach annotated
629///   response data to the end event.
630///
631/// # Returns
632/// A [`Result`] containing a boxed stream of JSON chunks.
633///
634/// # Errors
635/// Returns [`FlowError::GuardrailRejected`] when conditional-execution
636/// guardrails block the call, or any error raised by request intercepts,
637/// execution intercepts, stream callbacks, codecs, or the provider callback.
638///
639/// # Notes
640/// The LLM-start event is emitted before stream execution intercepts run.
641///
642/// The returned stream emits chunk-level results while the runtime defers the
643/// LLM-end event until the collector and finalizer complete.
644pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Result<LlmJsonStream> {
645    let LlmStreamCallExecuteParams {
646        name,
647        request,
648        func,
649        collector,
650        finalizer,
651        parent,
652        attributes,
653        data,
654        metadata,
655        model_name,
656        codec,
657        response_codec,
658    } = params;
659    ensure_runtime_owner()?;
660    {
661        let (entries, subscribers, parent_uuid, guardrail_metadata) = {
662            let scope_stack = current_scope_stack();
663            let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
664            let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
665                &registries.llm_conditional_execution_guardrails
666            });
667            let scope_subscribers = scope_guard.collect_scope_local_subscribers();
668            let context = global_context();
669            let state = context
670                .read()
671                .map_err(|error| FlowError::Internal(error.to_string()))?;
672            let entries = state.llm_conditional_execution_entries(&scope_locals);
673            let subscribers = state.collect_event_subscribers(&scope_subscribers);
674            (
675                entries,
676                subscribers,
677                resolve_parent_uuid(parent.as_ref()),
678                metadata.clone(),
679            )
680        };
681        if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
682            &request,
683            &entries,
684            &subscribers,
685            parent_uuid,
686            guardrail_metadata,
687        )? {
688            let mut rejection_data = json!({});
689            if let Some(object) = rejection_data.as_object_mut() {
690                object.insert("rejected".into(), json!(true));
691                object.insert("rejection_reason".into(), json!(&error));
692            }
693            let _ = event(
694                EmitMarkEventParams::builder()
695                    .name(&name)
696                    .parent_opt(parent.as_ref())
697                    .data(rejection_data)
698                    .metadata_opt(metadata.clone())
699                    .build(),
700            );
701            return Err(FlowError::GuardrailRejected(error));
702        }
703    }
704
705    let (intercepted_request, annotated_request) =
706        run_request_intercepts_with_codec(&name, request, codec)?;
707
708    let handle = create_llm_handle(
709        CreateLlmHandleParams::builder()
710            .name(name.as_str())
711            .parent_uuid_opt(resolve_parent_uuid(parent.as_ref()))
712            .attributes(attributes)
713            .data_opt(data.clone())
714            .metadata_opt(metadata.clone())
715            .model_name_opt(model_name)
716            .build(),
717    )?;
718    emit_llm_start(&handle, &intercepted_request, annotated_request)?;
719
720    let execution = {
721        let scope_stack = current_scope_stack();
722        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
723        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
724            &registries.llm_stream_execution_intercepts
725        });
726        let context = global_context();
727        let state = context
728            .read()
729            .map_err(|error| FlowError::Internal(error.to_string()))?;
730        state.llm_stream_build_execution_chain(&name, func, &scope_locals)
731    };
732
733    match execution(intercepted_request).await {
734        Ok(raw_stream) => {
735            let wrapper = LlmStreamWrapper::new(
736                raw_stream,
737                handle,
738                collector,
739                finalizer,
740                data,
741                metadata,
742                response_codec,
743            );
744            Ok(Box::pin(wrapper) as LlmJsonStream)
745        }
746        Err(error) => {
747            let _ = emit_llm_end_without_output(&handle, metadata);
748            Err(error)
749        }
750    }
751}
752
753/// Run only the LLM request-intercept chain.
754///
755/// This applies the currently active global and scope-local request intercepts
756/// without emitting lifecycle events or invoking provider execution.
757///
758/// # Parameters
759/// - `name`: Logical provider or model family name used when resolving the
760///   intercept chain.
761/// - `request`: Raw [`LlmRequest`] to transform.
762///
763/// # Returns
764/// A [`Result`] containing the transformed [`LlmRequest`].
765///
766/// # Errors
767/// Returns any error raised by the request-intercept chain.
768///
769/// # Notes
770/// Conditional guardrails, codecs, and execution intercepts are not run by
771/// this helper.
772pub fn llm_request_intercepts(name: &str, request: LlmRequest) -> Result<LlmRequest> {
773    ensure_runtime_owner()?;
774    let scope_stack = current_scope_stack();
775    let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
776    let scope_locals =
777        scope_guard.collect_scope_local_registries(|registries| &registries.llm_request_intercepts);
778    let context = global_context();
779    let state = context
780        .read()
781        .map_err(|error| FlowError::Internal(error.to_string()))?;
782    let (request, _) = state.llm_request_intercepts_chain(name, request, None, &scope_locals)?;
783    Ok(request)
784}
785
786/// Run only the LLM conditional-execution guardrail chain.
787///
788/// This evaluates whether an LLM call should be allowed to proceed without
789/// invoking request intercepts or execution. Each evaluated guardrail emits an
790/// automatic guardrail scope start/end pair for observability.
791///
792/// # Parameters
793/// - `request`: Raw [`LlmRequest`] to validate.
794///
795/// # Returns
796/// A [`Result`] that is `Ok(())` when all guardrails allow execution.
797///
798/// # Errors
799/// Returns [`FlowError::GuardrailRejected`] when a guardrail blocks execution,
800/// or any error raised by the guardrail chain itself.
801///
802/// # Notes
803/// This helper is useful for preflight checks when the caller needs the
804/// rejection result without starting an LLM span. Guardrail scopes are still
805/// emitted for the conditional checks themselves.
806pub fn llm_conditional_execution(request: &LlmRequest) -> Result<()> {
807    ensure_runtime_owner()?;
808    let (entries, subscribers, parent_uuid) = {
809        let scope_stack = current_scope_stack();
810        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
811        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
812            &registries.llm_conditional_execution_guardrails
813        });
814        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
815        let context = global_context();
816        let state = context
817            .read()
818            .map_err(|error| FlowError::Internal(error.to_string()))?;
819        let entries = state.llm_conditional_execution_entries(&scope_locals);
820        let subscribers = state.collect_event_subscribers(&scope_subscribers);
821        (entries, subscribers, resolve_parent_uuid(None))
822    };
823    if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
824        request,
825        &entries,
826        &subscribers,
827        parent_uuid,
828        None,
829    )? {
830        return Err(FlowError::GuardrailRejected(error));
831    }
832    Ok(())
833}
nemo_relay/api/llm.rs

nemo_relay/api/
llm.rs
