openai.rs 97.2 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

4
5
use std::{
    collections::HashSet,
6
    fmt::Display,
7
8
9
10
    sync::Arc,
    time::{SystemTime, UNIX_EPOCH},
};

11
use axum::{
12
    Json, Router,
13
    body::Body,
14
    extract::State,
15
    http::Request,
Ryan Olson's avatar
Ryan Olson committed
16
    http::{HeaderMap, StatusCode},
17
    middleware::{self, Next},
18
19
    response::{
        IntoResponse, Response,
20
        sse::{KeepAlive, Sse},
21
22
23
    },
    routing::{get, post},
};
24
use dynamo_runtime::config::environment_names::llm as env_llm;
Ryan Olson's avatar
Ryan Olson committed
25
26
27
28
use dynamo_runtime::{
    pipeline::{AsyncEngineContextProvider, Context},
    protocols::annotated::AnnotationsProvider,
};
29
use futures::{StreamExt, stream};
30
31
32
use serde::{Deserialize, Serialize};

use super::{
33
34
    RouteDoc,
    disconnect::{ConnectionHandle, create_connection_monitor, monitor_for_disconnects},
35
    error::HttpError,
36
37
38
39
    metrics::{
        Endpoint, EventConverter, process_response_and_observe_metrics,
        process_response_using_event_converter_and_observe_metrics,
    },
40
    service_v2,
41
};
42
use crate::engines::ValidateRequest;
43
use crate::protocols::openai::chat_completions::aggregator::ChatCompletionAggregator;
44
use crate::protocols::openai::nvext::apply_header_routing_overrides;
45
use crate::protocols::openai::{
46
47
48
49
    chat_completions::{
        NvCreateChatCompletionRequest, NvCreateChatCompletionResponse,
        NvCreateChatCompletionStreamResponse,
    },
50
51
    completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
    embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
52
    images::{NvCreateImageRequest, NvImagesResponse},
53
    responses::{NvCreateResponse, NvResponse, ResponseParams, chat_completion_to_response},
54
    videos::{NvCreateVideoRequest, NvVideosResponse},
55
};
56
use crate::request_template::RequestTemplate;
57
use crate::types::Annotated;
58
59
use dynamo_runtime::logging::get_distributed_tracing_context;
use tracing::Instrument;
60

Ryan Olson's avatar
Ryan Olson committed
61
62
63
64
65
pub const DYNAMO_REQUEST_ID_HEADER: &str = "x-dynamo-request-id";

/// Dynamo Annotation for the request ID
pub const ANNOTATION_REQUEST_ID: &str = "request_id";

66
67
const VALIDATION_PREFIX: &str = "Validation: ";

68
69
70
// Default axum max body limit without configuring is 2MB: https://docs.rs/axum/latest/axum/extract/struct.DefaultBodyLimit.html
/// Default body limit in bytes (45MB) to support 500k+ token payloads.
/// Can be configured at compile time using the DYN_FRONTEND_BODY_LIMIT_MB environment variable
71
pub(super) fn get_body_limit() -> usize {
72
    std::env::var(env_llm::DYN_HTTP_BODY_LIMIT_MB)
73
74
75
76
77
78
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .map(|mb| mb * 1024 * 1024)
        .unwrap_or(45 * 1024 * 1024)
}

Ryan Olson's avatar
Ryan Olson committed
79
80
pub type ErrorResponse = (StatusCode, Json<ErrorMessage>);

81
#[derive(Serialize, Deserialize, Debug)]
Ryan Olson's avatar
Ryan Olson committed
82
pub(crate) struct ErrorMessage {
83
84
85
86
87
88
89
90
91
92
93
    message: String,
    #[serde(rename = "type")]
    error_type: String,
    code: u16,
}

fn map_error_code_to_error_type(code: StatusCode) -> String {
    match code.canonical_reason() {
        Some(reason) => reason.to_string(),
        None => "UnknownError".to_string(),
    }
94
95
}

Ryan Olson's avatar
Ryan Olson committed
96
impl ErrorMessage {
97
    /// Not Found Error
Ryan Olson's avatar
Ryan Olson committed
98
    pub fn model_not_found() -> ErrorResponse {
99
100
        let code = StatusCode::NOT_FOUND;
        let error_type = map_error_code_to_error_type(code);
101
        (
102
            code,
Ryan Olson's avatar
Ryan Olson committed
103
            Json(ErrorMessage {
104
105
106
                message: "Model not found".to_string(),
                error_type,
                code: code.as_u16(),
107
108
109
110
111
112
            }),
        )
    }

    /// Service Unavailable
    /// This is returned when the service is live, but not ready.
Ryan Olson's avatar
Ryan Olson committed
113
    pub fn _service_unavailable() -> ErrorResponse {
114
115
        let code = StatusCode::SERVICE_UNAVAILABLE;
        let error_type = map_error_code_to_error_type(code);
116
        (
117
            code,
Ryan Olson's avatar
Ryan Olson committed
118
            Json(ErrorMessage {
119
120
121
                message: "Service is not ready".to_string(),
                error_type,
                code: code.as_u16(),
122
123
124
125
126
127
128
129
            }),
        )
    }

    /// Internal Service Error
    /// Return this error when the service encounters an internal error.
    /// We should return a generic message to the client instead of the real error.
    /// Internal Services errors are the result of misconfiguration or bugs in the service.
Ryan Olson's avatar
Ryan Olson committed
130
    pub fn internal_server_error(msg: &str) -> ErrorResponse {
131
        tracing::error!("Internal server error: {msg}");
132
133
        let code = StatusCode::INTERNAL_SERVER_ERROR;
        let error_type = map_error_code_to_error_type(code);
134
        (
135
            code,
Ryan Olson's avatar
Ryan Olson committed
136
            Json(ErrorMessage {
137
138
139
                message: msg.to_string(),
                error_type,
                code: code.as_u16(),
140
141
142
143
            }),
        )
    }

144
145
146
    /// Not Implemented Error
    /// Return this error when the client requests a feature that is not yet implemented.
    /// This should be used for features that are planned but not available.
147
    pub fn not_implemented_error<T: Display>(msg: T) -> ErrorResponse {
148
        tracing::error!("Not Implemented error: {msg}");
149
150
        let code = StatusCode::NOT_IMPLEMENTED;
        let error_type = map_error_code_to_error_type(code);
151
        (
152
            code,
Ryan Olson's avatar
Ryan Olson committed
153
            Json(ErrorMessage {
154
155
156
                message: msg.to_string(),
                error_type,
                code: code.as_u16(),
157
158
159
160
            }),
        )
    }

Neelay Shah's avatar
Neelay Shah committed
161
    /// The OAI endpoints call an [`dynamo.runtime::engine::AsyncEngine`] which are specialized to return
162
    /// an [`anyhow::Error`]. This method will convert the [`anyhow::Error`] into an [`HttpError`].
Ryan Olson's avatar
Ryan Olson committed
163
    /// If successful, it will return the [`HttpError`] as an [`ErrorMessage::internal_server_error`]
164
    /// with the details of the error.
Ryan Olson's avatar
Ryan Olson committed
165
    pub fn from_anyhow(err: anyhow::Error, alt_msg: &str) -> ErrorResponse {
166
167
168
        // First check for PipelineError::ServiceOverloaded
        if let Some(pipeline_err) =
            err.downcast_ref::<dynamo_runtime::pipeline::error::PipelineError>()
169
            && matches!(
170
171
                pipeline_err,
                dynamo_runtime::pipeline::error::PipelineError::ServiceOverloaded(_)
172
173
174
175
176
            )
        {
            return (
                StatusCode::SERVICE_UNAVAILABLE,
                Json(ErrorMessage {
177
178
179
                    message: pipeline_err.to_string(),
                    error_type: map_error_code_to_error_type(StatusCode::SERVICE_UNAVAILABLE),
                    code: StatusCode::SERVICE_UNAVAILABLE.as_u16(),
180
181
                }),
            );
182
183
184
        }

        // Then check for HttpError
185
        match err.downcast::<HttpError>() {
Ryan Olson's avatar
Ryan Olson committed
186
            Ok(http_error) => ErrorMessage::from_http_error(http_error),
187
            Err(err) => ErrorMessage::internal_server_error(&format!("{alt_msg}: {err:#}")),
188
189
190
191
        }
    }

    /// Implementers should only be able to throw 400-499 errors.
Ryan Olson's avatar
Ryan Olson committed
192
    pub fn from_http_error(err: HttpError) -> ErrorResponse {
193
        if err.code < 400 || err.code >= 500 {
Ryan Olson's avatar
Ryan Olson committed
194
            return ErrorMessage::internal_server_error(&err.message);
195
196
        }
        match StatusCode::from_u16(err.code) {
197
198
199
200
201
202
203
204
            Ok(code) => (
                code,
                Json(ErrorMessage {
                    message: err.message,
                    error_type: map_error_code_to_error_type(code),
                    code: code.as_u16(),
                }),
            ),
Ryan Olson's avatar
Ryan Olson committed
205
            Err(_) => ErrorMessage::internal_server_error(&err.message),
206
207
208
209
        }
    }
}

