Improve LLM

server-rs/src/utils/llm.rs

@@ -4,68 +4,62 @@ use tracing::warn;
 
 pub type LlmError = (StatusCode, String);
 
-/// Send a chat request to Ollama and return the parsed JSON response.
+/// Send a generateContent request to the Gemini API and return the parsed JSON response.
 ///
 /// Handles connection errors, non-success status codes, and JSON parse failures
 /// uniformly as `BAD_GATEWAY` errors.
-pub async fn ollama_chat(
+pub async fn gemini_chat(
     client: &reqwest::Client,
-    url: &str,
+    api_key: &str,
+    model: &str,
     body: &Value,
 ) -> Result<Value, LlmError> {
-    let response = client.post(url).json(body).send().await.map_err(|err| {
-        warn!(error = %err, "Failed to connect to Ollama");
+    let url = format!(
+        "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent?key={}",
+        model, api_key
+    );
+
+    let response = client.post(&url).json(body).send().await.map_err(|err| {
+        warn!(error = %err, "Failed to connect to Gemini API");
         (
             StatusCode::BAD_GATEWAY,
-            format!("Failed to connect to Ollama: {}", err),
+            format!("Failed to connect to Gemini API: {}", err),
         )
     })?;
 
     if !response.status().is_success() {
         let status = response.status();
         let body_text = response.text().await.unwrap_or_default();
-        warn!(status = %status, body = %body_text, "Ollama returned error");
+        warn!(status = %status, body = %body_text, "Gemini API returned error");
         return Err((
             StatusCode::BAD_GATEWAY,
-            format!("Ollama error {}: {}", status, body_text),
+            format!("Gemini API error {}: {}", status, body_text),
         ));
     }
 
     response.json().await.map_err(|err| {
-        warn!(error = %err, "Failed to parse Ollama response");
+        warn!(error = %err, "Failed to parse Gemini API response");
         (
             StatusCode::BAD_GATEWAY,
-            format!("Failed to parse Ollama response: {}", err),
+            format!("Failed to parse Gemini API response: {}", err),
         )
     })
 }
 
-/// Extract content from Ollama native response (`message.content`)
-pub fn extract_ollama_content(json: &Value) -> Result<&str, LlmError> {
-    json.get("message")
-        .and_then(|msg| msg.get("content"))
-        .and_then(|ct| ct.as_str())
+/// Extract text content from Gemini response (`candidates[0].content.parts[0].text`)
+pub fn extract_gemini_content(json: &Value) -> Result<&str, LlmError> {
+    json.get("candidates")
+        .and_then(|c| c.get(0))
+        .and_then(|c| c.get("content"))
+        .and_then(|c| c.get("parts"))
+        .and_then(|p| p.get(0))
+        .and_then(|p| p.get("text"))
+        .and_then(|t| t.as_str())
         .ok_or_else(|| {
-            warn!("Malformed Ollama response: missing message.content");
+            warn!("Malformed Gemini response: missing candidates[0].content.parts[0].text");
             (
                 StatusCode::BAD_GATEWAY,
-                "Malformed LLM response: missing message.content".into(),
+                "Malformed Gemini response: missing content".into(),
             )
         })
 }
-
-/// Strip `<think>...</think>` blocks from model output
-pub fn strip_think_blocks(text: &str) -> String {
-    let mut result = String::new();
-    let mut remaining = text;
-    while let Some(start) = remaining.find("<think>") {
-        result.push_str(&remaining[..start]);
-        if let Some(end) = remaining[start..].find("</think>") {
-            remaining = &remaining[start + end + 8..];
-        } else {
-            return result;
-        }
-    }
-    result.push_str(remaining);
-    result
-}