Skip to main content

adk_gemini/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!(chars <= 128, LongDisplayNameSnafu { display_name, chars });
48        self.display_name = Some(display_name);
49        Ok(self)
50    }
51
52    /// Set the system instruction for the cached content.
53    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    /// Add a user message to the cached content.
59    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    /// Add a model message to the cached content.
65    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    /// Add content directly to the cached content.
71    pub fn with_content(mut self, content: Content) -> Self {
72        self.contents.push(content);
73        self
74    }
75
76    /// Add multiple contents to the cached content.
77    pub fn with_contents(mut self, contents: Vec<Content>) -> Self {
78        self.contents.extend(contents);
79        self
80    }
81
82    /// Add a tool to the cached content.
83    pub fn with_tool(mut self, tool: Tool) -> Self {
84        self.tools.push(tool);
85        self
86    }
87
88    /// Add multiple tools to the cached content.
89    pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
90        self.tools.extend(tools);
91        self
92    }
93
94    /// Set the tool configuration.
95    pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
96        self.tool_config = Some(tool_config);
97        self
98    }
99
100    /// Set the TTL (Time To Live) for the cached content.
101    /// The cache will automatically expire after this duration.
102    pub fn with_ttl(mut self, ttl: Duration) -> Self {
103        self.expiration = Some(CacheExpirationRequest::from_ttl(ttl));
104        self
105    }
106
107    /// Set an explicit expiration time for the cached content.
108    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    /// Execute the cache creation request.
114    #[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}