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

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

11
12
13
14
15
16
17
18
19
20
use axum::{
    extract::State,
    http::StatusCode,
    response::{
        sse::{Event, KeepAlive, Sse},
        IntoResponse, Response,
    },
    routing::{get, post},
    Json, Router,
};
21
use dynamo_runtime::pipeline::{AsyncEngineContext, Context};
22
23
24
25
26
27
use futures::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use tokio_stream::wrappers::ReceiverStream;

use super::{
    error::HttpError,
28
    metrics::{Endpoint, InflightGuard, ResponseMetricCollector},
29
    service_v2, RouteDoc,
30
};
31
use crate::preprocessor::LLMMetricAnnotation;
32
use crate::protocols::openai::{
33
34
35
36
    chat_completions::{NvCreateChatCompletionRequest, NvCreateChatCompletionResponse},
    completions::{NvCreateCompletionRequest, NvCreateCompletionResponse},
    embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse},
    responses::{NvCreateResponse, NvResponse},
37
};
38
use crate::request_template::RequestTemplate;
39
use crate::types::Annotated;
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

#[derive(Serialize, Deserialize)]
pub(crate) struct ErrorResponse {
    error: String,
}

impl ErrorResponse {
    /// Not Found Error
    pub fn model_not_found() -> (StatusCode, Json<ErrorResponse>) {
        (
            StatusCode::NOT_FOUND,
            Json(ErrorResponse {
                error: "Model not found".to_string(),
            }),
        )
    }

    /// Service Unavailable
    /// This is returned when the service is live, but not ready.
    pub fn _service_unavailable() -> (StatusCode, Json<ErrorResponse>) {
        (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(ErrorResponse {
                error: "Service is not ready".to_string(),
            }),
        )
    }

    /// 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.
    pub fn internal_server_error(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
        tracing::error!("Internal server error: {msg}");
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(ErrorResponse {
                error: msg.to_string(),
            }),
        )
    }

82
83
84
85
86
87
88
89
90
91
92
93
94
    /// 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.
    pub fn not_implemented_error(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
        tracing::error!("Not Implemented error: {msg}");
        (
            StatusCode::NOT_IMPLEMENTED,
            Json(ErrorResponse {
                error: msg.to_string(),
            }),
        )
    }

Neelay Shah's avatar
Neelay Shah committed
95
    /// The OAI endpoints call an [`dynamo.runtime::engine::AsyncEngine`] which are specialized to return
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
    /// an [`anyhow::Error`]. This method will convert the [`anyhow::Error`] into an [`HttpError`].
    /// If successful, it will return the [`HttpError`] as an [`ErrorResponse::internal_server_error`]
    /// with the details of the error.
    pub fn from_anyhow(err: anyhow::Error, alt_msg: &str) -> (StatusCode, Json<ErrorResponse>) {
        match err.downcast::<HttpError>() {
            Ok(http_error) => ErrorResponse::from_http_error(http_error),
            Err(err) => ErrorResponse::internal_server_error(&format!("{alt_msg}: {err}")),
        }
    }

    /// Implementers should only be able to throw 400-499 errors.
    pub fn from_http_error(err: HttpError) -> (StatusCode, Json<ErrorResponse>) {
        if err.code < 400 || err.code >= 500 {
            return ErrorResponse::internal_server_error(&err.message);
        }
        match StatusCode::from_u16(err.code) {
            Ok(code) => (code, Json(ErrorResponse { error: err.message })),
            Err(_) => ErrorResponse::internal_server_error(&err.message),
        }
    }
}

impl From<HttpError> for ErrorResponse {
    fn from(err: HttpError) -> Self {
        ErrorResponse { error: err.message }
    }
}

