af-llm 0.2.0

//! 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::sync::Arc;
use std::time::Duration;

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(Debug, 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 LlmConfig {
    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>,
}

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().timeout(config.timeout).build()?;
        Ok(Self {
            http,
            base_url: config.base_url.trim_end_matches('/').into(),
            api_key: config.api_key.into(),
            breaker,
        })
    }

    /// 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,
        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 = self.stream_collect_once(&url, request, &mut on_delta).await;
        match &result {
            Ok(_) => self.breaker.record_success(),
            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 mut req = self.http.post(url).json(request);
        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() {
            return Err(LlmError::Api {
                status: status.as_u16(),
                body: response.text().await.unwrap_or_default(),
            });
        }

        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(),
        ))
    }
}

#[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);
    }
}