llmprism 0.4.0

//! [`Registry`] -- where providers are registered and looked up by name.

use std::collections::HashMap;
use std::sync::Arc;

use crate::audio::{AudioInput, PendingSpeechToTextRequest, PendingTextToSpeechRequest};
use crate::embeddings::PendingEmbeddingsRequest;
use crate::error::Error;
use crate::images::PendingImagesRequest;
use crate::middleware::{MiddlewareProvider, ProviderMiddleware};
use crate::moderation::PendingModerationRequest;
use crate::provider::Provider;
use crate::rerank::PendingRerankRequest;
use crate::schema::ObjectSchema;
use crate::structured::PendingStructuredRequest;
use crate::text::PendingTextRequest;

/// Maps provider names (like `"openai"`) to the [`Provider`] that handles them.
///
/// This is the front door of the crate: your application code asks the registry
/// for a provider by name rather than constructing a provider directly. That one
/// level of indirection is what lets tests swap in a
/// [`FakeProvider`](crate::testing::FakeProvider) under the same name with no
/// changes anywhere else -- the code under test never knows the difference.
///
/// # Example
///
/// ```
/// # #[cfg(feature = "openai")]
/// # {
/// use llmprism::providers::openai::OpenAiProvider;
/// use llmprism::Registry;
///
/// let mut registry = Registry::new();
/// registry.register("openai", OpenAiProvider::new("sk-..."));
///
/// // Later, anywhere in your application:
/// let builder = registry.text("openai", "gpt-4o-mini");
/// assert!(builder.is_ok());
/// # }
/// ```
#[derive(Default, Clone)]
pub struct Registry {
    providers: HashMap<String, Arc<dyn Provider>>,
}

impl Registry {
    /// Creates an empty registry with no providers. Useful in tests, where you
    /// typically only want to register a [`FakeProvider`](crate::testing::FakeProvider).
    pub fn new() -> Self {
        Self::default()
    }

    /// Registers `provider` under `name`. If a provider was already registered
    /// under that name, it's replaced -- this is exactly how tests swap in a fake:
    /// register it under the same name your application code already uses (e.g.
    /// `"openai"`).
    pub fn register(&mut self, name: impl Into<String>, provider: impl Provider + 'static) {
        self.providers.insert(name.into(), Arc::new(provider));
    }

    /// Like [`register`](Self::register), but for when you already have a shared
    /// `Arc<dyn Provider>` -- for example, one reused across multiple registries.
    pub fn register_arc(&mut self, name: impl Into<String>, provider: Arc<dyn Provider>) {
        self.providers.insert(name.into(), provider);
    }

    /// Wraps the provider registered under `name` with `middleware`, replacing
    /// its registration in place -- every capability call made through this
    /// registry against `name` now goes through `middleware` first. See
    /// [`ProviderMiddleware`] for what a middleware can do.
    ///
    /// Middlewares compose: calling `wrap` again on a name that's already
    /// wrapped nests the new middleware around the existing one, so the most
    /// recently attached middleware sees a call first.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if nothing is registered under `name`.
    ///
    /// # Example
    ///
    /// ```
    /// use async_trait::async_trait;
    /// use llmprism::middleware::ProviderMiddleware;
    /// use llmprism::testing::{FakeProvider, FakeTextResponse};
    /// use llmprism::text::{Step, TextRequest};
    /// use llmprism::{Error, Provider, Registry};
    ///
    /// struct AddSystemPrompt(&'static str);
    ///
    /// #[async_trait]
    /// impl ProviderMiddleware for AddSystemPrompt {
    ///     async fn text_step(&self, mut request: TextRequest, next: &dyn Provider) -> Result<Step, Error> {
    ///         request.system_prompts.insert(0, self.0.to_string());
    ///         next.text_step(&request).await
    ///     }
    /// }
    ///
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Error> {
    /// let mut registry = Registry::new();
    /// registry.register("fake", FakeProvider::new("fake").respond_with(FakeTextResponse::new("hi")));
    /// registry.wrap("fake", AddSystemPrompt("Be concise."))?;
    ///
    /// registry.text("fake", "test-model")?.with_prompt("hello").generate().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn wrap(
        &mut self,
        name: &str,
        middleware: impl ProviderMiddleware + 'static,
    ) -> Result<(), Error> {
        let inner = self.provider(name)?;
        self.register_arc(name, Arc::new(MiddlewareProvider::new(inner, middleware)));
        Ok(())
    }

