Skip to main content

daimon_core/
model.rs

1//! The [`Model`] trait — the plugin interface for LLM providers.
2
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6
7use crate::error::Result;
8use crate::stream::ResponseStream;
9use crate::types::{ChatRequest, ChatResponse};
10
11/// Trait for LLM providers. Supports both synchronous and streaming generation.
12///
13/// Implement this trait to add a new model provider to Daimon. The trait is
14/// object-safe via [`ErasedModel`], so providers can be stored behind
15/// `Arc<dyn ErasedModel>`.
16pub trait Model: Send + Sync {
17    /// Generates a complete response for the given request.
18    fn generate(&self, request: &ChatRequest) -> impl Future<Output = Result<ChatResponse>> + Send;
19
20    /// Returns a stream of response chunks for token-by-token output.
21    fn generate_stream(
22        &self,
23        request: &ChatRequest,
24    ) -> impl Future<Output = Result<ResponseStream>> + Send;
25}
26
27/// Shared ownership of a model via `Arc<dyn ErasedModel>`.
28pub type SharedModel = Arc<dyn ErasedModel>;
29
30/// Object-safe wrapper for [`Model`], enabling dynamic dispatch via `Arc<dyn ErasedModel>`.
31pub trait ErasedModel: Send + Sync {
32    /// Object-safe version of [`Model::generate`].
33    fn generate_erased<'a>(
34        &'a self,
35        request: &'a ChatRequest,
36    ) -> Pin<Box<dyn Future<Output = Result<ChatResponse>> + Send + 'a>>;
37
38    /// Object-safe version of [`Model::generate_stream`].
39    fn generate_stream_erased<'a>(
40        &'a self,
41        request: &'a ChatRequest,
42    ) -> Pin<Box<dyn Future<Output = Result<ResponseStream>> + Send + 'a>>;
43}
44
45impl<T: Model> ErasedModel for T {
46    fn generate_erased<'a>(
47        &'a self,
48        request: &'a ChatRequest,
49    ) -> Pin<Box<dyn Future<Output = Result<ChatResponse>> + Send + 'a>> {
50        Box::pin(self.generate(request))
51    }
52
53    fn generate_stream_erased<'a>(
54        &'a self,
55        request: &'a ChatRequest,
56    ) -> Pin<Box<dyn Future<Output = Result<ResponseStream>> + Send + 'a>> {
57        Box::pin(self.generate_stream(request))
58    }
59}