opencrates 3.0.1

//! Provider abstractions for external services

use crate::utils::error::OpenCratesError;
use crate::utils::openai_agents::Usage;
use crate::utils::templates::CrateSpec;
use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Represents a request to an LLM provider for code generation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationRequest {
    /// The specification for the crate to be generated.
    pub spec: CrateSpec,
    /// The prompt or instruction for the generation task.
    #[serde(default)]
    pub prompt: Option<String>,
    /// The maximum number of tokens to generate.
    #[serde(default)]
    pub max_tokens: Option<u32>,
    /// The specific model to use for generation.
    #[serde(default)]
    pub model: Option<String>,
    /// The sampling temperature for the generation.
    #[serde(default)]
    pub temperature: Option<f32>,
    /// Additional context to provide for the generation.
    #[serde(default)]
    pub context: Option<String>,
}

/// Represents a response from an LLM provider.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GenerationResponse {
    /// The generated code or text preview.
    pub preview: String,
    /// The reason why the generation finished.
    pub finish_reason: Option<String>,
    /// Metrics associated with the generation.
    pub metrics: Usage,
}

/// Represents a single search result from a search provider.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SearchResult {
    /// The title of the search result.
    pub title: String,
    /// A brief description or summary of the result.
    pub description: String,
    /// The URL of the search result.
    pub url: String,
    /// A score indicating the relevance of the result.
    pub relevance_score: f32,
}

/// Defines options for a search provider query.
#[derive(Debug, Clone)]
pub struct SearchOptions {
    /// The maximum number of results to return.
    pub limit: usize,
    /// The maximum number of results to fetch internally.
    pub max_results: usize,
    /// The timeout for the search request.
    pub timeout: Option<Duration>,
    /// Whether to include code in the search results.
    pub include_code: bool,
    /// Whether to include snippets in the search results.
    pub include_snippets: bool,
    /// A filter for the programming language of the results.
    pub language_filter: Option<String>,
    /// Whether to filter out duplicate results.
    pub filter_duplicates: bool,
}

impl Default for SearchOptions {
    fn default() -> Self {
        Self {
            limit: 10,
            max_results: 10,
            timeout: Some(Duration::from_secs(30)),
            include_code: false,
            include_snippets: true,
            language_filter: None,
            filter_duplicates: true,
        }
    }
}

/// A trait for providers that can search for information.
#[async_trait]
pub trait SearchProvider: Send + Sync {
    /// Performs a general search query.
    async fn search(
        &self,
        query: &str,
        options: &SearchOptions,
    ) -> Result<Vec<SearchResult>, OpenCratesError>;
    /// Performs a search specifically for code.
    async fn search_code(
        &self,
        query: &str,
        language: Option<&str>,
    ) -> Result<Vec<SearchResult>, OpenCratesError>;
    /// Retrieves documentation for a specific crate.
    async fn get_documentation(&self, crate_name: &str) -> Result<String, OpenCratesError>;
    /// Checks the health of the search provider.
    async fn health_check(&self) -> Result<bool, OpenCratesError>;
    /// Returns the name of the search provider.
    fn name(&self) -> &str;
}

/// A trait for Language Model (LLM) providers that can generate code.
#[async_trait]
pub trait LLMProvider: Send + Sync {
    /// Generates code based on a given request.
    async fn generate(
        &self,
        request: &GenerationRequest,
    ) -> Result<GenerationResponse, OpenCratesError>;
    /// Checks the health of the LLM provider.
    async fn health_check(&self) -> Result<bool, OpenCratesError>;
    /// Returns the name of the LLM provider.
    fn name(&self) -> &str;
    /// Returns the provider as a `dyn Any` for downcasting.
    fn as_any(&self) -> &dyn std::any::Any;
}

/// A registry for all available LLM providers.
pub struct ProviderRegistry {
    providers: std::collections::HashMap<String, Box<dyn LLMProvider>>,
}

impl Default for ProviderRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl ProviderRegistry {
    /// Creates a new, empty `ProviderRegistry`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            providers: std::collections::HashMap::new(),
        }
    }

    /// Registers a new LLM provider with the given name.
    pub fn register(&mut self, name: &str, provider: Box<dyn LLMProvider>) -> Result<()> {
        if self.providers.contains_key(name) {
            return Err(anyhow::anyhow!(
                "Provider with name '{}' already registered",
                name
            ));
        }
        self.providers.insert(name.to_string(), provider);
        Ok(())
    }

    /// Retrieves a provider from the registry by name.
    #[must_use]
    pub fn get_provider(&self, name: &str) -> Option<&dyn LLMProvider> {
        self.providers.get(name).map(std::convert::AsRef::as_ref)
    }
}

/// A utility for selecting a provider from the registry.
pub struct ProviderSelector;

impl ProviderSelector {
    /// Selects a provider from the registry by name.
    pub fn select<'a>(registry: &'a ProviderRegistry, name: &str) -> Result<&'a dyn LLMProvider> {
        registry
            .get_provider(name)
            .ok_or_else(|| anyhow::anyhow!("Provider '{}' not found", name))
    }
}

/// A trait for legacy LLM providers, for backward compatibility.
#[async_trait]
pub trait LegacyLLMProvider: Send + Sync + Clone {
    /// Generates code based on a given request.
    async fn generate(
        &self,
        request: GenerationRequest,
    ) -> Result<GenerationResponse, OpenCratesError>;
    /// Sets the model for the provider to use.
    async fn set_model(&self, model: &str) -> Result<(), OpenCratesError>;
    /// Validates the API key for the provider.
    async fn validate_api_key(&self) -> Result<bool, OpenCratesError>;
}

// Re-exports
pub use aichat::AIChatIntegration;
pub use aider::AiderIntegration;
pub use codex::CodexProvider;
pub use openai::OpenAIProvider;
pub use websearch::WebSearchProvider;

/// Integrates with the `aichat` command-line tool.
pub mod aichat;
/// Integrates with the `aider` command-line tool.
pub mod aider;
/// Provides integration with the Codex API.
pub mod codex;
/// Provides enhanced OpenAI integration.
pub mod enhanced_openai;
/// Contains types related to AI models.
pub mod model_types;
/// Provides integration with the OpenAI API.
pub mod openai;
/// Provides web search functionality.
pub mod websearch;