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

jev_bridge/
server.rs
