Skip to main content

adk_tool/builtin/
load_artifacts.rs

1use adk_core::{AdkError, Result, Tool, ToolContext};
2use async_trait::async_trait;
3use base64::{Engine as _, engine::general_purpose::STANDARD};
4use serde_json::{Value, json};
5use std::sync::Arc;
6
7/// A tool that loads binary artifacts by name and returns their content.
8pub struct LoadArtifactsTool {
9    name: String,
10    description: String,
11}
12
13impl LoadArtifactsTool {
14    /// Create a new `LoadArtifactsTool` with default name and description.
15    pub fn new() -> Self {
16        Self {
17            name: "load_artifacts".to_string(),
18            description: "Loads artifacts by name and returns their content. Accepts an array of artifact names.".to_string(),
19        }
20    }
21}
22
23impl Default for LoadArtifactsTool {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29#[async_trait]
30impl Tool for LoadArtifactsTool {
31    fn name(&self) -> &str {
32        &self.name
33    }
34
35    fn description(&self) -> &str {
36        &self.description
37    }
38
39    fn is_long_running(&self) -> bool {
40        false
41    }
42
43    fn parameters_schema(&self) -> Option<Value> {
44        Some(json!({
45            "type": "object",
46            "properties": {
47                "artifact_names": {
48                    "type": "array",
49                    "items": {
50                        "type": "string"
51                    },
52                    "description": "List of artifact names to load"
53                }
54            },
55            "required": ["artifact_names"]
56        }))
57    }
58
59    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
60        let artifact_service =
61            ctx.artifacts().ok_or_else(|| AdkError::tool("ArtifactService not available"))?;
62
63        let artifact_names = args["artifact_names"]
64            .as_array()
65            .ok_or_else(|| AdkError::tool("artifact_names must be an array"))?;
66
67        let mut results = Vec::new();
68
69        for name_value in artifact_names {
70            let name = name_value
71                .as_str()
72                .ok_or_else(|| AdkError::tool("artifact name must be a string"))?;
73
74            match artifact_service.load(name).await {
75                Ok(part) => {
76                    let content = match part {
77                        adk_core::Part::Text { text } => json!({
78                            "type": "text",
79                            "text": text,
80                        }),
81                        adk_core::Part::InlineData { mime_type, data } => {
82                            // Base64 encode binary data for JSON transport
83                            let encoded = STANDARD.encode(&data);
84                            json!({
85                                "type": "inline_data",
86                                "mime_type": mime_type,
87                                "data_base64": encoded,
88                                "size_bytes": data.len(),
89                            })
90                        }
91                        _ => json!({ "type": "unknown" }),
92                    };
93
94                    results.push(json!({
95                        "name": name,
96                        "content": content,
97                    }));
98                }
99                Err(_) => {
100                    results.push(json!({
101                        "name": name,
102                        "error": "Artifact not found",
103                    }));
104                }
105            }
106        }
107
108        Ok(json!({
109            "artifacts": results,
110        }))
111    }
112}