caliban_provider_google/
lib.rs1#![allow(clippy::missing_errors_doc)]
8#![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
30pub 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 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 pub fn from_transport(transport: T) -> Self {
57 Self { transport }
58 }
59}
60
61#[cfg(feature = "vertex")]
62impl GoogleProvider<crate::transport::vertex::VertexTransport> {
63 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}