polyc_llm/
lib.rs

1//! Provider-agnostic LLM trait and wire types for polychrome.
2//!
3//! This crate defines the [`LlmProvider`] trait that every concrete provider
4//! backend implements, and the language-shaped Rust types that flow through it.
5//!
6//! [`LlmProvider`] is the seam that lets the planner swap LLM backends
7//! without touching its own code: one impl crate per provider, dispatched
8//! behind a `dyn LlmProvider` trait object.
9//!
10//! # Modules
11//!
12//! - [`request`] — [`CompletionRequest`] and everything reachable from it
13//!   ([`Message`], [`Content`], [`ToolSpec`], [`ToolChoice`]).
14//! - [`chunk`] — [`Chunk`] (streaming response events) and friends
15//!   ([`Usage`], [`StopReason`]).
16//! - [`error`] — the [`LlmError`] trait bound that [`LlmProvider::Error`]
17//!   must satisfy.
18//!
19//! The trait itself lives in this crate root; the wire types live in the
20//! modules above and are re-exported here for convenience.
21
22pub mod chunk;
23pub mod erased;
24pub mod error;
25pub mod request;
26pub mod sse;
27pub mod turn;
28
29use async_trait::async_trait;
30use futures::stream::BoxStream;
31
32pub use chunk::{Chunk, StopReason, Usage};
33pub use erased::{BoxError, DynProvider, ErasedProvider, into_dyn};
34pub use error::{LlmError, LlmErrorKind, kind_from_http_status};
35pub use request::{
36    CompletionRequest, Content, ImageRef, JsonSchema, Message, Role, ToolCall, ToolChoice,
37    ToolResult, ToolSpec, humanize_tool_name,
38};
39
40/// The seam between the planner and any concrete LLM backend.
41///
42/// One implementation per backend, registered at startup and dispatched behind
43/// a trait object so the planner swaps backends without recompiling. A single
44/// method —
45/// [`complete`](LlmProvider::complete) — takes a [`CompletionRequest`] and
46/// returns a stream of [`Chunk`]s; non-streaming callers simply drain the
47/// stream.
48///
49/// The `'static` bound and [`Send`] + [`Sync`] make providers storable in the
50/// control plane's routing table (`arc-swap`'d) and shareable across tasks.
51/// [`Self::Error`] is bounded by [`LlmError`] so failures are uniform across
52/// providers while each keeps its own concrete error type.
53#[async_trait]
54pub trait LlmProvider: Send + Sync + 'static {
55    /// The provider's concrete error type. Bounded by [`LlmError`]
56    /// (`std::error::Error + Send + Sync + 'static`).
57    type Error: LlmError;
58
59    /// Runs a completion, returning a stream of [`Chunk`]s.
60    ///
61    /// The outer `Result` reports failures that occur before the stream opens
62    /// (auth, request validation, transport dial). Once the stream is live,
63    /// per-chunk failures surface as `Err` items within it — a stream can yield
64    /// several good chunks and then fault mid-flight.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`Self::Error`] if the request cannot be dispatched or the
69    /// provider rejects it before streaming begins.
70    async fn complete(
71        &self,
72        req: CompletionRequest,
73    ) -> Result<BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>;
74}
75
76#[cfg(test)]
77mod tests {
78    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]
79
80    use futures::{StreamExt, stream};
81
82    use super::{Chunk, CompletionRequest, LlmProvider, StopReason, Usage, error::DummyError};
83
84    /// Reference provider: echoes the first user message back as text, then a
85    /// usage tally and an end-of-turn stop. Proves the trait is implementable
86    /// and that its stream can be driven to completion.
87    struct EchoProvider;
88
89    #[async_trait::async_trait]
90    impl LlmProvider for EchoProvider {
91        type Error = DummyError;
92
93        async fn complete(
94            &self,
95            req: CompletionRequest,
96        ) -> Result<futures::stream::BoxStream<'static, Result<Chunk, Self::Error>>, Self::Error>
97        {
98            if req.messages.is_empty() {
99                return Err(DummyError::Other("no messages".to_owned()));
100            }
101            let chunks = vec![
102                Ok(Chunk::text_delta("echo")),
103                Ok(Chunk::Usage(Usage {
104                    input_tokens: 3,
105                    output_tokens: 1,
106                })),
107                Ok(Chunk::Stop(StopReason::EndTurn)),
108            ];
109            Ok(stream::iter(chunks).boxed())
110        }
111    }
112
113    #[tokio::test]
114    async fn provider_streams_chunks_to_completion() {
115        let provider = EchoProvider;
116        let mut req = CompletionRequest::new("test-model");
117        req.messages.push(super::Message::user("hi"));
118
119        let stream = provider.complete(req).await.expect("stream opens");
120        let collected: Vec<Chunk> = stream.map(Result::unwrap).collect().await;
121
122        assert_eq!(collected.len(), 3);
123        assert_eq!(collected[0], Chunk::text_delta("echo"));
124        assert!(matches!(collected[2], Chunk::Stop(StopReason::EndTurn)));
125    }
126
127    #[tokio::test]
128    async fn provider_reports_pre_stream_failure() {
129        let provider = EchoProvider;
130        let req = CompletionRequest::new("test-model"); // no messages
131
132        match provider.complete(req).await {
133            Err(DummyError::Other(_)) => {}
134            Err(other) => panic!("wrong error: {other}"),
135            Ok(_) => panic!("expected pre-stream rejection"),
136        }
137    }
138
139    #[tokio::test]
140    async fn usable_as_trait_object() {
141        // The PRD calls for `dyn LlmProvider` dispatch; confirm object safety.
142        let provider: Box<dyn LlmProvider<Error = DummyError>> = Box::new(EchoProvider);
143        let mut req = CompletionRequest::new("m");
144        req.messages.push(super::Message::user("yo"));
145        let stream = provider.complete(req).await.expect("stream opens");
146        let n = stream.count().await;
147        assert_eq!(n, 3);
148    }
149}
polyc_llm/lib.rs

polyc_llm/
lib.rs
