omniference 0.3.2

A multi-protocol inference engine with provider adapters
Documentation
use crate::types::{Modality, ModelCapabilities, ReasoningBudget};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
pub struct Limits {
	pub context: Option<u32>,
	pub input: Option<u32>,
	pub output: Option<u32>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ContextTier {
	pub min_context_tokens: u32,
	pub input: Option<f64>,
	pub output: Option<f64>,
	pub cache_read: Option<f64>,
	pub cache_write: Option<f64>,
	pub reasoning: Option<f64>,
	pub input_audio: Option<f64>,
	pub output_audio: Option<f64>,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct ModelPricing {
	pub input: f64,
	pub output: f64,
	pub cache_read: Option<f64>,
	pub cache_write: Option<f64>,
	pub reasoning: Option<f64>,
	pub input_audio: Option<f64>,
	pub output_audio: Option<f64>,
	#[serde(default)]
	pub tiers: Vec<ContextTier>,
}

#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
pub struct CatalogEntry {
	pub name: Option<String>,
	pub limits: Limits,
	pub input_modalities: Vec<Modality>,
	pub output_modalities: Vec<Modality>,
	pub capabilities: Vec<ModelCapabilities>,
	pub pricing: Option<ModelPricing>,
	pub reasoning_budget: Option<ReasoningBudget>,
	#[serde(default)]
	pub reasoning_disabled: bool,
	pub aliases: Vec<String>,
}

impl CatalogEntry {
	pub fn is_empty(&self) -> bool {
		self.name.is_none()
			&& self.limits == Limits::default()
			&& self.input_modalities.is_empty()
			&& self.output_modalities.is_empty()
			&& self.capabilities.is_empty()
			&& self.pricing.is_none()
			&& self.reasoning_budget.is_none()
			&& !self.reasoning_disabled
			&& self.aliases.is_empty()
	}

	pub fn merge(self, higher: CatalogEntry) -> CatalogEntry {
		let lower_reasoning_budget = (!self.reasoning_disabled)
			.then(|| {
				self.reasoning_budget
					.clone()
					.or_else(|| self.capabilities.iter().find_map(ReasoningBudget::from_legacy_capability))
			})
			.flatten();
		let higher_has_reasoning = higher.reasoning_budget.is_some() || higher.capabilities.iter().any(is_reasoning_capability);
		let reasoning_disabled = if higher.reasoning_disabled {
			true
		} else if higher_has_reasoning {
			false
		} else {
			self.reasoning_disabled
		};
		let mut aliases = self.aliases;
		for alias in higher.aliases {
			if !aliases.contains(&alias) {
				aliases.push(alias);
			}
		}

		let higher_only_updates_reasoning = !higher.capabilities.is_empty() && higher.capabilities.iter().all(is_reasoning_capability);
		let mut capabilities = if higher.capabilities.is_empty() {
			self.capabilities
		} else if higher_only_updates_reasoning {
			let mut capabilities = self.capabilities;
			capabilities.retain(|capability| !is_reasoning_capability(capability));
			capabilities.extend(higher.capabilities);
			capabilities
		} else {
			higher.capabilities
		};
		let reasoning_budget = if reasoning_disabled {
			capabilities.retain(|capability| !is_reasoning_capability(capability));
			None
		} else {
			higher
				.reasoning_budget
				.or_else(|| capabilities.iter().find_map(ReasoningBudget::from_legacy_capability))
				.or(lower_reasoning_budget)
		};
		capabilities.retain(|capability| ReasoningBudget::from_legacy_capability(capability).is_none());
		if let Some(capability) = reasoning_budget.as_ref().and_then(ReasoningBudget::legacy_capability) {
			capabilities.push(capability);
		}
		capabilities.sort_by_key(ModelCapabilities::as_str);
		capabilities.dedup();

		CatalogEntry {
			name: higher.name.or(self.name),
			limits: Limits {
				context: higher.limits.context.or(self.limits.context),
				input: higher.limits.input.or(self.limits.input),
				output: higher.limits.output.or(self.limits.output),
			},
			input_modalities: if higher.input_modalities.is_empty() {
				self.input_modalities
			} else {
				higher.input_modalities
			},
			output_modalities: if higher.output_modalities.is_empty() {
				self.output_modalities
			} else {
				higher.output_modalities
			},
			capabilities,
			pricing: higher.pricing.or(self.pricing),
			reasoning_budget,
			reasoning_disabled,
			aliases,
		}
	}
}

fn is_reasoning_capability(capability: &ModelCapabilities) -> bool {
	matches!(
		capability,
		ModelCapabilities::Reasoning
			| ModelCapabilities::ReasoningEffortNone
			| ModelCapabilities::ReasoningEffortMinimal
			| ModelCapabilities::ReasoningEffortLow
			| ModelCapabilities::ReasoningEffortMedium
			| ModelCapabilities::ReasoningEffortHigh
			| ModelCapabilities::ReasoningEffortXHigh
			| ModelCapabilities::ReasoningBudgetTokens_1024_32000
			| ModelCapabilities::ReasoningBudgetTokens_1024_64000
			| ModelCapabilities::ReasoningBudgetTokens_128_32768
			| ModelCapabilities::ReasoningBudgetTokens_128_24576
	)
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct CatalogKey {
	pub provider_slug: Option<String>,
	pub model_id: String,
}