zesdex/apps/infrastructure/src/llm/provider.rs

//! Async HTTP client for OpenAI/Anthropic-compatible chat completion APIs,
//! supporting both non-streaming and SSE-streaming requests with automatic retry.

use rand_core::RngCore;
use std::time::Duration;

use zesdex_application::ports::ProviderService;
use zesdex_domain::core::{
    ChatMessage, ChatRequest, ChatResponse, SseParser, StreamEvent, StreamOptions, ToolDef,
};

use zesdex_domain::agent::defaults::{DEFAULT_API_BASE, DEFAULT_MODEL};
const DEFAULT_BASE_URL: &str = DEFAULT_API_BASE;
pub const DEFAULT_API_KEY: &str = "";
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(600);

// ---------------------------------------------------------------------------
// Retry helpers
// ---------------------------------------------------------------------------

fn backoff_seconds(attempt: u32, cap: u64) -> Duration {
    let base = 2u64.pow(attempt.saturating_sub(1));
    let delay = std::cmp::min(base, cap);
    let jitter_factor = 0.75 + (rand_core::OsRng.next_u32() % 51) as f64 / 100.0;
    Duration::from_secs_f64(delay as f64 * jitter_factor)
}

pub fn is_auth_error(err_str: &str) -> bool {
    let err_lower = err_str.to_lowercase();
    (err_str.contains("API error 401")
        || err_str.contains("API error 402")
        || err_str.contains("API error 403"))
        || err_lower.contains("unauthorized")
        || err_lower.contains("forbidden")
        || err_lower.contains("authentication failed")
}

fn is_rate_limit(err_str: &str) -> bool {
    err_str.contains("API error 429") || err_str.to_lowercase().contains("rate limit")
}

fn backoff_for_error(attempt: u32, err_str: &str) -> Duration {
    if is_rate_limit(err_str) {
        backoff_seconds(attempt, 60)
    } else {
        backoff_seconds(attempt, 30)
    }
}

// ---------------------------------------------------------------------------
// Client
// ---------------------------------------------------------------------------

/// Async HTTP client for a single LLM provider endpoint.
pub struct LlmClient {
    pub client: reqwest::Client,
    pub api_key: String,
    pub base_url: String,
    pub model: String,
}

impl LlmClient {
    pub fn new(mut api_key: String, model: String, base_url: Option<String>) -> Self {
        if api_key.is_empty() {
            api_key = DEFAULT_API_KEY.to_string();
        }
        let model = if model.is_empty() {
            DEFAULT_MODEL.to_string()
        } else {
            model
        };
        let client = match reqwest::Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .connect_timeout(CONNECT_TIMEOUT)
            .build()
        {
            Ok(c) => c,
            Err(e) => {
                tracing::warn!(
                    "failed to build reqwest client with connect timeout: {}. \
                     retrying without connect timeout",
                    e,
                );
                match reqwest::Client::builder().timeout(REQUEST_TIMEOUT).build() {
                    Ok(c) => c,
                    Err(e2) => {
                        tracing::warn!("also failed: {e2}. using default client");
                        reqwest::Client::new()
                    }
                }
            }
        };
        LlmClient {
            client,
            api_key,
            base_url: base_url
                .filter(|s| !s.is_empty())
                .unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
            model,
        }
    }

    async fn try_stream_once(
        &self,
        req: &ChatRequest,
        url: &str,
        on_event: &mut (dyn FnMut(&StreamEvent) -> bool + Send),
    ) -> anyhow::Result<(ChatMessage, Option<(u64, u64)>)> {
        let mut http_req = self
            .client
            .post(url)
            .header("Content-Type", "application/json");
        if !self.api_key.is_empty() {
            http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
        }

        let mut resp = http_req.json(req).send().await.map_err(|e| {
            if e.is_timeout() {
                anyhow::anyhow!(
                    "API request timed out after {REQUEST_TIMEOUT:?}. \
                     Check your network or try again."
                )
            } else if e.is_connect() {
                anyhow::anyhow!(
                    "Could not connect to {}. \
                     Is the URL correct and is the service reachable?",
                    self.base_url
                )
            } else {
                anyhow::anyhow!("API request failed: {e}")
            }
        })?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
        }

