Skip to main content

sova_ai/
client.rs

1//! [`Ai`] builder and [`AiClient`] app state.
2
3use crate::error::AiError;
4use crate::fake::FakeAi;
5use crate::model::SharedModel;
6use aisdk::core::{LanguageModel, LanguageModelRequest};
7
8/// Plugin builder — install with `app.install(Ai::…)`.
9#[derive(Clone)]
10pub struct Ai {
11    model: Option<SharedModel>,
12    fake: Option<FakeAi>,
13    default_system: Option<String>,
14}
15
16impl Default for Ai {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl Ai {
23    pub fn new() -> Self {
24        Self {
25            model: None,
26            fake: None,
27            default_system: None,
28        }
29    }
30
31    /// Default language model (OpenAI, Anthropic, [`FakeAi`], …).
32    pub fn model<M: LanguageModel>(mut self, model: M) -> Self {
33        self.model = Some(SharedModel::wrap(model));
34        self.fake = None;
35        self
36    }
37
38    /// Install a recording fake (tests). Same as `.model(fake)` plus `client.fake()`.
39    pub fn fake(fake: FakeAi) -> Self {
40        Self {
41            model: Some(SharedModel::wrap(fake.clone())),
42            fake: Some(fake),
43            default_system: None,
44        }
45    }
46
47    /// Default system prompt prepended when a bound call does not set `.system()`.
48    pub fn system(mut self, system: impl Into<String>) -> Self {
49        self.default_system = Some(system.into());
50        self
51    }
52
53    pub(crate) fn default_system(&self) -> Option<&str> {
54        self.default_system.as_deref()
55    }
56
57    pub fn into_client(self) -> Result<AiClient, AiError> {
58        let model = self.model.ok_or(AiError::NoModel)?;
59        Ok(AiClient {
60            model,
61            fake: self.fake,
62            default_system: self.default_system,
63        })
64    }
65
66    /// Build the client without installing (tests / manual `app.state`).
67    pub fn client(self) -> Result<AiClient, AiError> {
68        self.into_client()
69    }
70}
71
72/// Shared AI client in app state (`req.state::<AiClient>()` / `req.ai()`).
73#[derive(Clone)]
74pub struct AiClient {
75    model: SharedModel,
76    fake: Option<FakeAi>,
77    default_system: Option<String>,
78}
79
80impl AiClient {
81    pub fn model(&self) -> SharedModel {
82        self.model.clone()
83    }
84
85    pub fn fake(&self) -> Option<&FakeAi> {
86        self.fake.as_ref()
87    }
88
89    pub fn default_system(&self) -> Option<&str> {
90        self.default_system.as_deref()
91    }
92
93    /// Escape hatch: full AISDK request builder pre-seeded with the installed model.
94    pub fn builder(
95        &self,
96    ) -> aisdk::core::language_model::request::LanguageModelRequestBuilder<
97        SharedModel,
98        aisdk::core::language_model::request::SystemStage,
99    > {
100        LanguageModelRequest::builder().model(self.model.clone())
101    }
102}