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