langchain_rust/language_models/
mod.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5pub mod llm;
6pub mod options;
7
8mod error;
9pub use error::*;
10
11//TODO: check if its this should have a data:serde::Value to save all other things, like OpenAI
12//function responses
13#[derive(Debug, Serialize, Deserialize, Clone, Default)]
14pub struct GenerateResult {
15    pub tokens: Option<TokenUsage>,
16    pub generation: String,
17}
18
19impl GenerateResult {
20    pub fn to_hashmap(&self) -> HashMap<String, String> {
21        let mut map = HashMap::new();
22
23        // Insert the 'generation' field into the hashmap
24        map.insert("generation".to_string(), self.generation.clone());
25
26        // Check if 'tokens' is Some and insert its fields into the hashmap
27        if let Some(ref tokens) = self.tokens {
28            map.insert(
29                "prompt_tokens".to_string(),
30                tokens.prompt_tokens.to_string(),
31            );
32            map.insert(
33                "completion_tokens".to_string(),
34                tokens.completion_tokens.to_string(),
35            );
36            map.insert("total_tokens".to_string(), tokens.total_tokens.to_string());
37        }
38
39        map
40    }
41}
42
43#[derive(Debug, Serialize, Deserialize, Clone, Default)]
44pub struct TokenUsage {
45    pub prompt_tokens: u32,
46    pub completion_tokens: u32,
47    pub total_tokens: u32,
48}
49
50impl TokenUsage {
51    pub fn sum(&self, other: &TokenUsage) -> TokenUsage {
52        TokenUsage {
53            prompt_tokens: self.prompt_tokens + other.prompt_tokens,
54            completion_tokens: self.completion_tokens + other.completion_tokens,
55            total_tokens: self.total_tokens + other.total_tokens,
56        }
57    }
58
59    pub fn add(&mut self, other: &TokenUsage) {
60        self.prompt_tokens += other.prompt_tokens;
61        self.completion_tokens += other.completion_tokens;
62        self.total_tokens += other.total_tokens;
63    }
64}
65
66impl TokenUsage {
67    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
68        Self {
69            prompt_tokens,
70            completion_tokens,
71            total_tokens: prompt_tokens + completion_tokens,
72        }
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn token_usage_new_sum_add() {
82        let a = TokenUsage::new(3, 4);
83        assert_eq!(a.total_tokens, 7);
84        let b = TokenUsage::new(1, 2);
85        let s = a.sum(&b);
86        assert_eq!(
87            (s.prompt_tokens, s.completion_tokens, s.total_tokens),
88            (4, 6, 10)
89        );
90        let mut c = a.clone();
91        c.add(&b);
92        assert_eq!(c.total_tokens, 10);
93        // sum does not mutate
94        assert_eq!(a.total_tokens, 7);
95    }
96
97    #[test]
98    fn generate_result_to_hashmap() {
99        let without = GenerateResult {
100            generation: "hi".into(),
101            tokens: None,
102        };
103        let m = without.to_hashmap();
104        assert_eq!(m.len(), 1);
105        assert_eq!(m["generation"], "hi");
106
107        let with = GenerateResult {
108            generation: "hi".into(),
109            tokens: Some(TokenUsage::new(1, 2)),
110        };
111        let m = with.to_hashmap();
112        assert_eq!(m["prompt_tokens"], "1");
113        assert_eq!(m["completion_tokens"], "2");
114        assert_eq!(m["total_tokens"], "3");
115    }
116}
langchain_rust/language_models/mod.rs

langchain_rust/language_models/
mod.rs
