Skip to main content

sova_ai/
model.rs

1//! Type-erased [`LanguageModel`](aisdk::core::LanguageModel) handle stored in [`crate::AiClient`].
2
3use aisdk::core::capabilities::{
4    ReasoningSupport, StructuredOutputSupport, TextInputSupport, TextOutputSupport,
5    ToolCallSupport,
6};
7use aisdk::core::language_model::{
8    LanguageModel, LanguageModelOptions, LanguageModelResponse, LanguageModelStreamChunk,
9};
10use aisdk::Result as AisdkResult;
11use async_trait::async_trait;
12use futures_util::Stream;
13use std::fmt;
14use std::pin::Pin;
15use std::sync::Arc;
16use tokio::sync::Mutex;
17
18type ProviderStream =
19    Pin<Box<dyn Stream<Item = AisdkResult<Vec<LanguageModelStreamChunk>>> + Send>>;
20
21#[async_trait]
22trait DynModel: Send + Sync + fmt::Debug {
23    fn name(&self) -> String;
24    async fn generate_text(&self, options: LanguageModelOptions)
25        -> AisdkResult<LanguageModelResponse>;
26    async fn stream_text(&self, options: LanguageModelOptions) -> AisdkResult<ProviderStream>;
27}
28
29struct Holding<M: LanguageModel> {
30    name: String,
31    inner: Mutex<M>,
32}
33
34impl<M: LanguageModel> fmt::Debug for Holding<M> {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.debug_struct("Holding").field("name", &self.name).finish()
37    }
38}
39
40#[async_trait]
41impl<M: LanguageModel> DynModel for Holding<M> {
42    fn name(&self) -> String {
43        self.name.clone()
44    }
45
46    async fn generate_text(
47        &self,
48        options: LanguageModelOptions,
49    ) -> AisdkResult<LanguageModelResponse> {
50        self.inner.lock().await.generate_text(options).await
51    }
52
53    async fn stream_text(&self, options: LanguageModelOptions) -> AisdkResult<ProviderStream> {
54        self.inner.lock().await.stream_text(options).await
55    }
56}
57
58/// Cloneable model handle used as the default for `LanguageModelRequest`.
59///
60/// Implements AISDK capability markers so the typed builder accepts prompts/tools
61/// (runtime model selection — same idea as aisdk `DynamicModel`).
62#[derive(Clone)]
63pub struct SharedModel {
64    inner: Arc<dyn DynModel>,
65}
66
67impl fmt::Debug for SharedModel {
68    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69        f.debug_struct("SharedModel")
70            .field("name", &self.inner.name())
71            .finish()
72    }
73}
74
75impl SharedModel {
76    /// Wrap any AISDK [`LanguageModel`] (provider or [`crate::FakeAi`]).
77    pub fn wrap<M: LanguageModel>(model: M) -> Self {
78        let name = model.name();
79        Self {
80            inner: Arc::new(Holding {
81                name,
82                inner: Mutex::new(model),
83            }),
84        }
85    }
86}
87
88impl TextInputSupport for SharedModel {}
89impl TextOutputSupport for SharedModel {}
90impl ToolCallSupport for SharedModel {}
91impl StructuredOutputSupport for SharedModel {}
92impl ReasoningSupport for SharedModel {}
93
94#[async_trait]
95impl LanguageModel for SharedModel {
96    fn name(&self) -> String {
97        self.inner.name()
98    }
99
100    async fn generate_text(
101        &mut self,
102        options: LanguageModelOptions,
103    ) -> AisdkResult<LanguageModelResponse> {
104        self.inner.generate_text(options).await
105    }
106
107    async fn stream_text(&mut self, options: LanguageModelOptions) -> AisdkResult<ProviderStream> {
108        self.inner.stream_text(options).await
109    }
110}