dynamo_llm/
protocols.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! # Dynamo LLM Protocols
5//!
6//! This module contains the protocols, i.e. messages formats, used to exchange requests and responses
7//! both publicly via the HTTP API and internally between Dynamo components.
8//!
9
10use futures::{Stream, StreamExt};
11use serde::{Deserialize, Serialize};
12
13pub mod codec;
14pub mod common;
15pub mod openai;
16pub mod tensor;
17
18/// The token ID type
19pub type TokenIdType = u32;
20pub use dynamo_runtime::engine::DataStream;
21
22// TODO: This is an awkward dependency that we need to address
23// Originally, all the Annotated/SSE Codec bits where in the LLM protocol module; however, [Annotated]
24// has become the common response envelope for dynamo.
25// We may want to move the original Annotated back here and has a Infallible conversion to the the
26// ResponseEnvelop in dynamo.
27pub use dynamo_runtime::protocols::annotated::Annotated;
28
29/// The LLM responses have multiple different fields and nests of objects to get to the actual
30/// text completion returned. This trait can be applied to the `choice` level objects to extract
31/// the completion text.
32///
33/// To avoid an optional, if no completion text is found, the [`ContentProvider::content`] should
34/// return an empty string.
35pub trait ContentProvider {
36    fn content(&self) -> String;
37}
38
39/// Converts of a stream of [codec::Message]s into a stream of [Annotated]s.
40pub fn convert_sse_stream<R>(
41    stream: impl Stream<Item = Result<codec::Message, codec::SseCodecError>>,
42) -> impl Stream<Item = Annotated<R>>
43where
44    R: for<'de> Deserialize<'de> + Serialize,
45{
46    stream.map(|message| match message {
47        Ok(message) => {
48            let delta = Annotated::<R>::try_from(message);
49            match delta {
50                Ok(delta) => delta,
51                Err(e) => Annotated::from_error(e.to_string()),
52            }
53        }
54        Err(e) => Annotated::from_error(e.to_string()),
55    })
56}