use crate::error::Result;
pub trait LlmCallback: Send + Sync {
fn generate(&self, prompt: &str, max_tokens: usize) -> Result<String>;
fn model_name(&self) -> &str {
"unknown"
}
fn is_available(&self) -> bool {
true
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockLlm {
response: String,
}
impl MockLlm {
fn new(response: &str) -> Self {
Self {
response: response.to_string(),
}
}
}
impl LlmCallback for MockLlm {
fn generate(&self, _prompt: &str, _max_tokens: usize) -> Result<String> {
Ok(self.response.clone())
}
fn model_name(&self) -> &str {
"mock"
}
}
#[test]
fn trait_object_works() {
let llm: Box<dyn LlmCallback> = Box::new(MockLlm::new("hello"));
assert_eq!(llm.generate("test", 100).unwrap(), "hello");
assert_eq!(llm.model_name(), "mock");
assert!(llm.is_available());
}
#[test]
fn default_methods() {
struct MinimalLlm;
impl LlmCallback for MinimalLlm {
fn generate(&self, _prompt: &str, _max_tokens: usize) -> Result<String> {
Ok("ok".into())
}
}
let llm = MinimalLlm;
assert_eq!(llm.model_name(), "unknown");
assert!(llm.is_available());
}
}