yomo 0.6.2

A QUIC-based runtime for AI-LLM tool routing and serverless execution
Documentation
use std::collections::HashMap;
use std::sync::Arc;

use crate::serve_config::ConfigError;
use crate::serve_config::ProviderConfig;
use crate::llm_provider::Provider;
use crate::llm_provider::openai::build_openai_provider;
use crate::llm_provider::selection::SelectionStrategy;

#[derive(Clone)]
pub struct ProviderEntry {
    pub provider_type: String,
    pub model_id: String,
    pub provider: Arc<dyn Provider>,
}

pub struct ProviderRegistry<M> {
    providers: HashMap<String, ProviderEntry>,
    strategy: Arc<dyn SelectionStrategy<M>>,
}

impl<M> ProviderRegistry<M> {
    pub fn from_providers(
        providers: &[ProviderConfig],
        strategy: Arc<dyn SelectionStrategy<M>>,
    ) -> Result<Self, ConfigError> {
        let mut registry: HashMap<String, ProviderEntry> = HashMap::new();
        let mut model_ids = std::collections::HashSet::new();
        for item in providers {
            if item.provider_type.trim().is_empty() {
                return Err(ConfigError::InvalidProvider(format!(
                    "provider type is required for {}",
                    item.model_id
                )));
            }
            if item.model_id.trim().is_empty() {
                return Err(ConfigError::InvalidProvider(
                    "model_id is required for provider".to_string(),
                ));
            }
            let normalized_model_id = item.model_id.to_ascii_lowercase();
            if !model_ids.insert(normalized_model_id) {
                return Err(ConfigError::InvalidProvider(format!(
                    "duplicate model_id: {}",
                    item.model_id
                )));
            }
        }

        for item in providers {
            let provider: Arc<dyn Provider> = match item.provider_type.as_str() {
                "openai" => Arc::new(build_openai_provider(&item.params)?),
                other => return Err(ConfigError::UnknownProviderType(other.to_string())),
            };

            let entry = ProviderEntry {
                provider_type: item.provider_type.clone(),
                model_id: item.model_id.clone(),
                provider,
            };
            registry.insert(item.model_id.clone(), entry);
        }

        Ok(Self::new(registry, strategy))
    }

    pub fn new(
        providers: HashMap<String, ProviderEntry>,
        strategy: Arc<dyn SelectionStrategy<M>>,
    ) -> Self {
        Self {
            providers,
            strategy,
        }
    }

    pub fn select(
        &self,
        model_id: Option<&str>,
        metadata: &M,
    ) -> Result<
        (crate::llm_provider::selection::SelectionResult, ProviderEntry),
        crate::llm_provider::selection::SelectionError,
    > {
        let selected = self
            .strategy
            .select(model_id, metadata)
            .map_err(|err| err)?;
        let provider = self
            .providers
            .values()
            .find(|provider| {
                provider.model_id.to_ascii_lowercase()
                    == selected.model_id.to_ascii_lowercase()
            })
            .cloned()
            .ok_or(crate::llm_provider::selection::SelectionError::ModelNotSupported)?;
        Ok((selected, provider))
    }

    pub fn providers(&self) -> &HashMap<String, ProviderEntry> {
        &self.providers
    }
}