nemo-flow 0.2.0

Core Rust SDK for NeMo Flow observability, scope management, and runtime instrumentation.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::sync::Arc;

use uuid::Uuid;

use crate::api::llm::LlmRequest;
use crate::api::runtime::EventSubscriberFn;
use crate::api::runtime::global_context;
use crate::api::runtime::{current_scope_stack, task_scope_top};
use crate::api::scope::ScopeHandle;
use crate::codec::request::AnnotatedLlmRequest;
use crate::codec::traits::LlmCodec;
use crate::error::{FlowError, Result};
use crate::shared_runtime::ensure_process_runtime_owner;

pub(crate) fn resolve_parent_uuid(parent: Option<&ScopeHandle>) -> Option<Uuid> {
    Some(
        parent
            .map(|handle| handle.uuid)
            .unwrap_or_else(|| task_scope_top().uuid),
    )
}

pub(crate) fn snapshot_event_subscribers(
    scope_local_subscribers: Vec<EventSubscriberFn>,
) -> Result<Vec<EventSubscriberFn>> {
    let context = global_context();
    let state = context
        .read()
        .map_err(|error| FlowError::Internal(error.to_string()))?;
    Ok(state.collect_event_subscribers(&scope_local_subscribers))
}

pub(crate) fn ensure_runtime_owner() -> Result<()> {
    ensure_process_runtime_owner()
}

pub(crate) fn run_request_intercepts_with_codec(
    name: &str,
    request: LlmRequest,
    codec: Option<Arc<dyn LlmCodec>>,
) -> Result<(LlmRequest, Option<Arc<AnnotatedLlmRequest>>)> {
    let scope_stack = current_scope_stack();
    let scope_guard = scope_stack.read().expect("scope stack lock poisoned");
    let scope_locals =
        scope_guard.collect_scope_local_registries(|registries| &registries.llm_request_intercepts);

    let context = global_context();
    let state = context
        .read()
        .map_err(|error| FlowError::Internal(error.to_string()))?;

    let original = request.clone();
    let annotated = match &codec {
        Some(codec) => Some(codec.decode(&request)?),
        None => None,
    };

    let (intercepted_request, intercepted_annotated) =
        state.llm_request_intercepts_chain(name, request, annotated, &scope_locals)?;

    match (codec, intercepted_annotated) {
        (Some(codec), Some(annotated)) => {
            let mut encoded = codec.encode(&annotated, &original)?;
            encoded.headers = intercepted_request.headers;
            Ok((encoded, Some(Arc::new(annotated))))
        }
        _ => Ok((intercepted_request, None)),
    }
}

#[cfg(test)]
#[path = "../../tests/unit/shared_tests.rs"]
mod tests;