llm_token_visualizer/
interop.rs

1//! Adapters for other crates, each behind its own optional feature.
2//!
3//! | Feature | Adds |
4//! |---|---|
5//! | `async-openai` | [`async_openai`]: tokens from an `async_openai` chat completion response |
6//!
7//! Any other client works through JSON: [`crate::detect::parse_logprobs`]
8//! reads OpenAI-style, legacy completions-style and Gemini-style responses.
9
10/// Adapter for [async-openai](https://crates.io/crates/async-openai) (types
11/// only, no HTTP client).
12#[cfg(feature = "async-openai")]
13#[cfg_attr(docsrs, doc(cfg(feature = "async-openai")))]
14pub mod async_openai {
15    use ::async_openai::types::chat::{ChatCompletionTokenLogprob, CreateChatCompletionResponse};
16    use anyhow::{bail, Result};
17
18    use crate::detect::{Alternative, LogprobToken};
19
20    impl From<&ChatCompletionTokenLogprob> for LogprobToken {
21        fn from(t: &ChatCompletionTokenLogprob) -> Self {
22            let mut top: Vec<Alternative> = t
23                .top_logprobs
24                .iter()
25                .map(|a| Alternative {
26                    token: a.token.clone(),
27                    logprob: f64::from(a.logprob),
28                })
29                .collect();
30            top.sort_by(|a, b| b.logprob.total_cmp(&a.logprob));
31            LogprobToken {
32                token: t.token.clone(),
33                logprob: f64::from(t.logprob),
34                top_logprobs: top,
35            }
36        }
37    }
38
39    /// The tokens of the first choice, ready for [`crate::detect::detect`].
40    /// Special tokens such as `<|eot_id|>` are dropped, as in
41    /// [`crate::detect::parse_logprobs`].
42    ///
43    /// # Errors
44    ///
45    /// When the response has no choices or no logprobs (request the
46    /// completion with `logprobs(true)` and `top_logprobs(3)`).
47    pub fn tokens_from_response(resp: &CreateChatCompletionResponse) -> Result<Vec<LogprobToken>> {
48        let Some(choice) = resp.choices.first() else {
49            bail!("the response has no choices");
50        };
51        let Some(content) = choice.logprobs.as_ref().and_then(|l| l.content.as_ref()) else {
52            bail!("no token logprobs; request the completion with logprobs(true)");
53        };
54        let tokens: Vec<LogprobToken> = content
55            .iter()
56            .map(LogprobToken::from)
57            .filter(|t| !(t.token.starts_with("<|") && t.token.ends_with("|>")))
58            .collect();
59        if tokens.is_empty() {
60            bail!("logprobs contain no tokens");
61        }
62        Ok(tokens)
63    }
64
65    #[cfg(test)]
66    mod tests {
67        use super::*;
68
69        #[test]
70        fn same_tokens_as_the_json_parser() {
71            let json = include_str!("../examples/logprobs/cuyp.json");
72            let resp: CreateChatCompletionResponse = serde_json::from_str(json).unwrap();
73            let a = tokens_from_response(&resp).unwrap();
74            let b = crate::detect::parse_logprobs(json).unwrap();
75            assert_eq!(a.len(), b.len());
76            for (x, y) in a.iter().zip(&b) {
77                assert_eq!(x.token, y.token);
78                // async-openai stores logprobs as f32.
79                assert!((x.logprob - y.logprob).abs() < 1e-6);
80            }
81            let r = crate::detect::detect(&a, 0.6);
82            assert!(r.spans.iter().any(|s| s.text == " Dordrecht"));
83        }
84    }
85}
llm_token_visualizer/interop.rs

llm_token_visualizer/
interop.rs
