adk-rs 0.6.0

Rust port of the Google Agent Development Kit (ADK).
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

//! The [`Model`] trait abstracts an LLM provider.

use std::collections::BTreeMap;
use std::sync::Arc;

use async_trait::async_trait;

use crate::core::llm_request::LlmRequest;
use crate::core::llm_response::LlmResponse;
use crate::core::stream::LlmResponseStream;
use crate::error::Result;

/// A unified LLM-provider abstraction.
///
/// Implementations live in `adk-providers-*` crates.
#[async_trait]
pub trait Model: Send + Sync + std::fmt::Debug + 'static {
    /// The canonical name of this model instance (e.g. `"gemini-2.5-flash"`).
    fn name(&self) -> &str;

    /// Glob-like name patterns this provider can serve. Used by
    /// [`ModelRegistry`] to dispatch requests by model name.
    fn supported_models(&self) -> &'static [&'static str];

    /// Single-shot generation.
    async fn generate_content(&self, req: LlmRequest) -> Result<LlmResponse>;

    /// Streaming generation. The default implementation yields the single
    /// non-streaming response wrapped in a one-element stream.
    async fn stream_generate_content(&self, req: LlmRequest) -> Result<LlmResponseStream> {
        use futures::stream;
        let resp = self.generate_content(req).await?;
        Ok(Box::pin(stream::once(async move { Ok(resp) })))
    }
}

/// A registry mapping model-name patterns → provider instances.
///
/// Lookups walk patterns in insertion order; the first pattern whose
/// glob matches the requested model name wins. There is no global state.
#[derive(Debug, Default, Clone)]
pub struct ModelRegistry {
    by_name: BTreeMap<String, Arc<dyn Model>>,
    patterns: Vec<(globset::GlobMatcher, Arc<dyn Model>)>,
}

impl ModelRegistry {
    /// Construct an empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register an instance under all of its `supported_models` patterns.
    pub fn register(&mut self, model: Arc<dyn Model>) {
        let patterns = model.supported_models();
        for &p in patterns {
            // Each pattern is a glob (e.g. `"gemini-*"`) or an exact name.
            let glob = match globset::Glob::new(p) {
                Ok(g) => g.compile_matcher(),
                Err(_) => continue,
            };
            self.patterns.push((glob, model.clone()));
        }
        self.by_name.insert(model.name().to_string(), model);
    }

    /// Look up a model by name.
    pub fn get(&self, name: &str) -> Option<Arc<dyn Model>> {
        if let Some(m) = self.by_name.get(name) {
            return Some(m.clone());
        }
        for (g, m) in &self.patterns {
            if g.is_match(name) {
                return Some(m.clone());
            }
        }
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::testing::MockModel;

    #[test]
    fn registry_dispatches_exact_match() {
        let mut r = ModelRegistry::new();
        let m = Arc::new(MockModel::new("test")) as Arc<dyn Model>;
        r.register(m.clone());
        assert!(r.get("test").is_some());
        assert!(r.get("missing").is_none());
    }
}