llm_token_visualizer/
lib.rs

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

llm_token_visualizer/
lib.rs
