"docs/source/features/tool_calling.md" did not exist on "6d525288c1a40ee70f9cff2fe08657f23bae88dc"
delta.rs 13.2 KB
Newer Older
1
2
3
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

4
use super::{NvCreateChatCompletionRequest, NvCreateChatCompletionStreamResponse};
Greg Clark's avatar
Greg Clark committed
5
use crate::{
6
    local_model::runtime_config::ModelRuntimeConfig,
Greg Clark's avatar
Greg Clark committed
7
8
9
    protocols::common::{self},
    types::TokenIdType,
};
10
use dynamo_parsers::{ParserResult, ReasoningParser, ReasoningParserType, ReasoningParserWrapper};
11

12
/// Provides a method for generating a [`DeltaGenerator`] from a chat completion request.
13
impl NvCreateChatCompletionRequest {
14
15
    /// Creates a [`DeltaGenerator`] instance based on the chat completion request.
    ///
16
17
18
    /// # Arguments
    /// * `request_id` - The request ID to use for the chat completion response ID.
    ///
19
20
    /// # Returns
    /// * [`DeltaGenerator`] configured with model name and response options.
21
    pub fn response_generator(&self, request_id: String) -> DeltaGenerator {
22
23
        let options = DeltaGeneratorOptions {
            enable_usage: true,
Greg Clark's avatar
Greg Clark committed
24
25
            enable_logprobs: self.inner.logprobs.unwrap_or(false)
                || self.inner.top_logprobs.unwrap_or(0) > 0,
26
            runtime_config: ModelRuntimeConfig::default(),
27
28
        };

29
        DeltaGenerator::new(self.inner.model.clone(), options, request_id)
30
31
32
    }
}

33
/// Configuration options for the [`DeltaGenerator`], controlling response behavior.
34
35
#[derive(Debug, Clone, Default)]
pub struct DeltaGeneratorOptions {
36
    /// Determines whether token usage statistics should be included in the response.
37
    pub enable_usage: bool,
38
    /// Determines whether log probabilities should be included in the response.
39
    pub enable_logprobs: bool,
40

41
    pub runtime_config: ModelRuntimeConfig,
42
43
}

44
/// Generates incremental chat completion responses in a streaming fashion.
45
#[derive(Debug)]
46
pub struct DeltaGenerator {
47
    /// Unique identifier for the chat completion session.
48
    id: String,
49
    /// Object type, representing a streamed chat completion response.
50
    object: String,
51
    /// Timestamp (Unix epoch) when the response was created.
Paul Hendricks's avatar
Paul Hendricks committed
52
    created: u32,
53
    model: String,
54
    /// Optional system fingerprint for version tracking.
55
    system_fingerprint: Option<String>,
56
    /// Optional service tier information for the response.
57
    service_tier: Option<dynamo_async_openai::types::ServiceTierResponse>,
58
    /// Tracks token usage for the completion request.
59
    usage: dynamo_async_openai::types::CompletionUsage,
60
    /// Counter tracking the number of messages issued.
61
    msg_counter: u64,
62
    /// Configuration options for response generation.
63
    options: DeltaGeneratorOptions,
64
65
66
67

    /// Reasoning Parser object
    /// This is used to parse reasoning content in the response.
    reasoning_parser: ReasoningParserWrapper,
68
69
70
}

