adk_gemini/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!(chars <= 128, LongDisplayNameSnafu { display_name, chars });
48 self.display_name = Some(display_name);
49 Ok(self)
50 }
51
52 pub fn with_system_instruction<S: Into<String>>(mut self, instruction: S) -> Self {
54 self.system_instruction = Some(Content::text(instruction.into()));
55 self
56 }
57
58 pub fn with_user_message<S: Into<String>>(mut self, message: S) -> Self {
60 self.contents.push(crate::Message::user(message.into()).content);
61 self
62 }
63
64 pub fn with_model_message<S: Into<String>>(mut self, message: S) -> Self {
66 self.contents.push(crate::Message::model(message.into()).content);
67 self
68 }
69
70 pub fn with_content(mut self, content: Content) -> Self {
72 self.contents.push(content);
73 self
74 }
75
76 pub fn with_contents(mut self, contents: Vec<Content>) -> Self {
78 self.contents.extend(contents);
79 self
80 }
81
82 pub fn with_tool(mut self, tool: Tool) -> Self {
84 self.tools.push(tool);
85 self
86 }
87
88 pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
90 self.tools.extend(tools);
91 self
92 }
93
94 pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
96 self.tool_config = Some(tool_config);
97 self
98 }
99
100 pub fn with_ttl(mut self, ttl: Duration) -> Self {
103 self.expiration = Some(CacheExpirationRequest::from_ttl(ttl));
104 self
105 }
106
107 pub fn with_expire_time(mut self, expire_time: time::OffsetDateTime) -> Self {
109 self.expiration = Some(CacheExpirationRequest::from_expire_time(expire_time));
110 self
111 }
112
113 #[instrument(skip_all, fields(
115 display.name = self.display_name,
116 messages.count = self.contents.len(),
117 tools.count = self.tools.len(),
118 system_instruction.present = self.system_instruction.is_some(),
119 ))]
120 pub async fn execute(self) -> Result<CachedContentHandle, Error> {
121 let model = self.client.model.clone();
122 let expiration = self.expiration.ok_or(Error::MissingExpiration)?;
123
124 let cached_content = CreateCachedContentRequest {
125 display_name: self.display_name,
126 model,
127 contents: if self.contents.is_empty() { None } else { Some(self.contents) },
128 tools: if self.tools.is_empty() { None } else { Some(self.tools) },
129 system_instruction: self.system_instruction,
130 tool_config: self.tool_config,
131 expiration,
132 };
133
134 let response = self
135 .client
136 .create_cached_content(cached_content)
137 .await
138 .map_err(Box::new)
139 .context(ClientSnafu)?;
140
141 let cache_name = response.name;
142
143 Ok(CachedContentHandle::new(cache_name, self.client))
144 }
145}