adk-tool 0.8.0

Tool system for Rust Agent Development Kit (ADK-Rust) agents (FunctionTool, MCP, Google Search)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! Spanner toolset structure and `Toolset` trait implementation.

use crate::spanner::tools::{SpannerExecuteSql, SpannerGetTableSchema, SpannerListTables};
use adk_core::{ReadonlyContext, Result, Tool, Toolset};
use async_trait::async_trait;
use std::sync::Arc;

/// How Spanner credentials are resolved at runtime.
#[derive(Debug, Clone)]
pub enum CredentialSource {
    /// Use Google Cloud Application Default Credentials.
    ApplicationDefault,
    /// A secret name resolved via `ToolContext::get_secret()` at execution time.
    SecretRef(String),
}

/// Native Spanner toolset providing SQL execution, schema inspection, and
/// table listing tools.
///
/// Authenticates using Google Cloud Application Default Credentials or a
/// service account key provided via the configured
/// [`SecretProvider`](adk_core::ToolContext::get_secret).
///
/// # Example
///
/// ```rust,ignore
/// use adk_tool::spanner::SpannerToolset;
///
/// // Application Default Credentials
/// let toolset = SpannerToolset::new("my-project", "my-instance", "my-database");
///
/// // Via secret provider (resolved at tool execution time)
/// let toolset = SpannerToolset::from_secret(
///     "my-project",
///     "my-instance",
///     "my-database",
///     "spanner-credentials",
/// );
/// ```
pub struct SpannerToolset {
    pub(crate) project_id: String,
    pub(crate) instance_id: String,
    pub(crate) database_id: String,
    pub(crate) credentials: CredentialSource,
}

impl SpannerToolset {
    /// Create a new `SpannerToolset` using Application Default Credentials.
    ///
    /// Requires a project ID, instance ID, and database ID to identify the
    /// target Spanner database.
    pub fn new(
        project_id: impl Into<String>,
        instance_id: impl Into<String>,
        database_id: impl Into<String>,
    ) -> Self {
        Self {
            project_id: project_id.into(),
            instance_id: instance_id.into(),
            database_id: database_id.into(),
            credentials: CredentialSource::ApplicationDefault,
        }
    }

    /// Create a new `SpannerToolset` that resolves credentials from the
    /// secret provider at execution time via `ctx.get_secret(secret_name)`.
    pub fn from_secret(
        project_id: impl Into<String>,
        instance_id: impl Into<String>,
        database_id: impl Into<String>,
        secret_name: impl Into<String>,
    ) -> Self {
        Self {
            project_id: project_id.into(),
            instance_id: instance_id.into(),
            database_id: database_id.into(),
            credentials: CredentialSource::SecretRef(secret_name.into()),
        }
    }
}

#[async_trait]
impl Toolset for SpannerToolset {
    fn name(&self) -> &str {
        "spanner"
    }

    async fn tools(&self, _ctx: Arc<dyn ReadonlyContext>) -> Result<Vec<Arc<dyn Tool>>> {
        let project_id = self.project_id.clone();
        let instance_id = self.instance_id.clone();
        let database_id = self.database_id.clone();
        let credentials = self.credentials.clone();

        Ok(vec![
            Arc::new(SpannerExecuteSql::new(
                project_id.clone(),
                instance_id.clone(),
                database_id.clone(),
                credentials.clone(),
            )),
            Arc::new(SpannerGetTableSchema::new(
                project_id.clone(),
                instance_id.clone(),
                database_id.clone(),
                credentials.clone(),
            )),
            Arc::new(SpannerListTables::new(project_id, instance_id, database_id, credentials)),
        ])
    }
}