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, metadata_with_otel_status, resolve_parent_uuid,
23    run_request_intercepts_with_codec, snapshot_event_subscribers,
24};
25use crate::codec::request::AnnotatedLlmRequest;
26use crate::codec::response::{AnnotatedLlmResponse, attach_estimated_cost_for_provider};
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    request_codec: Option<&dyn LlmCodec>,
280) -> Result<()> {
281    ensure_runtime_owner()?;
282    let (event, subscribers) = {
283        let scope_stack = current_scope_stack();
284        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
285        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
286            &registries.llm_sanitize_request_guardrails
287        });
288        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
289        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
290        let context = global_context();
291        let state = context
292            .read()
293            .map_err(|error| FlowError::Internal(error.to_string()))?;
294
295        let sanitized_request = state.llm_sanitize_request_chain(request.clone(), &scope_locals);
296        let annotated_request = match request_codec {
297            Some(codec)
298                if sanitized_request.headers != request.headers
299                    || sanitized_request.content != request.content =>
300            {
301                codec.decode(&sanitized_request).ok().map(Arc::new)
302            }
303            _ => annotated_request,
304        };
305        let input = serde_json::to_value(&sanitized_request).unwrap_or(Json::Null);
306        let event = state.build_llm_start_event(handle, Some(input), annotated_request);
307        (event, subscribers)
308    };
309    NemoRelayContextState::emit_event(&event, &subscribers);
310    Ok(())
311}
312
313/// Start a manual LLM lifecycle span.
314///
315/// This emits an LLM-start event after applying sanitize-request guardrails to
316/// the payload recorded for observability.
317///
318/// # Parameters
319/// - `name`: Logical provider or model family name recorded on the span.
320/// - `request`: Raw [`LlmRequest`] associated with the span.
321/// - `parent`: Optional explicit parent scope.
322/// - `attributes`: LLM attribute bitflags applied to the span.
323/// - `data`: Optional application payload stored on the returned handle. The
324///   emitted start event data is the sanitized `request` payload.
325/// - `metadata`: Optional JSON metadata recorded on the start event.
326/// - `model_name`: Optional normalized model name recorded separately from the
327///   request payload.
328/// - `annotated_request`: Optional normalized request annotation produced by a
329///   codec.
330/// - `timestamp`: Optional timestamp recorded as the handle start time and on
331///   the emitted start event. When `None`, the current UTC time is used.
332///
333/// # Returns
334/// A [`Result`] containing the created [`LlmHandle`].
335///
336/// # Errors
337/// Returns an error when the runtime owner check fails or when internal state
338/// cannot be read safely.
339///
340/// # Notes
341/// Sanitize-request guardrails affect only the emitted start-event payload, not
342/// the caller-owned [`LlmRequest`].
343pub fn llm_call(params: LlmCallParams<'_>) -> Result<LlmHandle> {
344    let handle_params = CreateLlmHandleParams::builder()
345        .name(params.name)
346        .parent_uuid_opt(resolve_parent_uuid(params.parent))
347        .attributes(params.attributes)
348        .data_opt(params.data)
349        .metadata_opt(params.metadata)
350        .model_name_opt(params.model_name)
351        .timestamp_opt(params.timestamp)
352        .build();
353    let handle = create_llm_handle(handle_params)?;
354    emit_llm_start(&handle, params.request, params.annotated_request, None)?;
355    Ok(handle)
356}
357
358#[derive(Clone, Copy)]
359struct LlmCallEndBehavior {
360    response_codec_errors_fatal: bool,
361    attach_estimated_cost: bool,
362}
363
364/// Finish a manual LLM lifecycle span.
365///
366/// This emits an LLM-end event for a handle previously returned by
367/// [`llm_call`].
368///
369/// # Parameters
370/// - `handle`: LLM handle to close.
371/// - `response`: Raw provider response associated with the end event.
372/// - `data`: Optional application payload retained for compatibility. The
373///   emitted end event data is the sanitized `response` unless it sanitizes to
374///   JSON null, in which case this payload is used.
375/// - `metadata`: Optional JSON metadata recorded on the end event.
376/// - `annotated_response`: Optional normalized response annotation produced by
377///   a response codec. When omitted and `response_codec` is supplied, the
378///   annotation is decoded from the sanitized end-event payload.
379/// - `response_codec`: Optional response codec used to produce a normalized
380///   response annotation from the sanitized end-event payload.
381/// - `timestamp`: Optional timestamp recorded on the emitted end event. When
382///   `None`, the runtime uses the current UTC time, or one microsecond after
383///   the handle start time if the current time is not later.
384///
385/// # Returns
386/// A [`Result`] that is `Ok(())` when the end event has been emitted.
387///
388/// # Errors
389/// Returns an error when the runtime owner check fails, internal state cannot be
390/// read safely, or response codec decoding fails.
391///
392/// # Notes
393/// Sanitize-response guardrails affect only the emitted end-event payload, not
394/// the caller-owned `response` value.
395pub fn llm_call_end(params: LlmCallEndParams<'_>) -> Result<()> {
396    llm_call_end_with_behavior(
397        params,
398        LlmCallEndBehavior {
399            response_codec_errors_fatal: true,
400            attach_estimated_cost: false,
401        },
402    )
403}
404
405fn llm_call_end_with_behavior(
406    params: LlmCallEndParams<'_>,
407    behavior: LlmCallEndBehavior,
408) -> Result<()> {
409    let LlmCallEndParams {
410        handle,
411        response,
412        data,
413        metadata,
414        annotated_response,
415        response_codec,
416        timestamp,
417    } = params;
418    ensure_runtime_owner()?;
419    let mut decode_error = None;
420    let (event, subscribers) = {
421        let scope_stack = current_scope_stack();
422        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
423        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
424            &registries.llm_sanitize_response_guardrails
425        });
426        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
427        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
428        let context = global_context();
429        let state = context
430            .read()
431            .map_err(|error| FlowError::Internal(error.to_string()))?;
432
433        let sanitized_response = state.llm_sanitize_response_chain(response, &scope_locals);
434        let data = if sanitized_response.is_null() {
435            data
436        } else {
437            Some(sanitized_response)
438        };
439        let annotated_response = match annotated_response {
440            Some(annotated_response) => Some(annotated_response),
441            None => match (response_codec.as_ref(), data.as_ref()) {
442                (Some(codec), Some(response)) => match codec.decode_response(response) {
443                    Ok(mut decoded) => {
444                        if behavior.attach_estimated_cost {
445                            attach_estimated_cost_for_provider(&mut decoded, Some(&handle.name));
446                        }
447                        Some(Arc::new(decoded))
448                    }
449                    Err(error) => {
450                        decode_error = Some(error);
451                        None
452                    }
453                },
454                _ => None,
455            },
456        };
457
458        let end_metadata = metadata_with_otel_status(metadata, "OK", None);
459        let event = state.build_llm_end_event(
460            EndLlmHandleParams::builder()
461                .handle(handle)
462                .data_opt(data)
463                .metadata_opt(end_metadata)
464                .annotated_response_opt(annotated_response)
465                .timestamp_opt(timestamp)
466                .build(),
467        );
468        (event, subscribers)
469    };
470    NemoRelayContextState::emit_event(&event, &subscribers);
471    if let Some(error) = decode_error
472        && behavior.response_codec_errors_fatal
473    {
474        Err(error)
475    } else {
476        Ok(())
477    }
478}
479
480fn emit_llm_end_without_output(handle: &LlmHandle, metadata: Option<Json>) -> Result<()> {
481    ensure_runtime_owner()?;
482    let (event, subscribers) = {
483        let scope_stack = current_scope_stack();
484        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
485        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
486        let subscribers = snapshot_event_subscribers(scope_subscribers)?;
487        let context = global_context();
488        let state = context
489            .read()
490            .map_err(|error| FlowError::Internal(error.to_string()))?;
491        let event = state.end_llm_handle(handle, handle.data.clone(), metadata, None);
492        (event, subscribers)
493    };
494    NemoRelayContextState::emit_event(&event, &subscribers);
495    Ok(())
496}
497
498/// Execute an LLM call through the managed middleware pipeline.
499///
500/// This runs conditional-execution guardrails, request intercepts, and
501/// sanitize-request guardrails, emits the LLM-start event, then runs execution
502/// intercepts, the provider callback when it is not replaced, and
503/// sanitize-response guardrails in the runtime-defined order.
504///
505/// # Parameters
506/// - `name`: Logical provider or model family name recorded on emitted events.
507/// - `request`: Raw [`LlmRequest`] passed into the managed pipeline.
508/// - `func`: Provider callback or execution continuation.
509/// - `parent`: Optional explicit parent scope for the emitted LLM span.
510/// - `attributes`: LLM attribute bitflags applied to the managed span.
511/// - `data`: Optional application payload stored on the managed LLM handle. It
512///   may be used on failure end events that have no output payload.
513/// - `metadata`: Optional JSON metadata recorded on emitted events.
514/// - `model_name`: Optional normalized model name for observability output.
515/// - `codec`: Optional request codec used to produce annotated request data for
516///   intercepts and events.
517/// - `response_codec`: Optional response codec used to attach annotated
518///   response data to the end event.
519///
520/// # Returns
521/// A [`Result`] containing the raw JSON response returned by the callback or
522/// an execution intercept.
523///
524/// # Errors
525/// Returns [`FlowError::GuardrailRejected`] when conditional-execution
526/// guardrails block the call, or any error raised by request intercepts,
527/// execution intercepts, codecs, or the callback itself.
528///
529/// # Notes
530/// The LLM-start event is emitted before execution intercepts run. When
531/// execution fails after that point, the runtime still emits an LLM-end event
532/// without an output payload.
533///
534/// Response codecs enrich observability output only and do not change the
535/// value returned to the caller.
536pub async fn llm_call_execute(params: LlmCallExecuteParams) -> Result<Json> {
537    let LlmCallExecuteParams {
538        name,
539        request,
540        func,
541        parent,
542        attributes,
543        data,
544        metadata,
545        model_name,
546        codec,
547        response_codec,
548    } = params;
549    ensure_runtime_owner()?;
550    {
551        let (entries, subscribers, parent_uuid, guardrail_metadata) = {
552            let scope_stack = current_scope_stack();
553            let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
554            let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
555                &registries.llm_conditional_execution_guardrails
556            });
557            let scope_subscribers = scope_guard.collect_scope_local_subscribers();
558            let context = global_context();
559            let state = context
560                .read()
561                .map_err(|error| FlowError::Internal(error.to_string()))?;
562            let entries = state.llm_conditional_execution_entries(&scope_locals);
563            let subscribers = state.collect_event_subscribers(&scope_subscribers);
564            (
565                entries,
566                subscribers,
567                resolve_parent_uuid(parent.as_ref()),
568                metadata.clone(),
569            )
570        };
571        if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
572            &request,
573            &entries,
574            &subscribers,
575            parent_uuid,
576            guardrail_metadata,
577        )? {
578            let mut rejection_data = json!({});
579            if let Some(object) = rejection_data.as_object_mut() {
580                object.insert("rejected".into(), json!(true));
581                object.insert("rejection_reason".into(), json!(&error));
582            }
583            let _ = event(
584                EmitMarkEventParams::builder()
585                    .name(&name)
586                    .parent_opt(parent.as_ref())
587                    .data(rejection_data)
588                    .metadata_opt(metadata.clone())
589                    .build(),
590            );
591            return Err(FlowError::GuardrailRejected(error));
592        }
593    }
594
595    let request_codec = codec.clone();
596    let (intercepted_request, annotated_request) =
597        run_request_intercepts_with_codec(&name, request, codec)?;
598
599    let handle = create_llm_handle(
600        CreateLlmHandleParams::builder()
601            .name(name.as_str())
602            .parent_uuid_opt(resolve_parent_uuid(parent.as_ref()))
603            .attributes(attributes)
604            .data_opt(data.clone())
605            .metadata_opt(metadata.clone())
606            .model_name_opt(model_name)
607            .build(),
608    )?;
609    emit_llm_start(
610        &handle,
611        &intercepted_request,
612        annotated_request.clone(),
613        request_codec.as_deref(),
614    )?;
615
616    let execution = {
617        let scope_stack = current_scope_stack();
618        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
619        let scope_locals = scope_guard
620            .collect_scope_local_registries(|registries| &registries.llm_execution_intercepts);
621        let context = global_context();
622        let state = context
623            .read()
624            .map_err(|error| FlowError::Internal(error.to_string()))?;
625        state.llm_build_execution_chain(&name, func, &scope_locals)
626    };
627
628    match execution(intercepted_request).await {
629        Ok(response) => {
630            llm_call_end_with_behavior(
631                LlmCallEndParams::builder()
632                    .handle(&handle)
633                    .response(response.clone())
634                    .data_opt(data)
635                    .metadata_opt(metadata)
636                    .response_codec_opt(response_codec)
637                    .build(),
638                LlmCallEndBehavior {
639                    response_codec_errors_fatal: false,
640                    attach_estimated_cost: true,
641                },
642            )?;
643            Ok(response)
644        }
645        Err(error) => {
646            let end_metadata =
647                metadata_with_otel_status(metadata, "ERROR", Some(error.to_string()));
648            let _ = emit_llm_end_without_output(&handle, end_metadata);
649            Err(error)
650        }
651    }
652}
653
654/// Execute a streaming LLM call through the managed middleware pipeline.
655///
656/// This runs the same pre-execution middleware as [`llm_call_execute`], emits
657/// the LLM-start event, and then wraps the provider stream so chunk callbacks
658/// and finalization can emit a single LLM-end event when streaming completes.
659///
660/// # Parameters
661/// - `name`: Logical provider or model family name recorded on emitted events.
662/// - `request`: Raw [`LlmRequest`] passed into the managed pipeline.
663/// - `func`: Streaming provider callback or execution continuation.
664/// - `collector`: Per-chunk collector callback used to accumulate stream state.
665/// - `finalizer`: Finalizer callback used to construct the completed response.
666/// - `parent`: Optional explicit parent scope for the emitted LLM span.
667/// - `attributes`: LLM attribute bitflags applied to the managed span.
668/// - `data`: Optional application payload stored on the managed LLM handle. It
669///   may be used on failure end events that have no output payload.
670/// - `metadata`: Optional JSON metadata recorded on emitted events.
671/// - `model_name`: Optional normalized model name for observability output.
672/// - `codec`: Optional request codec used to produce annotated request data for
673///   intercepts and events.
674/// - `response_codec`: Optional response codec used to attach annotated
675///   response data to the end event.
676///
677/// # Returns
678/// A [`Result`] containing a boxed stream of JSON chunks.
679///
680/// # Errors
681/// Returns [`FlowError::GuardrailRejected`] when conditional-execution
682/// guardrails block the call, or any error raised by request intercepts,
683/// execution intercepts, stream callbacks, codecs, or the provider callback.
684///
685/// # Notes
686/// The LLM-start event is emitted before stream execution intercepts run.
687///
688/// The returned stream emits chunk-level results while the runtime defers the
689/// LLM-end event until the collector and finalizer complete.
690pub async fn llm_stream_call_execute(params: LlmStreamCallExecuteParams) -> Result<LlmJsonStream> {
691    let LlmStreamCallExecuteParams {
692        name,
693        request,
694        func,
695        collector,
696        finalizer,
697        parent,
698        attributes,
699        data,
700        metadata,
701        model_name,
702        codec,
703        response_codec,
704    } = params;
705    ensure_runtime_owner()?;
706    {
707        let (entries, subscribers, parent_uuid, guardrail_metadata) = {
708            let scope_stack = current_scope_stack();
709            let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
710            let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
711                &registries.llm_conditional_execution_guardrails
712            });
713            let scope_subscribers = scope_guard.collect_scope_local_subscribers();
714            let context = global_context();
715            let state = context
716                .read()
717                .map_err(|error| FlowError::Internal(error.to_string()))?;
718            let entries = state.llm_conditional_execution_entries(&scope_locals);
719            let subscribers = state.collect_event_subscribers(&scope_subscribers);
720            (
721                entries,
722                subscribers,
723                resolve_parent_uuid(parent.as_ref()),
724                metadata.clone(),
725            )
726        };
727        if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
728            &request,
729            &entries,
730            &subscribers,
731            parent_uuid,
732            guardrail_metadata,
733        )? {
734            let mut rejection_data = json!({});
735            if let Some(object) = rejection_data.as_object_mut() {
736                object.insert("rejected".into(), json!(true));
737                object.insert("rejection_reason".into(), json!(&error));
738            }
739            let _ = event(
740                EmitMarkEventParams::builder()
741                    .name(&name)
742                    .parent_opt(parent.as_ref())
743                    .data(rejection_data)
744                    .metadata_opt(metadata.clone())
745                    .build(),
746            );
747            return Err(FlowError::GuardrailRejected(error));
748        }
749    }
750
751    let request_codec = codec.clone();
752    let (intercepted_request, annotated_request) =
753        run_request_intercepts_with_codec(&name, request, codec)?;
754
755    let handle = create_llm_handle(
756        CreateLlmHandleParams::builder()
757            .name(name.as_str())
758            .parent_uuid_opt(resolve_parent_uuid(parent.as_ref()))
759            .attributes(attributes)
760            .data_opt(data.clone())
761            .metadata_opt(metadata.clone())
762            .model_name_opt(model_name)
763            .build(),
764    )?;
765    emit_llm_start(
766        &handle,
767        &intercepted_request,
768        annotated_request,
769        request_codec.as_deref(),
770    )?;
771
772    let execution = {
773        let scope_stack = current_scope_stack();
774        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
775        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
776            &registries.llm_stream_execution_intercepts
777        });
778        let context = global_context();
779        let state = context
780            .read()
781            .map_err(|error| FlowError::Internal(error.to_string()))?;
782        state.llm_stream_build_execution_chain(&name, func, &scope_locals)
783    };
784
785    match execution(intercepted_request).await {
786        Ok(raw_stream) => {
787            let wrapper = LlmStreamWrapper::new(
788                raw_stream,
789                handle,
790                collector,
791                finalizer,
792                data,
793                metadata,
794                response_codec,
795            );
796            Ok(Box::pin(wrapper) as LlmJsonStream)
797        }
798        Err(error) => {
799            let end_metadata =
800                metadata_with_otel_status(metadata, "ERROR", Some(error.to_string()));
801            let _ = emit_llm_end_without_output(&handle, end_metadata);
802            Err(error)
803        }
804    }
805}
806
807/// Run only the LLM request-intercept chain.
808///
809/// This applies the currently active global and scope-local request intercepts
810/// without emitting lifecycle events or invoking provider execution.
811///
812/// # Parameters
813/// - `name`: Logical provider or model family name used when resolving the
814///   intercept chain.
815/// - `request`: Raw [`LlmRequest`] to transform.
816///
817/// # Returns
818/// A [`Result`] containing the transformed [`LlmRequest`].
819///
820/// # Errors
821/// Returns any error raised by the request-intercept chain.
822///
823/// # Notes
824/// Conditional guardrails, codecs, and execution intercepts are not run by
825/// this helper.
826pub fn llm_request_intercepts(name: &str, request: LlmRequest) -> Result<LlmRequest> {
827    ensure_runtime_owner()?;
828    let scope_stack = current_scope_stack();
829    let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
830    let scope_locals =
831        scope_guard.collect_scope_local_registries(|registries| &registries.llm_request_intercepts);
832    let context = global_context();
833    let state = context
834        .read()
835        .map_err(|error| FlowError::Internal(error.to_string()))?;
836    let (request, _) = state.llm_request_intercepts_chain(name, request, None, &scope_locals)?;
837    Ok(request)
838}
839
840/// Run only the LLM conditional-execution guardrail chain.
841///
842/// This evaluates whether an LLM call should be allowed to proceed without
843/// invoking request intercepts or execution. Each evaluated guardrail emits an
844/// automatic guardrail scope start/end pair for observability.
845///
846/// # Parameters
847/// - `request`: Raw [`LlmRequest`] to validate.
848///
849/// # Returns
850/// A [`Result`] that is `Ok(())` when all guardrails allow execution.
851///
852/// # Errors
853/// Returns [`FlowError::GuardrailRejected`] when a guardrail blocks execution,
854/// or any error raised by the guardrail chain itself.
855///
856/// # Notes
857/// This helper is useful for preflight checks when the caller needs the
858/// rejection result without starting an LLM span. Guardrail scopes are still
859/// emitted for the conditional checks themselves.
860pub fn llm_conditional_execution(request: &LlmRequest) -> Result<()> {
861    ensure_runtime_owner()?;
862    let (entries, subscribers, parent_uuid) = {
863        let scope_stack = current_scope_stack();
864        let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
865        let scope_locals = scope_guard.collect_scope_local_registries(|registries| {
866            &registries.llm_conditional_execution_guardrails
867        });
868        let scope_subscribers = scope_guard.collect_scope_local_subscribers();
869        let context = global_context();
870        let state = context
871            .read()
872            .map_err(|error| FlowError::Internal(error.to_string()))?;
873        let entries = state.llm_conditional_execution_entries(&scope_locals);
874        let subscribers = state.collect_event_subscribers(&scope_subscribers);
875        (entries, subscribers, resolve_parent_uuid(None))
876    };
877    if let Some(error) = NemoRelayContextState::llm_conditional_execution_snapshot_chain(
878        request,
879        &entries,
880        &subscribers,
881        parent_uuid,
882        None,
883    )? {
884        return Err(FlowError::GuardrailRejected(error));
885    }
886    Ok(())
887}
888
889#[cfg(test)]
890#[path = "../../tests/unit/llm_api_tests.rs"]
891mod tests;
nemo_relay/api/llm.rs

nemo_relay/api/
llm.rs