/// 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.
#[tracing::instrument(skip_all)]
async fn completions(
134
    State(state): State<Arc<service_v2::State>>,
135
    Json(request): Json<NvCreateCompletionRequest>,
136
137
138
139
140
141
142
143
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

    // todo - extract distributed tracing id and context id from headers
    let request_id = uuid::Uuid::new_v4().to_string();

    // todo - decide on default
144
    let streaming = request.inner.stream.unwrap_or(false);
145
146

    // update the request to always stream
147
    let inner = async_openai::types::CreateCompletionRequest {
148
        stream: Some(true),
149
        ..request.inner
150
151
    };

152
    let request = NvCreateCompletionRequest {
153
154
155
        inner,
        nvext: request.nvext,
    };
156

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

    // todo - error handling should be more robust
    let engine = state
163
        .manager()
164
165
166
        .get_completions_engine(model)
        .map_err(|_| ErrorResponse::model_not_found())?;

167
    let mut inflight_guard =
168
169
170
        state
            .metrics_clone()
            .create_inflight_guard(model, Endpoint::Completions, streaming);
171

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

174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
    // setup context
    // todo - inherit request_id from distributed trace details
    let request = Context::with_id(request, request_id.clone());

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

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

    // 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 {
191
192
193
194
        let stream = stream.map(move |response| {
            process_event_converter(EventConverter::from(response), &mut response_collector)
        });
        let stream = monitor_for_disconnects(stream.boxed(), ctx, inflight_guard).await;
195

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

198
        if let Some(keep_alive) = state.sse_keep_alive() {
199
200
201
202
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
203
    } else {
204
        // TODO: report ISL/OSL for non-streaming requests
205
        let response = NvCreateCompletionResponse::from_annotated_stream(stream.into())
206
207
208
209
210
211
212
213
214
215
            .await
            .map_err(|e| {
                tracing::error!(
                    "Failed to fold completions stream for {}: {:?}",
                    request_id,
                    e
                );
                ErrorResponse::internal_server_error("Failed to fold completions stream")
            })?;

216
        inflight_guard.mark_ok();
217
218
219
220
        Ok(Json(response).into_response())
    }
}

221
222
#[tracing::instrument(skip_all)]
async fn embeddings(
223
224
    State(state): State<Arc<service_v2::State>>,
    Json(request): Json<NvCreateEmbeddingRequest>,
225
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
    // return a 503 if the service is not ready
    check_ready(&state)?;

    // todo - extract distributed tracing id and context id from headers
    let request_id = uuid::Uuid::new_v4().to_string();

    // 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;

    // todo - error handling should be more robust
    let engine = state
        .manager()
        .get_embeddings_engine(model)
        .map_err(|_| ErrorResponse::model_not_found())?;

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

    // setup context
    // todo - inherit request_id from distributed trace details
    let request = Context::with_id(request, request_id.clone());

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

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

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

278
279
280
281
282
283
284
285
286
287
/// 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.
#[tracing::instrument(skip_all)]
async fn chat_completions(
288
    State((state, template)): State<(Arc<service_v2::State>, Option<RequestTemplate>)>,
289
    Json(mut request): Json<NvCreateChatCompletionRequest>,
290
291
292
293
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

294
295
296
297
298
299
300
301
302
303
304
305
306
307
    // 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);
        }
    }
    tracing::trace!("Received chat completions request: {:?}", request.inner);

308
309
310
311
    // todo - extract distributed tracing id and context id from headers
    let request_id = uuid::Uuid::new_v4().to_string();

    // todo - decide on default
Paul Hendricks's avatar
Paul Hendricks committed
312
    let streaming = request.inner.stream.unwrap_or(false);
313
314

    // update the request to always stream
Paul Hendricks's avatar
Paul Hendricks committed
315
    let inner_request = async_openai::types::CreateChatCompletionRequest {
316
        stream: Some(true),
Paul Hendricks's avatar
Paul Hendricks committed
317
318
        ..request.inner
    };
319

320
    let request = NvCreateChatCompletionRequest {
Paul Hendricks's avatar
Paul Hendricks committed
321
        inner: inner_request,
322
        nvext: request.nvext,
323
324
325
326
    };

    // todo - make the protocols be optional for model name
    // todo - when optional, if none, apply a default
Paul Hendricks's avatar
Paul Hendricks committed
327
    let model = &request.inner.model;
328
329
330
331
332

    // todo - determine the proper error code for when a request model is not present
    tracing::trace!("Getting chat completions engine for model: {}", model);

    let engine = state
333
        .manager()
334
335
336
        .get_chat_completions_engine(model)
        .map_err(|_| ErrorResponse::model_not_found())?;

337
    let mut inflight_guard =
338
339
340
        state
            .metrics_clone()
            .create_inflight_guard(model, Endpoint::ChatCompletions, streaming);
