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    #[serde(rename = "v-0")]
11    #[default]
12    VZero,
13
14    #[serde(rename = "custom")]
15    Custom {
16        name: String,
17        /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
18        display_name: Option<String>,
19        max_tokens: u64,
20        max_output_tokens: Option<u64>,
21        max_completion_tokens: Option<u64>,
22    },
23}
24
25impl Model {
26    pub fn default_fast() -> Self {
27        Self::VZero
28    }
29
30    pub fn from_id(id: &str) -> Result<Self> {
31        match id {
32            "v-0" => Ok(Self::VZero),
33            invalid_id => anyhow::bail!("invalid model id '{invalid_id}'"),
34        }
35    }
36
37    pub fn id(&self) -> &str {
38        match self {
39            Self::VZero => "v-0",
40            Self::Custom { name, .. } => name,
41        }
42    }
43
44    pub fn display_name(&self) -> &str {
45        match self {
46            Self::VZero => "Vercel v0",
47            Self::Custom {
48                name, display_name, ..
49            } => display_name.as_ref().unwrap_or(name),
50        }
51    }
52
53    pub fn max_token_count(&self) -> u64 {
54        match self {
55            Self::VZero => 128_000,
56            Self::Custom { max_tokens, .. } => *max_tokens,
57        }
58    }
59
60    pub fn max_output_tokens(&self) -> Option<u64> {
61        match self {
62            Self::Custom {
63                max_output_tokens, ..
64            } => *max_output_tokens,
65            Self::VZero => Some(32_768),
66        }
67    }
68
69    /// Returns whether the given model supports the `parallel_tool_calls` parameter.
70    ///
71    /// If the model does not support the parameter, do not pass it up, or the API will return an error.
72    pub fn supports_parallel_tool_calls(&self) -> bool {
73        match self {
74            Self::VZero => true,
75            Model::Custom { .. } => false,
76        }
77    }
78}