Ryan Olson's avatar
Ryan Olson committed
210
impl From<HttpError> for ErrorMessage {
211
    fn from(err: HttpError) -> Self {
212
213
214
215
216
217
218
        ErrorMessage {
            message: err.message,
            error_type: map_error_code_to_error_type(
                StatusCode::from_u16(err.code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
            ),
            code: err.code,
        }
219
220
221
    }
}

222
223
224
225
226
227
228
229
// Problem: Currently we are using JSON from axum as the request validator. Whenever there is an invalid JSON, it will return a 422.
// But all the downstream apps that relies on openai based APIs, expects to get 400 for all these cases otherwise they fail badly
// Solution: Intercept the response from handlers and convert ANY 422 status codes to 400 with the actual error message.
pub async fn smart_json_error_middleware(request: Request<Body>, next: Next) -> Response {
    let response = next.run(request).await;

    if response.status() == StatusCode::UNPROCESSABLE_ENTITY {
        let (_parts, body) = response.into_parts();
230
        let body_bytes = axum::body::to_bytes(body, get_body_limit())
231
232
233
234
235
236
            .await
            .unwrap_or_default();
        let error_message = String::from_utf8_lossy(&body_bytes).to_string();
        (
            StatusCode::BAD_REQUEST,
            Json(ErrorMessage {
237
238
239
                message: error_message,
                error_type: map_error_code_to_error_type(StatusCode::BAD_REQUEST),
                code: StatusCode::BAD_REQUEST.as_u16(),
240
241
242
243
244
245
246
247
248
            }),
        )
            .into_response()
    } else {
        // Pass through if it is not a 422
        response
    }
}

Ryan Olson's avatar
Ryan Olson committed
249
/// Get the request ID from a primary source, or next from the headers, or lastly create a new one if not present
250
// TODO: Similar function exists in lib/llm/src/grpc/service/openai.rs but with different signature and simpler logic
251
pub(super) fn get_or_create_request_id(primary: Option<&str>, headers: &HeaderMap) -> String {
252
    // Try to get request id from trace context
253
254
255
256
    if let Some(trace_context) = get_distributed_tracing_context()
        && let Some(x_dynamo_request_id) = trace_context.x_dynamo_request_id
    {
        return x_dynamo_request_id;
257
258
    }

Ryan Olson's avatar
Ryan Olson committed
259
    // Try to get the request ID from the primary source
260
261
262
263
    if let Some(primary) = primary
        && let Ok(uuid) = uuid::Uuid::parse_str(primary)
    {
        return uuid.to_string();
Ryan Olson's avatar
Ryan Olson committed
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    }

    // Try to get the request ID header as a string slice
    let request_id_opt = headers
        .get(DYNAMO_REQUEST_ID_HEADER)
        .and_then(|h| h.to_str().ok());

    // Try to parse the request ID as a UUID, or generate a new one if missing/invalid
    let uuid = match request_id_opt {
        Some(request_id) => {
            uuid::Uuid::parse_str(request_id).unwrap_or_else(|_| uuid::Uuid::new_v4())
        }
        None => uuid::Uuid::new_v4(),
    };

    uuid.to_string()
}

282
283
284
285
286
287
288
289
/// OpenAI Completions Request Handler
///
/// This method will handle the incoming request for the `/v1/completions endpoint`. The endpoint is a "source"
/// for an [`super::OpenAICompletionsStreamingEngine`] and will return a stream of
/// responses which will be forward to the client.
///
/// Note: For all requests, streaming or non-streaming, we always call the engine with streaming enabled. For
/// non-streaming requests, we will fold the stream into a single response as part of this handler.
Ryan Olson's avatar
Ryan Olson committed
290
async fn handler_completions(
291
    State(state): State<Arc<service_v2::State>>,
Ryan Olson's avatar
Ryan Olson committed
292
    headers: HeaderMap,
293
    Json(mut request): Json<NvCreateCompletionRequest>,
Ryan Olson's avatar
Ryan Olson committed
294
295
296
297
) -> Result<Response, ErrorResponse> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

298
299
    request.nvext = apply_header_routing_overrides(request.nvext.take(), &headers);

Ryan Olson's avatar
Ryan Olson committed
300
301
302
303
304
305
    // create the context for the request
    let request_id = get_or_create_request_id(request.inner.user.as_deref(), &headers);
    let request = Context::with_id(request, request_id);
    let context = request.context();

    // create the connection handles
306
307
    let (mut connection_handle, stream_handle) =
        create_connection_monitor(context.clone(), Some(state.metrics_clone())).await;
Ryan Olson's avatar
Ryan Olson committed
308
309
310

    // possibly long running task
    // if this returns a streaming response, the stream handle will be armed and captured by the response stream
311
    let response = tokio::spawn(completions(state, request, stream_handle).in_current_span())
Ryan Olson's avatar
Ryan Olson committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        .await
        .map_err(|e| {
            ErrorMessage::internal_server_error(&format!(
                "Failed to await chat completions task: {:?}",
                e,
            ))
        })?;

    // if we got here, then we will return a response and the potentially long running task has completed successfully
    // without need to be cancelled.
    connection_handle.disarm();

    response
}

#[tracing::instrument(skip_all)]
async fn completions(
    state: Arc<service_v2::State>,
    request: Context<NvCreateCompletionRequest>,
    stream_handle: ConnectionHandle,
) -> Result<Response, ErrorResponse> {
333
334
    use crate::protocols::openai::completions::get_prompt_batch_size;

335
336
337
    // return a 503 if the service is not ready
    check_ready(&state)?;

338
339
340
    // Validate stream_options is only used when streaming (NVBug 5662680)
    validate_completion_stream_options(&request)?;

341
342
    validate_completion_fields_generic(&request)?;

343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
    // Detect batch prompts
    let batch_size = get_prompt_batch_size(&request.inner.prompt);
    let n = request.inner.n.unwrap_or(1);

    // If single prompt or single-element batch, use original flow
    if batch_size == 1 {
        return completions_single(state, request, stream_handle).await;
    }

    // Batch processing: handle multiple prompts
    completions_batch(state, request, stream_handle, batch_size, n).await
}

