rerank_config.rs

// ABOUTME: Reranking provider configuration for CodeGraph
// ABOUTME: Supports Jina API-based reranking and Ollama chat-based reranking
use anyhow::Result;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Default)]
#[serde(rename_all = "snake_case")]
pub enum RerankProvider {
    /// Jina AI reranking API (jina-reranker-v3)
    Jina,
    /// Ollama chat-based reranking (e.g., Qwen3-Reranker)
    Ollama,
    /// No reranking (use HNSW scores directly)
    #[default]
    None,
}

impl std::fmt::Display for RerankProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Jina => write!(f, "jina"),
            Self::Ollama => write!(f, "ollama"),
            Self::None => write!(f, "none"),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct JinaRerankConfig {
    /// Jina reranking model
    pub model: String,
    /// Jina API key environment variable name
    pub api_key_env: String,
    /// Jina API base URL
    #[serde(default = "JinaRerankConfig::default_api_base")]
    pub api_base: String,
    /// Maximum number of retries for API requests
    #[serde(default = "JinaRerankConfig::default_max_retries")]
    pub max_retries: u32,
    /// Timeout for API requests (seconds)
    #[serde(default = "JinaRerankConfig::default_timeout_secs")]
    pub timeout_secs: u64,
}

impl JinaRerankConfig {
    fn default_api_base() -> String {
        "https://api.jina.ai/v1".to_string()
    }

    fn default_max_retries() -> u32 {
        3
    }

    fn default_timeout_secs() -> u64 {
        30
    }
}

impl Default for JinaRerankConfig {
    fn default() -> Self {
        Self {
            model: "jina-reranker-v3".to_string(),
            api_key_env: "JINA_API_KEY".to_string(),
            api_base: Self::default_api_base(),
            max_retries: Self::default_max_retries(),
            timeout_secs: Self::default_timeout_secs(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct OllamaRerankConfig {
    /// Ollama reranking model (e.g., "dengcao/Qwen3-Reranker-8B:Q3_K_M")
    /// Defaults to CODEGRAPH_OLLAMA_RERANK_MODEL or OLLAMA_RERANK_MODEL env var
    #[serde(default = "OllamaRerankConfig::default_model")]
    pub model: String,
    /// Ollama API base URL
    /// Defaults to CODEGRAPH_OLLAMA_URL or OLLAMA_URL env var
    #[serde(default = "OllamaRerankConfig::default_api_base")]
    pub api_base: String,
    /// Maximum number of retries for API requests
    #[serde(default = "OllamaRerankConfig::default_max_retries")]
    pub max_retries: u32,
    /// Timeout for API requests (seconds)
    #[serde(default = "OllamaRerankConfig::default_timeout_secs")]
    pub timeout_secs: u64,
    /// Temperature for chat completion (0.0 = deterministic)
    #[serde(default = "OllamaRerankConfig::default_temperature")]
    pub temperature: f32,
}

impl OllamaRerankConfig {
    fn default_model() -> String {
        std::env::var("CODEGRAPH_OLLAMA_RERANK_MODEL")
            .or_else(|_| std::env::var("OLLAMA_RERANK_MODEL"))
            .unwrap_or_else(|_| "dengcao/Qwen3-Reranker-8B:Q3_K_M".to_string())
    }

    fn default_api_base() -> String {
        std::env::var("CODEGRAPH_OLLAMA_URL")
            .or_else(|_| std::env::var("OLLAMA_URL"))
            .unwrap_or_else(|_| "http://localhost:11434".to_string())
    }

    fn default_max_retries() -> u32 {
        3
    }

    fn default_timeout_secs() -> u64 {
        30
    }

    fn default_temperature() -> f32 {
        0.0 // Deterministic for consistent reranking
    }
}

impl Default for OllamaRerankConfig {
    fn default() -> Self {
        Self {
            model: Self::default_model(),
            api_base: Self::default_api_base(),
            max_retries: Self::default_max_retries(),
            timeout_secs: Self::default_timeout_secs(),
            temperature: Self::default_temperature(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct RerankConfig {
    /// Reranking provider
    #[serde(default)]
    pub provider: RerankProvider,

    /// Number of top results to return after reranking
    #[serde(default = "RerankConfig::default_top_n")]
    pub top_n: usize,

    /// Jina-specific configuration
    #[serde(default)]
    pub jina: Option<JinaRerankConfig>,

    /// Ollama-specific configuration
    #[serde(default)]
    pub ollama: Option<OllamaRerankConfig>,
}

impl RerankConfig {
    fn default_top_n() -> usize {
        10
    }

    pub fn validate(&self) -> Result<()> {
        match &self.provider {
            RerankProvider::Jina => {
                anyhow::ensure!(
                    self.jina.is_some(),
                    "Jina configuration required when using Jina reranking provider"
                );
                if let Some(config) = &self.jina {
                    anyhow::ensure!(
                        !config.model.is_empty(),
                        "Jina reranking model name cannot be empty"
                    );
                    anyhow::ensure!(
                        !config.api_key_env.is_empty(),
                        "Jina API key environment variable name cannot be empty"
                    );
                }
            }
            RerankProvider::Ollama => {
                anyhow::ensure!(
                    self.ollama.is_some(),
                    "Ollama configuration required when using Ollama reranking provider"
                );
                if let Some(config) = &self.ollama {
                    anyhow::ensure!(
                        !config.model.is_empty(),
                        "Ollama reranking model name cannot be empty"
                    );
                }
            }
            RerankProvider::None => {
                // No validation needed for None
            }
        }

        Ok(())
    }

    pub fn for_jina(model: &str) -> Self {
        Self {
            provider: RerankProvider::Jina,
            top_n: Self::default_top_n(),
            jina: Some(JinaRerankConfig {
                model: model.to_string(),
                ..Default::default()
            }),
            ollama: None,
        }
    }

    pub fn for_ollama(model: &str) -> Self {
        Self {
            provider: RerankProvider::Ollama,
            top_n: Self::default_top_n(),
            jina: None,
            ollama: Some(OllamaRerankConfig {
                model: model.to_string(),
                ..Default::default()
            }),
        }
    }
}

impl Default for RerankConfig {
    fn default() -> Self {
        Self {
            provider: RerankProvider::default(),
            top_n: Self::default_top_n(),
            jina: None,
            ollama: None,
        }
    }
}

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

    #[test]
    fn test_default_rerank_config() {
        let config = RerankConfig::default();
        assert_eq!(config.provider, RerankProvider::None);
        assert_eq!(config.top_n, 10);
        assert!(config.jina.is_none());
        assert!(config.ollama.is_none());
    }

    #[test]
    fn test_jina_config_creation() {
        let config = RerankConfig::for_jina("jina-reranker-v3");
        assert_eq!(config.provider, RerankProvider::Jina);
        assert!(config.jina.is_some());
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_ollama_config_creation() {
        let config = RerankConfig::for_ollama("dengcao/Qwen3-Reranker-8B:Q3_K_M");
        assert_eq!(config.provider, RerankProvider::Ollama);
        assert!(config.ollama.is_some());
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_config_validation() {
        let mut config = RerankConfig {
            provider: RerankProvider::Jina,
            jina: None,
            ..RerankConfig::default()
        };
        assert!(config.validate().is_err());

        config.provider = RerankProvider::Ollama;
        config.ollama = None;
        assert!(config.validate().is_err());
    }
}