[v4/p7-model-kinds] feat(providers): the remaining upstream model kinds

aimux-core/Cargo.toml

@@ -27,6 +27,7 @@ rand = { workspace = true }
 base64 = "0.22"
 jsonschema = { workspace = true }
 regex = "1"
+indexmap = { version = "2", features = ["serde"] }
 
 [dev-dependencies]
 tokio = { workspace = true, features = ["macros", "rt", "test-util"] }

aimux-core/src/custom_provider.rs

@@ -0,0 +1,158 @@
+//! `custom_provider`: a provider assembled from alias tables, with an optional
+//! fallback provider.
+//!
+//! Mirrors the AI SDK's `customProvider` (`ai/src/registry/custom-provider.ts`).
+//! A lookup consults the alias table of its modality first, then the fallback
+//! provider; with neither, it is [`AiMuxError::NoSuchModel`]. There are no
+//! string aliases that resolve through a global default provider: Rust has
+//! none.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use crate::embedding_model::EmbeddingModel;
+use crate::error::AiMuxError;
+use crate::evaluation_model::EvaluationModel;
+use crate::files_model::Files;
+use crate::image_model::ImageModel;
+use crate::language_model::LanguageModel;
+use crate::provider::Provider;
+use crate::reranking_model::RerankingModel;
+use crate::search_model::SearchModel;
+use crate::skills_model::Skills;
+use crate::speech_model::SpeechModel;
+use crate::transcription_model::TranscriptionModel;
+use crate::video_model::VideoModel;
+
+/// Options of [`custom_provider`]; every field is optional.
+#[derive(Default)]
+pub struct CustomProviderOptions {
+    pub language_models: HashMap<String, Arc<dyn LanguageModel>>,
+    pub embedding_models: HashMap<String, Arc<dyn EmbeddingModel>>,
+    pub image_models: HashMap<String, Arc<dyn ImageModel>>,
+    pub transcription_models: HashMap<String, Arc<dyn TranscriptionModel>>,
+    pub speech_models: HashMap<String, Arc<dyn SpeechModel>>,
+    pub reranking_models: HashMap<String, Arc<dyn RerankingModel>>,
+    pub video_models: HashMap<String, Arc<dyn VideoModel>>,
+    pub search_models: HashMap<String, Arc<dyn SearchModel>>,
+    pub evaluation_models: HashMap<String, Arc<dyn EvaluationModel>>,
+    /// A files interface for uploading files.
+    pub files: Option<Arc<dyn Files>>,
+    /// A skills interface for uploading skills.
+    pub skills: Option<Arc<dyn Skills>>,
+    /// Consulted when a requested model is not in the alias tables.
+    pub fallback_provider: Option<Arc<dyn Provider>>,
+}
+
+/// A provider built by [`custom_provider`].
+pub struct CustomProvider {
+    options: CustomProviderOptions,
+}
+
+/// Create a custom provider from alias tables and an optional fallback.
+#[must_use]
+pub fn custom_provider(options: CustomProviderOptions) -> CustomProvider {
+    CustomProvider { options }
+}
+
+/// The aliased model, else what the fallback offers, else `NoSuchModel`.
+fn resolve<T: ?Sized>(
+    models: &HashMap<String, Arc<T>>,
+    id: &str,
+    model_type: &str,
+    fallback: impl FnOnce() -> Option<Result<Arc<T>, AiMuxError>>,
+) -> Result<Arc<T>, AiMuxError> {
+    models
+        .get(id)
+        .map(|model| Ok(Arc::clone(model)))
+        .or_else(fallback)
+        .unwrap_or_else(|| Err(AiMuxError::no_such_model(id, model_type)))
+}
+
+impl Provider for CustomProvider {
+    fn language_model(&self, id: &str) -> Result<Arc<dyn LanguageModel>, AiMuxError> {
+        let o = &self.options;
+        resolve(&o.language_models, id, "languageModel", || {
+            o.fallback_provider.as_ref().map(|p| p.language_model(id))
+        })
+    }
+
+    fn embedding_model(&self, id: &str) -> Result<Arc<dyn EmbeddingModel>, AiMuxError> {
+        let o = &self.options;
+        resolve(&o.embedding_models, id, "embeddingModel", || {
+            o.fallback_provider.as_ref().map(|p| p.embedding_model(id))
+        })
+    }
+
+    fn image_model(&self, id: &str) -> Result<Arc<dyn ImageModel>, AiMuxError> {
+        let o = &self.options;
+        resolve(&o.image_models, id, "imageModel", || {
+            o.fallback_provider.as_ref().map(|p| p.image_model(id))
+        })
+    }
+
+    // The optional modalities are always defined here, as upstream: an id
+    // neither table nor fallback serves is `NoSuchModel`, never `None`.
+
+    fn transcription_model(
+        &self,
+        id: &str,
+    ) -> Option<Result<Arc<dyn TranscriptionModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(
+            &o.transcription_models,
+            id,
+            "transcriptionModel",
+            || o.fallback_provider.as_ref()?.transcription_model(id),
+        ))
+    }
+
+    fn speech_model(&self, id: &str) -> Option<Result<Arc<dyn SpeechModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(&o.speech_models, id, "speechModel", || {
+            o.fallback_provider.as_ref()?.speech_model(id)
+        }))
+    }
+
+    fn reranking_model(&self, id: &str) -> Option<Result<Arc<dyn RerankingModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(&o.reranking_models, id, "rerankingModel", || {
+            o.fallback_provider.as_ref()?.reranking_model(id)
+        }))
+    }
+
+    fn video_model(&self, id: &str) -> Option<Result<Arc<dyn VideoModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(&o.video_models, id, "videoModel", || {
+            o.fallback_provider.as_ref()?.video_model(id)
+        }))
+    }
+
+    fn search_model(&self, id: &str) -> Option<Result<Arc<dyn SearchModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(&o.search_models, id, "searchModel", || {
+            o.fallback_provider.as_ref()?.search_model(id)
+        }))
+    }
+
+    fn evaluation_model(&self, id: &str) -> Option<Result<Arc<dyn EvaluationModel>, AiMuxError>> {
+        let o = &self.options;
+        Some(resolve(&o.evaluation_models, id, "evaluationModel", || {
+            o.fallback_provider.as_ref()?.evaluation_model(id)
+        }))
+    }
+
+    fn files(&self) -> Option<Arc<dyn Files>> {
+        let o = &self.options;
+        o.files
+            .clone()
+            .or_else(|| o.fallback_provider.as_ref()?.files())
+    }
+
+    fn skills(&self) -> Option<Arc<dyn Skills>> {
+        let o = &self.options;
+        o.skills
+            .clone()
+            .or_else(|| o.fallback_provider.as_ref()?.skills())
+    }
+}