impl DeltaGenerator {
71
72
73
74
75
    /// Creates a new [`DeltaGenerator`] instance with the specified model and options.
    ///
    /// # Arguments
    /// * `model` - The model name used for response generation.
    /// * `options` - Configuration options for enabling usage and log probabilities.
76
    /// * `request_id` - The request ID to use for the chat completion response.
77
78
79
    ///
    /// # Returns
    /// * A new instance of [`DeltaGenerator`].
80
    pub fn new(model: String, options: DeltaGeneratorOptions, request_id: String) -> Self {
81
82
83
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
84
85
86
87
88
            .as_secs();

        // SAFETY: Casting from `u64` to `u32` could lead to precision loss after `u32::MAX`,
        // but this will not be an issue until 2106.
        let now: u32 = now.try_into().expect("timestamp exceeds u32::MAX");
Paul Hendricks's avatar
Paul Hendricks committed
89

90
        let usage = dynamo_async_openai::types::CompletionUsage {
Paul Hendricks's avatar
Paul Hendricks committed
91
92
93
94
95
96
            prompt_tokens: 0,
            completion_tokens: 0,
            total_tokens: 0,
            prompt_tokens_details: None,
            completion_tokens_details: None,
        };
97

98
99
100
        // Reasoning parser type
        // This is hardcoded for now, but can be made configurable later.
        // TODO: Make parser type configurable once front-end integration is determined
101
        // Change to GptOss to test GptOSS parser
102
        // Reasoning parser wrapper
103
104
105
106
107
108
109
        let reasoning_parser = ReasoningParserType::get_reasoning_parser_from_name(
            options
                .runtime_config
                .reasoning_parser
                .as_deref()
                .unwrap_or("basic"),
        );
110

111
112
        let chatcmpl_id = format!("chatcmpl-{request_id}");

113
        Self {
114
            id: chatcmpl_id,
115
116
117
118
119
            object: "chat.completion.chunk".to_string(),
            created: now,
            model,
            system_fingerprint: None,
            service_tier: None,
Paul Hendricks's avatar
Paul Hendricks committed
120
            usage,
121
122
            msg_counter: 0,
            options,
123
            reasoning_parser,
124
125
126
        }
    }

127
128
129
130
    /// Updates the prompt token usage count.
    ///
    /// # Arguments
    /// * `isl` - The number of prompt tokens used.
Paul Hendricks's avatar
Paul Hendricks committed
131
    pub fn update_isl(&mut self, isl: u32) {
132
133
134
        self.usage.prompt_tokens = isl;
    }

Greg Clark's avatar
Greg Clark committed
135
136
137
    pub fn create_logprobs(
        &self,
        tokens: Vec<common::llm_backend::TokenType>,
138
        token_ids: &[TokenIdType],
Greg Clark's avatar
Greg Clark committed
139
140
        logprobs: Option<common::llm_backend::LogProbs>,
        top_logprobs: Option<common::llm_backend::TopLogprobs>,
141
    ) -> Option<dynamo_async_openai::types::ChatChoiceLogprobs> {
Greg Clark's avatar
Greg Clark committed
142
143
144
145
146
147
148
        if !self.options.enable_logprobs || logprobs.is_none() {
            return None;
        }

        let toks = tokens
            .into_iter()
            .zip(token_ids)
149
            .map(|(token, token_id)| (token.unwrap_or_default(), *token_id))
Greg Clark's avatar
Greg Clark committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
            .collect::<Vec<(String, TokenIdType)>>();
        let tok_lps = toks
            .iter()
            .zip(logprobs.unwrap())
            .map(|(_, lp)| lp as f32)
            .collect::<Vec<f32>>();

        let content = top_logprobs.map(|top_logprobs| {
            toks.iter()
                .zip(tok_lps)
                .zip(top_logprobs)
                .map(|(((t, tid), lp), top_lps)| {
                    let mut found_selected_token = false;
                    let mut converted_top_lps = top_lps
                        .iter()
                        .map(|top_lp| {
                            let top_t = top_lp.token.clone().unwrap_or_default();
                            let top_tid = top_lp.token_id;
                            found_selected_token = found_selected_token || top_tid == *tid;
169
                            dynamo_async_openai::types::TopLogprobs {
Greg Clark's avatar
Greg Clark committed
170
171
172
173
174
                                token: top_t,
                                logprob: top_lp.logprob as f32,
                                bytes: None,
                            }
                        })
175
                        .collect::<Vec<dynamo_async_openai::types::TopLogprobs>>();
Greg Clark's avatar
Greg Clark committed
176
177
                    if !found_selected_token {
                        // If the selected token is not in the top logprobs, add it
178
                        converted_top_lps.push(dynamo_async_openai::types::TopLogprobs {
Greg Clark's avatar
Greg Clark committed
179
180
181
182
183
                            token: t.clone(),
                            logprob: lp,
                            bytes: None,
                        });
                    }
184
                    dynamo_async_openai::types::ChatCompletionTokenLogprob {
Greg Clark's avatar
Greg Clark committed
185
186
187
188
189
190
191
192
193
                        token: t.clone(),
                        logprob: lp,
                        bytes: None,
                        top_logprobs: converted_top_lps,
                    }
                })
                .collect()
        });

194
        Some(dynamo_async_openai::types::ChatChoiceLogprobs {
Greg Clark's avatar
Greg Clark committed
195
196
197
198
199
            content,
            refusal: None,
        })
    }

200
201
202
203
204
205
206
207
208
    fn create_reasoning_content(
        &mut self,
        text: &Option<String>,
        token_ids: &[u32],
    ) -> Option<ParserResult> {
        let text_ref = text.as_deref().unwrap_or("");
        if text_ref.is_empty() && token_ids.is_empty() {
            return None;
        }
209
210
        let parser_result = self
            .reasoning_parser
211
            .parse_reasoning_streaming_incremental(text_ref, token_ids);
212
213
214
215

        Some(parser_result)
    }

216
217
218
219
220
221
222
223
224
    /// Creates a choice within a chat completion response.
    ///
    /// # Arguments
    /// * `index` - The index of the choice in the completion response.
    /// * `text` - The text content for the response.
    /// * `finish_reason` - The reason why the response finished (e.g., stop, length, etc.).
    /// * `logprobs` - Optional log probabilities of the generated tokens.
    ///
    /// # Returns
225
    /// * An [`dynamo_async_openai::types::CreateChatCompletionStreamResponse`] instance representing the choice.
Paul Hendricks's avatar
Paul Hendricks committed
226
    #[allow(deprecated)]
227
    pub fn create_choice(
228
        &mut self,
Paul Hendricks's avatar
Paul Hendricks committed
229
        index: u32,
230
        text: Option<String>,
231
        reasoning_content: Option<String>,
232
233
        finish_reason: Option<dynamo_async_openai::types::FinishReason>,
        logprobs: Option<dynamo_async_openai::types::ChatChoiceLogprobs>,
234
    ) -> NvCreateChatCompletionStreamResponse {
235
        let delta = dynamo_async_openai::types::ChatCompletionStreamResponseDelta {
236
            content: text,
237
238
            function_call: None,
            tool_calls: None,
239
            role: if self.msg_counter == 0 {
240
                Some(dynamo_async_openai::types::Role::Assistant)
241
242
243
            } else {
                None
            },
Paul Hendricks's avatar
Paul Hendricks committed
244
            refusal: None,
245
            reasoning_content,
246
247
        };

248
        let choice = dynamo_async_openai::types::ChatChoiceStream {
Paul Hendricks's avatar
Paul Hendricks committed
249
250
251
252
253
254
255
256
            index,
            delta,
            finish_reason,
            logprobs,
        };

        let choices = vec![choice];

257
258
259
260
261
        let mut usage = self.usage.clone();
        if self.options.enable_usage {
            usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
        }

262
        dynamo_async_openai::types::CreateChatCompletionStreamResponse {
263
264
265
266
267
            id: self.id.clone(),
            object: self.object.clone(),
            created: self.created,
            model: self.model.clone(),
            system_fingerprint: self.system_fingerprint.clone(),
Paul Hendricks's avatar
Paul Hendricks committed
268
            choices,
269
            usage: if self.options.enable_usage {
270
                Some(usage)
271
272
273
274
275
276
277
278
            } else {
                None
            },
            service_tier: self.service_tier.clone(),
        }
    }
}

