lash_core/
model.rs

1use std::num::NonZeroUsize;
2
3use crate::provider::ModelCapability;
4
5#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
6#[serde(deny_unknown_fields)]
7pub struct ModelSpec {
8    pub id: String,
9    #[serde(default, skip_serializing_if = "Option::is_none")]
10    pub variant: Option<String>,
11    pub limits: ModelLimits,
12    /// Host-supplied capability metadata: what effort levels the model exposes
13    /// and how they encode on the wire. Lash validates the requested variant
14    /// against this and threads it onto every request the provider receives.
15    #[serde(default, skip_serializing_if = "ModelCapability::is_empty")]
16    pub capability: ModelCapability,
17}
18
19impl ModelSpec {
20    pub fn new(id: impl Into<String>, context_window_tokens: NonZeroUsize) -> Self {
21        Self {
22            id: id.into(),
23            variant: None,
24            limits: ModelLimits {
25                context_window_tokens,
26                output_token_capacity: None,
27            },
28            capability: ModelCapability::default(),
29        }
30    }
31
32    pub fn with_limits(
33        id: impl Into<String>,
34        variant: Option<String>,
35        limits: ModelLimits,
36    ) -> Self {
37        Self {
38            id: id.into(),
39            variant,
40            limits,
41            capability: ModelCapability::default(),
42        }
43    }
44
45    pub fn with_variant(mut self, variant: Option<String>) -> Self {
46        self.variant = variant;
47        self
48    }
49
50    pub fn with_capability(mut self, capability: ModelCapability) -> Self {
51        self.capability = capability;
52        self
53    }
54
55    /// Build a spec from the prompt budget (`context_window_tokens` — the
56    /// maximum input the provider accepts for this model on this route) and the
57    /// optional output cap. The budget is what history pruning and the UI
58    /// measure against.
59    pub fn from_token_limits(
60        id: impl Into<String>,
61        variant: Option<String>,
62        context_window_tokens: usize,
63        output_token_capacity: Option<usize>,
64    ) -> Result<Self, String> {
65        Ok(Self::with_limits(
66            id,
67            variant,
68            ModelLimits::from_token_limits(context_window_tokens, output_token_capacity)?,
69        ))
70    }
71
72    pub fn context_window_tokens(&self) -> usize {
73        self.limits.context_window_tokens.get()
74    }
75}
76
77impl Default for ModelSpec {
78    fn default() -> Self {
79        Self::new(
80            String::new(),
81            NonZeroUsize::new(1).expect("one is non-zero"),
82        )
83    }
84}
85
86#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
87#[serde(deny_unknown_fields)]
88pub struct ModelLimits {
89    /// The prompt budget: the maximum input tokens the provider accepts for
90    /// this model on this route. History pruning and the UI measure against
91    /// this — not the model's total context (input + output), which would
92    /// over-budget by the whole response reservation.
93    pub context_window_tokens: NonZeroUsize,
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub output_token_capacity: Option<NonZeroUsize>,
96}
97
98impl ModelLimits {
99    pub fn from_token_limits(
100        context_window_tokens: usize,
101        output_token_capacity: Option<usize>,
102    ) -> Result<Self, String> {
103        Ok(Self {
104            context_window_tokens: nonzero_token_limit(
105                "context_window_tokens",
106                context_window_tokens,
107            )?,
108            output_token_capacity: optional_nonzero_token_limit(
109                "output_token_capacity",
110                output_token_capacity,
111            )?,
112        })
113    }
114}
115
116impl Default for ModelLimits {
117    fn default() -> Self {
118        Self {
119            context_window_tokens: NonZeroUsize::new(1).expect("one is non-zero"),
120            output_token_capacity: None,
121        }
122    }
123}
124
125fn nonzero_token_limit(name: &str, value: usize) -> Result<NonZeroUsize, String> {
126    NonZeroUsize::new(value).ok_or_else(|| format!("{name} must be greater than zero"))
127}
128
129fn optional_nonzero_token_limit(
130    name: &str,
131    value: Option<usize>,
132) -> Result<Option<NonZeroUsize>, String> {
133    value
134        .map(|value| nonzero_token_limit(name, value))
135        .transpose()
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn model_spec_constructors_preserve_identity_variant_and_limits() {
144        let limits = ModelLimits::from_token_limits(8_192, Some(1_024)).expect("valid limits");
145
146        let spec =
147            ModelSpec::with_limits("provider/model", Some("fast".to_string()), limits.clone());
148
149        assert_eq!(spec.id, "provider/model");
150        assert_eq!(spec.variant.as_deref(), Some("fast"));
151        assert_eq!(spec.limits, limits);
152
153        let changed = spec.clone().with_variant(Some("accurate".to_string()));
154        assert_eq!(changed.id, "provider/model");
155        assert_eq!(changed.variant.as_deref(), Some("accurate"));
156        assert_eq!(changed.limits, spec.limits);
157
158        let cleared = changed.with_variant(None);
159        assert_eq!(cleared.id, "provider/model");
160        assert_eq!(cleared.variant, None);
161        assert_eq!(cleared.context_window_tokens(), 8_192);
162    }
163
164    #[test]
165    fn model_token_limit_constructors_reject_zero_and_preserve_output_cap() {
166        let spec = ModelSpec::from_token_limits(
167            "provider/model",
168            Some("variant-a".to_string()),
169            200_000,
170            Some(4_096),
171        )
172        .expect("valid token limits");
173
174        assert_eq!(spec.id, "provider/model");
175        assert_eq!(spec.variant.as_deref(), Some("variant-a"));
176        assert_eq!(spec.context_window_tokens(), 200_000);
177        assert_eq!(
178            spec.limits.output_token_capacity.map(NonZeroUsize::get),
179            Some(4_096)
180        );
181
182        let context_error = ModelSpec::from_token_limits("bad-context", None, 0, Some(1))
183            .expect_err("zero context");
184        assert!(
185            context_error.contains("context_window_tokens"),
186            "context error should name the invalid field: {context_error}"
187        );
188
189        let output_error = ModelLimits::from_token_limits(1, Some(0)).expect_err("zero output cap");
190        assert!(
191            output_error.contains("output_token_capacity"),
192            "output error should name the invalid field: {output_error}"
193        );
194    }
195}
lash_core/model.rs

lash_core/
model.rs
