rig_core/providers/huggingface/

client.rs

use crate::client::{
    self, BearerAuth, Capabilities, Capable, DebugExt, Nothing, Provider, ProviderBuilder,
    ProviderClient,
};
use crate::http_client;
#[cfg(feature = "image")]
use crate::image_generation::ImageGenerationError;
use crate::transcription::TranscriptionError;
use std::fmt::Debug;
use std::fmt::Display;

#[derive(Debug, Clone, PartialEq, Default)]
pub enum SubProvider {
    #[default]
    HFInference,
    Together,
    SambaNova,
    Fireworks,
    Hyperbolic,
    Nebius,
    Novita,
    Custom(String),
}

impl SubProvider {
    /// Get the chat completion endpoint for the SubProvider
    /// Required because Huggingface Inference requires the model
    /// in the url and in the request body.
    pub fn completion_endpoint(&self, _model: &str) -> String {
        "v1/chat/completions".to_string()
    }

    /// Get the transcription endpoint for the SubProvider
    /// Required because Huggingface Inference requires the model
    /// in the url and in the request body.
    pub fn transcription_endpoint(&self, model: &str) -> Result<String, TranscriptionError> {
        match self {
            SubProvider::HFInference => Ok(format!("/{model}")),
            _ => Err(TranscriptionError::ProviderError(format!(
                "transcription endpoint is not supported yet for {self}"
            ))),
        }
    }

    /// Get the image generation endpoint for the SubProvider
    /// Required because Huggingface Inference requires the model
    /// in the url and in the request body.
    #[cfg(feature = "image")]
    pub fn image_generation_endpoint(&self, model: &str) -> Result<String, ImageGenerationError> {
        match self {
            SubProvider::HFInference => Ok(format!("/{model}")),
            _ => Err(ImageGenerationError::ProviderError(format!(
                "image generation endpoint is not supported yet for {self}"
            ))),
        }
    }

    pub fn model_identifier(&self, model: &str) -> String {
        match self {
            // Fireworks addresses models by a fully-qualified id. Guard against
            // re-prefixing an already-qualified id (e.g. a per-request model
            // override that is already fully qualified) — the generic path
            // applies this to the resolved request model unconditionally, so
            // without the guard a qualified override would become an invalid
            // `accounts/fireworks/models/accounts/fireworks/models/...` id.
            SubProvider::Fireworks => {
                const FIREWORKS_PREFIX: &str = "accounts/fireworks/models/";
                if model.starts_with(FIREWORKS_PREFIX) {
                    model.to_string()
                } else {
                    format!("{FIREWORKS_PREFIX}{model}")
                }
            }
            _ => model.to_string(),
        }
    }
}

impl From<&str> for SubProvider {
    fn from(s: &str) -> Self {
        SubProvider::Custom(s.to_string())
    }
}

impl From<String> for SubProvider {
    fn from(value: String) -> Self {
        SubProvider::Custom(value)
    }
}

impl Display for SubProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let route = match self {
            SubProvider::HFInference => "hf-inference/models".to_string(),
            SubProvider::Together => "together".to_string(),
            SubProvider::SambaNova => "sambanova".to_string(),
            SubProvider::Fireworks => "fireworks-ai".to_string(),
            SubProvider::Hyperbolic => "hyperbolic".to_string(),
            SubProvider::Nebius => "nebius".to_string(),
            SubProvider::Novita => "novita".to_string(),
            SubProvider::Custom(route) => route.clone(),
        };

        write!(f, "{route}")
    }
}

// ================================================================
// Main Huggingface Client
// ================================================================
const HUGGINGFACE_API_BASE_URL: &str = "https://router.huggingface.co";

#[derive(Debug, Default, Clone)]
pub struct HuggingFaceExt {
    subprovider: SubProvider,
}

#[derive(Debug, Default, Clone)]
pub struct HuggingFaceBuilder {
    subprovider: SubProvider,
}

type HuggingFaceApiKey = BearerAuth;

pub type Client<H = reqwest::Client> = client::Client<HuggingFaceExt, H>;
pub type ClientBuilder<H = crate::markers::Missing> =
    client::ClientBuilder<HuggingFaceBuilder, HuggingFaceApiKey, H>;

impl Provider for HuggingFaceExt {
    type Builder = HuggingFaceBuilder;

    const VERIFY_PATH: &'static str = "/api/whoami-v2";
}

impl crate::providers::openai::completion::OpenAICompatibleProvider for HuggingFaceExt {
    const PROVIDER_NAME: &'static str = "huggingface";

