perspt_core/

llm_provider.rs

//! # LLM Provider Module
//!
//! Thread-safe LLM provider abstraction for multi-agent use.
//! Wraps genai::Client with Arc<RwLock<>> for shared state.

use anyhow::{Context, Result};
use futures::StreamExt;
use genai::adapter::AdapterKind;
use genai::chat::{ChatMessage, ChatRequest, ChatStreamEvent};
use genai::resolver::{Endpoint, ProviderConfig, ServiceTargetResolver};
use genai::Client;
use genai::ServiceTarget;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::{mpsc, RwLock};

use crate::config::Config;

/// End of transmission signal
pub const EOT_SIGNAL: &str = "<|EOT|>";

/// Effective provider id and model after merging config, CLI flags, and env.
#[derive(Debug, Clone)]
pub struct ResolvedProvider {
    /// Provider id, e.g. `openai`, `ollama`.
    pub provider: String,
    /// Model name to use (passed to genai verbatim so namespacing works).
    pub model: String,
}

/// Detect the provider id and a sensible default model from environment keys.
///
/// Used as the fallback when no provider is configured. Falls back to a local
/// Ollama setup when no API keys are present.
pub fn detect_provider_from_env() -> (&'static str, &'static str) {
    if std::env::var("GEMINI_API_KEY").is_ok() {
        ("gemini", "gemini-3.1-flash-lite-preview")
    } else if std::env::var("OPENAI_API_KEY").is_ok() {
        ("openai", "gpt-4o-mini")
    } else if std::env::var("ANTHROPIC_API_KEY").is_ok() {
        ("anthropic", "claude-3-5-sonnet-20241022")
    } else if std::env::var("GROQ_API_KEY").is_ok() {
        ("groq", "llama-3.1-8b-instant")
    } else if std::env::var("COHERE_API_KEY").is_ok() {
        ("cohere", "command-r-plus")
    } else if std::env::var("XAI_API_KEY").is_ok() {
        ("xai", "grok-beta")
    } else if std::env::var("DEEPSEEK_API_KEY").is_ok() {
        ("deepseek", "deepseek-chat")
    } else {
        // Default to Ollama for local usage
        ("ollama", "llama3.2")
    }
}

/// Response from a non-streaming LLM call, carrying text and token usage.
#[derive(Debug, Clone)]
pub struct LlmResponse {
    pub text: String,
    pub tokens_in: Option<i32>,
    pub tokens_out: Option<i32>,
}

/// Shared state for rate limiting and token counting
#[derive(Default)]
struct SharedState {
    total_tokens_used: usize,
    request_count: usize,
}

/// Thread-safe LLM provider implementation using Arc<RwLock<>>.
///
/// This provider can be cheaply cloned and shared across multiple agents.
/// Each clone shares the same underlying client and rate limiting state.
#[derive(Clone)]
pub struct GenAIProvider {
    /// The underlying genai client
    client: Arc<Client>,
    /// Shared state for rate limiting and metrics
    shared: Arc<RwLock<SharedState>>,
}

impl GenAIProvider {
    /// Creates a new GenAI provider with automatic configuration.
    pub fn new() -> Result<Self> {
        let client = Client::default();
        Ok(Self::from_client(client))
    }

    /// Creates a new GenAI provider with explicit configuration.
    pub fn new_with_config(provider_type: Option<&str>, api_key: Option<&str>) -> Result<Self> {
        let adapter_kind = provider_type.and_then(|provider| match str_to_adapter_kind(provider) {
            Ok(adapter_kind) => Some(adapter_kind),
            Err(_) => {
                log::warn!("Unknown provider type for genai client: {provider}");
                None
            }
        });

        // Set environment variable if API key is provided
        if let (Some(provider), Some(key)) = (provider_type, api_key) {
            if let Some(env_var) = provider_api_key_env_var(provider) {
                log::info!("Setting {env_var} environment variable for genai client");
                std::env::set_var(env_var, key);
            } else if provider.eq_ignore_ascii_case("ollama") {
                log::info!("Ollama provider detected - no API key required for local setup");
            } else {
                log::warn!("Unknown provider type for API key: {provider}");
            }
        }

        let client = match adapter_kind {
            Some(adapter_kind) => build_bound_client(adapter_kind, provider_type),
            None => Client::default(),
        };

        Ok(Self::from_client(client))
    }

