File size: 2,845 Bytes
e5034c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | use derive_more::derive::Display;
use derive_setters::Setters;
use fake::Dummy;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use strum_macros::EnumString;
/// Represents input modalities that a model can accept
#[derive(
Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, EnumString, JsonSchema, Dummy,
)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase", ascii_case_insensitive)]
pub enum InputModality {
/// Text input (all models support this)
Text,
/// Image input (vision-capable models)
Image,
}
/// Default input modalities when not specified (text-only)
fn default_input_modalities() -> Vec<InputModality> {
vec![InputModality::Text]
}
#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Setters, JsonSchema, Dummy)]
#[setters(strip_option)]
pub struct Model {
pub id: ModelId,
pub name: Option<String>,
pub description: Option<String>,
pub context_length: Option<u64>,
// TODO: add provider information to the model
pub tools_supported: Option<bool>,
/// Whether the model supports parallel tool calls
pub supports_parallel_tool_calls: Option<bool>,
/// Whether the model supports reasoning
pub supports_reasoning: Option<bool>,
/// Input modalities supported by the model (defaults to text-only)
#[serde(default = "default_input_modalities")]
pub input_modalities: Vec<InputModality>,
}
#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct Parameters {
pub tool_supported: bool,
}
impl Parameters {
pub fn new(tool_supported: bool) -> Self {
Self { tool_supported }
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize, Hash, Eq, Display, JsonSchema, Dummy)]
#[serde(transparent)]
pub struct ModelId(String);
impl ModelId {
pub fn new<T: Into<String>>(id: T) -> Self {
Self(id.into())
}
}
impl Model {
/// Creates a new `Model` with the given id and default values for all other
/// fields.
pub fn new(id: impl Into<ModelId>) -> Self {
Self {
id: id.into(),
name: None,
description: None,
context_length: None,
tools_supported: None,
supports_parallel_tool_calls: None,
supports_reasoning: None,
input_modalities: default_input_modalities(),
}
}
}
impl From<String> for ModelId {
fn from(value: String) -> Self {
ModelId(value)
}
}
impl From<&str> for ModelId {
fn from(value: &str) -> Self {
ModelId(value.to_string())
}
}
impl ModelId {
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::str::FromStr for ModelId {
type Err = std::convert::Infallible;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Ok(ModelId(s.to_string()))
}
}
|