gemini_adk_rs/llm/
mod.rs

1//! LLM abstraction — decouples agents from specific model providers.
2//!
3//! The `BaseLlm` trait provides a unified interface for generating content
4//! from any LLM. The `GeminiLlm` implementation wraps gemini-live's `Client`
5//! for Gemini models.
6
7pub mod gemini;
8mod mock;
9pub mod registry;
10
11pub use gemini::{GeminiLlm, GeminiLlmParams};
12pub use mock::MockLlm;
13pub use registry::LlmRegistry;
14
15use async_trait::async_trait;
16use futures_util::StreamExt;
17use serde::{Deserialize, Serialize};
18
19use gemini_genai_rs::prelude::{Content, Part, Tool};
20
21/// Provides access tokens for VertexAI authentication.
22///
23/// Implement this trait to supply dynamically refreshed tokens.
24/// The default implementation reads `GOOGLE_ACCESS_TOKEN` from the environment.
25pub trait TokenProvider: Send + Sync {
26    /// Return a valid access token. Called before each `generate()` request
27    /// when using VertexAI variant.
28    fn token(&self) -> String;
29}
30
31/// Default token provider — reads `GOOGLE_ACCESS_TOKEN` from environment.
32pub struct EnvTokenProvider;
33
34impl TokenProvider for EnvTokenProvider {
35    fn token(&self) -> String {
36        std::env::var("GOOGLE_ACCESS_TOKEN").unwrap_or_default()
37    }
38}
39
40/// Token provider that shells out to `gcloud auth print-access-token`,
41/// caching the result with a configurable TTL.
42pub struct GcloudTokenProvider {
43    cache: parking_lot::Mutex<(String, std::time::Instant)>,
44    ttl: std::time::Duration,
45}
46
47impl GcloudTokenProvider {
48    /// Create a new provider with the given cache TTL (recommended: 45 minutes).
49    pub fn new(ttl: std::time::Duration) -> Self {
50        Self {
51            cache: parking_lot::Mutex::new((String::new(), std::time::Instant::now())),
52            ttl,
53        }
54    }
55}
56
57impl TokenProvider for GcloudTokenProvider {
58    fn token(&self) -> String {
59        let mut guard = self.cache.lock();
60        let (ref mut cached_token, ref mut fetched_at) = *guard;
61        if !cached_token.is_empty() && fetched_at.elapsed() < self.ttl {
62            return cached_token.clone();
63        }
64        // Shell out to gcloud
65        match std::process::Command::new("gcloud")
66            .args(["auth", "print-access-token"])
67            .output()
68        {
69            Ok(output) if output.status.success() => {
70                let token = String::from_utf8_lossy(&output.stdout).trim().to_string();
71                *cached_token = token.clone();
72                *fetched_at = std::time::Instant::now();
73                token
74            }
75            _ => {
76                // Fall back to env var
77                std::env::var("GOOGLE_ACCESS_TOKEN").unwrap_or_default()
78            }
79        }
80    }
81}
82
83/// Configuration for an LLM generation request.
84///
85/// Every field a provider can honour is here, and [`GeminiLlm`] sends every
86/// one of them. New fields may be added; build requests with
87/// `..Default::default()` or [`LlmRequest::from_text`].
88#[derive(Debug, Clone, Default, Serialize, Deserialize)]
89pub struct LlmRequest {
90    /// The messages/contents to send.
91    pub contents: Vec<Content>,
92    /// The model for this request, overriding the provider's default model
93    /// (e.g. `"gemini-2.5-pro"`). `None` uses [`BaseLlm::model_id`].
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub model: Option<String>,
96    /// System instruction.
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub system_instruction: Option<String>,
99    /// Available tools.
100    #[serde(skip_serializing_if = "Vec::is_empty", default)]
101    pub tools: Vec<Tool>,
102    /// Temperature for generation.
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub temperature: Option<f32>,
105    /// Maximum output tokens.
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pub max_output_tokens: Option<u32>,
108    /// Nucleus sampling threshold.
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub top_p: Option<f32>,
111    /// Number of highest-probability tokens to sample from.
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub top_k: Option<u32>,
114    /// Strings that end generation when produced.
115    #[serde(skip_serializing_if = "Vec::is_empty", default)]
116    pub stop_sequences: Vec<String>,
117    /// Token budget for the model's thinking, on models that think.
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub thinking_budget: Option<u32>,
120    /// MIME type for structured output (e.g., `"application/json"`).
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub response_mime_type: Option<String>,
123    /// JSON Schema for structured output. Requires `response_mime_type = "application/json"`.
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub response_json_schema: Option<serde_json::Value>,
126}
127
128impl LlmRequest {
129    /// Create a request from a single user message.
130    pub fn from_text(text: impl Into<String>) -> Self {
131        Self {
132            contents: vec![Content {
133                role: Some(gemini_genai_rs::prelude::Role::User),
134                parts: vec![Part::Text { text: text.into() }],
135            }],
136            ..Default::default()
137        }
138    }
139
140    /// Create a request from existing contents.
141    pub fn from_contents(contents: Vec<Content>) -> Self {
142        Self {
143            contents,
144            ..Default::default()
145        }
146    }
147}
148
149/// The response from an LLM generation request.
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct LlmResponse {
152    /// The generated content.
153    pub content: Content,
154    /// Finish reason (if available).
155    #[serde(skip_serializing_if = "Option::is_none")]
156    pub finish_reason: Option<String>,
157    /// Token usage (if available).
158    #[serde(skip_serializing_if = "Option::is_none")]
159    pub usage: Option<TokenUsage>,
160}
161
162impl LlmResponse {
163    /// A model turn that says `text` and finishes.
164    pub fn from_text(text: impl Into<String>) -> Self {
165        Self::from_parts(vec![Part::Text { text: text.into() }], Some("STOP"))
166    }
167
168    /// A model turn that calls one tool.
169    pub fn tool_call(name: impl Into<String>, args: serde_json::Value) -> Self {
170        Self::tool_calls([(name, args)])
171    }
172
173    /// A model turn that calls several tools at once, in order.
174    pub fn tool_calls<N: Into<String>>(
175        calls: impl IntoIterator<Item = (N, serde_json::Value)>,
176    ) -> Self {
177        let parts = calls
178            .into_iter()
179            .map(|(name, args)| Part::FunctionCall {
180                function_call: gemini_genai_rs::prelude::FunctionCall {
181                    name: name.into(),
182                    args,
183                    id: None,
184                },
185            })
186            .collect();
187        Self::from_parts(parts, None)
188    }
189
190    fn from_parts(parts: Vec<Part>, finish_reason: Option<&str>) -> Self {
191        Self {
192            content: Content {
193                role: Some(gemini_genai_rs::prelude::Role::Model),
194                parts,
195            },
196            finish_reason: finish_reason.map(str::to_owned),
197            usage: None,
198        }
199    }
200
201    /// Add a streamed chunk to this response: its parts are appended (adjacent
202    /// text merged), and its finish reason and usage, when present, replace
203    /// these — a provider's streamed usage is cumulative.
204    pub fn append(&mut self, chunk: LlmResponse) {
205        for part in chunk.content.parts {
206            match (self.content.parts.last_mut(), part) {
207                (Some(Part::Text { text }), Part::Text { text: more }) => text.push_str(&more),
208                (_, part) => self.content.parts.push(part),
209            }
210        }
211        if chunk.finish_reason.is_some() {
212            self.finish_reason = chunk.finish_reason;
213        }
214        if chunk.usage.is_some() {
215            self.usage = chunk.usage;
216        }
217    }
218
219    /// Attach token usage, as a provider reports it.
220    pub fn with_usage(mut self, prompt_tokens: u32, completion_tokens: u32) -> Self {
221        self.usage = Some(TokenUsage::new(prompt_tokens, completion_tokens));
222        self
223    }
224
225    /// Extract text from the response, concatenating all text parts.
226    pub fn text(&self) -> String {
227        self.content
228            .parts
229            .iter()
230            .filter_map(|p| match p {
231                Part::Text { text } => Some(text.as_str()),
232                _ => None,
233            })
234            .collect::<Vec<_>>()
235            .join("")
236    }
237
238    /// Extract function calls from the response.
239    pub fn function_calls(&self) -> Vec<&gemini_genai_rs::prelude::FunctionCall> {
240        self.content
241            .parts
242            .iter()
243            .filter_map(|p| match p {
244                Part::FunctionCall { function_call } => Some(function_call),
245                _ => None,
246            })
247            .collect()
248    }
249}
250
251/// Token usage statistics.
252#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
253pub struct TokenUsage {
254    /// Input/prompt tokens.
255    pub prompt_tokens: u32,
256    /// Output/completion tokens.
257    pub completion_tokens: u32,
258    /// Total tokens.
259    pub total_tokens: u32,
260}
261
262impl TokenUsage {
263    /// Usage for one call; the total is the sum of the two.
264    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
265        Self {
266            prompt_tokens,
267            completion_tokens,
268            total_tokens: prompt_tokens.saturating_add(completion_tokens),
269        }
270    }
271}
272
273impl std::ops::Add for TokenUsage {
274    type Output = Self;
275
276    fn add(self, other: Self) -> Self {
277        Self {
278            prompt_tokens: self.prompt_tokens.saturating_add(other.prompt_tokens),
279            completion_tokens: self
280                .completion_tokens
281                .saturating_add(other.completion_tokens),
282            total_tokens: self.total_tokens.saturating_add(other.total_tokens),
283        }
284    }
285}
286
287impl std::ops::AddAssign for TokenUsage {
288    fn add_assign(&mut self, other: Self) {
289        *self = *self + other;
290    }
291}
292
293/// Errors from LLM operations.
294///
295/// Errors keep what the provider said — the HTTP status, the reason content
296/// was blocked — so a caller can decide what to do without parsing a message:
297///
298/// ```
299/// use gemini_adk_rs::llm::LlmError;
300///
301/// let err = LlmError::Api { status: 429, message: "quota exceeded".into() };
302/// assert!(err.is_rate_limited() && err.is_retryable());
303/// assert_eq!(err.status(), Some(429));
304/// ```
305#[derive(Debug, thiserror::Error)]
306#[non_exhaustive]
307pub enum LlmError {
308    /// The provider answered with an error status.
309    #[error("the model API returned HTTP {status}: {message}")]
310    Api {
311        /// The HTTP status code.
312        status: u16,
313        /// The provider's error message.
314        message: String,
315    },
316    /// Credentials are missing or were rejected. The message says how to fix it.
317    #[error("{0}")]
318    Auth(String),
319    /// The provider could not be reached: connection, TLS or timeout.
320    #[error("could not reach the model API: {0}")]
321    Transport(String),
322    /// The client is configured in a way no request can succeed with. The
323    /// message says what to change.
324    #[error("{0}")]
325    Config(String),
326    /// The HTTP request to the LLM API failed for another reason.
327    #[error("LLM request failed: {0}")]
328    RequestFailed(String),
329    /// The requested model is not available.
330    #[error("Model not available: {0}")]
331    ModelNotAvailable(String),
332    /// The request was rate-limited by the provider.
333    #[error("Rate limited")]
334    RateLimited,
335    /// The prompt or the reply was blocked by content safety; carries the
336    /// provider's reason (e.g. `"SAFETY"`, `"PROHIBITED_CONTENT"`).
337    #[error("blocked by content safety: {0}")]
338    ContentFiltered(String),
339    /// A catch-all for other LLM errors.
340    #[error("{0}")]
341    Other(String),
342}
343
344impl LlmError {
345    /// The HTTP status the provider answered with, when there was one.
346    pub fn status(&self) -> Option<u16> {
347        match self {
348            Self::Api { status, .. } => Some(*status),
349            Self::RateLimited => Some(429),
350            _ => None,
351        }
352    }
353
354    /// The provider is rate-limiting or out of quota (HTTP 429).
355    pub fn is_rate_limited(&self) -> bool {
356        self.status() == Some(429)
357    }
358
359    /// Credentials are missing, invalid or lack permission (HTTP 401/403).
360    pub fn is_auth(&self) -> bool {
361        matches!(self, Self::Auth(_)) || matches!(self.status(), Some(401 | 403))
362    }
363
364    /// The prompt or reply was blocked by content safety.
365    pub fn is_content_filtered(&self) -> bool {
366        matches!(self, Self::ContentFiltered(_))
367    }
368
369    /// Trying the same request again later may succeed: rate limits, server
370    /// errors (5xx) and transport failures. Auth, configuration, content
371    /// safety and other client errors will fail the same way again.
372    pub fn is_retryable(&self) -> bool {
373        matches!(self, Self::Transport(_)) || matches!(self.status(), Some(429 | 500..=599))
374    }
375}
376
377/// Capability declaration for a model — what callers may rely on without
378/// probing. See [`BaseLlm::capabilities`].
379#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
380pub struct ModelCapabilities {
381    /// Accepts a thinking budget / emits thought summaries.
382    pub thinking: bool,
383    /// Serves bidiGenerateContent (Live) sessions.
384    pub live_bidi: bool,
385    /// Produces native audio output.
386    pub audio_output: bool,
387    /// Accepts inline image/media parts in requests (vision).
388    pub vision_input: bool,
389    /// Honors cached-content references.
390    pub context_caching: bool,
391}
392
393impl ModelCapabilities {
394    /// Conservative inference from a model id string. Known family
395    /// substrings light up capabilities; unknown ids report text-only.
396    pub fn infer_from_id(id: &str) -> Self {
397        let id = id.to_ascii_lowercase();
398        let gemini = id.contains("gemini");
399        let live = id.contains("live") || id.contains("native-audio");
400        Self {
401            thinking: gemini && (id.contains("2.5") || id.contains("thinking")),
402            live_bidi: live,
403            audio_output: live || id.contains("tts"),
404            vision_input: gemini,
405            context_caching: gemini,
406        }
407    }
408}
409
410/// A reply streamed in chunks; see [`BaseLlm::generate_stream`].
411pub type LlmStream = futures_util::stream::BoxStream<'static, Result<LlmResponse, LlmError>>;
412
413/// Trait for LLM providers — decouples agents from specific models.
414///
415/// Implementations must be `Send + Sync` for use across async tasks.
416#[async_trait]
417pub trait BaseLlm: Send + Sync {
418    /// The model identifier (e.g., "gemini-2.5-flash").
419    fn model_id(&self) -> &str;
420
421    /// What this model supports — the ADK model-capability-declaration
422    /// pattern: the model states its capabilities instead of every caller
423    /// re-inferring them from the id string. The default derives a
424    /// conservative estimate from [`model_id`](Self::model_id); back-ends
425    /// with authoritative knowledge should override.
426    fn capabilities(&self) -> ModelCapabilities {
427        ModelCapabilities::infer_from_id(self.model_id())
428    }
429
430    /// Generate content from the LLM.
431    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError>;
432
433    /// Generate content as a stream of chunks, in the order the model
434    /// produced them. [`LlmResponse::append`] folds them into the whole reply;
435    /// usage, when a chunk carries it, is cumulative.
436    ///
437    /// The default yields [`generate`](Self::generate)'s reply as one chunk,
438    /// so every provider supports it; `GeminiLlm` streams for real.
439    async fn generate_stream(&self, request: LlmRequest) -> Result<LlmStream, LlmError> {
440        let response = self.generate(request).await?;
441        Ok(futures_util::stream::once(async move { Ok(response) }).boxed())
442    }
443
444    /// Pre-warm the HTTP connection pool to avoid cold-start latency.
445    ///
446    /// The default implementation is a no-op. `GeminiLlm` overrides this to
447    /// establish the TCP+TLS connection so the first real `generate()` call
448    /// doesn't pay the ~100-300ms handshake penalty.
449    async fn warm_up(&self) -> Result<(), LlmError> {
450        Ok(())
451    }
452}
453
454/// A shared model is a model, so `Arc<GeminiLlm>`, `Arc<dyn BaseLlm>` and a
455/// bare `GeminiLlm` are interchangeable wherever a [`BaseLlm`] is accepted.
456#[async_trait]
457impl<L: BaseLlm + ?Sized> BaseLlm for std::sync::Arc<L> {
458    fn model_id(&self) -> &str {
459        (**self).model_id()
460    }
461
462    fn capabilities(&self) -> ModelCapabilities {
463        (**self).capabilities()
464    }
465
466    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
467        (**self).generate(request).await
468    }
469
470    async fn generate_stream(&self, request: LlmRequest) -> Result<LlmStream, LlmError> {
471        (**self).generate_stream(request).await
472    }
473
474    async fn warm_up(&self) -> Result<(), LlmError> {
475        (**self).warm_up().await
476    }
477}
478
479/// A boxed model is a model; see the `Arc` implementation.
480#[async_trait]
481impl<L: BaseLlm + ?Sized> BaseLlm for Box<L> {
482    fn model_id(&self) -> &str {
483        (**self).model_id()
484    }
485
486    fn capabilities(&self) -> ModelCapabilities {
487        (**self).capabilities()
488    }
489
490    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
491        (**self).generate(request).await
492    }
493
494    async fn generate_stream(&self, request: LlmRequest) -> Result<LlmStream, LlmError> {
495        (**self).generate_stream(request).await
496    }
497
498    async fn warm_up(&self) -> Result<(), LlmError> {
499        (**self).warm_up().await
500    }
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506
507    #[test]
508    fn llm_request_from_text() {
509        let req = LlmRequest::from_text("Hello!");
510        assert_eq!(req.contents.len(), 1);
511        assert!(req.system_instruction.is_none());
512        assert!(req.tools.is_empty());
513    }
514
515    #[test]
516    fn llm_request_from_contents() {
517        let contents = vec![Content {
518            role: Some(gemini_genai_rs::prelude::Role::User),
519            parts: vec![Part::Text {
520                text: "Hello".into(),
521            }],
522        }];
523        let req = LlmRequest::from_contents(contents);
524        assert_eq!(req.contents.len(), 1);
525    }
526
527    #[test]
528    fn llm_response_text() {
529        let resp = LlmResponse {
530            content: Content {
531                role: Some(gemini_genai_rs::prelude::Role::Model),
532                parts: vec![
533                    Part::Text {
534                        text: "Hello ".into(),
535                    },
536                    Part::Text {
537                        text: "world!".into(),
538                    },
539                ],
540            },
541            finish_reason: Some("STOP".into()),
542            usage: None,
543        };
544        assert_eq!(resp.text(), "Hello world!");
545    }
546
547    #[test]
548    fn llm_response_function_calls() {
549        let resp = LlmResponse {
550            content: Content {
551                role: Some(gemini_genai_rs::prelude::Role::Model),
552                parts: vec![Part::FunctionCall {
553                    function_call: gemini_genai_rs::prelude::FunctionCall {
554                        name: "get_weather".into(),
555                        args: serde_json::json!({"city": "London"}),
556                        id: None,
557                    },
558                }],
559            },
560            finish_reason: None,
561            usage: None,
562        };
563        let calls = resp.function_calls();
564        assert_eq!(calls.len(), 1);
565        assert_eq!(calls[0].name, "get_weather");
566    }
567
568    #[test]
569    fn errors_classify_by_status_not_by_message() {
570        let api = |status| LlmError::Api {
571            status,
572            message: String::new(),
573        };
574        assert!(api(429).is_rate_limited() && api(429).is_retryable());
575        assert!(api(503).is_retryable() && !api(503).is_rate_limited());
576        assert!(api(401).is_auth() && api(403).is_auth() && !api(401).is_retryable());
577        assert!(!api(400).is_retryable() && !api(404).is_auth());
578        assert!(LlmError::Transport("reset".into()).is_retryable());
579        assert!(LlmError::RateLimited.is_rate_limited());
580        assert!(LlmError::Auth("no key".into()).is_auth());
581        assert!(LlmError::ContentFiltered("SAFETY".into()).is_content_filtered());
582        assert_eq!(api(418).status(), Some(418));
583        assert_eq!(LlmError::Other("x".into()).status(), None);
584    }
585
586    #[test]
587    fn append_folds_streamed_chunks_into_one_reply() {
588        let mut whole = LlmResponse::from_parts(vec![Part::Text { text: "Hel".into() }], None);
589        whole.append(LlmResponse::from_parts(
590            vec![Part::Text { text: "lo".into() }],
591            None,
592        ));
593        whole.append(LlmResponse::tool_call("f", serde_json::json!({})).with_usage(3, 1));
594        whole.append(LlmResponse::from_text("!").with_usage(3, 2));
595        assert_eq!(whole.text(), "Hello!");
596        assert_eq!(whole.function_calls().len(), 1);
597        assert_eq!(whole.finish_reason.as_deref(), Some("STOP"));
598        assert_eq!(
599            whole.usage,
600            Some(TokenUsage::new(3, 2)),
601            "usage is cumulative, not summed"
602        );
603    }
604
605    #[tokio::test]
606    async fn every_model_can_stream() {
607        let chunks: Vec<_> = MockLlm::text("one two three")
608            .generate_stream(LlmRequest::from_text("x"))
609            .await
610            .unwrap()
611            .map(Result::unwrap)
612            .collect()
613            .await;
614        assert_eq!(chunks.len(), 3, "the mock streams one word at a time");
615        let mut whole = chunks[0].clone();
616        for chunk in chunks.into_iter().skip(1) {
617            whole.append(chunk);
618        }
619        assert_eq!(whole.text(), "one two three");
620    }
621
622    #[test]
623    fn base_llm_is_object_safe() {
624        fn _assert(_: &dyn BaseLlm) {}
625    }
626
627    #[test]
628    fn token_usage() {
629        let usage = TokenUsage::new(10, 20);
630        assert_eq!(usage.total_tokens, 30);
631        assert_eq!((usage + TokenUsage::new(1, 2)).total_tokens, 33);
632    }
633
634    #[test]
635    fn response_constructors_build_model_turns() {
636        let said = LlmResponse::from_text("hi").with_usage(3, 4);
637        assert_eq!(said.text(), "hi");
638        assert_eq!(
639            said.content.role,
640            Some(gemini_genai_rs::prelude::Role::Model)
641        );
642        assert_eq!(said.finish_reason.as_deref(), Some("STOP"));
643        assert_eq!(said.usage, Some(TokenUsage::new(3, 4)));
644
645        let called = LlmResponse::tool_calls([
646            ("a", serde_json::json!({})),
647            ("b", serde_json::json!({ "x": 1 })),
648        ]);
649        let names: Vec<_> = called
650            .function_calls()
651            .iter()
652            .map(|c| c.name.as_str())
653            .collect();
654        assert_eq!(names, ["a", "b"]);
655        assert!(called.text().is_empty());
656    }
657
658    /// `Arc<L>`, `Arc<dyn BaseLlm>` and `Box<L>` all satisfy a generic bound,
659    /// so no call site needs to know which one it was handed.
660    #[tokio::test]
661    async fn shared_and_boxed_models_are_models() {
662        async fn ask(llm: impl BaseLlm) -> String {
663            llm.generate(LlmRequest::from_text("q"))
664                .await
665                .unwrap()
666                .text()
667        }
668        let mock = MockLlm::text("a").with_model_id("m");
669        let shared: std::sync::Arc<dyn BaseLlm> = std::sync::Arc::new(mock.clone());
670        assert_eq!(shared.model_id(), "m");
671        assert_eq!(ask(shared).await, "a");
672        assert_eq!(ask(std::sync::Arc::new(mock.clone())).await, "a");
673        assert_eq!(ask(Box::new(mock.clone()) as Box<dyn BaseLlm>).await, "a");
674        assert_eq!(mock.call_count(), 3);
675    }
676}
gemini_adk_rs/llm/mod.rs

gemini_adk_rs/llm/
mod.rs
