use std::time::Duration;
use futures::StreamExt;
use open_agent::retry::{RetryConfig, retry_with_backoff_conditional};
use open_agent::{AgentOptions, ApiProtocol, ContentBlock, FinishReason, StreamEvent, query};
use crate::config::LlmConfig;
use crate::llm::error::LlmError;
use crate::llm::json_parsing::{Extracted, extract_json};
use crate::text::excerpt;
const RESPONSE_EXCERPT_MAX: usize = 200;
pub struct LlmClient {
pub(crate) base_url: String,
pub(crate) model: String,
pub(crate) api_key: String,
pub(crate) protocol: ApiProtocol,
pub(crate) temperature: Option<f32>,
pub(crate) max_tokens: Option<u32>,
pub(crate) timeout_secs: u64,
pub(crate) retry_config: RetryConfig,
}
impl std::fmt::Debug for LlmClient {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LlmClient")
.field("base_url", &self.base_url)
.field("model", &self.model)
.field("api_key", &"<redacted>")
.field("protocol", &self.protocol.as_str())
.field("temperature", &self.temperature)
.field("max_tokens", &self.max_tokens)
.field("timeout_secs", &self.timeout_secs)
.finish()
}
}
impl LlmClient {
pub fn model(&self) -> &str {
&self.model
}
pub fn endpoint(&self) -> &str {
&self.base_url
}
pub fn temperature(&self) -> Option<f32> {
self.temperature
}
pub fn protocol(&self) -> ApiProtocol {
self.protocol
}
pub fn new(cfg: &LlmConfig) -> Result<Self, LlmError> {
if !cfg.enabled {
return Err(LlmError::NotConfigured(
"LLM is disabled in config (set `enabled = true`)".to_string(),
));
}
let endpoint = cfg.endpoint.clone().ok_or_else(|| {
LlmError::NotConfigured("LLM endpoint is not set in config".to_string())
})?;
let model = cfg
.model
.clone()
.ok_or_else(|| LlmError::NotConfigured("LLM model is not set in config".to_string()))?;
let api_key = cfg
.api_key
.clone()
.unwrap_or_else(|| "not-needed".to_string());
let protocol = crate::config::parse_protocol(cfg.protocol.as_deref()).ok_or_else(|| {
LlmError::NotConfigured(format!(
"unknown protocol `{}`; expected `openai` or `anthropic`",
cfg.protocol.as_deref().unwrap_or_default()
))
})?;
let max_attempts = cfg.max_retries.max(1);
Ok(LlmClient {
base_url: endpoint,
model,
api_key,
protocol,
temperature: cfg.temperature,
max_tokens: cfg.max_tokens,
timeout_secs: cfg.timeout_secs,
retry_config: RetryConfig {
max_attempts,
initial_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(60),
backoff_multiplier: 2.0,
jitter_factor: 0.1,
},
})
}
pub async fn complete_json(
&self,
system_prompt: &str,
user_content: &str,
) -> Result<Extracted, LlmError> {
let builder = AgentOptions::builder()
.model(&self.model)
.base_url(&self.base_url)
.api_key(&self.api_key)
.system_prompt(system_prompt)
.protocol(self.protocol)
.timeout(self.timeout_secs);
let builder = match self.temperature {
Some(temperature) => builder.temperature(temperature),
None => builder,
};
let builder = match self.max_tokens {
Some(limit) => builder.max_tokens(limit),
None => builder,
};
let options = builder
.build()
.map_err(|e| LlmError::NotConfigured(format!("AgentOptions build failed: {e}")))?;
let prompt = user_content.to_string();
let mut last_body = String::new();
for _ in 0..NO_JSON_ATTEMPTS {
let result: open_agent::Result<Answer> =
retry_with_backoff_conditional(self.retry_config.clone(), || {
self.run_one_query(&prompt, &options)
})
.await;
match result {
Ok(Answer::Parsed(extracted)) => return Ok(extracted),
Ok(Answer::NoJson { text, finish }) if !worth_asking_again(&finish) => {
return Err(LlmError::ModelStopped {
finish: finish.as_str().to_owned(),
message: stopped_message(&finish, &text),
});
}
Ok(Answer::NoJson { text, .. }) => last_body = text,
Err(e) => {
let status = e.status_code();
let message = format!("{e}");
return Err(LlmError::Transport { status, message });
}
}
}
Err(LlmError::Unparseable(format!(
"no JSON in the response after {NO_JSON_ATTEMPTS} attempts; \
the model answered: {}",
excerpt(&last_body, RESPONSE_EXCERPT_MAX)
)))
}
async fn run_one_query(
&self,
prompt: &str,
options: &AgentOptions,
) -> open_agent::Result<Answer> {
let mut stream = query(prompt, options).await?;
let mut text = String::new();
let mut finish = FinishReason::Unspecified;
while let Some(event) = stream.next().await {
match event? {
StreamEvent::Block(ContentBlock::Text(t)) => text.push_str(&t.text),
StreamEvent::Finish(reason) => finish = reason,
_ => {}
}
}
if text.trim().is_empty() && worth_asking_again(&finish) {
return Err(open_agent::Error::stream(
"the model returned an empty response",
));
}
Ok(match extract_json(&text) {
Some(extracted) => Answer::Parsed(extracted),
None => Answer::NoJson { text, finish },
})
}
}
enum Answer {
Parsed(Extracted),
NoJson {
text: String,
finish: FinishReason,
},
}
pub const NO_JSON_ATTEMPTS: u32 = 3;
fn worth_asking_again(finish: &FinishReason) -> bool {
!matches!(
finish,
FinishReason::Length | FinishReason::ContentFilter
)
}
fn stopped_message(finish: &FinishReason, text: &str) -> String {
match finish {
FinishReason::Length => format!(
"the model hit its output token limit before producing any JSON. \
This file is too large for this model to review in one request - \
split it, or use a provider with a larger output budget. \
It managed: {}",
excerpt(text, RESPONSE_EXCERPT_MAX)
),
_ => format!(
"the model stopped ({}) before producing any JSON: {}",
finish.as_str(),
excerpt(text, RESPONSE_EXCERPT_MAX)
),
}
}
#[cfg(test)]
mod tests;