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
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 =
59            ctx.artifacts().ok_or_else(|| AdkError::tool("ArtifactService not available"))?;
60
61        let artifact_names = args["artifact_names"]
62            .as_array()
63            .ok_or_else(|| AdkError::tool("artifact_names must be an array"))?;
64
65        let mut results = Vec::new();
66
67        for name_value in artifact_names {
68            let name = name_value
69                .as_str()
70                .ok_or_else(|| AdkError::tool("artifact name must be a string"))?;
71
72            match artifact_service.load(name).await {
73                Ok(part) => {
74                    let content = match part {
75                        adk_core::Part::Text { text } => json!({
76                            "type": "text",
77                            "text": text,
78                        }),
79                        adk_core::Part::InlineData { mime_type, data } => {
80                            // Base64 encode binary data for JSON transport
81                            let encoded = STANDARD.encode(&data);
82                            json!({
83                                "type": "inline_data",
84                                "mime_type": mime_type,
85                                "data_base64": encoded,
86                                "size_bytes": data.len(),
87                            })
88                        }
89                        _ => json!({ "type": "unknown" }),
90                    };
91
92                    results.push(json!({
93                        "name": name,
94                        "content": content,
95                    }));
96                }
97                Err(_) => {
98                    results.push(json!({
99                        "name": name,
100                        "error": "Artifact not found",
101                    }));
102                }
103            }
104        }
105
106        Ok(json!({
107            "artifacts": results,
108        }))
109    }
110}