nutrition_ai/
lib.rs

1mod types;
2use anyhow::{Result, anyhow};
3use google_generative_ai_rs::v1::api::{Client, PostResult};
4use google_generative_ai_rs::v1::gemini::request::{InlineData, Request};
5use google_generative_ai_rs::v1::gemini::{Content, Model, Part, Role};
6use tokio::time::{Duration, sleep};
7pub use types::{GeminiRequest, MimeType};
8
9/// Analyze a food image using Gemini API and return nutrition markdown string
10pub async fn generate_answer(req: GeminiRequest) -> Result<String, anyhow::Error> {
11    // Validate base64 input (rough check)
12    if req.file_base64.trim().is_empty() {
13        return Err(anyhow!("File base64 string is empty"));
14    }
15
16    // Validate MIME type
17    let mime = req.file_mime_type.as_str();
18    if mime != "image/png" && mime != "image/jpeg" {
19        return Err(anyhow!(
20            "Unsupported MIME type: '{}'. Must be 'image/png' or 'image/jpeg'",
21            mime
22        ));
23    }
24
25    let prompt = include_str!("prompts/prompt.txt").to_string();
26
27    // Determine model
28    let model = match req.model.as_deref() {
29        Some("Gemini1_0Pro") => Model::Gemini1_0Pro,
30        Some("Gemini1_5Pro") => Model::Gemini1_5Pro,
31        Some("Gemini1_5Flash") => Model::Gemini1_5Flash,
32        Some("Gemini1_5Flash8B") => Model::Gemini1_5Flash8B,
33        Some("Gemini2_0Flash") => Model::Gemini2_0Flash,
34        Some(custom) => Model::Custom(custom.to_string()),
35        None => Model::Gemini1_5Flash,
36    };
37
38    // Check Google Key
39    if req.google_key.trim().is_empty() {
40        return Err(anyhow!("Google API key is empty"));
41    }
42
43    // Initialize client
44    let client = Client::new_from_model(model, req.google_key.clone());
45
46    let parts = vec![
47        Part {
48            text: Some(prompt),
49            inline_data: None,
50            file_data: None,
51            video_metadata: None,
52        },
53        Part {
54            text: None,
55            inline_data: Some(InlineData {
56                mime_type: mime.to_string(),
57                data: req.file_base64.clone(),
58            }),
59            file_data: None,
60            video_metadata: None,
61        },
62    ];
63
64    let contents = vec![Content {
65        role: Role::User,
66        parts,
67    }];
68
69    let request = Request {
70        contents,
71        tools: vec![],
72        safety_settings: vec![],
73        generation_config: None,
74        system_instruction: None,
75    };
76
77    let mut retries = 0;
78    let result: PostResult;
79
80    loop {
81        match client.post(30, &request).await {
82            Ok(res) => {
83                result = res;
84                break;
85            }
86            Err(e) => {
87                let err_msg = e.to_string();
88                if err_msg.contains("503") || err_msg.contains("overloaded") {
89                    retries += 1;
90                    if retries >= 3 {
91                        return Err(anyhow!(
92                            "Gemini model is overloaded after 3 attempts. Please try again later."
93                        ));
94                    } else {
95                        let delay = Duration::from_secs(2 * retries); // 2s, 4s, 6s
96                        eprintln!("Model overloaded. Retrying in {:?}...", delay);
97                        sleep(delay).await;
98                        continue;
99                    }
100                } else {
101                    return Err(anyhow!("Gemini API error: {}", err_msg));
102                }
103            }
104        }
105    }
106
107    if let Some(rest) = result.rest() {
108        if let Some(candidate) = rest.candidates.get(0) {
109            if let Some(part) = candidate.content.parts.get(0) {
110                if let Some(text) = &part.text {
111                    return Ok(text.clone());
112                } else {
113                    return Err(anyhow!("No text found in candidate part"));
114                }
115            } else {
116                return Err(anyhow!("No parts in candidate"));
117            }
118        } else {
119            return Err(anyhow!("No candidates returned by Gemini API"));
120        }
121    } else {
122        return Err(anyhow!("Unexpected response type from Gemini API"));
123    }
124}
nutrition_ai/lib.rs

nutrition_ai/
lib.rs
