adk_tool/builtin/
load_artifacts.rs

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