Skip to main content

chat_gemini/api/
completion.rs

1use crate::api::types::error::handle_gemini_error;
2use crate::api::types::request::GeminiRequest;
3use crate::api::types::response::GeminiCompletionResponse;
4use crate::client::GeminiClient;
5use chat_core::error::{ChatError, ChatFailure};
6use chat_core::traits::CompletionProvider;
7use chat_core::transport::Transport;
8use chat_core::types::provider_meta::ProviderMeta;
9use chat_core::types::messages::Messages;
10use chat_core::types::options::ChatOptions;
11use chat_core::types::response::ChatResponse;
12use chat_core::types::tools::ToolDeclarations;
13
14#[async_trait::async_trait]
15impl<T: Transport> CompletionProvider for GeminiClient<T> {
16    async fn complete(
17        &mut self,
18        messages: &mut Messages,
19        tool_declarations: Option<&dyn ToolDeclarations>,
20        options: Option<&ChatOptions>,
21        structured_output: Option<&schemars::Schema>,
22    ) -> Result<ChatResponse, ChatFailure> {
23        let path = format!(
24            "{}/models/{}:generateContent",
25            self.base_path, self.model_name
26        );
27
28        let request_body = GeminiRequest::from_core(
29            messages,
30            tool_declarations,
31            Some(self.native_tools.as_slice()),
32            self.function_config.as_ref(),
33            options,
34            structured_output,
35            self.include_thoughts,
36        )
37        .map_err(ChatFailure::from_err)?;
38
39        let body = serde_json::to_vec(&request_body)
40            .map_err(|e| ChatFailure::from_err(ChatError::InvalidResponse(e.to_string())))?;
41
42        let req = chat_core::transport::Request {
43            scheme: self.scheme.clone(),
44            host: self.host.clone(),
45            path,
46            headers: vec![
47                ("x-goog-api-key".into(), self.api_key.clone()),
48                ("Content-Type".into(), "application/json".into()),
49            ],
50            body,
51        };
52
53        let res = self
54            .transport
55            .send(req)
56            .await
57            .map_err(ChatFailure::from_err)?;
58
59        let res = handle_gemini_error(res)?;
60
61        let gemini_data: GeminiCompletionResponse = serde_json::from_slice(&res.body)
62            .map_err(|e| ChatFailure::from_err(ChatError::InvalidResponse(e.to_string())))?;
63
64        gemini_data
65            .into_core_chat_response()
66            .map_err(ChatFailure::from_err)
67    }
68
69    fn metadata(&self) -> Option<&ProviderMeta> {
70        Some(&self.meta)
71    }
72}