mirror of
https://github.com/R0m1k3/Loki.git
synced 2026-10-15 19:29:16 +02:00
Lot 5 : /v1 relayé par la façade directement au moteur
La surface compatible OpenAI ne fait plus le détour par le Go : la façade authentifie avec la clé des complétions (préflight CORS exempté), réécrit le corps vers le second slot quand il est en service (id_slot à 1, valeurs gardées telles quelles), retire les X-Forwarded-* et relaie au moteur sans tampon — Host d'origine, sauf pour Strata. Les décisions restent au Go, qui tient l'état de la conversation : la façade lui demande où router, ouvre un ticket « en vol » avant de relayer (refus pendant un benchmark ou une optimisation, avec Retry-After) et le rend quand le flux se termine, que le client parte ou que le moteur soit injoignable. Sans ticket obtenu, la requête est refusée : le Go pourrait sinon effacer le slot en plein calcul. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
451770bfcf
commit
c9c9ac1ddd
7 files changed
+646
-1
No files matched your search
@@ -12,7 +12,7 @@ http-body-util = "0.1.5"
|
||||
hyper = { version = "1.12.0", features = ["http1", "server", "client"] }
|
||||
hyper-util = { version = "0.1.21", features = ["tokio", "server", "client-legacy", "http1"] }
|
||||
serde = { version = "1.0.229", features = ["derive"] }
|
||||
serde_json = "1.0.154"
|
||||
serde_json = { version = "1.0.154", features = ["raw_value"] }
|
||||
sha2 = "0.11.1"
|
||||
tokio = { version = "1.53.2", features = ["rt-multi-thread", "macros", "net", "io-util", "time", "signal", "process", "sync", "fs"] }
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
//! - `engine` : supervision du moteur (`loki serve` et son arbre, lot 3) ;
|
||||
//! - `ui` : UI web statique embarquée (lot 2) ;
|
||||
//! - `auth` : authentification de l'API de pilotage (lot 2) ;
|
||||
//! - `oai` : surface compatible OpenAI `/v1` relayée au moteur (lot 5) ;
|
||||
//! - `update` : mise à jour des binaires au démarrage du conteneur.
|
||||
|
||||
pub mod api;
|
||||
@@ -13,6 +14,7 @@ pub mod child;
|
||||
pub mod engine;
|
||||
pub mod engstatus;
|
||||
pub mod gguf;
|
||||
pub mod oai;
|
||||
pub mod paths;
|
||||
pub mod presets;
|
||||
pub mod proxy;
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
//! Surface compatible OpenAI (`/v1/*`), relayée par la façade directement au
|
||||
//! moteur (lot 5) — port de `mountOAI`, `requireCompletionKey` et
|
||||
//! `oaiHandler` (`llm_oai.go`, `web_auth.go`).
|
||||
//!
|
||||
//! Les DÉCISIONS restent au Go (`web_internal_oai.go`), qui tient l'état de
|
||||
//! la conversation : il dit où router (`route`), compte la requête « en vol »
|
||||
//! et refuse pendant un benchmark ou une optimisation (`begin`), puis apprend
|
||||
//! sa fin (`end`) — à la fin du flux ou quand le client part.
|
||||
//!
|
||||
//! Comme le Go : on NE touche PAS à `Authorization` (llama-server revérifie
|
||||
//! la clé, `--api-key`), on ne bufferise rien (les jetons arrivent un par
|
||||
//! un), et les `X-Forwarded-*` entrants ne passent pas au moteur.
|
||||
|
||||
use bytes::Bytes;
|
||||
use http_body_util::{BodyExt, Full, combinators::BoxBody};
|
||||
use hyper::body::{Frame, Incoming, SizeHint};
|
||||
use hyper::header::{self, HeaderValue};
|
||||
use hyper::{Method, Request, Response, StatusCode, Uri};
|
||||
use hyper_util::client::legacy::{Client, connect::HttpConnector};
|
||||
use serde_json::{Value, json, value::RawValue};
|
||||
use std::collections::BTreeMap;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use crate::api::Ctx;
|
||||
use crate::auth;
|
||||
use crate::proxy::{Body, strip_hop_by_hop};
|
||||
use crate::store::Store;
|
||||
|
||||
pub type OutBody = BoxBody<Bytes, hyper::Error>;
|
||||
pub type EngineClient = Client<HttpConnector, OutBody>;
|
||||
|
||||
/// Réponse d'erreur au format OpenAI (`sendOAIError`) : les SDK lisent
|
||||
/// `body.error.message`.
|
||||
pub fn oai_error(status: StatusCode, msg: &str, typ: &str, code: &str) -> Response<Body> {
|
||||
auth::json(
|
||||
status,
|
||||
&json!({ "error": { "message": msg, "type": typ, "param": null, "code": code } }),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn serve(ctx: &Ctx, engine: &EngineClient, req: Request<Incoming>) -> Response<Body> {
|
||||
let method = req.method().clone();
|
||||
let path = req.uri().path().to_string();
|
||||
let route = match ctx
|
||||
.store
|
||||
.internal_json(
|
||||
Method::GET,
|
||||
&format!(
|
||||
"/internal/v1/oai/route?method={}&path={}",
|
||||
crate::store::enc(method.as_str()),
|
||||
crate::store::enc(&path)
|
||||
),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
eprintln!("[loki-rs] /v1 : route illisible : {e}");
|
||||
return oai_error(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"configuration illisible — réessaie dans un instant",
|
||||
"api_error",
|
||||
"service_unavailable",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// requireCompletionKey : un préflight CORS ne porte jamais d'Authorization.
|
||||
if method != Method::OPTIONS {
|
||||
if route["key_ok"] != Value::Bool(true) {
|
||||
return oai_error(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"configuration illisible — réessaie dans un instant",
|
||||
"api_error",
|
||||
"service_unavailable",
|
||||
);
|
||||
}
|
||||
let key = route["key"].as_str().unwrap_or("");
|
||||
if !key.is_empty() && !auth::check_bearer(req.headers(), key) {
|
||||
let mut r = oai_error(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"Clé API invalide ou absente — envoie l'en-tête Authorization: Bearer <clé>.",
|
||||
"invalid_request_error",
|
||||
"invalid_api_key",
|
||||
);
|
||||
r.headers_mut().insert(
|
||||
header::WWW_AUTHENTICATE,
|
||||
HeaderValue::from_static(r#"Bearer realm="loki""#),
|
||||
);
|
||||
return r;
|
||||
}
|
||||
}
|
||||
|
||||
let port = route["port"].as_u64().unwrap_or(8080);
|
||||
let strata = route["strata"] == Value::Bool(true);
|
||||
let side_live = route["side_live"] == Value::Bool(true);
|
||||
let max_body = route["max_body"].as_u64().unwrap_or(0) as usize;
|
||||
|
||||
// SIDE_SLOT : la complétion d'un client passe par le second slot.
|
||||
let (mut parts, body) = req.into_parts();
|
||||
let post_v1 = method == Method::POST && path.starts_with("/v1/");
|
||||
let (body, side): (OutBody, bool) = if post_v1 && side_live {
|
||||
match side_slot_body(body, max_body).await {
|
||||
Ok((b, side)) => {
|
||||
parts
|
||||
.headers
|
||||
.insert(header::CONTENT_LENGTH, HeaderValue::from(b.len()));
|
||||
parts.headers.remove(header::TRANSFER_ENCODING);
|
||||
(Full::new(b).map_err(|n| match n {}).boxed(), side)
|
||||
}
|
||||
Err(r) => return *r,
|
||||
}
|
||||
} else {
|
||||
(body.boxed(), false)
|
||||
};
|
||||
|
||||
let read = method == Method::GET || method == Method::HEAD;
|
||||
let begin = ctx
|
||||
.store
|
||||
.internal_json(
|
||||
Method::POST,
|
||||
"/internal/v1/oai/begin",
|
||||
Some(&json!({ "read": read, "side": side, "post_v1": post_v1 })),
|
||||
)
|
||||
.await;
|
||||
let ticket = match begin {
|
||||
Ok(v) => {
|
||||
if let Some(refus) = v.get("refuse").filter(|r| r.is_object()) {
|
||||
let status = StatusCode::from_u16(refus["status"].as_u64().unwrap_or(503) as u16)
|
||||
.unwrap_or(StatusCode::SERVICE_UNAVAILABLE);
|
||||
let mut r = oai_error(
|
||||
status,
|
||||
refus["message"].as_str().unwrap_or(""),
|
||||
refus["type"].as_str().unwrap_or("server_error"),
|
||||
refus["code"].as_str().unwrap_or(""),
|
||||
);
|
||||
if let Some(ra) = refus["retry_after"].as_str()
|
||||
&& let Ok(v) = HeaderValue::from_str(ra)
|
||||
{
|
||||
r.headers_mut().insert(header::RETRY_AFTER, v);
|
||||
}
|
||||
return r;
|
||||
}
|
||||
v["ticket"].as_u64()
|
||||
}
|
||||
Err(e) => {
|
||||
// Sans ticket, le Go ne saurait pas qu'une requête occupe le moteur :
|
||||
// il pourrait effacer le slot en plein calcul. On refuse.
|
||||
eprintln!("[loki-rs] /v1 : begin : {e}");
|
||||
return oai_error(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"configuration illisible — réessaie dans un instant",
|
||||
"api_error",
|
||||
"service_unavailable",
|
||||
);
|
||||
}
|
||||
};
|
||||
let guard = TicketGuard {
|
||||
store: ctx.store.clone(),
|
||||
ticket,
|
||||
};
|
||||
|
||||
// Vers le moteur : Host d'origine, sauf Strata qui refuse un Host inconnu.
|
||||
let engine_host = format!("127.0.0.1:{port}");
|
||||
let pq = parts
|
||||
.uri
|
||||
.path_and_query()
|
||||
.map(|p| p.as_str())
|
||||
.unwrap_or("/");
|
||||
parts.uri = match format!("http://{engine_host}{pq}").parse::<Uri>() {
|
||||
Ok(u) => u,
|
||||
Err(_) => {
|
||||
return oai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"URI invalide",
|
||||
"invalid_request_error",
|
||||
"bad_request",
|
||||
);
|
||||
}
|
||||
};
|
||||
strip_hop_by_hop(&mut parts.headers);
|
||||
for h in [
|
||||
"x-forwarded-for",
|
||||
"x-forwarded-host",
|
||||
"x-forwarded-proto",
|
||||
"forwarded",
|
||||
] {
|
||||
parts.headers.remove(h);
|
||||
}
|
||||
if strata && let Ok(v) = HeaderValue::from_str(&engine_host) {
|
||||
parts.headers.insert(header::HOST, v);
|
||||
}
|
||||
let out = Request::from_parts(parts, body);
|
||||
match engine.request(out).await {
|
||||
Ok(resp) => {
|
||||
let (mut rp, rb) = resp.into_parts();
|
||||
strip_hop_by_hop(&mut rp.headers);
|
||||
Response::from_parts(
|
||||
rp,
|
||||
Guarded {
|
||||
inner: rb,
|
||||
_guard: guard,
|
||||
}
|
||||
.boxed(),
|
||||
)
|
||||
}
|
||||
Err(e) => {
|
||||
drop(guard);
|
||||
oai_error(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
&format!("moteur injoignable : {e}"),
|
||||
"api_error",
|
||||
"upstream_unavailable",
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `sideSlotProxyRewrite` : lit le corps (borné), pose `id_slot: 1` si c'est
|
||||
/// un objet JSON. Les valeurs sont gardées telles quelles (RawValue) ; les
|
||||
/// clés ressortent triées, comme `json.Marshal` d'une map Go.
|
||||
async fn side_slot_body(body: Incoming, max: usize) -> Result<(Bytes, bool), Box<Response<Body>>> {
|
||||
let raw = match http_body_util::Limited::new(body, max + 1).collect().await {
|
||||
Ok(c) => c.to_bytes(),
|
||||
Err(e) => {
|
||||
let too_big = e
|
||||
.downcast_ref::<http_body_util::LengthLimitError>()
|
||||
.is_some();
|
||||
return Err(Box::new(if too_big {
|
||||
oai_error(
|
||||
StatusCode::PAYLOAD_TOO_LARGE,
|
||||
"corps trop gros pour être routé vers le slot des clients",
|
||||
"invalid_request_error",
|
||||
"side_slot",
|
||||
)
|
||||
} else {
|
||||
oai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
&format!("corps illisible : {e}"),
|
||||
"invalid_request_error",
|
||||
"side_slot",
|
||||
)
|
||||
}));
|
||||
}
|
||||
};
|
||||
Ok(match rewrite_id_slot(&raw) {
|
||||
Some(b) => (Bytes::from(b), true),
|
||||
None => (raw, false),
|
||||
})
|
||||
}
|
||||
|
||||
/// Corps avec `id_slot: 1`, ou `None` si ce n'est pas un objet JSON.
|
||||
pub fn rewrite_id_slot(raw: &[u8]) -> Option<Vec<u8>> {
|
||||
let mut obj: BTreeMap<String, Box<RawValue>> = serde_json::from_slice(raw).ok()?;
|
||||
obj.insert("id_slot".into(), RawValue::from_string("1".into()).ok()?);
|
||||
serde_json::to_vec(&obj).ok()
|
||||
}
|
||||
|
||||
/// Rend le ticket au Go quand la réponse est finie ou abandonnée.
|
||||
struct TicketGuard {
|
||||
store: Store,
|
||||
ticket: Option<u64>,
|
||||
}
|
||||
|
||||
impl Drop for TicketGuard {
|
||||
fn drop(&mut self) {
|
||||
let Some(ticket) = self.ticket.take() else {
|
||||
return;
|
||||
};
|
||||
let store = self.store.clone();
|
||||
if let Ok(rt) = tokio::runtime::Handle::try_current() {
|
||||
rt.spawn(async move {
|
||||
if let Err(e) = store
|
||||
.internal_json(
|
||||
Method::POST,
|
||||
"/internal/v1/oai/end",
|
||||
Some(&json!({ "ticket": ticket })),
|
||||
)
|
||||
.await
|
||||
{
|
||||
eprintln!("[loki-rs] /v1 : end : {e}");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Corps de réponse du moteur, relayé trame par trame ; le ticket est rendu
|
||||
/// quand il est lâché (flux terminé ou client parti).
|
||||
struct Guarded {
|
||||
inner: Incoming,
|
||||
_guard: TicketGuard,
|
||||
}
|
||||
|
||||
impl hyper::body::Body for Guarded {
|
||||
type Data = Bytes;
|
||||
type Error = hyper::Error;
|
||||
|
||||
fn poll_frame(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<Option<Result<Frame<Bytes>, hyper::Error>>> {
|
||||
Pin::new(&mut self.inner).poll_frame(cx)
|
||||
}
|
||||
fn is_end_stream(&self) -> bool {
|
||||
self.inner.is_end_stream()
|
||||
}
|
||||
fn size_hint(&self) -> SizeHint {
|
||||
self.inner.size_hint()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn id_slot_pose_cles_triees_valeurs_intactes() {
|
||||
let raw = br#"{"stream":true,"model":"x","temperature":0.10000000000000001,"id_slot":0}"#;
|
||||
let out = String::from_utf8(rewrite_id_slot(raw).unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
out,
|
||||
r#"{"id_slot":1,"model":"x","stream":true,"temperature":0.10000000000000001}"#
|
||||
);
|
||||
assert!(rewrite_id_slot(b"[1,2]").is_none());
|
||||
assert!(rewrite_id_slot(b"pas du json").is_none());
|
||||
}
|
||||
}
|
||||
@@ -30,6 +30,7 @@ use tokio::net::TcpListener;
|
||||
use crate::api;
|
||||
use crate::auth::{HostEnv, INTERNAL_HEADER, InternalClient};
|
||||
use crate::engine;
|
||||
use crate::oai;
|
||||
use crate::store::Store;
|
||||
use crate::ui;
|
||||
|
||||
@@ -62,6 +63,8 @@ pub struct Config {
|
||||
|
||||
struct Proxy {
|
||||
client: Client<HttpConnector, Incoming>,
|
||||
/// Vers le moteur (llama-server / Strata), pour `/v1`.
|
||||
engine: oai::EngineClient,
|
||||
backend: SocketAddr,
|
||||
ctx: api::Ctx,
|
||||
}
|
||||
@@ -75,6 +78,7 @@ pub async fn serve(
|
||||
let client = Client::builder(TokioExecutor::new()).build_http();
|
||||
let proxy = Arc::new(Proxy {
|
||||
client,
|
||||
engine: Client::builder(TokioExecutor::new()).build_http(),
|
||||
backend: cfg.backend,
|
||||
ctx: api::Ctx {
|
||||
store: Store::new(cfg.backend, cfg.internal_token.clone()),
|
||||
@@ -131,6 +135,10 @@ impl Proxy {
|
||||
if let Some(resp) = ui::serve(req.method(), path, req.headers()) {
|
||||
return resp;
|
||||
}
|
||||
// Surface OpenAI : relayée au moteur, décisions demandées au Go.
|
||||
if path.starts_with("/v1/") {
|
||||
return oai::serve(&self.ctx, &self.engine, req).await;
|
||||
}
|
||||
if api::handles(req.method(), path) {
|
||||
return api::serve(&self.ctx, req).await;
|
||||
}
|
||||
|
||||
@@ -62,6 +62,30 @@ impl Store {
|
||||
.map_err(|_| "store : délai dépassé".to_string())?
|
||||
}
|
||||
|
||||
/// Appel JSON à une autre route de l'API interne du Go (`/internal/v1/…`).
|
||||
/// Code HTTP autre que 2xx → `Err`.
|
||||
pub async fn internal_json(
|
||||
&self,
|
||||
method: Method,
|
||||
path_and_query: &str,
|
||||
body: Option<&serde_json::Value>,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let body = body
|
||||
.map(|b| serde_json::to_vec(b).unwrap_or_default())
|
||||
.unwrap_or_default();
|
||||
let (status, bytes) = self.call(method, path_and_query, body).await?;
|
||||
if !status.is_success() {
|
||||
return Err(format!(
|
||||
"{path_and_query} : {status} {}",
|
||||
String::from_utf8_lossy(&bytes)
|
||||
));
|
||||
}
|
||||
if bytes.is_empty() {
|
||||
return Ok(serde_json::Value::Null);
|
||||
}
|
||||
serde_json::from_slice(&bytes).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
/// Valeur brute ; `Ok(None)` = absente, `Err` = pas pu lire (à distinguer
|
||||
/// pour les secrets : comme `getBytesErr`).
|
||||
pub async fn get(&self, bucket: &str, key: &str) -> Result<Option<Vec<u8>>, String> {
|
||||
|
||||
@@ -24,6 +24,12 @@ pub struct FakeGo {
|
||||
pub broken: Mutex<bool>,
|
||||
/// Chemins relayés (hors API interne).
|
||||
pub relayed: Mutex<Vec<String>>,
|
||||
/// /v1 : port du faux moteur, second slot en service, refus à renvoyer
|
||||
/// par begin, journal des begin/end.
|
||||
pub engine_port: Mutex<u16>,
|
||||
pub side_live: Mutex<bool>,
|
||||
pub refuse: Mutex<Option<serde_json::Value>>,
|
||||
pub oai_log: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl FakeGo {
|
||||
@@ -95,6 +101,39 @@ async fn handle(st: Arc<FakeGo>, req: Request<Incoming>) -> Response<Full<Bytes>
|
||||
let key = st.get("state", "web_key").unwrap_or_default();
|
||||
resp(200, format!(r#"{{"web_key":"{key}","version":"0.15.0"}}"#))
|
||||
}
|
||||
"/internal/v1/oai/route" => {
|
||||
let key = st.get("state", "api_key").unwrap_or_default();
|
||||
let v = serde_json::json!({
|
||||
"port": *st.engine_port.lock().unwrap(),
|
||||
"strata": false,
|
||||
"key": key,
|
||||
"key_ok": true,
|
||||
"max_body": 1024,
|
||||
"side_live": *st.side_live.lock().unwrap(),
|
||||
});
|
||||
resp(200, v.to_string())
|
||||
}
|
||||
"/internal/v1/oai/begin" => {
|
||||
let body = req.into_body().collect().await.unwrap().to_bytes();
|
||||
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
st.oai_log.lock().unwrap().push(format!(
|
||||
"begin read={} side={} post_v1={}",
|
||||
v["read"], v["side"], v["post_v1"]
|
||||
));
|
||||
if let Some(r) = st.refuse.lock().unwrap().clone() {
|
||||
return resp(200, serde_json::json!({ "refuse": r }).to_string());
|
||||
}
|
||||
resp(200, r#"{"ticket":7}"#)
|
||||
}
|
||||
"/internal/v1/oai/end" => {
|
||||
let body = req.into_body().collect().await.unwrap().to_bytes();
|
||||
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
st.oai_log
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("end {}", v["ticket"]));
|
||||
resp(204, "")
|
||||
}
|
||||
"/internal/v1/store/all" => {
|
||||
let s = st.store.lock().unwrap();
|
||||
let m: BTreeMap<String, String> = s
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
//! Lot 5 : /v1 relayé par la façade directement au moteur, décisions au Go.
|
||||
|
||||
mod common;
|
||||
|
||||
use bytes::Bytes;
|
||||
use common::fakego;
|
||||
use http_body_util::{BodyExt, Full, StreamBody};
|
||||
use hyper::body::{Frame, Incoming};
|
||||
use hyper::{Request, Response, StatusCode, service::service_fn};
|
||||
use hyper_util::client::legacy::Client;
|
||||
use hyper_util::rt::{TokioExecutor, TokioIo};
|
||||
use std::convert::Infallible;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
type B = http_body_util::combinators::BoxBody<Bytes, Infallible>;
|
||||
|
||||
/// Faux moteur : /v1/models renvoie la requête vue (en-têtes + corps),
|
||||
/// /v1/chat/completions diffuse deux événements à 500 ms d'écart.
|
||||
async fn faux_moteur(vu: Arc<Mutex<Vec<String>>>) -> u16 {
|
||||
let l = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = l.local_addr().unwrap().port();
|
||||
tokio::spawn(async move {
|
||||
while let Ok((s, _)) = l.accept().await {
|
||||
let vu = vu.clone();
|
||||
tokio::spawn(async move {
|
||||
let svc = service_fn(move |req: Request<Incoming>| {
|
||||
let vu = vu.clone();
|
||||
async move {
|
||||
let path = req.uri().path().to_string();
|
||||
let mut lines: Vec<String> = req
|
||||
.headers()
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}: {}", v.to_str().unwrap()))
|
||||
.collect();
|
||||
let body = req.into_body().collect().await.unwrap().to_bytes();
|
||||
lines.push(format!("body: {}", String::from_utf8_lossy(&body)));
|
||||
vu.lock().unwrap().push(lines.join("\n"));
|
||||
let r: Response<B> = if path == "/v1/chat/completions" {
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<Bytes>(2);
|
||||
tokio::spawn(async move {
|
||||
tx.send(Bytes::from_static(b"data: a\n\n")).await.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
let _ = tx.send(Bytes::from_static(b"data: b\n\n")).await;
|
||||
});
|
||||
let s = futures_util::stream::unfold(rx, |mut rx| async move {
|
||||
rx.recv()
|
||||
.await
|
||||
.map(|b| (Ok::<_, Infallible>(Frame::data(b)), rx))
|
||||
});
|
||||
Response::new(BodyExt::boxed(StreamBody::new(s)))
|
||||
} else {
|
||||
Response::new(Full::new(Bytes::from_static(b"{\"data\":[]}")).boxed())
|
||||
};
|
||||
Ok::<_, Infallible>(r)
|
||||
}
|
||||
});
|
||||
let _ = hyper::server::conn::http1::Builder::new()
|
||||
.serve_connection(TokioIo::new(s), svc)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
});
|
||||
port
|
||||
}
|
||||
|
||||
async fn send(
|
||||
facade: std::net::SocketAddr,
|
||||
method: &str,
|
||||
path: &str,
|
||||
headers: &[(&str, &str)],
|
||||
body: &str,
|
||||
) -> Response<Incoming> {
|
||||
let mut req = Request::builder()
|
||||
.method(method)
|
||||
.uri(format!("http://{facade}{path}"));
|
||||
for (k, v) in headers {
|
||||
req = req.header(*k, *v);
|
||||
}
|
||||
let client: Client<_, Full<Bytes>> = Client::builder(TokioExecutor::new()).build_http();
|
||||
client
|
||||
.request(req.body(Full::new(Bytes::from(body.to_string()))).unwrap())
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn attendre_end(go: &fakego::FakeGo) {
|
||||
for _ in 0..50 {
|
||||
if go
|
||||
.oai_log
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|l| l.starts_with("end"))
|
||||
{
|
||||
return;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn relais_et_ticket() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
let vu = Arc::new(Mutex::new(Vec::new()));
|
||||
*go.engine_port.lock().unwrap() = faux_moteur(vu.clone()).await;
|
||||
let r = send(
|
||||
f,
|
||||
"GET",
|
||||
"/v1/models",
|
||||
&[("host", "loki.exemple.fr"), ("x-forwarded-for", "6.6.6.6")],
|
||||
"",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), StatusCode::OK);
|
||||
let body = r.into_body().collect().await.unwrap().to_bytes();
|
||||
assert_eq!(&body[..], b"{\"data\":[]}");
|
||||
attendre_end(&go).await;
|
||||
assert_eq!(
|
||||
*go.oai_log.lock().unwrap(),
|
||||
vec!["begin read=true side=false post_v1=false", "end 7"]
|
||||
);
|
||||
let vu = vu.lock().unwrap()[0].clone();
|
||||
assert!(vu.contains("host: loki.exemple.fr"), "{vu}");
|
||||
assert!(!vu.contains("x-forwarded-for"), "{vu}");
|
||||
assert!(go.relayed.lock().unwrap().is_empty(), "/v1 relayé au Go");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn flux_sans_tampon_ticket_rendu_a_la_fin() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
*go.engine_port.lock().unwrap() = faux_moteur(Arc::new(Mutex::new(Vec::new()))).await;
|
||||
let t0 = Instant::now();
|
||||
let r = send(f, "POST", "/v1/chat/completions", &[], r#"{"stream":true}"#).await;
|
||||
let mut body = r.into_body();
|
||||
let first = body.frame().await.unwrap().unwrap().into_data().unwrap();
|
||||
assert_eq!(&first[..], b"data: a\n\n");
|
||||
assert!(
|
||||
t0.elapsed() < Duration::from_millis(300),
|
||||
"{:?}",
|
||||
t0.elapsed()
|
||||
);
|
||||
// Flux en cours : pas encore de end.
|
||||
assert!(
|
||||
!go.oai_log
|
||||
.lock()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.any(|l| l.starts_with("end"))
|
||||
);
|
||||
while body.frame().await.is_some() {}
|
||||
attendre_end(&go).await;
|
||||
assert_eq!(
|
||||
*go.oai_log.lock().unwrap(),
|
||||
vec!["begin read=false side=false post_v1=true", "end 7"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn second_slot_reecrit_le_corps() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
let vu = Arc::new(Mutex::new(Vec::new()));
|
||||
*go.engine_port.lock().unwrap() = faux_moteur(vu.clone()).await;
|
||||
*go.side_live.lock().unwrap() = true;
|
||||
let r = send(f, "POST", "/v1/models", &[], r#"{"model":"m","id_slot":0}"#).await;
|
||||
assert_eq!(r.status(), StatusCode::OK);
|
||||
let vu = vu.lock().unwrap()[0].clone();
|
||||
assert!(vu.contains(r#"body: {"id_slot":1,"model":"m"}"#), "{vu}");
|
||||
assert!(go.oai_log.lock().unwrap()[0].contains("side=true"));
|
||||
// Corps trop gros pour le second slot (max_body = 1024 dans le faux Go).
|
||||
let gros = format!(r#"{{"x":"{}"}}"#, "a".repeat(2000));
|
||||
let r = send(f, "POST", "/v1/models", &[], &gros).await;
|
||||
assert_eq!(r.status(), StatusCode::PAYLOAD_TOO_LARGE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cle_des_completions() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
*go.engine_port.lock().unwrap() = faux_moteur(Arc::new(Mutex::new(Vec::new()))).await;
|
||||
go.set("state", "api_key", "cle-oai");
|
||||
let r = send(f, "GET", "/v1/models", &[], "").await;
|
||||
assert_eq!(r.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(r.headers()["www-authenticate"], r#"Bearer realm="loki""#);
|
||||
let body = r.into_body().collect().await.unwrap().to_bytes();
|
||||
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(v["error"]["code"], "invalid_api_key");
|
||||
assert!(
|
||||
go.oai_log.lock().unwrap().is_empty(),
|
||||
"begin sans authentification"
|
||||
);
|
||||
let r = send(
|
||||
f,
|
||||
"GET",
|
||||
"/v1/models",
|
||||
&[("authorization", "Bearer cle-oai")],
|
||||
"",
|
||||
)
|
||||
.await;
|
||||
assert_eq!(r.status(), StatusCode::OK);
|
||||
// Préflight CORS : jamais d'Authorization, jamais refusé.
|
||||
let r = send(f, "OPTIONS", "/v1/models", &[], "").await;
|
||||
assert_eq!(r.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refus_benchmark() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
*go.engine_port.lock().unwrap() = faux_moteur(Arc::new(Mutex::new(Vec::new()))).await;
|
||||
*go.refuse.lock().unwrap() = Some(serde_json::json!({
|
||||
"status": 503, "message": "benchmark en cours sur ce moteur — réessaie dans un moment",
|
||||
"type": "server_error", "code": "benchmark_running", "retry_after": "30"
|
||||
}));
|
||||
let r = send(f, "POST", "/v1/chat/completions", &[], "{}").await;
|
||||
assert_eq!(r.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(r.headers()["retry-after"], "30");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn moteur_absent() {
|
||||
let t = tempfile::tempdir().unwrap();
|
||||
let (f, go, _stop) = fakego::start(t.path().into()).await;
|
||||
let mort = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.unwrap()
|
||||
.local_addr()
|
||||
.unwrap()
|
||||
.port();
|
||||
*go.engine_port.lock().unwrap() = mort;
|
||||
let r = send(f, "GET", "/v1/models", &[], "").await;
|
||||
assert_eq!(r.status(), StatusCode::BAD_GATEWAY);
|
||||
attendre_end(&go).await;
|
||||
assert!(
|
||||
go.oai_log.lock().unwrap().iter().any(|l| l == "end 7"),
|
||||
"ticket non rendu"
|
||||
);
|
||||
}
|
||||
Reference in new issue
Block a user