1use anyhow::Result;
 2use serde::{Deserialize, Serialize};
 3use strum::EnumIter;
 4
 5pub const VERCEL_API_URL: &str = "https://api.v0.dev/v1";
 6
 7#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
 8#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, EnumIter)]
 9pub enum Model {
10    #[default]
11    #[serde(rename = "v0-1.5-md")]
12    VZeroOnePointFiveMedium,
13    #[serde(rename = "custom")]
14    Custom {
15        name: String,
16        /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
17        display_name: Option<String>,
18        max_tokens: u64,
19        max_output_tokens: Option<u64>,
20        max_completion_tokens: Option<u64>,
21    },
22}
23
24impl Model {
25    pub fn default_fast() -> Self {
26        Self::VZeroOnePointFiveMedium
27    }
28
29    pub fn from_id(id: &str) -> Result<Self> {
30        match id {
31            "v0-1.5-md" => Ok(Self::VZeroOnePointFiveMedium),
32            invalid_id => anyhow::bail!("invalid model id '{invalid_id}'"),
33        }
34    }
35
36    pub fn id(&self) -> &str {
37        match self {
38            Self::VZeroOnePointFiveMedium => "v0-1.5-md",
39            Self::Custom { name, .. } => name,
40        }
41    }
42
43    pub fn display_name(&self) -> &str {
44        match self {
45            Self::VZeroOnePointFiveMedium => "v0-1.5-md",
46            Self::Custom {
47                name, display_name, ..
48            } => display_name.as_ref().unwrap_or(name),
49        }
50    }
51
52    pub fn max_token_count(&self) -> u64 {
53        match self {
54            Self::VZeroOnePointFiveMedium => 128_000,
55            Self::Custom { max_tokens, .. } => *max_tokens,
56        }
57    }
58
59    pub fn max_output_tokens(&self) -> Option<u64> {
60        match self {
61            Self::VZeroOnePointFiveMedium => Some(32_000),
62            Self::Custom {
63                max_output_tokens, ..
64            } => *max_output_tokens,
65        }
66    }
67
68    pub fn supports_parallel_tool_calls(&self) -> bool {
69        match self {
70            Self::VZeroOnePointFiveMedium => true,
71            Model::Custom { .. } => false,
72        }
73    }
74
75    pub fn supports_prompt_cache_key(&self) -> bool {
76        false
77    }
78}