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

llm_token_visualizer/
live.rs