    /// Build a provider from a `Config`, merging in environment detection and an
    /// optional CLI model override, and return the effective provider/model.
    ///
    /// Precedence:
    ///   - provider: `config.provider` > environment detection
    ///   - model:    `cli_model` > `config.model` > provider default
    ///   - api_key:  `config.api_key` > ambient environment
    ///   - base_url: `config.base_url` > ambient environment
    ///
    /// The returned client is bound to the resolved adapter, so custom/local
    /// OpenAI-compatible model names (e.g. `phi-4-npu-ov`) route correctly while
    /// recognized names still resolve by prefix. Model names are passed through
    /// verbatim so genai namespacing (`openai::model`) keeps working.
    pub fn from_config(
        config: &Config,
        cli_model: Option<&str>,
    ) -> Result<(Self, ResolvedProvider)> {
        let (env_provider, env_model) = detect_provider_from_env();

        let provider = config
            .provider
            .clone()
            .unwrap_or_else(|| env_provider.to_string());

        let env_model_override = std::env::var("OPENAI_MODEL")
            .or_else(|_| std::env::var("MODEL"))
            .ok();

        let model = cli_model
            .map(str::to_string)
            .or_else(|| config.model.clone())
            .or(env_model_override)
            .unwrap_or_else(|| env_model.to_string());

        // Propagate a configured base URL into the env var that build_bound_client
        // reads, without clobbering an explicit ambient override.
        if let Some(base_url) = config.base_url.as_deref() {
            if let Some(env_var) = provider_base_url_env_var(&provider) {
                if std::env::var(env_var).is_err() {
                    std::env::set_var(env_var, base_url);
                }
            }
        }

        let provider_obj = Self::new_with_config(Some(&provider), config.api_key.as_deref())?;
        Ok((provider_obj, ResolvedProvider { provider, model }))
    }

    fn from_client(client: Client) -> Self {
        Self {
            client: Arc::new(client),
            shared: Arc::new(RwLock::new(SharedState::default())),
        }
    }

    /// Get total tokens used across all requests
    pub async fn get_total_tokens_used(&self) -> usize {
        self.shared.read().await.total_tokens_used
    }

    /// Get total request count
    pub async fn get_request_count(&self) -> usize {
        self.shared.read().await.request_count
    }

    /// Increment request counter (for metrics)
    async fn increment_request(&self) {
        let mut state = self.shared.write().await;
        state.request_count += 1;
    }

    /// Add tokens to the total count
    pub async fn add_tokens(&self, count: usize) {
        let mut state = self.shared.write().await;
        state.total_tokens_used += count;
    }

    /// Retrieves all available models for a specific provider.
    pub async fn get_available_models(&self, provider: &str) -> Result<Vec<String>> {
        let adapter_kind = str_to_adapter_kind(provider)?;
        let provider_config = provider_base_url_from_env(provider)
            .map(|base_url| {
                ProviderConfig::from_endpoint(Endpoint::from_owned(normalize_base_url(&base_url)))
            })
            .unwrap_or_default();

        let models = self
            .client
            .all_model_names(adapter_kind, provider_config)
            .await
            .context(format!("Failed to get models for provider: {provider}"))?;

        Ok(models)
    }

    /// Generates a simple text response without streaming.
    /// Includes exponential backoff retry for rate limits and transient errors.
    pub async fn generate_response_simple(&self, model: &str, prompt: &str) -> Result<LlmResponse> {
        self.generate_response_with_retry(model, prompt, 3).await
    }

