af-llm 0.7.1

//! Async LLM client — port of `agent_core/llm/client.py::llm_completion`.
//!
//! Talks to any OpenAI-compatible `/chat/completions` endpoint. In production
//! this points at the LiteLLM proxy, which owns model routing and fallback; the
//! client itself owns one bounded streaming provider attempt and a shared
//! circuit breaker. The event-sourced Agent runtime is the sole retry owner.

use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};

use futures::stream::StreamExt;

use crate::circuit_breaker::CircuitBreaker;
use crate::error::{LlmError, Result};
use crate::stream::{SseDecoder, StreamAssembler};
use crate::types::{CompletionRequest, CompletionResponse};

/// Default per-request timeout. DeepSeek's server-side timeout is 900s, which
/// caused 15-minute hangs that exhausted DB connections; 30s is generous for
/// summarization and agent flows can override per call.
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);

/// Configuration for [`LlmClient`].
#[derive(Clone)]
pub struct LlmConfig {
    /// Base URL of the OpenAI-compatible API, e.g. `http://litellm:4000/v1`.
    pub base_url: String,
    /// Bearer token. May be empty when the proxy handles auth upstream.
    pub api_key: String,
    /// Per-request timeout.
    pub timeout: Duration,
}

impl fmt::Debug for LlmConfig {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("LlmConfig")
            .field("base_url", &self.base_url)
            .field("api_key", &"[REDACTED]")
            .field("timeout", &self.timeout)
            .finish()
    }
}

impl LlmConfig {
    /// Client for an OpenAI-compatible `base_url` using bearer `api_key`, with default retry, timeout and breaker settings.
    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>) -> Self {
        Self {
            base_url: base_url.into(),
            api_key: api_key.into(),
            timeout: DEFAULT_TIMEOUT,
        }
    }
}

/// Cheap to clone (everything shareable lives behind `Arc`). Share one instance
/// across tasks so they all observe the same circuit-breaker state.
#[derive(Clone)]
pub struct LlmClient {
    http: reqwest::Client,
    base_url: Arc<str>,
    api_key: Arc<str>,
    breaker: Arc<CircuitBreaker>,
    timeout: Duration,
    images: Option<Arc<dyn crate::images::ImageResolver>>,
}

impl LlmClient {
    /// Build a client with a fresh, default circuit breaker.
    pub fn new(config: LlmConfig) -> Result<Self> {
        Self::with_breaker(config, Arc::new(CircuitBreaker::default()))
    }

    /// Build a client sharing an existing circuit breaker (e.g. one breaker
    /// across several clients hitting the same provider).
    pub fn with_breaker(config: LlmConfig, breaker: Arc<CircuitBreaker>) -> Result<Self> {
        let http = reqwest::Client::builder().build()?;
        Ok(Self {
            http,
            base_url: config.base_url.trim_end_matches('/').into(),
            api_key: config.api_key.into(),
            breaker,
            timeout: config.timeout,
            images: None,
        })
    }

    /// Bind a caller-aware image resolver. Text-only clients need no resolver.
    pub fn with_image_resolver(mut self, resolver: Arc<dyn crate::images::ImageResolver>) -> Self {
        self.images = Some(resolver);
        self
    }

    /// Access the shared circuit breaker (health checks, metrics).
    pub fn breaker(&self) -> &Arc<CircuitBreaker> {
        &self.breaker
    }

    /// Execute one auditable streaming provider attempt. Retry and durable
    /// attempt accounting belong exclusively to `af-agent-runtime`.
    pub async fn complete_stream_single_attempt<F>(
        &self,
        request: &CompletionRequest,
        on_delta: F,
    ) -> Result<CompletionResponse>
    where
        F: FnMut(&str, bool) + Send,
    {
        let cancellation = tokio_util::sync::CancellationToken::new();
        self.complete_stream_single_attempt_controlled(
            request,
            &cancellation,
            Instant::now() + self.timeout,
            on_delta,
        )
        .await
    }

