nemo_flow/api/runtime/
callbacks.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Callback type aliases used by the runtime middleware pipeline.
5//!
6//! The public middleware registration APIs accept boxed or shared closures with
7//! the signatures defined in this module. These aliases centralize those
8//! signatures so the runtime can compose tool and LLM middleware consistently
9//! across bindings.
10
11use std::future::Future;
12use std::pin::Pin;
13use std::sync::Arc;
14
15use tokio_stream::Stream;
16
17use crate::api::event::Event;
18use crate::api::llm::LlmRequest;
19use crate::codec::request::AnnotatedLlmRequest;
20use crate::error::Result;
21use crate::json::Json;
22
23/// Sanitize a tool request payload before the runtime records it.
24///
25/// Tool sanitize callbacks are used only for observability payloads. They can
26/// rewrite the JSON arguments recorded on tool-start events without changing
27/// the caller-owned request that is passed to the tool implementation.
28pub type ToolSanitizeFn = Box<dyn Fn(&str, Json) -> Json + Send + Sync>;
29/// Decide whether a tool call is allowed to continue.
30///
31/// The callback receives the tool name and the current argument payload. It can
32/// return `Ok(None)` to allow execution, `Ok(Some(reason))` to reject the call
33/// with a guardrail message, or an error to abort evaluation entirely.
34pub type ToolConditionalFn = Box<dyn Fn(&str, &Json) -> Result<Option<String>> + Send + Sync>;
35/// Rewrite tool arguments before execution.
36///
37/// Tool request intercepts run in priority order and can transform the JSON
38/// payload that is eventually passed into the tool execution callback.
39pub type ToolInterceptFn = Box<dyn Fn(&str, Json) -> Result<Json> + Send + Sync>;
40/// Continuation type invoked by tool execution intercepts.
41///
42/// Execution intercepts receive this callable as their `next` continuation and
43/// can call it with modified arguments, wrap it, or skip it entirely.
44pub type ToolExecutionNextFn =
45    Arc<dyn Fn(Json) -> Pin<Box<dyn Future<Output = Result<Json>> + Send>> + Send + Sync>;
46/// Wrap or replace tool execution.
47///
48/// A tool execution intercept receives the tool name, the current argument
49/// payload, and the continuation representing the rest of the chain.
50pub type ToolExecutionFn = Arc<
51    dyn Fn(&str, Json, ToolExecutionNextFn) -> Pin<Box<dyn Future<Output = Result<Json>> + Send>>
52        + Send
53        + Sync,
54>;
55
56/// Sanitize an LLM request before the runtime records it.
57///
58/// LLM request sanitizers affect the serialized request payload emitted on
59/// start events. They do not mutate the caller-owned [`LlmRequest`] unless a
60/// separate request intercept does so.
61pub type LlmSanitizeRequestFn = Box<dyn Fn(LlmRequest) -> LlmRequest + Send + Sync>;
62/// Sanitize an LLM response before the runtime records it.
63///
64/// These callbacks rewrite the JSON response payload captured on LLM-end
65/// events, which is useful for redaction or payload normalization.
66pub type LlmSanitizeResponseFn = Box<dyn Fn(Json) -> Json + Send + Sync>;
67/// Decide whether an LLM call is allowed to continue.
68///
69/// The callback receives the current [`LlmRequest`] and can allow execution,
70/// reject it with a guardrail reason, or return an error.
71pub type LlmConditionalFn = Box<dyn Fn(&LlmRequest) -> Result<Option<String>> + Send + Sync>;
72/// Rewrite or annotate an LLM request before execution.
73///
74/// Request intercepts can transform the wire request, attach or replace a
75/// normalized [`AnnotatedLlmRequest`], or both.
76pub type LlmRequestInterceptFn = Box<
77    dyn Fn(
78            &str,
79            LlmRequest,
80            Option<AnnotatedLlmRequest>,
81        ) -> Result<(LlmRequest, Option<AnnotatedLlmRequest>)>
82        + Send
83        + Sync,
84>;
85/// Continuation type invoked by non-streaming LLM execution intercepts.
86///
87/// Execution intercepts use this callable to continue the non-streaming LLM
88/// pipeline after applying their own logic.
89pub type LlmExecutionNextFn =
90    Arc<dyn Fn(LlmRequest) -> Pin<Box<dyn Future<Output = Result<Json>> + Send>> + Send + Sync>;
91/// Wrap or replace non-streaming LLM execution.
92///
93/// A non-streaming execution intercept receives the logical provider name, the
94/// current request, and the continuation representing the rest of the chain.
95pub type LlmExecutionFn = Arc<
96    dyn Fn(
97            &str,
98            LlmRequest,
99            LlmExecutionNextFn,
100        ) -> Pin<Box<dyn Future<Output = Result<Json>> + Send>>
101        + Send
102        + Sync,
103>;
104/// Stream of JSON chunks produced by the managed streaming LLM pipeline.
105pub type LlmJsonStream = Pin<Box<dyn Stream<Item = Result<Json>> + Send>>;
106/// Per-chunk collector used by the streaming LLM runtime.
107pub type LlmCollectorFn = Box<dyn FnMut(Json) -> Result<()> + Send>;
108/// Finalizer used to synthesize the aggregate streaming response payload.
109pub type LlmFinalizerFn = Box<dyn FnOnce() -> Json + Send>;
110/// Scope-local registry references passed into streaming execution-chain builders.
111pub type LlmStreamExecutionRegistryRef<'a> = &'a crate::registry::SortedRegistry<
112    crate::api::registry::ExecutionIntercept<LlmStreamExecutionFn>,
113>;
114/// Slice of scope-local streaming execution registries.
115pub type LlmStreamExecutionRegistryRefs<'a> = &'a [LlmStreamExecutionRegistryRef<'a>];
116
117/// Continuation type invoked by streaming LLM execution intercepts.
118///
119/// This callable represents the remainder of the streaming LLM execution chain
120/// and resolves to a stream of JSON response chunks.
121pub type LlmStreamExecutionNextFn = Arc<
122    dyn Fn(LlmRequest) -> Pin<Box<dyn Future<Output = Result<LlmJsonStream>> + Send>> + Send + Sync,
123>;
124/// Wrap or replace streaming LLM execution.
125///
126/// A streaming execution intercept can observe or modify the request before
127/// invoking the continuation, and it can also replace the returned stream.
128pub type LlmStreamExecutionFn = Arc<
129    dyn Fn(
130            &str,
131            LlmRequest,
132            LlmStreamExecutionNextFn,
133        ) -> Pin<Box<dyn Future<Output = Result<LlmJsonStream>> + Send>>
134        + Send
135        + Sync,
136>;
137
138/// Consume runtime lifecycle events after they are emitted.
139///
140/// Event subscribers are invoked for scope, tool, LLM, and mark events after
141/// the runtime has built the final event payload.
142pub type EventSubscriberFn = Arc<dyn Fn(&Event) + Send + Sync>;
nemo_flow/api/runtime/callbacks.rs

nemo_flow/api/runtime/
callbacks.rs
