feat(xtask): 实现反馈模型列表的接口

YdrMaster · YdrMaster · commit c331a1cb05da · 2025-06-12T18:21:11.000+08:00
Signed-off-by: YdrMaster &lt;ydrml@hotmail.com&gt;
diff --git a/xtask/src/service/cache_manager.rs b/xtask/src/service/cache_manager.rs
@@ -23,7 +23,7 @@ impl CacheManager {
         &mut self,
         tokens: Vec<utok>,
         sample_args: SampleArgs,
-        max_steps: usize,
+        max_tokens: usize,
     ) -> (SessionId, Vec<utok>) {
         static SESSION_ID: AtomicUsize = AtomicUsize::new(0);
         let id = SessionId(SESSION_ID.fetch_add(1, SeqCst));
@@ -51,7 +51,7 @@ impl CacheManager {
                 cache,
             },
             &tokens[pos..],
-            max_steps,
+            max_tokens,
         );
         (id, tokens)
     }
diff --git a/xtask/src/service/client.rs b/xtask/src/service/client.rs
@@ -1,4 +1,4 @@
-﻿use super::V1_CHAT_COMPLETIONS;
+﻿use super::openai::POST_CHAT_COMPLETIONS;
 use log::{info, trace, warn};
 use openai_struct::{
     ChatCompletionRequestMessage, CreateChatCompletionRequest, CreateChatCompletionStreamResponse,
@@ -74,7 +74,10 @@ async fn send_single_request(
     }
 
     let req = client
-        .post(format!("http://localhost:{port}{V1_CHAT_COMPLETIONS}"))
+        .post(format!(
+            "http://localhost:{port}{}",
+            POST_CHAT_COMPLETIONS.1
+        ))
         .headers(headers.clone())
         .body(req_body)
         .timeout(Duration::from_secs(100));
diff --git a/xtask/src/service/mod.rs b/xtask/src/service/mod.rs
@@ -4,20 +4,22 @@ mod model;
 mod openai;
 mod response;
 
-use crate::{parse_gpus, service::model::Model};
+use crate::parse_gpus;
 use error::*;
 use http_body_util::{BodyExt, combinators::BoxBody};
 use hyper::{
-    Method, Request, Response,
+    Request, Response,
     body::{Bytes, Incoming},
     server::conn::http1,
     service::Service as HyperService,
 };
 use hyper_util::rt::TokioIo;
 use log::{info, warn};
-use openai::V1_CHAT_COMPLETIONS;
+use model::Model;
+use openai::create_models;
 use openai_struct::CreateChatCompletionRequest;
 use response::error;
+use response::json;
 use std::collections::HashMap;
 use std::{ffi::c_int, fs::read_to_string, path::Path};
 use std::{
@@ -42,7 +44,7 @@ pub struct ServiceArgs {
     #[clap(long)]
     gpus: Option<String>,
     #[clap(long)]
-    max_steps: Option<usize>,
+    max_tokens: Option<usize>,
     #[clap(long)]
     think: bool,
 }
@@ -51,7 +53,7 @@ pub struct ServiceArgs {
 pub struct ModelConfig {
     pub path: String,
     pub gpus: Option<Box<[c_int]>>,
-    pub max_steps: Option<usize>,
+    pub max_tokens: Option<usize>,
     pub think: Option<bool>,
 }
 
@@ -63,7 +65,7 @@ impl ServiceArgs {
             no_cuda_graph,
             name,
             gpus,
-            max_steps,
+            max_tokens,
             think,
         } = self;
 
@@ -77,7 +79,7 @@ impl ServiceArgs {
                 ModelConfig {
                     path: file.clone(),
                     gpus: Some(parse_gpus(gpus.as_deref())),
-                    max_steps,
+                    max_tokens,
                     think: Some(think),
                 },
             )]
@@ -139,19 +141,25 @@ impl HyperService<Request<Incoming>> for App {
     type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
 
     fn call(&self, req: Request<Incoming>) -> Self::Future {
-        let models = self.0.clone();
         match (req.method(), req.uri().path()) {
-            (&Method::POST, V1_CHAT_COMPLETIONS) => Box::pin(async move {
-                let whole_body = req.collect().await?.to_bytes();
-                let req = serde_json::from_slice::<CreateChatCompletionRequest>(&whole_body);
-                Ok(match req {
-                    Ok(req) => match models.get(&req.model) {
-                        Some(model) => model.complete_chat(req),
-                        None => error(Error::ModelNotFound(req.model)),
-                    },
-                    Err(e) => error(Error::WrongJson(e)),
+            openai::GET_MODELS => {
+                let json = json(create_models(self.0.keys().cloned()));
+                Box::pin(async move { Ok(json) })
+            }
+            openai::POST_CHAT_COMPLETIONS => {
+                let models = self.0.clone();
+                Box::pin(async move {
+                    let whole_body = req.collect().await?.to_bytes();
+                    let req = serde_json::from_slice::<CreateChatCompletionRequest>(&whole_body);
+                    Ok(match req {
+                        Ok(req) => match models.get(&req.model) {
+                            Some(model) => model.complete_chat(req),
+                            None => error(Error::ModelNotFound(req.model)),
+                        },
+                        Err(e) => error(Error::WrongJson(e)),
+                    })
                 })
-            }),
+            }
             // Return 404 Not Found for other routes.
             (method, uri) => {
                 let msg = Error::not_found(method, uri);
diff --git a/xtask/src/service/model.rs b/xtask/src/service/model.rs
@@ -28,8 +28,8 @@ use tokio::{
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
 pub(super) struct Model {
+    max_tokens: usize,
     terminal: Terminal,
-    max_steps: usize,
     sessions: Mutex<BTreeMap<SessionId, SessionInfo>>,
     cache_manager: Mutex<CacheManager>,
 }
@@ -48,7 +48,7 @@ impl Model {
         let ModelConfig {
             path,
             gpus,
-            max_steps,
+            max_tokens,
             think,
         } = config;
 
@@ -67,26 +67,25 @@ impl Model {
             (utok::MAX, utok::MAX)
         };
 
-        let service_manager = Arc::new(Model {
+        let model = Arc::new(Model {
+            max_tokens: max_tokens.unwrap_or(2 << 10),
             terminal: service.terminal().clone(),
-            max_steps: max_steps.unwrap_or(2 << 10),
             sessions: Mutex::new(sessions),
             cache_manager: Mutex::new(CacheManager::new(service.terminal().clone())),
         });
 
-        let service_manager_for_recv = service_manager.clone();
-
+        let model_ = model.clone();
         let join_handle = tokio::task::spawn_blocking(move || {
             loop {
                 let Received { sessions, outputs } = service.recv(Duration::from_millis(10));
 
+                let mut sessions_guard = model_.sessions.lock().unwrap();
                 // 先处理输出
                 for (session_id, tokens) in outputs {
                     if tokens.is_empty() {
                         continue;
                     }
 
-                    let mut sessions_guard = service_manager_for_recv.sessions.lock().unwrap();
                     let session_info = sessions_guard.get_mut(&session_id).unwrap();
                     // 更新 session_info
                     session_info.tokens.extend(&tokens);
@@ -111,12 +110,8 @@ impl Model {
                         &[]
                     };
 
-                    let think = service_manager_for_recv
-                        .terminal
-                        .decode(think, &mut session_info.buf);
-                    let text = service_manager_for_recv
-                        .terminal
-                        .decode(tokens, &mut session_info.buf);
+                    let think = model_.terminal.decode(think, &mut session_info.buf);
+                    let text = model_.terminal.decode(tokens, &mut session_info.buf);
                     debug!("解码完成：{tokens:?} -> {think:?} | {text:?}");
 
                     let response = create_chat_completion_response(
@@ -131,16 +126,12 @@ impl Model {
 
                     if session_info.sender.send(message).is_err() {
                         info!("{session_id:?} 客户端连接已关闭");
-                        service_manager_for_recv.terminal.stop(session_id);
+                        model_.terminal.stop(session_id);
                     }
                 }
 
                 // 处理会话结束
                 if !sessions.is_empty() {
-                    let mut sessions_guard = service_manager_for_recv.sessions.lock().unwrap();
-                    let mut cache_manager_guard =
-                        service_manager_for_recv.cache_manager.lock().unwrap();
-
                     for (session, reason) in sessions {
                         let SessionInfo {
                             tokens,
@@ -152,7 +143,11 @@ impl Model {
                         let reason = match reason {
                             // 正常完成，插回cache
                             ReturnReason::Finish => {
-                                cache_manager_guard.insert(tokens, session.cache);
+                                model_
+                                    .cache_manager
+                                    .lock()
+                                    .unwrap()
+                                    .insert(tokens, session.cache);
                                 info!("{:?} 正常完成", session.id);
                                 FinishReason::Stop
                             }
@@ -177,12 +172,12 @@ impl Model {
             }
         });
 
-        (service_manager, join_handle)
+        (model, join_handle)
     }
 
     pub fn complete_chat(
         &self,
-        completions: CreateChatCompletionRequest,
+        req: CreateChatCompletionRequest,
     ) -> Response<BoxBody<Bytes, hyper::Error>> {
         let CreateChatCompletionRequest {
             model,
@@ -191,10 +186,10 @@ impl Model {
             temperature,
             top_p,
             ..
-        } = completions;
+        } = req;
         let (sender, receiver) = mpsc::unbounded_channel();
 
-        let max_steps = max_tokens.map_or(self.max_steps, |n| n as usize);
+        let max_tokens = max_tokens.map_or(self.max_tokens, |n| n as _);
         let sample_args =
             SampleArgs::new(temperature.unwrap_or(0.), top_p.unwrap_or(1.), usize::MAX).unwrap();
 
@@ -242,7 +237,7 @@ impl Model {
             .cache_manager
             .lock()
             .unwrap()
-            .send(tokens, sample_args, max_steps);
+            .send(tokens, sample_args, max_tokens);
 
         let session_info = SessionInfo {
             sender,
diff --git a/xtask/src/service/openai.rs b/xtask/src/service/openai.rs
@@ -1,11 +1,35 @@
-﻿use llama_cu::SessionId;
+﻿use hyper::Method;
+use llama_cu::SessionId;
 use openai_struct::{
     ChatCompletionStreamResponseDelta, CreateChatCompletionStreamResponse,
-    CreateChatCompletionStreamResponseChoices, FinishReason,
+    CreateChatCompletionStreamResponseChoices, FinishReason, Model,
 };
+use serde::Serialize;
 
 const CHAT_COMPLETION_OBJECT: &str = "chat.completion.chunk";
-pub(crate) const V1_CHAT_COMPLETIONS: &str = "/v1/chat/completions";
+pub(crate) const GET_MODELS: (&Method, &str) = (&Method::GET, "models");
+pub(crate) const POST_CHAT_COMPLETIONS: (&Method, &str) = (&Method::POST, "/chat/completions");
+
+pub(crate) fn create_models(models: impl IntoIterator<Item = String>) -> impl Serialize {
+    #[derive(Serialize)]
+    struct Response {
+        object: &'static str,
+        data: Vec<Model>,
+    }
+
+    Response {
+        object: "list",
+        data: models
+            .into_iter()
+            .map(|id| Model {
+                id,
+                object: "model".into(),
+                owned_by: "QYLab".into(),
+                created: 0,
+            })
+            .collect(),
+    }
+}
 
 pub(crate) fn create_chat_completion_response(
     id: SessionId,
diff --git a/xtask/src/service/response.rs b/xtask/src/service/response.rs
@@ -10,6 +10,7 @@ use hyper::{
         CACHE_CONTROL, CONNECTION, CONTENT_TYPE,
     },
 };
+use serde::Serialize;
 use tokio_stream::{Stream, StreamExt};
 
 pub fn text_stream(
@@ -27,6 +28,17 @@ pub fn text_stream(
         .unwrap()
 }
 
+pub fn json(json: impl Serialize) -> Response<BoxBody<Bytes, hyper::Error>> {
+    Response::builder()
+        .status(StatusCode::OK)
+        .header(CONTENT_TYPE, "application/json")
+        .header(ACCESS_CONTROL_ALLOW_ORIGIN, "*")
+        .header(ACCESS_CONTROL_ALLOW_METHODS, "GET,POST")
+        .header(ACCESS_CONTROL_ALLOW_HEADERS, "Content-Type")
+        .body(full(serde_json::to_string(&json).unwrap()))
+        .unwrap()
+}
+
 pub fn error(e: Error) -> Response<BoxBody<Bytes, hyper::Error>> {
     Response::builder()
         .status(e.status())
