gemini_rust/cache/
builder.rs

1use std::sync::Arc;
2use std::time::Duration;
3use tracing::instrument;
4
5use snafu::ResultExt;
6
7use crate::client::GeminiClient;
8use crate::models::Content;
9
10use super::handle::*;
11use super::model::*;
12use super::*;
13
14use crate::tools::Tool;
15use crate::tools::ToolConfig;
16
17/// Builder for creating cached content with a fluent API.
18#[derive(Clone)]
19pub struct CacheBuilder {
20    client: Arc<GeminiClient>,
21    display_name: Option<String>,
22    contents: Vec<Content>,
23    system_instruction: Option<Content>,
24    tools: Vec<Tool>,
25    tool_config: Option<ToolConfig>,
26    expiration: Option<CacheExpirationRequest>,
27}
28
29impl CacheBuilder {
30    /// Creates a new CacheBuilder instance.
31    pub(crate) fn new(client: Arc<GeminiClient>) -> Self {
32        Self {
33            client,
34            display_name: None,
35            contents: Vec::new(),
36            system_instruction: None,
37            tools: Vec::new(),
38            tool_config: None,
39            expiration: None,
40        }
41    }
42
43    /// Set a display name for the cached content.
44    /// Maximum 128 Unicode characters.
45    pub fn with_display_name<S: Into<String>>(mut self, display_name: S) -> Result<Self, Error> {
46        let display_name = display_name.into();
47        let chars = display_name.chars().count();
48        snafu::ensure!(
49            chars <= 128,
50            LongDisplayNameSnafu {
51                display_name,
52                chars
53            }
54        );
55        self.display_name = Some(display_name);
56        Ok(self)
57    }
58
59    /// Set the system instruction for the cached content.
60    pub fn with_system_instruction<S: Into<String>>(mut self, instruction: S) -> Self {
61        self.system_instruction = Some(Content::text(instruction.into()));
62        self
63    }
64
65    /// Add a user message to the cached content.
66    pub fn with_user_message<S: Into<String>>(mut self, message: S) -> Self {
67        self.contents
68            .push(crate::Message::user(message.into()).content);
69        self
70    }
71
72    /// Add a model message to the cached content.
73    pub fn with_model_message<S: Into<String>>(mut self, message: S) -> Self {
74        self.contents
75            .push(crate::Message::model(message.into()).content);
76        self
77    }
78
79    /// Add content directly to the cached content.
80    pub fn with_content(mut self, content: Content) -> Self {
81        self.contents.push(content);
82        self
83    }
84
85    /// Add multiple contents to the cached content.
86    pub fn with_contents(mut self, contents: Vec<Content>) -> Self {
87        self.contents.extend(contents);
88        self
89    }
90
91    /// Add a tool to the cached content.
92    pub fn with_tool(mut self, tool: Tool) -> Self {
93        self.tools.push(tool);
94        self
95    }
96
97    /// Add multiple tools to the cached content.
98    pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
99        self.tools.extend(tools);
100        self
101    }
102
103    /// Set the tool configuration.
104    pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
105        self.tool_config = Some(tool_config);
106        self
107    }
108
109    /// Set the TTL (Time To Live) for the cached content.
110    /// The cache will automatically expire after this duration.
111    pub fn with_ttl(mut self, ttl: Duration) -> Self {
112        self.expiration = Some(CacheExpirationRequest::from_ttl(ttl));
113        self
114    }
115
116    /// Set an explicit expiration time for the cached content.
117    pub fn with_expire_time(mut self, expire_time: time::OffsetDateTime) -> Self {
118        self.expiration = Some(CacheExpirationRequest::from_expire_time(expire_time));
119        self
120    }
121
122    /// Execute the cache creation request.
123    #[instrument(skip_all, fields(
124        display.name = self.display_name,
125        messages.count = self.contents.len(),
126        tools.count = self.tools.len(),
127        system_instruction.present = self.system_instruction.is_some(),
128    ))]
129    pub async fn execute(self) -> Result<CachedContentHandle, Error> {
130        let model = self.client.model.clone();
131        let expiration = self.expiration.ok_or(Error::MissingExpiration)?;
132
133        let cached_content = CreateCachedContentRequest {
134            display_name: self.display_name,
135            model,
136            contents: if self.contents.is_empty() {
137                None
138            } else {
139                Some(self.contents)
140            },
141            tools: if self.tools.is_empty() {
142                None
143            } else {
144                Some(self.tools)
145            },
146            system_instruction: self.system_instruction,
147            tool_config: self.tool_config,
148            expiration,
149        };
150
151        let response = self
152            .client
153            .create_cached_content(cached_content)
154            .await
155            .map_err(Box::new)
156            .context(ClientSnafu)?;
157
158        let cache_name = response.name;
159
160        Ok(CachedContentHandle::new(cache_name, self.client))
161    }
162}