/// Handle single prompt completions (original logic)
#[tracing::instrument(skip_all)]
async fn completions_single(
    state: Arc<service_v2::State>,
    request: Context<NvCreateCompletionRequest>,
    stream_handle: ConnectionHandle,
) -> Result<Response, ErrorResponse> {
363
    let request_id = request.id().to_string();
364
365

    // todo - decide on default
366
    let streaming = request.inner.stream.unwrap_or(false);
367

368
369
370
371
372
373
    // todo - make the protocols be optional for model name
    // todo - when optional, if none, apply a default
    let model = request.inner.model.clone();
    // Create http_queue_guard early - tracks time waiting to be processed
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

374
375
    // todo - error handling should be more robust
    let engine = state
376
        .manager()
377
        .get_completions_engine(&model)
Ryan Olson's avatar
Ryan Olson committed
378
        .map_err(|_| ErrorMessage::model_not_found())?;
379

380
    let parsing_options = state.manager().get_parsing_options(&model);
381

382
    let mut response_collector = state.metrics_clone().create_response_collector(&model);
383

Ryan Olson's avatar
Ryan Olson committed
384
385
    // prepare to process any annotations
    let annotations = request.annotations();
386

387
388
389
390
391
392
    // Create inflight_guard before calling engine to ensure errors are counted
    let mut inflight_guard =
        state
            .metrics_clone()
            .create_inflight_guard(&model, Endpoint::Completions, streaming);

393
394
395
396
    // issue the generate call on the engine
    let stream = engine
        .generate(request)
        .await
Ryan Olson's avatar
Ryan Olson committed
397
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate completions"))?;
398
399
400
401

    // capture the context to cancel the stream if the client disconnects
    let ctx = stream.context();

Ryan Olson's avatar
Ryan Olson committed
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
    let annotations = annotations.map_or(Vec::new(), |annotations| {
        annotations
            .iter()
            .filter_map(|annotation| {
                if annotation == ANNOTATION_REQUEST_ID {
                    Annotated::<NvCreateCompletionResponse>::from_annotation(
                        ANNOTATION_REQUEST_ID,
                        &request_id,
                    )
                    .ok()
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
    });

    // apply any annotations to the front of the stream
    let stream = stream::iter(annotations).chain(stream);
421
422

    if streaming {
423
424
        // For streaming, we'll drop the http_queue_guard on the first token
        let mut http_queue_guard = Some(http_queue_guard);
425
426
427
428
429
430
431
432
433
434
435
436
437
438
        let stream = stream
            .map(move |response| {
                // Calls observe_response() on each token
                process_response_using_event_converter_and_observe_metrics(
                    EventConverter::from(response),
                    &mut response_collector,
                    &mut http_queue_guard,
                )
            })
            .filter_map(|result| {
                use futures::future;
                // Transpose Result<Option<T>> -> Option<Result<T>>
                future::ready(result.transpose())
            });
Ryan Olson's avatar
Ryan Olson committed
439
        let stream = monitor_for_disconnects(stream, ctx, inflight_guard, stream_handle);
440

441
442
        let mut sse_stream = Sse::new(stream);

443
        if let Some(keep_alive) = state.sse_keep_alive() {
444
445
446
447
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
448
    } else {
449
        // Tap the stream to collect metrics for non-streaming requests without altering items
450
        let mut http_queue_guard = Some(http_queue_guard);
451
        let stream = stream.inspect(move |response| {
452
453
454
455
456
457
            // Calls observe_response() on each token - drops http_queue_guard on first token
            process_response_and_observe_metrics(
                response,
                &mut response_collector,
                &mut http_queue_guard,
            );
458
459
        });

460
        let response = NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
            .await
            .map_err(|e| {
                tracing::error!(
                    "Failed to fold completions stream for {}: {:?}",
                    request_id,
                    e
                );
                ErrorMessage::internal_server_error(&format!(
                    "Failed to fold completions stream for {}: {:?}",
                    request_id, e
                ))
            })?;

        inflight_guard.mark_ok();
        Ok(Json(response).into_response())
    }
}

/// Handle batch prompt completions (multiple prompts with n choices each)
#[tracing::instrument(skip_all)]
async fn completions_batch(
    state: Arc<service_v2::State>,
    request: Context<NvCreateCompletionRequest>,
    stream_handle: ConnectionHandle,
    batch_size: usize,
    n: u8,
) -> Result<Response, ErrorResponse> {
    use crate::protocols::openai::completions::extract_single_prompt;
    use futures::stream::{self, StreamExt};

    let request_id = request.id().to_string();
    let streaming = request.inner.stream.unwrap_or(false);
    let model = request.inner.model.clone();

    // Create http_queue_guard early - tracks time waiting to be processed
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

    let engine = state
        .manager()
        .get_completions_engine(&model)
        .map_err(|_| ErrorMessage::model_not_found())?;

    let parsing_options = state.manager().get_parsing_options(&model);

    let mut response_collector = state.metrics_clone().create_response_collector(&model);

    // prepare to process any annotations
    let annotations = request.annotations();

    // Create inflight_guard before calling engine to ensure errors are counted
    let mut inflight_guard =
        state
            .metrics_clone()
            .create_inflight_guard(&model, Endpoint::Completions, streaming);

    // Generate streams for each prompt in the batch
    let mut all_streams = Vec::new();
    let mut first_ctx = None;

    for prompt_idx in 0..batch_size {
        // Extract single prompt at this index
        let single_prompt = extract_single_prompt(&request.inner.prompt, prompt_idx);

        // Create a new request with this single prompt
        let mut single_request = request.content().clone();
        single_request.inner.prompt = single_prompt;

        // Generate unique request_id for each prompt: original_id-{prompt_idx}
        let unique_request_id = format!("{}-{}", request.id(), prompt_idx);
        let single_request_context = Context::with_id(single_request, unique_request_id);

        // Generate stream for this prompt
        let stream = engine
            .generate(single_request_context)
            .await
            .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate completions"))?;

        // Capture context from first stream
        if first_ctx.is_none() {
            first_ctx = Some(stream.context());
        }

        // Remap choice indices: choice.index += prompt_idx * n
        let prompt_idx_u32 = prompt_idx as u32;
        let n_u32 = n as u32;
        let remapped_stream = stream.map(move |mut response| {
            if let Some(ref mut data) = response.data {
                for choice in &mut data.inner.choices {
                    choice.index += prompt_idx_u32 * n_u32;
                }
            }
            response
        });

        all_streams.push(remapped_stream);
    }

    // Merge all streams
    let merged_stream = stream::select_all(all_streams);

    // capture the context to cancel the stream if the client disconnects
    let ctx = first_ctx.expect("At least one stream should be generated");

    let annotations_vec = annotations.map_or(Vec::new(), |annotations| {
        annotations
            .iter()
            .filter_map(|annotation| {
                if annotation == ANNOTATION_REQUEST_ID {
                    Annotated::<NvCreateCompletionResponse>::from_annotation(
                        ANNOTATION_REQUEST_ID,
                        &request_id,
                    )
                    .ok()
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
    });

    // apply any annotations to the front of the stream
    let merged_stream = stream::iter(annotations_vec).chain(merged_stream);

    if streaming {
        // For streaming, we'll drop the http_queue_guard on the first token
        let mut http_queue_guard = Some(http_queue_guard);
587
588
589
590
591
592
593
594
595
596
597
598
599
600
        let stream = merged_stream
            .map(move |response| {
                // Calls observe_response() on each token
                process_response_using_event_converter_and_observe_metrics(
                    EventConverter::from(response),
                    &mut response_collector,
                    &mut http_queue_guard,
                )
            })
            .filter_map(|result| {
                use futures::future;
                // Transpose Result<Option<T>> -> Option<Result<T>>
                future::ready(result.transpose())
            });
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
        let stream = monitor_for_disconnects(stream, ctx, inflight_guard, stream_handle);

        let mut sse_stream = Sse::new(stream);

        if let Some(keep_alive) = state.sse_keep_alive() {
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
    } else {
        // Tap the stream to collect metrics for non-streaming requests without altering items
        let mut http_queue_guard = Some(http_queue_guard);
        let stream = merged_stream.inspect(move |response| {
            // Calls observe_response() on each token - drops http_queue_guard on first token
            process_response_and_observe_metrics(
                response,
                &mut response_collector,
                &mut http_queue_guard,
            );
        });

        let response = NvCreateCompletionResponse::from_annotated_stream(stream, parsing_options)
623
624
625
626
627
628
629
            .await
            .map_err(|e| {
                tracing::error!(
                    "Failed to fold completions stream for {}: {:?}",
                    request_id,
                    e
                );
630
631
632
633
                ErrorMessage::internal_server_error(&format!(
                    "Failed to fold completions stream for {}: {:?}",
                    request_id, e
                ))
634
635
            })?;

636
        inflight_guard.mark_ok();
637
638
639
640
        Ok(Json(response).into_response())
    }
}

641
642
#[tracing::instrument(skip_all)]
async fn embeddings(
643
    State(state): State<Arc<service_v2::State>>,
644
    headers: HeaderMap,
645
    Json(request): Json<NvCreateEmbeddingRequest>,
Ryan Olson's avatar
Ryan Olson committed
646
) -> Result<Response, ErrorResponse> {
647
648
649
    // return a 503 if the service is not ready
    check_ready(&state)?;

650
651
652
    let request_id = get_or_create_request_id(request.inner.user.as_deref(), &headers);
    let request = Context::with_id(request, request_id);
    let request_id = request.id().to_string();
653
654
655
656
657
658
659
660

    // Embeddings are typically not streamed, so we default to non-streaming
    let streaming = false;

    // todo - make the protocols be optional for model name
    // todo - when optional, if none, apply a default
    let model = &request.inner.model;

661
662
663
    // Create http_queue_guard early - tracks time waiting to be processed
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(model);

664
665
666
667
    // todo - error handling should be more robust
    let engine = state
        .manager()
        .get_embeddings_engine(model)
Ryan Olson's avatar
Ryan Olson committed
668
        .map_err(|_| ErrorMessage::model_not_found())?;
669
670
671
672
673
674
675

    // this will increment the inflight gauge for the model
    let mut inflight =
        state
            .metrics_clone()
            .create_inflight_guard(model, Endpoint::Embeddings, streaming);

676
677
    let mut response_collector = state.metrics_clone().create_response_collector(model);

678
679
680
681
    // issue the generate call on the engine
    let stream = engine
        .generate(request)
        .await
Ryan Olson's avatar
Ryan Olson committed
682
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate embeddings"))?;
683

684
685
686
687
688
689
690
691
692
693
694
    // Process stream to collect metrics and drop http_queue_guard on first token
    let mut http_queue_guard = Some(http_queue_guard);
    let stream = stream.inspect(move |response| {
        // Calls observe_response() on each token - drops http_queue_guard on first token
        process_response_and_observe_metrics(
            response,
            &mut response_collector,
            &mut http_queue_guard,
        );
    });

695
696
    // Embeddings are typically returned as a single response (non-streaming)
    // so we fold the stream into a single response
Ryan Olson's avatar
Ryan Olson committed
697
    let response = NvCreateEmbeddingResponse::from_annotated_stream(stream)
698
699
700
701
702
703
704
        .await
        .map_err(|e| {
            tracing::error!(
                "Failed to fold embeddings stream for {}: {:?}",
                request_id,
                e
            );
Ryan Olson's avatar
Ryan Olson committed
705
            ErrorMessage::internal_server_error("Failed to fold embeddings stream")
706
707
708
709
        })?;

    inflight.mark_ok();
    Ok(Json(response).into_response())
710
711
}

Ryan Olson's avatar
Ryan Olson committed
712
713
714
async fn handler_chat_completions(
    State((state, template)): State<(Arc<service_v2::State>, Option<RequestTemplate>)>,
    headers: HeaderMap,
715
    Json(mut request): Json<NvCreateChatCompletionRequest>,
Ryan Olson's avatar
Ryan Olson committed
716
717
718
719
) -> Result<Response, ErrorResponse> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

720
721
    request.nvext = apply_header_routing_overrides(request.nvext.take(), &headers);

Ryan Olson's avatar
Ryan Olson committed
722
723
724
725
726
727
    // create the context for the request
    let request_id = get_or_create_request_id(request.inner.user.as_deref(), &headers);
    let request = Context::with_id(request, request_id);
    let context = request.context();

    // create the connection handles
728
729
    let (mut connection_handle, stream_handle) =
        create_connection_monitor(context.clone(), Some(state.metrics_clone())).await;
Ryan Olson's avatar
Ryan Olson committed
730

731
732
733
734
735
736
737
738
739
    let response =
        tokio::spawn(chat_completions(state, template, request, stream_handle).in_current_span())
            .await
            .map_err(|e| {
                ErrorMessage::internal_server_error(&format!(
                    "Failed to await chat completions task: {:?}",
                    e,
                ))
            })?;
Ryan Olson's avatar
Ryan Olson committed
740
741
742
743
744
745
746
747

    // if we got here, then we will return a response and the potentially long running task has completed successfully
    // without need to be cancelled.
    connection_handle.disarm();

    response
}

748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
/// Checks if an Annotated event represents a backend error and extracts error information.
/// Returns Some((message, status_code)) if it's an error, None otherwise.
fn extract_backend_error_if_present<T: serde::Serialize>(
    event: &Annotated<T>,
) -> Option<(String, StatusCode)> {
    #[derive(serde::Deserialize)]
    struct ErrorPayload {
        message: Option<String>,
        code: Option<u16>,
    }

    // Check if event type is "error" (from postprocessor when FinishReason::Error is encountered)
    if let Some(event_type) = &event.event
        && event_type == "error"
    {
        let comment_str = event
            .comment
            .as_ref()
            .map(|c| c.join(", "))
            .unwrap_or_else(|| "Unknown error".to_string());

        // Try to parse comment as error JSON to extract status code
        if let Ok(error_payload) = serde_json::from_str::<ErrorPayload>(&comment_str) {
            let code = error_payload
                .code
                .and_then(|c| StatusCode::from_u16(c).ok())
                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
            let message = error_payload.message.unwrap_or(comment_str);
            return Some((message, code));
        }

        return Some((comment_str, StatusCode::INTERNAL_SERVER_ERROR));
    }

    // Check if the data payload itself contains an error structure with code >= 400
    if let Some(data) = &event.data
        && let Ok(json_value) = serde_json::to_value(data)
        && let Ok(error_payload) = serde_json::from_value::<ErrorPayload>(json_value.clone())
        && let Some(code_num) = error_payload.code
        && code_num >= 400
    {
        let code = StatusCode::from_u16(code_num).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
        let message = error_payload
            .message
            .unwrap_or_else(|| json_value.to_string());
        return Some((message, code));
    }

    // Check if comment contains error information (without event: error)
    if let Some(comments) = &event.comment
        && !comments.is_empty()
    {
        let comment_str = comments.join(", ");

        // Try to parse comment as error JSON with code >= 400
        if let Ok(error_payload) = serde_json::from_str::<ErrorPayload>(&comment_str)
            && let Some(code_num) = error_payload.code
            && code_num >= 400
        {
            let code = StatusCode::from_u16(code_num).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
            let message = error_payload.message.unwrap_or(comment_str);
            return Some((message, code));
        }

        // Comments present with no data AND no event type indicates error
        // (events with event types like "request_id" or "event.dynamo.test.sentinel" are annotations)
        if event.data.is_none() && event.event.is_none() {
            return Some((comment_str, StatusCode::INTERNAL_SERVER_ERROR));
        }
    }

    None
}

/// Checks if the first event in the stream is a backend error.
/// Returns Err(ErrorResponse) if error detected, Ok(stream) otherwise.
824
pub(super) async fn check_for_backend_error(
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
    mut stream: impl futures::Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>>
    + Send
    + Unpin
    + 'static,
) -> Result<
    impl futures::Stream<Item = Annotated<NvCreateChatCompletionStreamResponse>> + Send,
    ErrorResponse,
> {
    use futures::stream::StreamExt;

    // Peek at the first event
    if let Some(first_event) = stream.next().await {
        // Check if it's an error event
        if let Some((error_msg, status_code)) = extract_backend_error_if_present(&first_event) {
            return Err((
                status_code,
                Json(ErrorMessage {
                    message: error_msg,
                    error_type: map_error_code_to_error_type(status_code),
                    code: status_code.as_u16(),
                }),
            ));
        }

        // Not an error - reconstruct stream with first event
        let reconstructed_stream = futures::stream::iter(vec![first_event]).chain(stream);
        Ok(reconstructed_stream)
    } else {
        // Empty stream - this shouldn't happen but handle gracefully
        Ok(futures::stream::iter(vec![]).chain(stream))
    }
}

858
859
860
861
862
863
864
865
866
/// OpenAI Chat Completions Request Handler
///
/// This method will handle the incoming request for the /v1/chat/completions endpoint. The endpoint is a "source"
/// for an [`super::OpenAIChatCompletionsStreamingEngine`] and will return a stream of responses which will be
/// forward to the client.
///
/// Note: For all requests, streaming or non-streaming, we always call the engine with streaming enabled. For
/// non-streaming requests, we will fold the stream into a single response as part of this handler.
async fn chat_completions(
Ryan Olson's avatar
Ryan Olson committed
867
868
869
870
871
    state: Arc<service_v2::State>,
    template: Option<RequestTemplate>,
    mut request: Context<NvCreateChatCompletionRequest>,
    mut stream_handle: ConnectionHandle,
) -> Result<Response, ErrorResponse> {
872
873
874
    // return a 503 if the service is not ready
    check_ready(&state)?;

Ryan Olson's avatar
Ryan Olson committed
875
876
    let request_id = request.id().to_string();

877
878
879
    // Handle unsupported fields - if Some(resp) is returned by
    // validate_chat_completion_unsupported_fields,
    // then a field was used that is unsupported. We will log an error message
880
    // and early return a 501 NOT_IMPLEMENTED status code. Otherwise, proceed.
Ryan Olson's avatar
Ryan Olson committed
881
    validate_chat_completion_unsupported_fields(&request)?;
882

883
884
885
    // Handle required fields like messages shouldn't be empty.
    validate_chat_completion_required_fields(&request)?;

886
887
888
    // Validate stream_options is only used when streaming (NVBug 5662680)
    validate_chat_completion_stream_options(&request)?;

889
890
891
    // Handle Rest of Validation Errors
    validate_chat_completion_fields_generic(&request)?;

892
893
894
895
896
897
898
899
900
901
902
903
    // Apply template values if present
    if let Some(template) = template {
        if request.inner.model.is_empty() {
            request.inner.model = template.model.clone();
        }
        if request.inner.temperature.unwrap_or(0.0) == 0.0 {
            request.inner.temperature = Some(template.temperature);
        }
        if request.inner.max_completion_tokens.unwrap_or(0) == 0 {
            request.inner.max_completion_tokens = Some(template.max_completion_tokens);
        }
    }
Ryan Olson's avatar
Ryan Olson committed
904
    tracing::trace!("Received chat completions request: {:?}", request.content());
905
906

    // todo - decide on default
Paul Hendricks's avatar
Paul Hendricks committed
907
    let streaming = request.inner.stream.unwrap_or(false);
908
909
910
911

    // todo - make the protocols be optional for model name
    // todo - when optional, if none, apply a default
    // todo - determine the proper error code for when a request model is not present
912
913
914
915
916
    let model = request.inner.model.clone();

    // Create HTTP queue guard after template resolution so labels are correct
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

917
918
919
    tracing::trace!("Getting chat completions engine for model: {}", model);

    let engine = state
920
        .manager()
921
        .get_chat_completions_engine(&model)
Ryan Olson's avatar
Ryan Olson committed
922
        .map_err(|_| ErrorMessage::model_not_found())?;
923

924
    let parsing_options = state.manager().get_parsing_options(&model);
925

926
927
928
929
930
    let mut response_collector = state.metrics_clone().create_response_collector(&model);

    let annotations = request.annotations();

    // Create inflight_guard before calling engine to ensure errors are counted
931
    let mut inflight_guard =
932
933
        state
            .metrics_clone()
934
            .create_inflight_guard(&model, Endpoint::ChatCompletions, streaming);
935
936
937
938
939

    // issue the generate call on the engine
    let stream = engine
        .generate(request)
        .await
Ryan Olson's avatar
Ryan Olson committed
940
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate completions"))?;
941
942
943
944

    // capture the context to cancel the stream if the client disconnects
    let ctx = stream.context();

Ryan Olson's avatar
Ryan Olson committed
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
    // prepare any requested annotations
    let annotations = annotations.map_or(Vec::new(), |annotations| {
        annotations
            .iter()
            .filter_map(|annotation| {
                if annotation == ANNOTATION_REQUEST_ID {
                    Annotated::from_annotation(ANNOTATION_REQUEST_ID, &request_id).ok()
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
    });

    // apply any annotations to the front of the stream
    let stream = stream::iter(annotations).chain(stream);

962
963
964
965
    // todo - tap the stream and propagate request level metrics
    // note - we might do this as part of the post processing set to make it more generic

    if streaming {
966
967
968
969
        // For streaming responses, we return HTTP 200 immediately without checking for errors.
        // Once HTTP 200 OK is sent, we cannot change the status code, so any backend errors
        // must be delivered as SSE events with `event: error` in the stream (handled by
        // EventConverter and monitor_for_disconnects). This is standard SSE behavior.
970
        stream_handle.arm(); // allows the system to detect client disconnects and cancel the LLM generation
Ryan Olson's avatar
Ryan Olson committed
971

972
        let mut http_queue_guard = Some(http_queue_guard);
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
        let stream = stream
            .map(move |response| {
                // Calls observe_response() on each token
                // EventConverter will detect `event: "error"` and convert to SSE error events
                process_response_using_event_converter_and_observe_metrics(
                    EventConverter::from(response),
                    &mut response_collector,
                    &mut http_queue_guard,
                )
            })
            .filter_map(|result| {
                use futures::future;
                // Transpose Result<Option<T>> -> Option<Result<T>>
                future::ready(result.transpose())
            });
Ryan Olson's avatar
Ryan Olson committed
988
        let stream = monitor_for_disconnects(stream, ctx, inflight_guard, stream_handle);
989

990
991
        let mut sse_stream = Sse::new(stream);

992
        if let Some(keep_alive) = state.sse_keep_alive() {
993
994
995
996
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
997
    } else {
998
999
1000
1001
1002
1003
1004
1005
1006
        // Check first event for backend errors before aggregating (non-streaming only)
        let stream_with_check =
            check_for_backend_error(stream)
                .await
                .map_err(|error_response| {
                    tracing::error!(request_id, "Backend error detected: {:?}", error_response);
                    error_response
                })?;

1007
        let mut http_queue_guard = Some(http_queue_guard);
1008
        let stream = stream_with_check.inspect(move |response| {
1009
1010
1011
1012
1013
1014
            // Calls observe_response() on each token - drops http_queue_guard on first token
            process_response_and_observe_metrics(
                response,
                &mut response_collector,
                &mut http_queue_guard,
            );
1015
1016
        });

1017
1018
1019
1020
1021
1022
        let response =
            NvCreateChatCompletionResponse::from_annotated_stream(stream, parsing_options.clone())
                .await
                .map_err(|e| {
                    tracing::error!(
                        request_id,
1023
                        "Failed to parse chat completion response: {:?}",
1024
1025
1026
                        e
                    );
                    ErrorMessage::internal_server_error(&format!(
1027
                        "Failed to parse chat completion response: {}",
1028
1029
1030
                        e
                    ))
                })?;
1031

1032
        inflight_guard.mark_ok();
1033
1034
1035
1036
        Ok(Json(response).into_response())
    }
}

1037
1038
1039
1040
1041
/// Checks for unsupported fields in the request.
/// Returns Some(response) if unsupported fields are present.
#[allow(deprecated)]
pub fn validate_chat_completion_unsupported_fields(
    request: &NvCreateChatCompletionRequest,
Ryan Olson's avatar
Ryan Olson committed
1042
) -> Result<(), ErrorResponse> {
1043
1044
1045
    let inner = &request.inner;

    if inner.function_call.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1046
        return Err(ErrorMessage::not_implemented_error(
1047
1048
            VALIDATION_PREFIX.to_string()
                + "`function_call` is deprecated. Please migrate to use `tool_choice` instead.",
1049
1050
1051
1052
        ));
    }

    if inner.functions.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1053
        return Err(ErrorMessage::not_implemented_error(
1054
1055
            VALIDATION_PREFIX.to_string()
                + "`functions` is deprecated. Please migrate to use `tools` instead.",
1056
1057
1058
        ));
    }

Ryan Olson's avatar
Ryan Olson committed
1059
    Ok(())
1060
1061
}

1062
1063
1064
1065
1066
1067
1068
1069
1070
/// Validates that required fields are present and valid in the chat completion request
pub fn validate_chat_completion_required_fields(
    request: &NvCreateChatCompletionRequest,
) -> Result<(), ErrorResponse> {
    let inner = &request.inner;

    if inner.messages.is_empty() {
        return Err(ErrorMessage::from_http_error(HttpError {
            code: 400,
1071
1072
            message: VALIDATION_PREFIX.to_string()
                + "The 'messages' field cannot be empty. At least one message is required.",
1073
1074
1075
1076
1077
1078
        }));
    }

    Ok(())
}

1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
/// Validates that stream_options is only used when stream=true for chat completions (NVBug 5662680)
pub fn validate_chat_completion_stream_options(
    request: &NvCreateChatCompletionRequest,
) -> Result<(), ErrorResponse> {
    let inner = &request.inner;
    let streaming = inner.stream.unwrap_or(false);
    if !streaming && inner.stream_options.is_some() {
        return Err(ErrorMessage::from_http_error(HttpError {
            code: 400,
            message: VALIDATION_PREFIX.to_string()
                + "The 'stream_options' field is only allowed when 'stream' is set to true.",
        }));
    }
    Ok(())
}

1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
/// Validates a chat completion request and returns an error response if validation fails.
///
/// This function calls the `validate` method implemented for `NvCreateChatCompletionRequest`.
/// If validation fails, it maps the error into an OpenAI-compatible error response.
pub fn validate_chat_completion_fields_generic(
    request: &NvCreateChatCompletionRequest,
) -> Result<(), ErrorResponse> {
    request.validate().map_err(|e| {
        ErrorMessage::from_http_error(HttpError {
            code: 400,
1105
            message: VALIDATION_PREFIX.to_string() + &e.to_string(),
1106
1107
1108
1109
        })
    })
}

1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
/// Validates that stream_options is only used when stream=true for completions (NVBug 5662680)
pub fn validate_completion_stream_options(
    request: &NvCreateCompletionRequest,
) -> Result<(), ErrorResponse> {
    let inner = &request.inner;
    let streaming = inner.stream.unwrap_or(false);
    if !streaming && inner.stream_options.is_some() {
        return Err(ErrorMessage::from_http_error(HttpError {
            code: 400,
            message: VALIDATION_PREFIX.to_string()
                + "The 'stream_options' field is only allowed when 'stream' is set to true.",
        }));
    }
    Ok(())
}

1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
/// Validates a completion request and returns an error response if validation fails.
///
/// This function calls the `validate` method implemented for `NvCreateCompletionRequest`.
/// If validation fails, it maps the error into an OpenAI-compatible error response.
pub fn validate_completion_fields_generic(
    request: &NvCreateCompletionRequest,
) -> Result<(), ErrorResponse> {
    request.validate().map_err(|e| {
        ErrorMessage::from_http_error(HttpError {
            code: 400,
1136
            message: VALIDATION_PREFIX.to_string() + &e.to_string(),
1137
1138
1139
1140
        })
    })
}

1141
1142
1143
/// OpenAI Responses Request Handler
///
/// This method will handle the incoming request for the /v1/responses endpoint.
Ryan Olson's avatar
Ryan Olson committed
1144
async fn handler_responses(
1145
    State((state, template)): State<(Arc<service_v2::State>, Option<RequestTemplate>)>,
Ryan Olson's avatar
Ryan Olson committed
1146
    headers: HeaderMap,
1147
    Json(mut request): Json<NvCreateResponse>,
Ryan Olson's avatar
Ryan Olson committed
1148
1149
1150
1151
) -> Result<Response, ErrorResponse> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

1152
1153
    request.nvext = apply_header_routing_overrides(request.nvext.take(), &headers);

Ryan Olson's avatar
Ryan Olson committed
1154
    // create the context for the request
1155
    let request_id = get_or_create_request_id(None, &headers);
Ryan Olson's avatar
Ryan Olson committed
1156
1157
1158
1159
    let request = Context::with_id(request, request_id);
    let context = request.context();

    // create the connection handles
1160
    let (mut connection_handle, stream_handle) =
1161
        create_connection_monitor(context.clone(), Some(state.metrics_clone())).await;
Ryan Olson's avatar
Ryan Olson committed
1162

1163
1164
1165
1166
1167
1168
1169
1170
1171
    let response =
        tokio::spawn(responses(state, template, request, stream_handle).in_current_span())
            .await
            .map_err(|e| {
                ErrorMessage::internal_server_error(&format!(
                    "Failed to await responses task: {:?}",
                    e,
                ))
            })?;
Ryan Olson's avatar
Ryan Olson committed
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184

    // if we got here, then we will return a response and the potentially long running task has completed successfully
    // without need to be cancelled.
    connection_handle.disarm();

    response
}

#[tracing::instrument(level = "debug", skip_all, fields(request_id = %request.id()))]
async fn responses(
    state: Arc<service_v2::State>,
    template: Option<RequestTemplate>,
    mut request: Context<NvCreateResponse>,
1185
    mut stream_handle: ConnectionHandle,
Ryan Olson's avatar
Ryan Olson committed
1186
) -> Result<Response, ErrorResponse> {
1187
1188
1189
    // return a 503 if the service is not ready
    check_ready(&state)?;

1190
    // Create http_queue_guard early - tracks time waiting to be processed
1191
1192
    // model is Option<String> in upstream; extract to String, defaulting to empty
    let model = request.inner.model.clone().unwrap_or_default();
1193
1194
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

1195
1196
    // Handle unsupported fields - if Some(resp) is returned by validate_unsupported_fields,
    // then a field was used that is unsupported. We will log an error message
1197
    // and early return a 501 NOT_IMPLEMENTED status code. Otherwise, proceed.
1198
    if let Some(resp) = validate_response_unsupported_fields(&request) {
1199
1200
1201
        return Ok(resp.into_response());
    }

1202
1203
1204
1205
1206
    // Apply template values if present, with sensible defaults for the Responses API.
    // Unlike chat completions where backends may have their own defaults, the Responses API
    // should provide a generous default to avoid truncated responses (especially with
    // reasoning models that emit <think> tokens).
    const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 4096;
1207
1208

    if let Some(template) = template {
1209
1210
        if request.inner.model.as_deref().unwrap_or("").is_empty() {
            request.inner.model = Some(template.model.clone());
1211
        }
1212
        if request.inner.temperature.is_none() {
1213
1214
            request.inner.temperature = Some(template.temperature);
        }
1215
        if request.inner.max_output_tokens.is_none() {
1216
1217
            request.inner.max_output_tokens = Some(template.max_completion_tokens);
        }
1218
1219
    } else if request.inner.max_output_tokens.is_none() {
        request.inner.max_output_tokens = Some(DEFAULT_MAX_OUTPUT_TOKENS);
1220
    }
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
    tracing::trace!("Received responses request: {:?}", request.inner);

    // Extract request parameters before into_parts() consumes the request.
    // These are echoed back in the Response object per the OpenAI spec.
    let response_params = ResponseParams {
        temperature: request.inner.temperature,
        top_p: request.inner.top_p,
        max_output_tokens: request.inner.max_output_tokens,
        store: request.inner.store,
        tools: request.inner.tools.clone(),
        tool_choice: request.inner.tool_choice.clone(),
        instructions: request.inner.instructions.clone(),
    };
1234

1235
    let streaming = request.inner.stream.unwrap_or(false);
Ryan Olson's avatar
Ryan Olson committed
1236
    let request_id = request.id().to_string();
1237
    let (orig_request, context) = request.into_parts();
1238

1239
1240
    let mut chat_request: NvCreateChatCompletionRequest =
        orig_request.try_into().map_err(|e: anyhow::Error| {
1241
1242
1243
1244
1245
1246
1247
            tracing::error!(
                request_id,
                error = %e,
                "Failed to convert NvCreateResponse to NvCreateChatCompletionRequest",
            );
            ErrorMessage::not_implemented_error(
                VALIDATION_PREFIX.to_string()
1248
                    + "Failed to convert responses request: "
1249
1250
1251
                    + &e.to_string(),
            )
        })?;
1252

1253
1254
1255
1256
1257
1258
1259
    // For non-streaming responses, we still use internal streaming for aggregation,
    // but we set the chat completion stream flag appropriately.
    if !streaming {
        chat_request.inner.stream = Some(true); // Internal streaming for aggregation
    }

    let request = context.map(|mut _req| chat_request);
Ryan Olson's avatar
Ryan Olson committed
1260

1261
1262
1263
1264
    tracing::trace!("Getting chat completions engine for model: {}", model);

    let engine = state
        .manager()
1265
        .get_chat_completions_engine(&model)
Ryan Olson's avatar
Ryan Olson committed
1266
        .map_err(|_| ErrorMessage::model_not_found())?;
1267

1268
    let parsing_options = state.manager().get_parsing_options(&model);
1269

1270
    let mut response_collector = state.metrics_clone().create_response_collector(&model);
1271

1272
    tracing::trace!("Issuing generate call for responses");
1273
1274

    // issue the generate call on the engine
1275
    let engine_stream = engine
1276
1277
        .generate(request)
        .await
Ryan Olson's avatar
Ryan Olson committed
1278
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate completions"))?;
1279

1280
1281
1282
    // Capture the context to cancel the stream if the client disconnects
    let ctx = engine_stream.context();

1283
1284
1285
1286
    // Create inflight_guard now that actual processing has begun
    let mut inflight_guard =
        state
            .metrics_clone()
1287
            .create_inflight_guard(&model, Endpoint::Responses, streaming);
1288

1289
1290
1291
1292
1293
    if streaming {
        // For streaming responses, we return HTTP 200 immediately without checking for errors.
        // Once HTTP 200 OK is sent, we cannot change the status code, so any backend errors
        // must be delivered as SSE events in the stream. This is standard SSE behavior.
        stream_handle.arm(); // allows the system to detect client disconnects and cancel the LLM generation
1294

1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
        // Streaming path: convert chat completion stream chunks to Responses API SSE events.
        // The engine yields Annotated<NvCreateChatCompletionStreamResponse>. We extract the
        // inner stream response data and convert it to Responses API events.
        use crate::protocols::openai::responses::stream_converter::ResponseStreamConverter;
        use std::sync::atomic::{AtomicBool, Ordering};

        let mut converter = ResponseStreamConverter::new(model.clone(), response_params);
        let start_events = converter.emit_start_events();

        // Use std::sync::Mutex (not tokio) since process_chunk/emit_end_events are
        // synchronous -- no .await while lock is held. Avoids async lock overhead per token.
        let converter = std::sync::Arc::new(std::sync::Mutex::new(converter));
        let converter_end = converter.clone();

        // Track whether the backend sent an error event during the stream.
        // Shared between event_stream (writer) and done_stream (reader).
        let saw_error = std::sync::Arc::new(AtomicBool::new(false));
        let saw_error_end = saw_error.clone();

        let mut http_queue_guard = Some(http_queue_guard);

        // Process each annotated chunk: extract the stream response data, convert to events
        let event_stream = engine_stream
            .inspect(move |response| {
                process_response_and_observe_metrics(
                    response,
                    &mut response_collector,
                    &mut http_queue_guard,
                );
            })
            .filter_map(move |annotated_chunk| {
                let converter = converter.clone();
                let saw_error = saw_error.clone();
                async move {
                    // Check for backend error before extracting data.
                    // Error events have data: None and event: Some("error").
                    if annotated_chunk.data.is_none() {
                        if annotated_chunk.event.as_deref() == Some("error") {
                            saw_error.store(true, Ordering::Release);
                        }
                        return None;
                    }
                    let stream_resp = annotated_chunk.data?;
                    let mut conv = converter.lock().expect("converter lock poisoned");
                    let events = conv.process_chunk(&stream_resp);
                    Some(stream::iter(events))
                }
            })
            .flatten();

        // Chain: start_events -> chunk_events -> end_events
        let start_stream = stream::iter(start_events);

        let done_stream = stream::once(async move {
            let mut conv = converter_end.lock().expect("converter lock poisoned");
            let end_events = if saw_error_end.load(Ordering::Acquire) {
                conv.emit_error_events()
            } else {
                conv.emit_end_events()
            };
            stream::iter(end_events)
        })
        .flatten();

        let full_stream = start_stream.chain(event_stream).chain(done_stream);

        let full_stream = full_stream.map(|result| result.map_err(axum::Error::new));

        // Wrap with disconnect monitoring: detects client disconnects, cancels generation,
        // and defers inflight_guard.mark_ok() until the stream completes.
        let stream = monitor_for_disconnects(full_stream, ctx, inflight_guard, stream_handle);

        let mut sse_stream = Sse::new(stream);
        if let Some(keep_alive) = state.sse_keep_alive() {
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
    } else {
        // Non-streaming path: aggregate stream into single response

        // Check first event for backend errors before aggregating (non-streaming only)
        let stream_with_check =
            check_for_backend_error(engine_stream)
                .await
                .map_err(|error_response| {
                    tracing::error!(request_id, "Backend error detected: {:?}", error_response);
                    error_response
                })?;

        let mut http_queue_guard = Some(http_queue_guard);
        let stream = stream_with_check.inspect(move |response| {
            process_response_and_observe_metrics(
                response,
                &mut response_collector,
                &mut http_queue_guard,
            );
        });

        let response =
            NvCreateChatCompletionResponse::from_annotated_stream(stream, parsing_options.clone())
                .await
                .map_err(|e| {
                    tracing::error!(request_id, "Failed to fold responses stream: {:?}", e);
                    ErrorMessage::internal_server_error(&format!(
                        "Failed to fold responses stream: {}",
                        e
                    ))
                })?;

        // Convert NvCreateChatCompletionResponse --> NvResponse
        let response: NvResponse = chat_completion_to_response(response, &response_params)
1407
1408
1409
            .map_err(|e| {
                tracing::error!(
                    request_id,
1410
                    "Failed to convert NvCreateChatCompletionResponse to NvResponse: {:?}",
1411
1412
                    e
                );
1413
                ErrorMessage::internal_server_error("Failed to convert internal response")
1414
            })?;
1415

1416
        inflight_guard.mark_ok();
1417

1418
        Ok(Json(response).into_response())
1419
1420
1421
1422
1423
    }
}

/// Checks for unsupported fields in the request.
/// Returns Some(response) if unsupported fields are present.
1424
1425
1426
pub fn validate_response_unsupported_fields(
    request: &NvCreateResponse,
) -> Option<impl IntoResponse> {
1427
1428
1429
    let inner = &request.inner;

    if inner.background == Some(true) {
Ryan Olson's avatar
Ryan Olson committed
1430
        return Some(ErrorMessage::not_implemented_error(
1431
            VALIDATION_PREFIX.to_string() + "`background: true` is not supported.",
1432
1433
1434
        ));
    }
    if inner.include.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1435
        return Some(ErrorMessage::not_implemented_error(
1436
            VALIDATION_PREFIX.to_string() + "`include` is not supported.",
1437
1438
1439
        ));
    }
    if inner.previous_response_id.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1440
        return Some(ErrorMessage::not_implemented_error(
1441
            VALIDATION_PREFIX.to_string() + "`previous_response_id` is not supported.",
1442
1443
1444
        ));
    }
    if inner.prompt.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1445
        return Some(ErrorMessage::not_implemented_error(
1446
            VALIDATION_PREFIX.to_string() + "`prompt` is not supported.",
1447
1448
1449
        ));
    }
    if inner.reasoning.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1450
        return Some(ErrorMessage::not_implemented_error(
1451
            VALIDATION_PREFIX.to_string() + "`reasoning` is not supported.",
1452
1453
1454
        ));
    }
    if inner.service_tier.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1455
        return Some(ErrorMessage::not_implemented_error(
1456
            VALIDATION_PREFIX.to_string() + "`service_tier` is not supported.",
1457
1458
1459
        ));
    }
    if inner.store == Some(true) {
Ryan Olson's avatar
Ryan Olson committed
1460
        return Some(ErrorMessage::not_implemented_error(
1461
            VALIDATION_PREFIX.to_string() + "`store: true` is not supported.",
1462
1463
1464
        ));
    }
    if inner.text.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1465
        return Some(ErrorMessage::not_implemented_error(
1466
            VALIDATION_PREFIX.to_string() + "`text` is not supported.",
1467
1468
1469
        ));
    }
    if inner.truncation.is_some() {
Ryan Olson's avatar
Ryan Olson committed
1470
        return Some(ErrorMessage::not_implemented_error(
1471
            VALIDATION_PREFIX.to_string() + "`truncation` is not supported.",
1472
1473
1474
1475
1476
        ));
    }
    None
}

1477
1478
// todo - abstract this to the top level lib.rs to be reused
// todo - move the service_observer to its own state/arc
Ryan Olson's avatar
Ryan Olson committed
1479
fn check_ready(_state: &Arc<service_v2::State>) -> Result<(), ErrorResponse> {
1480
    // if state.service_observer.stage() != ServiceStage::Ready {
Ryan Olson's avatar
Ryan Olson committed
1481
    //     return Err(ErrorMessage::service_unavailable());
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
    // }
    Ok(())
}

/// openai compatible format
/// Example:
/// {
///  "object": "list",
///  "data": [
///    {
///      "id": "model-id-0",
///      "object": "model",
///      "created": 1686935002,
///      "owned_by": "organization-owner"
///    },
///    ]
/// }
async fn list_models_openai(
1500
    State(state): State<Arc<service_v2::State>>,
Ryan Olson's avatar
Ryan Olson committed
1501
) -> Result<Response, ErrorResponse> {
1502
1503
1504
1505
1506
1507
1508
1509
    check_ready(&state)?;

    let created = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();
    let mut data = Vec::new();

1510
1511
    let models: HashSet<String> = state.manager().model_display_names();
    for model_name in models {
1512
        data.push(ModelListing {
1513
            id: model_name.clone(),
1514
1515
1516
            object: "model", // Per OpenAI spec, this should be "model"
            created,
            owned_by: "nvidia".to_string(),
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
        });
    }

    let out = ListModelOpenAI {
        object: "list",
        data,
    };
    Ok(Json(out).into_response())
}

#[derive(Serialize)]
struct ListModelOpenAI {
    object: &'static str, // always "list"
    data: Vec<ModelListing>,
}

#[derive(Serialize)]
struct ModelListing {
    id: String,
1536
1537
    object: &'static str, // always "model" per OpenAI spec
    created: u64,         // Seconds since epoch
1538
1539
1540
1541
1542
1543
    owned_by: String,
}

/// Create an Axum [`Router`] for the OpenAI API Completions endpoint
/// If not path is provided, the default path is `/v1/completions`
pub fn completions_router(
1544
    state: Arc<service_v2::State>,
1545
1546
1547
1548
1549
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/completions".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
Ryan Olson's avatar
Ryan Olson committed
1550
        .route(&path, post(handler_completions))
1551
        .layer(middleware::from_fn(smart_json_error_middleware))
1552
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
1553
1554
1555
1556
1557
1558
1559
        .with_state(state);
    (vec![doc], router)
}

/// Create an Axum [`Router`] for the OpenAI API Chat Completions endpoint
/// If not path is provided, the default path is `/v1/chat/completions`
pub fn chat_completions_router(
1560
    state: Arc<service_v2::State>,
1561
    template: Option<RequestTemplate>,
1562
1563
1564
1565
1566
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/chat/completions".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
Ryan Olson's avatar
Ryan Olson committed
1567
        .route(&path, post(handler_chat_completions))
1568
        .layer(middleware::from_fn(smart_json_error_middleware))
1569
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
1570
        .with_state((state, template));
1571
1572
1573
    (vec![doc], router)
}

1574
1575
1576
/// Create an Axum [`Router`] for the OpenAI API Embeddings endpoint
/// If not path is provided, the default path is `/v1/embeddings`
pub fn embeddings_router(
1577
    state: Arc<service_v2::State>,
1578
1579
1580
1581
1582
1583
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/embeddings".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(embeddings))
1584
        .layer(middleware::from_fn(smart_json_error_middleware))
1585
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
1586
1587
1588
1589
        .with_state(state);
    (vec![doc], router)
}

1590
1591
/// List Models
pub fn list_models_router(
1592
    state: Arc<service_v2::State>,
1593
1594
1595
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    // Standard OpenAI compatible list models endpoint
1596
    let openai_path = path.unwrap_or("/v1/models".to_string());
1597
1598
1599
1600
1601
1602
    let doc_for_openai = RouteDoc::new(axum::http::Method::GET, &openai_path);

    let router = Router::new()
        .route(&openai_path, get(list_models_openai))
        .with_state(state);

1603
    (vec![doc_for_openai], router)
1604
1605
}

1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
/// Create an Axum [`Router`] for the OpenAI API Responses endpoint
/// If not path is provided, the default path is `/v1/responses`
pub fn responses_router(
    state: Arc<service_v2::State>,
    template: Option<RequestTemplate>,
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/responses".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
Ryan Olson's avatar
Ryan Olson committed
1616
        .route(&path, post(handler_responses))
1617
        .layer(middleware::from_fn(smart_json_error_middleware))
1618
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
1619
1620
1621
1622
        .with_state((state, template));
    (vec![doc], router)
}

1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
async fn images(
    State(state): State<Arc<service_v2::State>>,
    headers: HeaderMap,
    Json(request): Json<NvCreateImageRequest>,
) -> Result<Response, ErrorResponse> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

    let request_id = get_or_create_request_id(request.inner.user.as_deref(), &headers);
    let request = Context::with_id(request, request_id);
    let request_id = request.id().to_string();

    // Images are typically not streamed, so we default to non-streaming
    let streaming = false;

    // Get the model name from the request (diffusion model)
    let model = request
        .inner
        .model
        .as_ref()
        .map(|m| match m {
            dynamo_async_openai::types::ImageModel::DallE2 => "dall-e-2".to_string(),
            dynamo_async_openai::types::ImageModel::DallE3 => "dall-e-3".to_string(),
            dynamo_async_openai::types::ImageModel::Other(s) => s.clone(),
        })
        .unwrap_or_else(|| "diffusion".to_string());

    // Create http_queue_guard early - tracks time waiting to be processed
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

    // Get the image generation engine
    let engine = state
        .manager()
        .get_images_engine(&model)
        .map_err(|_| ErrorMessage::model_not_found())?;

    // this will increment the inflight gauge for the model
    let mut inflight =
        state
            .metrics_clone()
            .create_inflight_guard(&model, Endpoint::Images, streaming);

    let mut response_collector = state.metrics_clone().create_response_collector(&model);

    // Issue the generate call on the engine
    // Note: This uses ServerStreamingEngine for internal routing/distribution,
    // NOT for client-facing SSE streaming. The stream is immediately folded into
    // a single response below.
    let stream = engine
        .generate(request)
        .await
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate images"))?;

    // Process stream to collect metrics and drop http_queue_guard on first response
    let mut http_queue_guard = Some(http_queue_guard);
    let stream = stream.inspect(move |response| {
        // Calls observe_response() on each item - drops http_queue_guard on first item
        process_response_and_observe_metrics(
            response,
            &mut response_collector,
            &mut http_queue_guard,
        );
    });

    // Images are returned as a single response (non-streaming to client)
    // Fold the internal stream into a single response
    let response = NvImagesResponse::from_annotated_stream(stream)
        .await
        .map_err(|e| {
            tracing::error!("Failed to fold images stream for {}: {:?}", request_id, e);
            ErrorMessage::internal_server_error("Failed to fold images stream")
        })?;

    inflight.mark_ok();
    Ok(Json(response).into_response())
}

/// Create an Axum [`Router`] for the OpenAI API Images endpoint
/// If not path is provided, the default path is `/v1/images/generations`
pub fn images_router(
    state: Arc<service_v2::State>,
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/images/generations".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(images))
        .layer(middleware::from_fn(smart_json_error_middleware))
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
        .with_state(state);
    (vec![doc], router)
}

1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
async fn videos(
    State(state): State<Arc<service_v2::State>>,
    headers: HeaderMap,
    Json(request): Json<NvCreateVideoRequest>,
) -> Result<Response, ErrorResponse> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

    let request_id = get_or_create_request_id(request.user.as_deref(), &headers);
    let request = Context::with_id(request, request_id);
    let request_id = request.id().to_string();

    // Videos are typically not streamed, so we default to non-streaming
    let streaming = false;

    // Get the model name from the request (video generation model)
    let model = request.model.clone();

    // Create http_queue_guard early - tracks time waiting to be processed
    let http_queue_guard = state.metrics_clone().create_http_queue_guard(&model);

    // Get the video generation engine
    let engine = state
        .manager()
        .get_videos_engine(&model)
        .map_err(|_| ErrorMessage::model_not_found())?;

    // this will increment the inflight gauge for the model
    let mut inflight =
        state
            .metrics_clone()
            .create_inflight_guard(&model, Endpoint::Videos, streaming);

    let mut response_collector = state.metrics_clone().create_response_collector(&model);

    // issue the generate call on the engine
    let stream = engine
        .generate(request)
        .await
        .map_err(|e| ErrorMessage::from_anyhow(e, "Failed to generate videos"))?;

    // Process stream to collect metrics and drop http_queue_guard on first token
    let mut http_queue_guard = Some(http_queue_guard);
    let stream = stream.inspect(move |response| {
        // Calls observe_response() on each token - drops http_queue_guard on first token
        process_response_and_observe_metrics(
            response,
            &mut response_collector,
            &mut http_queue_guard,
        );
    });

    // Videos are typically returned as a single response (non-streaming)
    // so we fold the stream into a single response
    let response = NvVideosResponse::from_annotated_stream(stream)
        .await
        .map_err(|e| {
            tracing::error!("Failed to fold videos stream for {}: {:?}", request_id, e);
            ErrorMessage::internal_server_error("Failed to fold videos stream")
        })?;

    inflight.mark_ok();
    Ok(Json(response).into_response())
}

/// Create an Axum [`Router`] for the OpenAI API Videos endpoint
/// If no path is provided, the default path is `/v1/videos`
pub fn videos_router(
    state: Arc<service_v2::State>,
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    let path = path.unwrap_or("/v1/videos".to_string());
    let doc = RouteDoc::new(axum::http::Method::POST, &path);
    let router = Router::new()
        .route(&path, post(videos))
        .layer(middleware::from_fn(smart_json_error_middleware))
        .layer(axum::extract::DefaultBodyLimit::max(get_body_limit()))
        .with_state(state);
    (vec![doc], router)
}

1797
1798
#[cfg(test)]
mod tests {
1799

1800
1801
1802
1803
1804
1805
    use super::*;
    use crate::discovery::ModelManagerError;
    use crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest;
    use crate::protocols::openai::common_ext::CommonExt;
    use crate::protocols::openai::completions::NvCreateCompletionRequest;
    use crate::protocols::openai::responses::NvCreateResponse;
1806
    use dynamo_async_openai::types::responses::{
1807
1808
        CreateResponse, IncludeEnum, Input, PromptConfig, ServiceTier, TextConfig,
        TextResponseFormat, Truncation,
1809
    };
1810
    use dynamo_async_openai::types::{
1811
1812
        ChatCompletionRequestMessage, ChatCompletionRequestUserMessage,
        ChatCompletionRequestUserMessageContent, CreateChatCompletionRequest,
1813
        CreateCompletionRequest,
1814
    };
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825

    const BACKUP_ERROR_MESSAGE: &str = "Failed to generate completions";

    fn http_error_from_engine(code: u16) -> Result<(), anyhow::Error> {
        Err(HttpError {
            code,
            message: "custom error message".to_string(),
        })?
    }

    fn other_error_from_engine() -> Result<(), anyhow::Error> {
1826
        Err(ModelManagerError::ModelNotFound("foo".to_string()))?
1827
1828
    }

1829
1830
1831
1832
    fn make_base_request() -> NvCreateResponse {
        NvCreateResponse {
            inner: CreateResponse {
                input: Input::Text("hello".into()),
1833
1834
                model: Some("test-model".into()),
                ..Default::default()
1835
1836
1837
1838
1839
            },
            nvext: None,
        }
    }

1840
1841
1842
    #[test]
    fn test_http_error_response_from_anyhow() {
        let err = http_error_from_engine(400).unwrap_err();
1843
1844
1845
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::BAD_REQUEST);
        assert_eq!(response.1.message, "custom error message");
1846
1847
1848
1849
1850
    }

    #[test]
    fn test_error_response_from_anyhow_out_of_range() {
        let err = http_error_from_engine(399).unwrap_err();
1851
1852
1853
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.1.message, "custom error message");
1854
1855

        let err = http_error_from_engine(500).unwrap_err();
1856
1857
1858
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.1.message, "custom error message");
1859
1860

        let err = http_error_from_engine(501).unwrap_err();
