polyc_llm/

erased.rs

//! Type erasure for [`LlmProvider`] so the control plane can hold a single
//! `Arc<dyn LlmProvider>` regardless of which backend is configured.
//!
//! The trait keeps a per-provider associated [`LlmError`] type (each backend
//! ships its own concrete error). A trait object must fix that associated
//! type, so this module supplies one uniform error — [`BoxError`] — and an
//! [`ErasedProvider`] adapter that maps any provider's error into it. The
//! result is [`DynProvider`], the single trait-object type callers store and
//! dispatch through.
//!
//! Adding a backend then costs one trait impl plus one [`into_dyn`] call at the
//! wiring boundary — no change to the dispatch site or the planner.

use std::sync::Arc;

use async_trait::async_trait;
use futures::stream::{BoxStream, StreamExt};

use crate::{
    Chunk, CompletionRequest, LlmProvider,
    error::{LlmError, LlmErrorKind},
};

/// A provider error erased to one concrete type, so backends with differing
/// associated `Error`s can be stored behind a single trait object.
///
/// Transparent wrapper: [`Display`](std::fmt::Display) delegates to the inner
/// error and [`source`](std::error::Error::source) exposes it, so logs and
/// error chains read exactly as the un-erased error did. The second field
/// preserves the original [`LlmErrorKind`] across erasure (the boxed `dyn Error`
/// alone could not be re-classified).
#[derive(Debug)]
pub struct BoxError(
    Box<dyn std::error::Error + Send + Sync + 'static>,
    LlmErrorKind,
);

impl BoxError {
    /// Erase any [`LlmError`] into a `BoxError`, capturing its
    /// [`kind`](LlmError::kind) so the classification survives erasure.
    #[must_use]
    pub fn new<E: LlmError>(err: E) -> Self {
        let kind = err.kind();
        Self(Box::new(err), kind)
    }
}

impl LlmError for BoxError {
    fn kind(&self) -> LlmErrorKind {
        self.1
    }
}

impl std::fmt::Display for BoxError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        std::fmt::Display::fmt(&self.0, f)
    }
}

impl std::error::Error for BoxError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(&*self.0)
    }
}

/// Adapter that wraps a concrete [`LlmProvider`] and erases its associated
/// error to [`BoxError`], so the wrapped value coerces to [`DynProvider`].
///
/// The completion stream is mapped lazily — each item's error is boxed as it
/// arrives, preserving the bytes-as-they-arrive latency of the inner provider.
pub struct ErasedProvider<P>(P);

#[async_trait]
impl<P: LlmProvider> LlmProvider for ErasedProvider<P> {
    type Error = BoxError;

    async fn complete(
        &self,
        req: CompletionRequest,
    ) -> Result<BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error> {
        let stream = self.0.complete(req).await.map_err(BoxError::new)?;
        Ok(stream.map(|item| item.map_err(BoxError::new)).boxed())
    }
}

/// The single trait-object provider type the control plane stores. Every
/// concrete backend is erased to this via [`into_dyn`].
pub type DynProvider = dyn LlmProvider<Error = BoxError>;

/// Erase a concrete provider and wrap it in an `Arc` as a [`DynProvider`].
///
/// The one wiring-boundary call that lets a concrete backend be dispatched
/// behind the runtime-swappable trait object.
#[must_use]
pub fn into_dyn<P: LlmProvider>(provider: P) -> Arc<DynProvider> {
    Arc::new(ErasedProvider(provider))
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use futures::{StreamExt, stream};

    use super::{BoxError, DynProvider, into_dyn};
    use crate::{Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError};

    /// Provider whose `Error` is `DummyError` — a different concrete type than
    /// `BoxError`, so erasing it actually exercises the conversion.
    struct DummyProvider;

    #[async_trait::async_trait]
    impl LlmProvider for DummyProvider {
        type Error = DummyError;

        async fn complete(
            &self,
            req: CompletionRequest,
        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
        {
            if req.messages.is_empty() {
                return Err(DummyError::Other("no messages".to_owned()));
            }
            let chunks = vec![
                Ok(Chunk::text_delta("hi")),
                Ok(Chunk::Usage(Usage {
                    input_tokens: 1,
                    output_tokens: 1,
                })),
                Ok(Chunk::Stop(StopReason::EndTurn)),
            ];
            Ok(stream::iter(chunks).boxed())
        }
    }

    #[tokio::test]
    async fn erased_provider_streams_to_completion() {
        let provider: std::sync::Arc<DynProvider> = into_dyn(DummyProvider);
        let mut req = CompletionRequest::new("m");
        req.messages.push(crate::Message::user("yo"));

        let stream = provider.complete(req).await.expect("stream opens");
        let n = stream.count().await;
        assert_eq!(n, 3);
    }

    #[tokio::test]
    async fn erased_pre_stream_error_is_preserved() {
        let provider: std::sync::Arc<DynProvider> = into_dyn(DummyProvider);
        let req = CompletionRequest::new("m"); // no messages → pre-stream error

        // `Ok` here is a stream (not `Debug`), so match rather than `expect_err`.
        let Err(err) = provider.complete(req).await else {
            panic!("expected pre-stream rejection");
        };
        // Display delegates to the inner DummyError's message.
        assert_eq!(format!("{err}"), "other: no messages");
        // The original error is exposed as the source of the chain.
        let src = std::error::Error::source(&err).expect("source present");
        assert_eq!(format!("{src}"), "other: no messages");
    }

    #[test]
    fn box_error_satisfies_llm_error() {
        fn require_llm_error<E: crate::error::LlmError>() {}
        require_llm_error::<BoxError>();
    }

    #[test]
    fn box_error_preserves_kind_through_erasure() {
        use crate::error::{DummyError, LlmError, LlmErrorKind};
        // A 429 provider error keeps its RateLimit kind after erasure.
        let boxed = BoxError::new(DummyError::Provider {
            status: 429,
            body: String::new(),
        });
        assert_eq!(boxed.kind(), LlmErrorKind::RateLimit);
        // The default (unclassified) kind also round-trips.
        let other = BoxError::new(DummyError::Other("x".to_owned()));
        assert_eq!(other.kind(), LlmErrorKind::Other);
    }
}