    type StreamingUsage = crate::providers::openai::Usage;

    // Structured-output support varies by sub-provider; keep the
    // pre-migration behavior of dropping `output_schema` with a warning.
    const SUPPORTS_RESPONSE_FORMAT: bool = false;

    type Response = crate::providers::openai::CompletionResponse;

    // Chat completions live under the router's `/v1` while verification,
    // transcription, and image generation use root-relative paths, so the
    // prefix cannot live in the client base URL.
    fn completion_path(&self, _model: &str) -> String {
        self.subprovider.completion_endpoint(_model)
    }

    fn prepare_request(
        &self,
        request: &mut crate::providers::openai::completion::CompletionRequest,
    ) -> Result<(), crate::completion::CompletionError> {
        // Some sub-providers (Fireworks) address models through a qualified
        // identifier in the request body.
        request.model = self.subprovider.model_identifier(&request.model);
        Ok(())
    }
}

impl<H> Capabilities<H> for HuggingFaceExt {
    type Completion = Capable<super::completion::CompletionModel<H>>;
    type Embeddings = Nothing;
    type Transcription = Capable<super::transcription::TranscriptionModel<H>>;
    type ModelListing = Nothing;
    #[cfg(feature = "image")]
    type ImageGeneration = Capable<super::image_generation::ImageGenerationModel<H>>;

    #[cfg(feature = "audio")]
    type AudioGeneration = Nothing;
    type Rerank = Nothing;
}

impl DebugExt for HuggingFaceExt {
    fn fields(&self) -> impl Iterator<Item = (&'static str, &dyn Debug)> {
        std::iter::once(("subprovider", (&self.subprovider as &dyn Debug)))
    }
}

impl ProviderBuilder for HuggingFaceBuilder {
    type Extension<H>
        = HuggingFaceExt
    where
        H: http_client::HttpClientExt;
    type ApiKey = HuggingFaceApiKey;

    const BASE_URL: &'static str = HUGGINGFACE_API_BASE_URL;

    fn build<H>(
        builder: &client::ClientBuilder<Self, Self::ApiKey, H>,
    ) -> http_client::Result<Self::Extension<H>>
    where
        H: http_client::HttpClientExt,
    {
        Ok(HuggingFaceExt {
            subprovider: builder.ext().subprovider.clone(),
        })
    }
}

impl ProviderClient for Client {
    type Input = String;
    type Error = crate::client::ProviderClientError;

    /// Create a new Huggingface client from the `HUGGINGFACE_API_KEY` environment variable.
    fn from_env() -> Result<Self, Self::Error> {
        let api_key = crate::client::required_env_var("HUGGINGFACE_API_KEY")?;

        Self::new(&api_key).map_err(Into::into)
    }

    fn from_val(input: Self::Input) -> Result<Self, Self::Error> {
        Self::new(&input).map_err(Into::into)
    }
}

impl<H> ClientBuilder<H> {
    pub fn subprovider(mut self, subprovider: SubProvider) -> Self {
        *self.ext_mut() = HuggingFaceBuilder { subprovider };
        self
    }
}

impl<H> Client<H> {
    pub(crate) fn subprovider(&self) -> &SubProvider {
        &self.ext().subprovider
    }
}
#[cfg(test)]
mod tests {
    use super::SubProvider;

    #[test]
    fn test_client_initialization() {
        let _client =
            crate::providers::huggingface::Client::new("dummy-key").expect("Client::new() failed");
        let _client_from_builder = crate::providers::huggingface::Client::builder()
            .api_key("dummy-key")
            .build()
            .expect("Client::builder() failed");
    }

    #[test]
    fn fireworks_model_identifier_is_idempotent() {
        // A bare id is qualified once...
        assert_eq!(
            SubProvider::Fireworks.model_identifier("deepseek-v3"),
            "accounts/fireworks/models/deepseek-v3"
        );
        // ...and an already-qualified id (e.g. a per-request model override)
        // is left untouched rather than double-prefixed.
        assert_eq!(
            SubProvider::Fireworks.model_identifier("accounts/fireworks/models/deepseek-v3"),
            "accounts/fireworks/models/deepseek-v3"
        );
        // Other sub-providers pass the id through verbatim.
        assert_eq!(
            SubProvider::HFInference.model_identifier("meta-llama/Llama-3.1-8B"),
            "meta-llama/Llama-3.1-8B"
        );
    }
}