weft_core/api/
openai_compat.rs

1use crate::config::AppConfig;
2use crate::pipeline::Pipeline;
3use crate::process::ProcessManager;
4use crate::types::{ApiError, ApiErrorBody, ChatRequest};
5use crate::vkeys::VirtualKeyStore;
6use axum::extract::State;
7use axum::http::StatusCode;
8use axum::response::sse::{Event, KeepAlive, Sse};
9use axum::response::{IntoResponse, Response};
10use axum::Json;
11use futures_util::StreamExt;
12use std::convert::Infallible;
13use std::sync::{Arc, Mutex as StdMutex};
14use tokio::sync::RwLock;
15
16pub type SharedConfig = Arc<RwLock<AppConfig>>;
17pub type SharedPipeline = Arc<Pipeline>;
18
19#[derive(Clone)]
20pub struct AppState {
21    pub config: SharedConfig,
22    pub config_path: std::path::PathBuf,
23    pub pipeline: SharedPipeline,
24    pub process_manager: Arc<ProcessManager>,
25    pub vkey_store: Arc<VirtualKeyStore>,
26    pub package_manager: Arc<RwLock<crate::package::PackageManager>>,
27    pub wasm_handle: Arc<RwLock<Option<crate::package::bridge::WasmHandle>>>,
28    pub native_handle: Arc<RwLock<Option<crate::package::NativeHandle>>>,
29    pub resolved_apps: Arc<RwLock<crate::app::ResolvedAppMap>>,
30    pub capability_registry: Arc<RwLock<crate::app::CapabilityRegistry>>,
31    pub active_profile: Arc<RwLock<crate::app::AppProfile>>,
32    pub core_policy: Arc<crate::app::CorePolicy>,
33    pub generation_store: Arc<RwLock<crate::app::GenerationStoreMap>>,
34    pub package_index: Arc<crate::app::PackageIndex>,
35    pub repo_root: std::path::PathBuf,
36    pub data_dir: std::path::PathBuf,
37    pub runtime_token: Option<String>,
38    pub runtime_token_path: Option<std::path::PathBuf>,
39    pub chat_providers: Arc<RwLock<Vec<ChatProviderInfo>>>,
40    pub shutdown_tx: Arc<StdMutex<Option<tokio::sync::oneshot::Sender<()>>>>,
41    /// Native stream buffer: session_id -> pending token chunks.
42    /// Written by host_chat_completion_stream, read by /api/stream/tokens.
43    pub stream_buffer: Arc<StdMutex<std::collections::HashMap<String, Vec<String>>>>,
44}
45
46#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
47pub struct ChatProviderInfo {
48    pub name: String,
49    pub endpoint: String,
50    pub description: String,
51}
52
53pub async fn chat_completions(
54    State(state): State<AppState>,
55    Json(request): Json<ChatRequest>,
56) -> Response {
57    if request.stream {
58        return chat_completions_stream(state, request).await;
59    }
60
61    let config = state.config.read().await;
62    match state.pipeline.execute(&request, &config).await {
63        Ok(resp) => Json(resp).into_response(),
64        Err(e) => {
65            let body = ApiError {
66                error: ApiErrorBody {
67                    message: e.to_string(),
68                    error_type: "proxy_error".into(),
69                    code: Some("bad_gateway".into()),
70                },
71            };
72            (StatusCode::BAD_GATEWAY, Json(body)).into_response()
73        }
74    }
75}
76
77async fn chat_completions_stream(state: AppState, request: ChatRequest) -> Response {
78    let config = state.config.read().await.clone();
79
80    let (provider_name, resp) = match state.pipeline.execute_stream(&request, &config).await {
81        Ok(r) => r,
82        Err(e) => {
83            let body = ApiError {
84                error: ApiErrorBody {
85                    message: e.to_string(),
86                    error_type: "proxy_error".into(),
87                    code: Some("bad_gateway".into()),
88                },
89            };
90            return (StatusCode::BAD_GATEWAY, Json(body)).into_response();
91        }
92    };
93
94    let provider = config
95        .providers
96        .iter()
97        .find(|p| p.name == provider_name)
98        .cloned();
99
100    let transforms = state.pipeline.transforms.clone();
101
102    let stream = async_stream::stream! {
103        let mut byte_stream = resp.bytes_stream();
104        // 字节级缓冲：SSE chunk 按字节到达，一个多字节 UTF-8 字符（如中文 3 字节）
105        // 可能跨两个 chunk。必须在字节层累积、按 \n 切行，只对完整行解码，
106        // 否则每个 chunk 单独 from_utf8_lossy 会把跨边界的半个字符变成 �（乱码）。
107        let mut buffer: Vec<u8> = Vec::new();
108
109        while let Some(chunk_result) = byte_stream.next().await {
110            let chunk = match chunk_result {
111                Ok(c) => c,
112                Err(e) => {
113                    tracing::error!("Stream read error: {}", e);
114                    break;
115                }
116            };
117
118            buffer.extend_from_slice(&chunk);
119
120            while let Some(newline_pos) = buffer.iter().position(|&b| b == b'\n') {
121                let line_bytes: Vec<u8> = buffer.drain(..=newline_pos).collect();
122                // 去掉行尾的 \n（及可能的 \r）后按完整字节序列解码。
123                let line = String::from_utf8_lossy(&line_bytes).to_string();
124
125                let line = line.trim().to_string();
126                if line.is_empty() {
127                    continue;
128                }
129
130                if let Some(ref prov) = provider {
131                    match transforms.for_format(&prov.format).transform_stream_chunk(&line, prov).await {
132                        Ok(Some(chunk)) => {
133                            let data = serde_json::to_string(&chunk).unwrap_or_default();
134                            yield Ok::<_, Infallible>(Event::default().data(data));
135                        }
136                        Ok(None) => {
137                            if line.contains("[DONE]") {
138                                yield Ok(Event::default().data("[DONE]"));
139                            }
140                        }
141                        Err(e) => {
142                            tracing::warn!("Stream chunk transform error: {}", e);
143                        }
144                    }
145                }
146            }
147        }
148    };
149
150    Sse::new(stream)
151        .keep_alive(KeepAlive::default())
152        .into_response()
153}
154
155pub async fn list_models(State(state): State<AppState>) -> Json<serde_json::Value> {
156    let config = state.config.read().await;
157    let models: Vec<serde_json::Value> = config
158        .providers
159        .iter()
160        .flat_map(|p| {
161            p.models.iter().map(move |m| {
162                serde_json::json!({
163                    "id": m,
164                    "object": "model",
165                    "owned_by": p.name,
166                })
167            })
168        })
169        .collect();
170
171    Json(serde_json::json!({
172        "object": "list",
173        "data": models,
174    }))
175}
weft_core/api/openai_compat.rs

weft_core/api/
openai_compat.rs
