Initial ForgeLM import

crates/inference-server/src/routes.rs

@@ -0,0 +1,463 @@
+use axum::{
+    extract::{Json, State},
+    http::{HeaderMap, StatusCode},
+    response::{IntoResponse, Response, Sse},
+    routing::{get, post},
+    Router,
+};
+use std::{sync::Arc, time::SystemTime};
+use tokio::sync::Mutex;
+
+use api_contracts::{
+    chat::{
+        ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, Choice,
+        ChunkChoice, Delta, FinishReason, Role, Usage,
+    },
+    error::{ApiError, ErrorResponse},
+    health::HealthResponse,
+    models::ModelList,
+};
+use inference_core::InferenceEngine;
+use metrics::counters::PerfCounters;
+
+use crate::compat::{
+    authorize_request, build_sampler, ensure_supported_mode, finalize_assistant_message,
+    prepare_prompt, tool_calls_to_delta, validate_request, FinalizedAssistant, ServerConfig,
+    StreamingMode,
+};
+
+#[derive(Clone)]
+pub struct AppState {
+    pub engine: Arc<Mutex<InferenceEngine>>,
+    pub counters: Arc<PerfCounters>,
+    pub model_name: Arc<String>,
+    pub server_config: Arc<ServerConfig>,
+    pub start_time: std::time::Instant,
+}
+
+pub fn build_router(state: AppState) -> Router {
+    let cors = state.server_config.cors_layer();
+    Router::new()
+        .route("/health", get(health_handler))
+        .route("/v1/models", get(models_handler))
+        .route("/v1/chat/completions", post(chat_completions_handler))
+        .layer(cors)
+        .with_state(state)
+}
+
+fn unix_now() -> i64 {
+    SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs() as i64
+}
+
+fn error_response(status: StatusCode, error: ApiError) -> Response {
+    (status, Json(ErrorResponse { error })).into_response()
+}
+
+async fn health_handler(State(state): State<AppState>) -> impl IntoResponse {
+    Json(HealthResponse {
+        status: "ok",
+        model_loaded: true,
+        model_id: Some(state.model_name.as_ref().clone()),
+        uptime_secs: state.start_time.elapsed().as_secs(),
+    })
+}
+
+async fn models_handler(State(state): State<AppState>, headers: HeaderMap) -> Response {
+    if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
+        return error_response(status, error);
+    }
+
+    Json(ModelList {
+        object: "list",
+        data: state.server_config.model_list(unix_now()),
+    })
+    .into_response()
+}
+
+async fn chat_completions_handler(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    Json(req): Json<ChatCompletionRequest>,
+) -> Response {
+    tracing::info!(
+        model = %req.model,
+        stream = req.stream,
+        messages = req.messages.len(),
+        max_tokens = ?req.max_tokens,
+        "chat_completions: request received"
+    );
+
+    if let Err((status, error)) = authorize_request(&headers, &state.server_config) {
+        tracing::warn!(status = %status, "chat_completions: authorization failed");
+        return error_response(status, error);
+    }
+
+    if let Err((status, error)) = validate_request(&req, &state.server_config) {
+        tracing::warn!(status = %status, "chat_completions: validation failed");
+        return error_response(status, error);
+    }
+
+    if let Err(err) = ensure_supported_mode(&req) {
+        tracing::warn!(err = %err, "chat_completions: unsupported mode");
+        return error_response(
+            StatusCode::BAD_REQUEST,
+            ApiError::new(err.to_string(), "invalid_request_error").with_code("unsupported_mode"),
+        );
+    }
+
+    let sampler = build_sampler(&req);
+    let prepared = prepare_prompt(&req, &state.server_config);
+    tracing::debug!(prompt_len = prepared.prompt.len(), "chat_completions: prompt prepared");
+
+    if req.stream {
+        tracing::info!("chat_completions: starting stream response");
+        stream_response(
+            state,
+            req,
+            sampler,
+            prepared.prompt,
+            prepared.streaming_mode,
+        )
+        .await
+    } else {
+        tracing::info!("chat_completions: starting blocking response");
+        blocking_response(state, req, sampler, prepared.prompt).await
+    }
+}
+
+async fn stream_response(
+    state: AppState,
+    req: ChatCompletionRequest,
+    sampler: inference_core::SamplerConfig,
+    prompt: String,
+    streaming_mode: StreamingMode,
+) -> Response {
+    use axum::response::sse::{Event, KeepAlive};
+
+    let id = uuid::Uuid::new_v4().to_string();
+    let created = unix_now();
+    let resolved_model = state
+        .server_config
+        .resolve_model(&req.model)
+        .map(|model| model.root.clone())
+        .unwrap_or_else(|| state.model_name.as_ref().clone());
+
+    let engine_arc = Arc::clone(&state.engine);
+    let counters_arc = Arc::clone(&state.counters);
+
+    let stream = async_stream::stream! {
+        let first_chunk = ChatCompletionChunk {
+            id: id.clone(),
+            object: "chat.completion.chunk",
+            created,
+            model: resolved_model.clone(),
+            choices: vec![ChunkChoice {
+                index: 0,
+                delta: Delta { role: Some(Role::Assistant), content: None, tool_calls: None },
+                finish_reason: None,
+            }],
+            usage: None,
+        };
+        yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&first_chunk).unwrap()));
+
+        match streaming_mode {
+            StreamingMode::LiveTokens => {
+                let (tx, mut rx) = tokio::sync::mpsc::channel::<StreamItem>(256);
+                let prompt = prompt.clone();
+
+                tokio::task::spawn_blocking(move || {
+                    let rt = tokio::runtime::Handle::current();
+                    tracing::debug!("stream: acquiring engine lock");
+                    let mut engine = rt.block_on(engine_arc.lock());
+                    tracing::debug!("stream: engine lock acquired, creating session");
+                    let mut record = metrics::recorder::GenerationRecord::new(0);
+
+                    let mut session = match engine.new_session(&prompt, sampler) {
+                        Ok(session) => {
+                            tracing::info!(prompt_tokens = session.prompt_tokens(), "stream: session created");
+                            session
+                        }
+                        Err(err) => {
+                            tracing::error!(err = %err, "Failed to create session");
+                            let _ = rt.block_on(tx.send(StreamItem::Error));
+                            return;
+                        }
+                    };
+
+                    while let Some(token) = session.next_token() {
+                        record.tick();
+                        counters_arc.inc_tokens(1);
+                        let _ = rt.block_on(tx.send(StreamItem::Token(token)));
+                    }
+
+                    let usage = Usage {
+                        prompt_tokens: session.prompt_tokens(),
+                        completion_tokens: record.completion_tokens,
+                        total_tokens: session.prompt_tokens() + record.completion_tokens,
+                        tokens_per_second: record.tokens_per_second(),
+                    };
+                    tracing::info!(
+                        prompt_tokens = usage.prompt_tokens,
+                        completion_tokens = usage.completion_tokens,
+                        tps = ?usage.tokens_per_second,
+                        "stream: generation complete"
+                    );
+                    counters_arc.inc_completed();
+                    let _ = rt.block_on(tx.send(StreamItem::Done {
+                        finish_reason: FinishReason::Stop,
+                        usage,
+                    }));
+                });
+
+                let mut streamed_tokens: u32 = 0;
+                while let Some(item) = rx.recv().await {
+                    match item {
+                        StreamItem::Token(token) => {
+                            streamed_tokens += 1;
+                            if streamed_tokens <= 3 || streamed_tokens % 50 == 0 {
+                                tracing::debug!(n = streamed_tokens, token = %token, "stream: token");
+                            }
+                            let chunk = ChatCompletionChunk {
+                                id: id.clone(),
+                                object: "chat.completion.chunk",
+                                created,
+                                model: resolved_model.clone(),
+                                choices: vec![ChunkChoice {
+                                    index: 0,
+                                    delta: Delta { role: None, content: Some(token), tool_calls: None },
+                                    finish_reason: None,
+                                }],
+                                usage: None,
+                            };
+                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
+                        }
+                        StreamItem::Done { finish_reason, usage } => {
+                            tracing::info!(streamed_tokens, "stream: sending final chunk with usage");
+                            let final_chunk = ChatCompletionChunk {
+                                id: id.clone(),
+                                object: "chat.completion.chunk",
+                                created,
+                                model: resolved_model.clone(),
+                                choices: vec![ChunkChoice {
+                                    index: 0,
+                                    delta: Delta { role: None, content: None, tool_calls: None },
+                                    finish_reason: Some(finish_reason),
+                                }],
+                                usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| usage),
+                            };
+                            yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
+                            yield Ok(Event::default().data("[DONE]"));
+                            break;
+                        }
+                        StreamItem::Error => {
+                            tracing::error!("stream: sending error chunk");
+                            let chunk = ChatCompletionChunk {
+                                id: id.clone(),
+                                object: "chat.completion.chunk",
+                                created,
+                                model: resolved_model.clone(),
+                                choices: vec![ChunkChoice {
+                                    index: 0,
+                                    delta: Delta { role: None, content: Some("internal error".to_owned()), tool_calls: None },
+                                    finish_reason: Some(FinishReason::Stop),
+                                }],
+                                usage: None,
+                            };
+                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
+                            yield Ok(Event::default().data("[DONE]"));
+                            break;
+                        }
+                    }
+                }
+            }
+            StreamingMode::Buffered => {
+                match generate_full_completion(engine_arc, counters_arc, prompt, sampler, req.clone()).await {
+                    Ok(result) => {
+                        if let Some(tool_calls) = result.assistant.tool_calls {
+                            let chunk = ChatCompletionChunk {
+                                id: id.clone(),
+                                object: "chat.completion.chunk",
+                                created,
+                                model: resolved_model.clone(),
+                                choices: vec![ChunkChoice {
+                                    index: 0,
+                                    delta: Delta {
+                                        role: None,
+                                        content: None,
+                                        tool_calls: Some(tool_calls_to_delta(&tool_calls)),
+                                    },
+                                    finish_reason: None,
+                                }],
+                                usage: None,
+                            };
+                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
+                        } else if let Some(content) = result.assistant.content {
+                            let chunk = ChatCompletionChunk {
+                                id: id.clone(),
+                                object: "chat.completion.chunk",
+                                created,
+                                model: resolved_model.clone(),
+                                choices: vec![ChunkChoice {
+                                    index: 0,
+                                    delta: Delta { role: None, content: Some(content), tool_calls: None },
+                                    finish_reason: None,
+                                }],
+                                usage: None,
+                            };
+                            yield Ok(Event::default().data(serde_json::to_string(&chunk).unwrap()));
+                        }
+
+                        let final_chunk = ChatCompletionChunk {
+                            id: id.clone(),
+                            object: "chat.completion.chunk",
+                            created,
+                            model: resolved_model.clone(),
+                            choices: vec![ChunkChoice {
+                                index: 0,
+                                delta: Delta { role: None, content: None, tool_calls: None },
+                                finish_reason: Some(result.assistant.finish_reason),
+                            }],
+                            usage: req.stream_options.as_ref().filter(|options| options.include_usage).map(|_| result.usage),
+                        };
+                        yield Ok(Event::default().data(serde_json::to_string(&final_chunk).unwrap()));
+                        yield Ok(Event::default().data("[DONE]"));
+                    }
+                    Err(response) => {
+                        yield Ok(Event::default().data("[DONE]"));
+                        let _ = response;
+                    }
+                }
+            }
+        }
+    };
+
+    Sse::new(stream)
+        .keep_alive(KeepAlive::default())
+        .into_response()
+}
+
+async fn blocking_response(
+    state: AppState,
+    req: ChatCompletionRequest,
+    sampler: inference_core::SamplerConfig,
+    prompt: String,
+) -> Response {
+    let id = uuid::Uuid::new_v4().to_string();
+    let model = state
+        .server_config
+        .resolve_model(&req.model)
+        .map(|descriptor| descriptor.root.clone())
+        .unwrap_or_else(|| state.model_name.as_ref().clone());
+
+    match generate_full_completion(
+        Arc::clone(&state.engine),
+        Arc::clone(&state.counters),
+        prompt,
+        sampler,
+        req,
+    )
+    .await
+    {
+        Ok(result) => Json(ChatCompletionResponse {
+            id,
+            object: "chat.completion",
+            created: unix_now(),
+            model,
+            choices: vec![Choice {
+                index: 0,
+                message: ChatMessage {
+                    role: Role::Assistant,
+                    content: result.assistant.content,
+                    name: None,
+                    tool_calls: result.assistant.tool_calls,
+                    tool_call_id: None,
+                },
+                finish_reason: result.assistant.finish_reason,
+            }],
+            usage: result.usage,
+        })
+        .into_response(),
+        Err(response) => response,
+    }
+}
+
+async fn generate_full_completion(
+    engine_arc: Arc<Mutex<InferenceEngine>>,
+    counters_arc: Arc<PerfCounters>,
+    prompt: String,
+    sampler: inference_core::SamplerConfig,
+    req: ChatCompletionRequest,
+) -> Result<GenerationResult, Response> {
+    tracing::debug!(prompt_len = prompt.len(), "generate_full_completion: starting");
+    let result = tokio::task::spawn_blocking(move || {
+        let rt = tokio::runtime::Handle::current();
+        tracing::debug!("generate_full_completion: acquiring engine lock");
+        let mut engine = rt.block_on(engine_arc.lock());
+        tracing::debug!("generate_full_completion: engine lock acquired");
+        let mut record = metrics::recorder::GenerationRecord::new(0);
+        let mut content = String::new();
+
+        let mut session = engine.new_session(&prompt, sampler)?;
+        tracing::info!(prompt_tokens = session.prompt_tokens(), "generate_full_completion: session created");
+
+        while let Some(tok) = session.next_token() {
+            record.tick();
+            content.push_str(&tok);
+        }
+
+        counters_arc.inc_tokens(record.completion_tokens as u64);
+        counters_arc.inc_completed();
+
+        let usage = Usage {
+            prompt_tokens: session.prompt_tokens(),
+            completion_tokens: record.completion_tokens,
+            total_tokens: session.prompt_tokens() + record.completion_tokens,
+            tokens_per_second: record.tokens_per_second(),
+        };
+        tracing::info!(
+            prompt_tokens = usage.prompt_tokens,
+            completion_tokens = usage.completion_tokens,
+            content_len = content.len(),
+            tps = ?usage.tokens_per_second,
+            "generate_full_completion: done"
+        );
+        let assistant = finalize_assistant_message(content, &req);
+        Ok::<_, anyhow::Error>((assistant, usage))
+    })
+    .await;
+
+    match result {
+        Ok(Ok((assistant, usage))) => Ok(GenerationResult { assistant, usage }),
+        Ok(Err(err)) => {
+            tracing::error!(err = %err, "Generation failed");
+            Err(error_response(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                ApiError::new("Generation failed", "server_error").with_code("generation_failed"),
+            ))
+        }
+        Err(err) => {
+            tracing::error!(err = %err, "Spawn blocking panicked");
+            Err(error_response(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                ApiError::new("Internal server error", "server_error").with_code("internal_error"),
+            ))
+        }
+    }
+}
+
+struct GenerationResult {
+    assistant: FinalizedAssistant,
+    usage: Usage,
+}
+
+enum StreamItem {
+    Token(String),
+    Done {
+        finish_reason: FinishReason,
+        usage: Usage,
+    },
+    Error,
+}