341

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

344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
    // setup context
    // todo - inherit request_id from distributed trace details
    let request = Context::with_id(request, request_id.clone());

    tracing::trace!("Issuing generate call for chat completions");

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

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

    // 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 {
363
364
365
366
        let stream = stream.map(move |response| {
            process_event_converter(EventConverter::from(response), &mut response_collector)
        });
        let stream = monitor_for_disconnects(stream.boxed(), ctx, inflight_guard).await;
367

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

370
        if let Some(keep_alive) = state.sse_keep_alive() {
371
372
373
374
            sse_stream = sse_stream.keep_alive(KeepAlive::default().interval(keep_alive));
        }

        Ok(sse_stream.into_response())
375
    } else {
376
        // TODO: report ISL/OSL for non-streaming requests
377
        let response = NvCreateChatCompletionResponse::from_annotated_stream(stream.into())
378
379
380
381
382
383
384
385
386
387
388
389
390
            .await
            .map_err(|e| {
                tracing::error!(
                    request_id,
                    "Failed to fold chat completions stream for: {:?}",
                    e
                );
                ErrorResponse::internal_server_error(&format!(
                    "Failed to fold chat completions stream: {}",
                    e
                ))
            })?;

391
        inflight_guard.mark_ok();
392
393
394
395
        Ok(Json(response).into_response())
    }
}

396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
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
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
/// OpenAI Responses Request Handler
///
/// This method will handle the incoming request for the /v1/responses endpoint.
#[tracing::instrument(skip_all)]
async fn responses(
    State((state, template)): State<(Arc<service_v2::State>, Option<RequestTemplate>)>,
    Json(mut request): Json<NvCreateResponse>,
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
    // return a 503 if the service is not ready
    check_ready(&state)?;

    // 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
    // and early return a 501 NOT_IMPLEMENTED status code. Otherwise, proceeed.
    if let Some(resp) = validate_unsupported_fields(&request) {
        return Ok(resp.into_response());
    }

    // Handle non-text (image, audio, file) inputs - if Some(resp) is returned by
    // validate_input_is_text_only, then we are handling something other than Input::Text(_).
    // We will log an error message and early return a 501 NOT_IMPLEMENTED status code.
    // Otherwise, proceeed.
    if let Some(resp) = validate_input_is_text_only(&request) {
        return Ok(resp.into_response());
    }

    // 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_output_tokens.unwrap_or(0) == 0 {
            request.inner.max_output_tokens = Some(template.max_completion_tokens);
        }
    }
    tracing::trace!("Received chat completions request: {:?}", request.inner);

    let request_id = uuid::Uuid::new_v4().to_string();

    // Convert NvCreateResponse --> NvCreateChatCompletionRequest
    let request: NvCreateChatCompletionRequest = request.try_into().map_err(|e| {
        tracing::error!(
            request_id,
            "Failed to convert NvCreateResponse to NvCreateChatCompletionRequest: {:?}",
            e
        );
        ErrorResponse::not_implemented_error(&format!(
            "Only Input::Text(_) is currently supported: {}",
            e
        ))
    })?;

    let model = &request.inner.model;

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

    let engine = state
        .manager()
        .get_chat_completions_engine(model)
        .map_err(|_| ErrorResponse::model_not_found())?;

    let mut inflight_guard =
        state
            .metrics_clone()
            .create_inflight_guard(model, Endpoint::Responses, false);

    let _response_collector = state.metrics_clone().create_response_collector(model);

    let request = Context::with_id(request, request_id.clone());

    tracing::trace!("Issuing generate call for chat completions");

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

    // TODO: handle streaming, currently just unary
    let response = NvCreateChatCompletionResponse::from_annotated_stream(stream.into())
        .await
        .map_err(|e| {
            tracing::error!(
                request_id,
                "Failed to fold chat completions stream for: {:?}",
                e
            );
            ErrorResponse::internal_server_error(&format!(
                "Failed to fold chat completions stream: {}",
                e
            ))
        })?;

    // Convert NvCreateChatCompletionResponse --> NvResponse
    let response: NvResponse = response.try_into().map_err(|e| {
        tracing::error!(
            request_id,
            "Failed to convert NvCreateChatCompletionResponse to NvResponse: {:?}",
            e
        );
        ErrorResponse::internal_server_error("Failed to convert internal response")
    })?;

    inflight_guard.mark_ok();

    Ok(Json(response).into_response())
}

