feat: streaming (SSE) end-to-end

- Provider::chat_stream; OpenAI provider streams via stream_options.include_usage
- pipeline.run_stream: shared route/select, failover before first byte
- /v1/chat/completions forwards SSE verbatim with X-Cortiq-* headers; usage tapped for stats
- admin /test/stream + live streaming in Playground (stream toggle, 7 langs)
This commit is contained in:
infosave2007
2026-06-30 16:37:35 +03:00
parent 3d71b03788
commit 70d8cde258
17 changed files with 483 additions and 64 deletions
+36
View File
@@ -165,6 +165,7 @@ pub fn api_routes(admin_token: String) -> Router<SharedState> {
.route("/admin/api/stats", get(get_stats))
.route("/admin/api/requests", get(get_requests))
.route("/admin/api/test", post(run_test))
.route("/admin/api/test/stream", post(run_test_stream))
.route_layer(middleware::from_fn(move |req: Request, next: Next| {
let expected = admin_token.clone();
async move { auth(expected, req, next).await }
@@ -753,3 +754,38 @@ async fn run_test(State(state): State<SharedState>, Json(b): Json<TestBody>) ->
Err(e) => ok(json!({ "ok": false, "latency_ms": latency_ms, "error": e.to_string() })),
}
}
/// Streaming playground: forwards provider SSE chunks; routing decision is exposed
/// via X-Cortiq-* response headers (the SPA reads them after the stream ends).
async fn run_test_stream(State(state): State<SharedState>, Json(b): Json<TestBody>) -> Response {
if b.messages.is_empty() {
return ApiError::bad("messages must not be empty").into_response();
}
let req = ChatRequest {
routing: parse_routing(&b.model),
messages: b.messages,
tools: vec![],
params: GenParams {
temperature: b.temperature,
max_tokens: b.max_tokens,
..Default::default()
},
stream: true,
meta: RequestMeta {
account: "playground".into(),
protocol: "playground".into(),
..Default::default()
},
};
match state.pipeline.run_stream(req, &state).await {
Ok((info, stream)) => {
let mut headers = crate::protocols::openai_chat::cortiq_headers(&info);
headers.insert(
header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("text/event-stream"),
);
(headers, axum::body::Body::from_stream(stream)).into_response()
}
Err(e) => ApiError::bad(e.to_string()).into_response(),
}
}
+224 -3
View File
@@ -79,9 +79,13 @@ impl Pipeline {
result
}
async fn run_inner(&self, req: ChatRequest, state: &SharedState) -> Result<ChatResponse> {
let live = state.live();
/// Stages 1–3: extract text → route → select & capability-filter candidates.
/// Shared by both the non-streaming and streaming paths.
async fn resolve(
&self,
req: &ChatRequest,
live: &crate::state::Live,
) -> Result<(RouteDecision, String, Vec<String>)> {
// 1. Extract text for classification
let text = self.extract_text(
&req.messages,
@@ -159,6 +163,12 @@ impl Pipeline {
.to_string(),
));
}
Ok((decision, route_source, final_candidates))
}
async fn run_inner(&self, req: ChatRequest, state: &SharedState) -> Result<ChatResponse> {
let live = state.live();
let (decision, route_source, final_candidates) = self.resolve(&req, &live).await?;
// 4. Attempt provider call with failover (fall back to alternatives on error)
let mut last_err = None;
@@ -210,6 +220,89 @@ impl Pipeline {
}))
}
/// Streaming variant: returns the routing metadata plus a stream of SSE byte
/// chunks (OpenAI wire format). Failover applies only before the first byte —
/// once the upstream stream has started it cannot be retried. Statistics are
/// recorded when the stream completes (tokens taken from the final usage chunk).
pub async fn run_stream(
&self,
req: ChatRequest,
state: &SharedState,
) -> Result<(RouteInfo, crate::providers::ChatStream)> {
let started = Instant::now();
let account = req.meta.account.clone();
let protocol = req.meta.protocol.clone();
let directive = match &req.routing {
RoutingDirective::Pinned { .. } => "pinned",
RoutingDirective::Auto { .. } => "auto",
}
.to_string();
let live = state.live();
let (decision, route_source, final_candidates) = match self.resolve(&req, &live).await {
Ok(t) => t,
Err(e) => {
record_error(state, &account, &protocol, &directive, started, &e);
return Err(e);
}
};
let mut last_err = None;
let mut attempt = 0usize;
for model_id in final_candidates {
let provider = match live.registry.get(&model_id) {
Some(p) => p,
None => continue,
};
tracing::info!(
"Streaming request (label: {}, tier: {}) to model '{}' via source '{}'",
decision.task_label,
decision.complexity_tier,
model_id,
route_source
);
match provider.chat_stream(req.clone()).await {
Ok(stream) => {
let info = RouteInfo {
task_label: decision.task_label.clone(),
complexity_score: decision.complexity_score,
complexity_tier: decision.complexity_tier.clone(),
selected_model: model_id,
route_source: route_source.clone(),
router_request_id: decision.router_request_id.clone(),
cost_usd: 0.0,
failover: attempt > 0,
};
let tapped = tap_stream(
stream,
state.stats.clone(),
provider.clone(),
info.clone(),
account,
protocol,
directive,
started,
);
return Ok((info, Box::pin(tapped)));
}
Err(err) => {
tracing::warn!(
"Provider '{}' stream failed: {:?}. Attempting failover...",
model_id,
err
);
last_err = Some(err);
attempt += 1;
}
}
}
let e = last_err.unwrap_or_else(|| {
GatewayError::UpstreamUnavailable("All candidate models failed".to_string())
});
record_error(state, &account, &protocol, &directive, started, &e);
Err(e)
}
fn extract_text(&self, messages: &[Message], strategy: &str, max_chars: usize) -> String {
let raw_text = match strategy {
"last_user" => messages
@@ -249,3 +342,131 @@ impl Pipeline {
}
}
}
fn now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
fn record_error(
state: &SharedState,
account: &str,
protocol: &str,
directive: &str,
started: Instant,
e: &GatewayError,
) {
state.stats.record(RequestRecord {
ts: now_secs(),
account: account.to_string(),
protocol: protocol.to_string(),
directive: directive.to_string(),
task_label: String::new(),
tier: String::new(),
score: 0.0,
model_id: String::new(),
route_source: String::new(),
prompt_tokens: 0,
completion_tokens: 0,
cost_usd: 0.0,
latency_ms: started.elapsed().as_millis() as u64,
outcome: "error".to_string(),
failover: false,
error: Some(e.to_string()),
});
}
/// Forward a provider SSE stream to the client verbatim while tapping the final
/// `usage` chunk; record statistics when the stream completes.
#[allow(clippy::too_many_arguments)]
fn tap_stream(
inner: crate::providers::ChatStream,
stats: std::sync::Arc<crate::stats::Stats>,
provider: std::sync::Arc<dyn crate::providers::Provider>,
info: RouteInfo,
account: String,
protocol: String,
directive: String,
started: Instant,
) -> impl futures::Stream<Item = Result<bytes::Bytes>> {
use futures::StreamExt;
async_stream::stream! {
let mut inner = inner;
let mut buf = String::new();
let mut prompt_tokens = 0u32;
let mut completion_tokens = 0u32;
let mut got_usage = false;
let mut delta_chars = 0usize;
let mut errored: Option<String> = None;
while let Some(item) = inner.next().await {
match &item {
Ok(bytes) => {
if let Ok(s) = std::str::from_utf8(bytes) {
buf.push_str(s);
while let Some(idx) = buf.find("\n\n") {
let evt: String = buf.drain(..idx + 2).collect();
for line in evt.lines() {
if let Some(data) = line.trim_start().strip_prefix("data:") {
let data = data.trim();
if data.is_empty() || data == "[DONE]" {
continue;
}
if let Ok(v) =
serde_json::from_str::<serde_json::Value>(data)
{
let u = &v["usage"];
if !u.is_null() {
prompt_tokens = u["prompt_tokens"]
.as_u64()
.unwrap_or(prompt_tokens as u64)
as u32;
completion_tokens = u["completion_tokens"]
.as_u64()
.unwrap_or(completion_tokens as u64)
as u32;
got_usage = true;
}
if let Some(c) =
v["choices"][0]["delta"]["content"].as_str()
{
delta_chars += c.chars().count();
}
}
}
}
}
}
}
Err(e) => errored = Some(e.to_string()),
}
yield item;
}
if !got_usage && completion_tokens == 0 {
// estimate when the server didn't return usage (~4 chars per token)
completion_tokens = (delta_chars / 4) as u32;
}
let cost = provider.price(prompt_tokens, completion_tokens);
stats.record(RequestRecord {
ts: now_secs(),
account,
protocol,
directive,
task_label: info.task_label.clone(),
tier: info.complexity_tier.clone(),
score: info.complexity_score,
model_id: info.selected_model.clone(),
route_source: info.route_source.clone(),
prompt_tokens,
completion_tokens,
cost_usd: cost,
latency_ms: started.elapsed().as_millis() as u64,
outcome: if errored.is_some() { "error" } else { "ok" }.to_string(),
failover: info.failover,
error: errored,
});
}
}
+42 -45
View File
@@ -3,11 +3,37 @@
//! to the pipeline, then converts the canonical [`ChatResponse`] back to OpenAI format.
use crate::error::{GatewayError, Result};
use crate::model::{ChatRequest, GenParams, Message, RequestMeta, RoutingDirective};
use crate::model::{ChatRequest, GenParams, Message, RequestMeta, RouteInfo, RoutingDirective};
use crate::state::SharedState;
use axum::response::{IntoResponse, Response};
use axum::{extract::State, routing::post, Json, Router};
use serde::Deserialize;
/// Build the `X-Cortiq-*` response headers from routing metadata.
pub(crate) fn cortiq_headers(c: &RouteInfo) -> axum::http::HeaderMap {
use axum::http::{HeaderMap, HeaderValue};
fn put(h: &mut HeaderMap, k: &'static str, v: &str) {
if let Ok(val) = HeaderValue::from_str(v) {
h.insert(k, val);
}
}
let mut h = HeaderMap::new();
put(&mut h, "X-Cortiq-Task-Label", &c.task_label);
put(
&mut h,
"X-Cortiq-Complexity-Score",
&c.complexity_score.to_string(),
);
put(&mut h, "X-Cortiq-Complexity-Tier", &c.complexity_tier);
put(&mut h, "X-Cortiq-Selected-Model", &c.selected_model);
put(&mut h, "X-Cortiq-Route-Source", &c.route_source);
put(&mut h, "X-Cortiq-Cost-Usd", &c.cost_usd.to_string());
if let Some(id) = &c.router_request_id {
put(&mut h, "X-Cortiq-Request-Id", id);
}
h
}
pub fn routes() -> Router<SharedState> {
Router::new().route("/v1/chat/completions", post(handler))
}
@@ -46,7 +72,7 @@ fn parse_routing(model: &str) -> RoutingDirective {
async fn handler(
State(state): State<SharedState>,
Json(req): Json<OpenAiChatRequest>,
) -> Result<impl axum::response::IntoResponse> {
) -> Result<Response> {
// hot protocol toggle: if the adapter is disabled in config — return 404
if !state.live().cfg.protocols.openai_chat {
return Err(GatewayError::InvalidRequest(
@@ -77,52 +103,23 @@ async fn handler(
},
};
// streaming: forward provider SSE chunks verbatim with X-Cortiq-* headers
if req.stream {
return Err(GatewayError::InvalidRequest(
"Streaming is not yet supported in this version".into(),
));
let (info, stream) = state.pipeline.run_stream(canonical, &state).await?;
let mut headers = cortiq_headers(&info);
headers.insert(
axum::http::header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("text/event-stream"),
);
headers.insert(
axum::http::header::CACHE_CONTROL,
axum::http::HeaderValue::from_static("no-cache"),
);
return Ok((headers, axum::body::Body::from_stream(stream)).into_response());
}
let resp = state.pipeline.run(canonical, &state).await?;
let mut headers = axum::http::HeaderMap::new();
headers.insert(
"X-Cortiq-Task-Label",
axum::http::HeaderValue::from_str(&resp.cortiq.task_label)
.unwrap_or(axum::http::HeaderValue::from_static("")),
);
headers.insert(
"X-Cortiq-Complexity-Score",
axum::http::HeaderValue::from_str(&resp.cortiq.complexity_score.to_string())
.unwrap_or(axum::http::HeaderValue::from_static("0.0")),
);
headers.insert(
"X-Cortiq-Complexity-Tier",
axum::http::HeaderValue::from_str(&resp.cortiq.complexity_tier)
.unwrap_or(axum::http::HeaderValue::from_static("")),
);
headers.insert(
"X-Cortiq-Selected-Model",
axum::http::HeaderValue::from_str(&resp.cortiq.selected_model)
.unwrap_or(axum::http::HeaderValue::from_static("")),
);
headers.insert(
"X-Cortiq-Route-Source",
axum::http::HeaderValue::from_str(&resp.cortiq.route_source)
.unwrap_or(axum::http::HeaderValue::from_static("")),
);
headers.insert(
"X-Cortiq-Cost-Usd",
axum::http::HeaderValue::from_str(&resp.cortiq.cost_usd.to_string())
.unwrap_or(axum::http::HeaderValue::from_static("0.0")),
);
if let Some(req_id) = &resp.cortiq.router_request_id {
headers.insert(
"X-Cortiq-Request-Id",
axum::http::HeaderValue::from_str(req_id)
.unwrap_or(axum::http::HeaderValue::from_static("")),
);
}
let headers = cortiq_headers(&resp.cortiq);
let mut body = serde_json::json!({
"id": resp.id,
@@ -164,5 +161,5 @@ async fn handler(
});
}
Ok((headers, Json(body)))
Ok((headers, Json(body)).into_response())
}
+17 -2
View File
@@ -7,8 +7,16 @@
pub mod anthropic;
pub mod openai;
use crate::error::GatewayError;
use crate::model::{ChatRequest, ChatResponse};
use async_trait::async_trait;
use bytes::Bytes;
use futures::stream::BoxStream;
/// A streaming chat response: raw Server-Sent-Events byte chunks in OpenAI wire
/// format (`data: {json}\n\n` … `data: [DONE]\n\n`). The gateway forwards these
/// to the client verbatim while tapping the final `usage` for statistics.
pub type ChatStream = BoxStream<'static, crate::error::Result<Bytes>>;
/// Model capabilities — checked at selection time (whether tools/vision are needed).
#[derive(Clone, Debug, Default)]
@@ -41,6 +49,13 @@ pub trait Provider: Send + Sync {
/// Non-streaming chat call.
async fn chat(&self, req: ChatRequest) -> crate::error::Result<ChatResponse>;
// TODO(v0.1): async fn chat_stream(...) -> Result<ChatStream> (SSE).
// TODO(v0.1): async fn embed(...) -> Result<EmbedResponse>.
/// Streaming chat call (SSE). Default: not supported (override per provider).
async fn chat_stream(&self, _req: ChatRequest) -> crate::error::Result<ChatStream> {
Err(GatewayError::UpstreamUnavailable(format!(
"streaming is not supported by provider '{}'",
self.id()
)))
}
// TODO(v0.2): async fn embed(...) -> Result<EmbedResponse>.
}
+31
View File
@@ -154,4 +154,35 @@ impl Provider for OpenAiProvider {
},
})
}
async fn chat_stream(&self, req: ChatRequest) -> Result<super::ChatStream> {
use futures::StreamExt;
let mut body = self.to_wire(&req);
body["stream"] = serde_json::Value::Bool(true);
// ask OpenAI-compatible servers to include token usage in the final chunk
body["stream_options"] = serde_json::json!({ "include_usage": true });
let mut r = self
.http
.post(format!("{}/chat/completions", self.base_url))
.json(&body);
if let Some(key) = &self.api_key {
r = r.bearer_auth(key);
}
let resp = r
.send()
.await
.map_err(|e| GatewayError::UpstreamUnavailable(e.to_string()))?;
if !resp.status().is_success() {
let status = resp.status();
let err_body = resp.text().await.unwrap_or_default();
return Err(GatewayError::UpstreamUnavailable(format!(
"provider {} returned HTTP {}: {}",
self.id, status, err_body
)));
}
let stream = resp
.bytes_stream()
.map(|res| res.map_err(|e| GatewayError::UpstreamUnavailable(e.to_string())));
Ok(Box::pin(stream))
}
}