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