    /// Generates a response with configurable retry count and exponential backoff.
    pub async fn generate_response_with_retry(
        &self,
        model: &str,
        prompt: &str,
        max_retries: usize,
    ) -> Result<LlmResponse> {
        self.increment_request().await;

        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));

        log::debug!(
            "Sending chat request to model: {model} with prompt length: {} chars",
            prompt.len()
        );

        let start_time = Instant::now();
        let mut last_error: Option<anyhow::Error> = None;
        let mut retry_count = 0;

        while retry_count <= max_retries {
            if retry_count > 0 {
                // Exponential backoff: 1s, 2s, 4s, 8s, ... (capped at 16s)
                let delay_secs = std::cmp::min(1u64 << (retry_count - 1), 16);
                log::warn!(
                    "Retry {}/{} for model {} after {}s delay (previous error: {:?})",
                    retry_count,
                    max_retries,
                    model,
                    delay_secs,
                    last_error.as_ref().map(|e| e.to_string())
                );
                println!(
                    "   ⏳ Rate limited, retrying in {}s (attempt {}/{})",
                    delay_secs, retry_count, max_retries
                );
                tokio::time::sleep(tokio::time::Duration::from_secs(delay_secs)).await;
            }

            match self.client.exec_chat(model, chat_req.clone(), None).await {
                Ok(chat_res) => {
                    let tokens_in = chat_res.usage.prompt_tokens;
                    let tokens_out = chat_res.usage.completion_tokens;
                    let content = chat_res
                        .first_text()
                        .context("No text content in response")?;
                    log::debug!(
                        "Received response with {} characters in {}ms (tokens: in={:?}, out={:?})",
                        content.len(),
                        start_time.elapsed().as_millis(),
                        tokens_in,
                        tokens_out,
                    );

                    // Update shared token counter with real values when available
                    let total = tokens_in.unwrap_or(0) + tokens_out.unwrap_or(0);
                    if total > 0 {
                        self.add_tokens(total as usize).await;
                    }

                    return Ok(LlmResponse {
                        text: content.to_string(),
                        tokens_in,
                        tokens_out,
                    });
                }
                Err(e) => {
                    let err_str = e.to_string();

                    // Check if it's a retryable error (rate limit, server error, network)
                    let is_retryable = err_str.contains("429")
                        || err_str.contains("rate limit")
                        || err_str.contains("Rate limit")
                        || err_str.contains("RESOURCE_EXHAUSTED")
                        || err_str.contains("500")
                        || err_str.contains("502")
                        || err_str.contains("503")
                        || err_str.contains("504")
                        || err_str.contains("timeout")
                        || err_str.contains("connection");

                    if is_retryable && retry_count < max_retries {
                        log::warn!("Retryable error for model {}: {}", model, err_str);
                        last_error = Some(anyhow::anyhow!("{}", err_str));
                        retry_count += 1;
                        continue;
                    } else {
                        return Err(anyhow::anyhow!(
                            "Failed to execute chat request for model {}: {}",
                            model,
                            err_str
                        ));
                    }
                }
            }
        }

        // Should not reach here, but handle gracefully
        Err(last_error
            .unwrap_or_else(|| anyhow::anyhow!("Unknown error after {} retries", max_retries)))
    }

    /// Generates a streaming response and sends chunks via mpsc channel.
    pub async fn generate_response_stream_to_channel(
        &self,
        model: &str,
        prompt: &str,
        tx: mpsc::UnboundedSender<String>,
    ) -> Result<()> {
        self.increment_request().await;

        let chat_req = ChatRequest::default().append_message(ChatMessage::user(prompt));

        log::debug!("Sending streaming chat request to model: {model} with prompt: {prompt}");

        let chat_res_stream = self
            .client
            .exec_chat_stream(model, chat_req, None)
            .await
            .context(format!(
                "Failed to execute streaming chat request for model: {model}"
            ))?;

        let mut stream = chat_res_stream.stream;
        let mut chunk_count = 0;
        let mut total_content_length = 0;
        let mut stream_ended_explicitly = false;
        let start_time = Instant::now();

        log::info!(
            "=== STREAM START === Model: {}, Prompt length: {} chars",
            model,
            prompt.len()
        );

        while let Some(chunk_result) = stream.next().await {
            let elapsed = start_time.elapsed();

            match chunk_result {
                Ok(ChatStreamEvent::Start) => {
                    log::info!(">>> STREAM STARTED for model: {model} at {elapsed:?}");
                }
                Ok(ChatStreamEvent::Chunk(chunk)) => {
                    chunk_count += 1;
                    total_content_length += chunk.content.len();

                    if chunk_count % 10 == 0 || chunk.content.len() > 100 {
                        log::info!(
                            "CHUNK #{}: {} chars, total: {} chars, elapsed: {:?}",
                            chunk_count,
                            chunk.content.len(),
                            total_content_length,
                            elapsed
                        );
                    }

                    if !chunk.content.is_empty() && tx.send(chunk.content.clone()).is_err() {
                        log::error!(
                            "!!! CHANNEL SEND FAILED for chunk #{chunk_count} - STOPPING STREAM !!!"
                        );
                        break;
                    }
                }
                Ok(ChatStreamEvent::ReasoningChunk(chunk)) => {
                    log::info!(
                        "REASONING CHUNK: {} chars at {:?}",
                        chunk.content.len(),
                        elapsed
                    );
                    if !chunk.content.is_empty() {
                        let _ = tx.send(format!("__PERSPT_REASONING__:{}", chunk.content));
                    }
                }
                Ok(ChatStreamEvent::End(_)) => {
                    log::info!(">>> STREAM ENDED EXPLICITLY for model: {model} after {chunk_count} chunks, {total_content_length} chars, {elapsed:?} elapsed");
                    stream_ended_explicitly = true;
                    break;
                }
                Ok(ChatStreamEvent::ToolCallChunk(_)) => {
                    log::debug!("Tool call chunk received (ignored)");
                }
                Ok(ChatStreamEvent::ThoughtSignatureChunk(_)) => {
                    log::debug!("Thought signature chunk received (ignored)");
                }
                Err(e) => {
                    log::error!(
                        "!!! STREAM ERROR after {chunk_count} chunks at {elapsed:?}: {e} !!!"
                    );
                    let error_msg = format!("Stream error: {e}");
                    let _ = tx.send(error_msg);
                    return Err(e.into());
                }
            }
        }

        let final_elapsed = start_time.elapsed();
        if !stream_ended_explicitly {
            log::warn!("!!! STREAM ENDED IMPLICITLY (exhausted) for model: {model} after {chunk_count} chunks, {total_content_length} chars, {final_elapsed:?} elapsed !!!");
        }

        log::info!(
            "=== STREAM COMPLETE === Model: {model}, Final: {chunk_count} chunks, {total_content_length} chars, {final_elapsed:?} elapsed"
        );

        // Add approximate token count
        self.add_tokens(total_content_length / 4).await; // Rough estimate

        if tx.send(EOT_SIGNAL.to_string()).is_err() {
            log::error!("!!! FAILED TO SEND EOT SIGNAL - channel may be closed !!!");
            return Err(anyhow::anyhow!("Channel closed during EOT signal send"));
        }

        log::info!(">>> EOT SIGNAL SENT for model: {model} <<<");
        Ok(())
    }

    /// Get a list of supported providers
    pub fn get_supported_providers() -> Vec<&'static str> {
        vec![
            "openai",
            "anthropic",
            "gemini",
            "groq",
            "cohere",
            "ollama",
            "xai",
            "deepseek",
        ]
    }

    /// Get all available providers
    pub async fn get_available_providers(&self) -> Result<Vec<String>> {
        Ok(Self::get_supported_providers()
            .iter()
            .map(|s| s.to_string())
            .collect())
    }

    /// Test if a model is available and working
    pub async fn test_model(&self, model: &str) -> Result<bool> {
        match self.generate_response_simple(model, "Hello").await {
            Ok(_) => {
                log::info!("Model {model} is available and working");
                Ok(true)
            }
            Err(e) => {
                log::warn!("Model {model} test failed: {e}");
                Ok(false)
            }
        }
    }

    /// Validate and get the best available model for a provider
    pub async fn validate_model(&self, model: &str, provider_type: Option<&str>) -> Result<String> {
        if self.test_model(model).await? {
            return Ok(model.to_string());
        }

        if let Some(provider) = provider_type {
            if let Ok(models) = self.get_available_models(provider).await {
                if !models.is_empty() {
                    log::info!("Model {} not available, using {} instead", model, models[0]);
                    return Ok(models[0].clone());
                }
            }
        }

        log::warn!("Could not validate model {model}, proceeding anyway");
        Ok(model.to_string())
    }
}