279
/// Implements the [`crate::protocols::openai::DeltaGeneratorExt`] trait for [`DeltaGenerator`], allowing
280
/// it to transform backend responses into OpenAI-style streaming responses.
281
282
283
impl crate::protocols::openai::DeltaGeneratorExt<NvCreateChatCompletionStreamResponse>
    for DeltaGenerator
{
284
285
286
287
288
289
290
291
    /// Converts a backend response into a structured OpenAI-style streaming response.
    ///
    /// # Arguments
    /// * `delta` - The backend response containing generated text and metadata.
    ///
    /// # Returns
    /// * `Ok(NvCreateChatCompletionStreamResponse)` if conversion succeeds.
    /// * `Err(anyhow::Error)` if an error occurs.
292
293
294
    fn choice_from_postprocessor(
        &mut self,
        delta: crate::protocols::common::llm_backend::BackendOutput,
295
    ) -> anyhow::Result<NvCreateChatCompletionStreamResponse> {
296
        // Aggregate token usage if enabled.
297
        if self.options.enable_usage {
298
299
300
301
302
303
304
305
306
            // SAFETY: Casting from `usize` to `u32` could lead to precision loss after `u32::MAX`,
            // but this will not be an issue until context lengths exceed 4_294_967_295.
            let token_length: u32 = delta
                .token_ids
                .len()
                .try_into()
                .expect("token_ids length exceeds u32::MAX");

            self.usage.completion_tokens += token_length;
307
308
        }

Greg Clark's avatar
Greg Clark committed
309
310
        let logprobs = self.create_logprobs(
            delta.tokens,
311
            &delta.token_ids,
Greg Clark's avatar
Greg Clark committed
312
313
314
            delta.log_probs,
            delta.top_logprobs,
        );
315

316
        // Map backend finish reasons to OpenAI's finish reasons.
317
        let finish_reason = match delta.finish_reason {
318
319
320
321
322
323
324
325
326
327
            Some(common::FinishReason::EoS) => Some(dynamo_async_openai::types::FinishReason::Stop),
            Some(common::FinishReason::Stop) => {
                Some(dynamo_async_openai::types::FinishReason::Stop)
            }
            Some(common::FinishReason::Length) => {
                Some(dynamo_async_openai::types::FinishReason::Length)
            }
            Some(common::FinishReason::Cancelled) => {
                Some(dynamo_async_openai::types::FinishReason::Stop)
            }
328
            Some(common::FinishReason::ContentFilter) => {
329
                Some(dynamo_async_openai::types::FinishReason::ContentFilter)
330
            }
331
332
333
334
335
336
            Some(common::FinishReason::Error(err_msg)) => {
                return Err(anyhow::anyhow!(err_msg));
            }
            None => None,
        };

337
338
339
340
341
342
343
344
345
        let reasoning_parser_result = self
            .create_reasoning_content(&delta.text, &delta.token_ids)
            .unwrap_or_default();

        let (normal_text, reasoning_content) = (
            reasoning_parser_result.get_some_normal_text(),
            reasoning_parser_result.get_some_reasoning(),
        );

346
        // Create the streaming response.
347
        let index = 0;
348
349
350
351
352
353
354
        let stream_response = self.create_choice(
            index,
            normal_text,
            reasoning_content,
            finish_reason,
            logprobs,
        );
Paul Hendricks's avatar
Paul Hendricks committed
355

356
        Ok(stream_response)
357
    }
358
359
360
361

    fn get_isl(&self) -> Option<u32> {
        Some(self.usage.prompt_tokens)
    }
362
}