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        response::{ChatResponse, EmbeddingsResponse},
11    },
12};
13use async_trait::async_trait;
14#[cfg(feature = "stream")]
15use futures::stream::BoxStream;
16
17use tools_rs::ToolCollection;
18
19#[async_trait]
20pub trait CompletionProvider: Send + Sync {
21    async fn complete(
22        &mut self,
23        messages: &mut Messages,
24        tools: Option<&ToolCollection>,
25        options: Option<&ChatOptions>,
26        structured_output: Option<&schemars::Schema>,
27    ) -> Result<ChatResponse, ChatFailure>;
28}
29
30#[cfg(feature = "stream")]
31#[async_trait]
32pub trait StreamProvider: Send + Sync {
33    async fn stream(
34        &mut self,
35        messages: &mut Messages,
36        tools: Option<&ToolCollection>,
37        options: Option<&ChatOptions>,
38    ) -> Result<BoxStream<'static, Result<StreamEvent, ChatError>>, ChatError>;
39
40    /// Called after the stream has been fully consumed with the final response.
41    /// Providers can override this to store state from the completed stream.
42    fn on_stream_done(&mut self, _response: &ChatResponse) {}
43}
44
45#[async_trait]
46pub trait EmbeddingsProvider: Send + Sync {
47    async fn embed(&self, messages: &mut Messages) -> Result<EmbeddingsResponse, ChatFailure>;
48}