1861
1862
1863
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.1.message, "custom error message");
1864
1865
1866
1867
1868
    }

    #[test]
    fn test_other_error_response_from_anyhow() {
        let err = other_error_from_engine().unwrap_err();
1869
1870
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::INTERNAL_SERVER_ERROR);
1871
        assert_eq!(
1872
            response.1.message,
1873
1874
1875
1876
1877
1878
1879
            format!(
                "{}: {}",
                BACKUP_ERROR_MESSAGE,
                other_error_from_engine().unwrap_err()
            )
        );
    }
1880

1881
1882
1883
1884
1885
1886
1887
1888
    #[test]
    fn test_service_overloaded_error_response_from_anyhow() {
        use dynamo_runtime::pipeline::error::PipelineError;

        let err: anyhow::Error = PipelineError::ServiceOverloaded(
            "All workers are busy, please retry later".to_string(),
        )
        .into();
1889
1890
        let response = ErrorMessage::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(response.0, StatusCode::SERVICE_UNAVAILABLE);
1891
        assert_eq!(
1892
            response.1.message,
1893
1894
1895
1896
            "Service temporarily unavailable: All workers are busy, please retry later"
        );
    }

1897
1898
1899
    #[test]
    fn test_validate_unsupported_fields_accepts_clean_request() {
        let request = make_base_request();
1900
        let result = validate_response_unsupported_fields(&request);
1901
1902
1903
        assert!(result.is_none());
    }

