smartui/
gemini.rs

1use serde::{Deserialize, Serialize};
2use std::env;
3use crate::error::{GmError, Result};
4use async_trait::async_trait;
5
6#[derive(Debug, Serialize, Deserialize)]
7struct GeminiRequest {
8    contents: Vec<Content>,
9}
10
11#[derive(Debug, Serialize, Deserialize)]
12struct Content {
13    parts: Vec<Part>,
14}
15
16#[derive(Debug, Serialize, Deserialize)]
17struct Part {
18    text: String,
19}
20
21#[derive(Debug, Deserialize)]
22struct GeminiResponse {
23    candidates: Vec<Candidate>,
24}
25
26#[derive(Debug, Deserialize)]
27struct Candidate {
28    content: Content,
29}
30
31#[async_trait]
32pub trait GeminiClientTrait {
33    async fn generate_content(&self, prompt: String) -> Result<String>;
34}
35
36pub struct RealGeminiClient {
37    client: reqwest::Client,
38    api_key: String,
39    base_url: String,
40}
41
42impl RealGeminiClient {
43    pub fn new() -> Result<Self> {
44        let api_key = env::var("GEMINI_API_KEY")
45            .map_err(|_| GmError::ApiKeyNotFound)?;
46
47        Ok(Self {
48            client: reqwest::Client::new(),
49            api_key,
50            base_url: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent".to_string(),
51        })
52    }
53}
54
55#[async_trait]
56impl GeminiClientTrait for RealGeminiClient {
57    async fn generate_content(&self, prompt: String) -> Result<String> {
58        let request = GeminiRequest {
59            contents: vec![Content {
60                parts: vec![Part {
61                    text: prompt,
62                }],
63            }],
64        };
65
66        let response = self.client
67            .post(&self.base_url)
68            .header("Content-Type", "application/json")
69            .query(&[("key", &self.api_key)])
70            .json(&request)
71            .send()
72            .await
73            .map_err(GmError::GeminiApi)?;
74
75        if !response.status().is_success() {
76            return Err(GmError::InvalidResponse(format!(
77                "API returned status code: {}",
78                response.status()
79            )));
80        }
81
82        let response_body: GeminiResponse = response
83            .json()
84            .await
85            .map_err(GmError::GeminiApi)?;
86
87        if let Some(candidate) = response_body.candidates.first() {
88            if let Some(part) = candidate.content.parts.first() {
89                return Ok(part.text.clone());
90            }
91        }
92
93        Err(GmError::InvalidResponse("No content in response".to_string()))
94    }
95}
96
97#[async_trait]
98impl GeminiClientTrait for &RealGeminiClient {
99    async fn generate_content(&self, prompt: String) -> Result<String> {
100        (*self).generate_content(prompt).await
101    }
102}