fn build_bound_client(adapter_kind: AdapterKind, provider_type: Option<&str>) -> Client {
    let mut builder = Client::builder().with_adapter_kind(adapter_kind);

    if let Some(base_url) = provider_type.and_then(provider_base_url_from_env) {
        let endpoint = normalize_base_url(&base_url);
        let target_resolver = ServiceTargetResolver::from_resolver_fn(
            move |mut service_target: ServiceTarget| -> genai::resolver::Result<ServiceTarget> {
                if service_target.model.adapter_kind == adapter_kind {
                    service_target.endpoint = Endpoint::from_owned(endpoint.clone());
                }
                Ok(service_target)
            },
        );
        builder = builder.with_service_target_resolver(target_resolver);
    }

    builder.build()
}

fn provider_base_url_env_var(provider: &str) -> Option<&'static str> {
    match provider.to_lowercase().as_str() {
        "openai" => Some("OPENAI_BASE_URL"),
        "anthropic" => Some("ANTHROPIC_BASE_URL"),
        "gemini" | "google" => Some("GEMINI_BASE_URL"),
        "groq" => Some("GROQ_BASE_URL"),
        "cohere" => Some("COHERE_BASE_URL"),
        "ollama" => Some("OLLAMA_BASE_URL"),
        "xai" => Some("XAI_BASE_URL"),
        "deepseek" => Some("DEEPSEEK_BASE_URL"),
        _ => None,
    }
}