pub fn validate_input_is_text_only(request: &NvCreateResponse) -> Option<impl IntoResponse> {
    match &request.inner.input {
        async_openai::types::responses::Input::Text(_) => None,
        _ => Some(ErrorResponse::not_implemented_error("Only `Input::Text` is supported. Structured, multimedia, or custom input types are not yet implemented.")),
    }
}

/// Checks for unsupported fields in the request.
/// Returns Some(response) if unsupported fields are present.
pub fn validate_unsupported_fields(request: &NvCreateResponse) -> Option<impl IntoResponse> {
    let inner = &request.inner;

    if inner.background == Some(true) {
        return Some(ErrorResponse::not_implemented_error(
            "`background: true` is not supported.",
        ));
    }
    if inner.include.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`include` is not supported.",
        ));
    }
    if inner.instructions.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`instructions` is not supported.",
        ));
    }
    if inner.max_tool_calls.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`max_tool_calls` is not supported.",
        ));
    }
    if inner.metadata.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`metadata` is not supported.",
        ));
    }
    if inner.parallel_tool_calls == Some(true) {
        return Some(ErrorResponse::not_implemented_error(
            "`parallel_tool_calls: true` is not supported.",
        ));
    }
    if inner.previous_response_id.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`previous_response_id` is not supported.",
        ));
    }
    if inner.prompt.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`prompt` is not supported.",
        ));
    }
    if inner.reasoning.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`reasoning` is not supported.",
        ));
    }
    if inner.service_tier.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`service_tier` is not supported.",
        ));
    }
    if inner.store == Some(true) {
        return Some(ErrorResponse::not_implemented_error(
            "`store: true` is not supported.",
        ));
    }
    if inner.stream == Some(true) {
        return Some(ErrorResponse::not_implemented_error(
            "`stream: true` is not supported.",
        ));
    }
    if inner.text.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`text` is not supported.",
        ));
    }
    if inner.tool_choice.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`tool_choice` is not supported.",
        ));
    }
    if inner.tools.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`tools` is not supported.",
        ));
    }
    if inner.truncation.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`truncation` is not supported.",
        ));
    }
    if inner.user.is_some() {
        return Some(ErrorResponse::not_implemented_error(
            "`user` is not supported.",
        ));
    }

    None
}

