Skip to main content

starweaver_model/
profile.rs

1//! Model capability profiles and protocol families.
2
3use std::collections::BTreeSet;
4
5use serde::{Deserialize, Serialize};
6use serde_json::{Map, Value};
7
8/// Supported provider protocol families.
9#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
10#[serde(rename_all = "snake_case")]
11pub enum ProtocolFamily {
12    /// `OpenAI` Chat Completions protocol.
13    OpenAiChatCompletions,
14    /// `OpenAI` Responses protocol.
15    OpenAiResponses,
16    /// Anthropic Messages protocol.
17    AnthropicMessages,
18    /// Gemini generateContent protocol.
19    GeminiGenerateContent,
20    /// Bedrock Converse protocol.
21    BedrockConverse,
22}
23
24/// Provider-executed native tool families.
25#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
26#[serde(rename_all = "snake_case")]
27pub enum NativeToolKind {
28    /// Provider web search tool.
29    WebSearch,
30    /// Provider code execution tool.
31    CodeExecution,
32    /// Provider file search tool.
33    FileSearch,
34    /// Provider image generation tool.
35    ImageGeneration,
36    /// Provider-hosted MCP server tool.
37    McpServer,
38    /// Provider tool-search/discovery tool.
39    ToolSearch,
40    /// Gemini Google Search tool.
41    GoogleSearch,
42}
43
44impl NativeToolKind {
45    /// Return the provider-neutral native tool type string used in `NativeToolDefinition`.
46    #[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    /// Resolve a provider-neutral native tool type string.
60    #[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/// JSON schema normalization strategy for provider-specific schema subsets.
85#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
86#[serde(rename_all = "snake_case")]
87pub enum JsonSchemaTransformer {
88    /// Inline local `$defs` / `definitions` references and remove the top-level definition map.
89    InlineDefinitions,
90}
91
92impl JsonSchemaTransformer {
93    /// Transform a JSON schema value according to this strategy.
94    #[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/// Capability metadata and request-shaping policy.
183#[allow(clippy::struct_excessive_bools)]
184#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
185pub struct ModelProfile {
186    /// Protocol family used by the adapter.
187    pub protocol: ProtocolFamily,
188    /// Function/tool calling support.
189    pub supports_tools: bool,
190    /// Native support for tool return schemas.
191    pub supports_tool_return_schema: bool,
192    /// Native JSON schema output support.
193    pub supports_json_schema_output: bool,
194    /// JSON object mode support.
195    pub supports_json_object_output: bool,
196    /// Image return support.
197    pub supports_image_output: bool,
198    /// Image URL input support.
199    pub supports_image_input: bool,
200    /// Video URL input support.
201    pub supports_video_input: bool,
202    /// Audio URL input support.
203    pub supports_audio_input: bool,
204    /// Document URL input support.
205    pub supports_document_input: bool,
206    /// Flexible system prompt placement support.
207    pub supports_inline_system_prompts: bool,
208    /// Configurable thinking support.
209    pub supports_thinking: bool,
210    /// Drop sampling settings such as temperature/top-p/top-k/penalties/logit-bias when reasoning is enabled.
211    #[serde(default)]
212    pub drop_sampling_parameters_when_reasoning: bool,
213    /// Reasoning is always enabled.
214    pub thinking_always_enabled: bool,
215    /// Tags used by text-stream adapters to split thinking parts from text.
216    #[serde(default = "default_thinking_tags")]
217    pub thinking_tags: (String, String),
218    /// Ignore leading whitespace-only stream content before semantic parts.
219    pub ignore_streamed_leading_whitespace: bool,
220    /// Default structured output mode.
221    pub default_structured_output_mode: StructuredOutputMode,
222    /// Prompt template used for prompted structured output.
223    #[serde(default = "default_prompted_output_template")]
224    pub prompted_output_template: String,
225    /// Whether native structured output still needs schema instructions in the prompt.
226    pub native_output_requires_schema_in_instructions: bool,
227    /// Provider-neutral native tool families supported by this profile.
228    #[serde(default = "all_native_tools")]
229    pub supported_native_tools: BTreeSet<NativeToolKind>,
230    /// JSON schema transformer applied to tools and structured output schemas.
231    #[serde(default, skip_serializing_if = "Option::is_none")]
232    pub json_schema_transformer: Option<JsonSchemaTransformer>,
233    /// Message normalization policy.
234    pub message_normalization: MessageNormalization,
235}
236
237impl ModelProfile {
238    /// Create a default profile for a protocol family.
239    #[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/// Structured output strategy.
343#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
344#[serde(rename_all = "snake_case")]
345pub enum StructuredOutputMode {
346    /// Use provider-native JSON schema mode.
347    NativeJsonSchema,
348    /// Use JSON object mode.
349    NativeJsonObject,
350    /// Use tool/function call output.
351    Tool,
352    /// Use prompt-level instructions.
353    Prompted,
354}
355
356/// Provider message normalization strategy.
357#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
358#[serde(rename_all = "snake_case")]
359pub enum MessageNormalization {
360    /// Merge adjacent messages with the same provider role.
361    MergeAdjacentSameRole,
362    /// Preserve item boundaries.
363    PreserveItems,
364    /// Move system prompts to a top-level system field.
365    SystemField,
366    /// Move system prompts to a top-level system instruction object.
367    SystemInstruction,
368    /// Wrap system fragments into tagged user content.
369    WrapInlineSystem,
370}