gemini_rust/cache/
builder.rs1use 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
17pub 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 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 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 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 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 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 pub fn with_content(mut self, content: Content) -> Self {
80 self.contents.push(content);
81 self
82 }
83
84 pub fn with_contents(mut self, contents: Vec<Content>) -> Self {
86 self.contents.extend(contents);
87 self
88 }
89
90 pub fn with_tool(mut self, tool: Tool) -> Self {
92 self.tools.push(tool);
93 self
94 }
95
96 pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
98 self.tools.extend(tools);
99 self
100 }
101
102 pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
104 self.tool_config = Some(tool_config);
105 self
106 }
107
108 pub fn with_ttl(mut self, ttl: Duration) -> Self {
111 self.expiration = Some(CacheExpirationRequest::from_ttl(ttl));
112 self
113 }
114
115 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 #[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}