608
609
// todo - abstract this to the top level lib.rs to be reused
// todo - move the service_observer to its own state/arc
610
fn check_ready(_state: &Arc<service_v2::State>) -> Result<(), (StatusCode, Json<ErrorResponse>)> {
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    // if state.service_observer.stage() != ServiceStage::Ready {
    //     return Err(ErrorResponse::service_unavailable());
    // }
    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(
631
    State(state): State<Arc<service_v2::State>>,
632
633
634
635
636
637
638
639
640
) -> Result<Response, (StatusCode, Json<ErrorResponse>)> {
    check_ready(&state)?;

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

641
642
    let models: HashSet<String> = state.manager().model_display_names();
    for model_name in models {
643
        data.push(ModelListing {
644
            id: model_name.clone(),
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
            object: "object",
            created,                        // Where would this come from? The GGUF?
            owned_by: "nvidia".to_string(), // Get organization from GGUF
        });
    }

    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,
    object: &'static str, // always "object"
    created: u64,         //  Seconds since epoch
    owned_by: String,
}

/// This method will consume a stream of SSE events and forward them to a new stream defined by a tokio channel.
/// In this way, if the downstream is dropped, then the upstream will be unable to send any more events. This is
/// how we can monitor for disconnects and stop the generation of completions.
///
/// If a disconnect is detected, then the context will issue a `stop_generating` call to the context which will
/// propagate the cancellation signal to the backend.
async fn monitor_for_disconnects(
    stream: Pin<
        Box<dyn Stream<Item = Result<axum::response::sse::Event, axum::Error>> + std::marker::Send>,
    >,
    context: Arc<dyn AsyncEngineContext>,
683
    mut inflight_guard: InflightGuard,
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
) -> ReceiverStream<Result<Event, axum::Error>> {
    let (tx, rx) = tokio::sync::mpsc::channel(8);

    tokio::spawn(async move {
        let mut stream = stream;
        while let Some(event) = stream.next().await {
            let event = match event {
                Ok(event) => Ok(event),
                Err(err) => Ok(Event::default().event("error").comment(err.to_string())),
            };

            if (tx.send(event).await).is_err() {
                tracing::trace!("Forwarding SSE stream was dropped; breaking loop");
                context.stop_generating();
                break;
            }
        }

702
        // Stream completed successfully - mark as ok
703
        if tx.send(Ok(Event::default().data("[DONE]"))).await.is_ok() {
704
            inflight_guard.mark_ok();
705
706
707
708
709
710
711
712
713
714
715
716
717
718
        }
    });

    ReceiverStream::new(rx)
}

struct EventConverter<T>(Annotated<T>);

impl<T> From<Annotated<T>> for EventConverter<T> {
    fn from(annotated: Annotated<T>) -> Self {
        EventConverter(annotated)
    }
}

719
720
721
722
fn process_event_converter<T: Serialize>(
    annotated: EventConverter<T>,
    response_collector: &mut ResponseMetricCollector,
) -> Result<Event, axum::Error> {
723
    let mut annotated = annotated.0;
724

725
726
727
728
    // update metrics
    if let Ok(Some(metrics)) = LLMMetricAnnotation::from_annotation(&annotated) {
        response_collector.observe_current_osl(metrics.output_tokens);
        response_collector.observe_response(metrics.input_tokens, metrics.chunk_tokens);
729
730
731
732
733
734
735

        // Chomp the LLMMetricAnnotation so it's not returned in the response stream
        // TODO: add a flag to control what is returned in the SSE stream
        if annotated.event.as_deref() == Some(crate::preprocessor::ANNOTATION_LLM_METRICS) {
            annotated.event = None;
            annotated.comment = None;
        }
736
737
    }

738
    let mut event = Event::default();
739

740
741
742
    if let Some(data) = annotated.data {
        event = event.json_data(data)?;
    }
743

744
745
746
747
748
749
    if let Some(msg) = annotated.event {
        if msg == "error" {
            let msgs = annotated
                .comment
                .unwrap_or_else(|| vec!["unspecified error".to_string()]);
            return Err(axum::Error::new(msgs.join(" -- ")));
750
        }
751
752
        event = event.event(msg);
    }
753

754
755
756
757
    if let Some(comments) = annotated.comment {
        for comment in comments {
            event = event.comment(comment);
        }
758
    }
759
760

    Ok(event)
761
762
763
764
765
}

/// 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(
766
    state: Arc<service_v2::State>,
767
768
769
770
771
772
773
774
775
776
777
778
779
    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()
        .route(&path, post(completions))
        .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(
780
    state: Arc<service_v2::State>,
781
    template: Option<RequestTemplate>,
782
783
784
785
786
787
    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()
        .route(&path, post(chat_completions))
788
        .with_state((state, template));
789
790
791
    (vec![doc], router)
}

792
793
794
/// 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(
795
    state: Arc<service_v2::State>,
796
797
798
799
800
801
802
803
804
805
    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))
        .with_state(state);
    (vec![doc], router)
}

806
807
/// List Models
pub fn list_models_router(
808
    state: Arc<service_v2::State>,
809
810
811
    path: Option<String>,
) -> (Vec<RouteDoc>, Router) {
    // Standard OpenAI compatible list models endpoint
812
    let openai_path = path.unwrap_or("/v1/models".to_string());
813
814
815
816
817
818
    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);

819
    (vec![doc_for_openai], router)
820
821
}

822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
/// 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()
        .route(&path, post(responses))
        .with_state((state, template));
    (vec![doc], router)
}

837
838
#[cfg(test)]
mod tests {
839
840
841
842
843
844
845
    use std::collections::HashMap;

    use async_openai::types::responses::{
        CreateResponse, Input, InputContent, InputItem, InputMessage, PromptConfig,
        Role as ResponseRole, ServiceTier, TextConfig, TextResponseFormat, ToolChoice,
        ToolChoiceMode, Truncation,
    };
846
847