    /// Looks up the provider registered under `name`.
    ///
    /// You'll rarely need to call this directly -- prefer a capability method like
    /// [`text`](Self::text), which calls this for you and hands back a ready-to-use
    /// request builder.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if nothing is registered under `name`.
    /// Every name a provider is currently registered under, in no particular
    /// order -- useful for diagnostics, like confirming what `from_env()`
    /// actually found configured. (Plain text, not a doc link: `from_env`
    /// only exists with the `http` feature on, but this method doesn't
    /// require it.)
    pub fn provider_names(&self) -> Vec<&str> {
        self.providers.keys().map(String::as_str).collect()
    }

    pub fn provider(&self, name: &str) -> Result<Arc<dyn Provider>, Error> {
        self.providers
            .get(name)
            .cloned()
            .ok_or_else(|| Error::UnknownProvider(name.to_string()))
    }

    /// Starts a text-generation request against the provider registered under
    /// `provider_name`, targeting `model`. Chain `.with_*()` calls on the returned
    /// builder to describe the request, then call `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .text("openai", "gpt-4o-mini")?
    ///     .with_prompt("Say hello in one word.")
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn text(
        &self,
        provider_name: &str,
        model: impl Into<String>,
    ) -> Result<PendingTextRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingTextRequest::new(provider, model))
    }

    /// Starts a structured-output request against the provider registered
    /// under `provider_name`, targeting `model` and requiring the reply to
    /// match `schema`. Chain `.with_*()` calls on the returned builder to
    /// describe the conversation, then call `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::schema::{ObjectSchema, Schema, StringSchema};
    /// use llmprism::Registry;
    ///
    /// let schema = ObjectSchema::new("greeting")
    ///     .with_property(Schema::String(StringSchema::new("message")), true);
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .structured("openai", "gpt-4o-mini", schema)?
    ///     .with_prompt("Greet the user warmly.")
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn structured(
        &self,
        provider_name: &str,
        model: impl Into<String>,
        schema: ObjectSchema,
    ) -> Result<PendingStructuredRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingStructuredRequest::new(provider, model, schema))
    }

    /// Starts a moderation request against the provider registered under
    /// `provider_name`, targeting `model`. Add one or more `.with_input(...)`
    /// calls on the returned builder, then call `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't registered.
    /// If the provider itself has no moderation endpoint (Anthropic, for
    /// instance), `.generate()` on the returned builder will fail with
    /// [`Error::Unsupported`] instead.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .moderation("openai", "omni-moderation-latest")?
    ///     .with_input("Some user-submitted text.")
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn moderation(
        &self,
        provider_name: &str,
        model: impl Into<String>,
    ) -> Result<PendingModerationRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingModerationRequest::new(provider, model))
    }

    /// Starts an embeddings request against the provider registered under
    /// `provider_name`, targeting `model`. Add one or more
    /// `.with_input(...)` calls on the returned builder, then call
    /// `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't
    /// registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .embeddings("openai", "text-embedding-3-small")?
    ///     .with_input("The quick brown fox.")
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn embeddings(
        &self,
        provider_name: &str,
        model: impl Into<String>,
    ) -> Result<PendingEmbeddingsRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingEmbeddingsRequest::new(provider, model))
    }

    /// Starts a rerank request against the provider registered under
    /// `provider_name`, targeting `model`, scoring documents against
    /// `query`. Add one or more `.with_document(...)` calls on the returned
    /// builder, then call `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't
    /// registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "voyageai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .rerank("voyageai", "rerank-2.5", "What's the capital of France?")?
    ///     .with_document("Paris is the capital of France.")
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn rerank(
        &self,
        provider_name: &str,
        model: impl Into<String>,
        query: impl Into<String>,
    ) -> Result<PendingRerankRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingRerankRequest::new(provider, model, query))
    }

    /// Starts an image-generation request against the provider registered
    /// under `provider_name`, targeting `model` with the given `prompt`.
    /// Optionally chain `.with_*()` calls on the returned builder, then call
    /// `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't
    /// registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .images("openai", "dall-e-3", "A watercolor painting of a lighthouse.")?
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn images(
        &self,
        provider_name: &str,
        model: impl Into<String>,
        prompt: impl Into<String>,
    ) -> Result<PendingImagesRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingImagesRequest::new(provider, model, prompt))
    }

    /// Starts a text-to-speech request against the provider registered under
    /// `provider_name`, targeting `model`, speaking `input`. Optionally
    /// chain `.with_voice(...)` on the returned builder, then call
    /// `.generate()` to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't
    /// registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), llmprism::Error> {
    /// use llmprism::Registry;
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .text_to_speech("openai", "tts-1", "Hello, world!")?
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn text_to_speech(
        &self,
        provider_name: &str,
        model: impl Into<String>,
        input: impl Into<String>,
    ) -> Result<PendingTextToSpeechRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingTextToSpeechRequest::new(provider, model, input))
    }

    /// Starts a speech-to-text request against the provider registered under
    /// `provider_name`, targeting `model`, transcribing `audio`. Call
    /// `.generate()` on the returned builder to run it.
    ///
    /// # Errors
    ///
    /// Returns [`Error::UnknownProvider`] if `provider_name` isn't
    /// registered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # #[cfg(feature = "openai")]
    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
    /// use llmprism::audio::AudioInput;
    /// use llmprism::Registry;
    ///
    /// let audio = AudioInput::new(std::fs::read("recording.mp3")?, "audio/mpeg");
    ///
    /// let registry = Registry::from_env();
    /// let response = registry
    ///     .speech_to_text("openai", "whisper-1", audio)?
    ///     .generate()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn speech_to_text(
        &self,
        provider_name: &str,
        model: impl Into<String>,
        audio: AudioInput,
    ) -> Result<PendingSpeechToTextRequest, Error> {
        let provider = self.provider(provider_name)?;
        Ok(PendingSpeechToTextRequest::new(provider, model, audio))
    }

    /// Builds a registry from the first-party providers that have their API
    /// key set in the environment -- `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`,
    /// `GROQ_API_KEY`, `DEEPSEEK_API_KEY`, `MISTRAL_API_KEY`, `XAI_API_KEY`,
    /// `OPENROUTER_API_KEY`, `PERPLEXITY_API_KEY`, `ZAI_API_KEY`,
    /// `VOYAGEAI_API_KEY`, `ELEVENLABS_API_KEY`, `GEMINI_API_KEY` -- and whose
    /// Cargo feature is enabled. A provider
    /// whose key isn't set is simply left out, rather than causing an error;
    /// if you then try to use it, you'll get [`Error::UnknownProvider`] at
    /// the point of use, which is usually clearer than failing at startup.
    /// **Ollama is the one exception**: since it's typically a local server
    /// with no API key at all, it's always registered when the `ollama`
    /// feature is enabled (optionally pointed elsewhere with
    /// `OLLAMA_BASE_URL`), not gated on an env var being set.
    ///
    /// This is the fastest way to get started; reach for [`Registry::new`] plus
    /// manual [`register`](Self::register) calls when you want more control (a
    /// custom base URL, a provider that reads its key from somewhere other than an
    /// env var, and so on).
    #[cfg(feature = "http")]
    pub fn from_env() -> Self {
        // `mut` goes unused with `http` enabled but no actual provider
        // feature on top of it -- every block below is `#[cfg]`-gated on one,
        // so there's nothing to register in that combination. Real (if
        // pointless) to build; not worth a warning.
        #[allow(unused_mut)]
        let mut registry = Self::new();

        #[cfg(feature = "openai")]
        if let Ok(api_key) = std::env::var("OPENAI_API_KEY") {
            registry.register(
                "openai",
                crate::providers::openai::OpenAiProvider::new(api_key),
            );
        }

        #[cfg(feature = "anthropic")]
        if let Ok(api_key) = std::env::var("ANTHROPIC_API_KEY") {
            registry.register(
                "anthropic",
                crate::providers::anthropic::AnthropicProvider::new(api_key),
            );
        }

        #[cfg(feature = "groq")]
        if let Ok(api_key) = std::env::var("GROQ_API_KEY") {
            registry.register(
                "groq",
                crate::providers::openai_compatible::GroqProvider::new(api_key),
            );
        }

        #[cfg(feature = "deepseek")]
        if let Ok(api_key) = std::env::var("DEEPSEEK_API_KEY") {
            registry.register(
                "deepseek",
                crate::providers::openai_compatible::DeepSeekProvider::new(api_key),
            );
        }

        #[cfg(feature = "mistral")]
        if let Ok(api_key) = std::env::var("MISTRAL_API_KEY") {
            registry.register(
                "mistral",
                crate::providers::openai_compatible::MistralProvider::new(api_key),
            );
        }

        #[cfg(feature = "xai")]
        if let Ok(api_key) = std::env::var("XAI_API_KEY") {
            registry.register(
                "xai",
                crate::providers::openai_compatible::XaiProvider::new(api_key),
            );
        }

        #[cfg(feature = "openrouter")]
        if let Ok(api_key) = std::env::var("OPENROUTER_API_KEY") {
            registry.register(
                "openrouter",
                crate::providers::openai_compatible::OpenRouterProvider::new(api_key),
            );
        }

        #[cfg(feature = "perplexity")]
        if let Ok(api_key) = std::env::var("PERPLEXITY_API_KEY") {
            registry.register(
                "perplexity",
                crate::providers::openai_compatible::PerplexityProvider::new(api_key),
            );
        }

        #[cfg(feature = "zai")]
        if let Ok(api_key) = std::env::var("ZAI_API_KEY") {
            registry.register(
                "zai",
                crate::providers::openai_compatible::ZaiProvider::new(api_key),
            );
        }

        #[cfg(feature = "voyageai")]
        if let Ok(api_key) = std::env::var("VOYAGEAI_API_KEY") {
            registry.register(
                "voyageai",
                crate::providers::voyageai::VoyageAiProvider::new(api_key),
            );
        }

        // Unlike every provider above, Ollama needs no API key at all (it's
        // typically a local server), so it's always registered when the
        // feature is enabled rather than gated on an env var being set --
        // `OLLAMA_BASE_URL` only overrides *where* it points, not *whether*
        // it gets registered.
        #[cfg(feature = "ollama")]
        {
            let provider = match std::env::var("OLLAMA_BASE_URL") {
                Ok(base_url) => crate::providers::ollama::OllamaProvider::with_base_url(base_url),
                Err(_) => crate::providers::ollama::OllamaProvider::new(),
            };
            registry.register("ollama", provider);
        }

        #[cfg(feature = "elevenlabs")]
        if let Ok(api_key) = std::env::var("ELEVENLABS_API_KEY") {
            registry.register(
                "elevenlabs",
                crate::providers::elevenlabs::ElevenLabsProvider::new(api_key),
            );
        }

        #[cfg(feature = "gemini")]
        if let Ok(api_key) = std::env::var("GEMINI_API_KEY") {
            registry.register(
                "gemini",
                crate::providers::gemini::GeminiProvider::new(api_key),
            );
        }

        registry
    }
}