1904
1905
1906
1907
1908
1909
1910
1911
    #[test]
    fn test_validate_unsupported_fields_accepts_parallel_tool_calls() {
        let mut request = make_base_request();
        request.inner.parallel_tool_calls = Some(true);
        let result = validate_response_unsupported_fields(&request);
        assert!(result.is_none(), "parallel_tool_calls should be supported");
    }

1912
1913
1914
1915
1916
1917
1918
    #[test]
    fn test_validate_unsupported_fields_detects_flags() {
        #[allow(clippy::type_complexity)]
        let unsupported_cases: Vec<(&str, Box<dyn FnOnce(&mut CreateResponse)>)> = vec![
            ("background", Box::new(|r| r.background = Some(true))),
            (
                "include",
1919
                Box::new(|r| r.include = Some(vec![IncludeEnum::FileSearchCallResults])),
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
            ),
            (
                "previous_response_id",
                Box::new(|r| r.previous_response_id = Some("prev-id".into())),
            ),
            (
                "prompt",
                Box::new(|r| {
                    r.prompt = Some(PromptConfig {
                        id: "template-id".into(),
                        version: None,
                        variables: None,
                    })
                }),
            ),
            (
                "reasoning",
                Box::new(|r| r.reasoning = Some(Default::default())),
            ),
            (
                "service_tier",
                Box::new(|r| r.service_tier = Some(ServiceTier::Auto)),
            ),
            ("store", Box::new(|r| r.store = Some(true))),
            (
                "text",
                Box::new(|r| {
                    r.text = Some(TextConfig {
                        format: TextResponseFormat::Text,
1949
                        verbosity: None,
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
                    })
                }),
            ),
            (
                "truncation",
                Box::new(|r| r.truncation = Some(Truncation::Auto)),
            ),
        ];

        for (field, set_field) in unsupported_cases {
            let mut req = make_base_request();
            (set_field)(&mut req.inner);
1962
            let result = validate_response_unsupported_fields(&req);
1963
1964
1965
            assert!(result.is_some(), "Expected rejection for `{field}`");
        }
    }
