From 8ad75828b4e1d30b4da05121b933cb0578e4e4a1 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Thu, 25 Jun 2026 13:11:44 -0400 Subject: [PATCH 1/2] feat(types): add shared relay DTO crate Signed-off-by: Will Killian --- Cargo.lock | 14 + Cargo.toml | 2 + crates/core/Cargo.toml | 3 +- crates/core/src/api/event.rs | 813 +----------------- crates/core/src/api/llm.rs | 21 +- crates/core/src/api/scope.rs | 73 +- crates/core/src/api/tool.rs | 10 +- crates/core/src/codec/request.rs | 308 +------ crates/core/src/codec/response.rs | 362 +------- crates/core/src/json.rs | 4 +- .../core/src/observability/openinference.rs | 2 +- crates/core/src/observability/otel.rs | 3 +- crates/core/src/plugin.rs | 30 +- .../tests/integration/api_surface_tests.rs | 10 + .../core/tests/unit/codec/response_tests.rs | 3 +- crates/core/tests/unit/types_tests.rs | 6 +- crates/types/Cargo.toml | 26 + crates/types/src/api/event.rs | 798 +++++++++++++++++ crates/types/src/api/llm.rs | 29 + crates/types/src/api/mod.rs | 13 + crates/types/src/api/scope.rs | 79 ++ crates/types/src/api/tool.rs | 16 + crates/types/src/codec/mod.rs | 9 + crates/types/src/codec/request.rs | 310 +++++++ crates/types/src/codec/response.rs | 369 ++++++++ crates/types/src/lib.rs | 21 + crates/types/src/plugin.rs | 35 + crates/types/tests/serialization_tests.rs | 80 ++ 28 files changed, 1849 insertions(+), 1600 deletions(-) create mode 100644 crates/types/Cargo.toml create mode 100644 crates/types/src/api/event.rs create mode 100644 crates/types/src/api/llm.rs create mode 100644 crates/types/src/api/mod.rs create mode 100644 crates/types/src/api/scope.rs create mode 100644 crates/types/src/api/tool.rs create mode 100644 crates/types/src/codec/mod.rs create mode 100644 crates/types/src/codec/request.rs create mode 100644 crates/types/src/codec/response.rs create mode 100644 crates/types/src/lib.rs create mode 100644 crates/types/src/plugin.rs create mode 100644 crates/types/tests/serialization_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 066876da2..b3bcaf1a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1337,6 +1337,7 @@ dependencies = [ "futures-util", "getrandom 0.3.4", "js-sys", + "nemo-relay-types", "object_store", "openinference-semantic-conventions", "opentelemetry", @@ -1484,6 +1485,19 @@ dependencies = [ "uuid", ] +[[package]] +name = "nemo-relay-types" +version = "0.5.0" +dependencies = [ + "bitflags", + "chrono", + "schemars", + "serde", + "serde_json", + "typed-builder", + "uuid", +] + [[package]] name = "nemo-relay-wasm" version = "0.5.0" diff --git a/Cargo.toml b/Cargo.toml index 660593cee..f4ecf9232 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ [workspace] members = [ "crates/core", + "crates/types", "crates/adaptive", "crates/pii-redaction", "crates/cli", @@ -24,6 +25,7 @@ repository = "https://github.com/NVIDIA/NeMo-Relay" [workspace.dependencies] nemo-relay = { version = "0.5.0", path = "crates/core", default-features = false } +nemo-relay-types = { version = "0.5.0", path = "crates/types" } nemo-relay-adaptive = { version = "0.5.0", path = "crates/adaptive" } nemo-relay-pii-redaction = { version = "0.5.0", path = "crates/pii-redaction" } nemo-relay-ffi = { version = "0.5.0", path = "crates/ffi" } diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 5b9ca03c9..aebca07f8 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -30,7 +30,7 @@ atof-streaming = [ "tokio/rt-multi-thread", "tokio/time", ] -schema = ["dep:schemars"] +schema = ["dep:schemars", "nemo-relay-types/schema"] guardrails-remote = [ "dep:reqwest", "dep:rustls", @@ -75,6 +75,7 @@ openinference = [ ] [dependencies] +nemo-relay-types.workspace = true uuid = { workspace = true, features = ["v7", "serde"] } serde = { version = "1", features = ["derive", "rc"] } serde_json = "1" diff --git a/crates/core/src/api/event.rs b/crates/core/src/api/event.rs index e228f42f7..8d0b1e188 100644 --- a/crates/core/src/api/event.rs +++ b/crates/core/src/api/event.rs @@ -1,616 +1,22 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. //! Event types for Agent Trajectory Observability Format (ATOF) runtime events. use std::borrow::Cow; -use std::collections::BTreeMap; -use std::sync::Arc; -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use typed_builder::TypedBuilder; -use uuid::Uuid; +pub use nemo_relay_types::api::event::*; +use nemo_relay_types::api::llm::LlmRequest; +use nemo_relay_types::codec::request::AnnotatedLlmRequest; +use nemo_relay_types::codec::response::AnnotatedLlmResponse; -use crate::api::llm::{LlmAttributes, LlmRequest}; -use crate::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; -use crate::api::tool::ToolAttributes; -use crate::codec::request::AnnotatedLlmRequest; use crate::codec::resolve; -use crate::codec::response::AnnotatedLlmResponse; -use crate::json::Json; -/// Agent Trajectory Observability Format (ATOF) protocol version emitted by this runtime. -pub const ATOF_VERSION: &str = "0.1"; - -/// Identifier for the schema that describes an event's opaque `data` payload. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TypedBuilder)] -#[builder(field_defaults(setter(into)))] -pub struct DataSchema { - /// Schema name. - pub name: String, - /// Schema version. - pub version: String, -} - -/// Semantic category carried by ATOF `category`. +/// Core-only normalized LLM accessors for ATOF events. /// -/// This is intentionally string-backed so consumers can preserve category -/// values from newer producers without failing deserialization. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] -#[serde(transparent)] -pub struct EventCategory(String); - -impl EventCategory { - /// Top-level agent or workflow scope. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `agent`. - pub fn agent() -> Self { - Self("agent".into()) - } - - /// Generic function or application step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `function`. - pub fn function() -> Self { - Self("function".into()) - } - - /// LLM call. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `llm`. - pub fn llm() -> Self { - Self("llm".into()) - } - - /// Tool invocation. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `tool`. - pub fn tool() -> Self { - Self("tool".into()) - } - - /// Retrieval step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `retriever`. - pub fn retriever() -> Self { - Self("retriever".into()) - } - - /// Embedding-generation step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `embedder`. - pub fn embedder() -> Self { - Self("embedder".into()) - } - - /// Result reranking step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `reranker`. - pub fn reranker() -> Self { - Self("reranker".into()) - } - - /// Guardrail or validation step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `guardrail`. - pub fn guardrail() -> Self { - Self("guardrail".into()) - } - - /// Evaluation or scoring step. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `evaluator`. - pub fn evaluator() -> Self { - Self("evaluator".into()) - } - - /// Vendor-defined custom category. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `custom`. - pub fn custom() -> Self { - Self("custom".into()) - } - - /// Unknown or unclassified work. - /// - /// # Returns - /// An [`EventCategory`] with the wire value `unknown`. - pub fn unknown() -> Self { - Self("unknown".into()) - } - - /// Create a category from an arbitrary producer-provided string. - /// - /// # Parameters - /// - `value`: Wire category value to preserve. - /// - /// # Returns - /// An [`EventCategory`] containing `value`. - pub fn new(value: impl Into) -> Self { - Self(value.into()) - } - - /// Return the string form serialized on the wire. - /// - /// # Returns - /// The category value as a string slice. - pub fn as_str(&self) -> &str { - self.0.as_str() - } - - /// Convert this category to the closest legacy scope type for internal - /// adapters that still need span-kind classification. - /// - /// # Returns - /// The closest matching [`ScopeType`], or [`ScopeType::Unknown`] when the - /// category has no legacy equivalent. - pub fn to_scope_type(&self) -> ScopeType { - match self.as_str() { - "agent" => ScopeType::Agent, - "function" => ScopeType::Function, - "tool" => ScopeType::Tool, - "llm" => ScopeType::Llm, - "retriever" => ScopeType::Retriever, - "embedder" => ScopeType::Embedder, - "reranker" => ScopeType::Reranker, - "guardrail" => ScopeType::Guardrail, - "evaluator" => ScopeType::Evaluator, - "custom" => ScopeType::Custom, - _ => ScopeType::Unknown, - } - } -} - -impl From for EventCategory { - fn from(value: ScopeType) -> Self { - match value { - ScopeType::Agent => Self::agent(), - ScopeType::Function => Self::function(), - ScopeType::Tool => Self::tool(), - ScopeType::Llm => Self::llm(), - ScopeType::Retriever => Self::retriever(), - ScopeType::Embedder => Self::embedder(), - ScopeType::Reranker => Self::reranker(), - ScopeType::Guardrail => Self::guardrail(), - ScopeType::Evaluator => Self::evaluator(), - ScopeType::Custom => Self::custom(), - ScopeType::Unknown => Self::unknown(), - } - } -} - -impl From<&EventCategory> for ScopeType { - fn from(value: &EventCategory) -> Self { - value.to_scope_type() - } -} - -/// Agent Trajectory Observability Format (ATOF) lifecycle phase for a scope event. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum ScopeCategory { - /// Scope was entered. - Start, - /// Scope was exited. - End, -} - -/// Category-specific profile data. -/// -/// Unknown wire keys are preserved in `extra`. LLM annotations are serialized -/// under `category_profile` when a codec captures them. -#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TypedBuilder)] -#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] -pub struct CategoryProfile { - /// Normalized model identifier for LLM events. - #[builder(default)] - #[serde(default, skip_serializing_if = "Option::is_none")] - pub model_name: Option, - - /// LLM-provider correlation ID for Tool events. - #[builder(default)] - #[serde(default, skip_serializing_if = "Option::is_none")] - pub tool_call_id: Option, - - /// Vendor subtype required when `category == "custom"`. - #[builder(default)] - #[serde(default, skip_serializing_if = "Option::is_none")] - pub subtype: Option, - - /// Unknown category-profile keys preserved from newer producers. - #[builder(default)] - #[serde(flatten)] - pub extra: BTreeMap, - - /// Normalized request annotation for LLM start events. - #[builder(default)] - #[serde(default, skip_serializing_if = "Option::is_none")] - pub annotated_request: Option>, - - /// Normalized response annotation for LLM end events. - #[builder(default)] - #[serde(default, skip_serializing_if = "Option::is_none")] - pub annotated_response: Option>, -} - -impl CategoryProfile { - /// Return true when the profile has no wire-serialized fields. - /// - /// # Returns - /// `true` when no profile fields would be serialized on the wire. - pub fn is_wire_empty(&self) -> bool { - self.model_name.is_none() - && self.tool_call_id.is_none() - && self.subtype.is_none() - && self.annotated_request.is_none() - && self.annotated_response.is_none() - && self.extra.is_empty() - } -} - -/// Shared event metadata carried by every ATOF event. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] -#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] -pub struct BaseEvent { - /// ATOF protocol version. - #[builder(default = ATOF_VERSION.to_string())] - pub atof_version: String, - /// UUID of the parent scope, if any. - #[builder(default)] - pub parent_uuid: Option, - /// Unique identifier for the event or span. - #[builder(default = Uuid::now_v7())] - pub uuid: Uuid, - /// Event timestamp in UTC. - #[builder(default = Utc::now())] - #[serde(with = "timestamp")] - pub timestamp: DateTime, - /// Human-readable event name. - pub name: String, - /// Application-defined payload. - #[builder(default)] - pub data: Option, - /// Optional schema identifier for `data`. - #[builder(default)] - pub data_schema: Option, - /// Optional tracing/correlation metadata. - #[builder(default)] - pub metadata: Option, -} - -/// ATOF scope lifecycle event. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] -#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] -pub struct ScopeEvent { - /// Shared ATOF envelope. - #[serde(flatten)] - #[builder(setter(skip), default = BaseEvent::builder().name("").build())] - pub base: BaseEvent, - /// Scope lifecycle phase. - pub scope_category: ScopeCategory, - /// Canonical lowercase behavioral flags. - #[builder(default)] - pub attributes: Vec, - /// Semantic category of work. - pub category: EventCategory, - /// Category-specific typed fields. - #[builder(default)] - pub category_profile: Option, -} - -impl ScopeEvent { - /// Construct a scope event from a base envelope and ATOF-specific fields. - /// - /// # Parameters - /// - `base`: Shared ATOF event envelope. - /// - `scope_category`: Lifecycle phase for the scope event. - /// - `attributes`: Scope attributes to canonicalize and attach. - /// - `category`: Semantic event category. - /// - `category_profile`: Optional category-specific profile data. - /// - /// # Returns - /// A [`ScopeEvent`] containing the provided fields. - pub fn new( - base: BaseEvent, - scope_category: ScopeCategory, - attributes: Vec, - category: EventCategory, - category_profile: Option, - ) -> Self { - Self { - base, - scope_category, - attributes: canonicalize_attributes(attributes), - category, - category_profile, - } - } -} - -/// ATOF point-in-time mark event. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] -#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] -pub struct MarkEvent { - /// Shared ATOF envelope. - #[serde(flatten)] - #[builder(setter(skip), default = BaseEvent::builder().name("").build())] - pub base: BaseEvent, - /// Optional semantic category for the checkpoint. - #[builder(default)] - pub category: Option, - /// Optional category-specific typed fields. - #[builder(default)] - pub category_profile: Option, -} - -impl MarkEvent { - /// Construct a mark event from a base envelope and optional category data. - /// - /// # Parameters - /// - `base`: Shared ATOF event envelope. - /// - `category`: Optional semantic event category. - /// - `category_profile`: Optional category-specific profile data. - /// - /// # Returns - /// A [`MarkEvent`] containing the provided fields. - pub fn new( - base: BaseEvent, - category: Option, - category_profile: Option, - ) -> Self { - Self { - base, - category, - category_profile, - } - } -} - -/// Tagged union covering the two ATOF event kinds emitted by the runtime. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "lowercase")] -pub enum Event { - /// Scope lifecycle event. - Scope(ScopeEvent), - /// Point-in-time checkpoint event. - Mark(MarkEvent), -} - -impl Event { - /// Return the ATOF event kind. - /// - /// # Returns - /// `"scope"` for [`Event::Scope`] and `"mark"` for [`Event::Mark`]. - pub fn kind(&self) -> &'static str { - match self { - Self::Scope(_) => "scope", - Self::Mark(_) => "mark", - } - } - - /// Try to return this event as the canonical JSON object delivered by - /// language bindings to subscriber callbacks and ATOF exporters. - pub fn try_to_json_value(&self) -> serde_json::Result { - serde_json::to_value(self) - } - - /// Return this event as the canonical JSON object delivered by language - /// bindings to subscriber callbacks. - pub fn to_json_value(&self) -> Json { - self.try_to_json_value() - .expect("serializing an ATOF event to JSON should not fail") - } - - /// Return this event as canonical JSON. - pub fn to_json_string(&self) -> serde_json::Result { - serde_json::to_string(&self.try_to_json_value()?) - } - - /// Return the lifecycle phase for scope events. - /// - /// # Returns - /// `Some` lifecycle phase for scope events, otherwise `None`. - pub fn scope_category(&self) -> Option { - match self { - Self::Scope(event) => Some(event.scope_category), - Self::Mark(_) => None, - } - } - - /// Return the semantic category if present. - /// - /// # Returns - /// `Some` category for scope events and categorized mark events, otherwise - /// `None`. - pub fn category(&self) -> Option<&EventCategory> { - match self { - Self::Scope(event) => Some(&event.category), - Self::Mark(event) => event.category.as_ref(), - } - } - - /// Return the category-specific profile if present. - /// - /// # Returns - /// `Some` profile when category-specific fields are present. - pub fn category_profile(&self) -> Option<&CategoryProfile> { - match self { - Self::Scope(event) => event.category_profile.as_ref(), - Self::Mark(event) => event.category_profile.as_ref(), - } - } - - /// Return the mutable category-specific profile if present. - /// - /// # Returns - /// `Some` mutable profile when category-specific fields are present. - pub fn category_profile_mut(&mut self) -> Option<&mut CategoryProfile> { - match self { - Self::Scope(event) => event.category_profile.as_mut(), - Self::Mark(event) => event.category_profile.as_mut(), - } - } - - /// Return the parent scope UUID, if the event is nested under a scope. - /// - /// # Returns - /// `Some` parent UUID when the event has a parent scope, otherwise `None`. - pub fn parent_uuid(&self) -> Option { - self.base().parent_uuid - } - - /// Return the unique event or span UUID. - /// - /// # Returns - /// The event UUID. - pub fn uuid(&self) -> Uuid { - self.base().uuid - } - - /// Return the event timestamp. - /// - /// # Returns - /// The UTC event timestamp. - pub fn timestamp(&self) -> &DateTime { - &self.base().timestamp - } - - /// Return the human-readable event name. - /// - /// # Returns - /// The event name. - pub fn name(&self) -> &str { - self.base().name.as_str() - } - - /// Return the optional application payload attached to the event. - /// - /// # Returns - /// `Some` payload when event data is present, otherwise `None`. - pub fn data(&self) -> Option<&Json> { - self.base().data.as_ref() - } - - /// Return the optional data schema. - /// - /// # Returns - /// `Some` schema when the event payload declares one, otherwise `None`. - pub fn data_schema(&self) -> Option<&DataSchema> { - self.base().data_schema.as_ref() - } - - /// Return the optional metadata attached to the event. - /// - /// # Returns - /// `Some` metadata when present, otherwise `None`. - pub fn metadata(&self) -> Option<&Json> { - self.base().metadata.as_ref() - } - - /// Return attributes for scope events. - /// - /// # Returns - /// `Some` attributes for scope events, otherwise `None`. - pub fn attributes(&self) -> Option<&[String]> { - match self { - Self::Scope(event) => Some(event.attributes.as_slice()), - Self::Mark(_) => None, - } - } - - /// Return the semantic scope category for scope events. - /// - /// # Returns - /// `Some` legacy [`ScopeType`] when the event has a category. - pub fn scope_type(&self) -> Option { - self.category().map(EventCategory::to_scope_type) - } - - /// Return the semantic input payload for start events. - /// - /// # Returns - /// `Some` payload for scope-start events with data, otherwise `None`. - pub fn input(&self) -> Option<&Json> { - match self { - Self::Scope(event) if event.scope_category == ScopeCategory::Start => { - event.base.data.as_ref() - } - _ => None, - } - } - - /// Return the semantic output payload for end events. - /// - /// # Returns - /// `Some` payload for scope-end events with data, otherwise `None`. - pub fn output(&self) -> Option<&Json> { - match self { - Self::Scope(event) if event.scope_category == ScopeCategory::End => { - event.base.data.as_ref() - } - _ => None, - } - } - - /// Return the normalized model name for LLM events. - /// - /// # Returns - /// `Some` model name when the event profile includes one. - pub fn model_name(&self) -> Option<&str> { - self.category_profile() - .and_then(|profile| profile.model_name.as_deref()) - } - - /// Return the provider-specific tool-call correlation identifier. - /// - /// # Returns - /// `Some` tool call identifier when the event profile includes one. - pub fn tool_call_id(&self) -> Option<&str> { - self.category_profile() - .and_then(|profile| profile.tool_call_id.as_deref()) - } - - /// Return the runtime-only annotated LLM request. - /// - /// # Returns - /// `Some` annotated request when the event profile includes one. - pub fn annotated_request(&self) -> Option<&Arc> { - self.category_profile() - .and_then(|profile| profile.annotated_request.as_ref()) - } - - /// Return the runtime-only annotated LLM response. - /// - /// # Returns - /// `Some` annotated response when the event profile includes one. - pub fn annotated_response(&self) -> Option<&Arc> { - self.category_profile() - .and_then(|profile| profile.annotated_response.as_ref()) - } - +/// These helpers use built-in codec resolution, so they live in the runtime +/// crate rather than the shared DTO crate. +pub trait EventNormalizationExt { /// Normalized LLM request: the codec annotation when present, otherwise a /// best-effort decode of the start-event input payload. /// @@ -619,7 +25,16 @@ impl Event { /// LLM pipeline; events whose input is a bare payload or a non-LLM shape /// yield `None`. #[must_use] - pub fn normalized_llm_request(&self) -> Option> { + fn normalized_llm_request(&self) -> Option>; + + /// Normalized LLM response: the codec annotation when present, otherwise a + /// best-effort decode of the end-event output payload. + #[must_use] + fn normalized_llm_response(&self) -> Option>; +} + +impl EventNormalizationExt for Event { + fn normalized_llm_request(&self) -> Option> { if let Some(annotated) = self.annotated_request() { return Some(Cow::Borrowed(annotated.as_ref())); } @@ -627,200 +42,10 @@ impl Event { resolve::normalize_request(&request).map(Cow::Owned) } - /// Normalized LLM response: the codec annotation when present, otherwise a - /// best-effort decode of the end-event output payload. - #[must_use] - pub fn normalized_llm_response(&self) -> Option> { + fn normalized_llm_response(&self) -> Option> { if let Some(annotated) = self.annotated_response() { return Some(Cow::Borrowed(annotated.as_ref())); } resolve::normalize_response(self.output()?).map(Cow::Owned) } - - /// Return true for scope-start events. - /// - /// # Returns - /// `true` when the event is a scope-start event. - pub fn is_scope_start(&self) -> bool { - matches!( - self, - Self::Scope(ScopeEvent { - scope_category: ScopeCategory::Start, - .. - }) - ) - } - - /// Return true for scope-end events. - /// - /// # Returns - /// `true` when the event is a scope-end event. - pub fn is_scope_end(&self) -> bool { - matches!( - self, - Self::Scope(ScopeEvent { - scope_category: ScopeCategory::End, - .. - }) - ) - } - - fn base(&self) -> &BaseEvent { - match self { - Self::Scope(event) => &event.base, - Self::Mark(event) => &event.base, - } - } -} - -/// Convert handle bitflags into ATOF attributes. -/// -/// # Parameters -/// - `attributes`: Handle-specific attribute bitflags. -/// -/// # Returns -/// Canonical lowercase ATOF attribute strings for the provided bitflags. -pub fn attributes_from_handle(attributes: HandleAttributes) -> Vec { - match attributes { - HandleAttributes::Scope(attributes) => scope_attributes_to_strings(attributes), - HandleAttributes::Tool(attributes) => tool_attributes_to_strings(attributes), - HandleAttributes::Llm(attributes) => llm_attributes_to_strings(attributes), - } -} - -/// Convert scope bitflags into ATOF attributes. -/// -/// # Parameters -/// - `attributes`: Scope attribute bitflags. -/// -/// # Returns -/// Canonical lowercase ATOF attribute strings for the provided bitflags. -pub fn scope_attributes_to_strings(attributes: ScopeAttributes) -> Vec { - let mut values = Vec::new(); - if attributes.contains(ScopeAttributes::PARALLEL) { - values.push("parallel".to_string()); - } - if attributes.contains(ScopeAttributes::RELOCATABLE) { - values.push("relocatable".to_string()); - } - values -} - -/// Convert tool bitflags into ATOF attributes. -/// -/// # Parameters -/// - `attributes`: Tool attribute bitflags. -/// -/// # Returns -/// Canonical lowercase ATOF attribute strings for the provided bitflags. -pub fn tool_attributes_to_strings(attributes: ToolAttributes) -> Vec { - let mut values = Vec::new(); - if attributes.contains(ToolAttributes::REMOTE) { - values.push("remote".to_string()); - } - values -} - -/// Convert LLM bitflags into ATOF attributes. -/// -/// # Parameters -/// - `attributes`: LLM attribute bitflags. -/// -/// # Returns -/// Canonical lowercase ATOF attribute strings for the provided bitflags. -pub fn llm_attributes_to_strings(attributes: LlmAttributes) -> Vec { - let mut values = Vec::new(); - if attributes.contains(LlmAttributes::STATEFUL) { - values.push("stateful".to_string()); - } - if attributes.contains(LlmAttributes::STREAMING) { - values.push("streaming".to_string()); - } - values -} - -fn canonicalize_attributes(mut attributes: Vec) -> Vec { - attributes.sort(); - attributes.dedup(); - attributes -} - -mod timestamp { - use chrono::{DateTime, Utc}; - use serde::{ - Deserializer, Serializer, - de::{self, Visitor}, - }; - use std::fmt; - - /// Serialize a UTC timestamp as RFC 3339. - /// - /// # Parameters - /// - `value`: Timestamp to serialize. - /// - `serializer`: Serde serializer receiving the string value. - /// - /// # Returns - /// The serializer's success value. - /// - /// # Errors - /// Returns any error produced by the serializer. - pub fn serialize(value: &DateTime, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_str(&value.to_rfc3339()) - } - - /// Deserialize a UTC timestamp from an RFC 3339 string. - /// - /// # Parameters - /// - `deserializer`: Serde deserializer providing the timestamp value. - /// - /// # Returns - /// Parsed UTC timestamp. - /// - /// # Errors - /// Returns a serde error when the input is not a valid RFC 3339 timestamp. - pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> - where - D: Deserializer<'de>, - { - deserializer.deserialize_any(TimestampVisitor) - } - - struct TimestampVisitor; - - impl<'de> Visitor<'de> for TimestampVisitor { - type Value = DateTime; - - fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("an RFC 3339 timestamp string or epoch microseconds integer") - } - - fn visit_str(self, value: &str) -> Result - where - E: de::Error, - { - DateTime::parse_from_rfc3339(value) - .map(|timestamp| timestamp.with_timezone(&Utc)) - .map_err(E::custom) - } - - fn visit_i64(self, value: i64) -> Result - where - E: de::Error, - { - DateTime::::from_timestamp_micros(value) - .ok_or_else(|| E::custom("epoch microseconds value is out of range")) - } - - fn visit_u64(self, value: u64) -> Result - where - E: de::Error, - { - let value = i64::try_from(value) - .map_err(|_| E::custom("epoch microseconds value is out of range"))?; - self.visit_i64(value) - } - } } diff --git a/crates/core/src/api/llm.rs b/crates/core/src/api/llm.rs index 5647b9cb1..0af799d04 100644 --- a/crates/core/src/api/llm.rs +++ b/crates/core/src/api/llm.rs @@ -3,7 +3,6 @@ use std::sync::Arc; -use bitflags::bitflags; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -29,16 +28,7 @@ use crate::error::{FlowError, Result}; use crate::json::Json; use crate::stream::LlmStreamWrapper; -bitflags! { - /// Bitflags that modify LLM-call behavior and observability. - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] - pub struct LlmAttributes: u32 { - /// Marks the request as stateful from the runtime's perspective. - const STATEFUL = 0b01; - /// Marks the request as streaming. - const STREAMING = 0b10; - } -} +pub use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest}; /// Runtime-owned handle identifying an active or completed LLM call. #[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)] @@ -70,15 +60,6 @@ pub struct LlmHandle { pub model_name: Option, } -/// JSON-shaped LLM request payload passed through the runtime. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct LlmRequest { - /// Provider-specific request headers. - pub headers: serde_json::Map, - /// Provider-specific request body. - pub content: Json, -} - /// Builder parameters for [`NemoRelayContextState::create_llm_handle`]. #[derive(Debug, Clone, TypedBuilder)] #[builder(field_defaults(setter(strip_option(ignore_invalid, fallback_suffix = "_opt"))))] diff --git a/crates/core/src/api/scope.rs b/crates/core/src/api/scope.rs index 31d60d4c6..3a93c8fa1 100644 --- a/crates/core/src/api/scope.rs +++ b/crates/core/src/api/scope.rs @@ -10,83 +10,12 @@ use crate::api::runtime::{ use crate::api::shared::{ensure_runtime_owner, resolve_parent_uuid, snapshot_event_subscribers}; use crate::error::{FlowError, Result}; use crate::json::Json; -use bitflags::bitflags; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use typed_builder::TypedBuilder; use uuid::Uuid; -use crate::api::llm::LlmAttributes; -use crate::api::tool::ToolAttributes; - -bitflags! { - /// Bitflags that modify scope behavior and observability. - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] - pub struct ScopeAttributes: u32 { - /// Marks the scope as running in parallel with sibling work. - const PARALLEL = 0b01; - /// Marks the scope as safe to move across execution contexts. - const RELOCATABLE = 0b10; - } -} - -/// Semantic category attached to a scope lifecycle span. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum ScopeType { - /// A top-level agent or workflow scope. - Agent, - /// A generic function or application step. - Function, - /// A tool lifecycle scope. - Tool, - /// An LLM lifecycle scope. - Llm, - /// A retrieval step such as document search. - Retriever, - /// An embedding generation step. - Embedder, - /// A reranking step. - Reranker, - /// A guardrail or validation step. - Guardrail, - /// An evaluation or scoring step. - Evaluator, - /// A caller-defined custom scope category. - Custom, - /// A fallback for unknown or unsupported scope categories. - Unknown, -} - -impl ScopeType { - /// Return the stable lowercase string form used for encoded scope types. - pub const fn as_str(self) -> &'static str { - match self { - Self::Agent => "agent", - Self::Function => "function", - Self::Tool => "tool", - Self::Llm => "llm", - Self::Retriever => "retriever", - Self::Embedder => "embedder", - Self::Reranker => "reranker", - Self::Guardrail => "guardrail", - Self::Evaluator => "evaluator", - Self::Custom => "custom", - Self::Unknown => "unknown", - } - } -} - -/// Attribute bitflags attached to a concrete handle kind. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum HandleAttributes { - /// Scope-specific attributes. - Scope(ScopeAttributes), - /// Tool-specific attributes. - Tool(ToolAttributes), - /// LLM-specific attributes. - Llm(LlmAttributes), -} +pub use nemo_relay_types::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; /// Runtime-owned handle identifying an active or completed scope. #[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)] diff --git a/crates/core/src/api/tool.rs b/crates/core/src/api/tool.rs index 810380518..8e6512893 100644 --- a/crates/core/src/api/tool.rs +++ b/crates/core/src/api/tool.rs @@ -15,20 +15,12 @@ use crate::api::shared::{ }; use crate::error::{FlowError, Result}; use crate::json::Json; -use bitflags::bitflags; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use typed_builder::TypedBuilder; use uuid::Uuid; -bitflags! { - /// Bitflags that modify tool-call behavior and observability. - #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] - pub struct ToolAttributes: u32 { - /// Marks the tool as executing out-of-process. - const REMOTE = 0b01; - } -} +pub use nemo_relay_types::api::tool::ToolAttributes; /// Runtime-owned handle identifying an active or completed tool call. #[derive(Debug, Clone, Serialize, Deserialize, TypedBuilder)] diff --git a/crates/core/src/codec/request.rs b/crates/core/src/codec/request.rs index e1a499e0a..f2a896f80 100644 --- a/crates/core/src/codec/request.rs +++ b/crates/core/src/codec/request.rs @@ -1,313 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! LLM request codec types and trait. -//! -//! This module defines the [`AnnotatedLlmRequest`] type system for structured -//! LLM request representation and the [`crate::codec::traits::LlmCodec`] trait -//! for bidirectional translation between opaque [`crate::api::llm::LlmRequest`] -//! payloads and typed form. +//! Re-exported normalized LLM request data types. -use serde::{Deserialize, Serialize}; - -use crate::json::Json; - -// --------------------------------------------------------------------------- -// AnnotatedLlmRequest type hierarchy -// --------------------------------------------------------------------------- - -/// Structured view of an LLM request, produced by a Codec from opaque -/// [`LlmRequest`](crate::api::llm::LlmRequest) content. -/// -/// The `extra` field captures any provider-specific keys not modeled by the -/// known fields, ensuring lossless round-trip through `decode`/`encode`. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct AnnotatedLlmRequest { - /// Parsed conversation messages. - pub messages: Vec, - /// Model identifier (e.g., `"gpt-4"`, `"claude-sonnet-4-20250514"`). - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, - /// Common generation parameters, normalized. - #[serde(skip_serializing_if = "Option::is_none")] - pub params: Option, - /// Tool definitions (function schemas) available to the model. - #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, - /// Tool choice control. - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_choice: Option, - /// OpenAI Responses: whether to persist response state server-side. - #[serde(skip_serializing_if = "Option::is_none")] - pub store: Option, - /// OpenAI Responses: prior response to continue from. - #[serde(skip_serializing_if = "Option::is_none")] - pub previous_response_id: Option, - /// OpenAI Responses: context truncation behavior. - #[serde(skip_serializing_if = "Option::is_none")] - pub truncation: Option, - /// OpenAI Responses: reasoning configuration object. - #[serde(skip_serializing_if = "Option::is_none")] - pub reasoning: Option, - /// OpenAI Responses: include filter for additional output/state items. - #[serde(skip_serializing_if = "Option::is_none")] - pub include: Option, - /// OpenAI user identifier. - #[serde(skip_serializing_if = "Option::is_none")] - pub user: Option, - /// OpenAI metadata map/object. - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, - /// OpenAI service tier preference. - #[serde(skip_serializing_if = "Option::is_none")] - pub service_tier: Option, - /// OpenAI tool parallelism toggle. - #[serde(skip_serializing_if = "Option::is_none")] - pub parallel_tool_calls: Option, - /// OpenAI Responses max output token limit. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_output_tokens: Option, - /// OpenAI Responses max tool calls. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_tool_calls: Option, - /// OpenAI logprob fanout count. - #[serde(skip_serializing_if = "Option::is_none")] - pub top_logprobs: Option, - /// OpenAI streaming toggle. - #[serde(skip_serializing_if = "Option::is_none")] - pub stream: Option, - /// Extensible key-value pairs for unmodeled provider-specific fields. - /// Merged back into the request body during encode via `serde(flatten)`. - #[serde(flatten)] - pub extra: serde_json::Map, -} - -/// A single message in a conversation, tagged by role. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(tag = "role", rename_all = "lowercase")] -pub enum Message { - /// A system instruction message. - System { - /// The message content. - content: MessageContent, - /// Optional sender name. - #[serde(skip_serializing_if = "Option::is_none")] - name: Option, - }, - /// A user message. - User { - /// The message content. - content: MessageContent, - /// Optional sender name. - #[serde(skip_serializing_if = "Option::is_none")] - name: Option, - }, - /// An assistant response, optionally containing tool calls. - Assistant { - /// The message content (optional — may be absent when tool calls are present). - #[serde(skip_serializing_if = "Option::is_none")] - content: Option, - /// Tool calls requested by the assistant. - #[serde(skip_serializing_if = "Option::is_none")] - tool_calls: Option>, - /// Optional sender name. - #[serde(skip_serializing_if = "Option::is_none")] - name: Option, - }, - /// A tool result message. - Tool { - /// The tool execution result. - content: MessageContent, - /// The ID of the tool call this result corresponds to. - tool_call_id: String, - }, -} - -/// Message content: either a plain string or multimodal parts array. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(untagged)] -pub enum MessageContent { - /// Plain text content. - Text(String), - /// Multimodal content parts. - Parts(Vec), -} - -/// A single content part within a multimodal message. -/// -/// v1 supports text only. Future versions may add `ImageUrl`, `Audio`, etc. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -pub enum ContentPart { - /// A text content part. - Text { - /// The text content. - text: String, - }, - /// An image URL content part. - ImageUrl { - /// Image URL payload. - image_url: OpenAiImageUrl, - }, -} - -/// OpenAI image URL payload. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct OpenAiImageUrl { - /// URL for the image. - pub url: String, - /// Optional provider-specific detail hint. - #[serde(skip_serializing_if = "Option::is_none")] - pub detail: Option, -} - -/// A tool call requested by the assistant. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolCall { - /// Unique identifier for this tool call. - pub id: String, - /// The type of tool call (typically `"function"`). - #[serde(rename = "type")] - pub call_type: String, - /// The function to call. - pub function: FunctionCall, -} - -/// A function call within a tool call. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct FunctionCall { - /// The name of the function to call. - pub name: String, - /// The function arguments as a JSON string (per OpenAI convention). - pub arguments: String, -} - -/// A tool definition (function schema) available to the model. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolDefinition { - /// The type of tool (typically `"function"`). - #[serde(rename = "type")] - pub tool_type: String, - /// The function definition. - pub function: FunctionDefinition, -} - -/// A function definition within a tool definition. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct FunctionDefinition { - /// The name of the function. - pub name: String, - /// A description of what the function does. - #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, - /// The JSON Schema for the function parameters. - #[serde(skip_serializing_if = "Option::is_none")] - pub parameters: Option, -} - -/// Tool choice control: how the model should use available tools. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum ToolChoice { - /// Let the model decide whether to call a tool. - Auto, - /// Do not call any tools. - None, - /// The model must call at least one tool. - Required, - /// Force a specific function by name. - #[serde(untagged)] - Specific(ToolChoiceFunction), -} - -/// A specific tool choice that forces a named function. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolChoiceFunction { - /// The type (typically `"function"`). - #[serde(rename = "type")] - pub choice_type: String, - /// The function to call. - pub function: ToolChoiceFunctionName, -} - -/// The name component of a specific tool choice. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ToolChoiceFunctionName { - /// The function name. - pub name: String, -} - -/// Normalized generation parameters across providers. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] -pub struct GenerationParams { - /// Sampling temperature. - #[serde(skip_serializing_if = "Option::is_none")] - pub temperature: Option, - /// Maximum number of tokens to generate. - #[serde(skip_serializing_if = "Option::is_none")] - pub max_tokens: Option, - /// Nucleus sampling probability. - #[serde(skip_serializing_if = "Option::is_none")] - pub top_p: Option, - /// Stop sequences. - #[serde(skip_serializing_if = "Option::is_none")] - pub stop: Option>, -} - -// --------------------------------------------------------------------------- -// Helper methods -// --------------------------------------------------------------------------- - -impl AnnotatedLlmRequest { - /// Extract the text content of the first system message, if any. - /// - /// For [`MessageContent::Text`], returns the string directly. - /// For [`MessageContent::Parts`], returns the text of the first - /// [`ContentPart::Text`] part. - pub fn system_prompt(&self) -> Option<&str> { - self.messages.iter().find_map(|m| match m { - Message::System { content, .. } => match content { - MessageContent::Text(s) => Some(s.as_str()), - MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, - }), - }, - _ => None, - }) - } - - /// Get the text content of the last user message, if any. - /// - /// Searches messages in reverse order and returns the first user - /// message found. For [`MessageContent::Parts`], returns the text of - /// the first [`ContentPart::Text`] part. - pub fn last_user_message(&self) -> Option<&str> { - self.messages.iter().rev().find_map(|m| match m { - Message::User { content, .. } => match content { - MessageContent::Text(s) => Some(s.as_str()), - MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - ContentPart::Text { text } => Some(text.as_str()), - ContentPart::ImageUrl { .. } => None, - }), - }, - _ => None, - }) - } - - /// Check if any assistant message in the conversation contains tool calls. - /// - /// Returns `true` if at least one [`Message::Assistant`] variant has a - /// non-empty `tool_calls` field. - pub fn has_tool_calls(&self) -> bool { - self.messages.iter().any(|m| { - matches!( - m, - Message::Assistant { tool_calls: Some(calls), .. } if !calls.is_empty() - ) - }) - } -} +pub use nemo_relay_types::codec::request::*; #[cfg(test)] #[path = "../../tests/unit/codec/request_tests.rs"] diff --git a/crates/core/src/codec/response.rs b/crates/core/src/codec/response.rs index 8769d3904..f00e6ea93 100644 --- a/crates/core/src/codec/response.rs +++ b/crates/core/src/codec/response.rs @@ -1,14 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Normalized LLM response types produced by response codecs. -//! -//! This module defines [`AnnotatedLlmResponse`] and its supporting types -//! for structured, API-agnostic access to LLM response data. +//! Re-exported normalized LLM response data types and core pricing helpers. -use serde::{Deserialize, Serialize}; +use serde::Deserialize; -use crate::json::Json; +pub use nemo_relay_types::codec::response::*; pub use super::pricing::{ CacheReadAccounting, ModelPricing, PricingCatalog, PricingCatalogError, PricingConfig, @@ -19,169 +16,6 @@ pub use super::pricing::{ pricing_for_model, pricing_for_provider, reset_active_pricing_resolver, set_active_pricing_resolver, }; -use super::request::MessageContent; - -// --------------------------------------------------------------------------- -// AnnotatedLlmResponse type hierarchy -// --------------------------------------------------------------------------- - -/// Structured view of an LLM response, produced by a response codec from -/// raw JSON API output. -/// -/// The `extra` field captures any top-level keys not modeled by the known -/// fields, ensuring lossless round-trip through serde. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct AnnotatedLlmResponse { - /// Response ID from the API (e.g., "chatcmpl-abc123", "resp_abc123", "msg_abc123"). - #[serde(skip_serializing_if = "Option::is_none")] - pub id: Option, - - /// The model that actually served the request (may differ from requested model). - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, - - /// The assistant's response content, reusing [`MessageContent`] from request types. - #[serde(skip_serializing_if = "Option::is_none")] - pub message: Option, - - /// Tool calls requested by the model, normalized across APIs. - /// - /// Uses [`ResponseToolCall`] (arguments as [`Json`]) NOT the request-side - /// `ToolCall` (arguments as `String`). - #[serde(skip_serializing_if = "Option::is_none")] - pub tool_calls: Option>, - - /// Why generation stopped, normalized across APIs. - #[serde(skip_serializing_if = "Option::is_none")] - pub finish_reason: Option, - - /// Token usage statistics. - #[serde(skip_serializing_if = "Option::is_none")] - pub usage: Option, - - /// API-specific response data that cannot be normalized across providers. - #[serde(skip_serializing_if = "Option::is_none")] - pub api_specific: Option, - - /// Catch-all for unmodeled top-level fields, ensuring lossless round-trip. - #[serde(flatten)] - pub extra: serde_json::Map, -} - -// --------------------------------------------------------------------------- -// Usage -// --------------------------------------------------------------------------- - -/// Token usage statistics from an LLM API response. -/// -/// All fields are `Option` because not every provider supplies every -/// field. For example, cache token counts are only available from providers -/// that support prompt caching. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] -pub struct Usage { - /// Tokens consumed by the prompt/input. - #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_tokens: Option, - /// Tokens generated in the completion/output. - #[serde(skip_serializing_if = "Option::is_none")] - pub completion_tokens: Option, - /// Total tokens (prompt + completion). - #[serde(skip_serializing_if = "Option::is_none")] - pub total_tokens: Option, - /// Tokens served from prompt cache (read). - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read_tokens: Option, - /// Tokens written to prompt cache. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_write_tokens: Option, - /// Optional cost reported by provider data or estimated from Relay pricing. - #[serde(skip_serializing_if = "Option::is_none")] - pub cost: Option, -} - -/// Source of a normalized cost value. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum CostSource { - /// Cost was estimated by applying Relay's model pricing table to usage. - ModelPricing, - /// Cost was reported directly by a provider or framework payload. - ProviderReported, -} - -/// Normalized LLM response cost. -/// -/// Provider-reported cost is preserved as-is. Model-pricing estimates include -/// source and as-of metadata so downstream systems can audit stale pricing -/// tables without losing a usable estimate. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct CostEstimate { - /// Total cost in `currency`. - #[serde(skip_serializing_if = "Option::is_none")] - pub total: Option, - /// ISO 4217 currency code for the cost fields. - #[serde(default = "default_cost_currency")] - pub currency: String, - /// Uncached prompt/input token cost in `currency`. - #[serde(skip_serializing_if = "Option::is_none")] - pub input: Option, - /// Completion/output token cost in `currency`. - #[serde(skip_serializing_if = "Option::is_none")] - pub output: Option, - /// Prompt cache read cost in `currency`. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_read: Option, - /// Prompt cache write cost in `currency`. - #[serde(skip_serializing_if = "Option::is_none")] - pub cache_write: Option, - /// Origin of this cost value. - pub source: CostSource, - /// Provider associated with the cost or pricing estimate, if known. - #[serde(skip_serializing_if = "Option::is_none")] - pub pricing_provider: Option, - /// Model ID associated with the cost or pricing estimate, if known. - #[serde(skip_serializing_if = "Option::is_none")] - pub pricing_model: Option, - /// Date the pricing value was last verified, if known. - #[serde(skip_serializing_if = "Option::is_none")] - pub pricing_as_of: Option, - /// Source URL or label for the pricing value, if known. - #[serde(skip_serializing_if = "Option::is_none")] - pub pricing_source: Option, -} - -impl CostEstimate { - /// Returns the explicit total, or the sum of component costs when no total was supplied. - #[must_use] - pub fn total_or_component_sum(&self) -> Option { - self.total.or_else(|| { - let (has_component, total) = - [self.input, self.output, self.cache_read, self.cache_write] - .into_iter() - .flatten() - .fold((false, 0.0), |(_, total), value| (true, total + value)); - has_component.then_some(total) - }) - } - - /// Returns the total only when it is denominated in the requested currency. - #[must_use] - pub fn total_for_currency(&self, currency: &str) -> Option { - self.currency - .eq_ignore_ascii_case(currency) - .then_some(self.total) - .flatten() - } - - /// Returns the explicit or component-derived total in the requested currency. - #[must_use] - pub fn total_or_component_sum_for_currency(&self, currency: &str) -> Option { - self.currency - .eq_ignore_ascii_case(currency) - .then(|| self.total_or_component_sum()) - .flatten() - } -} /// Provider/framework cost object accepted by built-in response codecs. #[derive(Debug, Clone, Default, Deserialize)] @@ -272,196 +106,6 @@ fn default_cost_currency() -> String { "USD".into() } -// --------------------------------------------------------------------------- -// FinishReason -// --------------------------------------------------------------------------- - -/// Normalized reason why the model stopped generating. -/// -/// Maps from provider-specific stop reasons: -/// - **Complete**: OpenAI Chat `"stop"`, Anthropic `"end_turn"`, Responses `"completed"` -/// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens -/// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"` -/// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter -/// - **Unknown**: Forward-compatible catch-all for unrecognized reasons -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum FinishReason { - /// Model naturally completed its response. - Complete, - /// Maximum token limit reached. - Length, - /// Model requested a tool call. - ToolUse, - /// Content was filtered by safety systems. - ContentFilter, - /// Unknown or forward-compatible reason. - Unknown(String), -} - -impl FinishReason { - /// Returns `true` if the model naturally completed its response. - /// - /// Only the [`FinishReason::Complete`] variant returns `true`. - #[must_use] - pub fn is_complete(&self) -> bool { - matches!(self, FinishReason::Complete) - } -} - -// --------------------------------------------------------------------------- -// ResponseToolCall -// --------------------------------------------------------------------------- - -/// A tool call requested by the model in its response. -/// -/// Unlike the request-side `ToolCall` (which stores arguments as a JSON -/// string per OpenAI convention), response tool calls store arguments as -/// parsed [`Json`]. Codecs parse OpenAI's string arguments during decode; -/// Anthropic's `input` is already parsed JSON. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct ResponseToolCall { - /// Unique identifier for this tool call. - pub id: String, - /// The function/tool name. - pub name: String, - /// The arguments as parsed JSON (not a string). - pub arguments: Json, -} - -// --------------------------------------------------------------------------- -// ApiSpecificResponse -// --------------------------------------------------------------------------- - -/// API-specific response data that cannot be normalized across providers. -/// -/// Each variant captures fields unique to a particular LLM API, stored via -/// internal tagging on the `"api"` key. -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -#[serde(tag = "api")] -pub enum ApiSpecificResponse { - /// OpenAI Chat Completions-specific fields. - #[serde(rename = "openai_chat")] - OpenAIChat { - /// Token-level log probabilities (raw JSON, too complex to normalize). - #[serde(skip_serializing_if = "Option::is_none")] - logprobs: Option, - /// System fingerprint for reproducibility. - #[serde(skip_serializing_if = "Option::is_none")] - system_fingerprint: Option, - /// Processing tier used (e.g., "default"). - #[serde(skip_serializing_if = "Option::is_none")] - service_tier: Option, - }, - - /// OpenAI Responses API-specific fields. - #[serde(rename = "openai_responses")] - OpenAIResponses { - /// Full output items array for direct access. - #[serde(skip_serializing_if = "Option::is_none")] - output_items: Option>, - /// Response status (e.g., "completed", "incomplete"). - #[serde(skip_serializing_if = "Option::is_none")] - status: Option, - /// Details about why the response is incomplete. - #[serde(skip_serializing_if = "Option::is_none")] - incomplete_details: Option, - /// Echoed previous response ID for conversation continuation. - #[serde(skip_serializing_if = "Option::is_none")] - previous_response_id: Option, - /// Whether this response is marked for server-side storage. - #[serde(skip_serializing_if = "Option::is_none")] - store: Option, - /// Service tier used for the response. - #[serde(skip_serializing_if = "Option::is_none")] - service_tier: Option, - /// Truncation behavior metadata. - #[serde(skip_serializing_if = "Option::is_none")] - truncation: Option, - /// Reasoning configuration/result metadata. - #[serde(skip_serializing_if = "Option::is_none")] - reasoning: Option, - /// Raw input token details payload. - #[serde(skip_serializing_if = "Option::is_none")] - input_tokens_details: Option, - /// Raw output token details payload. - #[serde(skip_serializing_if = "Option::is_none")] - output_tokens_details: Option, - }, - - /// Anthropic Messages API-specific fields. - #[serde(rename = "anthropic_messages")] - AnthropicMessages { - /// Anthropic object type (typically `"message"`). - #[serde(skip_serializing_if = "Option::is_none")] - object_type: Option, - /// Anthropic response role (typically `"assistant"`). - #[serde(skip_serializing_if = "Option::is_none")] - role: Option, - /// Raw Anthropic stop_reason. - #[serde(skip_serializing_if = "Option::is_none")] - stop_reason: Option, - /// Which stop sequence was matched (if any). - #[serde(skip_serializing_if = "Option::is_none")] - stop_sequence: Option, - /// Anthropic response service tier when present. - #[serde(skip_serializing_if = "Option::is_none")] - service_tier: Option, - /// Anthropic container payload when present. - #[serde(skip_serializing_if = "Option::is_none")] - container: Option, - /// Full content blocks array for direct access. - #[serde(skip_serializing_if = "Option::is_none")] - content_blocks: Option>, - }, - - /// Custom/unknown API -- catch-all for user-implemented codecs. - #[serde(rename = "custom")] - Custom { - /// API identifier. - api_name: String, - /// Opaque API-specific data. - data: Json, - }, -} - -// --------------------------------------------------------------------------- -// Helper methods -// --------------------------------------------------------------------------- - -impl AnnotatedLlmResponse { - /// Extract the text content of the response message. - /// - /// For [`MessageContent::Text`], returns the string directly. - /// For [`MessageContent::Parts`], returns the text of the first - /// [`super::request::ContentPart::Text`] part. - /// Returns `None` if `message` is `None`. - #[must_use] - pub fn response_text(&self) -> Option<&str> { - match self.message.as_ref()? { - MessageContent::Text(s) => Some(s.as_str()), - MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { - super::request::ContentPart::Text { text } => Some(text.as_str()), - super::request::ContentPart::ImageUrl { .. } => None, - }), - } - } - - /// Check if the response contains any tool calls. - /// - /// Returns `true` if `tool_calls` is `Some` with at least one element. - #[must_use] - pub fn has_tool_calls(&self) -> bool { - self.tool_calls - .as_ref() - .is_some_and(|calls| !calls.is_empty()) - } -} - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - #[cfg(test)] #[path = "../../tests/unit/codec/response_tests.rs"] mod tests; diff --git a/crates/core/src/json.rs b/crates/core/src/json.rs index ec42a1f1c..3a76630da 100644 --- a/crates/core/src/json.rs +++ b/crates/core/src/json.rs @@ -7,9 +7,7 @@ //! throughout the crate, and a [`merge_json`] helper for shallow-merging //! optional JSON values. -/// Type alias for [`serde_json::Value`], used as the universal JSON -/// representation throughout the NeMo Relay runtime. -pub type Json = serde_json::Value; +pub use nemo_relay_types::Json; /// Shallow-merge two optional JSON values. /// diff --git a/crates/core/src/observability/openinference.rs b/crates/core/src/observability/openinference.rs index c694166bc..1e7abf3ff 100644 --- a/crates/core/src/observability/openinference.rs +++ b/crates/core/src/observability/openinference.rs @@ -21,7 +21,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use super::{estimate_cost_for_response_or_requested_model, manual}; -use crate::api::event::{Event, ScopeCategory}; +use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; diff --git a/crates/core/src/observability/otel.rs b/crates/core/src/observability/otel.rs index d85a7f48a..612fb1af1 100644 --- a/crates/core/src/observability/otel.rs +++ b/crates/core/src/observability/otel.rs @@ -21,8 +21,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use super::{estimate_cost_for_response_or_requested_model, manual}; -use crate::api::event::Event; -use crate::api::event::ScopeCategory; +use crate::api::event::{Event, EventNormalizationExt, ScopeCategory}; use crate::api::runtime::EventSubscriberFn; use crate::api::scope::ScopeType; use crate::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; diff --git a/crates/core/src/plugin.rs b/crates/core/src/plugin.rs index 6ddbd8195..94420f6df 100644 --- a/crates/core/src/plugin.rs +++ b/crates/core/src/plugin.rs @@ -38,6 +38,7 @@ use crate::api::runtime::{ ToolExecutionFn, ToolInterceptFn, ToolSanitizeFn, }; use crate::api::subscriber::{deregister_subscriber, register_subscriber}; +pub use nemo_relay_types::plugin::{ConfigDiagnostic, DiagnosticLevel}; pub mod dynamic; pub use dynamic::*; @@ -151,35 +152,6 @@ impl ConfigReport { } } -/// One validation or compatibility diagnostic. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -pub struct ConfigDiagnostic { - /// Severity level for the diagnostic. - pub level: DiagnosticLevel, - /// Stable diagnostic code suitable for machine checks. - pub code: String, - /// Optional component identifier associated with the diagnostic. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub component: Option, - /// Optional field path associated with the diagnostic. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub field: Option, - /// Human-readable diagnostic message. - pub message: String, -} - -/// Diagnostic severity. -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[serde(rename_all = "lowercase")] -pub enum DiagnosticLevel { - /// Non-fatal compatibility or validation issue. - Warning, - /// Fatal validation issue that blocks initialization. - Error, -} - /// Policy for how unsupported plugin/runtime config is handled. #[derive(Debug, Clone, Copy, Serialize, Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] diff --git a/crates/core/tests/integration/api_surface_tests.rs b/crates/core/tests/integration/api_surface_tests.rs index 23cd5a361..e751c15b3 100644 --- a/crates/core/tests/integration/api_surface_tests.rs +++ b/crates/core/tests/integration/api_surface_tests.rs @@ -119,6 +119,16 @@ fn expect_not_found(error: FlowError, needle: &str) { } } +#[test] +fn shared_type_reexports_keep_existing_core_paths() { + let request: LlmRequest = make_llm_request(json!({ "prompt": "hello" })); + let _shared_request: nemo_relay_types::api::llm::LlmRequest = request; + let scope_type: ScopeType = nemo_relay_types::api::scope::ScopeType::Agent; + assert_eq!(scope_type.as_str(), "agent"); + let attributes: ToolAttributes = nemo_relay_types::api::tool::ToolAttributes::REMOTE; + assert!(attributes.contains(ToolAttributes::REMOTE)); +} + #[test] fn test_manual_lifecycle_timestamp_overrides() { let _lock = TEST_MUTEX.lock().unwrap(); diff --git a/crates/core/tests/unit/codec/response_tests.rs b/crates/core/tests/unit/codec/response_tests.rs index c51ed246c..396196bec 100644 --- a/crates/core/tests/unit/codec/response_tests.rs +++ b/crates/core/tests/unit/codec/response_tests.rs @@ -8,10 +8,11 @@ use serde_json::{Value, json}; use std::fs; use std::time::{SystemTime, UNIX_EPOCH}; -use super::super::request::ContentPart; +use super::super::request::{ContentPart, MessageContent}; use super::super::traits::LlmResponseCodec; use crate::codec::pricing::pricing_test_mutex; use crate::error::FlowError; +use crate::json::Json; use crate::plugin::{ DiagnosticLevel, PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins, validate_plugin_config, diff --git a/crates/core/tests/unit/types_tests.rs b/crates/core/tests/unit/types_tests.rs index 46b6e4b75..0ce2ba22d 100644 --- a/crates/core/tests/unit/types_tests.rs +++ b/crates/core/tests/unit/types_tests.rs @@ -10,9 +10,9 @@ use serde_json::{Map, json}; use uuid::{Uuid, Version}; use crate::api::event::{ - BaseEvent, CategoryProfile, DataSchema, Event, EventCategory, MarkEvent, ScopeCategory, - ScopeEvent, attributes_from_handle, llm_attributes_to_strings, scope_attributes_to_strings, - tool_attributes_to_strings, + BaseEvent, CategoryProfile, DataSchema, Event, EventCategory, EventNormalizationExt, MarkEvent, + ScopeCategory, ScopeEvent, attributes_from_handle, llm_attributes_to_strings, + scope_attributes_to_strings, tool_attributes_to_strings, }; use crate::api::llm::{LlmAttributes, LlmHandle, LlmRequest}; use crate::api::scope::{HandleAttributes, ScopeAttributes, ScopeHandle, ScopeType}; diff --git a/crates/types/Cargo.toml b/crates/types/Cargo.toml new file mode 100644 index 000000000..b967cf972 --- /dev/null +++ b/crates/types/Cargo.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "nemo-relay-types" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "Shared serializable data model types for NeMo Relay." + +[lints] +workspace = true + +[features] +default = [] +schema = ["dep:schemars"] + +[dependencies] +bitflags = { version = "2", features = ["serde"] } +chrono = { version = "0.4", features = ["serde"] } +schemars = { version = "0.8", optional = true } +serde = { version = "1", features = ["derive", "rc"] } +serde_json = "1" +typed-builder = "0.23.2" +uuid = { workspace = true, features = ["v7", "serde"] } diff --git a/crates/types/src/api/event.rs b/crates/types/src/api/event.rs new file mode 100644 index 000000000..af17546ae --- /dev/null +++ b/crates/types/src/api/event.rs @@ -0,0 +1,798 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Event types for Agent Trajectory Observability Format (ATOF) runtime events. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use typed_builder::TypedBuilder; +use uuid::Uuid; + +use crate::Json; +use crate::api::llm::LlmAttributes; +use crate::api::scope::{HandleAttributes, ScopeAttributes, ScopeType}; +use crate::api::tool::ToolAttributes; +use crate::codec::request::AnnotatedLlmRequest; +use crate::codec::response::AnnotatedLlmResponse; + +/// Agent Trajectory Observability Format (ATOF) protocol version emitted by this runtime. +pub const ATOF_VERSION: &str = "0.1"; + +/// Identifier for the schema that describes an event's opaque `data` payload. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into)))] +pub struct DataSchema { + /// Schema name. + pub name: String, + /// Schema version. + pub version: String, +} + +/// Semantic category carried by ATOF `category`. +/// +/// This is intentionally string-backed so consumers can preserve category +/// values from newer producers without failing deserialization. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct EventCategory(String); + +impl EventCategory { + /// Top-level agent or workflow scope. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `agent`. + pub fn agent() -> Self { + Self("agent".into()) + } + + /// Generic function or application step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `function`. + pub fn function() -> Self { + Self("function".into()) + } + + /// LLM call. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `llm`. + pub fn llm() -> Self { + Self("llm".into()) + } + + /// Tool invocation. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `tool`. + pub fn tool() -> Self { + Self("tool".into()) + } + + /// Retrieval step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `retriever`. + pub fn retriever() -> Self { + Self("retriever".into()) + } + + /// Embedding-generation step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `embedder`. + pub fn embedder() -> Self { + Self("embedder".into()) + } + + /// Result reranking step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `reranker`. + pub fn reranker() -> Self { + Self("reranker".into()) + } + + /// Guardrail or validation step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `guardrail`. + pub fn guardrail() -> Self { + Self("guardrail".into()) + } + + /// Evaluation or scoring step. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `evaluator`. + pub fn evaluator() -> Self { + Self("evaluator".into()) + } + + /// Vendor-defined custom category. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `custom`. + pub fn custom() -> Self { + Self("custom".into()) + } + + /// Unknown or unclassified work. + /// + /// # Returns + /// An [`EventCategory`] with the wire value `unknown`. + pub fn unknown() -> Self { + Self("unknown".into()) + } + + /// Create a category from an arbitrary producer-provided string. + /// + /// # Parameters + /// - `value`: Wire category value to preserve. + /// + /// # Returns + /// An [`EventCategory`] containing `value`. + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + /// Return the string form serialized on the wire. + /// + /// # Returns + /// The category value as a string slice. + pub fn as_str(&self) -> &str { + self.0.as_str() + } + + /// Convert this category to the closest legacy scope type for internal + /// adapters that still need span-kind classification. + /// + /// # Returns + /// The closest matching [`ScopeType`], or [`ScopeType::Unknown`] when the + /// category has no legacy equivalent. + pub fn to_scope_type(&self) -> ScopeType { + match self.as_str() { + "agent" => ScopeType::Agent, + "function" => ScopeType::Function, + "tool" => ScopeType::Tool, + "llm" => ScopeType::Llm, + "retriever" => ScopeType::Retriever, + "embedder" => ScopeType::Embedder, + "reranker" => ScopeType::Reranker, + "guardrail" => ScopeType::Guardrail, + "evaluator" => ScopeType::Evaluator, + "custom" => ScopeType::Custom, + _ => ScopeType::Unknown, + } + } +} + +impl From for EventCategory { + fn from(value: ScopeType) -> Self { + match value { + ScopeType::Agent => Self::agent(), + ScopeType::Function => Self::function(), + ScopeType::Tool => Self::tool(), + ScopeType::Llm => Self::llm(), + ScopeType::Retriever => Self::retriever(), + ScopeType::Embedder => Self::embedder(), + ScopeType::Reranker => Self::reranker(), + ScopeType::Guardrail => Self::guardrail(), + ScopeType::Evaluator => Self::evaluator(), + ScopeType::Custom => Self::custom(), + ScopeType::Unknown => Self::unknown(), + } + } +} + +impl From<&EventCategory> for ScopeType { + fn from(value: &EventCategory) -> Self { + value.to_scope_type() + } +} + +/// Agent Trajectory Observability Format (ATOF) lifecycle phase for a scope event. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ScopeCategory { + /// Scope was entered. + Start, + /// Scope was exited. + End, +} + +/// Category-specific profile data. +/// +/// Unknown wire keys are preserved in `extra`. LLM annotations are serialized +/// under `category_profile` when a codec captures them. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct CategoryProfile { + /// Normalized model identifier for LLM events. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model_name: Option, + + /// LLM-provider correlation ID for Tool events. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, + + /// Vendor subtype required when `category == "custom"`. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subtype: Option, + + /// Unknown category-profile keys preserved from newer producers. + #[builder(default)] + #[serde(flatten)] + pub extra: BTreeMap, + + /// Normalized request annotation for LLM start events. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub annotated_request: Option>, + + /// Normalized response annotation for LLM end events. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub annotated_response: Option>, +} + +impl CategoryProfile { + /// Return true when the profile has no wire-serialized fields. + /// + /// # Returns + /// `true` when no profile fields would be serialized on the wire. + pub fn is_wire_empty(&self) -> bool { + self.model_name.is_none() + && self.tool_call_id.is_none() + && self.subtype.is_none() + && self.annotated_request.is_none() + && self.annotated_response.is_none() + && self.extra.is_empty() + } +} + +/// Shared event metadata carried by every ATOF event. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct BaseEvent { + /// ATOF protocol version. + #[builder(default = ATOF_VERSION.to_string())] + pub atof_version: String, + /// UUID of the parent scope, if any. + #[builder(default)] + pub parent_uuid: Option, + /// Unique identifier for the event or span. + #[builder(default = Uuid::now_v7())] + pub uuid: Uuid, + /// Event timestamp in UTC. + #[builder(default = Utc::now())] + #[serde(with = "timestamp")] + pub timestamp: DateTime, + /// Human-readable event name. + pub name: String, + /// Application-defined payload. + #[builder(default)] + pub data: Option, + /// Optional schema identifier for `data`. + #[builder(default)] + pub data_schema: Option, + /// Optional tracing/correlation metadata. + #[builder(default)] + pub metadata: Option, +} + +/// ATOF scope lifecycle event. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct ScopeEvent { + /// Shared ATOF envelope. + #[serde(flatten)] + #[builder(setter(skip), default = BaseEvent::builder().name("").build())] + pub base: BaseEvent, + /// Scope lifecycle phase. + pub scope_category: ScopeCategory, + /// Canonical lowercase behavioral flags. + #[builder(default)] + pub attributes: Vec, + /// Semantic category of work. + pub category: EventCategory, + /// Category-specific typed fields. + #[builder(default)] + pub category_profile: Option, +} + +impl ScopeEvent { + /// Construct a scope event from a base envelope and ATOF-specific fields. + /// + /// # Parameters + /// - `base`: Shared ATOF event envelope. + /// - `scope_category`: Lifecycle phase for the scope event. + /// - `attributes`: Scope attributes to canonicalize and attach. + /// - `category`: Semantic event category. + /// - `category_profile`: Optional category-specific profile data. + /// + /// # Returns + /// A [`ScopeEvent`] containing the provided fields. + pub fn new( + base: BaseEvent, + scope_category: ScopeCategory, + attributes: Vec, + category: EventCategory, + category_profile: Option, + ) -> Self { + Self { + base, + scope_category, + attributes: canonicalize_attributes(attributes), + category, + category_profile, + } + } +} + +/// ATOF point-in-time mark event. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[builder(field_defaults(setter(into, strip_option(ignore_invalid, fallback_suffix = "_opt"))))] +pub struct MarkEvent { + /// Shared ATOF envelope. + #[serde(flatten)] + #[builder(setter(skip), default = BaseEvent::builder().name("").build())] + pub base: BaseEvent, + /// Optional semantic category for the checkpoint. + #[builder(default)] + pub category: Option, + /// Optional category-specific typed fields. + #[builder(default)] + pub category_profile: Option, +} + +impl MarkEvent { + /// Construct a mark event from a base envelope and optional category data. + /// + /// # Parameters + /// - `base`: Shared ATOF event envelope. + /// - `category`: Optional semantic event category. + /// - `category_profile`: Optional category-specific profile data. + /// + /// # Returns + /// A [`MarkEvent`] containing the provided fields. + pub fn new( + base: BaseEvent, + category: Option, + category_profile: Option, + ) -> Self { + Self { + base, + category, + category_profile, + } + } +} + +/// Tagged union covering the two ATOF event kinds emitted by the runtime. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "lowercase")] +pub enum Event { + /// Scope lifecycle event. + Scope(ScopeEvent), + /// Point-in-time checkpoint event. + Mark(MarkEvent), +} + +impl Event { + /// Return the ATOF event kind. + /// + /// # Returns + /// `"scope"` for [`Event::Scope`] and `"mark"` for [`Event::Mark`]. + pub fn kind(&self) -> &'static str { + match self { + Self::Scope(_) => "scope", + Self::Mark(_) => "mark", + } + } + + /// Try to return this event as the canonical JSON object delivered by + /// language bindings to subscriber callbacks and ATOF exporters. + pub fn try_to_json_value(&self) -> serde_json::Result { + serde_json::to_value(self) + } + + /// Return this event as the canonical JSON object delivered by language + /// bindings to subscriber callbacks. + pub fn to_json_value(&self) -> Json { + self.try_to_json_value() + .expect("serializing an ATOF event to JSON should not fail") + } + + /// Return this event as canonical JSON. + pub fn to_json_string(&self) -> serde_json::Result { + serde_json::to_string(&self.try_to_json_value()?) + } + + /// Return the lifecycle phase for scope events. + /// + /// # Returns + /// `Some` lifecycle phase for scope events, otherwise `None`. + pub fn scope_category(&self) -> Option { + match self { + Self::Scope(event) => Some(event.scope_category), + Self::Mark(_) => None, + } + } + + /// Return the semantic category if present. + /// + /// # Returns + /// `Some` category for scope events and categorized mark events, otherwise + /// `None`. + pub fn category(&self) -> Option<&EventCategory> { + match self { + Self::Scope(event) => Some(&event.category), + Self::Mark(event) => event.category.as_ref(), + } + } + + /// Return the category-specific profile if present. + /// + /// # Returns + /// `Some` profile when category-specific fields are present. + pub fn category_profile(&self) -> Option<&CategoryProfile> { + match self { + Self::Scope(event) => event.category_profile.as_ref(), + Self::Mark(event) => event.category_profile.as_ref(), + } + } + + /// Return the mutable category-specific profile if present. + /// + /// # Returns + /// `Some` mutable profile when category-specific fields are present. + pub fn category_profile_mut(&mut self) -> Option<&mut CategoryProfile> { + match self { + Self::Scope(event) => event.category_profile.as_mut(), + Self::Mark(event) => event.category_profile.as_mut(), + } + } + + /// Return the parent scope UUID, if the event is nested under a scope. + /// + /// # Returns + /// `Some` parent UUID when the event has a parent scope, otherwise `None`. + pub fn parent_uuid(&self) -> Option { + self.base().parent_uuid + } + + /// Return the unique event or span UUID. + /// + /// # Returns + /// The event UUID. + pub fn uuid(&self) -> Uuid { + self.base().uuid + } + + /// Return the event timestamp. + /// + /// # Returns + /// The UTC event timestamp. + pub fn timestamp(&self) -> &DateTime { + &self.base().timestamp + } + + /// Return the human-readable event name. + /// + /// # Returns + /// The event name. + pub fn name(&self) -> &str { + self.base().name.as_str() + } + + /// Return the optional application payload attached to the event. + /// + /// # Returns + /// `Some` payload when event data is present, otherwise `None`. + pub fn data(&self) -> Option<&Json> { + self.base().data.as_ref() + } + + /// Return the optional data schema. + /// + /// # Returns + /// `Some` schema when the event payload declares one, otherwise `None`. + pub fn data_schema(&self) -> Option<&DataSchema> { + self.base().data_schema.as_ref() + } + + /// Return the optional metadata attached to the event. + /// + /// # Returns + /// `Some` metadata when present, otherwise `None`. + pub fn metadata(&self) -> Option<&Json> { + self.base().metadata.as_ref() + } + + /// Return attributes for scope events. + /// + /// # Returns + /// `Some` attributes for scope events, otherwise `None`. + pub fn attributes(&self) -> Option<&[String]> { + match self { + Self::Scope(event) => Some(event.attributes.as_slice()), + Self::Mark(_) => None, + } + } + + /// Return the semantic scope category for scope events. + /// + /// # Returns + /// `Some` legacy [`ScopeType`] when the event has a category. + pub fn scope_type(&self) -> Option { + self.category().map(EventCategory::to_scope_type) + } + + /// Return the semantic input payload for start events. + /// + /// # Returns + /// `Some` payload for scope-start events with data, otherwise `None`. + pub fn input(&self) -> Option<&Json> { + match self { + Self::Scope(event) if event.scope_category == ScopeCategory::Start => { + event.base.data.as_ref() + } + _ => None, + } + } + + /// Return the semantic output payload for end events. + /// + /// # Returns + /// `Some` payload for scope-end events with data, otherwise `None`. + pub fn output(&self) -> Option<&Json> { + match self { + Self::Scope(event) if event.scope_category == ScopeCategory::End => { + event.base.data.as_ref() + } + _ => None, + } + } + + /// Return the normalized model name for LLM events. + /// + /// # Returns + /// `Some` model name when the event profile includes one. + pub fn model_name(&self) -> Option<&str> { + self.category_profile() + .and_then(|profile| profile.model_name.as_deref()) + } + + /// Return the provider-specific tool-call correlation identifier. + /// + /// # Returns + /// `Some` tool call identifier when the event profile includes one. + pub fn tool_call_id(&self) -> Option<&str> { + self.category_profile() + .and_then(|profile| profile.tool_call_id.as_deref()) + } + + /// Return the runtime-only annotated LLM request. + /// + /// # Returns + /// `Some` annotated request when the event profile includes one. + pub fn annotated_request(&self) -> Option<&Arc> { + self.category_profile() + .and_then(|profile| profile.annotated_request.as_ref()) + } + + /// Return the runtime-only annotated LLM response. + /// + /// # Returns + /// `Some` annotated response when the event profile includes one. + pub fn annotated_response(&self) -> Option<&Arc> { + self.category_profile() + .and_then(|profile| profile.annotated_response.as_ref()) + } + + /// Return true for scope-start events. + /// + /// # Returns + /// `true` when the event is a scope-start event. + pub fn is_scope_start(&self) -> bool { + matches!( + self, + Self::Scope(ScopeEvent { + scope_category: ScopeCategory::Start, + .. + }) + ) + } + + /// Return true for scope-end events. + /// + /// # Returns + /// `true` when the event is a scope-end event. + pub fn is_scope_end(&self) -> bool { + matches!( + self, + Self::Scope(ScopeEvent { + scope_category: ScopeCategory::End, + .. + }) + ) + } + + fn base(&self) -> &BaseEvent { + match self { + Self::Scope(event) => &event.base, + Self::Mark(event) => &event.base, + } + } +} + +/// Convert handle bitflags into ATOF attributes. +/// +/// # Parameters +/// - `attributes`: Handle-specific attribute bitflags. +/// +/// # Returns +/// Canonical lowercase ATOF attribute strings for the provided bitflags. +pub fn attributes_from_handle(attributes: HandleAttributes) -> Vec { + match attributes { + HandleAttributes::Scope(attributes) => scope_attributes_to_strings(attributes), + HandleAttributes::Tool(attributes) => tool_attributes_to_strings(attributes), + HandleAttributes::Llm(attributes) => llm_attributes_to_strings(attributes), + } +} + +/// Convert scope bitflags into ATOF attributes. +/// +/// # Parameters +/// - `attributes`: Scope attribute bitflags. +/// +/// # Returns +/// Canonical lowercase ATOF attribute strings for the provided bitflags. +pub fn scope_attributes_to_strings(attributes: ScopeAttributes) -> Vec { + let mut values = Vec::new(); + if attributes.contains(ScopeAttributes::PARALLEL) { + values.push("parallel".to_string()); + } + if attributes.contains(ScopeAttributes::RELOCATABLE) { + values.push("relocatable".to_string()); + } + values +} + +/// Convert tool bitflags into ATOF attributes. +/// +/// # Parameters +/// - `attributes`: Tool attribute bitflags. +/// +/// # Returns +/// Canonical lowercase ATOF attribute strings for the provided bitflags. +pub fn tool_attributes_to_strings(attributes: ToolAttributes) -> Vec { + let mut values = Vec::new(); + if attributes.contains(ToolAttributes::REMOTE) { + values.push("remote".to_string()); + } + values +} + +/// Convert LLM bitflags into ATOF attributes. +/// +/// # Parameters +/// - `attributes`: LLM attribute bitflags. +/// +/// # Returns +/// Canonical lowercase ATOF attribute strings for the provided bitflags. +pub fn llm_attributes_to_strings(attributes: LlmAttributes) -> Vec { + let mut values = Vec::new(); + if attributes.contains(LlmAttributes::STATEFUL) { + values.push("stateful".to_string()); + } + if attributes.contains(LlmAttributes::STREAMING) { + values.push("streaming".to_string()); + } + values +} + +fn canonicalize_attributes(mut attributes: Vec) -> Vec { + attributes.sort(); + attributes.dedup(); + attributes +} + +mod timestamp { + use chrono::{DateTime, Utc}; + use serde::{ + Deserializer, Serializer, + de::{self, Visitor}, + }; + use std::fmt; + + /// Serialize a UTC timestamp as RFC 3339. + /// + /// # Parameters + /// - `value`: Timestamp to serialize. + /// - `serializer`: Serde serializer receiving the string value. + /// + /// # Returns + /// The serializer's success value. + /// + /// # Errors + /// Returns any error produced by the serializer. + pub fn serialize(value: &DateTime, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&value.to_rfc3339()) + } + + /// Deserialize a UTC timestamp from an RFC 3339 string. + /// + /// # Parameters + /// - `deserializer`: Serde deserializer providing the timestamp value. + /// + /// # Returns + /// Parsed UTC timestamp. + /// + /// # Errors + /// Returns a serde error when the input is not a valid RFC 3339 timestamp. + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + deserializer.deserialize_any(TimestampVisitor) + } + + struct TimestampVisitor; + + impl<'de> Visitor<'de> for TimestampVisitor { + type Value = DateTime; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an RFC 3339 timestamp string or epoch microseconds integer") + } + + fn visit_str(self, value: &str) -> Result + where + E: de::Error, + { + DateTime::parse_from_rfc3339(value) + .map(|timestamp| timestamp.with_timezone(&Utc)) + .map_err(E::custom) + } + + fn visit_i64(self, value: i64) -> Result + where + E: de::Error, + { + DateTime::::from_timestamp_micros(value) + .ok_or_else(|| E::custom("epoch microseconds value is out of range")) + } + + fn visit_u64(self, value: u64) -> Result + where + E: de::Error, + { + let value = i64::try_from(value) + .map_err(|_| E::custom("epoch microseconds value is out of range"))?; + self.visit_i64(value) + } + } +} diff --git a/crates/types/src/api/llm.rs b/crates/types/src/api/llm.rs new file mode 100644 index 000000000..caa917a65 --- /dev/null +++ b/crates/types/src/api/llm.rs @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared LLM data types. + +use bitflags::bitflags; +use serde::{Deserialize, Serialize}; + +use crate::Json; + +bitflags! { + /// Bitflags that modify LLM-call behavior and observability. + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] + pub struct LlmAttributes: u32 { + /// Marks the request as stateful from the runtime's perspective. + const STATEFUL = 0b01; + /// Marks the request as streaming. + const STREAMING = 0b10; + } +} + +/// JSON-shaped LLM request payload passed through the runtime. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LlmRequest { + /// Provider-specific request headers. + pub headers: serde_json::Map, + /// Provider-specific request body. + pub content: Json, +} diff --git a/crates/types/src/api/mod.rs b/crates/types/src/api/mod.rs new file mode 100644 index 000000000..2a660243d --- /dev/null +++ b/crates/types/src/api/mod.rs @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared API data types. + +/// Lifecycle event DTOs. +pub mod event; +/// LLM DTOs and attributes. +pub mod llm; +/// Scope DTOs and attributes. +pub mod scope; +/// Tool DTOs and attributes. +pub mod tool; diff --git a/crates/types/src/api/scope.rs b/crates/types/src/api/scope.rs new file mode 100644 index 000000000..349d5b600 --- /dev/null +++ b/crates/types/src/api/scope.rs @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared scope data types. + +use bitflags::bitflags; +use serde::{Deserialize, Serialize}; + +use crate::api::llm::LlmAttributes; +use crate::api::tool::ToolAttributes; + +bitflags! { + /// Bitflags that modify scope behavior and observability. + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] + pub struct ScopeAttributes: u32 { + /// Marks the scope as running in parallel with sibling work. + const PARALLEL = 0b01; + /// Marks the scope as safe to move across execution contexts. + const RELOCATABLE = 0b10; + } +} + +/// Semantic category attached to a scope lifecycle span. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ScopeType { + /// A top-level agent or workflow scope. + Agent, + /// A generic function or application step. + Function, + /// A tool lifecycle scope. + Tool, + /// An LLM lifecycle scope. + Llm, + /// A retrieval step such as document search. + Retriever, + /// An embedding generation step. + Embedder, + /// A reranking step. + Reranker, + /// A guardrail or validation step. + Guardrail, + /// An evaluation or scoring step. + Evaluator, + /// A caller-defined custom scope category. + Custom, + /// A fallback for unknown or unsupported scope categories. + Unknown, +} + +impl ScopeType { + /// Return the stable lowercase string form used for encoded scope types. + pub const fn as_str(self) -> &'static str { + match self { + Self::Agent => "agent", + Self::Function => "function", + Self::Tool => "tool", + Self::Llm => "llm", + Self::Retriever => "retriever", + Self::Embedder => "embedder", + Self::Reranker => "reranker", + Self::Guardrail => "guardrail", + Self::Evaluator => "evaluator", + Self::Custom => "custom", + Self::Unknown => "unknown", + } + } +} + +/// Attribute bitflags attached to a concrete handle kind. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum HandleAttributes { + /// Scope-specific attributes. + Scope(ScopeAttributes), + /// Tool-specific attributes. + Tool(ToolAttributes), + /// LLM-specific attributes. + Llm(LlmAttributes), +} diff --git a/crates/types/src/api/tool.rs b/crates/types/src/api/tool.rs new file mode 100644 index 000000000..db4f59803 --- /dev/null +++ b/crates/types/src/api/tool.rs @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared tool data types. + +use bitflags::bitflags; +use serde::{Deserialize, Serialize}; + +bitflags! { + /// Bitflags that modify tool-call behavior and observability. + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] + pub struct ToolAttributes: u32 { + /// Marks the tool as executing out-of-process. + const REMOTE = 0b01; + } +} diff --git a/crates/types/src/codec/mod.rs b/crates/types/src/codec/mod.rs new file mode 100644 index 000000000..5ff75f032 --- /dev/null +++ b/crates/types/src/codec/mod.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared normalized LLM request and response data types. + +/// Normalized LLM request data types. +pub mod request; +/// Normalized LLM response data types. +pub mod response; diff --git a/crates/types/src/codec/request.rs b/crates/types/src/codec/request.rs new file mode 100644 index 000000000..f54c0d267 --- /dev/null +++ b/crates/types/src/codec/request.rs @@ -0,0 +1,310 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! LLM request codec types and trait. +//! +//! This module defines the [`AnnotatedLlmRequest`] type system for structured +//! LLM request representation and the [`crate::codec::traits::LlmCodec`] trait +//! for bidirectional translation between opaque [`crate::api::llm::LlmRequest`] +//! payloads and typed form. + +use serde::{Deserialize, Serialize}; + +use crate::Json; + +// --------------------------------------------------------------------------- +// AnnotatedLlmRequest type hierarchy +// --------------------------------------------------------------------------- + +/// Structured view of an LLM request, produced by a Codec from opaque +/// [`LlmRequest`](crate::api::llm::LlmRequest) content. +/// +/// The `extra` field captures any provider-specific keys not modeled by the +/// known fields, ensuring lossless round-trip through `decode`/`encode`. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct AnnotatedLlmRequest { + /// Parsed conversation messages. + pub messages: Vec, + /// Model identifier (e.g., `"gpt-4"`, `"claude-sonnet-4-20250514"`). + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + /// Common generation parameters, normalized. + #[serde(skip_serializing_if = "Option::is_none")] + pub params: Option, + /// Tool definitions (function schemas) available to the model. + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, + /// Tool choice control. + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + /// OpenAI Responses: whether to persist response state server-side. + #[serde(skip_serializing_if = "Option::is_none")] + pub store: Option, + /// OpenAI Responses: prior response to continue from. + #[serde(skip_serializing_if = "Option::is_none")] + pub previous_response_id: Option, + /// OpenAI Responses: context truncation behavior. + #[serde(skip_serializing_if = "Option::is_none")] + pub truncation: Option, + /// OpenAI Responses: reasoning configuration object. + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, + /// OpenAI Responses: include filter for additional output/state items. + #[serde(skip_serializing_if = "Option::is_none")] + pub include: Option, + /// OpenAI user identifier. + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, + /// OpenAI metadata map/object. + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option, + /// OpenAI service tier preference. + #[serde(skip_serializing_if = "Option::is_none")] + pub service_tier: Option, + /// OpenAI tool parallelism toggle. + #[serde(skip_serializing_if = "Option::is_none")] + pub parallel_tool_calls: Option, + /// OpenAI Responses max output token limit. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + /// OpenAI Responses max tool calls. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tool_calls: Option, + /// OpenAI logprob fanout count. + #[serde(skip_serializing_if = "Option::is_none")] + pub top_logprobs: Option, + /// OpenAI streaming toggle. + #[serde(skip_serializing_if = "Option::is_none")] + pub stream: Option, + /// Extensible key-value pairs for unmodeled provider-specific fields. + /// Merged back into the request body during encode via `serde(flatten)`. + #[serde(flatten)] + pub extra: serde_json::Map, +} + +/// A single message in a conversation, tagged by role. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "role", rename_all = "lowercase")] +pub enum Message { + /// A system instruction message. + System { + /// The message content. + content: MessageContent, + /// Optional sender name. + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, + /// A user message. + User { + /// The message content. + content: MessageContent, + /// Optional sender name. + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, + /// An assistant response, optionally containing tool calls. + Assistant { + /// The message content (optional — may be absent when tool calls are present). + #[serde(skip_serializing_if = "Option::is_none")] + content: Option, + /// Tool calls requested by the assistant. + #[serde(skip_serializing_if = "Option::is_none")] + tool_calls: Option>, + /// Optional sender name. + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, + /// A tool result message. + Tool { + /// The tool execution result. + content: MessageContent, + /// The ID of the tool call this result corresponds to. + tool_call_id: String, + }, +} + +/// Message content: either a plain string or multimodal parts array. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MessageContent { + /// Plain text content. + Text(String), + /// Multimodal content parts. + Parts(Vec), +} + +/// A single content part within a multimodal message. +/// +/// v1 supports text only. Future versions may add `ImageUrl`, `Audio`, etc. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ContentPart { + /// A text content part. + Text { + /// The text content. + text: String, + }, + /// An image URL content part. + ImageUrl { + /// Image URL payload. + image_url: OpenAiImageUrl, + }, +} + +/// OpenAI image URL payload. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct OpenAiImageUrl { + /// URL for the image. + pub url: String, + /// Optional provider-specific detail hint. + #[serde(skip_serializing_if = "Option::is_none")] + pub detail: Option, +} + +/// A tool call requested by the assistant. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolCall { + /// Unique identifier for this tool call. + pub id: String, + /// The type of tool call (typically `"function"`). + #[serde(rename = "type")] + pub call_type: String, + /// The function to call. + pub function: FunctionCall, +} + +/// A function call within a tool call. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct FunctionCall { + /// The name of the function to call. + pub name: String, + /// The function arguments as a JSON string (per OpenAI convention). + pub arguments: String, +} + +/// A tool definition (function schema) available to the model. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolDefinition { + /// The type of tool (typically `"function"`). + #[serde(rename = "type")] + pub tool_type: String, + /// The function definition. + pub function: FunctionDefinition, +} + +/// A function definition within a tool definition. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct FunctionDefinition { + /// The name of the function. + pub name: String, + /// A description of what the function does. + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + /// The JSON Schema for the function parameters. + #[serde(skip_serializing_if = "Option::is_none")] + pub parameters: Option, +} + +/// Tool choice control: how the model should use available tools. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ToolChoice { + /// Let the model decide whether to call a tool. + Auto, + /// Do not call any tools. + None, + /// The model must call at least one tool. + Required, + /// Force a specific function by name. + #[serde(untagged)] + Specific(ToolChoiceFunction), +} + +/// A specific tool choice that forces a named function. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolChoiceFunction { + /// The type (typically `"function"`). + #[serde(rename = "type")] + pub choice_type: String, + /// The function to call. + pub function: ToolChoiceFunctionName, +} + +/// The name component of a specific tool choice. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ToolChoiceFunctionName { + /// The function name. + pub name: String, +} + +/// Normalized generation parameters across providers. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +pub struct GenerationParams { + /// Sampling temperature. + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + /// Maximum number of tokens to generate. + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + /// Nucleus sampling probability. + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + /// Stop sequences. + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option>, +} + +// --------------------------------------------------------------------------- +// Helper methods +// --------------------------------------------------------------------------- + +impl AnnotatedLlmRequest { + /// Extract the text content of the first system message, if any. + /// + /// For [`MessageContent::Text`], returns the string directly. + /// For [`MessageContent::Parts`], returns the text of the first + /// [`ContentPart::Text`] part. + pub fn system_prompt(&self) -> Option<&str> { + self.messages.iter().find_map(|m| match m { + Message::System { content, .. } => match content { + MessageContent::Text(s) => Some(s.as_str()), + MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { + ContentPart::Text { text } => Some(text.as_str()), + ContentPart::ImageUrl { .. } => None, + }), + }, + _ => None, + }) + } + + /// Get the text content of the last user message, if any. + /// + /// Searches messages in reverse order and returns the first user + /// message found. For [`MessageContent::Parts`], returns the text of + /// the first [`ContentPart::Text`] part. + pub fn last_user_message(&self) -> Option<&str> { + self.messages.iter().rev().find_map(|m| match m { + Message::User { content, .. } => match content { + MessageContent::Text(s) => Some(s.as_str()), + MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { + ContentPart::Text { text } => Some(text.as_str()), + ContentPart::ImageUrl { .. } => None, + }), + }, + _ => None, + }) + } + + /// Check if any assistant message in the conversation contains tool calls. + /// + /// Returns `true` if at least one [`Message::Assistant`] variant has a + /// non-empty `tool_calls` field. + pub fn has_tool_calls(&self) -> bool { + self.messages.iter().any(|m| { + matches!( + m, + Message::Assistant { tool_calls: Some(calls), .. } if !calls.is_empty() + ) + }) + } +} diff --git a/crates/types/src/codec/response.rs b/crates/types/src/codec/response.rs new file mode 100644 index 000000000..12beafc5c --- /dev/null +++ b/crates/types/src/codec/response.rs @@ -0,0 +1,369 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Normalized LLM response types produced by response codecs. +//! +//! This module defines [`AnnotatedLlmResponse`] and its supporting types +//! for structured, API-agnostic access to LLM response data. + +use serde::{Deserialize, Serialize}; + +use crate::Json; + +use super::request::MessageContent; + +// --------------------------------------------------------------------------- +// AnnotatedLlmResponse type hierarchy +// --------------------------------------------------------------------------- + +/// Structured view of an LLM response, produced by a response codec from +/// raw JSON API output. +/// +/// The `extra` field captures any top-level keys not modeled by the known +/// fields, ensuring lossless round-trip through serde. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct AnnotatedLlmResponse { + /// Response ID from the API (e.g., "chatcmpl-abc123", "resp_abc123", "msg_abc123"). + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, + + /// The model that actually served the request (may differ from requested model). + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// The assistant's response content, reusing [`MessageContent`] from request types. + #[serde(skip_serializing_if = "Option::is_none")] + pub message: Option, + + /// Tool calls requested by the model, normalized across APIs. + /// + /// Uses [`ResponseToolCall`] (arguments as [`Json`]) NOT the request-side + /// `ToolCall` (arguments as `String`). + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + + /// Why generation stopped, normalized across APIs. + #[serde(skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, + + /// Token usage statistics. + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, + + /// API-specific response data that cannot be normalized across providers. + #[serde(skip_serializing_if = "Option::is_none")] + pub api_specific: Option, + + /// Catch-all for unmodeled top-level fields, ensuring lossless round-trip. + #[serde(flatten)] + pub extra: serde_json::Map, +} + +// --------------------------------------------------------------------------- +// Usage +// --------------------------------------------------------------------------- + +/// Token usage statistics from an LLM API response. +/// +/// All fields are `Option` because not every provider supplies every +/// field. For example, cache token counts are only available from providers +/// that support prompt caching. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +pub struct Usage { + /// Tokens consumed by the prompt/input. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_tokens: Option, + /// Tokens generated in the completion/output. + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_tokens: Option, + /// Total tokens (prompt + completion). + #[serde(skip_serializing_if = "Option::is_none")] + pub total_tokens: Option, + /// Tokens served from prompt cache (read). + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read_tokens: Option, + /// Tokens written to prompt cache. + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_write_tokens: Option, + /// Optional cost reported by provider data or estimated from Relay pricing. + #[serde(skip_serializing_if = "Option::is_none")] + pub cost: Option, +} + +/// Source of a normalized cost value. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CostSource { + /// Cost was estimated by applying Relay's model pricing table to usage. + ModelPricing, + /// Cost was reported directly by a provider or framework payload. + ProviderReported, +} + +/// Normalized LLM response cost. +/// +/// Provider-reported cost is preserved as-is. Model-pricing estimates include +/// source and as-of metadata so downstream systems can audit stale pricing +/// tables without losing a usable estimate. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct CostEstimate { + /// Total cost in `currency`. + #[serde(skip_serializing_if = "Option::is_none")] + pub total: Option, + /// ISO 4217 currency code for the cost fields. + #[serde(default = "default_cost_currency")] + pub currency: String, + /// Uncached prompt/input token cost in `currency`. + #[serde(skip_serializing_if = "Option::is_none")] + pub input: Option, + /// Completion/output token cost in `currency`. + #[serde(skip_serializing_if = "Option::is_none")] + pub output: Option, + /// Prompt cache read cost in `currency`. + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_read: Option, + /// Prompt cache write cost in `currency`. + #[serde(skip_serializing_if = "Option::is_none")] + pub cache_write: Option, + /// Origin of this cost value. + pub source: CostSource, + /// Provider associated with the cost or pricing estimate, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub pricing_provider: Option, + /// Model ID associated with the cost or pricing estimate, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub pricing_model: Option, + /// Date the pricing value was last verified, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub pricing_as_of: Option, + /// Source URL or label for the pricing value, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub pricing_source: Option, +} + +impl CostEstimate { + /// Returns the explicit total, or the sum of component costs when no total was supplied. + #[must_use] + pub fn total_or_component_sum(&self) -> Option { + self.total.or_else(|| { + let (has_component, total) = + [self.input, self.output, self.cache_read, self.cache_write] + .into_iter() + .flatten() + .fold((false, 0.0), |(_, total), value| (true, total + value)); + has_component.then_some(total) + }) + } + + /// Returns the total only when it is denominated in the requested currency. + #[must_use] + pub fn total_for_currency(&self, currency: &str) -> Option { + self.currency + .eq_ignore_ascii_case(currency) + .then_some(self.total) + .flatten() + } + + /// Returns the explicit or component-derived total in the requested currency. + #[must_use] + pub fn total_or_component_sum_for_currency(&self, currency: &str) -> Option { + self.currency + .eq_ignore_ascii_case(currency) + .then(|| self.total_or_component_sum()) + .flatten() + } +} + +fn default_cost_currency() -> String { + "USD".into() +} + +// --------------------------------------------------------------------------- +// FinishReason +// --------------------------------------------------------------------------- + +/// Normalized reason why the model stopped generating. +/// +/// Maps from provider-specific stop reasons: +/// - **Complete**: OpenAI Chat `"stop"`, Anthropic `"end_turn"`, Responses `"completed"` +/// - **Length**: OpenAI Chat `"length"`, Anthropic `"max_tokens"`, Responses incomplete+max_output_tokens +/// - **ToolUse**: OpenAI Chat `"tool_calls"`, Anthropic `"tool_use"` +/// - **ContentFilter**: OpenAI Chat `"content_filter"`, Responses incomplete+content_filter +/// - **Unknown**: Forward-compatible catch-all for unrecognized reasons +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FinishReason { + /// Model naturally completed its response. + Complete, + /// Maximum token limit reached. + Length, + /// Model requested a tool call. + ToolUse, + /// Content was filtered by safety systems. + ContentFilter, + /// Unknown or forward-compatible reason. + Unknown(String), +} + +impl FinishReason { + /// Returns `true` if the model naturally completed its response. + /// + /// Only the [`FinishReason::Complete`] variant returns `true`. + #[must_use] + pub fn is_complete(&self) -> bool { + matches!(self, FinishReason::Complete) + } +} + +// --------------------------------------------------------------------------- +// ResponseToolCall +// --------------------------------------------------------------------------- + +/// A tool call requested by the model in its response. +/// +/// Unlike the request-side `ToolCall` (which stores arguments as a JSON +/// string per OpenAI convention), response tool calls store arguments as +/// parsed [`Json`]. Codecs parse OpenAI's string arguments during decode; +/// Anthropic's `input` is already parsed JSON. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ResponseToolCall { + /// Unique identifier for this tool call. + pub id: String, + /// The function/tool name. + pub name: String, + /// The arguments as parsed JSON (not a string). + pub arguments: Json, +} + +// --------------------------------------------------------------------------- +// ApiSpecificResponse +// --------------------------------------------------------------------------- + +/// API-specific response data that cannot be normalized across providers. +/// +/// Each variant captures fields unique to a particular LLM API, stored via +/// internal tagging on the `"api"` key. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "api")] +pub enum ApiSpecificResponse { + /// OpenAI Chat Completions-specific fields. + #[serde(rename = "openai_chat")] + OpenAIChat { + /// Token-level log probabilities (raw JSON, too complex to normalize). + #[serde(skip_serializing_if = "Option::is_none")] + logprobs: Option, + /// System fingerprint for reproducibility. + #[serde(skip_serializing_if = "Option::is_none")] + system_fingerprint: Option, + /// Processing tier used (e.g., "default"). + #[serde(skip_serializing_if = "Option::is_none")] + service_tier: Option, + }, + + /// OpenAI Responses API-specific fields. + #[serde(rename = "openai_responses")] + OpenAIResponses { + /// Full output items array for direct access. + #[serde(skip_serializing_if = "Option::is_none")] + output_items: Option>, + /// Response status (e.g., "completed", "incomplete"). + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + /// Details about why the response is incomplete. + #[serde(skip_serializing_if = "Option::is_none")] + incomplete_details: Option, + /// Echoed previous response ID for conversation continuation. + #[serde(skip_serializing_if = "Option::is_none")] + previous_response_id: Option, + /// Whether this response is marked for server-side storage. + #[serde(skip_serializing_if = "Option::is_none")] + store: Option, + /// Service tier used for the response. + #[serde(skip_serializing_if = "Option::is_none")] + service_tier: Option, + /// Truncation behavior metadata. + #[serde(skip_serializing_if = "Option::is_none")] + truncation: Option, + /// Reasoning configuration/result metadata. + #[serde(skip_serializing_if = "Option::is_none")] + reasoning: Option, + /// Raw input token details payload. + #[serde(skip_serializing_if = "Option::is_none")] + input_tokens_details: Option, + /// Raw output token details payload. + #[serde(skip_serializing_if = "Option::is_none")] + output_tokens_details: Option, + }, + + /// Anthropic Messages API-specific fields. + #[serde(rename = "anthropic_messages")] + AnthropicMessages { + /// Anthropic object type (typically `"message"`). + #[serde(skip_serializing_if = "Option::is_none")] + object_type: Option, + /// Anthropic response role (typically `"assistant"`). + #[serde(skip_serializing_if = "Option::is_none")] + role: Option, + /// Raw Anthropic stop_reason. + #[serde(skip_serializing_if = "Option::is_none")] + stop_reason: Option, + /// Which stop sequence was matched (if any). + #[serde(skip_serializing_if = "Option::is_none")] + stop_sequence: Option, + /// Anthropic response service tier when present. + #[serde(skip_serializing_if = "Option::is_none")] + service_tier: Option, + /// Anthropic container payload when present. + #[serde(skip_serializing_if = "Option::is_none")] + container: Option, + /// Full content blocks array for direct access. + #[serde(skip_serializing_if = "Option::is_none")] + content_blocks: Option>, + }, + + /// Custom/unknown API -- catch-all for user-implemented codecs. + #[serde(rename = "custom")] + Custom { + /// API identifier. + api_name: String, + /// Opaque API-specific data. + data: Json, + }, +} + +// --------------------------------------------------------------------------- +// Helper methods +// --------------------------------------------------------------------------- + +impl AnnotatedLlmResponse { + /// Extract the text content of the response message. + /// + /// For [`MessageContent::Text`], returns the string directly. + /// For [`MessageContent::Parts`], returns the text of the first + /// [`super::request::ContentPart::Text`] part. + /// Returns `None` if `message` is `None`. + #[must_use] + pub fn response_text(&self) -> Option<&str> { + match self.message.as_ref()? { + MessageContent::Text(s) => Some(s.as_str()), + MessageContent::Parts(parts) => parts.iter().find_map(|p| match p { + crate::codec::request::ContentPart::Text { text } => Some(text.as_str()), + crate::codec::request::ContentPart::ImageUrl { .. } => None, + }), + } + } + + /// Check if the response contains any tool calls. + /// + /// Returns `true` if `tool_calls` is `Some` with at least one element. + #[must_use] + pub fn has_tool_calls(&self) -> bool { + self.tool_calls + .as_ref() + .is_some_and(|calls| !calls.is_empty()) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs new file mode 100644 index 000000000..f36793fb7 --- /dev/null +++ b/crates/types/src/lib.rs @@ -0,0 +1,21 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)] + +//! Shared serializable data model types for NeMo Relay. +//! +//! This crate contains DTOs and wire-compatible type definitions used by the +//! runtime and native plugin SDK. It intentionally does not contain runtime +//! registries, dynamic loading, codecs, exporters, or process-global state. + +/// Public runtime API data types. +pub mod api; +/// Normalized LLM request and response data types. +pub mod codec; +/// Plugin configuration diagnostic data types. +pub mod plugin; + +/// Type alias for [`serde_json::Value`], used as the universal JSON +/// representation throughout NeMo Relay APIs. +pub type Json = serde_json::Value; diff --git a/crates/types/src/plugin.rs b/crates/types/src/plugin.rs new file mode 100644 index 000000000..eeb3be024 --- /dev/null +++ b/crates/types/src/plugin.rs @@ -0,0 +1,35 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared plugin diagnostic data types. + +use serde::{Deserialize, Serialize}; + +/// Diagnostic severity returned by plugin validation. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "lowercase")] +pub enum DiagnosticLevel { + /// Non-fatal compatibility or validation issue. + Warning, + /// Fatal validation issue that blocks initialization. + Error, +} + +/// Structured validation diagnostic for plugin validation. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ConfigDiagnostic { + /// Severity level for the diagnostic. + pub level: DiagnosticLevel, + /// Stable diagnostic code suitable for machine checks. + pub code: String, + /// Optional component identifier associated with the diagnostic. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub component: Option, + /// Optional field path associated with the diagnostic. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub field: Option, + /// Human-readable diagnostic message. + pub message: String, +} diff --git a/crates/types/tests/serialization_tests.rs b/crates/types/tests/serialization_tests.rs new file mode 100644 index 000000000..9e899ff7d --- /dev/null +++ b/crates/types/tests/serialization_tests.rs @@ -0,0 +1,80 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Serialization compatibility tests for shared NeMo Relay DTOs. + +use std::sync::Arc; + +use nemo_relay_types::api::event::{ + BaseEvent, CategoryProfile, Event, EventCategory, ScopeCategory, ScopeEvent, + llm_attributes_to_strings, +}; +use nemo_relay_types::api::llm::{LlmAttributes, LlmRequest}; +use nemo_relay_types::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; +use nemo_relay_types::codec::response::AnnotatedLlmResponse; +use serde_json::{Map, json}; + +#[test] +fn event_round_trips_with_annotated_llm_profiles() { + let request = AnnotatedLlmRequest { + messages: vec![Message::User { + content: MessageContent::Text("hello".into()), + name: None, + }], + model: Some("model".into()), + params: None, + tools: None, + tool_choice: None, + store: None, + previous_response_id: None, + truncation: None, + reasoning: None, + include: None, + user: None, + metadata: None, + service_tier: None, + parallel_tool_calls: None, + max_output_tokens: None, + max_tool_calls: None, + top_logprobs: None, + stream: None, + extra: Map::new(), + }; + let response = AnnotatedLlmResponse { + id: Some("resp_1".into()), + model: Some("model".into()), + message: Some(MessageContent::Text("world".into())), + tool_calls: None, + finish_reason: None, + usage: None, + api_specific: None, + extra: Map::new(), + }; + let event = Event::Scope(ScopeEvent::new( + BaseEvent::builder() + .name("llm") + .data(json!(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "hello" }), + })) + .build(), + ScopeCategory::Start, + llm_attributes_to_strings(LlmAttributes::STATEFUL), + EventCategory::llm(), + Some(CategoryProfile { + annotated_request: Some(Arc::new(request)), + annotated_response: Some(Arc::new(response)), + ..CategoryProfile::default() + }), + )); + + let encoded = serde_json::to_value(&event).expect("event should serialize"); + let decoded: Event = serde_json::from_value(encoded).expect("event should deserialize"); + assert_eq!(decoded.name(), "llm"); + assert_eq!( + decoded + .annotated_response() + .and_then(|response| response.id.as_deref()), + Some("resp_1") + ); +} From cfdfb308c291f3dec7d561c53919fa39a969e787 Mon Sep 17 00:00:00 2001 From: Alex Fournier Date: Thu, 25 Jun 2026 14:46:35 -0700 Subject: [PATCH 2/2] fix: import event normalization helpers in observability Signed-off-by: Alex Fournier --- crates/core/src/observability/atif.rs | 2 +- crates/core/src/observability/mod.rs | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index 62544aadc..67dedff5a 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -35,7 +35,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; -use crate::api::event::Event; +use crate::api::event::{Event, EventNormalizationExt}; use crate::api::runtime::EventSubscriberFn; use crate::api::subscriber::flush_subscribers; use crate::codec::request::{AnnotatedLlmRequest, Message, MessageContent}; diff --git a/crates/core/src/observability/mod.rs b/crates/core/src/observability/mod.rs index 685a6d2bf..447fc4ddb 100644 --- a/crates/core/src/observability/mod.rs +++ b/crates/core/src/observability/mod.rs @@ -3,6 +3,8 @@ //! Optional observability integrations for NeMo Relay Core. +use crate::api::event::EventNormalizationExt; + #[cfg(test)] use std::sync::Mutex;