From 4fcf2795f3d9a0c5a89fcf0608a5602650979975 Mon Sep 17 00:00:00 2001
From: Monius <m1105930979@gmail.com>
Date: Thu, 16 May 2024 01:21:00 +0800
Subject: [PATCH] v0.1.10 apply request

---
 rim-cli/Cargo.toml    |  2 +-
 rim/src/client.rs     | 18 ++++++++++++++++--
 rim/src/lib.rs        | 41 ++---------------------------------------
 rim/src/llm/google.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+), 42 deletions(-)

diff --git a/rim-cli/Cargo.toml b/rim-cli/Cargo.toml
index 2d6b76b..2d73d4a 100644
--- a/rim-cli/Cargo.toml
+++ b/rim-cli/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rim-cli"
-version = "0.1.9"
+version = "0.1.10"
 edition = "2021"
 
 [dependencies]
diff --git a/rim/src/client.rs b/rim/src/client.rs
index 02df1d1..81d8c15 100644
--- a/rim/src/client.rs
+++ b/rim/src/client.rs
@@ -18,8 +18,22 @@ impl RimClient {
         Self::new(model)
     }
 
-    pub async fn generate_caption(&self, content: String) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
-        Ok("GPT4V Caption".to_string())
+    pub async fn generate_caption(&self, data: String) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
+        let api = self.model.get_api();
+        let payload = self.model.payload(data);
+
+        let client = reqwest::Client::builder()
+            .pool_idle_timeout(tokio::time::Duration::from_secs(1))
+            .build()?;
+        let response = client.post(api)
+            .json(&payload)
+            .send()
+            .await?
+            .json::<serde_json::Value>()
+            .await?;
+
+        let raw = &response["candidates"][0]["content"]["parts"][0]["text"];
+        Ok(raw.to_string())
     }
 
     pub fn log_prompt(&self) {
diff --git a/rim/src/lib.rs b/rim/src/lib.rs
index e3006c5..124dce3 100644
--- a/rim/src/lib.rs
+++ b/rim/src/lib.rs
@@ -1,45 +1,8 @@
 pub mod llm;
 pub mod client;
 pub mod modality;
-use futures::StreamExt;
-
-// pub fn single_cap(f: &str) {
-//     let start_time = std::time::Instant::now();
-//     println!("Processing file: {:?}", f);
-
-//     if let Ok(x) = _rt() {
-//         x.block_on(async {
-//             let _b64 = modality::image::async_base64(f.into()).await;
-//         });
-//     }
-
-//     let elapsed_time = start_time.elapsed();
-//     println!("Processing time: {:?}", elapsed_time);
-// }
-
-// pub fn batch_cap(d: &str) {
-//     let start_time = std::time::Instant::now();
-//     println!("Processing directory: {:?}", d);
-
-//     let rt = tokio::runtime::Runtime::new().unwrap();
-
-//     rt.block_on(async {
-//         let mut tasks: Vec<_> = std::fs::read_dir(d)
-//             .unwrap()
-//             .filter_map(Result::ok)
-//             .map(|entry| entry.path())
-//             .filter(|path| path.extension().unwrap_or_default() == "png")
-//             .map(|f| tokio::spawn(async move { modality::image::async_base64_log(f).await; }))
-//             .collect();
-//         for task in tasks {
-//             task.await.unwrap();
-//         }
-//     });
-
-//     let elapsed_time = start_time.elapsed();
-//     println!("Processing time: {:?}", elapsed_time);
-// }
 
+use futures::StreamExt;
 
 async fn caption(img: modality::Image, clt: &client::RimClient) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     let _b64 = img._base64().await?;
@@ -82,7 +45,7 @@ pub fn _rt(pth: &str, keys: Vec<String>, prompt: String) -> Result<(), Box<dyn s
         clients.push(_client);
     }
 
-    let mut images: Vec<modality::Image> = std::fs::read_dir(pth)
+    let images: Vec<modality::Image> = std::fs::read_dir(pth)
         .unwrap()
         .filter_map(Result::ok)
         .map(|entry| entry.path().display().to_string())
diff --git a/rim/src/llm/google.rs b/rim/src/llm/google.rs
index d37f078..70dcf3a 100644
--- a/rim/src/llm/google.rs
+++ b/rim/src/llm/google.rs
@@ -1,6 +1,8 @@
 const GEMINI_FLASH: &str = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=";
 const GEMINI_PRO: &str = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=";
 
+use serde_json::json;
+
 #[derive(Debug)]
 pub struct Gemini {
     prompt: String,
@@ -25,4 +27,44 @@ impl Gemini {
     pub fn get_api(&self) -> String{
         self.api.clone()
     }
+
+    pub fn payload(&self, data: String) -> serde_json::Value{
+        let payload = json!({
+            "contents": [
+                {   
+                    "role": "user",
+                    "parts": [
+                        {"text": self.prompt.clone()},
+                        {"inlineData": { "mimeType": "image/png", "data": data } },
+                    ]
+                }
+            ],
+            "generationConfig": {
+                "temperature": 1,
+                "topK": 64,
+                "topP": 0.95,
+                "maxOutputTokens": 8192,
+                "stopSequences": []
+            },
+            "safetySettings": [
+                {
+                "category": "HARM_CATEGORY_HARASSMENT",
+                "threshold": "BLOCK_NONE"
+                },
+                {
+                "category": "HARM_CATEGORY_HATE_SPEECH",
+                "threshold": "BLOCK_NONE"
+                },
+                {
+                "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+                "threshold": "BLOCK_NONE"
+                },
+                {
+                "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+                "threshold": "BLOCK_NONE"
+                }
+            ]
+        });
+        payload
+    }
 }
\ No newline at end of file