Skip to main content

llm_token_visualizer/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2//! Flag the words an LLM answer was unsure about, from the token log
3//! probabilities (logprobs) that OpenAI-compatible APIs return.
4//!
5//! Each token's probability is `exp(logprob)`. A word is flagged when any of
6//! its tokens with letters or digits falls below the threshold, neighbouring
7//! flagged words merge into one span, and each span keeps the alternatives the
8//! model was weighing at its weakest token.
9//!
10//! ```
11//! use llm_token_visualizer::detect::{detect, parse_logprobs};
12//!
13//! // Three real tokens from a Llama 3.1 8B answer ("... in Dordrecht").
14//! let json = r#"[
15//!   {"token": " D", "logprob": -0.0153},
16//!   {"token": "ord", "logprob": -0.5620, "top_logprobs": [
17//!     {"token": "ord", "logprob": -0.5620},
18//!     {"token": "üsseldorf", "logprob": -0.9370},
19//!     {"token": "elf", "logprob": -3.5620}]},
20//!   {"token": "recht", "logprob": -0.0004}
21//! ]"#;
22//!
23//! let tokens = parse_logprobs(json)?;
24//! let report = detect(&tokens, 0.6);
25//!
26//! assert_eq!(report.spans.len(), 1);
27//! assert_eq!(report.spans[0].text, " Dordrecht");
28//! assert_eq!(
29//!     report.spans[0].describe(),
30//!     r#"p=0.57 at "ord"; model also considered "üsseldorf" (0.39), "elf" (0.03)"#
31//! );
32//! # Ok::<(), anyhow::Error>(())
33//! ```
34//!
35//! [`report`] renders the same result as a terminal heatmap, a self-contained
36//! HTML page or Markdown. The command-line tool is `llm-token-visualizer`
37//! (`cargo install llm-token-visualizer`).
38
39#![deny(missing_docs)]
40
41/// The README's Rust examples, compiled and run as doctests.
42#[cfg(doctest)]
43#[doc = include_str!("../README.md")]
44pub struct ReadmeDoctests;
45
46/// Input types for the visualizer renderers.
47pub mod data;
48pub mod detect;
49pub mod interop;
50#[cfg(feature = "live")]
51pub mod live;
52/// Terminal, HTML and Markdown renderers for [`TokenAnalysis`] input (visualizer mode).
53pub mod renderer;
54pub mod report;
55/// Helpers for the visualizer: a simple tokenizer, metrics and issue detection.
56pub mod utils;
57
58pub use data::{
59    ConfidenceLevel, FlagType, TokenAnalysis, TokenFlag, TokenInfo, VisualizationConfig,
60};
61pub use renderer::{HtmlRenderer, MarkdownRenderer, Renderer, TerminalRenderer};
62#[allow(deprecated)]
63pub use utils::create_mock_analysis;
64pub use utils::{detect_issues, simple_tokenize, AnalysisMetrics};
65
66use anyhow::Result;
67
68/// Main visualization function that can be used by other applications
69pub fn visualize_tokens(
70    text: &str,
71    analysis: &TokenAnalysis,
72    format: &str,
73    config: Option<VisualizationConfig>,
74) -> Result<String> {
75    let config = config.unwrap_or_default();
76
77    match format {
78        "terminal" => {
79            let renderer = TerminalRenderer::new();
80            renderer.render(text, analysis, &config)
81        }
82        "html" => {
83            let renderer = HtmlRenderer::new();
84            renderer.render(text, analysis, &config)
85        }
86        "markdown" => {
87            let renderer = MarkdownRenderer::new();
88            renderer.render(text, analysis, &config)
89        }
90        _ => Err(anyhow::anyhow!("Unsupported format: {}", format)),
91    }
92}
93
94/// Renders made-up scores (see [`utils::create_mock_analysis`]); the
95/// confidence values are a hash of each word, not anything a model said.
96#[deprecated(
97    since = "0.5.0",
98    note = "renders made-up scores; use detect::parse_logprobs and detect::detect on a real response"
99)]
100#[allow(deprecated)]
101pub fn quick_analyze(text: &str, format: &str) -> Result<String> {
102    let analysis = utils::create_mock_analysis(text);
103    let config = VisualizationConfig::default();
104    visualize_tokens(text, &analysis, format, Some(config))
105}
106
107/// Comprehensive analysis with issue detection
108pub fn analyze_with_issues(analysis: &TokenAnalysis) -> (utils::AnalysisMetrics, Vec<String>) {
109    let metrics = utils::AnalysisMetrics::from_analysis(analysis);
110    let issues = utils::detect_issues(analysis);
111    (metrics, issues)
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_visualize_tokens() {
120        let analysis = TokenAnalysis {
121            tokens: vec![
122                TokenInfo {
123                    text: "Hello".to_string(),
124                    confidence: 0.9,
125                },
126                TokenInfo {
127                    text: " world".to_string(),
128                    confidence: 0.8,
129                },
130            ],
131            flags: vec![],
132        };
133
134        let result = visualize_tokens("Hello world", &analysis, "markdown", None);
135        assert!(result.is_ok());
136        assert!(result.unwrap().contains("Hello"));
137    }
138
139    #[test]
140    #[allow(deprecated)]
141    fn test_quick_analyze() {
142        let result = quick_analyze("This is a test", "markdown");
143        assert!(result.is_ok());
144        assert!(result.unwrap().contains("Token Analysis"));
145    }
146
147    #[test]
148    fn test_analyze_with_issues() {
149        let analysis = TokenAnalysis {
150            tokens: vec![
151                TokenInfo {
152                    text: "good".to_string(),
153                    confidence: 0.9,
154                },
155                TokenInfo {
156                    text: "bad".to_string(),
157                    confidence: 0.1,
158                },
159            ],
160            flags: vec![],
161        };
162
163        let (metrics, issues) = analyze_with_issues(&analysis);
164        assert_eq!(metrics.total_tokens, 2);
165        assert!(!issues.is_empty());
166    }
167}