Skip to main content

tea_model/
request.rs

1use std::collections::BTreeSet;
2
3use serde_json::Value;
4use tea_protocol::{
5    CanonicalMessage, ContentBlock, ModelId, ProtocolMetadata, ReasoningEffort, TokenCount,
6};
7use thiserror::Error;
8
9use crate::{HostedToolKind, HostedToolOptions, ModelSpec};
10
11/// Maximum UTF-8 bytes in a system prompt.
12pub const MAX_SYSTEM_PROMPT_BYTES: usize = 1024 * 1024;
13/// Maximum canonical messages in one model request.
14pub const MAX_REQUEST_MESSAGES: usize = 4096;
15/// Maximum model-visible tools in one request.
16pub const MAX_MODEL_TOOLS: usize = 256;
17/// Maximum UTF-8 bytes in one model-visible tool description.
18pub const MAX_TOOL_DESCRIPTION_BYTES: usize = 16 * 1024;
19/// Maximum encoded JSON bytes in one tool input schema.
20pub const MAX_TOOL_SCHEMA_BYTES: usize = 256 * 1024;
21/// Maximum JSON nesting depth in one tool input schema.
22pub const MAX_TOOL_SCHEMA_DEPTH: usize = 32;
23
24/// Provider-neutral reasoning request.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct ReasoningOptions {
27    effort: ReasoningEffort,
28    budget_tokens: Option<TokenCount>,
29}
30
31impl ReasoningOptions {
32    /// Creates reasoning options with provider/model default token budgeting.
33    #[must_use]
34    pub const fn new(effort: ReasoningEffort) -> Self {
35        Self {
36            effort,
37            budget_tokens: None,
38        }
39    }
40
41    /// Adds a requested reasoning-token budget.
42    #[must_use]
43    pub const fn with_budget(mut self, budget_tokens: TokenCount) -> Self {
44        self.budget_tokens = Some(budget_tokens);
45        self
46    }
47
48    /// Returns the requested effort.
49    #[must_use]
50    pub const fn effort(self) -> ReasoningEffort {
51        self.effort
52    }
53
54    /// Returns the optional reasoning-token budget.
55    #[must_use]
56    pub const fn budget_tokens(self) -> Option<TokenCount> {
57        self.budget_tokens
58    }
59}
60
61/// Model-visible client function tool.
62#[derive(Debug, Clone, PartialEq)]
63pub struct FunctionToolDefinition {
64    name: String,
65    description: String,
66    input_schema: Value,
67}
68
69impl FunctionToolDefinition {
70    /// Creates a validated model-visible tool definition.
71    ///
72    /// # Errors
73    ///
74    /// Returns an error for invalid names/descriptions or non-object,
75    /// oversized, or excessively nested schemas.
76    pub fn new(
77        name: impl Into<String>,
78        description: impl Into<String>,
79        input_schema: Value,
80    ) -> Result<Self, ModelRequestError> {
81        let name = name.into();
82        let description = description.into();
83        validate_tool_contract(&name, &description, &input_schema)?;
84        Ok(Self {
85            name,
86            description,
87            input_schema,
88        })
89    }
90
91    /// Returns the canonical tool name.
92    #[must_use]
93    pub fn name(&self) -> &str {
94        &self.name
95    }
96
97    /// Returns the model-visible description.
98    #[must_use]
99    pub fn description(&self) -> &str {
100        &self.description
101    }
102
103    /// Returns the provider-neutral input JSON Schema.
104    #[must_use]
105    pub const fn input_schema(&self) -> &Value {
106        &self.input_schema
107    }
108}
109
110/// Model-visible provider-hosted tool.
111#[derive(Debug, Clone, PartialEq)]
112pub struct HostedToolDefinition {
113    name: String,
114    description: String,
115    input_schema: Value,
116    options: HostedToolOptions,
117}
118
119impl HostedToolDefinition {
120    fn new(
121        description: impl Into<String>,
122        input_schema: Value,
123        options: HostedToolOptions,
124    ) -> Result<Self, ModelRequestError> {
125        let name = options.kind().name().to_owned();
126        let description = description.into();
127        validate_tool_contract(&name, &description, &input_schema)?;
128        Ok(Self {
129            name,
130            description,
131            input_schema,
132            options,
133        })
134    }
135
136    /// Returns the canonical hosted tool name.
137    #[must_use]
138    pub fn name(&self) -> &str {
139        &self.name
140    }
141
142    /// Returns the model-visible description.
143    #[must_use]
144    pub fn description(&self) -> &str {
145        &self.description
146    }
147
148    /// Returns the stable input schema used by client fallback and accounting.
149    #[must_use]
150    pub const fn input_schema(&self) -> &Value {
151        &self.input_schema
152    }
153
154    /// Returns the required hosted capability kind.
155    #[must_use]
156    pub const fn kind(&self) -> HostedToolKind {
157        self.options.kind()
158    }
159
160    /// Returns portable hosted-tool options.
161    #[must_use]
162    pub const fn options(&self) -> &HostedToolOptions {
163        &self.options
164    }
165}
166
167/// One model-visible client function or provider-hosted tool definition.
168#[derive(Debug, Clone, PartialEq)]
169pub enum ModelToolDefinition {
170    /// A function call that the client must execute.
171    Function(FunctionToolDefinition),
172    /// A tool executed inside the provider response lifecycle.
173    Hosted(HostedToolDefinition),
174}
175
176impl ModelToolDefinition {
177    /// Creates a validated client function definition.
178    ///
179    /// # Errors
180    ///
181    /// Returns an error for invalid names/descriptions or non-object,
182    /// oversized, or excessively nested schemas.
183    pub fn new(
184        name: impl Into<String>,
185        description: impl Into<String>,
186        input_schema: Value,
187    ) -> Result<Self, ModelRequestError> {
188        FunctionToolDefinition::new(name, description, input_schema).map(Self::Function)
189    }
190
191    /// Creates a validated provider-hosted definition.
192    ///
193    /// # Errors
194    ///
195    /// Returns an error when the common model contract violates request bounds.
196    pub fn hosted(
197        description: impl Into<String>,
198        input_schema: Value,
199        options: HostedToolOptions,
200    ) -> Result<Self, ModelRequestError> {
201        HostedToolDefinition::new(description, input_schema, options).map(Self::Hosted)
202    }
203
204    /// Returns the canonical model-visible name.
205    #[must_use]
206    pub fn name(&self) -> &str {
207        match self {
208            Self::Function(tool) => &tool.name,
209            Self::Hosted(tool) => &tool.name,
210        }
211    }
212
213    /// Returns the model-visible description.
214    #[must_use]
215    pub fn description(&self) -> &str {
216        match self {
217            Self::Function(tool) => &tool.description,
218            Self::Hosted(tool) => &tool.description,
219        }
220    }
221
222    /// Returns the provider-neutral input JSON Schema.
223    #[must_use]
224    pub const fn input_schema(&self) -> &Value {
225        match self {
226            Self::Function(tool) => &tool.input_schema,
227            Self::Hosted(tool) => &tool.input_schema,
228        }
229    }
230
231    /// Returns this definition as a client function.
232    #[must_use]
233    pub const fn as_function(&self) -> Option<&FunctionToolDefinition> {
234        match self {
235            Self::Function(tool) => Some(tool),
236            Self::Hosted(_) => None,
237        }
238    }
239
240    /// Returns this definition as a hosted tool.
241    #[must_use]
242    pub const fn as_hosted(&self) -> Option<&HostedToolDefinition> {
243        match self {
244            Self::Function(_) => None,
245            Self::Hosted(tool) => Some(tool),
246        }
247    }
248
249    /// Returns the hosted capability kind, if any.
250    #[must_use]
251    pub const fn hosted_kind(&self) -> Option<HostedToolKind> {
252        match self {
253            Self::Function(_) => None,
254            Self::Hosted(tool) => Some(tool.kind()),
255        }
256    }
257}
258
259/// Immutable provider-neutral snapshot for one model request.
260#[derive(Debug, Clone, PartialEq)]
261pub struct ModelRequest {
262    model_id: ModelId,
263    system_prompt: Option<String>,
264    messages: Vec<CanonicalMessage>,
265    tools: Vec<ModelToolDefinition>,
266    allow_parallel_tool_calls: bool,
267    reasoning: Option<ReasoningOptions>,
268    max_output_tokens: Option<TokenCount>,
269    metadata: ProtocolMetadata,
270}
271
272impl ModelRequest {
273    /// Creates a request for a non-empty canonical transcript.
274    ///
275    /// # Errors
276    ///
277    /// Returns an error when messages are empty, too numerous, or invalid at
278    /// the canonical protocol wire boundary.
279    pub fn new(
280        model_id: ModelId,
281        messages: Vec<CanonicalMessage>,
282    ) -> Result<Self, ModelRequestError> {
283        validate_messages(&messages)?;
284        Ok(Self {
285            model_id,
286            system_prompt: None,
287            messages,
288            tools: Vec::new(),
289            allow_parallel_tool_calls: false,
290            reasoning: None,
291            max_output_tokens: None,
292            metadata: ProtocolMetadata::default(),
293        })
294    }
295
296    /// Adds a bounded non-empty system prompt.
297    ///
298    /// # Errors
299    ///
300    /// Returns an error when the prompt is empty, oversized, or contains a
301    /// null character.
302    pub fn with_system_prompt(
303        mut self,
304        system_prompt: impl Into<String>,
305    ) -> Result<Self, ModelRequestError> {
306        let system_prompt = system_prompt.into();
307        if system_prompt.is_empty()
308            || system_prompt.len() > MAX_SYSTEM_PROMPT_BYTES
309            || system_prompt.contains('\0')
310        {
311            return Err(ModelRequestError::InvalidSystemPrompt);
312        }
313        self.system_prompt = Some(system_prompt);
314        Ok(self)
315    }
316
317    /// Adds model-visible tools in deterministic source order.
318    ///
319    /// # Errors
320    ///
321    /// Returns an error for an oversized list or duplicate tool names.
322    pub fn with_tools(
323        mut self,
324        tools: Vec<ModelToolDefinition>,
325        allow_parallel: bool,
326    ) -> Result<Self, ModelRequestError> {
327        if tools.len() > MAX_MODEL_TOOLS {
328            return Err(ModelRequestError::TooManyTools);
329        }
330        let mut names = BTreeSet::new();
331        if tools.iter().any(|tool| !names.insert(tool.name())) {
332            return Err(ModelRequestError::DuplicateToolName);
333        }
334        self.tools = tools;
335        self.allow_parallel_tool_calls = allow_parallel;
336        Ok(self)
337    }
338
339    /// Adds provider-neutral reasoning options.
340    #[must_use]
341    pub const fn with_reasoning(mut self, reasoning: ReasoningOptions) -> Self {
342        self.reasoning = Some(reasoning);
343        self
344    }
345
346    /// Adds a requested output-token limit.
347    #[must_use]
348    pub const fn with_max_output_tokens(mut self, max_output_tokens: TokenCount) -> Self {
349        self.max_output_tokens = Some(max_output_tokens);
350        self
351    }
352
353    /// Adds bounded namespaced request metadata.
354    #[must_use]
355    pub fn with_metadata(mut self, metadata: ProtocolMetadata) -> Self {
356        self.metadata = metadata;
357        self
358    }
359
360    /// Validates this request against one advertised model.
361    ///
362    /// # Errors
363    ///
364    /// Returns an error when model identity, capability, or output limits do
365    /// not satisfy the request.
366    pub fn validate_for(&self, model: &ModelSpec) -> Result<(), ModelRequestError> {
367        if self.model_id != *model.model_id() {
368            return Err(ModelRequestError::ModelMismatch);
369        }
370        validate_messages(&self.messages)?;
371        let capabilities = model.capabilities();
372        if request_contains_image(&self.messages) && !capabilities.accepts_images() {
373            return Err(ModelRequestError::ImageInputUnsupported);
374        }
375        if let Some(reasoning) = self.reasoning {
376            let Some(profile) = model.reasoning_profile() else {
377                return Err(ModelRequestError::ReasoningUnsupported);
378            };
379            if !profile.supported_efforts().contains(&reasoning.effort()) {
380                return Err(ModelRequestError::ReasoningEffortUnsupported);
381            }
382        }
383        if self.tools.iter().any(|tool| tool.as_function().is_some())
384            && !capabilities.supports_tools()
385        {
386            return Err(ModelRequestError::ToolsUnsupported);
387        }
388        if self.tools.iter().any(|tool| {
389            tool.hosted_kind()
390                .is_some_and(|kind| !capabilities.supports_hosted_tool(kind))
391        }) {
392            return Err(ModelRequestError::HostedToolUnsupported);
393        }
394        if self.allow_parallel_tool_calls
395            && self.tools.iter().any(|tool| tool.as_function().is_some())
396            && !capabilities.supports_parallel_tool_calls()
397        {
398            return Err(ModelRequestError::ParallelToolsUnsupported);
399        }
400        let output_limit = self
401            .max_output_tokens
402            .unwrap_or_else(|| model.max_output_tokens());
403        if output_limit.get() == 0 || output_limit > model.max_output_tokens() {
404            return Err(ModelRequestError::OutputLimitUnsupported);
405        }
406        if self
407            .reasoning
408            .and_then(ReasoningOptions::budget_tokens)
409            .is_some_and(|budget| budget.get() == 0 || budget > output_limit)
410        {
411            return Err(ModelRequestError::ReasoningBudgetUnsupported);
412        }
413        Ok(())
414    }
415
416    /// Returns the selected model.
417    #[must_use]
418    pub const fn model_id(&self) -> &ModelId {
419        &self.model_id
420    }
421
422    /// Returns the optional system prompt.
423    #[must_use]
424    pub fn system_prompt(&self) -> Option<&str> {
425        self.system_prompt.as_deref()
426    }
427
428    /// Returns canonical messages in source order.
429    #[must_use]
430    pub fn messages(&self) -> &[CanonicalMessage] {
431        &self.messages
432    }
433
434    /// Returns model-visible tools in source order.
435    #[must_use]
436    pub fn tools(&self) -> &[ModelToolDefinition] {
437        &self.tools
438    }
439
440    /// Returns whether the model may request several tools in one response.
441    #[must_use]
442    pub const fn allow_parallel_tool_calls(&self) -> bool {
443        self.allow_parallel_tool_calls
444    }
445
446    /// Returns reasoning options.
447    #[must_use]
448    pub const fn reasoning(&self) -> Option<ReasoningOptions> {
449        self.reasoning
450    }
451
452    /// Returns the request-specific output limit.
453    #[must_use]
454    pub const fn max_output_tokens(&self) -> Option<TokenCount> {
455        self.max_output_tokens
456    }
457
458    /// Returns bounded request metadata.
459    #[must_use]
460    pub const fn metadata(&self) -> &ProtocolMetadata {
461        &self.metadata
462    }
463}
464
465/// Error returned when building or validating a model request.
466#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
467pub enum ModelRequestError {
468    /// At least one canonical message is required.
469    #[error("model request requires at least one message")]
470    EmptyMessages,
471    /// The canonical message count exceeds the request limit.
472    #[error("model request contains too many messages")]
473    TooManyMessages,
474    /// A message fails its protocol wire invariant.
475    #[error("model request contains an invalid canonical message")]
476    InvalidMessage,
477    /// System prompt is empty, oversized, or contains a null character.
478    #[error("system prompt is invalid")]
479    InvalidSystemPrompt,
480    /// Tool name is not canonical.
481    #[error("tool name is invalid")]
482    InvalidToolName,
483    /// Tool description is empty, oversized, or contains a null character.
484    #[error("tool description is invalid")]
485    InvalidToolDescription,
486    /// A portable web-search domain is invalid.
487    #[error("web-search domain must be a canonical lowercase hostname")]
488    InvalidWebSearchDomain,
489    /// A portable web-search policy exceeds the domain limit.
490    #[error("web-search domain policy contains too many domains")]
491    TooManyWebSearchDomains,
492    /// Portable allow and block domain policies cannot be combined.
493    #[error("web-search allowed and blocked domains are mutually exclusive")]
494    ConflictingWebSearchDomainFilters,
495    /// An approximate web-search location field is invalid.
496    #[error("web-search location is invalid")]
497    InvalidWebSearchLocation,
498    /// Tool schema root must be a JSON object value.
499    #[error("tool input schema must be a JSON object")]
500    ToolSchemaMustBeObject,
501    /// Tool schema must explicitly declare `type: object`.
502    #[error("tool input schema must declare object type")]
503    ToolSchemaMustDeclareObject,
504    /// Tool schema exceeds encoded-byte or nesting limits.
505    #[error("tool input schema exceeds supported bounds")]
506    ToolSchemaOutOfBounds,
507    /// Request contains too many tool definitions.
508    #[error("model request contains too many tools")]
509    TooManyTools,
510    /// Request contains duplicate tool names.
511    #[error("model request contains a duplicate tool name")]
512    DuplicateToolName,
513    /// Request selects a different model than the specification.
514    #[error("model request does not match model specification")]
515    ModelMismatch,
516    /// Request includes an image but model accepts text only.
517    #[error("model does not support image input")]
518    ImageInputUnsupported,
519    /// Request asks for reasoning from a model without reasoning support.
520    #[error("model does not support reasoning")]
521    ReasoningUnsupported,
522    /// Request asks for an effort not supported by the selected model.
523    #[error("model does not support the requested reasoning effort")]
524    ReasoningEffortUnsupported,
525    /// Request includes tools for a model without tool support.
526    #[error("model does not support tools")]
527    ToolsUnsupported,
528    /// Request allows parallel tools for a serial-tool model.
529    #[error("model does not support parallel tool calls")]
530    ParallelToolsUnsupported,
531    /// Request contains a hosted tool unsupported by the selected model.
532    #[error("model does not support a requested hosted tool")]
533    HostedToolUnsupported,
534    /// Request output limit is zero or exceeds the model limit.
535    #[error("requested output limit is unsupported")]
536    OutputLimitUnsupported,
537    /// Reasoning budget is zero or exceeds the selected output limit.
538    #[error("requested reasoning budget is unsupported")]
539    ReasoningBudgetUnsupported,
540}
541
542fn validate_messages(messages: &[CanonicalMessage]) -> Result<(), ModelRequestError> {
543    if messages.is_empty() {
544        return Err(ModelRequestError::EmptyMessages);
545    }
546    if messages.len() > MAX_REQUEST_MESSAGES {
547        return Err(ModelRequestError::TooManyMessages);
548    }
549    if messages
550        .iter()
551        .any(|message| serde_json::to_value(message).is_err())
552    {
553        return Err(ModelRequestError::InvalidMessage);
554    }
555    Ok(())
556}
557
558fn request_contains_image(messages: &[CanonicalMessage]) -> bool {
559    messages.iter().any(|message| {
560        let content = match message {
561            CanonicalMessage::User { content, .. }
562            | CanonicalMessage::Assistant { content, .. }
563            | CanonicalMessage::ToolResult { content, .. } => content,
564        };
565        content
566            .iter()
567            .any(|block| matches!(block, ContentBlock::Image { .. }))
568    })
569}
570
571fn validate_tool_name(value: &str) -> Result<(), ModelRequestError> {
572    let mut bytes = value.bytes();
573    if value.len() > 128
574        || !bytes.next().is_some_and(|byte| byte.is_ascii_lowercase())
575        || !bytes.all(|byte| {
576            byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'_' | b'-' | b'.')
577        })
578    {
579        return Err(ModelRequestError::InvalidToolName);
580    }
581    Ok(())
582}
583
584fn validate_tool_contract(
585    name: &str,
586    description: &str,
587    input_schema: &Value,
588) -> Result<(), ModelRequestError> {
589    validate_tool_name(name)?;
590    if description.is_empty()
591        || description.len() > MAX_TOOL_DESCRIPTION_BYTES
592        || description.contains('\0')
593    {
594        return Err(ModelRequestError::InvalidToolDescription);
595    }
596    let object = input_schema
597        .as_object()
598        .ok_or(ModelRequestError::ToolSchemaMustBeObject)?;
599    if object.get("type").and_then(Value::as_str) != Some("object") {
600        return Err(ModelRequestError::ToolSchemaMustDeclareObject);
601    }
602    validate_schema_bounds(input_schema)
603}
604
605fn validate_schema_bounds(value: &Value) -> Result<(), ModelRequestError> {
606    if serde_json::to_vec(value)
607        .map_err(|_| ModelRequestError::ToolSchemaOutOfBounds)?
608        .len()
609        > MAX_TOOL_SCHEMA_BYTES
610        || json_depth(value) > MAX_TOOL_SCHEMA_DEPTH
611    {
612        return Err(ModelRequestError::ToolSchemaOutOfBounds);
613    }
614    Ok(())
615}
616
617fn json_depth(value: &Value) -> usize {
618    match value {
619        Value::Array(values) => 1 + values.iter().map(json_depth).max().unwrap_or(0),
620        Value::Object(values) => 1 + values.values().map(json_depth).max().unwrap_or(0),
621        _ => 1,
622    }
623}