1use std::collections::BTreeSet;
4
5use serde::{Deserialize, Serialize};
6use serde_json::{Map, Value};
7
8#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
10#[serde(rename_all = "snake_case")]
11pub enum ProtocolFamily {
12 OpenAiChatCompletions,
14 OpenAiResponses,
16 AnthropicMessages,
18 GeminiGenerateContent,
20 BedrockConverse,
22}
23
24#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
26#[serde(rename_all = "snake_case")]
27pub enum NativeToolKind {
28 WebSearch,
30 CodeExecution,
32 FileSearch,
34 ImageGeneration,
36 McpServer,
38 ToolSearch,
40 GoogleSearch,
42}
43
44impl NativeToolKind {
45 #[must_use]
47 pub const fn as_str(self) -> &'static str {
48 match self {
49 Self::WebSearch => "web_search",
50 Self::CodeExecution => "code_execution",
51 Self::FileSearch => "file_search",
52 Self::ImageGeneration => "image_generation",
53 Self::McpServer => "mcp_server",
54 Self::ToolSearch => "tool_search",
55 Self::GoogleSearch => "google_search",
56 }
57 }
58
59 #[must_use]
61 pub fn from_tool_type(tool_type: &str) -> Option<Self> {
62 match tool_type {
63 "web_search" | "web_search_preview" => Some(Self::WebSearch),
64 "code_execution" | "code_interpreter" => Some(Self::CodeExecution),
65 "file_search" => Some(Self::FileSearch),
66 "image_generation" => Some(Self::ImageGeneration),
67 "mcp_server" => Some(Self::McpServer),
68 "tool_search" => Some(Self::ToolSearch),
69 "google_search" => Some(Self::GoogleSearch),
70 _ => None,
71 }
72 }
73}
74
75fn default_prompted_output_template() -> String {
76 "Always respond with a JSON object that matches this schema:\n\n{schema}\n\nDon't include any text or Markdown fencing before or after."
77 .to_string()
78}
79
80fn default_thinking_tags() -> (String, String) {
81 ("<think>".to_string(), "</think>".to_string())
82}
83
84#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
86#[serde(rename_all = "snake_case")]
87pub enum JsonSchemaTransformer {
88 InlineDefinitions,
90}
91
92impl JsonSchemaTransformer {
93 #[must_use]
95 pub fn transform_schema(self, schema: &Value) -> Value {
96 match self {
97 Self::InlineDefinitions => inline_schema_definitions(schema),
98 }
99 }
100}
101
102fn inline_schema_definitions(schema: &Value) -> Value {
103 let mut schema = schema.clone();
104 let definitions = schema
105 .get("$defs")
106 .or_else(|| schema.get("definitions"))
107 .and_then(Value::as_object)
108 .cloned()
109 .unwrap_or_default();
110 if !definitions.is_empty() {
111 inline_schema_refs(&mut schema, &definitions);
112 }
113 if let Value::Object(object) = &mut schema {
114 object.remove("$defs");
115 object.remove("definitions");
116 }
117 schema
118}
119
120fn inline_schema_refs(value: &mut Value, definitions: &Map<String, Value>) {
121 match value {
122 Value::Object(object) => {
123 if let Some(definition) = object
124 .get("$ref")
125 .and_then(Value::as_str)
126 .and_then(|reference| schema_ref_name(reference, definitions))
127 .and_then(|name| definitions.get(name))
128 .cloned()
129 {
130 object.remove("$ref");
131 if object.is_empty() {
132 *value = definition;
133 inline_schema_refs(value, definitions);
134 return;
135 }
136 if let Value::Object(definition_object) = definition {
137 for (key, nested) in definition_object {
138 object.entry(key).or_insert(nested);
139 }
140 }
141 }
142 object.remove("$defs");
143 object.remove("definitions");
144 for nested in object.values_mut() {
145 inline_schema_refs(nested, definitions);
146 }
147 }
148 Value::Array(items) => {
149 for item in items {
150 inline_schema_refs(item, definitions);
151 }
152 }
153 Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
154 }
155}
156
157fn schema_ref_name<'a>(reference: &'a str, definitions: &Map<String, Value>) -> Option<&'a str> {
158 let name = reference
159 .strip_prefix("#/$defs/")
160 .or_else(|| reference.strip_prefix("#/definitions/"))?;
161 definitions.contains_key(name).then_some(name)
162}
163
164fn all_native_tools() -> BTreeSet<NativeToolKind> {
165 [
166 NativeToolKind::WebSearch,
167 NativeToolKind::CodeExecution,
168 NativeToolKind::FileSearch,
169 NativeToolKind::ImageGeneration,
170 NativeToolKind::McpServer,
171 NativeToolKind::ToolSearch,
172 NativeToolKind::GoogleSearch,
173 ]
174 .into_iter()
175 .collect()
176}
177
178const fn no_native_tools() -> BTreeSet<NativeToolKind> {
179 BTreeSet::new()
180}
181
182#[allow(clippy::struct_excessive_bools)]
184#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
185pub struct ModelProfile {
186 pub protocol: ProtocolFamily,
188 pub supports_tools: bool,
190 pub supports_tool_return_schema: bool,
192 pub supports_json_schema_output: bool,
194 pub supports_json_object_output: bool,
196 pub supports_image_output: bool,
198 pub supports_image_input: bool,
200 pub supports_video_input: bool,
202 pub supports_audio_input: bool,
204 pub supports_document_input: bool,
206 pub supports_inline_system_prompts: bool,
208 pub supports_thinking: bool,
210 #[serde(default)]
212 pub drop_sampling_parameters_when_reasoning: bool,
213 pub thinking_always_enabled: bool,
215 #[serde(default = "default_thinking_tags")]
217 pub thinking_tags: (String, String),
218 pub ignore_streamed_leading_whitespace: bool,
220 pub default_structured_output_mode: StructuredOutputMode,
222 #[serde(default = "default_prompted_output_template")]
224 pub prompted_output_template: String,
225 pub native_output_requires_schema_in_instructions: bool,
227 #[serde(default = "all_native_tools")]
229 pub supported_native_tools: BTreeSet<NativeToolKind>,
230 #[serde(default, skip_serializing_if = "Option::is_none")]
232 pub json_schema_transformer: Option<JsonSchemaTransformer>,
233 pub message_normalization: MessageNormalization,
235}
236
237impl ModelProfile {
238 #[must_use]
240 #[allow(clippy::too_many_lines)]
241 pub fn for_protocol(protocol: ProtocolFamily) -> Self {
242 let base = Self {
243 protocol,
244 supports_tools: true,
245 supports_tool_return_schema: false,
246 supports_json_schema_output: false,
247 supports_json_object_output: false,
248 supports_image_output: false,
249 supports_image_input: false,
250 supports_video_input: false,
251 supports_audio_input: false,
252 supports_document_input: false,
253 supports_inline_system_prompts: false,
254 supports_thinking: false,
255 drop_sampling_parameters_when_reasoning: false,
256 thinking_always_enabled: false,
257 thinking_tags: default_thinking_tags(),
258 ignore_streamed_leading_whitespace: false,
259 default_structured_output_mode: StructuredOutputMode::Tool,
260 prompted_output_template: default_prompted_output_template(),
261 native_output_requires_schema_in_instructions: false,
262 supported_native_tools: no_native_tools(),
263 json_schema_transformer: None,
264 message_normalization: MessageNormalization::MergeAdjacentSameRole,
265 };
266 match protocol {
267 ProtocolFamily::OpenAiChatCompletions => Self {
268 supports_json_schema_output: true,
269 supports_json_object_output: true,
270 supports_image_input: true,
271 supports_inline_system_prompts: true,
272 supports_thinking: true,
273 drop_sampling_parameters_when_reasoning: true,
274 default_structured_output_mode: StructuredOutputMode::NativeJsonSchema,
275 supported_native_tools: no_native_tools(),
276 message_normalization: MessageNormalization::MergeAdjacentSameRole,
277 ..base
278 },
279 ProtocolFamily::OpenAiResponses => Self {
280 supports_json_schema_output: true,
281 supports_json_object_output: true,
282 supports_image_input: true,
283 supports_inline_system_prompts: true,
284 supports_thinking: true,
285 drop_sampling_parameters_when_reasoning: true,
286 default_structured_output_mode: StructuredOutputMode::NativeJsonSchema,
287 supported_native_tools: [
288 NativeToolKind::WebSearch,
289 NativeToolKind::CodeExecution,
290 NativeToolKind::FileSearch,
291 NativeToolKind::ImageGeneration,
292 NativeToolKind::McpServer,
293 NativeToolKind::ToolSearch,
294 ]
295 .into_iter()
296 .collect(),
297 message_normalization: MessageNormalization::PreserveItems,
298 ..base
299 },
300 ProtocolFamily::AnthropicMessages => Self {
301 supports_json_schema_output: true,
302 supports_json_object_output: false,
303 supports_image_input: true,
304 supports_document_input: true,
305 supports_thinking: true,
306 drop_sampling_parameters_when_reasoning: true,
307 default_structured_output_mode: StructuredOutputMode::NativeJsonSchema,
308 supported_native_tools: no_native_tools(),
309 message_normalization: MessageNormalization::SystemField,
310 ..base
311 },
312 ProtocolFamily::GeminiGenerateContent => Self {
313 supports_tool_return_schema: true,
314 supports_json_schema_output: true,
315 supports_json_object_output: true,
316 supports_image_input: true,
317 supports_video_input: true,
318 supports_audio_input: true,
319 supports_document_input: true,
320 supports_thinking: true,
321 default_structured_output_mode: StructuredOutputMode::NativeJsonSchema,
322 supported_native_tools: [
323 NativeToolKind::GoogleSearch,
324 NativeToolKind::CodeExecution,
325 ]
326 .into_iter()
327 .collect(),
328 message_normalization: MessageNormalization::SystemInstruction,
329 ..base
330 },
331 ProtocolFamily::BedrockConverse => Self {
332 supports_image_input: true,
333 supports_document_input: true,
334 default_structured_output_mode: StructuredOutputMode::Tool,
335 message_normalization: MessageNormalization::SystemField,
336 ..base
337 },
338 }
339 }
340}
341
342#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
344#[serde(rename_all = "snake_case")]
345pub enum StructuredOutputMode {
346 NativeJsonSchema,
348 NativeJsonObject,
350 Tool,
352 Prompted,
354}
355
356#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
358#[serde(rename_all = "snake_case")]
359pub enum MessageNormalization {
360 MergeAdjacentSameRole,
362 PreserveItems,
364 SystemField,
366 SystemInstruction,
368 WrapInlineSystem,
370}