1966
1967
1968
1969
1970
1971
1972
1973
1974

    #[test]
    fn test_validate_chat_completion_required_fields_empty_messages() {
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![],
                ..Default::default()
            },
1975
            common: Default::default(),
1976
            nvext: None,
1977
            chat_template_args: None,
1978
            media_io_kwargs: None,
1979
            unsupported_fields: Default::default(),
1980
1981
1982
        };
        let result = validate_chat_completion_required_fields(&request);
        assert!(result.is_err());
1983
1984
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
1985
            assert_eq!(
1986
                error_response.1.message,
1987
1988
1989
                format!(
                    "{VALIDATION_PREFIX}The 'messages' field cannot be empty. At least one message is required."
                )
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
            );
        }
    }

    #[test]
    fn test_validate_chat_completion_required_fields_with_messages() {
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                ..Default::default()
            },
2007
            common: Default::default(),
2008
            nvext: None,
2009
            chat_template_args: None,
2010
            media_io_kwargs: None,
2011
            unsupported_fields: Default::default(),
2012
2013
2014
2015
        };
        let result = validate_chat_completion_required_fields(&request);
        assert!(result.is_ok());
    }
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031

    #[test]
    // Test for all Bad Requests Example for Chat Completion
    // 1. Echo:  Should be a boolean : Not Done
    // 2. Frequency Penalty: Should be a float between -2.0 and 2.0 : Done
    // 3. logprobs: Done
    // 4. Model Format: Should be a string : Not Done
    // 5. Prompt or Messages Validation
    // 6. Max Tokens: Should be a positive integer
    // 7. Presence Penalty: Should be a float between -2.0 and 2.0 : Done
    // 8. Stop : Should be a string or an array of strings : Not Done
    // 9. Invalid or Out of range temperature: Done
    // 10.Invalid or out of range top_p: Done
    // 11. Repetition Penalty: Should be a float between 0.0 and 2.0 : Done
    // 12. Logprobs: Should be a positive integer between 0 and 5 : Done
    // invalid or non existing user : Only empty string is not allowed validation is there. How can we check non-extisting user ?
