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