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}