rskit-ai 0.2.0-alpha.2

Shared AI vocabulary for rskit AI/ML crates
Documentation
//! Provider and model identity vocabulary.

use serde::{Deserialize, Serialize};

/// Provider identifier for a model.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum Provider {
    /// OpenAI.
    OpenAI,
    /// Anthropic.
    Anthropic,
    /// Google Gemini/Vertex.
    Google,
    /// Cohere.
    Cohere,
    /// Mistral.
    Mistral,
    /// Meta-hosted or Meta-native model family.
    Meta,
    /// AWS Bedrock.
    AWSBedrock,
    /// Azure OpenAI.
    AzureOpenAI,
    /// Ollama.
    Ollama,
    /// NVIDIA Triton.
    Triton,
    /// vLLM.
    Vllm,
    /// Hugging Face Text Generation Inference.
    Tgi,
    /// Unknown or private provider name.
    Custom(String),
}

/// Model capability declaration.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Capabilities {
    /// Whether streaming responses are supported.
    pub streaming: bool,
    /// Whether image inputs are supported.
    pub vision: bool,
    /// Whether audio inputs are supported.
    pub audio: bool,
    /// Whether tool use/function calling is supported.
    pub tool_use: bool,
    /// Whether JSON-mode/structured generation is supported.
    pub json_mode: bool,
    /// Whether reasoning token accounting is supported.
    pub reasoning_tokens: bool,
    /// Maximum input tokens accepted by the model.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_input_tokens: Option<u64>,
    /// Maximum output tokens generated by the model.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_output_tokens: Option<u64>,
}

/// Canonical model identity.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Model {
    /// Model name/identifier.
    pub name: String,
    /// Provider that serves the model.
    pub provider: Provider,
    /// Optional provider model version.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,
    /// Model capabilities.
    pub capabilities: Capabilities,
}