v0.1.10 apply request

AUTOM77 · May 15, 2024 · 4fcf279 · 4fcf279
1 parent 8e420a9
commit 4fcf279
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 42 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
+    }
 }