aimux-core/src/evaluation_model.rs

@@ -0,0 +1,217 @@
+//! The `EvaluationModel` trait — the provider-facing interface for the
+//! experimental Choice / Score / Boolean evaluation.
+//!
+//! Aligned with Vercel AI SDK `EvaluationModelV4`
+//! (`provider/src/evaluation-model/v4/`).
+
+use std::collections::HashMap;
+
+use async_trait::async_trait;
+use indexmap::IndexMap;
+use serde::{Deserialize, Deserializer, Serialize};
+use serde_json::Value;
+use ts_rs::TS;
+
+use crate::AbortSignal;
+use crate::error::AiMuxError;
+use crate::shared::{
+    JsonObject, ResponseInfo, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning,
+};
+
+/// Shared state or structured instructions for an evaluation
+/// (`EvaluationModelV4Input`: a string, a JSON object or a JSON array).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+#[serde(untagged)]
+#[ts(export)]
+pub enum EvaluationInput {
+    Text(String),
+    #[ts(type = "Record<string, unknown>")]
+    Object(JsonObject),
+    #[ts(type = "Array<unknown>")]
+    Array(Vec<Value>),
+}
+
+/// The kind of a question (`EvaluationModelV4Question['type']`).
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
+#[serde(rename_all = "lowercase")]
+#[ts(export)]
+pub enum EvaluationQuestionType {
+    Choice,
+    Score,
+    Boolean,
+}
+
+impl EvaluationQuestionType {
+    /// The wire name (`"choice"`, `"score"`, `"boolean"`).
+    #[must_use]
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Choice => "choice",
+            Self::Score => "score",
+            Self::Boolean => "boolean",
+        }
+    }
+}
+
+/// The optional descriptions of a Boolean question's outcomes. An absent key
+/// and an explicit `null` differ: `Some(None)` is `null` (no description).
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TS)]
+#[ts(export)]
+pub struct BooleanCriteria {
+    #[serde(
+        rename = "true",
+        default,
+        deserialize_with = "present",
+        skip_serializing_if = "Option::is_none"
+    )]
+    #[ts(optional)]
+    pub true_: Option<Option<EvaluationInput>>,
+    #[serde(
+        rename = "false",
+        default,
+        deserialize_with = "present",
+        skip_serializing_if = "Option::is_none"
+    )]
+    #[ts(optional)]
+    pub false_: Option<Option<EvaluationInput>>,
+}
+
+fn present<'de, D, T>(deserializer: D) -> Result<Option<Option<T>>, D::Error>
+where
+    D: Deserializer<'de>,
+    T: Deserialize<'de>,
+{
+    Option::<T>::deserialize(deserializer).map(Some)
+}
+
+/// A judgment to make about the shared state (`EvaluationModelV4Question`).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+#[serde(tag = "type", rename_all = "lowercase")]
+#[ts(export)]
+pub enum EvaluationQuestion {
+    Choice {
+        instructions: EvaluationInput,
+        /// Nonempty map of option names to descriptions; `None` means no
+        /// description. Insertion order is the option order.
+        #[ts(as = "std::collections::HashMap<String, Option<EvaluationInput>>")]
+        criteria: IndexMap<String, Option<EvaluationInput>>,
+    },
+    Score {
+        instructions: EvaluationInput,
+        /// At least two ordered levels, indexed from zero.
+        criteria: Vec<Option<EvaluationInput>>,
+    },
+    Boolean {
+        instructions: EvaluationInput,
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        criteria: Option<BooleanCriteria>,
+    },
+}
+
+impl EvaluationQuestion {
+    /// The kind of this question.
+    #[must_use]
+    pub fn question_type(&self) -> EvaluationQuestionType {
+        match self {
+            Self::Choice { .. } => EvaluationQuestionType::Choice,
+            Self::Score { .. } => EvaluationQuestionType::Score,
+            Self::Boolean { .. } => EvaluationQuestionType::Boolean,
+        }
+    }
+}
+
+/// Options passed to [`EvaluationModel::do_evaluate`]
+/// (`EvaluationModelV4CallOptions`).
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+#[ts(export)]
+pub struct EvaluationCallOptions {
+    /// One shared state, even when the value is an array.
+    pub state: EvaluationInput,
+    /// Questions by id; insertion order is the question order.
+    #[ts(as = "std::collections::HashMap<String, EvaluationQuestion>")]
+    pub questions: IndexMap<String, EvaluationQuestion>,
+    /// Abort signal for cancelling the evaluation.
+    #[serde(skip)]
+    #[ts(skip)]
+    pub abort_signal: Option<AbortSignal>,
+    pub headers: Option<SharedHeaders>,
+    pub provider_options: Option<SharedProviderOptions>,
+}
+
+/// The answer to one question (`EvaluationModelV4Answer`).
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+#[serde(tag = "type", rename_all = "lowercase")]
+#[ts(export)]
+pub enum EvaluationAnswer {
+    Choice {
+        /// The selected option, with maximal probability when a distribution
+        /// exists.
+        choice: String,
+        /// Complete distribution over the question's options, when available.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        probabilities: Option<HashMap<String, f64>>,
+    },
+    Score {
+        /// Fractional position in `[0, levels - 1]`.
+        score: f64,
+        /// Complete distribution keyed by zero-based level indices as
+        /// strings; when supplied, `score` is its probability-weighted mean.
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        probabilities: Option<HashMap<String, f64>>,
+    },
+    Boolean {
+        /// Model-estimated P(true), in `[0, 1]`. Not confidence in either
+        /// outcome.
+        probability: f64,
+    },
+}
+
+/// Decimal places the provider rounded its output to; omit for full precision.
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TS)]
+#[ts(export)]
+pub struct EvaluationRounding {
+    pub probability_decimals: Option<u32>,
+    pub score_decimals: Option<u32>,
+}
+
+/// Token usage of an evaluation.
+#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, TS)]
+#[ts(export)]
+pub struct EvaluationUsage {
+    pub input_tokens: Option<u32>,
+    pub output_tokens: Option<u32>,
+}
+
+/// The result of [`EvaluationModel::do_evaluate`] (`EvaluationModelV4Result`).
+#[derive(Debug, Clone, Serialize, Deserialize, TS)]
+#[ts(export)]
+pub struct EvaluationResult {
+    /// Exactly one answer per question, under the original question ids.
+    #[ts(as = "std::collections::HashMap<String, EvaluationAnswer>")]
+    pub answers: IndexMap<String, EvaluationAnswer>,
+    pub rounding: Option<EvaluationRounding>,
+    pub usage: Option<EvaluationUsage>,
+    pub warnings: Vec<Warning>,
+    pub provider_metadata: Option<SharedProviderMetadata>,
+    pub response: Option<ResponseInfo>,
+}
+
+/// The experimental evaluation model contract. May change in patch releases.
+#[async_trait]
+pub trait EvaluationModel: Send + Sync {
+    /// Provider name, e.g. `"openai.evaluation"`.
+    fn provider(&self) -> &str;
+
+    /// Provider-specific model ID.
+    fn model_id(&self) -> &str;
+
+    /// Supported question types, used to reject unsupported calls before any
+    /// I/O.
+    fn supported_question_types(&self) -> &[EvaluationQuestionType];
+
+    /// Evaluate every question against the same state. No partial results.
+    async fn do_evaluate(
+        &self,
+        options: &EvaluationCallOptions,
+    ) -> Result<EvaluationResult, AiMuxError>;
+}