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<Msg>,
temperature: f32,
max_tokens: i32,
#[serde(skip_serializing_if = "Option::is_none")]
reasoning_budget: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
stream: Option<bool>,
}
#[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<Choice>,
}
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<String> {
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<i32>) -> 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::<String>()));
tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await;
continue;
}
match serde_json::from_str::<Resp>(&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("</think>") {
clean = clean[idx + "</think>".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("</think>") {
from_rc = from_rc[idx + "</think>".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::<String>()));
}
}
}
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<i32>) -> Result<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::<String>()));
tokio::time::sleep(Duration::from_secs(1 + attempt as u64)).await;
continue;
}
match serde_json::from_str::<Resp>(&text) {
Ok(r) => {
if let Some(c) = r.choices.into_iter().next() {
let mut content = c.message.content.clone();
// strip <think>...</think> blocks
if let Some(idx) = content.rfind("</think>") {
content = content[idx + "</think>".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("</think>") {
from_rc = from_rc[idx + "</think>".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::<String>()));
}
}
}
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<F>(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, on_chunk: F) -> Result<String>
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<F>(&self, system: &str, user: &str, max_tokens: i32, temperature: f32, budget: Option<i32>, mut on_chunk: F) -> Result<String>
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::<String>()));
}
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::<serde_json::Value>(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("</think>") {
full = full[idx + "</think>".len()..].trim().to_string();
}
Ok(full)
}
}