polyc_llm/

lib.rs

//! Provider-agnostic LLM trait and wire types for polychrome.
//!
//! This crate defines the [`LlmProvider`] trait that every concrete provider
//! backend implements, and the language-shaped Rust types that flow through it.
//!
//! [`LlmProvider`] is the seam that lets the planner swap LLM backends
//! without touching its own code: one impl crate per provider, dispatched
//! behind a `dyn LlmProvider` trait object.
//!
//! # Modules
//!
//! - [`request`] — [`CompletionRequest`] and everything reachable from it
//!   ([`Message`], [`Content`], [`ToolSpec`], [`ToolChoice`]).
//! - [`chunk`] — [`Chunk`] (streaming response events) and friends
//!   ([`Usage`], [`StopReason`]).
//! - [`error`] — the [`LlmError`] trait bound that [`LlmProvider::Error`]
//!   must satisfy.
//!
//! The trait itself lives in this crate root; the wire types live in the
//! modules above and are re-exported here for convenience.

pub mod chunk;
pub mod erased;
pub mod error;
mod metrics;
pub mod model_info;
pub mod preflight;
pub mod request;
pub mod sse;
pub mod turn;

use async_trait::async_trait;
use futures::stream::BoxStream;

pub use chunk::{Chunk, StopReason, Usage};
pub use erased::{BoxError, DynProvider, ErasedProvider, into_dyn};
pub use error::{LlmError, LlmErrorKind, kind_from_http_status, parse_retry_after};
pub use model_info::{FALLBACK_CONTEXT_WINDOW, ModelInfo, lookup_model, lookup_model_or_fallback};
pub use preflight::{PreflightReport, ProbeOutcome, preflight};
pub use request::{
    APPROVAL_STATUS_GROUND_RULE, CacheHint, CompletionRequest, Content, GATED_TOOL_APPROVAL_NOTE,
    ImageRef, JsonSchema, Message, Role, ToolCall, ToolChoice, ToolResult, ToolSpec,
};

/// Force-register this crate's Prometheus call-latency histogram.
///
/// Makes it appear in a `/metrics` scrape immediately — before any provider
/// call has completed. Idempotent (backed by a `OnceLock`); call once at
/// process startup, alongside any other crate's own `init_metrics`.
pub fn init_metrics() {
    metrics::force();
}

/// The seam between the planner and any concrete LLM backend.
///
/// One implementation per backend, registered at startup and dispatched behind
/// a trait object so the planner swaps backends without recompiling. A single
/// method —
/// [`complete`](LlmProvider::complete) — takes a [`CompletionRequest`] and
/// returns a stream of [`Chunk`]s; non-streaming callers simply drain the
/// stream.
///
/// The `'static` bound and [`Send`] + [`Sync`] make providers storable in the
/// control plane's routing table (`arc-swap`'d) and shareable across tasks.
/// [`Self::Error`] is bounded by [`LlmError`] so failures are uniform across
/// providers while each keeps its own concrete error type.
#[async_trait]
pub trait LlmProvider: Send + Sync + 'static {
    /// The provider's concrete error type. Bounded by [`LlmError`]
    /// (`std::error::Error + Send + Sync + 'static`).
    type Error: LlmError;

    /// Runs a completion, returning a stream of [`Chunk`]s.
    ///
    /// The outer `Result` reports failures that occur before the stream opens
    /// (auth, request validation, transport dial). Once the stream is live,
    /// per-chunk failures surface as `Err` items within it — a stream can yield
    /// several good chunks and then fault mid-flight.
    ///
    /// # Errors
    ///
    /// Returns [`Self::Error`] if the request cannot be dispatched or the
    /// provider rejects it before streaming begins.
    async fn complete(
        &self,
        req: CompletionRequest,
    ) -> Result<BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>;
}

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

    use futures::{StreamExt, stream};

    use super::{Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError};

    /// Reference provider: echoes the first user message back as text, then a
    /// usage tally and an end-of-turn stop. Proves the trait is implementable
    /// and that its stream can be driven to completion.
    struct EchoProvider;

    #[async_trait::async_trait]
    impl LlmProvider for EchoProvider {
        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("echo")),
                Ok(Chunk::Usage(Usage {
                    input_tokens: 3,
                    output_tokens: 1,
                    ..Default::default()
                })),
                Ok(Chunk::Stop(StopReason::EndTurn)),
            ];
            Ok(stream::iter(chunks).boxed())
        }
    }

    #[tokio::test]
    async fn provider_streams_chunks_to_completion() {
        let provider = EchoProvider;
        let mut req = CompletionRequest::new("test-model");
        req.messages.push(super::Message::user("hi"));

        let stream = provider.complete(req).await.expect("stream opens");
        let collected: Vec<Chunk> = stream.map(Result::unwrap).collect().await;

        assert_eq!(collected.len(), 3);
        assert_eq!(collected[0], Chunk::text_delta("echo"));
        assert!(matches!(collected[2], Chunk::Stop(StopReason::EndTurn)));
    }

    #[tokio::test]
    async fn provider_reports_pre_stream_failure() {
        let provider = EchoProvider;
        let req = CompletionRequest::new("test-model"); // no messages

        match provider.complete(req).await {
            Err(DummyError::Other(_)) => {}
            Err(other) => panic!("wrong error: {other}"),
            Ok(_) => panic!("expected pre-stream rejection"),
        }
    }

    #[tokio::test]
    async fn usable_as_trait_object() {
        // The PRD calls for `dyn LlmProvider` dispatch; confirm object safety.
        let provider: Box<dyn LlmProvider<Error = DummyError>> = Box::new(EchoProvider);
        let mut req = CompletionRequest::new("m");
        req.messages.push(super::Message::user("yo"));
        let stream = provider.complete(req).await.expect("stream opens");
        let n = stream.count().await;
        assert_eq!(n, 3);
    }
}