pub mod anthropic;
pub mod null;
pub mod openai;
use crate::config::{LlmConfig, LlmProvider};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
pub use anthropic::AnthropicLlmClient;
pub use null::NullLlmClient;
pub use openai::OpenAiLlmClient;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LlmVerdict {
pub match_spec: bool,
pub reason: String,
}
pub trait LlmClient: Send + Sync {
fn evaluate(&self, system_prompt: &str, user_prompt: &str) -> Option<LlmVerdict>;
fn complete(&self, system_prompt: &str, user_prompt: &str) -> Option<String>;
}
pub struct BudgetedClient {
inner: Arc<dyn LlmClient>,
remaining: AtomicU32,
}
impl BudgetedClient {
pub fn new(inner: Arc<dyn LlmClient>, budget: u32) -> Self {
Self {
inner,
remaining: AtomicU32::new(budget),
}
}
}
impl LlmClient for BudgetedClient {
fn evaluate(&self, system_prompt: &str, user_prompt: &str) -> Option<LlmVerdict> {
loop {
let current = self.remaining.load(Ordering::Acquire);
if current == 0 {
return None;
}
if self
.remaining
.compare_exchange(current, current - 1, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
break;
}
}
self.inner.evaluate(system_prompt, user_prompt)
}
fn complete(&self, system_prompt: &str, user_prompt: &str) -> Option<String> {
loop {
let current = self.remaining.load(Ordering::Acquire);
if current == 0 {
return None;
}
if self
.remaining
.compare_exchange(current, current - 1, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
break;
}
}
self.inner.complete(system_prompt, user_prompt)
}
}
pub fn build_client(cfg: &LlmConfig, force_off: bool) -> Arc<dyn LlmClient> {
if force_off || !cfg.enabled {
return Arc::new(NullLlmClient);
}
let inner: Arc<dyn LlmClient> = match cfg.provider {
LlmProvider::Anthropic => {
match AnthropicLlmClient::from_env(cfg.model.clone(), cfg.timeout_s) {
Some(c) => Arc::new(c),
None => {
eprintln!(
"spec-drift: [llm] enabled but ANTHROPIC_API_KEY is not set; \
skipping LLM rules for this run."
);
Arc::new(NullLlmClient)
}
}
}
LlmProvider::OpenAi | LlmProvider::Local => {
match OpenAiLlmClient::new(cfg.provider, cfg.model.clone(), cfg.timeout_s) {
Some(c) => Arc::new(c),
None => {
if cfg.provider == LlmProvider::OpenAi {
eprintln!(
"spec-drift: [llm] enabled with openai provider but OPENAI_API_KEY is not set; \
skipping LLM rules for this run."
);
} else {
eprintln!(
"spec-drift: [llm] enabled with local provider but creation failed; \
skipping LLM rules for this run."
);
}
Arc::new(NullLlmClient)
}
}
}
};
Arc::new(BudgetedClient::new(inner, cfg.max_calls))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicU32, Ordering};
struct CountingClient {
calls: AtomicU32,
}
impl LlmClient for CountingClient {
fn evaluate(&self, _: &str, _: &str) -> Option<LlmVerdict> {
self.calls.fetch_add(1, Ordering::AcqRel);
Some(LlmVerdict {
match_spec: true,
reason: "ok".into(),
})
}
fn complete(&self, _: &str, _: &str) -> Option<String> {
self.calls.fetch_add(1, Ordering::AcqRel);
Some("ok".into())
}
}
#[test]
fn budget_exhausts_and_stops_calling_inner() {
let inner = Arc::new(CountingClient {
calls: AtomicU32::new(0),
});
let client = BudgetedClient::new(inner.clone(), 2);
assert!(client.evaluate("s", "u").is_some());
assert!(client.evaluate("s", "u").is_some());
assert!(client.evaluate("s", "u").is_none());
assert!(client.evaluate("s", "u").is_none());
assert_eq!(inner.calls.load(Ordering::Acquire), 2);
}
#[test]
fn force_off_always_yields_null_client() {
let cfg = LlmConfig {
enabled: true,
..LlmConfig::default()
};
let client = build_client(&cfg, true);
assert!(client.evaluate("s", "u").is_none());
}
#[test]
fn disabled_config_yields_null_client() {
let cfg = LlmConfig::default(); let client = build_client(&cfg, false);
assert!(client.evaluate("s", "u").is_none());
}
}