        struct StreamedTurn {
            content: String,
            tool_calls: Vec<zesdex_domain::core::ToolCall>,
            done_received: bool,
        }

        impl StreamedTurn {
            fn new() -> Self {
                StreamedTurn {
                    content: String::new(),
                    tool_calls: Vec::new(),
                    done_received: false,
                }
            }

            fn apply_event(&mut self, event: &StreamEvent) {
                match event {
                    StreamEvent::Token(t) => self.content.push_str(t),
                    StreamEvent::Reasoning(_) => {}
                    StreamEvent::ToolCallDelta {
                        index,
                        id,
                        name,
                        arguments_delta,
                    } => {
                        const MAX_TOOL_CALLS: usize = 64;
                        let index = usize::min(*index, MAX_TOOL_CALLS.saturating_sub(1));

                        while self.tool_calls.len() <= index {
                            self.tool_calls.push(zesdex_domain::core::ToolCall {
                                id: String::new(),
                                type_: "function".to_string(),
                                function: zesdex_domain::core::ToolFunction {
                                    name: String::new(),
                                    arguments: serde_json::Value::String(String::new()),
                                },
                            });
                        }

                        let tc = &mut self.tool_calls[index];

                        if let Some(ref id_val) = id {
                            tc.id = id_val.clone();
                        }
                        if let Some(ref name_val) = name {
                            tc.function.name = name_val.clone();
                        }
                        if let serde_json::Value::String(ref mut args) = tc.function.arguments {
                            args.push_str(arguments_delta);
                        }
                    }
                    _ => {
                        tracing::debug!("unhandled stream event in apply_event: {event:?}");
                    }
                }
            }

            fn build_assistant_message(self) -> ChatMessage {
                ChatMessage {
                    role: zesdex_domain::core::Role::Assistant,
                    content: if self.content.is_empty() {
                        None
                    } else {
                        Some(self.content)
                    },
                    tool_calls: if self.tool_calls.is_empty() {
                        None
                    } else {
                        Some(self.tool_calls)
                    },
                    tool_call_id: None,
                    name: None,
                }
            }
        }

        let mut turn = StreamedTurn::new();
        let mut usage: Option<(u64, u64)> = None;
        let mut parser = SseParser::new();
        let mut byte_buf: Vec<u8> = Vec::new();

        while let Some(chunk) = resp.chunk().await? {
            byte_buf.extend_from_slice(&chunk);

            let valid_len = match std::str::from_utf8(&byte_buf) {
                Ok(s) => s.len(),
                Err(e) => {
                    let n = e.valid_up_to();
                    if n == 0 {
                        byte_buf.drain(..1);
                        continue;
                    }
                    n
                }
            };
            let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
            byte_buf.drain(..valid_len);

            for event in parser.feed(&text) {
                if !on_event(&event) {
                    anyhow::bail!("aborted");
                }
                match &event {
                    StreamEvent::Usage {
                        prompt_tokens,
                        completion_tokens,
                        ..
                    } => {
                        usage = Some((*prompt_tokens, *completion_tokens));
                    }
                    StreamEvent::Error(msg) => {
                        anyhow::bail!("stream error: {msg}");
                    }
                    StreamEvent::Done => {
                        turn.apply_event(&event);
                        turn.done_received = true;
                        return Ok((turn.build_assistant_message(), usage));
                    }
                    other => {
                        turn.apply_event(other);
                    }
                }
            }
        }

        Ok((turn.build_assistant_message(), usage))
    }
}

