polyc_llm/
erased.rs

1//! Type erasure for [`LlmProvider`] so the control plane can hold a single
2//! `Arc<dyn LlmProvider>` regardless of which backend is configured.
3//!
4//! The trait keeps a per-provider associated [`LlmError`] type (each backend
5//! ships its own concrete error). A trait object must fix that associated
6//! type, so this module supplies one uniform error — [`BoxError`] — and an
7//! [`ErasedProvider`] adapter that maps any provider's error into it. The
8//! result is [`DynProvider`], the single trait-object type callers store and
9//! dispatch through.
10//!
11//! Adding a backend then costs one trait impl plus one [`into_dyn`] call at the
12//! wiring boundary — no change to the dispatch site or the planner.
13
14use std::sync::Arc;
15
16use async_trait::async_trait;
17use futures::stream::{BoxStream, StreamExt};
18
19use crate::{
20    Chunk, CompletionRequest, LlmProvider,
21    error::{LlmError, LlmErrorKind},
22};
23
24/// A provider error erased to one concrete type, so backends with differing
25/// associated `Error`s can be stored behind a single trait object.
26///
27/// Transparent wrapper: [`Display`](std::fmt::Display) delegates to the inner
28/// error and [`source`](std::error::Error::source) exposes it, so logs and
29/// error chains read exactly as the un-erased error did. The second field
30/// preserves the original [`LlmErrorKind`] across erasure (the boxed `dyn Error`
31/// alone could not be re-classified).
32#[derive(Debug)]
33pub struct BoxError(
34    Box<dyn std::error::Error + Send + Sync + 'static>,
35    LlmErrorKind,
36    Option<std::time::Duration>,
37);
38
39impl BoxError {
40    /// Erase any [`LlmError`] into a `BoxError`, capturing its
41    /// [`kind`](LlmError::kind) and [`retry_after`](LlmError::retry_after) so
42    /// the classification and rate-limit hint survive erasure.
43    #[must_use]
44    pub fn new<E: LlmError>(err: E) -> Self {
45        let kind = err.kind();
46        let retry_after = err.retry_after();
47        Self(Box::new(err), kind, retry_after)
48    }
49}
50
51impl LlmError for BoxError {
52    fn kind(&self) -> LlmErrorKind {
53        self.1
54    }
55
56    fn retry_after(&self) -> Option<std::time::Duration> {
57        self.2
58    }
59}
60
61impl std::fmt::Display for BoxError {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        std::fmt::Display::fmt(&self.0, f)
64    }
65}
66
67impl std::error::Error for BoxError {
68    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
69        Some(&*self.0)
70    }
71}
72
73/// Adapter that wraps a concrete [`LlmProvider`] and erases its associated
74/// error to [`BoxError`], so the wrapped value coerces to [`DynProvider`].
75///
76/// The completion stream is mapped lazily — each item's error is boxed as it
77/// arrives, preserving the bytes-as-they-arrive latency of the inner provider.
78pub struct ErasedProvider<P>(P);
79
80#[async_trait]
81impl<P: LlmProvider> LlmProvider for ErasedProvider<P> {
82    type Error = BoxError;
83
84    async fn complete(
85        &self,
86        req: CompletionRequest,
87    ) -> Result<BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
88        let start = std::time::Instant::now();
89        match self.0.complete(req).await {
90            Ok(stream) => {
91                crate::metrics::record_call(Ok(()), start.elapsed());
92                Ok(stream.map(|item| item.map_err(BoxError::new)).boxed())
93            }
94            Err(err) => {
95                let err = BoxError::new(err);
96                crate::metrics::record_call(Err(err.kind()), start.elapsed());
97                Err(err)
98            }
99        }
100    }
101}
102
103/// The single trait-object provider type the control plane stores. Every
104/// concrete backend is erased to this via [`into_dyn`].
105pub type DynProvider = dyn LlmProvider<Error = BoxError>;
106
107/// Erase a concrete provider and wrap it in an `Arc` as a [`DynProvider`].
108///
109/// The one wiring-boundary call that lets a concrete backend be dispatched
110/// behind the runtime-swappable trait object.
111#[must_use]
112pub fn into_dyn<P: LlmProvider>(provider: P) -> Arc<DynProvider> {
113    Arc::new(ErasedProvider(provider))
114}
115
116#[cfg(test)]
117mod tests {
118    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]
119
120    use futures::{StreamExt, stream};
121
122    use super::{BoxError, DynProvider, into_dyn};
123    use crate::{Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError};
124
125    /// Provider whose `Error` is `DummyError` — a different concrete type than
126    /// `BoxError`, so erasing it actually exercises the conversion.
127    struct DummyProvider;
128
129    #[async_trait::async_trait]
130    impl LlmProvider for DummyProvider {
131        type Error = DummyError;
132
133        async fn complete(
134            &self,
135            req: CompletionRequest,
136        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
137        {
138            if req.messages.is_empty() {
139                return Err(DummyError::Other("no messages".to_owned()));
140            }
141            let chunks = vec![
142                Ok(Chunk::text_delta("hi")),
143                Ok(Chunk::Usage(Usage {
144                    input_tokens: 1,
145                    output_tokens: 1,
146                    ..Default::default()
147                })),
148                Ok(Chunk::Stop(StopReason::EndTurn)),
149            ];
150            Ok(stream::iter(chunks).boxed())
151        }
152    }
153
154    #[tokio::test]
155    async fn erased_provider_streams_to_completion() {
156        let provider: std::sync::Arc<DynProvider> = into_dyn(DummyProvider);
157        let mut req = CompletionRequest::new("m");
158        req.messages.push(crate::Message::user("yo"));
159
160        let stream = provider.complete(req).await.expect("stream opens");
161        let n = stream.count().await;
162        assert_eq!(n, 3);
163    }
164
165    #[tokio::test]
166    async fn erased_pre_stream_error_is_preserved() {
167        let provider: std::sync::Arc<DynProvider> = into_dyn(DummyProvider);
168        let req = CompletionRequest::new("m"); // no messages → pre-stream error
169
170        // `Ok` here is a stream (not `Debug`), so match rather than `expect_err`.
171        let Err(err) = provider.complete(req).await else {
172            panic!("expected pre-stream rejection");
173        };
174        // Display delegates to the inner DummyError's message.
175        assert_eq!(format!("{err}"), "other: no messages");
176        // The original error is exposed as the source of the chain.
177        let src = std::error::Error::source(&err).expect("source present");
178        assert_eq!(format!("{src}"), "other: no messages");
179    }
180
181    #[test]
182    fn box_error_satisfies_llm_error() {
183        fn require_llm_error<E: crate::error::LlmError>() {}
184        require_llm_error::<BoxError>();
185    }
186
187    #[test]
188    fn box_error_preserves_kind_through_erasure() {
189        use crate::error::{DummyError, LlmError, LlmErrorKind};
190        // A 429 provider error keeps its RateLimit kind after erasure.
191        let boxed = BoxError::new(DummyError::Provider {
192            status: 429,
193            body: String::new(),
194        });
195        assert_eq!(boxed.kind(), LlmErrorKind::RateLimit);
196        // The default (unclassified) kind also round-trips.
197        let other = BoxError::new(DummyError::Other("x".to_owned()));
198        assert_eq!(other.kind(), LlmErrorKind::Other);
199    }
200}
polyc_llm/erased.rs

polyc_llm/
erased.rs
