langchain_rust/language_models/
mod.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::schemas::ToolCall;
6
7pub mod llm;
8pub mod options;
9
10mod error;
11pub use error::*;
12
13/// The result of an LLM call.
14///
15/// When the model asks for tools, `tool_calls` holds them (same type for every provider)
16/// and `generation` holds the same calls as an OpenAI-style JSON array, so code that
17/// reads `generation` keeps working. Use [`GenerateResult::with_tool_calls`] to build one.
18#[derive(Debug, Serialize, Deserialize, Clone, Default)]
19pub struct GenerateResult {
20    pub tokens: Option<TokenUsage>,
21    pub generation: String,
22    #[serde(default, skip_serializing_if = "Vec::is_empty")]
23    pub tool_calls: Vec<ToolCall>,
24}
25
26impl GenerateResult {
27    /// Sets `tool_calls` and mirrors them into `generation` as JSON. Empty `calls`
28    /// leave the result untouched.
29    pub fn with_tool_calls(mut self, calls: Vec<ToolCall>) -> Self {
30        if !calls.is_empty() {
31            self.generation = ToolCall::list_to_string(&calls);
32            self.tool_calls = calls;
33        }
34        self
35    }
36
37    pub fn has_tool_calls(&self) -> bool {
38        !self.tool_calls.is_empty()
39    }
40
41    pub fn to_hashmap(&self) -> HashMap<String, String> {
42        let mut map = HashMap::new();
43
44        // Insert the 'generation' field into the hashmap
45        map.insert("generation".to_string(), self.generation.clone());
46
47        // Check if 'tokens' is Some and insert its fields into the hashmap
48        if let Some(ref tokens) = self.tokens {
49            map.insert(
50                "prompt_tokens".to_string(),
51                tokens.prompt_tokens.to_string(),
52            );
53            map.insert(
54                "completion_tokens".to_string(),
55                tokens.completion_tokens.to_string(),
56            );
57            map.insert("total_tokens".to_string(), tokens.total_tokens.to_string());
58        }
59
60        map
61    }
62}
63
64#[derive(Debug, Serialize, Deserialize, Clone, Default)]
65pub struct TokenUsage {
66    pub prompt_tokens: u32,
67    pub completion_tokens: u32,
68    pub total_tokens: u32,
69}
70
71impl TokenUsage {
72    pub fn sum(&self, other: &TokenUsage) -> TokenUsage {
73        TokenUsage {
74            prompt_tokens: self.prompt_tokens + other.prompt_tokens,
75            completion_tokens: self.completion_tokens + other.completion_tokens,
76            total_tokens: self.total_tokens + other.total_tokens,
77        }
78    }
79
80    pub fn add(&mut self, other: &TokenUsage) {
81        self.prompt_tokens += other.prompt_tokens;
82        self.completion_tokens += other.completion_tokens;
83        self.total_tokens += other.total_tokens;
84    }
85}
86
87impl TokenUsage {
88    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
89        Self {
90            prompt_tokens,
91            completion_tokens,
92            total_tokens: prompt_tokens + completion_tokens,
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn token_usage_new_sum_add() {
103        let a = TokenUsage::new(3, 4);
104        assert_eq!(a.total_tokens, 7);
105        let b = TokenUsage::new(1, 2);
106        let s = a.sum(&b);
107        assert_eq!(
108            (s.prompt_tokens, s.completion_tokens, s.total_tokens),
109            (4, 6, 10)
110        );
111        let mut c = a.clone();
112        c.add(&b);
113        assert_eq!(c.total_tokens, 10);
114        // sum does not mutate
115        assert_eq!(a.total_tokens, 7);
116    }
117
118    #[test]
119    fn generate_result_to_hashmap() {
120        let without = GenerateResult {
121            generation: "hi".into(),
122            ..Default::default()
123        };
124        let m = without.to_hashmap();
125        assert_eq!(m.len(), 1);
126        assert_eq!(m["generation"], "hi");
127
128        let with = GenerateResult {
129            generation: "hi".into(),
130            tokens: Some(TokenUsage::new(1, 2)),
131            ..Default::default()
132        };
133        let m = with.to_hashmap();
134        assert_eq!(m["prompt_tokens"], "1");
135        assert_eq!(m["completion_tokens"], "2");
136        assert_eq!(m["total_tokens"], "3");
137    }
138
139    #[test]
140    fn with_tool_calls_mirrors_into_generation() {
141        let calls = vec![ToolCall::new("c1", "weather", r#"{"city":"Lima"}"#)];
142        let r = GenerateResult {
143            generation: "let me check".into(),
144            ..Default::default()
145        }
146        .with_tool_calls(calls.clone());
147        assert!(r.has_tool_calls());
148        assert_eq!(r.tool_calls, calls);
149        assert_eq!(ToolCall::parse_list(&r.generation).unwrap(), calls);
150
151        let plain = GenerateResult {
152            generation: "hi".into(),
153            ..Default::default()
154        }
155        .with_tool_calls(vec![]);
156        assert!(!plain.has_tool_calls());
157        assert_eq!(plain.generation, "hi");
158    }
159
160    #[test]
161    fn generate_result_serde_is_backward_compatible() {
162        // Old payloads without `tool_calls` still deserialize.
163        let old: GenerateResult =
164            serde_json::from_str(r#"{"tokens":null,"generation":"hi"}"#).unwrap();
165        assert!(old.tool_calls.is_empty());
166        // No tool calls -> field omitted.
167        let v = serde_json::to_value(&old).unwrap();
168        assert!(v.get("tool_calls").is_none());
169
170        let with = GenerateResult::default().with_tool_calls(vec![ToolCall::new("a", "f", "{}")]);
171        let v = serde_json::to_value(&with).unwrap();
172        assert_eq!(v["tool_calls"][0]["function"]["name"], "f");
173        let back: GenerateResult = serde_json::from_value(v).unwrap();
174        assert_eq!(back.tool_calls, with.tool_calls);
175    }
176}
langchain_rust/language_models/mod.rs

langchain_rust/language_models/
mod.rs