impl ProviderService for LlmClient {
    async fn chat(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
        max_tokens: Option<u32>,
        temperature: Option<f32>,
    ) -> anyhow::Result<(ChatMessage, Option<(u64, u64)>)> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(max_tokens.unwrap_or(4096)),
            temperature: Some(temperature.unwrap_or(0.7)),
            tools,
            stream: Some(false),
            stop: None,
            stream_options: None,
            tool_choice: None,
            top_p: None,
        };

        let url = format!("{}/chat/completions", self.base_url);
        let max_retries = 10;
        let mut attempt = 0u32;

        loop {
            attempt += 1;

            let mut http_req = self
                .client
                .post(&url)
                .header("Content-Type", "application/json");

            if !self.api_key.is_empty() {
                http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
            }

            let result = async {
                let resp = http_req.json(&req).send().await.map_err(|e| {
                    if e.is_timeout() {
                        anyhow::anyhow!(
                            "API request timed out after {REQUEST_TIMEOUT:?}. \
                             Check your network or try again."
                        )
                    } else if e.is_connect() {
                        anyhow::anyhow!(
                            "Could not connect to {}. \
                             Is the URL correct and is the service reachable?",
                            self.base_url
                        )
                    } else {
                        anyhow::anyhow!("API request failed: {e}")
                    }
                })?;

                if !resp.status().is_success() {
                    let status = resp.status();
                    let body = resp.text().await.unwrap_or_default();
                    anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
                }

                let data: ChatResponse = resp.json().await?;
                let usage = data
                    .usage
                    .map(|u| (u64::from(u.prompt_tokens), u64::from(u.completion_tokens)));
                let message = data
                    .choices
                    .into_iter()
                    .next()
                    .and_then(|c| c.message)
                    .ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
                Ok((message, usage))
            }
            .await;

            match result {
                Ok((msg, usage)) => return Ok((msg, usage)),
                Err(e) => {
                    let err_str = e.to_string();
                    if attempt >= max_retries || is_auth_error(&err_str) {
                        return Err(e);
                    }
                    let delay = backoff_for_error(attempt, &err_str);
                    tokio::time::sleep(delay).await;
                }
            }
        }
    }

    async fn chat_stream(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
        max_tokens: Option<u32>,
        temperature: Option<f32>,
        mut on_event: Box<dyn FnMut(&StreamEvent) -> bool + Send>,
    ) -> anyhow::Result<(ChatMessage, Option<(u64, u64)>)> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(max_tokens.unwrap_or(4096)),
            temperature: Some(temperature.unwrap_or(0.7)),
            tools,
            stream: Some(true),
            stop: None,
            stream_options: Some(StreamOptions {
                include_usage: true,
            }),
            tool_choice: None,
            top_p: None,
        };

        let url = format!("{}/chat/completions", self.base_url);
        let max_retries_stream = 10;
        let mut attempt = 0u32;

        loop {
            attempt += 1;
            let mut captured_content = false;
            let mut wrapped = |event: &StreamEvent| -> bool {
                match event {
                    StreamEvent::Token(_)
                    | StreamEvent::Reasoning(_)
                    | StreamEvent::ToolCallDelta { .. } => {
                        captured_content = true;
                    }
                    _ => {}
                }
                on_event(event)
            };

            match self.try_stream_once(&req, &url, &mut wrapped).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    let err_str = e.to_string();
                    if is_auth_error(&err_str) || captured_content {
                        return Err(e);
                    }
                    if attempt >= max_retries_stream {
                        return Err(e);
                    }
                    let delay = backoff_for_error(attempt, &err_str);
                    tokio::time::sleep(delay).await;
                }
            }
        }
    }
}

pub fn resolve_api_key(
    settings: &zesdex_domain::cms::Settings,
    app_config: &zesdex_domain::cms::AppConfig,
) -> String {
    let provider = &settings.provider;

    let mut api_key = settings.api_keys.get(provider).cloned().unwrap_or_default();

    if api_key.is_empty() {
        if let Some(provider_cfg) = app_config.providers.get(provider) {
            api_key = provider_cfg
                .api_key_env
                .as_ref()
                .and_then(|env| std::env::var(env).ok())
                .or_else(|| provider_cfg.default_api_key.clone())
                .unwrap_or_default();
        }
    }

    api_key
}