jev_bridge/
server.rs

1//! The bridge: one resident upstream model behind a Jev-style scoring API.
2
3use std::sync::Arc;
4
5use anyhow::{bail, Context, Result};
6use axum::{
7    body::Bytes,
8    extract::State,
9    http::{header::AUTHORIZATION, HeaderMap, StatusCode},
10    response::{IntoResponse, Response},
11    routing::{get, post},
12    Json, Router,
13};
14use serde::Serialize;
15use serde_json::{json, Value};
16use tokio::sync::Mutex;
17
18use crate::prompt::{direct_messages, prompt_sha256, softmax, ChatTemplate, LocalComponents, RuntimeClient, LETTERS, PROMPT_VERSION};
19use crate::strategy::{
20    build_body, candidates_restricted_to, input_tokens, parse_candidates, parse_logprobs,
21    request_url, truncated, CompletionResponse, Probe,
22};
23use crate::wire::{
24    request_rows, response_from_results, OrderedMap, ReadoutStatus, Row, ScoredAnswer,
25    PROBABILITY_STATUS,
26};
27
28/// One configured bridge over a single upstream model.
29pub struct Bridge {
30    client: reqwest::Client,
31    template: ChatTemplate,
32    probe: Probe,
33    upstream_model: String,
34    upstream_key: Option<String>,
35    openai_base: String,
36    native_base: String,
37    slots: Vec<u32>,
38    served_model: String,
39    description: String,
40    release_date: String,
41    max_input_tokens: Option<usize>,
42    /// What the deployment declared about the readout channel on this model.
43    readout: ReadoutStatus,
44    /// Serialize upstream work so one resident model is not driven concurrently.
45    lock: Mutex<()>,
46}
47
48/// Everything [`Bridge::connect`] needs to reach one upstream model.
49pub struct BridgeConfig {
50    /// OpenAI-compatible root, for example `http://127.0.0.1:8080/v1`.
51    pub openai_base: String,
52    /// Native root for llama.cpp's `/completion`.
53    pub native_base: String,
54    /// Model name sent upstream.
55    pub upstream_model: String,
56    /// Bearer token for the upstream service, when it needs one.
57    pub upstream_key: Option<String>,
58    /// Model id this bridge accepts from clients.
59    pub served_model: String,
60    /// Description returned by `GET /v1/models`.
61    pub description: String,
62    /// ISO release date returned by `GET /v1/models`.
63    pub release_date: String,
64    /// Template variables, for example `{"enable_thinking": false}`.
65    pub chat_template_kwargs: Value,
66    /// Reject rows above this token count instead of truncating them.
67    pub max_input_tokens: Option<usize>,
68    /// Render the chat template and/or tokenize in-process instead of asking
69    /// the runtime.
70    pub local: LocalComponents,
71    /// What the deployment knows about the readout channel on this model.
72    ///
73    /// Defaults to `unvalidated`: the bridge cannot prove readout quality at
74    /// runtime, so a deployment that checked it (for example with
75    /// [`crate::evaluate`] over an alignment run) states so explicitly.
76    pub readout: ReadoutStatus,
77}
78
79/// A scored row plus the evidence needed to compare a bridged run against a
80/// reference run row by row.
81#[derive(Debug, Clone)]
82pub struct DetailedScore {
83    /// The normalized answer handed to the wire layer.
84    pub answer: ScoredAnswer,
85    /// Raw option log probabilities, before the softmax.
86    pub option_logprobs: Vec<f64>,
87    /// SHA-256 of the rendered prompt, for comparing runs row by row.
88    pub prompt_sha256: String,
89}
90
91impl Bridge {
92    /// Resolve the prompt contract and pick a transport, or fail with reasons.
93    pub async fn connect(client: reqwest::Client, mut config: BridgeConfig) -> Result<Self> {
94        let runtime = RuntimeClient::new(
95            client.clone(),
96            config.native_base.clone(),
97            config.upstream_model.clone(),
98            config.upstream_key.clone(),
99        );
100        let template = ChatTemplate::new(
101            runtime,
102            std::mem::take(&mut config.local),
103            config.chat_template_kwargs.clone(),
104        );
105
106        let slots = template
107            .resolve_slots()
108            .await
109            .context("resolving the single-token answer slots failed")?;
110
111        // A sixteen-option probe exercises the widest contract this bridge
112        // serves, so a transport that passes it will hold for real requests.
113        let probe_options: Vec<crate::wire::RowOption> = LETTERS
114            .chars()
115            .map(|letter| crate::wire::RowOption {
116                id: letter.to_string(),
117                description: format!("Probe option {letter}"),
118            })
119            .collect();
120        let messages = direct_messages(&json!("probe"), "probe", &probe_options);
121        let prompt = template
122            .render(&messages)
123            .await
124            .context("rendering the probe prompt failed")?;
125        template
126            .verify_boundary(&prompt, &slots)
127            .await
128            .context("verifying the answer-token boundary failed")?;
129
130        let letters: Vec<char> = LETTERS.chars().collect();
131        let mut chosen = None;
132        let mut failures = Vec::new();
133        for probe in crate::strategy::CANDIDATES {
134            match try_probe(
135                &client,
136                probe,
137                &config,
138                &prompt,
139                &slots,
140                &letters,
141            )
142            .await
143            {
144                Ok(()) => {
145                    chosen = Some(probe);
146                    break;
147                }
148                Err(error) => failures.push(format!("{}: {error:#}", probe.name())),
149            }
150        }
151        let probe = chosen.ok_or_else(|| {
152            anyhow::anyhow!(
153                "no supported transport answered the probe:\n  {}",
154                failures.join("\n  ")
155            )
156        })?;
157
158        Ok(Self {
159            client,
160            template,
161            probe,
162            upstream_model: config.upstream_model,
163            upstream_key: config.upstream_key,
164            openai_base: config.openai_base,
165            native_base: config.native_base,
166            slots,
167            served_model: config.served_model,
168            description: config.description,
169            release_date: config.release_date,
170            max_input_tokens: config.max_input_tokens,
171            readout: config.readout,
172            lock: Mutex::new(()),
173        })
174    }
175
176    /// The transport negotiated at startup.
177    pub fn probe(&self) -> Probe {
178        self.probe
179    }
180
181    /// What the deployment declared about the readout channel on this model.
182    pub fn readout(&self) -> &ReadoutStatus {
183        &self.readout
184    }
185
186    /// Resolved token id for each of the sixteen answer letters.
187    pub fn slots(&self) -> &[u32] {
188        &self.slots
189    }
190
191    /// The model name this bridge sends upstream.
192    pub fn upstream_model(&self) -> &str {
193        &self.upstream_model
194    }
195
196    /// Template variables forwarded on every render.
197    pub fn chat_template_kwargs(&self) -> Value {
198        self.template.chat_template_kwargs()
199    }
200
201    /// Whether the prompt is rendered in-process rather than by the runtime.
202    pub fn renders_locally(&self) -> bool {
203        self.template.renders_locally()
204    }
205
206    /// Whether tokenization happens in-process rather than by the runtime.
207    pub fn tokenizes_locally(&self) -> bool {
208        self.template.tokenizes_locally()
209    }
210
211    /// Score every row of one request against the same upstream model.
212    pub async fn score_rows(&self, rows: &[Row]) -> Result<Vec<DetailedScore>> {
213        let _guard = self.lock.lock().await;
214        let mut results = Vec::with_capacity(rows.len());
215        for row in rows {
216            results.push(self.score_row(row).await?);
217        }
218        Ok(results)
219    }
220
221    async fn score_row(&self, row: &Row) -> Result<DetailedScore> {
222        let count = row.options.len();
223        let messages = direct_messages(&row.state, &row.question, &row.options);
224        let prompt = self
225            .template
226            .render(&messages)
227            .await
228            .with_context(|| format!("row {:?}", row.id))?;
229
230        let slots = &self.slots[..count];
231        let letters: Vec<char> = LETTERS.chars().take(count).collect();
232        let body = build_body(self.probe, &self.upstream_model, &prompt, slots);
233        let url = request_url(self.probe, &self.openai_base, &self.native_base);
234        let response: CompletionResponse =
235            post_json(&self.client, &url, &body, self.upstream_key.as_deref()).await?;
236
237        if truncated(self.probe, &response) {
238            bail!(
239                "row {:?}: the upstream runtime truncated the prompt; shorten the state or raise \
240                 its context size",
241                row.id
242            );
243        }
244        let input_tokens = input_tokens(self.probe, &response);
245        if let Some(limit) = self.max_input_tokens
246            && input_tokens > limit as u64
247        {
248            bail!(
249                "row {:?}: {input_tokens} input tokens exceed limit {limit}; no truncation allowed",
250                row.id
251            );
252        }
253
254        let logprobs = parse_logprobs(self.probe, &response, slots, &letters)
255            .with_context(|| format!("row {:?}", row.id))?;
256        let probabilities = softmax(&logprobs);
257
258        Ok(DetailedScore {
259            answer: ScoredAnswer {
260                id: row.id.clone(),
261                option_ids: row.options.iter().map(|option| option.id.clone()).collect(),
262                probabilities,
263                input_tokens,
264                prompt_version: Some(PROMPT_VERSION.to_string()),
265            },
266            option_logprobs: logprobs,
267            prompt_sha256: prompt_sha256(&prompt),
268        })
269    }
270}
271
272async fn try_probe(
273    client: &reqwest::Client,
274    probe: Probe,
275    config: &BridgeConfig,
276    prompt: &str,
277    slots: &[u32],
278    letters: &[char],
279) -> Result<()> {
280    let body = build_body(probe, &config.upstream_model, prompt, slots);
281    let url = request_url(probe, &config.openai_base, &config.native_base);
282    let response: CompletionResponse = post_json(client, &url, &body, config.upstream_key.as_deref()).await?;
283    let logprobs = parse_logprobs(probe, &response, slots, letters)?;
284    crate::strategy::validate_probe(&logprobs)?;
285    if probe.restricts_distribution() {
286        let candidates = parse_candidates(probe, &response)?;
287        if !candidates_restricted_to(&candidates, slots, letters) {
288            bail!(
289                "the server ignored the distribution restriction and returned foreign tokens, so \
290                 this transport would not be what it claims"
291            );
292        }
293    }
294    Ok(())
295}
296
297/// POST a typed request body and decode a typed response.
298///
299/// The bridge always knows which shape it expects, so the response is decoded
300/// into a struct instead of being poked at as a generic document; a server that
301/// answers with something else fails here rather than later.
302pub async fn post_json<T: serde::de::DeserializeOwned>(
303    client: &reqwest::Client,
304    url: &str,
305    body: &(impl serde::Serialize + ?Sized),
306    api_key: Option<&str>,
307) -> Result<T> {
308    let mut request = client.post(url).json(body);
309    if let Some(key) = api_key {
310        request = request.bearer_auth(key);
311    }
312    let response = request
313        .send()
314        .await
315        .with_context(|| format!("POST {url} failed"))?;
316    let status = response.status();
317    let text = response
318        .text()
319        .await
320        .with_context(|| format!("reading the {url} response failed"))?;
321    if !status.is_success() {
322        bail!(
323            "POST {url} returned HTTP {status}: {}",
324            text.chars().take(400).collect::<String>()
325        );
326    }
327    serde_json::from_str(&text).with_context(|| {
328        format!(
329            "POST {url} returned an unexpected body: {}",
330            text.chars().take(400).collect::<String>()
331        )
332    })
333}
334
335/// Shared HTTP state.
336pub struct AppState {
337    /// The configured bridge, shared by every handler.
338    pub bridge: Bridge,
339    /// When set, clients must present this bearer token.
340    pub api_key: Option<String>,
341}
342
343/// The HTTP surface: `/v1/systemone`, `/v1/models` and `/health`.
344pub fn router(state: Arc<AppState>) -> Router {
345    Router::new()
346        .route("/v1/systemone", post(systemone))
347        .route("/v1/models", get(models))
348        .route("/health", get(health))
349        .with_state(state)
350}
351
352async fn systemone(State(state): State<Arc<AppState>>, headers: HeaderMap, body: Bytes) -> Response {
353    if let Err(response) = authorize(&state, &headers) {
354        return response;
355    }
356    let payload: Value = match serde_json::from_slice(&body) {
357        Ok(value) => value,
358        Err(error) => {
359            return error_response(
360                StatusCode::BAD_REQUEST,
361                format!("body: must be a JSON object ({error})"),
362            )
363        }
364    };
365    let (specs, rows) = match request_rows(&payload, &state.bridge.served_model) {
366        Ok(parsed) => parsed,
367        Err(error) => return error_response(StatusCode::BAD_REQUEST, error.0),
368    };
369    let results = match state.bridge.score_rows(&rows).await {
370        Ok(results) => results,
371        Err(error) => {
372            return error_response(StatusCode::BAD_GATEWAY, format!("{error:#}"));
373        }
374    };
375    let answers: Vec<ScoredAnswer> = results.into_iter().map(|scored| scored.answer).collect();
376    match response_from_results(
377        &state.bridge.served_model,
378        &specs,
379        &answers,
380        state.bridge.readout(),
381    ) {
382        Ok(response) => Json(response).into_response(),
383        Err(error) => error_response(StatusCode::INTERNAL_SERVER_ERROR, error.0),
384    }
385}
386
387/// One entry in `GET /v1/models`.
388#[derive(Debug, Serialize)]
389struct ModelEntry {
390    name: String,
391    description: String,
392    release_date: String,
393}
394
395#[derive(Debug, Serialize)]
396struct ModelsResponse {
397    models: Vec<ModelEntry>,
398}
399
400/// What the bridge negotiated, exposed so an operator can see it rather than
401/// guess which transport is in use.
402#[derive(Debug, Serialize)]
403struct HealthResponse {
404    status: &'static str,
405    served_model: String,
406    upstream_model: String,
407    probe: &'static str,
408    probe_note: &'static str,
409    endpoint: &'static str,
410    prompt_version: &'static str,
411    /// The resolved token id of every answer letter.
412    answer_slots: OrderedMap<u32>,
413    chat_template_kwargs: Value,
414    renders_locally: bool,
415    tokenizes_locally: bool,
416    probability_status: &'static str,
417    /// Whether the readout channel was validated on the served model.
418    readout: ReadoutStatus,
419}
420
421async fn models(State(state): State<Arc<AppState>>, headers: HeaderMap) -> Response {
422    if let Err(response) = authorize(&state, &headers) {
423        return response;
424    }
425    Json(ModelsResponse {
426        models: vec![ModelEntry {
427            name: state.bridge.served_model.clone(),
428            description: state.bridge.description.clone(),
429            release_date: state.bridge.release_date.clone(),
430        }],
431    })
432    .into_response()
433}
434
435async fn health(State(state): State<Arc<AppState>>, headers: HeaderMap) -> Response {
436    if let Err(response) = authorize(&state, &headers) {
437        return response;
438    }
439    let mut answer_slots = OrderedMap::new();
440    for (letter, id) in LETTERS.chars().zip(state.bridge.slots()) {
441        answer_slots.insert(letter.to_string(), *id);
442    }
443    Json(HealthResponse {
444        status: "ok",
445        served_model: state.bridge.served_model.clone(),
446        upstream_model: state.bridge.upstream_model().to_string(),
447        probe: state.bridge.probe().name(),
448        probe_note: state.bridge.probe().note(),
449        endpoint: state.bridge.probe().endpoint(),
450        prompt_version: PROMPT_VERSION,
451        answer_slots,
452        chat_template_kwargs: state.bridge.chat_template_kwargs(),
453        renders_locally: state.bridge.renders_locally(),
454        tokenizes_locally: state.bridge.tokenizes_locally(),
455        probability_status: PROBABILITY_STATUS,
456        readout: state.bridge.readout().clone(),
457    })
458    .into_response()
459}
460
461fn error_response(status: StatusCode, message: String) -> Response {
462    (status, Json(json!({"error": message}))).into_response()
463}
464
465/// Rejecting a request returns a full HTTP response, which is larger than the
466/// success value; boxing it would only add an allocation on the error path.
467#[allow(clippy::result_large_err)]
468fn authorize(state: &AppState, headers: &HeaderMap) -> Result<(), Response> {
469    let Some(expected) = &state.api_key else {
470        return Ok(());
471    };
472    let presented = headers
473        .get(AUTHORIZATION)
474        .and_then(|value| value.to_str().ok())
475        .and_then(|value| value.strip_prefix("Bearer "))
476        .unwrap_or_default();
477    if constant_time_eq(presented.as_bytes(), expected.as_bytes()) {
478        Ok(())
479    } else {
480        Err(error_response(
481            StatusCode::UNAUTHORIZED,
482            "missing or invalid bearer token".to_string(),
483        ))
484    }
485}
486
487fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
488    if left.len() != right.len() {
489        return false;
490    }
491    left.iter()
492        .zip(right)
493        .fold(0u8, |accumulator, (a, b)| accumulator | (a ^ b))
494        == 0
495}
jev_bridge/server.rs

jev_bridge/
server.rs
