Skip to main content

chat_core/
traits.rs

1#[cfg(feature = "stream")]
2use crate::error::ChatError;
3#[cfg(feature = "stream")]
4use crate::types::response::StreamEvent;
5use crate::{
6    error::ChatFailure,
7    types::{
8        messages::Messages,
9        options::ChatOptions,
10        provider_meta::ProviderMeta,
11        response::{ChatResponse, EmbeddingsResponse},
12        tools::ToolDeclarations,
13    },
14};
15use async_trait::async_trait;
16#[cfg(feature = "stream")]
17use futures::stream::BoxStream;
18
19/// Tool declarations are the one and only view providers have into the
20/// user's tool collections. The trait intentionally exposes nothing
21/// about execution, metadata, or strategy — providers translate
22/// declarations to their wire format and nothing else. For the
23/// builder-side collection API, see [`crate::types::tools`].
24#[async_trait]
25pub trait CompletionProvider: Send + Sync {
26    /// Run one completion step.
27    ///
28    /// `tool_declarations`, when `Some`, is a view onto every scoped
29    /// tool collection the chat loop holds. Call `.json()` on it to get
30    /// the JSON array of function declarations that providers splice
31    /// into their native request format.
32    async fn complete(
33        &mut self,
34        messages: &mut Messages,
35        tool_declarations: Option<&dyn ToolDeclarations>,
36        options: Option<&ChatOptions>,
37        structured_output: Option<&schemars::Schema>,
38    ) -> Result<ChatResponse, ChatFailure>;
39
40    fn metadata(&self) -> Option<&ProviderMeta> {
41        None
42    }
43}
44
45#[cfg(feature = "stream")]
46#[async_trait]
47pub trait StreamProvider: Send + Sync {
48    async fn stream(
49        &mut self,
50        messages: &mut Messages,
51        tool_declarations: Option<&dyn ToolDeclarations>,
52        options: Option<&ChatOptions>,
53    ) -> Result<BoxStream<'static, Result<StreamEvent, ChatError>>, ChatError>;
54
55    /// Called after the stream has been fully consumed with the final response.
56    /// Providers can override this to store state from the completed stream.
57    fn on_stream_done(&mut self, _response: &ChatResponse) {}
58}
59
60/// Combined supertrait for providers that support both completion and streaming.
61/// All providers that implement both `CompletionProvider` and `StreamProvider`
62/// automatically implement this trait via the blanket impl.
63#[cfg(feature = "stream")]
64pub trait ChatProvider: CompletionProvider + StreamProvider {}
65
66#[cfg(feature = "stream")]
67impl<T: CompletionProvider + StreamProvider> ChatProvider for T {}
68
69#[async_trait]
70pub trait EmbeddingsProvider: Send + Sync {
71    async fn embed(&self, messages: &mut Messages) -> Result<EmbeddingsResponse, ChatFailure>;
72}