gemini_rust/cache/
builder.rs

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