use std::sync::Arc;
use std::time::Instant;
use anyhow::Result;
use super::{Backend, GenAnswer, GenRequest, Generator};
use crate::viz::claude::{
api_key_present, default_transport, ClaudeMessage, ClaudeRequest, ClaudeTransport, DEFAULT_MODEL,
};
pub struct ClaudeGenerator {
id: String,
model: String,
transport: Arc<dyn ClaudeTransport>,
}
impl ClaudeGenerator {
pub fn new(model: &str) -> Self {
Self::with_transport_impl(model, default_transport())
}
pub fn with_transport(model: &str, transport: Arc<dyn ClaudeTransport>) -> Self {
Self::with_transport_impl(model, transport)
}
fn with_transport_impl(model: &str, transport: Arc<dyn ClaudeTransport>) -> Self {
let model = if model.trim().is_empty() { DEFAULT_MODEL.to_string() } else { model.trim().to_string() };
Self { id: format!("claude:{model}"), model, transport }
}
}
impl Backend for ClaudeGenerator {
fn id(&self) -> &str {
&self.id
}
fn available(&self) -> bool {
api_key_present()
}
}
impl Generator for ClaudeGenerator {
fn complete(&self, req: &GenRequest) -> Result<GenAnswer> {
let prompt = match &req.system {
Some(s) if !s.trim().is_empty() => format!("{s}\n\n{}", req.prompt),
_ => req.prompt.clone(),
};
let mut creq = ClaudeRequest::new(self.model.as_str(), vec![ClaudeMessage::user(prompt)]);
if req.max_tokens > 0 {
creq.max_tokens = req.max_tokens.min(u32::MAX as usize) as u32;
}
let t0 = Instant::now();
let text = self
.transport
.send(&creq)
.map_err(|e| anyhow::anyhow!("claude generate: {e}"))?;
let latency_ms = t0.elapsed().as_secs_f64() * 1000.0;
Ok(GenAnswer {
text: text.clone(),
tokens_in: (req.prompt.len() / 4) as i64,
tokens_out: (text.len() / 4) as i64,
tokens_per_s: 0.0,
latency_ms,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::viz::claude::{ClaudeError, FakeTransport};
#[test]
fn completes_via_injected_transport() {
let t = Arc::new(FakeTransport::replying("fn add(a: i32, b: i32) -> i32 { a + b }"));
let g = ClaudeGenerator::with_transport("claude-opus-4-8", t);
assert_eq!(g.id(), "claude:claude-opus-4-8");
let ans = g.complete(&GenRequest::new("implement add")).unwrap();
assert!(ans.text.contains("fn add"), "returns the transport's completion");
assert!(ans.tokens_out > 0, "token economics estimated from the text");
}
#[test]
fn empty_model_defaults_and_system_folds_into_the_prompt() {
let t = Arc::new(FakeTransport::replying("ok"));
let g = ClaudeGenerator::with_transport("", t.clone());
assert_eq!(g.id(), format!("claude:{DEFAULT_MODEL}"), "empty model ⇒ default");
let req = GenRequest::new("do the thing").with_system("You are a Rust expert.");
g.complete(&req).unwrap();
let seen = t.requests();
assert_eq!(seen.len(), 1);
let sent = &seen[0].messages[0].content;
assert!(sent.contains("Rust expert"), "system prompt folded in");
assert!(sent.contains("do the thing"), "user prompt present");
}
#[test]
fn transport_error_surfaces_not_panics() {
let t = Arc::new(FakeTransport::failing(ClaudeError::NoCredential));
let g = ClaudeGenerator::with_transport("m", t);
assert!(g.complete(&GenRequest::new("x")).is_err(), "no key ⇒ Err, never a panic");
}
}