2032
    // Unknown fields : Done (rejected via extra_fields catch-all)
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
    // guided_whitespace_pattern null or invalid : Not Done
    // "response_format": { "type": "invalid_format" } : Not Done
    // "logit_bias": { "invalid_token": "not_a_number" }, : Partial Validation is already there
    fn test_bad_base_request_for_completion() {
        // Frequency Penalty: Should be a float between -2.0 and 2.0
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                frequency_penalty: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2047
            metadata: None,
2048
            unsupported_fields: Default::default(),
2049
2050
2051
2052
        };

        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2053
2054
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2055
            assert_eq!(
2056
                error_response.1.message,
2057
                format!("{VALIDATION_PREFIX}Frequency penalty must be between -2 and 2, got -3")
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
            );
        }

        // Presence Penalty: Should be a float between -2.0 and 2.0
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                presence_penalty: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2071
            metadata: None,
2072
            unsupported_fields: Default::default(),
2073
2074
2075
        };
        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2076
2077
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2078
            assert_eq!(
2079
                error_response.1.message,
2080
                format!("{VALIDATION_PREFIX}Presence penalty must be between -2 and 2, got -3")
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
            );
        }

        // Temperature: Should be a float between 0.0 and 2.0
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                temperature: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2094
            metadata: None,
2095
            unsupported_fields: Default::default(),
2096
2097
2098
        };
        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2099
2100
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2101
            assert_eq!(
2102
                error_response.1.message,
2103
                format!("{VALIDATION_PREFIX}Temperature must be between 0 and 2, got -3")
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
            );
        }

        // Top P: Should be a float between 0.0 and 1.0
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                top_p: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2117
            metadata: None,
2118
            unsupported_fields: Default::default(),
2119
2120
2121
        };
        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2122
2123
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2124
            assert_eq!(
2125
                error_response.1.message,
2126
                format!("{VALIDATION_PREFIX}Top_p must be between 0 and 1, got -3")
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
            );
        }

        // Repetition Penalty: Should be a float between 0.0 and 2.0
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                ..Default::default()
            },
            common: CommonExt::builder()
                .repetition_penalty(-3.0)
                .build()
                .unwrap(),
            nvext: None,
