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
11pub const MAX_SYSTEM_PROMPT_BYTES: usize = 1024 * 1024;
13pub const MAX_REQUEST_MESSAGES: usize = 4096;
15pub const MAX_MODEL_TOOLS: usize = 256;
17pub const MAX_TOOL_DESCRIPTION_BYTES: usize = 16 * 1024;
19pub const MAX_TOOL_SCHEMA_BYTES: usize = 256 * 1024;
21pub const MAX_TOOL_SCHEMA_DEPTH: usize = 32;
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct ReasoningOptions {
27 effort: ReasoningEffort,
28 budget_tokens: Option<TokenCount>,
29}
30
31impl ReasoningOptions {
32 #[must_use]
34 pub const fn new(effort: ReasoningEffort) -> Self {
35 Self {
36 effort,
37 budget_tokens: None,
38 }
39 }
40
41 #[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 #[must_use]
50 pub const fn effort(self) -> ReasoningEffort {
51 self.effort
52 }
53
54 #[must_use]
56 pub const fn budget_tokens(self) -> Option<TokenCount> {
57 self.budget_tokens
58 }
59}
60
61#[derive(Debug, Clone, PartialEq)]
63pub struct FunctionToolDefinition {
64 name: String,
65 description: String,
66 input_schema: Value,
67}
68
69impl FunctionToolDefinition {
70 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 #[must_use]
93 pub fn name(&self) -> &str {
94 &self.name
95 }
96
97 #[must_use]
99 pub fn description(&self) -> &str {
100 &self.description
101 }
102
103 #[must_use]
105 pub const fn input_schema(&self) -> &Value {
106 &self.input_schema
107 }
108}
109
110#[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 #[must_use]
138 pub fn name(&self) -> &str {
139 &self.name
140 }
141
142 #[must_use]
144 pub fn description(&self) -> &str {
145 &self.description
146 }
147
148 #[must_use]
150 pub const fn input_schema(&self) -> &Value {
151 &self.input_schema
152 }
153
154 #[must_use]
156 pub const fn kind(&self) -> HostedToolKind {
157 self.options.kind()
158 }
159
160 #[must_use]
162 pub const fn options(&self) -> &HostedToolOptions {
163 &self.options
164 }
165}
166
167#[derive(Debug, Clone, PartialEq)]
169pub enum ModelToolDefinition {
170 Function(FunctionToolDefinition),
172 Hosted(HostedToolDefinition),
174}
175
176impl ModelToolDefinition {
177 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 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 #[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 #[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 #[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 #[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 #[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 #[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#[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 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 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 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 #[must_use]
341 pub const fn with_reasoning(mut self, reasoning: ReasoningOptions) -> Self {
342 self.reasoning = Some(reasoning);
343 self
344 }
345
346 #[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 #[must_use]
355 pub fn with_metadata(mut self, metadata: ProtocolMetadata) -> Self {
356 self.metadata = metadata;
357 self
358 }
359
360 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 #[must_use]
418 pub const fn model_id(&self) -> &ModelId {
419 &self.model_id
420 }
421
422 #[must_use]
424 pub fn system_prompt(&self) -> Option<&str> {
425 self.system_prompt.as_deref()
426 }
427
428 #[must_use]
430 pub fn messages(&self) -> &[CanonicalMessage] {
431 &self.messages
432 }
433
434 #[must_use]
436 pub fn tools(&self) -> &[ModelToolDefinition] {
437 &self.tools
438 }
439
440 #[must_use]
442 pub const fn allow_parallel_tool_calls(&self) -> bool {
443 self.allow_parallel_tool_calls
444 }
445
446 #[must_use]
448 pub const fn reasoning(&self) -> Option<ReasoningOptions> {
449 self.reasoning
450 }
451
452 #[must_use]
454 pub const fn max_output_tokens(&self) -> Option<TokenCount> {
455 self.max_output_tokens
456 }
457
458 #[must_use]
460 pub const fn metadata(&self) -> &ProtocolMetadata {
461 &self.metadata
462 }
463}
464
465#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
467pub enum ModelRequestError {
468 #[error("model request requires at least one message")]
470 EmptyMessages,
471 #[error("model request contains too many messages")]
473 TooManyMessages,
474 #[error("model request contains an invalid canonical message")]
476 InvalidMessage,
477 #[error("system prompt is invalid")]
479 InvalidSystemPrompt,
480 #[error("tool name is invalid")]
482 InvalidToolName,
483 #[error("tool description is invalid")]
485 InvalidToolDescription,
486 #[error("web-search domain must be a canonical lowercase hostname")]
488 InvalidWebSearchDomain,
489 #[error("web-search domain policy contains too many domains")]
491 TooManyWebSearchDomains,
492 #[error("web-search allowed and blocked domains are mutually exclusive")]
494 ConflictingWebSearchDomainFilters,
495 #[error("web-search location is invalid")]
497 InvalidWebSearchLocation,
498 #[error("tool input schema must be a JSON object")]
500 ToolSchemaMustBeObject,
501 #[error("tool input schema must declare object type")]
503 ToolSchemaMustDeclareObject,
504 #[error("tool input schema exceeds supported bounds")]
506 ToolSchemaOutOfBounds,
507 #[error("model request contains too many tools")]
509 TooManyTools,
510 #[error("model request contains a duplicate tool name")]
512 DuplicateToolName,
513 #[error("model request does not match model specification")]
515 ModelMismatch,
516 #[error("model does not support image input")]
518 ImageInputUnsupported,
519 #[error("model does not support reasoning")]
521 ReasoningUnsupported,
522 #[error("model does not support the requested reasoning effort")]
524 ReasoningEffortUnsupported,
525 #[error("model does not support tools")]
527 ToolsUnsupported,
528 #[error("model does not support parallel tool calls")]
530 ParallelToolsUnsupported,
531 #[error("model does not support a requested hosted tool")]
533 HostedToolUnsupported,
534 #[error("requested output limit is unsupported")]
536 OutputLimitUnsupported,
537 #[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}