1use crate::error::AiError;
4use crate::fake::FakeAi;
5use crate::model::SharedModel;
6use aisdk::core::{LanguageModel, LanguageModelRequest};
7
8#[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 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 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 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 pub fn client(self) -> Result<AiClient, AiError> {
68 self.into_client()
69 }
70}
71
72#[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 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}