pub mod ollama;
pub mod prompt;
use anyhow::Result;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ChatMessage {
pub role: Role,
pub content: String,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
}
pub trait LlmProvider: Send + Sync {
fn generate(&self, messages: &[ChatMessage]) -> Result<String>;
fn is_available(&self) -> bool;
fn model_name(&self) -> &str;
}
pub fn strip_thinking(text: &str) -> String {
let mut result = String::with_capacity(text.len());
let mut rest = text;
while let Some(start) = rest.find("<think>") {
result.push_str(&rest[..start]);
if let Some(end) = rest[start..].find("</think>") {
rest = &rest[start + end + "</think>".len()..];
} else {
return result.trim().to_string();
}
}
result.push_str(rest);
result.trim().to_string()
}
pub fn split_thinking(text: &str) -> (Option<String>, String) {
let mut thinking = String::new();
let mut answer = String::with_capacity(text.len());
let mut rest = text;
while let Some(start) = rest.find("<think>") {
answer.push_str(&rest[..start]);
if let Some(end) = rest[start..].find("</think>") {
let think_content = &rest[start + "<think>".len()..start + end];
thinking.push_str(think_content.trim());
rest = &rest[start + end + "</think>".len()..];
} else {
break;
}
}
answer.push_str(rest);
let answer = answer.trim().to_string();
if thinking.is_empty() {
(None, answer)
} else {
(Some(thinking), answer)
}
}