use anyhow::{anyhow, Result}; use serde::{Deserialize, Serialize}; use std::time::Duration; #[derive(Serialize)] struct Msg { role: String, content: String, } #[derive(Serialize)] struct Req { model: String, messages: Vec, temperature: f32, max_tokens: i32, #[serde(skip_serializing_if = "Option::is_none")] reasoning_budget: Option, #[serde(skip_serializing_if = "Option::is_none")] stream: Option, } #[derive(Deserialize)] struct Choice { message: ChoiceMsg, } #[derive(Deserialize)] struct ChoiceMsg { #[serde(default)] content: String, #[serde(default)] reasoning_content: String, } #[derive(Deserialize)] struct Resp { choices: Vec, } pub struct Llm { pub base_url: String, pub model: String, pub client: reqwest::Client, } impl Llm { pub fn new(base_url: String, model: String) -> Self { let client = reqwest::Client::builder() .timeout(Duration::from_secs(600)) .build() .unwrap(); Self { base_url, model, client } } pub async fn chat(&self, system: &str, user: &str, max_tokens: i32, temperature: f32) -> Result { self.chat_with_budget(system, user, max_tokens, temperature, Some(256)).await } pub async fn chat_full(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, budget: Option) -> Result<(String, String)> { let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/')); let budgets = [budget]; let mut last_err = None; for budget in budgets { let body = Req { model: self.model.clone(), messages: vec![ Msg { role: "system".into(), content: system.into() }, Msg { role: "user".into(), content: user.into() }, ], temperature, max_tokens, reasoning_budget: budget, stream: None, }; for attempt in 0..2 { match self.client.post(&url).json(&body).send().await { Ok(resp) => { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); if !status.is_success() { last_err = Some(anyhow!("HTTP {}: {}", status, text.chars().take(500).collect::())); tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await; continue; } match serde_json::from_str::(&text) { Ok(r) => { if let Some(c) = r.choices.into_iter().next() { let content = c.message.content.clone(); let reasoning = c.message.reasoning_content.clone(); // strip thinking from content let mut clean = content.clone(); if let Some(idx) = clean.rfind("") { clean = clean[idx + "".len()..].trim().to_string(); } // If content is empty, use reasoning as the answer // (the model put its entire response in thinking) if clean.trim().is_empty() && !reasoning.trim().is_empty() { let mut from_rc = reasoning.clone(); if let Some(idx) = from_rc.rfind("") { from_rc = from_rc[idx + "".len()..].trim().to_string(); } if !from_rc.trim().is_empty() { // Use reasoning as content, but don't duplicate in thinking return Ok((from_rc.clone(), String::new())); } } if !clean.trim().is_empty() { return Ok((clean, reasoning)); } last_err = Some(anyhow!("empty response from model")); } else { last_err = Some(anyhow!("no choices in response")); } } Err(e) => { last_err = Some(anyhow!("parse error: {} body={}", e, text.chars().take(300).collect::())); } } } Err(e) => { last_err = Some(anyhow!("request error: {}", e)); } } tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await; } } Err(last_err.unwrap_or_else(|| anyhow!("unknown error"))) } pub async fn chat_with_budget(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, budget: Option) -> Result { let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/')); let budgets = [budget]; let mut last_err = None; for budget in budgets { let body = Req { model: self.model.clone(), messages: vec![ Msg { role: "system".into(), content: system.into() }, Msg { role: "user".into(), content: user.into() }, ], temperature, max_tokens, reasoning_budget: budget, stream: None, }; for attempt in 0..2 { match self.client.post(&url).json(&body).send().await { Ok(resp) => { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); if !status.is_success() { last_err = Some(anyhow!("HTTP {}: {}", status, text.chars().take(500).collect::())); tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await; continue; } match serde_json::from_str::(&text) { Ok(r) => { if let Some(c) = r.choices.into_iter().next() { let mut content = c.message.content.clone(); // strip ... blocks if let Some(idx) = content.rfind("") { content = content[idx + "".len()..].trim().to_string(); } // Only use reasoning_content if it contains a real answer (bullet or NONE) if content.trim().is_empty() { let rc = c.message.reasoning_content.clone(); let mut from_rc = rc.clone(); if let Some(idx) = from_rc.rfind("") { from_rc = from_rc[idx + "".len()..].trim().to_string(); } if from_rc.contains("- ") || from_rc.trim().eq_ignore_ascii_case("NONE") { content = from_rc; } } if !content.trim().is_empty() { return Ok(content); } last_err = Some(anyhow!("empty response from model")); } else { last_err = Some(anyhow!("no choices in response")); } } Err(e) => { last_err = Some(anyhow!("parse error: {} body={}", e, text.chars().take(300).collect::())); } } } Err(e) => { last_err = Some(anyhow!("request error: {}", e)); } } tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await; } } Err(last_err.unwrap_or_else(|| anyhow!("unknown error"))) } /// Streaming chat. Calls `on_chunk` with each text delta as it arrives. pub async fn chat_stream(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, on_chunk: F) -> Result where F: FnMut(&str), { self.chat_stream_with_budget(system, user, max_tokens, temperature, Some(256), on_chunk).await } pub async fn chat_stream_with_budget(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, budget: Option, mut on_chunk: F) -> Result where F: FnMut(&str), { use futures_util::StreamExt; let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/')); let body = Req { model: self.model.clone(), messages: vec![ Msg { role: "system".into(), content: system.into() }, Msg { role: "user".into(), content: user.into() }, ], temperature, max_tokens, reasoning_budget: budget, stream: Some(true), }; let resp = self.client.post(&url).json(&body).send().await?; if !resp.status().is_success() { let text = resp.text().await.unwrap_or_default(); return Err(anyhow!("HTTP error: {}", text.chars().take(300).collect::())); } let mut stream = resp.bytes_stream(); let mut full = String::new(); let mut buf = String::new(); while let Some(chunk) = stream.next().await { let bytes = chunk.map_err(|e| anyhow!("stream error: {}", e))?; buf.push_str(&String::from_utf8_lossy(&bytes)); // process complete SSE lines while let Some(nl) = buf.find('\n') { let line = buf[..nl].trim().to_string(); buf = buf[nl + 1..].to_string(); if line.is_empty() || line == "data: [DONE]" { continue; } if let Some(data) = line.strip_prefix("data: ") { if let Ok(v) = serde_json::from_str::(data) { if let Some(delta) = v.get("choices") .and_then(|c| c.as_array()) .and_then(|a| a.first()) .and_then(|c| c.get("delta")) .and_then(|d| d.get("content")) .and_then(|x| x.as_str()) { full.push_str(delta); on_chunk(delta); } // also handle reasoning_content delta (stream thinking too) if let Some(delta) = v.get("choices") .and_then(|c| c.as_array()) .and_then(|a| a.first()) .and_then(|c| c.get("delta")) .and_then(|d| d.get("reasoning_content")) .and_then(|x| x.as_str()) { // Stream thinking to the callback so the user can see it on_chunk(delta); } } } } } // strip thinking blocks if let Some(idx) = full.rfind("") { full = full[idx + "".len()..].trim().to_string(); } Ok(full) } }