lc_providers/
wrapper.rs

1// lc-providers/src/wrapper.rs
2//! Wrapper that normalizes any `BaseChatModel` to use `ProviderError`.
3//!
4//! This enables `Arc<dyn BaseChatModel<Error = ProviderError>>` —
5//! a single trait-object type that works with every provider.
6
7use crate::error::ProviderError;
8use async_trait::async_trait;
9use futures_util::{Stream, StreamExt};
10use lc_core::language_models::{BaseChatModel, BaseLanguageModel, LLMResult, StreamChunk};
11use lc_core::runnables::Runnable;
12use lc_core::tools::ToolDefinition;
13use lc_core::RunnableConfig;
14use lc_schema::Message;
15use std::pin::Pin;
16use std::sync::Arc;
17
18/// Wrapper that converts any `BaseChatModel`'s error into `ProviderError`.
19///
20/// This allows heterogeneous LLM providers to be stored behind a single
21/// `Arc<dyn BaseChatModel<Error = ProviderError>>` trait object.
22pub struct ChatModelWrapper<L> {
23    inner: L,
24}
25
26impl<L> ChatModelWrapper<L> {
27    /// Create a new wrapper around the given LLM.
28    pub fn new(llm: L) -> Self {
29        Self { inner: llm }
30    }
31}
32
33#[async_trait]
34impl<L> Runnable<Vec<Message>, LLMResult> for ChatModelWrapper<L>
35where
36    L: Runnable<Vec<Message>, LLMResult> + Send + Sync,
37    L::Error: Into<ProviderError>,
38{
39    type Error = ProviderError;
40
41    async fn invoke(
42        &self,
43        input: Vec<Message>,
44        config: Option<RunnableConfig>,
45    ) -> Result<LLMResult, ProviderError> {
46        self.inner.invoke(input, config).await.map_err(Into::into)
47    }
48
49    async fn batch(
50        &self,
51        inputs: Vec<Vec<Message>>,
52        config: Option<RunnableConfig>,
53    ) -> Result<Vec<LLMResult>, ProviderError> {
54        self.inner.batch(inputs, config).await.map_err(Into::into)
55    }
56
57    async fn stream(
58        &self,
59        input: Vec<Message>,
60        config: Option<RunnableConfig>,
61    ) -> Result<Pin<Box<dyn Stream<Item = Result<LLMResult, ProviderError>> + Send>>, ProviderError>
62    {
63        let stream = self.inner.stream(input, config).await.map_err(Into::into)?;
64        Ok(Box::pin(stream.map(|r| r.map_err(Into::into))))
65    }
66}
67
68#[async_trait]
69impl<L> BaseLanguageModel<Vec<Message>, LLMResult> for ChatModelWrapper<L>
70where
71    L: BaseLanguageModel<Vec<Message>, LLMResult> + Send + Sync,
72    L::Error: Into<ProviderError>,
73{
74    fn model_name(&self) -> &str {
75        self.inner.model_name()
76    }
77
78    fn get_num_tokens(&self, text: &str) -> usize {
79        self.inner.get_num_tokens(text)
80    }
81
82    fn temperature(&self) -> Option<f32> {
83        self.inner.temperature()
84    }
85
86    fn max_tokens(&self) -> Option<usize> {
87        self.inner.max_tokens()
88    }
89
90    fn with_temperature(self, temp: f32) -> Self
91    where
92        Self: Sized,
93    {
94        Self {
95            inner: self.inner.with_temperature(temp),
96        }
97    }
98
99    fn with_max_tokens(self, max: usize) -> Self
100    where
101        Self: Sized,
102    {
103        Self {
104            inner: self.inner.with_max_tokens(max),
105        }
106    }
107}
108
109#[async_trait]
110impl<L> BaseChatModel for ChatModelWrapper<L>
111where
112    L: BaseChatModel + Send + Sync,
113    L::Error: Into<ProviderError>,
114{
115    async fn chat(
116        &self,
117        messages: Vec<Message>,
118        config: Option<RunnableConfig>,
119    ) -> Result<LLMResult, ProviderError> {
120        self.inner.chat(messages, config).await.map_err(Into::into)
121    }
122
123    async fn stream_chat(
124        &self,
125        messages: Vec<Message>,
126        config: Option<RunnableConfig>,
127    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, ProviderError>> + Send>>, ProviderError>
128    {
129        let stream = self
130            .inner
131            .stream_chat(messages, config)
132            .await
133            .map_err(Into::into)?;
134        Ok(Box::pin(stream.map(|r| r.map_err(Into::into))))
135    }
136
137    fn bind_tools(
138        &self,
139        tools: Vec<ToolDefinition>,
140    ) -> Option<Box<dyn BaseChatModel<Error = ProviderError> + Send + Sync>> {
141        // Delegate to the inner model, then re-normalize the bound model's
142        // error to `ProviderError` so it fits the unified trait-object type
143        // (Q1). `bind_tools` returns a model with the same error type `L::Error`;
144        // `BoundModel` maps it via `Into<ProviderError>`.
145        self.inner.bind_tools(tools).map(|bound| {
146            Box::new(BoundModel { inner: bound })
147                as Box<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
148        })
149    }
150}
151
152/// Adapter normalizing a boxed `BaseChatModel`'s error to `ProviderError`.
153///
154/// `BaseChatModel::bind_tools` returns `Box<dyn BaseChatModel<Error = Self::Error>>`,
155/// which cannot itself implement `BaseChatModel` (the orphan rule forbids
156/// implementing a foreign trait for `Box<dyn ...>`). This local adapter re-errors
157/// the box so it can be stored behind the unified
158/// `Box<dyn BaseChatModel<Error = ProviderError>>` (providers Q1).
159struct BoundModel<E> {
160    inner: Box<dyn BaseChatModel<Error = E> + Send + Sync>,
161}
162
163#[async_trait]
164impl<E> Runnable<Vec<Message>, LLMResult> for BoundModel<E>
165where
166    E: Into<ProviderError> + std::error::Error + Send + Sync + 'static,
167{
168    type Error = ProviderError;
169
170    async fn invoke(
171        &self,
172        input: Vec<Message>,
173        config: Option<RunnableConfig>,
174    ) -> Result<LLMResult, ProviderError> {
175        self.inner.invoke(input, config).await.map_err(Into::into)
176    }
177
178    async fn batch(
179        &self,
180        inputs: Vec<Vec<Message>>,
181        config: Option<RunnableConfig>,
182    ) -> Result<Vec<LLMResult>, ProviderError> {
183        self.inner.batch(inputs, config).await.map_err(Into::into)
184    }
185
186    async fn stream(
187        &self,
188        input: Vec<Message>,
189        config: Option<RunnableConfig>,
190    ) -> Result<Pin<Box<dyn Stream<Item = Result<LLMResult, ProviderError>> + Send>>, ProviderError>
191    {
192        let stream = self.inner.stream(input, config).await.map_err(Into::into)?;
193        Ok(Box::pin(stream.map(|r| r.map_err(Into::into))))
194    }
195}
196
197#[async_trait]
198impl<E> BaseLanguageModel<Vec<Message>, LLMResult> for BoundModel<E>
199where
200    E: Into<ProviderError> + std::error::Error + Send + Sync + 'static,
201{
202    fn model_name(&self) -> &str {
203        self.inner.model_name()
204    }
205
206    fn get_num_tokens(&self, text: &str) -> usize {
207        self.inner.get_num_tokens(text)
208    }
209
210    fn temperature(&self) -> Option<f32> {
211        self.inner.temperature()
212    }
213
214    fn max_tokens(&self) -> Option<usize> {
215        self.inner.max_tokens()
216    }
217
218    // `with_temperature`/`with_max_tokens` are `Self: Sized` and unreachable on
219    // a trait object, so returning `self` is honest for this private adapter;
220    // sampling overrides flow through `RunnableConfig` (providers Q2).
221    fn with_temperature(self, _temp: f32) -> Self
222    where
223        Self: Sized,
224    {
225        self
226    }
227
228    fn with_max_tokens(self, _max: usize) -> Self
229    where
230        Self: Sized,
231    {
232        self
233    }
234}
235
236#[async_trait]
237impl<E> BaseChatModel for BoundModel<E>
238where
239    E: Into<ProviderError> + std::error::Error + Send + Sync + 'static,
240{
241    async fn chat(
242        &self,
243        messages: Vec<Message>,
244        config: Option<RunnableConfig>,
245    ) -> Result<LLMResult, ProviderError> {
246        self.inner.chat(messages, config).await.map_err(Into::into)
247    }
248
249    async fn stream_chat(
250        &self,
251        messages: Vec<Message>,
252        config: Option<RunnableConfig>,
253    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, ProviderError>> + Send>>, ProviderError>
254    {
255        let stream = self
256            .inner
257            .stream_chat(messages, config)
258            .await
259            .map_err(Into::into)?;
260        Ok(Box::pin(stream.map(|r| r.map_err(Into::into))))
261    }
262}
263
264/// Wrap any `BaseChatModel` into an `Arc<dyn BaseChatModel<Error = ProviderError>>`.
265///
266/// This is the primary way to create a uniform trait object from any provider.
267///
268/// # Example
269///
270/// ```ignore
271/// use lc_providers::wrap_chat_model;
272///
273/// let openai = OpenAIChat::new(config);
274/// let llm: Arc<dyn BaseChatModel<Error = ProviderError>> = wrap_chat_model(openai);
275/// ```
276pub fn wrap_chat_model<L>(llm: L) -> Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
277where
278    L: BaseChatModel + Send + Sync + 'static,
279    L::Error: Into<ProviderError>,
280{
281    Arc::new(ChatModelWrapper::new(llm))
282}
lc_providers/wrapper.rs

lc_providers/
wrapper.rs
