Skip to main content

llm_token_visualizer/
lib.rs

1pub mod data;
2pub mod detect;
3#[cfg(feature = "live")]
4pub mod live;
5pub mod renderer;
6pub mod utils;
7
8pub use data::{
9    ConfidenceLevel, FlagType, TokenAnalysis, TokenFlag, TokenInfo, VisualizationConfig,
10};
11pub use renderer::{HtmlRenderer, MarkdownRenderer, Renderer, TerminalRenderer};
12pub use utils::{create_mock_analysis, detect_issues, simple_tokenize, AnalysisMetrics};
13
14use anyhow::Result;
15
16/// Main visualization function that can be used by other applications
17pub fn visualize_tokens(
18    text: &str,
19    analysis: &TokenAnalysis,
20    format: &str,
21    config: Option<VisualizationConfig>,
22) -> Result<String> {
23    let config = config.unwrap_or_default();
24
25    match format {
26        "terminal" => {
27            let renderer = TerminalRenderer::new();
28            renderer.render(text, analysis, &config)
29        }
30        "html" => {
31            let renderer = HtmlRenderer::new();
32            renderer.render(text, analysis, &config)
33        }
34        "markdown" => {
35            let renderer = MarkdownRenderer::new();
36            renderer.render(text, analysis, &config)
37        }
38        _ => Err(anyhow::anyhow!("Unsupported format: {}", format)),
39    }
40}
41
42/// Quick analysis function for testing - creates mock data and visualizes
43pub fn quick_analyze(text: &str, format: &str) -> Result<String> {
44    let analysis = utils::create_mock_analysis(text);
45    let config = VisualizationConfig::default();
46    visualize_tokens(text, &analysis, format, Some(config))
47}
48
49/// Comprehensive analysis with issue detection
50pub fn analyze_with_issues(analysis: &TokenAnalysis) -> (utils::AnalysisMetrics, Vec<String>) {
51    let metrics = utils::AnalysisMetrics::from_analysis(analysis);
52    let issues = utils::detect_issues(analysis);
53    (metrics, issues)
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn test_visualize_tokens() {
62        let analysis = TokenAnalysis {
63            tokens: vec![
64                TokenInfo {
65                    text: "Hello".to_string(),
66                    confidence: 0.9,
67                },
68                TokenInfo {
69                    text: " world".to_string(),
70                    confidence: 0.8,
71                },
72            ],
73            flags: vec![],
74        };
75
76        let result = visualize_tokens("Hello world", &analysis, "markdown", None);
77        assert!(result.is_ok());
78        assert!(result.unwrap().contains("Hello"));
79    }
80
81    #[test]
82    fn test_quick_analyze() {
83        let result = quick_analyze("This is a test", "markdown");
84        assert!(result.is_ok());
85        assert!(result.unwrap().contains("Token Analysis"));
86    }
87
88    #[test]
89    fn test_analyze_with_issues() {
90        let analysis = TokenAnalysis {
91            tokens: vec![
92                TokenInfo {
93                    text: "good".to_string(),
94                    confidence: 0.9,
95                },
96                TokenInfo {
97                    text: "bad".to_string(),
98                    confidence: 0.1,
99                },
100            ],
101            flags: vec![],
102        };
103
104        let (metrics, issues) = analyze_with_issues(&analysis);
105        assert_eq!(metrics.total_tokens, 2);
106        assert!(!issues.is_empty());
107    }
108}