fn provider_base_url_from_env(provider: &str) -> Option<String> {
    let env_var = provider_base_url_env_var(provider)?;

    std::env::var(env_var)
        .ok()
        .map(|value| value.trim().to_string())
        .filter(|value| !value.is_empty())
}

fn provider_api_key_env_var(provider: &str) -> Option<&'static str> {
    match provider.to_lowercase().as_str() {
        "openai" => Some("OPENAI_API_KEY"),
        "anthropic" => Some("ANTHROPIC_API_KEY"),
        "gemini" | "google" => Some("GEMINI_API_KEY"),
        "groq" => Some("GROQ_API_KEY"),
        "cohere" => Some("COHERE_API_KEY"),
        "xai" => Some("XAI_API_KEY"),
        "deepseek" => Some("DEEPSEEK_API_KEY"),
        _ => None,
    }
}

fn normalize_base_url(base_url: &str) -> String {
    if base_url.ends_with('/') {
        base_url.to_string()
    } else {
        format!("{base_url}/")
    }
}

/// Convert a provider string to genai AdapterKind
fn str_to_adapter_kind(provider: &str) -> Result<AdapterKind> {
    match provider.to_lowercase().as_str() {
        "openai" => Ok(AdapterKind::OpenAI),
        "anthropic" => Ok(AdapterKind::Anthropic),
        "gemini" | "google" => Ok(AdapterKind::Gemini),
        "groq" => Ok(AdapterKind::Groq),
        "cohere" => Ok(AdapterKind::Cohere),
        "ollama" => Ok(AdapterKind::Ollama),
        "xai" => Ok(AdapterKind::Xai),
        "deepseek" => Ok(AdapterKind::DeepSeek),
        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
    }
}

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

    #[test]
    fn test_str_to_adapter_kind() {
        assert!(str_to_adapter_kind("openai").is_ok());
        assert!(str_to_adapter_kind("anthropic").is_ok());
        assert!(str_to_adapter_kind("gemini").is_ok());
        assert!(str_to_adapter_kind("google").is_ok());
        assert!(str_to_adapter_kind("groq").is_ok());
        assert!(str_to_adapter_kind("cohere").is_ok());
        assert!(str_to_adapter_kind("ollama").is_ok());
        assert!(str_to_adapter_kind("xai").is_ok());
        assert!(str_to_adapter_kind("deepseek").is_ok());
        assert!(str_to_adapter_kind("invalid").is_err());
    }

    #[tokio::test]
    async fn test_provider_creation() {
        let provider = GenAIProvider::new();
        assert!(provider.is_ok());
    }

    #[tokio::test]
    async fn test_configured_provider_binds_adapter_for_custom_model_names() {
        let provider = GenAIProvider::new_with_config(Some("openai"), None).unwrap();
        let target = provider
            .client
            .resolve_service_target("gemma4-32b-it")
            .await
            .unwrap();

        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
    }

    #[tokio::test]
    async fn test_namespaced_model_resolves_on_unbound_client() {
        // genai-native namespacing must work without a bound client.
        let provider = GenAIProvider::new().unwrap();
        let target = provider
            .client
            .resolve_service_target("openai::phi-4-npu-ov")
            .await
            .unwrap();

        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
    }

    #[tokio::test]
    async fn test_from_config_binds_adapter_for_custom_model() {
        let config = Config {
            provider: Some("openai".to_string()),
            model: Some("phi-4-npu-ov".to_string()),
            ..Default::default()
        };
        let (provider, resolved) = GenAIProvider::from_config(&config, None).unwrap();
        assert_eq!(resolved.provider, "openai");
        assert_eq!(resolved.model, "phi-4-npu-ov");

        let target = provider
            .client
            .resolve_service_target(&resolved.model)
            .await
            .unwrap();
        assert_eq!(target.model.adapter_kind, AdapterKind::OpenAI);
    }

    #[test]
    fn test_from_config_model_precedence() {
        let config = Config {
            provider: Some("openai".to_string()),
            model: Some("config-model".to_string()),
            ..Default::default()
        };
        // CLI override wins over config model.
        let (_p, resolved) = GenAIProvider::from_config(&config, Some("cli-model")).unwrap();
        assert_eq!(resolved.model, "cli-model");
    }

    #[tokio::test]
    async fn test_openai_base_url_overrides_bound_provider_endpoint() {
        let previous = std::env::var("OPENAI_BASE_URL").ok();
        std::env::set_var("OPENAI_BASE_URL", "https://custom.example/v1");

        let provider = GenAIProvider::new_with_config(Some("openai"), None).unwrap();
        let target = provider
            .client
            .resolve_service_target("gemma4-32b-it")
            .await
            .unwrap();

        assert_eq!(target.endpoint.base_url(), "https://custom.example/v1/");

        match previous {
            Some(value) => std::env::set_var("OPENAI_BASE_URL", value),
            None => std::env::remove_var("OPENAI_BASE_URL"),
        }
    }

    #[test]
    fn test_normalize_base_url() {
        assert_eq!(
            normalize_base_url("https://custom.example/v1"),
            "https://custom.example/v1/"
        );
        assert_eq!(
            normalize_base_url("https://custom.example/v1/"),
            "https://custom.example/v1/"
        );
    }

    #[tokio::test]
    async fn test_provider_is_clonable() {
        let provider = GenAIProvider::new().unwrap();
        let _clone1 = provider.clone();
        let _clone2 = provider.clone();
        // All clones share the same underlying state
    }
}