langchain_rust/language_models/
mod.rs

1use std::{collections::HashMap, pin::Pin};
2
3use futures::Stream;
4use serde::{Deserialize, Serialize};
5
6use crate::schemas::ToolCall;
7
8pub mod llm;
9pub mod options;
10
11mod error;
12pub use error::*;
13
14/// The result of an LLM call.
15///
16/// When the model asks for tools, `tool_calls` holds them (same type for every provider)
17/// and `generation` holds the same calls as an OpenAI-style JSON array, so code that
18/// reads `generation` keeps working. Use [`GenerateResult::with_tool_calls`] to build one.
19#[derive(Debug, Serialize, Deserialize, Clone, Default)]
20pub struct GenerateResult {
21    pub tokens: Option<TokenUsage>,
22    pub generation: String,
23    #[serde(default, skip_serializing_if = "Vec::is_empty")]
24    pub tool_calls: Vec<ToolCall>,
25}
26
27impl GenerateResult {
28    /// Sets `tool_calls` and mirrors them into `generation` as JSON. Empty `calls`
29    /// leave the result untouched.
30    pub fn with_tool_calls(mut self, calls: Vec<ToolCall>) -> Self {
31        if !calls.is_empty() {
32            self.generation = ToolCall::list_to_string(&calls);
33            self.tool_calls = calls;
34        }
35        self
36    }
37
38    pub fn has_tool_calls(&self) -> bool {
39        !self.tool_calls.is_empty()
40    }
41
42    pub fn to_hashmap(&self) -> HashMap<String, String> {
43        let mut map = HashMap::new();
44
45        // Insert the 'generation' field into the hashmap
46        map.insert("generation".to_string(), self.generation.clone());
47
48        // Check if 'tokens' is Some and insert its fields into the hashmap
49        if let Some(ref tokens) = self.tokens {
50            map.insert(
51                "prompt_tokens".to_string(),
52                tokens.prompt_tokens.to_string(),
53            );
54            map.insert(
55                "completion_tokens".to_string(),
56                tokens.completion_tokens.to_string(),
57            );
58            map.insert("total_tokens".to_string(), tokens.total_tokens.to_string());
59        }
60
61        map
62    }
63}
64
65/// An event from [`llm::LLM::generate_stream`] (and `Chain::call_stream`).
66///
67/// Zero or more `Text` events arrive as the model writes, then exactly one `Done` with
68/// the complete result: full text, typed `tool_calls` and token usage when the
69/// provider reports it.
70#[derive(Debug, Clone)]
71#[non_exhaustive]
72pub enum LLMEvent {
73    /// A chunk of text, as the model produces it.
74    Text(String),
75    /// The complete response. Always the last event.
76    Done(GenerateResult),
77}
78
79/// Stream returned by [`llm::LLM::generate_stream`].
80pub type LLMEventStream<'a> = Pin<Box<dyn Stream<Item = Result<LLMEvent, LLMError>> + Send + 'a>>;
81
82/// Boxes a stream of LLM events; also pins down the error type for `try_stream!` blocks.
83pub(crate) fn box_llm_events<'a, S>(stream: S) -> LLMEventStream<'a>
84where
85    S: Stream<Item = Result<LLMEvent, LLMError>> + Send + 'a,
86{
87    Box::pin(stream)
88}
89
90/// The events for a result that was not streamed: its text in one piece (unless the
91/// model asked for tools, in which case `generation` is just the JSON mirror), then `Done`.
92pub(crate) fn result_events(result: GenerateResult) -> Vec<LLMEvent> {
93    let mut events = Vec::with_capacity(2);
94    if !result.has_tool_calls() && !result.generation.is_empty() {
95        events.push(LLMEvent::Text(result.generation.clone()));
96    }
97    events.push(LLMEvent::Done(result));
98    events
99}
100
101#[derive(Debug, Serialize, Deserialize, Clone, Default, PartialEq, Eq)]
102pub struct TokenUsage {
103    pub prompt_tokens: u32,
104    pub completion_tokens: u32,
105    pub total_tokens: u32,
106}
107
108impl TokenUsage {
109    pub fn sum(&self, other: &TokenUsage) -> TokenUsage {
110        TokenUsage {
111            prompt_tokens: self.prompt_tokens + other.prompt_tokens,
112            completion_tokens: self.completion_tokens + other.completion_tokens,
113            total_tokens: self.total_tokens + other.total_tokens,
114        }
115    }
116
117    pub fn add(&mut self, other: &TokenUsage) {
118        self.prompt_tokens += other.prompt_tokens;
119        self.completion_tokens += other.completion_tokens;
120        self.total_tokens += other.total_tokens;
121    }
122}
123
124impl TokenUsage {
125    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
126        Self {
127            prompt_tokens,
128            completion_tokens,
129            total_tokens: prompt_tokens + completion_tokens,
130        }
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    #[test]
139    fn token_usage_new_sum_add() {
140        let a = TokenUsage::new(3, 4);
141        assert_eq!(a.total_tokens, 7);
142        let b = TokenUsage::new(1, 2);
143        let s = a.sum(&b);
144        assert_eq!(
145            (s.prompt_tokens, s.completion_tokens, s.total_tokens),
146            (4, 6, 10)
147        );
148        let mut c = a.clone();
149        c.add(&b);
150        assert_eq!(c.total_tokens, 10);
151        // sum does not mutate
152        assert_eq!(a.total_tokens, 7);
153    }
154
155    #[test]
156    fn generate_result_to_hashmap() {
157        let without = GenerateResult {
158            generation: "hi".into(),
159            ..Default::default()
160        };
161        let m = without.to_hashmap();
162        assert_eq!(m.len(), 1);
163        assert_eq!(m["generation"], "hi");
164
165        let with = GenerateResult {
166            generation: "hi".into(),
167            tokens: Some(TokenUsage::new(1, 2)),
168            ..Default::default()
169        };
170        let m = with.to_hashmap();
171        assert_eq!(m["prompt_tokens"], "1");
172        assert_eq!(m["completion_tokens"], "2");
173        assert_eq!(m["total_tokens"], "3");
174    }
175
176    #[test]
177    fn with_tool_calls_mirrors_into_generation() {
178        let calls = vec![ToolCall::new("c1", "weather", r#"{"city":"Lima"}"#)];
179        let r = GenerateResult {
180            generation: "let me check".into(),
181            ..Default::default()
182        }
183        .with_tool_calls(calls.clone());
184        assert!(r.has_tool_calls());
185        assert_eq!(r.tool_calls, calls);
186        assert_eq!(ToolCall::parse_list(&r.generation).unwrap(), calls);
187
188        let plain = GenerateResult {
189            generation: "hi".into(),
190            ..Default::default()
191        }
192        .with_tool_calls(vec![]);
193        assert!(!plain.has_tool_calls());
194        assert_eq!(plain.generation, "hi");
195    }
196
197    #[test]
198    fn generate_result_serde_is_backward_compatible() {
199        // Old payloads without `tool_calls` still deserialize.
200        let old: GenerateResult =
201            serde_json::from_str(r#"{"tokens":null,"generation":"hi"}"#).unwrap();
202        assert!(old.tool_calls.is_empty());
203        // No tool calls -> field omitted.
204        let v = serde_json::to_value(&old).unwrap();
205        assert!(v.get("tool_calls").is_none());
206
207        let with = GenerateResult::default().with_tool_calls(vec![ToolCall::new("a", "f", "{}")]);
208        let v = serde_json::to_value(&with).unwrap();
209        assert_eq!(v["tool_calls"][0]["function"]["name"], "f");
210        let back: GenerateResult = serde_json::from_value(v).unwrap();
211        assert_eq!(back.tool_calls, with.tool_calls);
212    }
213
214    #[test]
215    fn result_events_skip_text_for_tool_calls_and_empty_output() {
216        let text = |events: &[LLMEvent]| {
217            events
218                .iter()
219                .filter_map(|e| match e {
220                    LLMEvent::Text(t) => Some(t.clone()),
221                    _ => None,
222                })
223                .collect::<Vec<_>>()
224        };
225        let plain = result_events(GenerateResult {
226            generation: "hi".into(),
227            ..Default::default()
228        });
229        assert_eq!(text(&plain), vec!["hi"]);
230        assert!(matches!(plain.last(), Some(LLMEvent::Done(r)) if r.generation == "hi"));
231
232        let tools = result_events(
233            GenerateResult::default().with_tool_calls(vec![ToolCall::new("a", "f", "{}")]),
234        );
235        assert!(text(&tools).is_empty(), "tool-call JSON is not answer text");
236        assert!(matches!(tools.as_slice(), [LLMEvent::Done(r)] if r.has_tool_calls()));
237
238        assert!(matches!(
239            result_events(GenerateResult::default()).as_slice(),
240            [LLMEvent::Done(_)]
241        ));
242    }
243
244    /// An LLM that only implements `generate` still streams through the default method.
245    #[derive(Clone)]
246    struct OnlyGenerate(Result<&'static str, &'static str>);
247
248    #[async_trait::async_trait]
249    impl llm::LLM for OnlyGenerate {
250        async fn generate(
251            &self,
252            _: &[crate::schemas::Message],
253        ) -> Result<GenerateResult, LLMError> {
254            match self.0 {
255                Ok(s) => Ok(GenerateResult {
256                    generation: s.into(),
257                    ..Default::default()
258                }),
259                Err(e) => Err(LLMError::OtherError(e.into())),
260            }
261        }
262        async fn stream(
263            &self,
264            _: &[crate::schemas::Message],
265        ) -> Result<
266            Pin<Box<dyn Stream<Item = Result<crate::schemas::StreamData, LLMError>> + Send>>,
267            LLMError,
268        > {
269            unimplemented!()
270        }
271    }
272
273    #[tokio::test]
274    async fn default_generate_stream_wraps_generate() {
275        use futures::StreamExt;
276        use llm::LLM;
277
278        let events: Vec<_> = OnlyGenerate(Ok("hello"))
279            .generate_stream(&[])
280            .map(Result::unwrap)
281            .collect()
282            .await;
283        assert!(matches!(&events[0], LLMEvent::Text(t) if t == "hello"));
284        assert!(matches!(&events[1], LLMEvent::Done(r) if r.generation == "hello"));
285        assert_eq!(events.len(), 2);
286
287        let events: Vec<_> = OnlyGenerate(Err("down"))
288            .generate_stream(&[])
289            .collect()
290            .await;
291        assert_eq!(events.len(), 1);
292        assert!(events[0].as_ref().unwrap_err().to_string().contains("down"));
293    }
294}
langchain_rust/language_models/mod.rs

langchain_rust/language_models/
mod.rs
