1use crate::client::AiClient;
4use crate::error::AiError;
5use crate::stream::stream_to_response;
6use aisdk::core::{GenerateTextResponse, LanguageModelRequest, Messages};
7use sova_core::{Request, Response};
8use std::sync::Arc;
9
10pub trait AiExt {
12 fn ai(&self) -> AiBound;
13}
14
15impl AiExt for Request {
16 fn ai(&self) -> AiBound {
17 let client = self.state::<AiClient>();
18 AiBound::new(client)
19 }
20}
21
22pub struct AiBound {
24 client: Arc<AiClient>,
25 system: Option<String>,
26 prompt: Option<String>,
27 messages: Option<Messages>,
28}
29
30impl AiBound {
31 pub(crate) fn new(client: Arc<AiClient>) -> Self {
32 let system = client.default_system().map(str::to_owned);
33 Self {
34 client,
35 system,
36 prompt: None,
37 messages: None,
38 }
39 }
40
41 pub fn system(mut self, system: impl Into<String>) -> Self {
42 self.system = Some(system.into());
43 self
44 }
45
46 pub fn prompt(mut self, prompt: impl Into<String>) -> Self {
47 self.prompt = Some(prompt.into());
48 self
49 }
50
51 pub fn messages(mut self, messages: Messages) -> Self {
52 self.messages = Some(messages);
53 self
54 }
55
56 pub fn builder(
58 self,
59 ) -> aisdk::core::language_model::request::LanguageModelRequestBuilder<
60 crate::SharedModel,
61 aisdk::core::language_model::request::SystemStage,
62 > {
63 self.client.builder()
64 }
65
66 fn build_request(self) -> Result<LanguageModelRequest<crate::SharedModel>, AiError> {
67 let AiBound {
68 client,
69 system,
70 prompt,
71 messages,
72 } = self;
73
74 let b = LanguageModelRequest::builder().model(client.model());
75
76 Ok(match (system, prompt, messages) {
77 (Some(sys), Some(p), None) => b.system(sys).prompt(p).build(),
78 (Some(sys), None, Some(m)) => b.system(sys).messages(m).build(),
79 (None, Some(p), None) => b.prompt(p).build(),
80 (None, None, Some(m)) => b.messages(m).build(),
81 _ => return Err(AiError::EmptyPrompt),
82 })
83 }
84
85 pub async fn generate(self) -> Result<GenerateTextResponse, AiError> {
87 let mut req = self.build_request()?;
88 Ok(req.generate_text().await?)
89 }
90
91 pub async fn text(self) -> Result<String, AiError> {
93 let out = self.generate().await?;
94 Ok(out.text().unwrap_or_default())
95 }
96
97 pub async fn stream_response(self) -> Result<Response, AiError> {
99 let mut req = self.build_request()?;
100 let stream = req.stream_text().await?;
101 Ok(stream_to_response(stream))
102 }
103}