langchain_rust/language_models/
llm.rs

1use std::pin::Pin;
2
3use async_trait::async_trait;
4use futures::Stream;
5
6use crate::schemas::{Message, StreamData};
7
8use super::{
9    box_llm_events, options::CallOptions, result_events, GenerateResult, LLMError, LLMEventStream,
10};
11
12#[async_trait]
13pub trait LLM: Sync + Send + LLMClone {
14    async fn generate(&self, messages: &[Message]) -> Result<GenerateResult, LLMError>;
15    async fn invoke(&self, prompt: &str) -> Result<String, LLMError> {
16        self.generate(&[Message::new_human_message(prompt)])
17            .await
18            .map(|res| res.generation)
19    }
20    async fn stream(
21        &self,
22        _messages: &[Message],
23    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, LLMError>> + Send>>, LLMError>;
24
25    /// Streams the response as [`LLMEvent`](super::LLMEvent)s: `Text` chunks while the
26    /// model writes, then one `Done` with the complete [`GenerateResult`] (including typed
27    /// `tool_calls`). Errors, including the initial request failing, come as stream items.
28    ///
29    /// The default implementation calls [`generate`](LLM::generate) and emits its result
30    /// as a single chunk, so every LLM supports it; providers override it to stream for real.
31    ///
32    /// ```rust,ignore
33    /// let mut events = llm.generate_stream(&messages);
34    /// while let Some(event) = events.next().await {
35    ///     match event? {
36    ///         LLMEvent::Text(chunk) => print!("{chunk}"),
37    ///         LLMEvent::Done(result) => println!("\n{:?}", result.tokens),
38    ///         _ => {}
39    ///     }
40    /// }
41    /// ```
42    fn generate_stream<'a>(&'a self, messages: &'a [Message]) -> LLMEventStream<'a> {
43        box_llm_events(async_stream::try_stream! {
44            let result = self.generate(messages).await?;
45            for event in result_events(result) {
46                yield event;
47            }
48        })
49    }
50
51    /// This is usefull when you want to create a chain and override
52    /// LLM options
53    fn add_options(&mut self, _options: CallOptions) {
54        // No action taken
55    }
56    //This is usefull when using non chat models
57    fn messages_to_string(&self, messages: &[Message]) -> String {
58        Message::messages_to_string(messages)
59    }
60}
61
62pub trait LLMClone {
63    fn clone_box(&self) -> Box<dyn LLM>;
64}
65
66impl<T> LLMClone for T
67where
68    T: 'static + LLM + Clone,
69{
70    fn clone_box(&self) -> Box<dyn LLM> {
71        Box::new(self.clone())
72    }
73}
74
75impl<L> From<L> for Box<dyn LLM>
76where
77    L: 'static + LLM,
78{
79    fn from(llm: L) -> Self {
80        Box::new(llm)
81    }
82}
langchain_rust/language_models/llm.rs

langchain_rust/language_models/
llm.rs
