jev_bridge/
strategy.rs

1//! Runtime probing of a generic OpenAI-compatible inference service.
2//!
3//! The bridge never guesses: every candidate transport is exercised with a
4//! real probe request before it is adopted, and the chosen probe is reported
5//! by `GET /health` so an operator can see exactly which contract is in use.
6//!
7//! Probe results on llama.cpp (b11096) shaped this module:
8//!   * `/v1/completions` answers with chat-shaped `logprobs.content[0].top_logprobs`
9//!     entries that carry an explicit `id`, and ignores `allowed_token_ids`.
10//!   * `/completion` exposes `completion_probabilities[0].top_logprobs` and
11//!     leaves `top_probs` null.
12//!   * An unrestricted top-k needs k >= 50 before all sixteen answer letters
13//!     appear, so unrestricted probes request a generous k.
14
15use anyhow::{anyhow, bail, Context, Result};
16use serde::{Deserialize, Serialize};
17use serde_json::Value;
18
19/// How the target service exposes option-token log probabilities.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Probe {
22    /// vLLM: restrict the distribution to the option tokens and ask for their
23    /// exact token ids. This is the same contract fastjev's own vLLM backend
24    /// uses through `SamplingParams`.
25    LogprobTokenIds,
26    /// vLLM without `logprob_token_ids`: restrict the distribution and read
27    /// the top-k, which must then contain every option token.
28    AllowedTokenIdsTopK,
29    /// OpenAI-compatible completions: read the top-k of the unrestricted
30    /// next-token distribution.
31    CompletionsTopK,
32    /// llama.cpp native `/completion` with `n_probs`.
33    LlamaCppNProbs,
34}
35
36/// The transports to try, in the order they are probed.
37pub const CANDIDATES: [Probe; 4] = [
38    Probe::LogprobTokenIds,
39    Probe::AllowedTokenIdsTopK,
40    Probe::CompletionsTopK,
41    Probe::LlamaCppNProbs,
42];
43
44impl Probe {
45    /// Stable identifier reported by `GET /health`.
46    pub fn name(self) -> &'static str {
47        match self {
48            Probe::LogprobTokenIds => "vllm-logprob-token-ids",
49            Probe::AllowedTokenIdsTopK => "vllm-allowed-token-ids-top-k",
50            Probe::CompletionsTopK => "openai-completions-top-k",
51            Probe::LlamaCppNProbs => "llamacpp-native-n-probs",
52        }
53    }
54
55    /// One line describing what this transport relies on.
56    pub fn note(self) -> &'static str {
57        match self {
58            Probe::LogprobTokenIds => {
59                "option distribution is restricted to the declared answer tokens; exact token ids requested"
60            }
61            Probe::AllowedTokenIdsTopK => {
62                "option distribution is restricted to the declared answer tokens; top-k must cover them"
63            }
64            Probe::CompletionsTopK => {
65                "unrestricted next-token distribution; every answer token must fall inside top-k"
66            }
67            Probe::LlamaCppNProbs => {
68                "llama.cpp native completion probabilities; every answer token must fall inside n_probs"
69            }
70        }
71    }
72
73    /// Whether the request is expected to restrict the distribution to the
74    /// answer tokens. A probe that claims restriction but returns foreign
75    /// candidates is rejected rather than silently trusted.
76    pub fn restricts_distribution(self) -> bool {
77        matches!(self, Probe::LogprobTokenIds | Probe::AllowedTokenIdsTopK)
78    }
79
80    /// The path this transport posts to.
81    pub fn endpoint(self) -> &'static str {
82        match self {
83            Probe::LlamaCppNProbs => "/completion",
84            _ => "/completions",
85        }
86    }
87
88    fn url(self, openai_base: &str, native_base: &str) -> String {
89        match self {
90            Probe::LlamaCppNProbs => format!("{}/completion", native_base.trim_end_matches('/')),
91            _ => format!("{}/completions", openai_base.trim_end_matches('/')),
92        }
93    }
94}
95
96/// Candidate count requested from a top-k endpoint.
97///
98/// A restricted distribution only ever contains the answer tokens, so k equal
99/// to the option count suffices. An unrestricted distribution also contains
100/// ordinary vocabulary, and measurement on a 2B llama.cpp model showed four of
101/// sixteen answer letters falling outside k=20, so unrestricted probes ask for
102/// a wide margin instead.
103fn top_k(probe: Probe, slot_count: usize) -> usize {
104    match probe {
105        Probe::LogprobTokenIds => 1,
106        Probe::AllowedTokenIdsTopK => slot_count.max(4),
107        Probe::CompletionsTopK | Probe::LlamaCppNProbs => (slot_count * 4).clamp(32, 128),
108    }
109}
110
111/// The request body for an OpenAI-compatible `/v1/completions` call.
112///
113/// Optional fields are omitted rather than sent as `null`: a runtime that does
114/// not know `allowed_token_ids` is likelier to answer a request without it, and
115/// the probe already checks whether the field had any effect.
116#[derive(Debug, Serialize)]
117pub struct OpenAiCompletionRequest {
118    model: String,
119    prompt: String,
120    /// One token is enough: only the first position's distribution is read.
121    max_tokens: u32,
122    /// Left at 1.0 so the returned log probabilities are the model's own.
123    temperature: f32,
124    /// How many candidates to return from the head of the distribution.
125    logprobs: usize,
126    /// Ask vLLM to label candidates by token id instead of token text, which
127    /// removes any ambiguity about leading spaces.
128    #[serde(skip_serializing_if = "Option::is_none")]
129    return_tokens_as_token_ids: Option<bool>,
130    /// vLLM extension: restrict sampling to these token ids, which turns the
131    /// returned log probabilities into the conditional option distribution.
132    #[serde(skip_serializing_if = "Option::is_none")]
133    allowed_token_ids: Option<Vec<u32>>,
134    /// vLLM extension: report these exact token ids even when they fall
135    /// outside the top-k.
136    #[serde(skip_serializing_if = "Option::is_none")]
137    logprob_token_ids: Option<Vec<u32>>,
138}
139
140/// The request body for llama.cpp's native `/completion`.
141///
142/// Every sampler knob is set explicitly and neutrally. llama.cpp reports the
143/// distribution *after* the sampler chain, so inheriting a server preset such
144/// as `repeat-penalty = 1.05` or `top-p = 0.85` would return a distorted
145/// distribution instead of the model's own.
146#[derive(Debug, Serialize)]
147pub struct LlamaCppCompletionRequest {
148    model: String,
149    prompt: String,
150    n_predict: u32,
151    /// How many candidates to return, matching `logprobs` above.
152    n_probs: usize,
153    cache_prompt: bool,
154    temperature: f32,
155    /// 0 disables top-k truncation.
156    top_k: i32,
157    top_p: f32,
158    min_p: f32,
159    repeat_penalty: f32,
160    presence_penalty: f32,
161    frequency_penalty: f32,
162    typical_p: f32,
163}
164
165/// A scoring request in whichever dialect the chosen probe speaks.
166#[derive(Debug, Serialize)]
167#[serde(untagged)]
168pub enum RequestBody {
169    /// `/v1/completions` in the OpenAI dialect.
170    OpenAi(OpenAiCompletionRequest),
171    /// llama.cpp's native `/completion`.
172    LlamaCpp(LlamaCppCompletionRequest),
173}
174
175/// Build the scoring request for one probe.
176pub fn build_body(probe: Probe, model: &str, prompt: &str, slots: &[u32]) -> RequestBody {
177    let candidates = top_k(probe, slots.len());
178    match probe {
179        Probe::LogprobTokenIds => RequestBody::OpenAi(OpenAiCompletionRequest {
180            model: model.to_string(),
181            prompt: prompt.to_string(),
182            max_tokens: 1,
183            temperature: 1.0,
184            logprobs: candidates,
185            return_tokens_as_token_ids: Some(true),
186            allowed_token_ids: Some(slots.to_vec()),
187            logprob_token_ids: Some(slots.to_vec()),
188        }),
189        Probe::AllowedTokenIdsTopK => RequestBody::OpenAi(OpenAiCompletionRequest {
190            model: model.to_string(),
191            prompt: prompt.to_string(),
192            max_tokens: 1,
193            temperature: 1.0,
194            logprobs: candidates,
195            return_tokens_as_token_ids: Some(true),
196            allowed_token_ids: Some(slots.to_vec()),
197            logprob_token_ids: None,
198        }),
199        Probe::CompletionsTopK => RequestBody::OpenAi(OpenAiCompletionRequest {
200            model: model.to_string(),
201            prompt: prompt.to_string(),
202            max_tokens: 1,
203            temperature: 1.0,
204            logprobs: candidates,
205            return_tokens_as_token_ids: Some(true),
206            allowed_token_ids: None,
207            logprob_token_ids: None,
208        }),
209        Probe::LlamaCppNProbs => RequestBody::LlamaCpp(LlamaCppCompletionRequest {
210            model: model.to_string(),
211            prompt: prompt.to_string(),
212            n_predict: 1,
213            n_probs: candidates,
214            cache_prompt: true,
215            temperature: 1.0,
216            top_k: 0,
217            top_p: 1.0,
218            min_p: 0.0,
219            repeat_penalty: 1.0,
220            presence_penalty: 0.0,
221            frequency_penalty: 0.0,
222            typical_p: 1.0,
223        }),
224    }
225}
226
227/// The full URL a probe posts to.
228pub fn request_url(probe: Probe, openai_base: &str, native_base: &str) -> String {
229    probe.url(openai_base, native_base)
230}
231
232/// One candidate token in a returned distribution.
233#[derive(Debug, Clone, PartialEq)]
234pub struct Candidate {
235    /// Token id, when the runtime reported one.
236    pub id: Option<u32>,
237    /// Token text, or `token_id:<n>` in the OpenAI map shape.
238    pub token: String,
239    /// Natural-log probability.
240    pub logprob: f64,
241}
242
243/// One candidate as a runtime serialised it.
244///
245/// llama.cpp sends an array of these on both endpoints, with an explicit `id`
246/// and a `logprob`; older builds send `prob` instead. Everything is optional
247/// because the bridge must report a *missing* candidate rather than default one.
248#[derive(Debug, Clone, Deserialize)]
249struct CandidateEntry {
250    #[serde(default)]
251    id: Option<u32>,
252    #[serde(default)]
253    token: String,
254    #[serde(default)]
255    logprob: Option<f64>,
256    #[serde(default)]
257    prob: Option<f64>,
258}
259
260/// A position's candidates, in whichever shape the runtime uses.
261///
262/// OpenAI and vLLM return a map from token text — or `token_id:<n>` when
263/// `return_tokens_as_token_ids` is honoured — to a log probability. llama.cpp
264/// returns an array of entries. Both are accepted, and the map's values stay
265/// dynamic because they may be a bare number or an object with a `logprob`.
266#[derive(Debug, Deserialize)]
267#[serde(untagged)]
268enum CandidateList {
269    Entries(Vec<CandidateEntry>),
270    Map(serde_json::Map<String, Value>),
271}
272
273/// One generated position on llama.cpp's native `/completion`.
274#[derive(Debug, Deserialize)]
275struct NativePosition {
276    #[serde(default)]
277    top_logprobs: Option<CandidateList>,
278    /// Only consulted when `top_logprobs` is absent: a probability carries less
279    /// precision than a log probability.
280    #[serde(default)]
281    top_probs: Option<CandidateList>,
282}
283
284/// One generated position in an OpenAI-shaped choice.
285#[derive(Debug, Deserialize)]
286struct ChoiceLogprobs {
287    /// OpenAI completions shape: one entry per generated position.
288    #[serde(default)]
289    top_logprobs: Vec<CandidateList>,
290    /// llama.cpp answers the completions endpoint with chat-shaped content.
291    #[serde(default)]
292    content: Vec<NativePosition>,
293}
294
295#[derive(Debug, Deserialize)]
296struct Choice {
297    #[serde(default)]
298    finish_reason: Option<String>,
299    #[serde(default)]
300    logprobs: Option<ChoiceLogprobs>,
301}
302
303#[derive(Debug, Deserialize)]
304struct Usage {
305    #[serde(default)]
306    prompt_tokens: Option<u64>,
307}
308
309/// Every field the bridge reads from a scoring response.
310///
311/// One struct covers both endpoints: a runtime fills the fields its dialect
312/// uses and leaves the rest absent, and the probe that was chosen already
313/// decided which fields to read.
314#[derive(Debug, Deserialize)]
315pub struct CompletionResponse {
316    #[serde(default)]
317    choices: Vec<Choice>,
318    #[serde(default)]
319    usage: Option<Usage>,
320    /// llama.cpp native `/completion`.
321    #[serde(default)]
322    completion_probabilities: Vec<NativePosition>,
323    #[serde(default)]
324    tokens_evaluated: Option<u64>,
325    #[serde(default)]
326    truncated: bool,
327}
328
329fn candidate_logprob(entry: &CandidateEntry) -> Option<f64> {
330    entry
331        .logprob
332        .or_else(|| entry.prob.map(f64::ln))
333        .filter(|value| value.is_finite())
334}
335
336fn map_logprob(value: &Value) -> Option<f64> {
337    if let Some(number) = value.as_f64() {
338        return Some(number);
339    }
340    value
341        .get("logprob")
342        .and_then(Value::as_f64)
343        .or_else(|| value.get("prob").and_then(Value::as_f64).map(f64::ln))
344}
345
346/// Normalize the several shapes servers use for one position's candidates.
347fn collect_candidates(list: &CandidateList) -> Vec<Candidate> {
348    match list {
349        CandidateList::Entries(entries) => entries
350            .iter()
351            .filter_map(|entry| {
352                Some(Candidate {
353                    id: entry.id,
354                    token: entry.token.clone(),
355                    logprob: candidate_logprob(entry)?,
356                })
357            })
358            .collect(),
359        CandidateList::Map(map) => map
360            .iter()
361            .filter_map(|(key, value)| {
362                let logprob = map_logprob(value)?;
363                let id = key
364                    .strip_prefix("token_id:")
365                    .and_then(|text| text.parse::<u32>().ok());
366                Some(Candidate {
367                    id,
368                    token: key.clone(),
369                    logprob,
370                })
371            })
372            .collect(),
373    }
374}
375
376/// Locate the candidate list for the first generated position.
377///
378/// The order within each probe is deliberate: the shape that probe is expected
379/// to receive is tried first, and the alternative is only a fallback for
380/// runtimes that answer in the other dialect.
381fn candidate_position(
382    probe: Probe,
383    response: &CompletionResponse,
384) -> Result<&CandidateList> {
385    let choice = response.choices.first();
386    let logprobs = choice.and_then(|choice| choice.logprobs.as_ref());
387    match probe {
388        // llama.cpp answers `/v1/completions` with chat-shaped logprobs.
389        Probe::CompletionsTopK => logprobs
390            .and_then(|logprobs| logprobs.content.first())
391            .and_then(|position| position.top_logprobs.as_ref())
392            .or_else(|| logprobs.and_then(|logprobs| logprobs.top_logprobs.first()))
393            .ok_or_else(|| anyhow!("response has no candidates under choices[0].logprobs")),
394        Probe::LogprobTokenIds | Probe::AllowedTokenIdsTopK => logprobs
395            .and_then(|logprobs| logprobs.top_logprobs.first())
396            .or_else(|| {
397                logprobs
398                    .and_then(|logprobs| logprobs.content.first())
399                    .and_then(|position| position.top_logprobs.as_ref())
400            })
401            .ok_or_else(|| anyhow!("response has no candidates under choices[0].logprobs")),
402        Probe::LlamaCppNProbs => response
403            .completion_probabilities
404            .first()
405            .and_then(|position| {
406                position
407                    .top_logprobs
408                    .as_ref()
409                    .or(position.top_probs.as_ref())
410            })
411            .ok_or_else(|| anyhow!("response has no candidates under completion_probabilities[0]")),
412    }
413}
414
415fn find_candidate(candidates: &[Candidate], slot: u32, letter: char) -> Option<f64> {
416    let by_id = candidates.iter().find(|candidate| candidate.id == Some(slot));
417    let text = letter.to_string();
418    by_id
419        .or_else(|| candidates.iter().find(|candidate| candidate.token == text))
420        .map(|candidate| candidate.logprob)
421}
422
423/// Decode a scoring response into the struct the bridge reads.
424pub fn parse_response(body: &Value) -> Result<CompletionResponse> {
425    serde_json::from_value(body.clone())
426        .context("the scoring response did not match any known shape")
427}
428
429/// Extract one log probability per requested slot, in slot order.
430///
431/// A missing slot is an error rather than a zero: silently dropping an option
432/// would change the meaning of the returned distribution.
433pub fn parse_logprobs(
434    probe: Probe,
435    response: &CompletionResponse,
436    slots: &[u32],
437    letters: &[char],
438) -> Result<Vec<f64>> {
439    if let Some(finish) = response
440        .choices
441        .first()
442        .and_then(|choice| choice.finish_reason.as_deref())
443        && finish == "error"
444    {
445        bail!("service reported finish_reason=error");
446    }
447    let candidates = collect_candidates(candidate_position(probe, response)?);
448    if candidates.is_empty() {
449        bail!("returned distribution contains no usable candidates");
450    }
451    let mut result = Vec::with_capacity(slots.len());
452    for (slot, letter) in slots.iter().zip(letters) {
453        let logprob = find_candidate(&candidates, *slot, *letter).ok_or_else(|| {
454            anyhow!(
455                "option token {letter:?} (id {slot}) is absent from the returned distribution; \
456                 raise the requested top-k or restrict the distribution to the answer tokens"
457            )
458        })?;
459        if !logprob.is_finite() {
460            bail!("logprob for option token {letter:?} is not finite");
461        }
462        result.push(logprob);
463    }
464    Ok(result)
465}
466
467/// Whether every returned candidate belongs to the declared answer tokens.
468///
469/// A probe that asks for a restricted distribution but receives foreign
470/// candidates has been silently ignored by the server, which would make the
471/// reported transport a lie.
472pub fn candidates_restricted_to(candidates: &[Candidate], slots: &[u32], letters: &[char]) -> bool {
473    candidates.iter().all(|candidate| {
474        candidate
475            .id
476            .map(|id| slots.contains(&id))
477            .unwrap_or_else(|| candidate.token.chars().count() == 1 && letters.contains(&candidate.token.chars().next().unwrap()))
478    })
479}
480
481/// Decode the returned candidates without requiring any slot to be present.
482pub fn parse_candidates(probe: Probe, response: &CompletionResponse) -> Result<Vec<Candidate>> {
483    Ok(collect_candidates(candidate_position(probe, response)?))
484}
485
486/// The prompt length the runtime reported, or zero when it reported none.
487pub fn input_tokens(probe: Probe, response: &CompletionResponse) -> u64 {
488    match probe {
489        Probe::LlamaCppNProbs => response.tokens_evaluated.unwrap_or(0),
490        _ => response
491            .usage
492            .as_ref()
493            .and_then(|usage| usage.prompt_tokens)
494            .unwrap_or(0),
495    }
496}
497
498/// Whether the runtime silently shortened the prompt.
499///
500/// A truncated prompt would score a decision the caller never asked for, so the
501/// bridge turns this into an error instead of a slightly wrong answer.
502pub fn truncated(probe: Probe, response: &CompletionResponse) -> bool {
503    match probe {
504        Probe::LlamaCppNProbs => response.truncated,
505        _ => false,
506    }
507}
508
509/// Validate that a probe response is usable before adopting its transport.
510pub fn validate_probe(logprobs: &[f64]) -> Result<()> {
511    if logprobs.len() < 2 {
512        bail!("probe returned fewer than two option scores");
513    }
514    if logprobs.iter().any(|value| !value.is_finite()) {
515        bail!("probe returned a non-finite option score");
516    }
517    if logprobs.iter().any(|value| *value > 1e-6) {
518        bail!("probe returned a positive log probability, which cannot be a log softmax value");
519    }
520    Ok(())
521}
522
523#[cfg(test)]
524mod tests {
525    use super::*;
526    use serde_json::json;
527
528    #[test]
529    fn logprob_token_ids_body_restricts_and_selects() {
530        let RequestBody::OpenAi(body) = build_body(Probe::LogprobTokenIds, "m", "prompt", &[32, 33])
531        else {
532            panic!("this probe speaks the OpenAI dialect");
533        };
534        assert_eq!(body.allowed_token_ids.as_deref(), Some(&[32u32, 33][..]));
535        assert_eq!(body.logprob_token_ids.as_deref(), Some(&[32u32, 33][..]));
536        assert_eq!(body.max_tokens, 1);
537        assert_eq!(body.logprobs, 1, "only the requested ids need reporting");
538    }
539
540    #[test]
541    fn unrestricted_probes_request_a_wide_top_k() {
542        assert_eq!(top_k(Probe::CompletionsTopK, 16), 64);
543        assert_eq!(top_k(Probe::LlamaCppNProbs, 16), 64);
544        assert_eq!(top_k(Probe::CompletionsTopK, 2), 32);
545        assert_eq!(top_k(Probe::AllowedTokenIdsTopK, 16), 16);
546        assert_eq!(top_k(Probe::LogprobTokenIds, 16), 1);
547    }
548
549    #[test]
550    fn llamacpp_body_neutralizes_the_sampler() {
551        let RequestBody::LlamaCpp(body) = build_body(Probe::LlamaCppNProbs, "m", "prompt", &[54, 55])
552        else {
553            panic!("this probe speaks the llama.cpp dialect");
554        };
555        assert_eq!(body.n_predict, 1);
556        assert_eq!(body.n_probs, 32);
557        assert_eq!(body.temperature, 1.0);
558        assert_eq!(body.top_k, 0, "top-k truncation would hide low-ranked options");
559        assert_eq!(body.repeat_penalty, 1.0, "a penalty would distort the distribution");
560    }
561
562    #[test]
563    fn an_unrestricted_probe_omits_the_vllm_only_fields() {
564        let RequestBody::OpenAi(body) = build_body(Probe::CompletionsTopK, "m", "prompt", &[54, 55])
565        else {
566            panic!("this probe speaks the OpenAI dialect");
567        };
568        assert!(body.allowed_token_ids.is_none());
569        assert!(body.logprob_token_ids.is_none());
570        assert_eq!(body.logprobs, 32);
571    }
572
573    #[test]
574    fn vllm_map_response_is_parsed_by_token_id() {
575        let body = json!({"choices": [{"logprobs": {"top_logprobs": [{
576            "token_id:32": -0.51, "token_id:33": -1.2, "token_id:99": -4.0
577        }]}}]});
578        let values = parse_logprobs(Probe::LogprobTokenIds, &parse_response(&body).unwrap(), &[32, 33], &['A', 'B']).unwrap();
579        assert_eq!(values, vec![-0.51, -1.2]);
580    }
581
582    #[test]
583    fn vllm_map_response_falls_back_to_token_text() {
584        let body = json!({"choices": [{"logprobs": {"top_logprobs": [{"A": -0.1, "B": -2.0}]}}]});
585        let values = parse_logprobs(Probe::CompletionsTopK, &parse_response(&body).unwrap(), &[32, 33], &['A', 'B']).unwrap();
586        assert_eq!(values, vec![-0.1, -2.0]);
587    }
588
589    #[test]
590    fn llamacpp_completions_response_is_parsed_by_id() {
591        // Shape observed on llama.cpp b11096: chat-style logprobs on the
592        // completions endpoint, with an explicit token id.
593        let body = json!({"choices": [{"finish_reason": "length", "logprobs": {"content": [{
594            "id": 55, "token": "B", "logprob": -2.1e-05,
595            "top_logprobs": [
596                {"id": 55, "token": "B", "logprob": -2.1e-05},
597                {"id": 54, "token": "A", "logprob": -12.32},
598                {"id": 56, "token": "C", "logprob": -15.61}
599            ]}]}}]});
600        let values =
601            parse_logprobs(Probe::CompletionsTopK, &parse_response(&body).unwrap(), &[54, 55, 56], &['A', 'B', 'C']).unwrap();
602        assert_eq!(values, vec![-12.32, -2.1e-05, -15.61]);
603    }
604
605    #[test]
606    fn llamacpp_native_response_prefers_top_logprobs() {
607        let body = json!({"completion_probabilities": [{"id": 32, "token": "A", "prob": 0.6,
608            "top_probs": [{"id": 32, "token": "A", "prob": 0.6}],
609            "top_logprobs": [{"id": 32, "token": "A", "logprob": -0.51},
610                             {"id": 33, "token": "B", "logprob": -0.92}]}]});
611        let values = parse_logprobs(Probe::LlamaCppNProbs, &parse_response(&body).unwrap(), &[32, 33], &['A', 'B']).unwrap();
612        assert_eq!(values, vec![-0.51, -0.92]);
613    }
614
615    #[test]
616    fn llamacpp_native_response_falls_back_to_probabilities() {
617        let body = json!({"completion_probabilities": [{"id": 32, "token": "A",
618            "top_probs": [{"id": 32, "token": "A", "prob": 0.5}, {"id": 33, "token": "B", "prob": 0.5}]}]});
619        let values = parse_logprobs(Probe::LlamaCppNProbs, &parse_response(&body).unwrap(), &[32, 33], &['A', 'B']).unwrap();
620        assert!((values[0] - 0.5f64.ln()).abs() < 1e-12);
621    }
622
623    #[test]
624    fn missing_option_is_an_error_not_a_zero() {
625        let body = json!({"choices": [{"logprobs": {"top_logprobs": [{"token_id:32": -0.1}]}}]});
626        let error =
627            parse_logprobs(Probe::CompletionsTopK, &parse_response(&body).unwrap(), &[32, 33], &['A', 'B']).unwrap_err();
628        assert!(error.to_string().contains("absent from the returned distribution"));
629    }
630
631    #[test]
632    fn foreign_candidates_reveal_an_unrestricted_distribution() {
633        let restricted = vec![Candidate { id: Some(54), token: "A".into(), logprob: -0.1 }];
634        let unrestricted = vec![
635            Candidate { id: Some(54), token: "A".into(), logprob: -0.1 },
636            Candidate { id: Some(608), token: "The".into(), logprob: -9.0 },
637        ];
638        assert!(candidates_restricted_to(&restricted, &[54, 55], &['A', 'B']));
639        assert!(!candidates_restricted_to(&unrestricted, &[54, 55], &['A', 'B']));
640    }
641
642    #[test]
643    fn probe_validation_rejects_positive_logprobs() {
644        assert!(validate_probe(&[-0.5, -1.5]).is_ok());
645        assert!(validate_probe(&[0.5, -1.5]).is_err());
646    }
647
648    #[test]
649    fn urls_split_between_openai_and_native_bases() {
650        assert_eq!(
651            request_url(Probe::CompletionsTopK, "http://host:8000/v1/", "http://host:8000"),
652            "http://host:8000/v1/completions"
653        );
654        assert_eq!(
655            request_url(Probe::LlamaCppNProbs, "http://host:8080/v1", "http://host:8080/"),
656            "http://host:8080/completion"
657        );
658    }
659}
jev_bridge/strategy.rs

jev_bridge/
strategy.rs
