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// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! # Dynamo LLM Protocols
17//!
18//! This module contains the protocols, i.e. messages formats, used to exchange requests and responses
19//! both publicly via the HTTP API and internally between Dynamo components.
20//!
21
22use futures::{Stream, StreamExt};
23use serde::{Deserialize, Serialize};
24use std::pin::Pin;
25
26pub mod codec;
27pub mod common;
28pub mod openai;
29
30/// The token ID type
31pub type TokenIdType = u32;
32pub type DataStream<T> = Pin<Box<dyn Stream<Item = T> + Send + Sync>>;
33
34// TODO: This is an awkward dependency that we need to address
35// Originally, all the Annotated/SSE Codec bits where in the LLM protocol module; however, [Annotated]
36// has become the common response envelope for dynamo.
37// We may want to move the original Annotated back here and has a Infallible conversion to the the
38// ResponseEnvelop in dynamo.
39pub use dynamo_runtime::protocols::annotated::Annotated;
40
41/// The LLM responses have multiple different fields and nests of objects to get to the actual
42/// text completion returned. This trait can be applied to the `choice` level objects to extract
43/// the completion text.
44///
45/// To avoid an optional, if no completion text is found, the [`ContentProvider::content`] should
46/// return an empty string.
47pub trait ContentProvider {
48    fn content(&self) -> String;
49}
50
51#[derive(Serialize, Deserialize, Debug, Clone)]
52pub struct Usage {
53    pub prompt_tokens: i32,
54    pub completion_tokens: i32,
55    pub total_tokens: i32,
56}
57
58/// Converts of a stream of [codec::Message]s into a stream of [Annotated]s.
59pub fn convert_sse_stream<R>(
60    stream: DataStream<Result<codec::Message, codec::SseCodecError>>,
61) -> DataStream<Annotated<R>>
62where
63    R: for<'de> Deserialize<'de> + Serialize,
64{
65    let stream = stream.map(|message| match message {
66        Ok(message) => {
67            let delta = Annotated::<R>::try_from(message);
68            match delta {
69                Ok(delta) => delta,
70                Err(e) => Annotated::from_error(e.to_string()),
71            }
72        }
73        Err(e) => Annotated::from_error(e.to_string()),
74    });
75    Box::pin(stream)
76}