use crate::primitives::{Address, Hash, Timestamp};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelInfo {
pub model_id: String,
pub name: String,
pub version: String,
pub description: String,
pub modality: ModelModality,
pub architecture: String,
pub provider: Address,
pub model_hash: Hash,
pub parameters: ModelParameters,
pub pricing: PricingConfig,
pub status: ModelStatus,
pub metadata: HashMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub moe: Option<MoeMetadata>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timeseries: Option<TimeseriesParameters>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub vision: Option<VisionParameters>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub audio: Option<AudioParameters>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub video: Option<VideoParameters>,
}
impl ModelInfo {
pub fn new(
model_id: String,
name: String,
version: String,
modality: ModelModality,
provider: Address,
) -> Self {
Self {
model_id,
name,
version,
description: String::new(),
modality,
architecture: String::new(),
provider,
model_hash: Hash::zero(),
parameters: ModelParameters::default(),
pricing: PricingConfig::default(),
status: ModelStatus::Pending,
metadata: HashMap::new(),
moe: None,
timeseries: None,
vision: None,
audio: None,
video: None,
}
}
pub fn with_moe(mut self, moe: MoeMetadata) -> Self {
self.moe = Some(moe);
self
}
pub fn with_timeseries(mut self, params: TimeseriesParameters) -> Self {
self.timeseries = Some(params);
self
}
pub fn with_vision(mut self, params: VisionParameters) -> Self {
self.vision = Some(params);
self
}
pub fn with_audio(mut self, params: AudioParameters) -> Self {
self.audio = Some(params);
self
}
pub fn with_video(mut self, params: VideoParameters) -> Self {
self.video = Some(params);
self
}
pub fn is_moe(&self) -> bool {
self.moe.is_some()
}
pub fn with_description(mut self, description: String) -> Self {
self.description = description;
self
}
pub fn with_architecture(mut self, architecture: String) -> Self {
self.architecture = architecture;
self
}
pub fn with_hash(mut self, hash: Hash) -> Self {
self.model_hash = hash;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub enum ModelModality {
#[default]
Text,
Image,
Audio,
Timeseries,
Video,
TextImage,
TextAudio,
Multimodal,
}
impl ModelModality {
pub fn supports(&self, requested: ModelModality) -> bool {
if *self == requested {
return true;
}
match *self {
ModelModality::Multimodal => !matches!(requested, ModelModality::Timeseries),
ModelModality::TextImage => matches!(requested, ModelModality::Text | ModelModality::Image),
ModelModality::TextAudio => matches!(requested, ModelModality::Text | ModelModality::Audio),
_ => false,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelParameters {
pub parameter_count: Option<u64>,
pub context_window: u32,
pub max_output_tokens: u32,
pub input_formats: Vec<String>,
pub output_formats: Vec<String>,
pub capabilities: Vec<String>,
}
impl Default for ModelParameters {
fn default() -> Self {
Self {
parameter_count: None,
context_window: 4096,
max_output_tokens: 2048,
input_formats: vec!["text".to_string()],
output_formats: vec!["text".to_string()],
capabilities: Vec::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MoeMetadata {
pub num_experts: u32,
pub experts_per_token: u8,
#[serde(default)]
pub shared_experts: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub params_per_expert_x10: Option<u32>,
pub routing_strategy: MoeRoutingStrategy,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub load_balance_coef_x10000: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub attention_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expert_specialization: Option<Vec<String>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub capacity_factor_x100: Option<u32>,
}
impl MoeMetadata {
pub fn new(
num_experts: u32,
experts_per_token: u8,
routing_strategy: MoeRoutingStrategy,
) -> Self {
Self {
num_experts,
experts_per_token,
shared_experts: 0,
params_per_expert_x10: None,
routing_strategy,
load_balance_coef_x10000: None,
attention_type: None,
expert_specialization: None,
capacity_factor_x100: None,
}
}
pub fn with_shared_experts(mut self, shared: u32) -> Self {
self.shared_experts = shared;
self
}
pub fn with_params_per_expert_x10(mut self, params_x10: u32) -> Self {
self.params_per_expert_x10 = Some(params_x10);
self
}
pub fn with_attention_type(mut self, attn: impl Into<String>) -> Self {
self.attention_type = Some(attn.into());
self
}
pub fn with_expert_specialization(mut self, labels: Vec<String>) -> Self {
self.expert_specialization = Some(labels);
self
}
pub fn with_capacity_factor_x100(mut self, cf_x100: u32) -> Self {
self.capacity_factor_x100 = Some(cf_x100);
self
}
pub fn active_experts_per_token(&self) -> u32 {
self.experts_per_token as u32 + self.shared_experts
}
pub fn total_routed_params_x10(&self) -> Option<u64> {
self.params_per_expert_x10
.map(|p| p as u64 * self.num_experts as u64)
}
pub fn active_params_per_token_x10(&self) -> Option<u64> {
self.params_per_expert_x10
.map(|p| p as u64 * self.active_experts_per_token() as u64)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MoeRoutingStrategy {
TopK,
TopP,
ExpertChoice,
Switch,
Soft,
Sinkhorn,
Hash,
Custom,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TimeseriesParameters {
pub context_length: u32,
pub max_horizon: u32,
pub n_quantiles: u32,
pub num_features: u32,
}
impl TimeseriesParameters {
pub fn univariate(context_length: u32, max_horizon: u32) -> Self {
Self {
context_length,
max_horizon,
n_quantiles: 1,
num_features: 1,
}
}
pub fn with_quantiles(mut self, n_quantiles: u32) -> Self {
self.n_quantiles = n_quantiles;
self
}
pub fn with_features(mut self, num_features: u32) -> Self {
self.num_features = num_features;
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct VisionParameters {
pub input_size: u32,
pub embedding_dim: u32,
pub normalization: String,
pub image_formats: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AudioParameters {
pub sample_rate: u32,
pub encoder_filename: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub decoder_filename: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub joiner_filename: Option<String>,
pub max_audio_seconds: u32,
#[serde(default)]
pub languages: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct VideoParameters {
pub frame_size: u32,
pub num_frames: u32,
pub fps: u32,
pub embedding_dim: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelStatus {
Pending,
Active,
Inactive,
Deprecated,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InferenceRequest {
pub request_id: String,
pub model_id: String,
pub requester: Address,
pub input: Vec<u8>,
pub parameters: InferenceParameters,
pub max_price: u64,
pub timestamp: Timestamp,
pub callback: Option<Address>,
}
impl InferenceRequest {
pub fn new(
model_id: String,
requester: Address,
input: Vec<u8>,
max_price: u64,
) -> Self {
Self {
request_id: uuid::Uuid::new_v4().to_string(),
model_id,
requester,
input,
parameters: InferenceParameters::default(),
max_price,
timestamp: Timestamp::now(),
callback: None,
}
}
pub fn with_parameters(mut self, parameters: InferenceParameters) -> Self {
self.parameters = parameters;
self
}
pub fn with_callback(mut self, callback: Address) -> Self {
self.callback = Some(callback);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InferenceParameters {
pub temperature: Option<u32>, pub top_p: Option<u32>, pub top_k: Option<u32>,
pub max_tokens: Option<u32>,
pub stop_sequences: Vec<String>,
pub custom: HashMap<String, String>,
}
impl Default for InferenceParameters {
fn default() -> Self {
Self {
temperature: Some(100), top_p: None,
top_k: None,
max_tokens: None,
stop_sequences: Vec::new(),
custom: HashMap::new(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InferenceResponse {
pub request_id: String,
pub response_id: String,
pub model_id: String,
pub provider: Address,
pub output: Vec<u8>,
pub metadata: InferenceMetadata,
pub price: u64,
pub timestamp: Timestamp,
#[serde(default = "default_synthetic_content")]
pub synthetic_content: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provenance: Option<ProvenanceManifest>,
}
fn default_synthetic_content() -> bool {
true
}
impl InferenceResponse {
pub fn new(
request_id: String,
model_id: String,
provider: Address,
output: Vec<u8>,
price: u64,
) -> Self {
Self {
request_id,
response_id: uuid::Uuid::new_v4().to_string(),
model_id,
provider,
output,
metadata: InferenceMetadata::default(),
price,
timestamp: Timestamp::now(),
synthetic_content: true,
provenance: None,
}
}
pub fn with_provenance(mut self, manifest: ProvenanceManifest) -> Self {
self.provenance = Some(manifest);
self
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProvenanceManifest {
pub content_hash: Hash,
pub model_id: String,
pub provider: Address,
pub signed_at: Timestamp,
pub assertion: String,
pub signer_public_key: Vec<u8>,
pub signature: Vec<u8>,
pub algorithm: String,
}
impl ProvenanceManifest {
pub fn canonical_preimage(&self) -> Vec<u8> {
let mut buf = Vec::with_capacity(
self.content_hash.0.len()
+ self.model_id.len()
+ self.provider.0.len()
+ 8
+ self.assertion.len(),
);
buf.extend_from_slice(&self.content_hash.0);
buf.extend_from_slice(self.model_id.as_bytes());
buf.extend_from_slice(&self.provider.0);
buf.extend_from_slice(&self.signed_at.as_millis().to_le_bytes());
buf.extend_from_slice(self.assertion.as_bytes());
buf
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct InferenceMetadata {
pub input_tokens: u32,
pub output_tokens: u32,
pub latency_ms: u64,
pub model_version: Option<String>,
pub finish_reason: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InferenceProvider {
pub address: Address,
pub name: String,
pub endpoint_url: Option<String>,
pub models: Vec<String>,
pub capacity: ProviderCapacity,
pub pricing: PricingConfig,
pub reputation: u64,
pub total_inferences: u64,
pub status: ProviderStatus,
pub registered_at: Timestamp,
}
impl InferenceProvider {
pub fn new(address: Address, name: String) -> Self {
Self {
address,
name,
endpoint_url: None,
models: Vec::new(),
capacity: ProviderCapacity::default(),
pricing: PricingConfig::default(),
reputation: 0,
total_inferences: 0,
status: ProviderStatus::Pending,
registered_at: Timestamp::now(),
}
}
pub fn with_endpoint_url(mut self, url: impl Into<String>) -> Self {
self.endpoint_url = Some(url.into());
self
}
pub fn add_model(&mut self, model_id: String) {
if !self.models.contains(&model_id) {
self.models.push(model_id);
}
}
pub fn serves_model(&self, model_id: &str) -> bool {
self.models.iter().any(|m| m == model_id)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProviderCapacity {
pub max_concurrent_requests: u32,
pub active_requests: u32,
pub requests_per_second: u32,
pub max_batch_size: u32,
}
impl Default for ProviderCapacity {
fn default() -> Self {
Self {
max_concurrent_requests: 10,
active_requests: 0,
requests_per_second: 100,
max_batch_size: 1,
}
}
}
impl ProviderCapacity {
pub fn has_capacity(&self) -> bool {
self.active_requests < self.max_concurrent_requests
}
pub fn utilization(&self) -> u8 {
if self.max_concurrent_requests == 0 {
0
} else {
((self.active_requests as f64 / self.max_concurrent_requests as f64) * 100.0) as u8
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProviderStatus {
Pending,
Active,
Inactive,
Suspended,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PricingConfig {
pub price_per_input_token: u64,
pub price_per_output_token: u64,
pub minimum_price: u64,
pub pricing_model: PricingModel,
}
impl Default for PricingConfig {
fn default() -> Self {
Self {
price_per_input_token: 10,
price_per_output_token: 20,
minimum_price: 100,
pricing_model: PricingModel::PerToken,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PricingModel {
PerToken,
PerRequest,
PerComputeTime,
Dynamic,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelServiceInstance {
pub instance_id: String,
pub model_id: String,
pub model_name: String,
pub provider_address: Address,
pub provider_name: String,
pub location: ModelLocation,
pub api_endpoint: String,
pub mcp_endpoint: String,
pub status: ServiceStatus,
pub parameters: String,
pub pricing: PricingConfig,
pub created_at: u64,
#[serde(default)]
pub last_seen: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub load_info: Option<ModelLoadInfo>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelLocation {
Local,
Network,
}
impl std::fmt::Display for ModelLocation {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Local => write!(f, "local"),
Self::Network => write!(f, "network"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ServiceStatus {
Online,
Offline,
Degraded,
}
impl std::fmt::Display for ServiceStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Online => write!(f, "online"),
Self::Offline => write!(f, "offline"),
Self::Degraded => write!(f, "degraded"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelLoadInfo {
pub active_requests: u32,
pub max_concurrent: u32,
pub utilization_percent: u8,
pub load_level: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
Text {
text: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<CacheControl>,
},
Thinking { thinking: String },
ToolUse {
id: String,
name: String,
input: serde_json::Value,
},
ToolResult {
tool_use_id: String,
content: ToolResultContent,
#[serde(default, skip_serializing_if = "Option::is_none")]
is_error: Option<bool>,
},
Image { source: ImageSource },
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CacheControl {
Ephemeral,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ToolResultContent {
Text(String),
Blocks(Vec<ContentBlock>),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ImageSource {
Base64 {
media_type: String,
data: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RichChatMessage {
pub role: String,
pub content: MessageContent,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessageContent {
Text(String),
Blocks(Vec<ContentBlock>),
}
impl MessageContent {
pub fn into_blocks(self) -> Vec<ContentBlock> {
match self {
MessageContent::Text(s) => vec![ContentBlock::Text {
text: s,
cache_control: None,
}],
MessageContent::Blocks(b) => b,
}
}
pub fn as_blocks(&self) -> std::borrow::Cow<'_, [ContentBlock]> {
match self {
MessageContent::Text(s) => std::borrow::Cow::Owned(vec![ContentBlock::Text {
text: s.clone(),
cache_control: None,
}]),
MessageContent::Blocks(b) => std::borrow::Cow::Borrowed(b),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolSchema {
pub name: String,
pub description: String,
pub input_schema: serde_json::Value,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
Low,
#[default]
Medium,
High,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum SystemPrompt {
Text(String),
Blocks(Vec<ContentBlock>),
}
impl SystemPrompt {
pub fn as_text(&self) -> String {
match self {
SystemPrompt::Text(s) => s.clone(),
SystemPrompt::Blocks(blocks) => blocks
.iter()
.filter_map(|b| match b {
ContentBlock::Text { text, .. } => Some(text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join(""),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
EndTurn,
MaxTokens,
StopSequence,
ToolUse,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RichUsage {
pub input_tokens: u32,
pub output_tokens: u32,
#[serde(default)]
pub cache_creation_input_tokens: u32,
#[serde(default)]
pub cache_read_input_tokens: u32,
}
#[cfg(test)]
mod rich_chat_tests {
use super::*;
#[test]
fn text_block_roundtrip() {
let b = ContentBlock::Text {
text: "hello".to_string(),
cache_control: None,
};
let json = serde_json::to_string(&b).unwrap();
assert_eq!(json, r#"{"type":"text","text":"hello"}"#);
let decoded: ContentBlock = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, b);
}
#[test]
fn thinking_block_roundtrip() {
let b = ContentBlock::Thinking {
thinking: "let me check".to_string(),
};
let json = serde_json::to_string(&b).unwrap();
assert!(json.contains(r#""type":"thinking""#));
let decoded: ContentBlock = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, b);
}
#[test]
fn tool_use_block_roundtrip() {
let b = ContentBlock::ToolUse {
id: "tu_01".to_string(),
name: "get_price".to_string(),
input: serde_json::json!({"pair": "TNZO/USD"}),
};
let json = serde_json::to_string(&b).unwrap();
assert!(json.contains(r#""type":"tool_use""#));
let decoded: ContentBlock = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, b);
}
#[test]
fn tool_result_string_content() {
let b = ContentBlock::ToolResult {
tool_use_id: "tu_01".to_string(),
content: ToolResultContent::Text("0.42".to_string()),
is_error: None,
};
let json = serde_json::to_string(&b).unwrap();
let decoded: ContentBlock = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, b);
}
#[test]
fn message_content_string_normalizes_to_text_block() {
let mc = MessageContent::Text("hello".to_string());
let blocks = mc.into_blocks();
assert_eq!(blocks.len(), 1);
match &blocks[0] {
ContentBlock::Text { text, .. } => assert_eq!(text, "hello"),
_ => panic!("expected text block"),
}
}
#[test]
fn message_content_accepts_string_or_blocks() {
let s: MessageContent = serde_json::from_str(r#""hello""#).unwrap();
assert!(matches!(s, MessageContent::Text(_)));
let b: MessageContent =
serde_json::from_str(r#"[{"type":"text","text":"hi"}]"#).unwrap();
assert!(matches!(b, MessageContent::Blocks(_)));
}
#[test]
fn stop_reason_serializes_snake_case() {
assert_eq!(serde_json::to_string(&StopReason::EndTurn).unwrap(), r#""end_turn""#);
assert_eq!(serde_json::to_string(&StopReason::ToolUse).unwrap(), r#""tool_use""#);
assert_eq!(
serde_json::to_string(&StopReason::MaxTokens).unwrap(),
r#""max_tokens""#
);
}
#[test]
fn reasoning_effort_serializes_lowercase() {
assert_eq!(serde_json::to_string(&ReasoningEffort::Low).unwrap(), r#""low""#);
assert_eq!(serde_json::to_string(&ReasoningEffort::High).unwrap(), r#""high""#);
}
#[test]
fn system_prompt_blocks_concatenate() {
let sp = SystemPrompt::Blocks(vec![
ContentBlock::Text {
text: "you are ".to_string(),
cache_control: None,
},
ContentBlock::Text {
text: "helpful".to_string(),
cache_control: Some(CacheControl::Ephemeral),
},
]);
assert_eq!(sp.as_text(), "you are helpful");
}
#[test]
fn full_rich_request_roundtrip() {
let json = r#"{
"role": "user",
"content": [
{"type": "text", "text": "What is TNZO trading at?"}
]
}"#;
let msg: RichChatMessage = serde_json::from_str(json).unwrap();
assert_eq!(msg.role, "user");
assert_eq!(msg.content.as_blocks().len(), 1);
}
#[test]
fn assistant_with_thinking_and_tool_use() {
let json = r#"{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "I should query the price oracle."},
{"type": "tool_use", "id": "tu_01", "name": "get_price", "input": {"pair": "TNZO/USD"}}
]
}"#;
let msg: RichChatMessage = serde_json::from_str(json).unwrap();
let blocks = msg.content.as_blocks();
assert_eq!(blocks.len(), 2);
assert!(matches!(&blocks[0], ContentBlock::Thinking { .. }));
assert!(matches!(&blocks[1], ContentBlock::ToolUse { .. }));
}
}
#[cfg(test)]
mod moe_tests {
use super::*;
#[test]
fn moe_metadata_mixtral_8x7b() {
let moe = MoeMetadata::new(8, 2, MoeRoutingStrategy::TopK)
.with_params_per_expert_x10(70) .with_attention_type("gqa");
assert_eq!(moe.num_experts, 8);
assert_eq!(moe.experts_per_token, 2);
assert_eq!(moe.shared_experts, 0);
assert_eq!(moe.active_experts_per_token(), 2);
assert_eq!(moe.total_routed_params_x10(), Some(560)); assert_eq!(moe.active_params_per_token_x10(), Some(140)); }
#[test]
fn moe_metadata_deepseek_shared_experts() {
let moe = MoeMetadata::new(64, 6, MoeRoutingStrategy::TopK)
.with_shared_experts(2)
.with_params_per_expert_x10(3); assert_eq!(moe.active_experts_per_token(), 8); assert_eq!(moe.active_params_per_token_x10(), Some(24)); }
#[test]
fn moe_metadata_specialization_roundtrip() {
let labels = vec!["math".to_string(), "code".to_string(), "reasoning".to_string()];
let moe = MoeMetadata::new(3, 1, MoeRoutingStrategy::Switch)
.with_expert_specialization(labels.clone());
assert_eq!(moe.expert_specialization.as_ref(), Some(&labels));
}
#[test]
fn model_info_moe_wiring() {
let info = ModelInfo::new(
"mixtral-8x7b".to_string(),
"Mixtral".to_string(),
"0.1".to_string(),
ModelModality::Text,
Address::zero(),
);
assert!(!info.is_moe());
let info = info.with_moe(MoeMetadata::new(8, 2, MoeRoutingStrategy::TopK));
assert!(info.is_moe());
assert_eq!(info.moe.as_ref().unwrap().num_experts, 8);
}
#[test]
fn moe_metadata_serde_json_omits_when_absent() {
let info = ModelInfo::new(
"dense-7b".to_string(),
"Dense".to_string(),
"0.1".to_string(),
ModelModality::Text,
Address::zero(),
);
let json = serde_json::to_string(&info).unwrap();
assert!(!json.contains("\"moe\""));
}
#[test]
fn moe_metadata_serde_roundtrip() {
let info = ModelInfo::new(
"mixtral-8x7b".to_string(),
"Mixtral".to_string(),
"0.1".to_string(),
ModelModality::Text,
Address::zero(),
)
.with_moe(
MoeMetadata::new(8, 2, MoeRoutingStrategy::TopK)
.with_params_per_expert_x10(70)
.with_attention_type("gqa")
.with_capacity_factor_x100(125),
);
let json = serde_json::to_string(&info).unwrap();
let decoded: ModelInfo = serde_json::from_str(&json).unwrap();
assert_eq!(decoded.moe, info.moe);
}
}