Skip to main content

caliban_provider_google/
lib.rs

1//! Google Gemini schema family for the caliban agent harness.
2//!
3//! Provides [`GoogleProvider<T: Transport>`] generic over its transport.
4//! The AI Studio transport is always available. The Vertex AI transport
5//! is gated behind the `vertex` cargo feature (added in B.10).
6
7#![allow(clippy::missing_errors_doc)]
8// Transitive dependencies pull in multiple versions of some crates.
9#![allow(clippy::multiple_crate_versions)]
10
11pub mod config;
12pub mod error;
13pub mod ir_convert;
14pub mod models;
15pub mod schema;
16pub mod transport;
17
18mod stream_parse;
19
20use async_trait::async_trait;
21use caliban_provider::{
22    Capabilities, CompletionRequest, CompletionResponse, Error, MessageStream, ModelInfo, Provider,
23    Result,
24};
25
26use crate::config::AIStudioConfig;
27use crate::transport::Transport;
28use crate::transport::ai_studio::AIStudioTransport;
29
30/// Google Gemini provider, generic over its transport.
31pub struct GoogleProvider<T: Transport> {
32    transport: T,
33}
34
35impl<T: Transport> std::fmt::Debug for GoogleProvider<T> {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("GoogleProvider").finish_non_exhaustive()
38    }
39}
40
41impl GoogleProvider<AIStudioTransport> {
42    /// Construct a `GoogleProvider` using the Google AI Studio HTTPS transport.
43    ///
44    /// # Errors
45    ///
46    /// Returns `Err` if the underlying `reqwest` client cannot be built.
47    pub fn ai_studio(cfg: AIStudioConfig) -> Result<Self> {
48        AIStudioTransport::new(cfg)
49            .map(|t| Self { transport: t })
50            .map_err(Error::adapter)
51    }
52}
53
54impl<T: Transport> GoogleProvider<T> {
55    /// Construct a `GoogleProvider` from an arbitrary `Transport`.
56    pub fn from_transport(transport: T) -> Self {
57        Self { transport }
58    }
59}
60
61#[cfg(feature = "vertex")]
62impl GoogleProvider<crate::transport::vertex::VertexTransport> {
63    /// Construct a `GoogleProvider` using the Google Vertex AI transport.
64    ///
65    /// # Errors
66    ///
67    /// Returns `Err` if the underlying `reqwest` client cannot be built.
68    pub fn vertex(cfg: crate::config::VertexConfig) -> Result<Self> {
69        crate::transport::vertex::VertexTransport::new(cfg)
70            .map(|t| Self { transport: t })
71            .map_err(Error::adapter)
72    }
73}
74
75#[async_trait]
76impl<T: Transport> Provider for GoogleProvider<T> {
77    async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse> {
78        req.validate()?;
79        let canonical_model = req.model.clone();
80        let wire_model = self.transport.wire_model_id(&canonical_model);
81        let allow_urls = self.transport.supports_url_images();
82        let native = ir_convert::ir_to_native_request(req, allow_urls)?;
83        let native_resp = self
84            .transport
85            .send(&wire_model, &native)
86            .await
87            .map_err(Error::from)?;
88        ir_convert::native_response_to_ir(native_resp)
89    }
90
91    async fn stream(&self, req: CompletionRequest) -> Result<MessageStream> {
92        req.validate()?;
93        let canonical_model = req.model.clone();
94        let wire_model = self.transport.wire_model_id(&canonical_model);
95        let allow_urls = self.transport.supports_url_images();
96        let native = ir_convert::ir_to_native_request(req, allow_urls)?;
97        let bytes_stream = self
98            .transport
99            .stream(&wire_model, &native)
100            .await
101            .map_err(Error::from)?;
102        Ok(stream_parse::map_gemini_sse_to_events(bytes_stream))
103    }
104
105    fn capabilities(&self, model: &str) -> Capabilities {
106        models::capabilities_for(model)
107    }
108
109    fn list_models(&self) -> Vec<ModelInfo> {
110        models::models()
111    }
112
113    fn name(&self) -> &'static str {
114        "google"
115    }
116}