Skip to main content

sova_ai/
bound.rs

1//! `req.ai()` fluent generate / stream helpers.
2
3use 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
10/// `req.ai()` — bound to the installed [`AiClient`].
11pub 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
22/// Fluent call builder (`prompt` / `system` / `generate` / `stream_response`).
23pub 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    /// Full AISDK builder (model already set).
57    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    /// Non-streaming generation.
86    pub async fn generate(self) -> Result<GenerateTextResponse, AiError> {
87        let mut req = self.build_request()?;
88        Ok(req.generate_text().await?)
89    }
90
91    /// Convenience: generate and return assistant text (empty string if none).
92    pub async fn text(self) -> Result<String, AiError> {
93        let out = self.generate().await?;
94        Ok(out.text().unwrap_or_default())
95    }
96
97    /// Stream model output as `text/event-stream` [`Response`].
98    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}