    use super::*;
848
849
    use crate::discovery::ModelManagerError;
    use crate::protocols::openai::responses::NvCreateResponse;
850
851
852
853
854
855
856
857
858
859
860

    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> {
861
        Err(ModelManagerError::ModelNotFound("foo".to_string()))?
862
863
    }

864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
    fn make_base_request() -> NvCreateResponse {
        NvCreateResponse {
            inner: CreateResponse {
                input: Input::Text("hello".into()),
                model: "test-model".into(),
                background: None,
                include: None,
                instructions: None,
                max_output_tokens: None,
                max_tool_calls: None,
                metadata: None,
                parallel_tool_calls: None,
                previous_response_id: None,
                prompt: None,
                reasoning: None,
                service_tier: None,
                store: None,
                stream: None,
                text: None,
                tool_choice: None,
                tools: None,
                truncation: None,
                user: None,
                temperature: None,
                top_logprobs: None,
                top_p: None,
            },
            nvext: None,
        }
    }

895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
    #[test]
    fn test_http_error_response_from_anyhow() {
        let err = http_error_from_engine(400).unwrap_err();
        let (status, response) = ErrorResponse::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(status, StatusCode::BAD_REQUEST);
        assert_eq!(response.error, "custom error message");
    }

    #[test]
    fn test_error_response_from_anyhow_out_of_range() {
        let err = http_error_from_engine(399).unwrap_err();
        let (status, response) = ErrorResponse::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.error, "custom error message");

        let err = http_error_from_engine(500).unwrap_err();
        let (status, response) = ErrorResponse::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.error, "custom error message");

        let err = http_error_from_engine(501).unwrap_err();
        let (status, response) = ErrorResponse::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(response.error, "custom error message");
    }

    #[test]
    fn test_other_error_response_from_anyhow() {
        let err = other_error_from_engine().unwrap_err();
        let (status, response) = ErrorResponse::from_anyhow(err, BACKUP_ERROR_MESSAGE);
        assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR);
        assert_eq!(
            response.error,
            format!(
                "{}: {}",
                BACKUP_ERROR_MESSAGE,
                other_error_from_engine().unwrap_err()
            )
        );
    }
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031

    #[test]
    fn test_validate_input_is_text_only_accepts_text() {
        let request = make_base_request();
        let result = validate_input_is_text_only(&request);
        assert!(result.is_none());
    }

    #[test]
    fn test_validate_input_is_text_only_rejects_items() {
        let mut request = make_base_request();
        request.inner.input = Input::Items(vec![InputItem::Message(InputMessage {
            kind: Default::default(),
            role: ResponseRole::User,
            content: InputContent::TextInput("structured".into()),
        })]);
        let result = validate_input_is_text_only(&request);
        assert!(result.is_some());
    }

    #[test]
    fn test_validate_unsupported_fields_accepts_clean_request() {
        let request = make_base_request();
        let result = validate_unsupported_fields(&request);
        assert!(result.is_none());
    }

    #[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",
                Box::new(|r| r.include = Some(vec!["file_search_call.results".into()])),
            ),
            (
                "instructions",
                Box::new(|r| r.instructions = Some("System prompt".into())),
            ),
            ("max_tool_calls", Box::new(|r| r.max_tool_calls = Some(3))),
            ("metadata", Box::new(|r| r.metadata = Some(HashMap::new()))),
            (
                "parallel_tool_calls",
                Box::new(|r| r.parallel_tool_calls = Some(true)),
            ),
            (
                "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))),
            ("stream", Box::new(|r| r.stream = Some(true))),
            (
                "text",
                Box::new(|r| {
                    r.text = Some(TextConfig {
                        format: TextResponseFormat::Text,
                    })
                }),
            ),
            (
                "tool_choice",
                Box::new(|r| r.tool_choice = Some(ToolChoice::Mode(ToolChoiceMode::Required))),
            ),
            ("tools", Box::new(|r| r.tools = Some(vec![]))),
            (
                "truncation",
                Box::new(|r| r.truncation = Some(Truncation::Auto)),
            ),
            ("user", Box::new(|r| r.user = Some("user-id".into()))),
        ];

        for (field, set_field) in unsupported_cases {
            let mut req = make_base_request();
            (set_field)(&mut req.inner);
            let result = validate_unsupported_fields(&req);
            assert!(result.is_some(), "Expected rejection for `{field}`");
        }
    }
1032
}