    /// Execute one provider attempt bounded by caller cancellation and deadline.
    pub async fn complete_stream_single_attempt_controlled<F>(
        &self,
        request: &CompletionRequest,
        cancellation: &tokio_util::sync::CancellationToken,
        deadline: Instant,
        mut on_delta: F,
    ) -> Result<CompletionResponse>
    where
        F: FnMut(&str, bool) + Send,
    {
        if self.breaker.is_open() {
            return Err(LlmError::CircuitOpen);
        }
        let url = format!("{}/chat/completions", self.base_url);
        let result = tokio::select! {
            biased;
            _ = cancellation.cancelled() => Err(LlmError::Cancelled),
            _ = tokio::time::sleep_until(deadline.into()) => Err(LlmError::DeadlineExceeded),
            result = self.stream_collect_once(&url, request, &mut on_delta) => result,
        };
        match &result {
            Ok(_) => self.breaker.record_success(),
            Err(LlmError::Cancelled | LlmError::DeadlineExceeded | LlmError::InvalidInput(_)) => {}
            Err(_) => self.breaker.record_failure(),
        }
        result
    }

    async fn stream_collect_once<F>(
        &self,
        url: &str,
        request: &CompletionRequest,
        on_delta: &mut F,
    ) -> Result<CompletionResponse>
    where
        F: FnMut(&str, bool) + Send,
    {
        let payload = crate::images::provider_request(request, self.images.as_deref()).await?;
        let mut req = self.http.post(url).json(&payload);
        if let Some(attempt_id) = &request.provider_attempt_id {
            req = req
                .header("Idempotency-Key", attempt_id)
                .header("X-Agent-Factory-Attempt-Id", attempt_id);
        }
        if !self.api_key.is_empty() {
            req = req.bearer_auth(self.api_key.as_ref());
        }
        let response = req.send().await?;
        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            return Err(LlmError::Api {
                status: status.as_u16(),
                body: if request
                    .messages
                    .iter()
                    .any(|message| !message.images.is_empty())
                {
                    "[multimodal provider error redacted]".into()
                } else {
                    redact_and_truncate(&body, self.api_key.as_ref())
                },
            });
        }

        let mut assembler = StreamAssembler::default();
        let mut decoder = SseDecoder::default();
        let mut bytes = response.bytes_stream();
        while let Some(chunk) = bytes.next().await {
            for data in decoder.push(&chunk?)? {
                if data == "[DONE]" {
                    decoder.finish()?;
                    return assembler.finish();
                }
                if let Some(delta) = assembler.apply_json(&data)? {
                    on_delta(&delta.content, delta.has_tool_calls);
                }
            }
        }
        decoder.finish()?;
        Err(LlmError::StreamProtocol(
            "stream ended before [DONE]".into(),
        ))
    }
}

const MAX_ERROR_BODY_BYTES: usize = 4 * 1024;

fn redact_and_truncate(body: &str, secret: &str) -> String {
    let redacted = if secret.is_empty() {
        body.to_owned()
    } else {
        body.replace(secret, "[REDACTED]")
    };
    if redacted.len() <= MAX_ERROR_BODY_BYTES {
        return redacted;
    }
    let mut end = MAX_ERROR_BODY_BYTES;
    while !redacted.is_char_boundary(end) {
        end -= 1;
    }
    format!("{}...[truncated]", &redacted[..end])
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn prepared_streaming_request_is_wire_stable() {
        let mut request = CompletionRequest::new("model", vec![]).stream(true);
        request.stream_options = Some(crate::StreamOptions {
            include_usage: true,
        });
        let wire = serde_json::to_value(&request).unwrap();
        assert_eq!(wire["stream"], true);
        assert_eq!(wire["stream_options"]["include_usage"], true);
    }

    #[test]
    fn config_debug_and_provider_errors_hide_credentials() {
        let config = LlmConfig::new("https://provider.invalid/v1", "top-secret");
        assert!(!format!("{config:?}").contains("top-secret"));

        let body = format!("token=top-secret {}", "界".repeat(MAX_ERROR_BODY_BYTES));
        let safe = redact_and_truncate(&body, "top-secret");
        assert!(!safe.contains("top-secret"));
        assert!(safe.len() <= MAX_ERROR_BODY_BYTES + "...[truncated]".len());
    }

    #[tokio::test]
    async fn caller_cancellation_wins_before_network_io() {
        let client =
            LlmClient::new(LlmConfig::new("https://provider.invalid/v1", "secret")).unwrap();
        let cancellation = tokio_util::sync::CancellationToken::new();
        cancellation.cancel();
        let result = client
            .complete_stream_single_attempt_controlled(
                &CompletionRequest::new("model", vec![]),
                &cancellation,
                Instant::now() + Duration::from_secs(1),
                |_, _| {},
            )
            .await;
        assert!(matches!(result, Err(LlmError::Cancelled)));
        assert_eq!(client.breaker().status().failure_count, 0);
    }
}