2142
            metadata: None,
2143
            unsupported_fields: Default::default(),
2144
2145
2146
        };
        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2147
2148
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2149
            assert_eq!(
2150
                error_response.1.message,
2151
                format!("{VALIDATION_PREFIX}Repetition penalty must be between 0 and 2, got -3")
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
            );
        }

        // Logprobs: Should be a positive integer between 0 and 5
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                logprobs: Some(6),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2165
            metadata: None,
2166
            unsupported_fields: Default::default(),
2167
2168
2169
        };
        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
2170
2171
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2172
            assert_eq!(
2173
                error_response.1.message,
2174
                format!("{VALIDATION_PREFIX}Logprobs must be between 0 and 5, got 6")
2175
2176
2177
2178
            );
        }
    }

2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
    #[test]
    fn test_metadata_field_nested() {
        use serde_json::json;

        // Test metadata field with nested object
        let request = NvCreateCompletionRequest {
            inner: CreateCompletionRequest {
                model: "test-model".to_string(),
                prompt: "Hello".into(),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
            metadata: json!({
                "user": {"id": 1, "name": "user-1"},
                "session": {"id": "session-1", "timestamp": 1640995200}
            })
            .into(),
2197
            unsupported_fields: Default::default(),
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
        };

        let result = validate_completion_fields_generic(&request);
        assert!(result.is_ok());

        // Verify metadata is accessible
        assert!(request.metadata.is_some());
        assert_eq!(request.metadata.as_ref().unwrap()["user"]["id"], 1);
    }

2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
    #[test]
    fn test_bad_base_request_for_chatcompletion() {
        // Frequency Penalty: Should be a float between -2.0 and 2.0
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                frequency_penalty: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2225
            chat_template_args: None,
2226
            media_io_kwargs: None,
2227
            unsupported_fields: Default::default(),
2228
2229
2230
2231
        };

        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2232
2233
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2234
            assert_eq!(
2235
                error_response.1.message,
2236
                format!("{VALIDATION_PREFIX}Frequency penalty must be between -2 and 2, got -3")
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
            );
        }

        // Presence Penalty: Should be a float between -2.0 and 2.0
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                presence_penalty: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2255
            chat_template_args: None,
2256
            media_io_kwargs: None,
2257
            unsupported_fields: Default::default(),
2258
2259
2260
        };
        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2261
2262
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2263
            assert_eq!(
2264
                error_response.1.message,
2265
                format!("{VALIDATION_PREFIX}Presence penalty must be between -2 and 2, got -3")
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
            );
        }

        // Temperature: Should be a float between 0.0 and 2.0
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                temperature: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2284
            chat_template_args: None,
2285
            media_io_kwargs: None,
2286
            unsupported_fields: Default::default(),
2287
2288
2289
        };
        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2290
2291
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2292
            assert_eq!(
2293
                error_response.1.message,
2294
                format!("{VALIDATION_PREFIX}Temperature must be between 0 and 2, got -3")
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
            );
        }

        // Top P: Should be a float between 0.0 and 1.0
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                top_p: Some(-3.0),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2313
            chat_template_args: None,
2314
            media_io_kwargs: None,
2315
            unsupported_fields: Default::default(),
2316
2317
2318
        };
        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2319
2320
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2321
            assert_eq!(
2322
                error_response.1.message,
2323
                format!("{VALIDATION_PREFIX}Top_p must be between 0 and 1, got -3")
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
            );
        }

        // Repetition Penalty: Should be a float between 0.0 and 2.0
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                ..Default::default()
            },
            common: CommonExt::builder()
                .repetition_penalty(-3.0)
                .build()
                .unwrap(),
            nvext: None,
2344
            chat_template_args: None,
2345
            media_io_kwargs: None,
2346
            unsupported_fields: Default::default(),
2347
2348
2349
        };
        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2350
2351
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2352
            assert_eq!(
2353
                error_response.1.message,
2354
                format!("{VALIDATION_PREFIX}Repetition penalty must be between 0 and 2, got -3")
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
            );
        }

        // Top Logprobs: Should be a positive integer between 0 and 20
        let request = NvCreateChatCompletionRequest {
            inner: CreateChatCompletionRequest {
                model: "test-model".to_string(),
                messages: vec![ChatCompletionRequestMessage::User(
                    ChatCompletionRequestUserMessage {
                        content: ChatCompletionRequestUserMessageContent::Text("Hello".to_string()),
                        name: None,
                    },
                )],
                top_logprobs: Some(25),
                ..Default::default()
            },
            common: Default::default(),
            nvext: None,
2373
            chat_template_args: None,
2374
            media_io_kwargs: None,
2375
            unsupported_fields: Default::default(),
2376
2377
2378
        };
        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
2379
2380
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
2381
            assert_eq!(
2382
                error_response.1.message,
2383
                format!("{VALIDATION_PREFIX}Top_logprobs must be between 0 and 20, got 25")
2384
2385
2386
            );
        }
    }
2387
2388

    #[test]
2389
2390
    fn test_chat_completions_unknown_fields_rejected() {
        // Test that known unsupported fields are rejected and all shown in error message
2391
2392
2393
2394
2395
        let json = r#"{
            "messages": [{"role": "user", "content": "Hello"}],
            "model": "test-model",
            "add_special_tokens": true,
            "documents": ["doc1"],
2396
            "chat_template": "custom"
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
        }"#;

        let request: NvCreateChatCompletionRequest = serde_json::from_str(json).unwrap();

        // Verify all unsupported fields were captured
        assert!(
            request
                .unsupported_fields
                .contains_key("add_special_tokens")
        );
        assert!(request.unsupported_fields.contains_key("documents"));
        assert!(request.unsupported_fields.contains_key("chat_template"));

        let result = validate_chat_completion_fields_generic(&request);
        assert!(result.is_err());
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
            let msg = &error_response.1.message;
            assert!(msg.contains("Unsupported parameter"));
            // Verify all fields appear in the error message
            assert!(msg.contains("add_special_tokens"));
            assert!(msg.contains("documents"));
            assert!(msg.contains("chat_template"));
        }
    }
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453

    #[test]
    fn test_completions_unsupported_fields_rejected() {
        // Test that known unsupported fields are rejected and all shown in error message
        let json = r#"{
            "model": "test-model",
            "prompt": "Hello",
            "add_special_tokens": true,
            "response_format": {"type": "json_object"}
        }"#;

        let request: NvCreateCompletionRequest = serde_json::from_str(json).unwrap();

        // Verify both unsupported fields were captured
        assert!(
            request
                .unsupported_fields
                .contains_key("add_special_tokens")
        );
        assert!(request.unsupported_fields.contains_key("response_format"));

        let result = validate_completion_fields_generic(&request);
        assert!(result.is_err());
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::BAD_REQUEST);
            let msg = &error_response.1.message;
            assert!(msg.contains("Unsupported parameter"));
            // Verify both fields appear in error message
            assert!(msg.contains("add_special_tokens"));
            assert!(msg.contains("response_format"));
        }
    }
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585

    #[tokio::test]
    async fn test_check_for_backend_error_with_error_event() {
        use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
        use futures::stream;

        // Create an error event
        let error_event = Annotated::<NvCreateChatCompletionStreamResponse> {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: Some(vec!["Backend service unavailable".to_string()]),
        };

        let test_stream = stream::iter(vec![error_event]);
        let result = check_for_backend_error(test_stream).await;

        // Should return an error
        assert!(result.is_err());
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::INTERNAL_SERVER_ERROR);
            assert_eq!(error_response.1.message, "Backend service unavailable");
        }
    }

    #[tokio::test]
    async fn test_check_for_backend_error_with_json_error_and_code() {
        use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
        use futures::stream;

        // Create an error event with JSON payload containing error code in comment
        let error_json =
            r#"{"message":"prompt > max_seq_len","type":"Internal Server Error","code":500}"#;
        let error_event = Annotated::<NvCreateChatCompletionStreamResponse> {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: Some(vec![error_json.to_string()]),
        };

        let test_stream = stream::iter(vec![error_event]);
        let result = check_for_backend_error(test_stream).await;

        // Should return an error with correct status code extracted from JSON
        assert!(result.is_err());
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::INTERNAL_SERVER_ERROR);
            assert_eq!(error_response.1.message, "prompt > max_seq_len");
            assert_eq!(error_response.1.code, 500);
        }
    }

    #[tokio::test]
    async fn test_check_for_backend_error_with_normal_event() {
        use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
        use dynamo_async_openai::types::CreateChatCompletionStreamResponse;
        use futures::stream::{self, StreamExt};

        // Create a normal data event
        let normal_event = Annotated::<NvCreateChatCompletionStreamResponse> {
            data: Some(CreateChatCompletionStreamResponse {
                id: "test-id".to_string(),
                choices: vec![],
                created: 0,
                model: "test-model".to_string(),
                system_fingerprint: None,
                object: "chat.completion.chunk".to_string(),
                service_tier: None,
                usage: None,
                nvext: None,
            }),
            id: Some("msg-1".to_string()),
            event: None,
            comment: None,
        };

        let test_stream = stream::iter(vec![normal_event.clone()]);
        let result = check_for_backend_error(test_stream).await;

        // Should return Ok with the stream
        assert!(result.is_ok());
        let mut returned_stream = result.unwrap();

        // Verify we can read the event back from the stream
        let first = returned_stream.next().await;
        assert!(first.is_some());
        let first_event = first.unwrap();
        assert_eq!(first_event.id, Some("msg-1".to_string()));
    }

    #[tokio::test]
    async fn test_check_for_backend_error_with_empty_stream() {
        use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
        use futures::stream::{self, StreamExt};

        // Create an empty stream
        let test_stream =
            stream::iter::<Vec<Annotated<NvCreateChatCompletionStreamResponse>>>(vec![]);
        let result = check_for_backend_error(test_stream).await;

        // Should return Ok with an empty stream
        assert!(result.is_ok());
        let mut returned_stream = result.unwrap();

        // Verify stream is empty
        let first = returned_stream.next().await;
        assert!(first.is_none());
    }

    #[tokio::test]
    async fn test_check_for_backend_error_with_comment_but_no_event_type() {
        use crate::types::openai::chat_completions::NvCreateChatCompletionStreamResponse;
        use futures::stream;

        // Create an event with comment but no event type and no data (error indicator)
        let error_event = Annotated::<NvCreateChatCompletionStreamResponse> {
            data: None,
            id: None,
            event: None,
            comment: Some(vec!["Connection timeout".to_string()]),
        };

        let test_stream = stream::iter(vec![error_event]);
        let result = check_for_backend_error(test_stream).await;

        // Should return an error based on is_backend_error_event logic
        assert!(result.is_err());
        if let Err(error_response) = result {
            assert_eq!(error_response.0, StatusCode::INTERNAL_SERVER_ERROR);
            assert_eq!(error_response.1.message, "Connection timeout");
        }
    }
2586
}