llm_token_visualizer/
live.rs

1//! Fetch a completion with token logprobs from an OpenAI-compatible API.
2//!
3//! Pick a provider preset with [`Provider`]; each one knows its default base
4//! URL, the environment variable that holds its API key (if it needs one), a
5//! default model where one is stable enough to hardcode, and how that API
6//! wants logprobs requested.
7//!
8//! Without a preset the behaviour is the same as before 0.4.0: OpenAI, with
9//! `OPENAI_API_KEY` (required) and `OPENAI_BASE_URL` (optional). Any server
10//! that implements Chat Completions with `logprobs` works through `base_url`.
11
12use anyhow::{bail, Context, Result};
13use serde_json::{json, Value};
14
15pub const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
16pub const DEFAULT_MODEL: &str = "gpt-4o-mini";
17
18/// Number of alternatives requested per token.
19pub const TOP_LOGPROBS: u32 = 3;
20
21/// How a provider wants logprobs requested.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum LogprobsParam {
24    /// OpenAI style: `"logprobs": true, "top_logprobs": N`.
25    BoolWithTop,
26    /// Together style: `"logprobs": N` (an integer, the number of alternatives).
27    Integer,
28}
29
30/// A provider preset.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Provider {
33    OpenAi,
34    OpenRouter,
35    Together,
36    Vllm,
37    Ollama,
38}
39
40impl Provider {
41    pub const ALL: [Provider; 5] = [
42        Provider::OpenAi,
43        Provider::OpenRouter,
44        Provider::Together,
45        Provider::Vllm,
46        Provider::Ollama,
47    ];
48
49    pub fn name(self) -> &'static str {
50        match self {
51            Provider::OpenAi => "openai",
52            Provider::OpenRouter => "openrouter",
53            Provider::Together => "together",
54            Provider::Vllm => "vllm",
55            Provider::Ollama => "ollama",
56        }
57    }
58
59    pub fn from_name(name: &str) -> Option<Provider> {
60        let n = name.trim().to_ascii_lowercase();
61        Provider::ALL.into_iter().find(|p| p.name() == n)
62    }
63
64    /// Base URL used when `--base-url` is not given.
65    pub fn default_base_url(self) -> &'static str {
66        match self {
67            Provider::OpenAi => DEFAULT_BASE_URL,
68            Provider::OpenRouter => "https://openrouter.ai/api/v1",
69            Provider::Together => "https://api.together.ai/v1",
70            Provider::Vllm => "http://localhost:8000/v1",
71            Provider::Ollama => "http://localhost:11434/v1",
72        }
73    }
74
75    /// Environment variable holding the API key, and whether it is required.
76    /// vLLM only needs a key when the server was started with `--api-key`.
77    pub fn key_env(self) -> Option<(&'static str, bool)> {
78        match self {
79            Provider::OpenAi => Some(("OPENAI_API_KEY", true)),
80            Provider::OpenRouter => Some(("OPENROUTER_API_KEY", true)),
81            Provider::Together => Some(("TOGETHER_API_KEY", true)),
82            Provider::Vllm => Some(("VLLM_API_KEY", false)),
83            Provider::Ollama => None,
84        }
85    }
86
87    /// Model used when `--model` is not given. `None` means the user must
88    /// pick one, because the provider's catalog (or the local server) decides.
89    pub fn default_model(self) -> Option<&'static str> {
90        match self {
91            Provider::OpenAi => Some(DEFAULT_MODEL),
92            Provider::OpenRouter => Some("openai/gpt-4o-mini"),
93            Provider::Together | Provider::Vllm | Provider::Ollama => None,
94        }
95    }
96
97    /// What to suggest when no model was given.
98    pub fn model_hint(self) -> &'static str {
99        match self {
100            Provider::Together => {
101                "a Together model id (GET https://api.together.ai/v1/models lists them)"
102            }
103            Provider::Vllm => "the model the server was started with (see GET /v1/models)",
104            Provider::Ollama => {
105                "a model you have pulled, as listed by `ollama list`, e.g. qwen2.5-coder:14b"
106            }
107            _ => "a model id",
108        }
109    }
110
111    pub fn logprobs_param(self) -> LogprobsParam {
112        match self {
113            Provider::Together => LogprobsParam::Integer,
114            _ => LogprobsParam::BoolWithTop,
115        }
116    }
117}
118
119/// Resolved settings for one request.
120#[derive(Debug, Clone, PartialEq)]
121pub struct Target {
122    pub provider: Provider,
123    pub base_url: String,
124    pub model: String,
125    pub api_key: Option<String>,
126}
127
128impl Target {
129    /// Resolve a provider, optional overrides and an environment lookup into a
130    /// request target. `env` is a parameter so this can be tested without
131    /// touching the process environment.
132    ///
133    /// For the OpenAI preset, `OPENAI_BASE_URL` still applies (backwards
134    /// compatible with 0.3); `base_url` beats it.
135    pub fn resolve(
136        provider: Provider,
137        base_url: Option<&str>,
138        model: Option<&str>,
139        env: impl Fn(&str) -> Option<String>,
140    ) -> Result<Target> {
141        let base_url = match base_url {
142            Some(u) => u.to_string(),
143            None if provider == Provider::OpenAi => {
144                env("OPENAI_BASE_URL").unwrap_or_else(|| DEFAULT_BASE_URL.to_string())
145            }
146            None => provider.default_base_url().to_string(),
147        };
148        let base_url = base_url.trim().trim_end_matches('/').to_string();
149        if base_url.is_empty() {
150            bail!("--base-url is empty");
151        }
152
153        let model = match model.or(provider.default_model()) {
154            Some(m) => m.to_string(),
155            None => bail!(
156                "--provider {} needs --model: {}",
157                provider.name(),
158                provider.model_hint()
159            ),
160        };
161
162        let api_key = match provider.key_env() {
163            Some((var, required)) => match env(var).filter(|k| !k.trim().is_empty()) {
164                Some(k) => Some(k),
165                None if required => bail!(
166                    "--provider {} needs {} in the environment",
167                    provider.name(),
168                    var
169                ),
170                None => None,
171            },
172            None => None,
173        };
174
175        Ok(Target {
176            provider,
177            base_url,
178            model,
179            api_key,
180        })
181    }
182
183    pub fn url(&self) -> String {
184        format!("{}/chat/completions", self.base_url)
185    }
186
187    /// Build the request body. Temperature 0 so reruns are comparable.
188    pub fn body(&self, prompt: &str, max_tokens: u32) -> Value {
189        let mut b = json!({
190            "model": self.model,
191            "messages": [{"role": "user", "content": prompt}],
192            "max_tokens": max_tokens,
193            "temperature": 0,
194        });
195        match self.provider.logprobs_param() {
196            LogprobsParam::BoolWithTop => {
197                b["logprobs"] = json!(true);
198                b["top_logprobs"] = json!(TOP_LOGPROBS);
199            }
200            LogprobsParam::Integer => {
201                b["logprobs"] = json!(TOP_LOGPROBS);
202            }
203        }
204        if self.provider == Provider::OpenRouter {
205            // Only route to upstream providers that support every parameter
206            // in the request, logprobs included.
207            b["provider"] = json!({"require_parameters": true});
208        }
209        b
210    }
211}
212
213/// Build an OpenAI-style request body (kept for library callers of 0.3).
214pub fn request_body(prompt: &str, model: &str, max_tokens: u32) -> Value {
215    Target {
216        provider: Provider::OpenAi,
217        base_url: DEFAULT_BASE_URL.to_string(),
218        model: model.to_string(),
219        api_key: None,
220    }
221    .body(prompt, max_tokens)
222}
223
224/// Check that a successful response actually carries token logprobs, and
225/// explain clearly when it does not (instead of an empty report).
226pub fn check_has_logprobs(raw: &str, target: &Target) -> Result<()> {
227    let v: Value = serde_json::from_str(raw).context("the API response is not JSON")?;
228    if let Some(err) = v.get("error") {
229        bail!("{} returned an error: {}", target.url(), err);
230    }
231    let choice = v.pointer("/choices/0");
232    let lp = choice.and_then(|c| c.get("logprobs"));
233    let has = match lp {
234        Some(Value::Object(o)) => {
235            o.get("content")
236                .and_then(Value::as_array)
237                .is_some_and(|a| !a.is_empty())
238                || o.get("tokens")
239                    .and_then(Value::as_array)
240                    .is_some_and(|a| !a.is_empty())
241        }
242        Some(Value::Array(a)) => !a.is_empty(),
243        _ => false,
244    };
245    if has {
246        return Ok(());
247    }
248    let answer = choice
249        .and_then(|c| c.pointer("/message/content"))
250        .and_then(Value::as_str)
251        .unwrap_or("");
252    let mut msg = format!(
253        "{} (model {}) answered, but returned no token logprobs, so there is nothing to analyze.",
254        target.provider.name(),
255        target.model
256    );
257    msg.push_str(match target.provider {
258        Provider::Ollama => {
259            " Ollama returns logprobs on its OpenAI-compatible endpoint in recent versions; update Ollama and try again."
260        }
261        Provider::OpenRouter => {
262            " Pick a model whose upstream provider supports logprobs (the model's page on openrouter.ai lists supported parameters)."
263        }
264        Provider::OpenAi => " Not every model returns logprobs; OpenAI's logprobs examples use gpt-4o-mini.",
265        _ => " This model or server may not support logprobs.",
266    });
267    if !answer.is_empty() {
268        let short: String = answer.chars().take(200).collect();
269        msg.push_str(&format!("\nThe answer was: {short}"));
270    }
271    bail!(msg)
272}
273
274/// Send the prompt and return the raw JSON response text.
275pub fn fetch_target(target: &Target, prompt: &str, max_tokens: u32) -> Result<String> {
276    let url = target.url();
277    let mut req = ureq::post(&url).set("Content-Type", "application/json");
278    if let Some(key) = &target.api_key {
279        req = req.set("Authorization", &format!("Bearer {}", key));
280    }
281    if target.provider == Provider::OpenRouter {
282        req = req
283            .set(
284                "HTTP-Referer",
285                "https://gitlab.com/mattbusel/LLM-Hallucination-Detection-Script",
286            )
287            .set("X-Title", "llm-token-visualizer");
288    }
289    let resp = req.send_string(&target.body(prompt, max_tokens).to_string());
290
291    match resp {
292        Ok(r) => r.into_string().context("could not read the API response"),
293        Err(ureq::Error::Status(code, r)) => {
294            let body = r.into_string().unwrap_or_default();
295            bail!("{} returned HTTP {}: {}", url, code, body.trim())
296        }
297        Err(e) => {
298            let local = matches!(target.provider, Provider::Ollama | Provider::Vllm);
299            if local {
300                bail!(
301                    "request to {} failed: {} (is the {} server running?)",
302                    url,
303                    e,
304                    target.provider.name()
305                )
306            }
307            Err(anyhow::anyhow!("request to {} failed: {}", url, e))
308        }
309    }
310}
311
312/// OpenAI with `OPENAI_API_KEY` and optional `OPENAI_BASE_URL` (0.3 behaviour).
313pub fn fetch(prompt: &str, model: &str, max_tokens: u32) -> Result<String> {
314    let t = Target::resolve(Provider::OpenAi, None, Some(model), |k| {
315        std::env::var(k).ok()
316    })?;
317    fetch_target(&t, prompt, max_tokens)
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use std::collections::HashMap;
324
325    fn env(pairs: &[(&str, &str)]) -> impl Fn(&str) -> Option<String> {
326        let m: HashMap<String, String> = pairs
327            .iter()
328            .map(|(k, v)| (k.to_string(), v.to_string()))
329            .collect();
330        move |k| m.get(k).cloned()
331    }
332
333    #[test]
334    fn request_asks_for_logprobs() {
335        let b = request_body("hi", "m", 50);
336        assert_eq!(b["logprobs"], true);
337        assert_eq!(b["top_logprobs"], 3);
338        assert_eq!(b["messages"][0]["content"], "hi");
339        assert_eq!(b["model"], "m");
340        assert_eq!(b["temperature"], 0);
341        assert!(b.get("provider").is_none());
342    }
343
344    #[test]
345    fn provider_names_round_trip() {
346        for p in Provider::ALL {
347            assert_eq!(Provider::from_name(p.name()), Some(p));
348        }
349        assert_eq!(
350            Provider::from_name("OpenRouter"),
351            Some(Provider::OpenRouter)
352        );
353        assert_eq!(Provider::from_name("anthropic"), None);
354    }
355
356    #[test]
357    fn openai_default_is_backwards_compatible() {
358        let t = Target::resolve(
359            Provider::OpenAi,
360            None,
361            None,
362            env(&[("OPENAI_API_KEY", "sk")]),
363        )
364        .unwrap();
365        assert_eq!(t.url(), "https://api.openai.com/v1/chat/completions");
366        assert_eq!(t.model, "gpt-4o-mini");
367        assert_eq!(t.api_key.as_deref(), Some("sk"));
368
369        let t = Target::resolve(
370            Provider::OpenAi,
371            None,
372            None,
373            env(&[
374                ("OPENAI_API_KEY", "sk"),
375                ("OPENAI_BASE_URL", "https://x.test/v1/"),
376            ]),
377        )
378        .unwrap();
379        assert_eq!(t.url(), "https://x.test/v1/chat/completions");
380
381        let e = Target::resolve(Provider::OpenAi, None, None, env(&[])).unwrap_err();
382        assert!(e.to_string().contains("OPENAI_API_KEY"), "{e}");
383    }
384
385    #[test]
386    fn base_url_flag_beats_env_and_preset() {
387        let t = Target::resolve(
388            Provider::OpenAi,
389            Some("http://h:1/v1/"),
390            None,
391            env(&[
392                ("OPENAI_API_KEY", "sk"),
393                ("OPENAI_BASE_URL", "https://x.test/v1"),
394            ]),
395        )
396        .unwrap();
397        assert_eq!(t.url(), "http://h:1/v1/chat/completions");
398        let t = Target::resolve(
399            Provider::Ollama,
400            Some("http://box:11434/v1"),
401            Some("m"),
402            env(&[]),
403        )
404        .unwrap();
405        assert_eq!(t.url(), "http://box:11434/v1/chat/completions");
406    }
407
408    #[test]
409    fn openrouter_request() {
410        let t = Target::resolve(
411            Provider::OpenRouter,
412            None,
413            None,
414            env(&[
415                ("OPENROUTER_API_KEY", "or"),
416                ("OPENAI_BASE_URL", "https://ignored"),
417            ]),
418        )
419        .unwrap();
420        assert_eq!(t.url(), "https://openrouter.ai/api/v1/chat/completions");
421        assert_eq!(t.model, "openai/gpt-4o-mini");
422        assert_eq!(t.api_key.as_deref(), Some("or"));
423        let b = t.body("q", 10);
424        assert_eq!(b["logprobs"], true);
425        assert_eq!(b["top_logprobs"], 3);
426        assert_eq!(b["provider"]["require_parameters"], true);
427        let e = Target::resolve(Provider::OpenRouter, None, None, env(&[])).unwrap_err();
428        assert!(e.to_string().contains("OPENROUTER_API_KEY"), "{e}");
429    }
430
431    #[test]
432    fn together_request_uses_integer_logprobs() {
433        let t = Target::resolve(
434            Provider::Together,
435            None,
436            Some("some/model"),
437            env(&[("TOGETHER_API_KEY", "tg")]),
438        )
439        .unwrap();
440        assert_eq!(t.url(), "https://api.together.ai/v1/chat/completions");
441        let b = t.body("q", 10);
442        assert_eq!(b["logprobs"], 3);
443        assert!(b.get("top_logprobs").is_none());
444        let e = Target::resolve(
445            Provider::Together,
446            None,
447            None,
448            env(&[("TOGETHER_API_KEY", "tg")]),
449        )
450        .unwrap_err();
451        assert!(e.to_string().contains("--model"), "{e}");
452    }
453
454    #[test]
455    fn local_providers_need_no_key() {
456        let t =
457            Target::resolve(Provider::Ollama, None, Some("qwen2.5-coder:14b"), env(&[])).unwrap();
458        assert_eq!(t.url(), "http://localhost:11434/v1/chat/completions");
459        assert_eq!(t.api_key, None);
460        assert_eq!(t.body("q", 5)["logprobs"], true);
461
462        let t = Target::resolve(Provider::Vllm, None, Some("m"), env(&[])).unwrap();
463        assert_eq!(t.url(), "http://localhost:8000/v1/chat/completions");
464        assert_eq!(t.api_key, None);
465        let t = Target::resolve(
466            Provider::Vllm,
467            None,
468            Some("m"),
469            env(&[("VLLM_API_KEY", "v")]),
470        )
471        .unwrap();
472        assert_eq!(t.api_key.as_deref(), Some("v"));
473
474        let e = Target::resolve(Provider::Ollama, None, None, env(&[])).unwrap_err();
475        assert!(e.to_string().contains("ollama list"), "{e}");
476    }
477
478    fn target(p: Provider) -> Target {
479        Target {
480            provider: p,
481            base_url: p.default_base_url().into(),
482            model: "m".into(),
483            api_key: None,
484        }
485    }
486
487    #[test]
488    fn missing_logprobs_is_explained() {
489        let raw = r#"{"choices":[{"message":{"content":"Rembrandt."},"logprobs":null}]}"#;
490        let e = check_has_logprobs(raw, &target(Provider::Ollama))
491            .unwrap_err()
492            .to_string();
493        assert!(e.contains("no token logprobs"), "{e}");
494        assert!(e.contains("Rembrandt."), "{e}");
495        let raw = r#"{"choices":[{"message":{"content":"x"}}]}"#;
496        assert!(check_has_logprobs(raw, &target(Provider::OpenRouter)).is_err());
497        let raw = r#"{"choices":[{"logprobs":{"content":[]}}]}"#;
498        assert!(check_has_logprobs(raw, &target(Provider::OpenAi)).is_err());
499    }
500
501    #[test]
502    fn present_logprobs_pass_the_check() {
503        let raw = r#"{"choices":[{"logprobs":{"content":[{"token":"a","logprob":-0.1}]}}]}"#;
504        assert!(check_has_logprobs(raw, &target(Provider::OpenAi)).is_ok());
505        let raw = r#"{"choices":[{"logprobs":{"tokens":["a"],"token_logprobs":[-0.1]}}]}"#;
506        assert!(check_has_logprobs(raw, &target(Provider::Together)).is_ok());
507    }
508}
llm_token_visualizer/live.rs

llm_token_visualizer/
live.rs
