Skip to main content

gglib_core/domain/
capabilities.rs

1//! Model capability detection, inference, and request transformation.
2//!
3//! This module owns two orthogonal pipelines that operate on different phases
4//! of a model request:
5//!
6//! ## 1. Request-side capability pipeline
7//!
8//! Before a chat-completion request is forwarded to llama-server the proxy
9//! consults the stored [`ModelCapabilities`] flags to decide whether to rewrite
10//! the message list.  Flags are inferred at import time and stored in the
11//! database; they can also be overridden at any time via the API or CLI.
12//!
13//! | Layer | Function | When it fires |
14//! |---|---|---|
15//! | Template analysis | [`infer_from_chat_template`] | At model import — reads `tokenizer.chat_template` from the GGUF |
16//! | Architecture registry | [`capabilities_from_architecture`] | At model import — reads `general.architecture` as a backstop when the GGUF ships without a chat template |
17//! | Request rewriting | [`transform_messages_for_capabilities`] | At proxy time — merges consecutive same-role messages for models that require strict turn alternation |
18//!
19//! The result of Layer 1 and Layer 2 is **OR-combined** and stored in
20//! `Model.capabilities`.  The proxy reads this value once per request via a
21//! single catalog lookup.
22//!
23//! ## 2. Response-side normalization pipeline
24//!
25//! Separate from request rewriting, some models (e.g., Qwen) embed tool-call
26//! JSON inside XML tags in the response text.  This is handled by the
27//! `format:*` tag pipeline in `gglib-proxy::normalize`, which is entirely
28//! independent from `ModelCapabilities`.
29//!
30//! ## Template analysis — positive vs. negative signals
31//!
32//! [`infer_from_chat_template`] uses two kinds of signals for system-role
33//! detection, evaluated in priority order:
34//!
35//! | Priority | Signal | Example pattern | Conclusion |
36//! |---|---|---|---|
37//! | **1 (positive)** | `[SYSTEM_PROMPT]` in template | Mistral v7 | `SUPPORTS_SYSTEM_ROLE` set |
38//! | **1 (positive)** | `[AVAILABLE_TOOLS]` in template | Mistral v3/v3-tekken | `SUPPORTS_SYSTEM_ROLE` set |
39//! | **2 (negative)** | `"Only user, assistant and tool roles…"` | Old Mistral v1/v2 | `SUPPORTS_SYSTEM_ROLE` not set |
40//! | **2 (negative)** | `"got system"` / `"Raise exception"` | Other strict models | `SUPPORTS_SYSTEM_ROLE` not set |
41//! | **default** | No signal found | Generic template | `SUPPORTS_SYSTEM_ROLE` set |
42//!
43//! Positive evidence takes precedence: if `[SYSTEM_PROMPT]` or `[AVAILABLE_TOOLS]`
44//! appears, the negative patterns are ignored for system-role purposes.  This
45//! matters because some Jinja templates contain both an error-raise branch for
46//! unknown roles AND a valid system branch guarded by `[SYSTEM_PROMPT]`.
47//!
48//! ## Architecture registry
49//!
50//! [`capabilities_from_architecture`] maps GGUF `general.architecture` strings
51//! to [`ModelCapabilities`] flags.  This is the **backstop** for models whose
52//! quantized builds strip the `tokenizer.chat_template` section, making
53//! `infer_from_chat_template` return `empty()`.
54//!
55//! | Architecture string | Models | Flags |
56//! |---|---|---|
57//! | `"mistral"` | Mistral v1/v2 (old) | `REQUIRES_STRICT_TURNS` |
58//! | `"mistral3"` | Devstral, Ministral, Mistral Small 3 | `REQUIRES_STRICT_TURNS \| SUPPORTS_SYSTEM_ROLE` |
59//!
60//! **To add a new architecture:**
61//!
62//! 1. Add a match arm in [`capabilities_from_architecture`] mapping the
63//!    architecture string to the appropriate flags.
64//! 2. Add a unit test in the `#[cfg(test)]` block at the bottom of this file.
65//! 3. If the architecture also needs **response-side** normalization (XML tool
66//!    calls, custom reasoning tags, etc.), follow the steps in `CONTRIBUTING.md`
67//!    under "Adding a new model architecture" to add a `format:*` parser as well.
68//! 4. No other files need touching — all call sites already use these functions.
69//!
70//! **Note on Qwen:** Qwen is intentionally absent from the registry.  Qwen's
71//! quantized builds always ship a full chat template, so
72//! [`infer_from_chat_template`] handles the request side.  Its response-side
73//! `<tool_call>` XML is handled by the `format:qwen-xml` tag pipeline.
74
75use bitflags::bitflags;
76use serde::{Deserialize, Serialize};
77
78bitflags! {
79    /// Model capabilities inferred from chat template analysis.
80    ///
81    /// These flags describe what the model's chat template can handle.
82    /// Absence means "we don't know" or "not needed", not "forbidden".
83    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
84    #[repr(transparent)]
85    pub struct ModelCapabilities: u32 {
86        /// Model supports system role natively in its chat template.
87        ///
88        /// When set: system messages can be passed through unchanged.
89        /// When unset: system messages must be converted to user messages.
90        const SUPPORTS_SYSTEM_ROLE    = 0b0000_0001;
91
92        /// Model requires strict user/assistant alternation.
93        ///
94        /// When set: consecutive messages of same role must be merged.
95        /// When unset: message order can be arbitrary (OpenAI-style).
96        const REQUIRES_STRICT_TURNS   = 0b0000_0010;
97
98        /// Model supports tool/function calling.
99        ///
100        /// When set: tool_calls and tool role messages are supported.
101        /// When unset: tool functionality should not be used.
102        const SUPPORTS_TOOL_CALLS     = 0b0000_0100;
103
104        /// Model has reasoning/thinking capability.
105        ///
106        /// When set: model may produce <think> tags or reasoning_content.
107        /// When unset: model produces only standard responses.
108        const SUPPORTS_REASONING      = 0b0000_1000;
109    }
110}
111
112impl Default for ModelCapabilities {
113    /// Default capabilities represent "unknown" state.
114    ///
115    /// Models start with empty capabilities and must be explicitly inferred.
116    /// This prevents incorrect assumptions about model constraints.
117    fn default() -> Self {
118        Self::empty()
119    }
120}
121
122impl Serialize for ModelCapabilities {
123    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
124    where
125        S: serde::Serializer,
126    {
127        self.bits().serialize(serializer)
128    }
129}
130
131impl<'de> Deserialize<'de> for ModelCapabilities {
132    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
133    where
134        D: serde::Deserializer<'de>,
135    {
136        let bits = u32::deserialize(deserializer)?;
137        Ok(Self::from_bits_truncate(bits))
138    }
139}
140
141impl ModelCapabilities {
142    /// Check if model supports system role.
143    pub const fn supports_system_role(self) -> bool {
144        self.contains(Self::SUPPORTS_SYSTEM_ROLE)
145    }
146
147    /// Check if model requires strict user/assistant alternation.
148    pub const fn requires_strict_turns(self) -> bool {
149        self.contains(Self::REQUIRES_STRICT_TURNS)
150    }
151
152    /// Check if model supports tool/function calls.
153    pub const fn supports_tool_calls(self) -> bool {
154        self.contains(Self::SUPPORTS_TOOL_CALLS)
155    }
156
157    /// Check if model supports reasoning phases.
158    pub const fn supports_reasoning(self) -> bool {
159        self.contains(Self::SUPPORTS_REASONING)
160    }
161}
162
163/// Infer model capabilities from chat template Jinja source and model name.
164///
165/// Uses string heuristics to detect template constraints. Returns safe
166/// defaults if template is missing or unparseable.
167///
168/// # Detection Strategy
169///
170/// Two-layer approach:
171/// - **Layer 1 (Metadata)**: Check chat template for reliable signals (preferred)
172/// - **Layer 2 (Name Heuristics)**: Use model name patterns as fallback when metadata is missing
173///
174/// # Capabilities Detected
175///
176/// - **System role**: Positive signals (`[SYSTEM_PROMPT]`, `[AVAILABLE_TOOLS]`) take precedence over
177///   negative signals (explicit rejection messages).  Generic templates with neither signal default
178///   to `SUPPORTS_SYSTEM_ROLE` set.
179/// - **Strict turns**: Looks for alternation enforcement logic (`ns.index % 2`,
180///   `conversation roles must alternate`, etc.)
181/// - **Tool calling**: Checks for `<tool_call>`, `if tools`, `function_call` patterns (metadata);
182///   falls back to model name patterns like "hermes", "functionary" (heuristic)
183/// - **Reasoning**: Checks for `<think>`, `<reasoning>`, `enable_thinking` (metadata);
184///   falls back to model name patterns like "deepseek-r1", "qwq", "o1" (heuristic)
185///
186/// # Fallback Behavior
187///
188/// Missing or unparseable templates default to empty capabilities (unknown state).
189pub fn infer_from_chat_template(
190    template: Option<&str>,
191    model_name: Option<&str>,
192) -> ModelCapabilities {
193    let mut caps = ModelCapabilities::empty();
194
195    // ─────────────────────────────────────────────────────────────────────────────
196    // Layer 1: Metadata-based detection (chat template analysis)
197    // ─────────────────────────────────────────────────────────────────────────────
198
199    let mut tool_detected_from_metadata = false;
200    let mut reasoning_detected_from_metadata = false;
201
202    if let Some(template) = template {
203        // ── System role detection ───────────────────────────────────────────
204        //
205        // Positive evidence (Mistral v7 / v3-tekken) takes precedence over any
206        // negative error-raise patterns.  Some templates contain both a
207        // `[SYSTEM_PROMPT]` branch AND a generic "unsupported role" catch-all,
208        // so we must check positive signals first.
209        //
210        // Sources:
211        //   llama.cpp/src/llama-chat.cpp — `tmpl_contains("[SYSTEM_PROMPT]")` →
212        //     LLM_CHAT_TEMPLATE_MISTRAL_V7; system role handled natively.
213        //   `[AVAILABLE_TOOLS]` → LLM_CHAT_TEMPLATE_MISTRAL_V3; system prepended inline.
214        let supports_system_positive =
215            template.contains("[SYSTEM_PROMPT]") || template.contains("[AVAILABLE_TOOLS]");
216
217        let forbids_system = !supports_system_positive
218            && (template.contains("Only user, assistant and tool roles are supported")
219                || template.contains("got system")
220                || template.contains("Raise exception for unsupported roles"));
221
222        if !forbids_system {
223            caps |= ModelCapabilities::SUPPORTS_SYSTEM_ROLE;
224        }
225
226        // Check for strict alternation requirements
227        // Mistral-style templates enforce user/assistant alternation with modulo checks
228        let requires_alternation = template.contains("must alternate user and assistant")
229            || template.contains("conversation roles must alternate")
230            || template.contains("ns.index % 2");
231
232        if requires_alternation {
233            caps |= ModelCapabilities::REQUIRES_STRICT_TURNS;
234        }
235
236        // Detect tool calling support from template
237        let has_tool_patterns = template.contains("<tool_call>")
238            || template.contains("<|python_tag|>")
239            || template.contains("if tools")
240            || template.contains("tools is defined")
241            || template.contains("tool_calls")
242            || template.contains("function_call");
243
244        if has_tool_patterns {
245            caps |= ModelCapabilities::SUPPORTS_TOOL_CALLS;
246            tool_detected_from_metadata = true;
247        }
248
249        // Detect reasoning/thinking support from template
250        let has_reasoning_patterns = template.contains("<think>")
251            || template.contains("</think>")
252            || template.contains("<reasoning>")
253            || template.contains("</reasoning>")
254            || template.contains("enable_thinking")
255            || template.contains("thinking_forced_open")
256            || template.contains("reasoning_content");
257
258        if has_reasoning_patterns {
259            caps |= ModelCapabilities::SUPPORTS_REASONING;
260            reasoning_detected_from_metadata = true;
261        }
262    }
263
264    // ─────────────────────────────────────────────────────────────────────────────
265    // Layer 2: Name-based heuristic fallback (when metadata is inconclusive)
266    // ─────────────────────────────────────────────────────────────────────────────
267    //
268    // Only use name patterns when chat template didn't provide clear evidence.
269    // This is less reliable but helps with models that have incomplete metadata.
270
271    if let Some(name) = model_name {
272        let name_lower = name.to_lowercase();
273
274        // Heuristic: Tool calling support based on model name
275        if !tool_detected_from_metadata {
276            let has_tool_name = name_lower.contains("hermes")
277                || name_lower.contains("functionary")
278                || name_lower.contains("firefunction")
279                || name_lower.contains("gorilla");
280
281            if has_tool_name {
282                caps |= ModelCapabilities::SUPPORTS_TOOL_CALLS;
283            }
284        }
285
286        // Heuristic: Reasoning support based on model name
287        if !reasoning_detected_from_metadata {
288            let has_reasoning_name = name_lower.contains("deepseek-r1")
289                || name_lower.contains("qwq")
290                || name_lower.contains("-r1-")
291                || name_lower.contains("o1");
292
293            if has_reasoning_name {
294                caps |= ModelCapabilities::SUPPORTS_REASONING;
295            }
296        }
297    }
298
299    caps
300}
301
302/// Map a GGUF `general.architecture` value to its inherent [`ModelCapabilities`].
303///
304/// This is the **single source of truth** for architecture-level behavioural
305/// constraints that apply to the **request** side (message preprocessing).
306/// It is consulted during model registration alongside
307/// [`infer_from_chat_template`] — the two results are `OR`-ed together so that
308/// either signal is sufficient.
309///
310/// # Scope: request preprocessing only
311///
312/// This registry governs `ModelCapabilities` flags (strict-turn coalescing,
313/// system-role conversion, etc.).  It does **not** handle response-stream
314/// dialect normalization — that is a separate concern handled by the
315/// `GgufCapabilities.extensions` → `format:*` tag → `get_parser()` pipeline
316/// in `gglib-core::normalize::registry`.
317///
318/// For example:
319/// - **Qwen** tool-call XML normalization already flows through
320///   `detect_tool_support()` → `extensions.insert("format:qwen-xml")` →
321///   `to_tags()` → `get_parser()` → `QwenXmlParser`.  Qwen's chat template
322///   always contains `<tool_call>` patterns, so `infer_from_chat_template`
323///   (Layer 1) sets `SUPPORTS_TOOL_CALLS` reliably.  No architecture entry
324///   is needed here for Qwen.
325/// - **Mistral** does need an entry: its templates enforce strict alternation,
326///   but many quantised builds ship with the tokenizer section stripped, so
327///   the template layer produces no signal.  `general.architecture = "mistral"`
328///   is always present and provides the necessary backstop.
329///
330/// # Rationale
331///
332/// Some models ship without a parseable `tokenizer.chat_template` in the GGUF
333/// (stripped quantisation builds, partial uploads).  The chat-template layer
334/// then returns `ModelCapabilities::empty()`, silently leaving constraints
335/// unapplied.  Reading `general.architecture` from the GGUF gives us a
336/// ground-truth signal that is always present and never varies by quantisation.
337///
338/// # Adding a new architecture
339///
340/// 1. Add a new `"arch_name" => { … }` arm below.
341/// 2. Add a corresponding unit test in the `#[cfg(test)]` block.
342/// 3. No other file needs touching — all call sites use this function.
343///
344/// # Arguments
345///
346/// * `arch` — value of the `general.architecture` GGUF key
347///   (e.g. `"mistral"`, `"llama"`, `"qwen2"`).  `None` means the key was
348///   absent; returns `empty()` so the model gets pass-through treatment.
349#[must_use]
350pub fn capabilities_from_architecture(arch: Option<&str>) -> ModelCapabilities {
351    let Some(arch) = arch else {
352        return ModelCapabilities::empty();
353    };
354
355    match arch {
356        // Old Mistral v1/v2 — strict alternation, no system role.
357        // Many quantised builds strip the tokenizer section, so the template
358        // layer is blind; this entry is the request-side backstop.
359        "mistral" => ModelCapabilities::REQUIRES_STRICT_TURNS,
360
361        // Newer Mistral-family models (Devstral, Ministral, Mistral Small 3).
362        // Architecture string changed from `"mistral"` to `"mistral3"` when
363        // Mistral adopted mistral-common / Tekken tokeniser.  These models
364        // support system role via `[SYSTEM_PROMPT]…[/SYSTEM_PROMPT]` tokens
365        // (Mistral v7 chat template) but still require strict alternation.
366        "mistral3" => {
367            ModelCapabilities::REQUIRES_STRICT_TURNS | ModelCapabilities::SUPPORTS_SYSTEM_ROLE
368        }
369
370        // All other architectures: no request-side constraints inferred from
371        // architecture alone.  Chat-template analysis may still set flags,
372        // and response-stream normalization is handled by the format:* tag
373        // pipeline independently.
374        _ => ModelCapabilities::empty(),
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    #[test]
383    fn test_default_capabilities() {
384        let caps = ModelCapabilities::default();
385        // Default is "unknown" - no capabilities set
386        assert!(caps.is_empty());
387        assert!(!caps.supports_system_role());
388        assert!(!caps.requires_strict_turns());
389        assert!(!caps.supports_tool_calls());
390        assert!(!caps.supports_reasoning());
391    }
392
393    #[test]
394    fn test_infer_openai_style() {
395        let template = r"
396            {% for message in messages %}
397                {{ message.role }}: {{ message.content }}
398            {% endfor %}
399        ";
400        let caps = infer_from_chat_template(Some(template), None);
401        assert!(caps.supports_system_role());
402        assert!(!caps.requires_strict_turns());
403    }
404
405    #[test]
406    fn test_infer_mistral_style() {
407        let template = r"
408            {% if message.role == 'system' %}
409                {{ raise_exception('Only user, assistant and tool roles are supported, got system.') }}
410            {% endif %}
411            {% if (message['role'] == 'user') != (ns.index % 2 == 0) %}
412                {{ raise_exception('conversation roles must alternate user and assistant') }}
413            {% endif %}
414        ";
415        let caps = infer_from_chat_template(Some(template), None);
416        assert!(!caps.supports_system_role());
417        assert!(caps.requires_strict_turns());
418    }
419
420    #[test]
421    fn test_infer_missing_template() {
422        let caps = infer_from_chat_template(None, None);
423        // Missing template means unknown capabilities - no assumptions made
424        assert!(caps.is_empty());
425        assert!(!caps.supports_system_role());
426    }
427
428    #[test]
429    fn test_tool_calling_from_template() {
430        let template = r"
431            {% if tools %}
432                <tool_call>{{ message.tool_calls }}</tool_call>
433            {% endif %}
434        ";
435        let caps = infer_from_chat_template(Some(template), None);
436        assert!(caps.supports_tool_calls());
437    }
438
439    #[test]
440    fn test_reasoning_from_template() {
441        let template = r"
442            {% if enable_thinking %}
443                <think>{{ message.thinking }}</think>
444            {% endif %}
445        ";
446        let caps = infer_from_chat_template(Some(template), None);
447        assert!(caps.supports_reasoning());
448    }
449
450    #[test]
451    fn test_tool_calling_name_fallback() {
452        // No template, but model name suggests tool support
453        let caps = infer_from_chat_template(None, Some("hermes-2-pro-7b"));
454        assert!(caps.supports_tool_calls());
455    }
456
457    #[test]
458    fn test_reasoning_name_fallback() {
459        // No template, but model name suggests reasoning support
460        let caps = infer_from_chat_template(None, Some("deepseek-r1-lite"));
461        assert!(caps.supports_reasoning());
462    }
463
464    #[test]
465    fn test_metadata_plus_name_fallback() {
466        // Template present but has no tool markers - should still use name fallback
467        let template = "simple template with no tool markers";
468        let caps = infer_from_chat_template(Some(template), Some("hermes-model"));
469        // Name fallback should kick in because metadata didn't detect tools
470        assert!(caps.supports_tool_calls());
471    }
472
473    #[test]
474    fn test_metadata_detected_skips_name_fallback() {
475        // When metadata detects capability, name pattern is ignored
476        let template = "<tool_call>detected</tool_call>";
477        let caps = infer_from_chat_template(Some(template), Some("not-a-tool-model"));
478        // Metadata detected it, so tool support is enabled regardless of name
479        assert!(caps.supports_tool_calls());
480    }
481
482    #[test]
483    fn test_combined_detections() {
484        let template = r"
485            {% if tools %}<tool_call>{{ tool }}</tool_call>{% endif %}
486            <think>{{ reasoning }}</think>
487        ";
488        let caps = infer_from_chat_template(Some(template), None);
489        assert!(caps.supports_tool_calls());
490        assert!(caps.supports_reasoning());
491    }
492
493    // ─── capabilities_from_architecture ─────────────────────────────────────
494
495    #[test]
496    fn test_arch_none_returns_empty() {
497        assert!(capabilities_from_architecture(None).is_empty());
498    }
499
500    #[test]
501    fn test_arch_mistral_requires_strict_turns() {
502        let caps = capabilities_from_architecture(Some("mistral"));
503        assert!(caps.requires_strict_turns());
504    }
505
506    #[test]
507    fn test_arch_llama_returns_empty() {
508        assert!(capabilities_from_architecture(Some("llama")).is_empty());
509    }
510
511    #[test]
512    fn test_arch_unknown_returns_empty() {
513        assert!(capabilities_from_architecture(Some("future-arch-xyz")).is_empty());
514    }
515
516    #[test]
517    fn test_arch_mistral3_strict_turns_and_system_role() {
518        let caps = capabilities_from_architecture(Some("mistral3"));
519        assert!(
520            caps.requires_strict_turns(),
521            "mistral3 must enforce strict turns"
522        );
523        assert!(
524            caps.supports_system_role(),
525            "mistral3 supports system via [SYSTEM_PROMPT]"
526        );
527    }
528
529    #[test]
530    fn test_infer_mistral_v7_supports_system() {
531        // Mistral v7 Jinja template: contains [SYSTEM_PROMPT] token.
532        // This is positive evidence — system role IS supported natively.
533        let template = r"
534            {% if messages[0].role == 'system' %}
535                [SYSTEM_PROMPT]{{ messages[0].content }}[/SYSTEM_PROMPT]
536            {% endif %}
537            {% for message in messages %}
538                {% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
539                    {{ raise_exception('conversation roles must alternate') }}
540                {% endif %}
541            {% endfor %}
542        ";
543        let caps = infer_from_chat_template(Some(template), None);
544        assert!(
545            caps.supports_system_role(),
546            "[SYSTEM_PROMPT] is positive evidence"
547        );
548        assert!(caps.requires_strict_turns(), "still enforces alternation");
549    }
550
551    #[test]
552    fn test_infer_mistral_v3_supports_system() {
553        // Mistral v3 / v3-tekken template: contains [AVAILABLE_TOOLS] token.
554        // llama.cpp prepends system content to the first user turn for these.
555        let template = r"
556            {% if tools is defined %}[AVAILABLE_TOOLS]{{ tools | tojson }}[/AVAILABLE_TOOLS]{% endif %}
557            {% for message in messages %}
558                {% if message.role == 'user' %}[INST]{{ message.content }}[/INST]
559                {% elif message.role == 'assistant' %}{{ message.content }}</s>
560                {% endif %}
561            {% endfor %}
562        ";
563        let caps = infer_from_chat_template(Some(template), None);
564        assert!(
565            caps.supports_system_role(),
566            "[AVAILABLE_TOOLS] is positive evidence"
567        );
568    }
569
570    #[test]
571    fn test_infer_mistral_v1_forbids_system() {
572        // Old Mistral v1/v2 template: no positive tokens, explicit rejection.
573        // Must NOT set SUPPORTS_SYSTEM_ROLE.
574        let template = r"
575            {% if message.role == 'system' %}
576                {{ raise_exception('Only user, assistant and tool roles are supported, got system.') }}
577            {% endif %}
578        ";
579        let caps = infer_from_chat_template(Some(template), None);
580        assert!(
581            !caps.supports_system_role(),
582            "v1/v2 genuinely rejects system role"
583        );
584    }
585
586    #[test]
587    fn test_arch_or_template_additive() {
588        // Template detects tool calls; architecture adds strict turns.
589        // The two are ORed so both flags appear in the result.
590        let template = "<tool_call>{{ tool }}</tool_call>";
591        let from_template = infer_from_chat_template(Some(template), None);
592        let from_arch = capabilities_from_architecture(Some("mistral"));
593        let combined = from_template | from_arch;
594        assert!(combined.supports_tool_calls(), "tool calls from template");
595        assert!(combined.requires_strict_turns(), "strict turns from arch");
596    }
597}
598
599// ─────────────────────────────────────────────────────────────────────────────
600// Message Transformation
601// ─────────────────────────────────────────────────────────────────────────────
602
603/// The content of a chat message.
604///
605/// The `OpenAI` API allows `content` to be either a plain string or a structured
606/// array of typed content parts (text blocks, image URLs, tool results, etc.).
607/// Both forms are preserved faithfully through serialize/deserialize
608/// round-trips so the proxy never re-shapes data it did not need to touch.
609///
610/// # Serde behaviour
611///
612/// Uses `#[serde(untagged)]`, so the wire representation is unchanged:
613/// - `Text("hello")` → `"hello"` (JSON string)
614/// - `Parts([…])` → `[{"type":"text","text":"…"},…]` (JSON array)
615///
616/// A JSON `null` or missing `content` field is handled by the surrounding
617/// `Option<MessageContent>` with `#[serde(default)]`.
618#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
619#[serde(untagged)]
620pub enum MessageContent {
621    /// Plain UTF-8 text.
622    Text(String),
623    /// Structured content parts (text, `image_url`, `tool_result`, …).
624    ///
625    /// Individual part shapes are defined by the `OpenAI` API spec and
626    /// validated by the model, not here.
627    Parts(Vec<serde_json::Value>),
628}
629
630impl MessageContent {
631    /// Borrow the inner string slice when this is plain-text content.
632    pub fn as_str(&self) -> Option<&str> {
633        match self {
634            Self::Text(s) => Some(s),
635            Self::Parts(_) => None,
636        }
637    }
638
639    /// Consume into a single flat `String`.
640    ///
641    /// For [`Text`] the string is returned as-is.  For [`Parts`] all
642    /// `{"type":"text","text":"…"}` entries are concatenated; other part
643    /// types (images, etc.) are omitted — callers should only use this when
644    /// a plain-text representation is required (e.g. the `[System]: ` prefix
645    /// during system-message conversion).
646    ///
647    /// [`Text`]: Self::Text
648    /// [`Parts`]: Self::Parts
649    pub fn into_string(self) -> String {
650        match self {
651            Self::Text(s) => s,
652            Self::Parts(parts) => parts
653                .iter()
654                .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
655                .collect::<Vec<_>>()
656                .join(""),
657        }
658    }
659
660    /// Merge `other` into `self`, producing a single combined [`MessageContent`].
661    ///
662    /// | `self`  | `other` | result |
663    /// |---------|---------|--------|
664    /// | Text    | Text    | Text joined with `"\n\n"` |
665    /// | Parts   | Parts   | Parts arrays concatenated |
666    /// | Text    | Parts   | Parts with a leading text block |
667    /// | Parts   | Text    | Parts with a trailing text block |
668    ///
669    /// Empty strings are handled gracefully (no `"\n\n"` separator when
670    /// either side is empty).
671    fn merge_with(self, other: Self) -> Self {
672        match (self, other) {
673            (Self::Text(mut a), Self::Text(b)) => {
674                if a.is_empty() {
675                    return Self::Text(b);
676                }
677                if b.is_empty() {
678                    return Self::Text(a);
679                }
680                a.push_str("\n\n");
681                a.push_str(&b);
682                Self::Text(a)
683            }
684            (Self::Parts(mut a), Self::Parts(b)) => {
685                a.extend(b);
686                Self::Parts(a)
687            }
688            (Self::Text(a), Self::Parts(b)) => {
689                let mut parts = vec![serde_json::json!({"type": "text", "text": a})];
690                parts.extend(b);
691                Self::Parts(parts)
692            }
693            (Self::Parts(mut a), Self::Text(b)) => {
694                a.push(serde_json::json!({"type": "text", "text": b}));
695                Self::Parts(a)
696            }
697        }
698    }
699}
700
701impl From<String> for MessageContent {
702    fn from(s: String) -> Self {
703        Self::Text(s)
704    }
705}
706
707impl From<&str> for MessageContent {
708    fn from(s: &str) -> Self {
709        Self::Text(s.to_string())
710    }
711}
712
713/// A chat message for transformation.
714///
715/// `content` uses [`MessageContent`] which accepts both a plain JSON string
716/// and a JSON array of content-part objects during deserialization, preserving
717/// the original form during serialization.
718///
719/// # Lossless round-trip
720///
721/// [`transform_messages_for_capabilities`] is reached by deserializing a
722/// client's `messages` array into this type and re-serializing the result, so
723/// any field this struct does not model would be **deleted** on the way
724/// through.  Three fields are named because the transform reads them;
725/// everything else — `tool_call_id`, `name`, vendor extensions — is carried
726/// verbatim in [`extra`](Self::extra) and never interpreted.
727///
728/// This is not hypothetical tidiness.  `tool_call_id` is required by the Jinja
729/// templates of exactly the models that set
730/// [`REQUIRES_STRICT_TURNS`](ModelCapabilities::REQUIRES_STRICT_TURNS) — the
731/// Mistral family — so dropping it would break tool calling on the models the
732/// transform exists to serve.
733#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
734pub struct ChatMessage {
735    pub role: String,
736    #[serde(default, skip_serializing_if = "Option::is_none")]
737    pub content: Option<MessageContent>,
738    #[serde(default, skip_serializing_if = "Option::is_none")]
739    pub tool_calls: Option<serde_json::Value>,
740    /// Every other key the message carried, passed through untouched.
741    ///
742    /// Flattened, so these sit at the message's top level on the wire exactly
743    /// where they came from.  An empty map serializes to nothing.
744    #[serde(flatten)]
745    pub extra: serde_json::Map<String, serde_json::Value>,
746}
747
748impl ChatMessage {
749    /// Merge `other` into `self` in-place.
750    ///
751    /// Used during strict-turn coalescing to combine consecutive same-role
752    /// messages.  Content is merged via [`MessageContent::merge_with`]; tool
753    /// calls are concatenated as JSON arrays.
754    ///
755    /// For [`extra`](Self::extra) the surviving message's keys win and only
756    /// absent ones are adopted from `other`.  Only `user` / `assistant`
757    /// messages are ever merged and those rarely carry extras, so this is a
758    /// tie-break rule rather than a load-bearing one — but silently preferring
759    /// the *later* message's `name` over the one already in the merged text
760    /// would be the surprising choice.
761    fn merge_into(&mut self, other: Self) {
762        self.content = match (self.content.take(), other.content) {
763            (None, b) => b,
764            (a, None) => a,
765            (Some(a), Some(b)) => Some(a.merge_with(b)),
766        };
767        match (self.tool_calls.as_mut(), other.tool_calls) {
768            (_, None) => {}
769            (None, tc) => self.tool_calls = tc,
770            (Some(last_tc), Some(msg_tc)) => {
771                if let (Some(la), Some(ma)) = (last_tc.as_array_mut(), msg_tc.as_array()) {
772                    la.extend_from_slice(ma);
773                }
774            }
775        }
776        for (key, value) in other.extra {
777            self.extra.entry(key).or_insert(value);
778        }
779    }
780}
781
782/// Merge consecutive system messages into a single message.
783///
784/// This is universally safe because:
785/// - No model template requires multiple system messages
786/// - Merging preserves all content with clear separation
787/// - It prevents errors in strict-turn templates (e.g., gemma3/medgemma)
788///
789/// # Arguments
790///
791/// * `messages` - The input chat messages
792///
793/// # Returns
794///
795/// Messages with consecutive system messages merged
796fn merge_consecutive_system_messages(messages: Vec<ChatMessage>) -> Vec<ChatMessage> {
797    if messages.is_empty() {
798        return messages;
799    }
800
801    let mut result: Vec<ChatMessage> = Vec::with_capacity(messages.len());
802
803    for msg in messages {
804        let is_system_merge = result
805            .last()
806            .is_some_and(|last| last.role == "system" && msg.role == "system");
807        if is_system_merge {
808            let last = result.last_mut().unwrap();
809            last.content = match (last.content.take(), msg.content) {
810                (None, b) => b,
811                (a, None) => a,
812                (Some(a), Some(b)) => Some(a.merge_with(b)),
813            };
814        } else {
815            result.push(msg);
816        }
817    }
818
819    result
820}
821
822/// Transform chat messages based on model capabilities.
823///
824/// This is a pure function that applies capability-aware transformations:
825/// - Merges consecutive system messages (always, for all models)
826/// - Converts system messages to user messages when model doesn't support system role
827/// - Merges consecutive same-role messages when model requires strict alternation
828///
829/// # Invariant
830///
831/// Consecutive system messages are ALWAYS merged, regardless of capabilities.
832/// This prevents Jinja template errors in models with strict role alternation.
833///
834/// When capabilities are unknown (empty), only system message merging is applied.
835/// This prevents degrading standard models while ensuring universal compatibility.
836///
837/// # Arguments
838///
839/// * `messages` - The input chat messages to transform
840/// * `capabilities` - The model's capability flags
841///
842/// # Returns
843///
844/// Transformed messages suitable for the model's constraints
845pub fn transform_messages_for_capabilities(
846    mut messages: Vec<ChatMessage>,
847    capabilities: ModelCapabilities,
848) -> Vec<ChatMessage> {
849    // STEP 0 (ALWAYS): Merge consecutive system messages.
850    // This is safe for ALL models and prevents Jinja template errors
851    // in models with strict role alternation (e.g., gemma3/medgemma).
852    // Must run BEFORE the capabilities check to protect unknown models.
853    messages = merge_consecutive_system_messages(messages);
854
855    // Pass through if capabilities are unknown
856    if capabilities.is_empty() {
857        return messages;
858    }
859
860    // STEP 1: Transform system messages if the model doesn't support them
861    if !capabilities.contains(ModelCapabilities::SUPPORTS_SYSTEM_ROLE) {
862        for msg in &mut messages {
863            if msg.role == "system" {
864                msg.role = "user".to_string();
865                if let Some(content) = msg.content.take() {
866                    msg.content = Some(MessageContent::Text(format!(
867                        "[System]: {}",
868                        content.into_string()
869                    )));
870                }
871            }
872        }
873    }
874
875    // STEP 2: Merge consecutive same-role messages if strict turns are required
876    if capabilities.contains(ModelCapabilities::REQUIRES_STRICT_TURNS) {
877        let mut merged: Vec<ChatMessage> = Vec::new();
878        for msg in messages {
879            let is_mergeable = msg.role == "user" || msg.role == "assistant";
880            let same_role_as_last = merged.last().is_some_and(|last| last.role == msg.role);
881            if is_mergeable && same_role_as_last {
882                merged.last_mut().unwrap().merge_into(msg);
883            } else {
884                merged.push(msg);
885            }
886        }
887        return merged;
888    }
889
890    messages
891}
892
893#[cfg(test)]
894mod transform_tests {
895    use super::*;
896
897    #[test]
898    fn test_transform_unknown_passes_through_non_system() {
899        // Non-system messages pass through unchanged with empty capabilities
900        let messages = vec![
901            ChatMessage {
902                role: "user".to_string(),
903                content: Some(MessageContent::Text("Hello".to_string())),
904                tool_calls: None,
905                ..Default::default()
906            },
907            ChatMessage {
908                role: "assistant".to_string(),
909                content: Some(MessageContent::Text("Hi there".to_string())),
910                tool_calls: None,
911                ..Default::default()
912            },
913        ];
914        let original = messages.clone();
915        let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
916        assert_eq!(result, original);
917    }
918
919    #[test]
920    fn test_merges_consecutive_system_messages_always() {
921        // Even with empty capabilities, consecutive system messages should merge
922        let messages = vec![
923            ChatMessage {
924                role: "system".to_string(),
925                content: Some(MessageContent::Text(
926                    "You are a helpful assistant.".to_string(),
927                )),
928                tool_calls: None,
929                ..Default::default()
930            },
931            ChatMessage {
932                role: "system".to_string(),
933                content: Some(MessageContent::Text(
934                    "WORKING_MEMORY:\n- task1 (ok): done".to_string(),
935                )),
936                tool_calls: None,
937                ..Default::default()
938            },
939            ChatMessage {
940                role: "user".to_string(),
941                content: Some(MessageContent::Text("Hello".to_string())),
942                tool_calls: None,
943                ..Default::default()
944            },
945        ];
946        let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
947
948        assert_eq!(result.len(), 2);
949        assert_eq!(result[0].role, "system");
950        assert_eq!(
951            result[0].content.as_ref().and_then(|c| c.as_str()),
952            Some("You are a helpful assistant.\n\nWORKING_MEMORY:\n- task1 (ok): done")
953        );
954        assert_eq!(result[1].role, "user");
955    }
956
957    #[test]
958    fn test_merges_three_consecutive_system_messages() {
959        let messages = vec![
960            ChatMessage {
961                role: "system".to_string(),
962                content: Some(MessageContent::Text("First.".to_string())),
963                tool_calls: None,
964                ..Default::default()
965            },
966            ChatMessage {
967                role: "system".to_string(),
968                content: Some(MessageContent::Text("Second.".to_string())),
969                tool_calls: None,
970                ..Default::default()
971            },
972            ChatMessage {
973                role: "system".to_string(),
974                content: Some(MessageContent::Text("Third.".to_string())),
975                tool_calls: None,
976                ..Default::default()
977            },
978        ];
979        let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
980
981        assert_eq!(result.len(), 1);
982        assert_eq!(
983            result[0].content.as_ref().and_then(|c| c.as_str()),
984            Some("First.\n\nSecond.\n\nThird.")
985        );
986    }
987
988    #[test]
989    fn test_handles_empty_system_content() {
990        let messages = vec![
991            ChatMessage {
992                role: "system".to_string(),
993                content: Some(MessageContent::Text(String::new())),
994                tool_calls: None,
995                ..Default::default()
996            },
997            ChatMessage {
998                role: "system".to_string(),
999                content: Some(MessageContent::Text("Actual content".to_string())),
1000                tool_calls: None,
1001                ..Default::default()
1002            },
1003        ];
1004        let result = transform_messages_for_capabilities(messages, ModelCapabilities::empty());
1005
1006        assert_eq!(result.len(), 1);
1007        assert_eq!(
1008            result[0].content.as_ref().and_then(|c| c.as_str()),
1009            Some("Actual content")
1010        );
1011    }
1012
1013    #[test]
1014    fn test_transform_system_to_user() {
1015        let messages = vec![
1016            ChatMessage {
1017                role: "system".to_string(),
1018                content: Some(MessageContent::Text("You are helpful".to_string())),
1019                tool_calls: None,
1020                ..Default::default()
1021            },
1022            ChatMessage {
1023                role: "user".to_string(),
1024                content: Some(MessageContent::Text("Hello".to_string())),
1025                tool_calls: None,
1026                ..Default::default()
1027            },
1028        ];
1029        // Use REQUIRES_STRICT_TURNS which doesn't support system but doesn't merge different roles
1030        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1031        let result = transform_messages_for_capabilities(messages, caps);
1032        // System becomes user, both messages remain separate (user + user but different content)
1033        assert_eq!(result.len(), 1); // They get merged because both are now "user"
1034        assert_eq!(result[0].role, "user");
1035        let content_str = result[0].content.as_ref().and_then(|c| c.as_str()).unwrap();
1036        assert!(content_str.contains("[System]: You are helpful"));
1037        assert!(content_str.contains("Hello"));
1038    }
1039
1040    #[test]
1041    fn test_transform_preserves_system_when_supported() {
1042        let messages = vec![ChatMessage {
1043            role: "system".to_string(),
1044            content: Some(MessageContent::Text("You are helpful".to_string())),
1045            tool_calls: None,
1046            ..Default::default()
1047        }];
1048        let caps = ModelCapabilities::SUPPORTS_SYSTEM_ROLE;
1049        let result = transform_messages_for_capabilities(messages, caps);
1050        assert_eq!(result[0].role, "system");
1051        assert_eq!(
1052            result[0].content,
1053            Some(MessageContent::Text("You are helpful".to_string()))
1054        );
1055    }
1056
1057    #[test]
1058    fn test_transform_merges_consecutive_user_messages() {
1059        let messages = vec![
1060            ChatMessage {
1061                role: "user".to_string(),
1062                content: Some(MessageContent::Text("First".to_string())),
1063                tool_calls: None,
1064                ..Default::default()
1065            },
1066            ChatMessage {
1067                role: "user".to_string(),
1068                content: Some(MessageContent::Text("Second".to_string())),
1069                tool_calls: None,
1070                ..Default::default()
1071            },
1072        ];
1073        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1074        let result = transform_messages_for_capabilities(messages, caps);
1075        assert_eq!(result.len(), 1);
1076        assert_eq!(
1077            result[0].content,
1078            Some(MessageContent::Text("First\n\nSecond".to_string()))
1079        );
1080    }
1081
1082    #[test]
1083    fn test_transform_does_not_merge_tool_messages() {
1084        let messages = vec![
1085            ChatMessage {
1086                role: "tool".to_string(),
1087                content: Some(MessageContent::Text("Result 1".to_string())),
1088                tool_calls: None,
1089                ..Default::default()
1090            },
1091            ChatMessage {
1092                role: "tool".to_string(),
1093                content: Some(MessageContent::Text("Result 2".to_string())),
1094                tool_calls: None,
1095                ..Default::default()
1096            },
1097        ];
1098        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1099        let result = transform_messages_for_capabilities(messages, caps);
1100        assert_eq!(result.len(), 2); // Should not merge tool messages
1101    }
1102
1103    #[test]
1104    fn test_transform_combined_system_and_merge() {
1105        let messages = vec![
1106            ChatMessage {
1107                role: "system".to_string(),
1108                content: Some(MessageContent::Text("Be helpful".to_string())),
1109                tool_calls: None,
1110                ..Default::default()
1111            },
1112            ChatMessage {
1113                role: "user".to_string(),
1114                content: Some(MessageContent::Text("First".to_string())),
1115                tool_calls: None,
1116                ..Default::default()
1117            },
1118            ChatMessage {
1119                role: "user".to_string(),
1120                content: Some(MessageContent::Text("Second".to_string())),
1121                tool_calls: None,
1122                ..Default::default()
1123            },
1124        ];
1125        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS; // No system support + strict turns
1126        let result = transform_messages_for_capabilities(messages, caps);
1127        assert_eq!(result.len(), 1); // System→user + merge
1128        assert_eq!(result[0].role, "user");
1129        let content_str = result[0].content.as_ref().and_then(|c| c.as_str()).unwrap();
1130        assert!(content_str.contains("[System]: Be helpful"));
1131        assert!(content_str.contains("First"));
1132        assert!(content_str.contains("Second"));
1133    }
1134
1135    #[test]
1136    fn test_merge_consecutive_assistant_with_tool_calls() {
1137        // This is the main bug fix: consecutive assistant messages with tool_calls
1138        // should be merged for models requiring strict turns
1139        let tool_call_1 = serde_json::json!([
1140            {
1141                "id": "call_1",
1142                "type": "function",
1143                "function": {
1144                    "name": "get_weather",
1145                    "arguments": "{\"location\":\"Paris\"}"
1146                }
1147            }
1148        ]);
1149        let tool_call_2 = serde_json::json!([
1150            {
1151                "id": "call_2",
1152                "type": "function",
1153                "function": {
1154                    "name": "get_time",
1155                    "arguments": "{\"timezone\":\"UTC\"}"
1156                }
1157            }
1158        ]);
1159
1160        let messages = vec![
1161            ChatMessage {
1162                role: "user".to_string(),
1163                content: Some(MessageContent::Text("What's the weather?".to_string())),
1164                tool_calls: None,
1165                ..Default::default()
1166            },
1167            ChatMessage {
1168                role: "assistant".to_string(),
1169                content: Some(MessageContent::Text("Let me check...".to_string())),
1170                tool_calls: Some(tool_call_1),
1171                ..Default::default()
1172            },
1173            ChatMessage {
1174                role: "assistant".to_string(),
1175                content: Some(MessageContent::Text("And the time...".to_string())),
1176                tool_calls: Some(tool_call_2),
1177                ..Default::default()
1178            },
1179        ];
1180
1181        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1182        let result = transform_messages_for_capabilities(messages, caps);
1183
1184        // Should merge into 2 messages: user + merged assistant
1185        assert_eq!(result.len(), 2);
1186        assert_eq!(result[0].role, "user");
1187        assert_eq!(result[1].role, "assistant");
1188
1189        // Content should be merged
1190        assert_eq!(
1191            result[1].content,
1192            Some(MessageContent::Text(
1193                "Let me check...\n\nAnd the time...".to_string()
1194            ))
1195        );
1196
1197        // Tool calls should be concatenated
1198        let merged_tool_calls = result[1].tool_calls.as_ref().unwrap();
1199        let tool_calls_array = merged_tool_calls.as_array().unwrap();
1200        assert_eq!(tool_calls_array.len(), 2);
1201        assert_eq!(tool_calls_array[0]["id"], "call_1");
1202        assert_eq!(tool_calls_array[1]["id"], "call_2");
1203    }
1204
1205    #[test]
1206    fn test_merge_assistant_messages_only_first_has_content() {
1207        // First message has content, second has only tool_calls
1208        let tool_call = serde_json::json!([
1209            {
1210                "id": "call_1",
1211                "type": "function",
1212                "function": {
1213                    "name": "get_weather",
1214                    "arguments": "{}"
1215                }
1216            }
1217        ]);
1218
1219        let messages = vec![
1220            ChatMessage {
1221                role: "assistant".to_string(),
1222                content: Some(MessageContent::Text("Let me check...".to_string())),
1223                tool_calls: None,
1224                ..Default::default()
1225            },
1226            ChatMessage {
1227                role: "assistant".to_string(),
1228                content: None,
1229                tool_calls: Some(tool_call),
1230                ..Default::default()
1231            },
1232        ];
1233
1234        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1235        let result = transform_messages_for_capabilities(messages, caps);
1236
1237        assert_eq!(result.len(), 1);
1238        assert_eq!(
1239            result[0].content,
1240            Some(MessageContent::Text("Let me check...".to_string()))
1241        );
1242        assert!(result[0].tool_calls.is_some());
1243    }
1244
1245    #[test]
1246    fn test_merge_assistant_messages_only_second_has_content() {
1247        // First message has only tool_calls, second has content
1248        let tool_call = serde_json::json!([
1249            {
1250                "id": "call_1",
1251                "type": "function",
1252                "function": {
1253                    "name": "get_weather",
1254                    "arguments": "{}"
1255                }
1256            }
1257        ]);
1258
1259        let messages = vec![
1260            ChatMessage {
1261                role: "assistant".to_string(),
1262                content: None,
1263                tool_calls: Some(tool_call),
1264                ..Default::default()
1265            },
1266            ChatMessage {
1267                role: "assistant".to_string(),
1268                content: Some(MessageContent::Text("Result received".to_string())),
1269                tool_calls: None,
1270                ..Default::default()
1271            },
1272        ];
1273
1274        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1275        let result = transform_messages_for_capabilities(messages, caps);
1276
1277        assert_eq!(result.len(), 1);
1278        assert_eq!(
1279            result[0].content,
1280            Some(MessageContent::Text("Result received".to_string()))
1281        );
1282        assert!(result[0].tool_calls.is_some());
1283    }
1284
1285    #[test]
1286    fn test_merge_assistant_messages_neither_has_content() {
1287        // Both messages have only tool_calls, no content
1288        let tool_call_1 = serde_json::json!([
1289            {
1290                "id": "call_1",
1291                "type": "function",
1292                "function": {"name": "tool1", "arguments": "{}"}
1293            }
1294        ]);
1295        let tool_call_2 = serde_json::json!([
1296            {
1297                "id": "call_2",
1298                "type": "function",
1299                "function": {"name": "tool2", "arguments": "{}"}
1300            }
1301        ]);
1302
1303        let messages = vec![
1304            ChatMessage {
1305                role: "assistant".to_string(),
1306                content: None,
1307                tool_calls: Some(tool_call_1),
1308                ..Default::default()
1309            },
1310            ChatMessage {
1311                role: "assistant".to_string(),
1312                content: None,
1313                tool_calls: Some(tool_call_2),
1314                ..Default::default()
1315            },
1316        ];
1317
1318        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1319        let result = transform_messages_for_capabilities(messages, caps);
1320
1321        assert_eq!(result.len(), 1);
1322        assert!(result[0].content.is_none());
1323
1324        let merged_tool_calls = result[0].tool_calls.as_ref().unwrap();
1325        let tool_calls_array = merged_tool_calls.as_array().unwrap();
1326        assert_eq!(tool_calls_array.len(), 2);
1327    }
1328
1329    #[test]
1330    fn test_no_merge_without_strict_turns_capability() {
1331        // Even with consecutive assistant messages, don't merge if capability not set
1332        let messages = vec![
1333            ChatMessage {
1334                role: "assistant".to_string(),
1335                content: Some(MessageContent::Text("First".to_string())),
1336                tool_calls: None,
1337                ..Default::default()
1338            },
1339            ChatMessage {
1340                role: "assistant".to_string(),
1341                content: Some(MessageContent::Text("Second".to_string())),
1342                tool_calls: None,
1343                ..Default::default()
1344            },
1345        ];
1346
1347        let caps = ModelCapabilities::empty();
1348        let result = transform_messages_for_capabilities(messages, caps);
1349
1350        // Should NOT merge without REQUIRES_STRICT_TURNS capability
1351        assert_eq!(result.len(), 2);
1352    }
1353
1354    #[test]
1355    fn test_merge_preserves_different_role_boundaries() {
1356        // Don't merge across different roles
1357        let tool_call = serde_json::json!([
1358            {
1359                "id": "call_1",
1360                "type": "function",
1361                "function": {"name": "tool1", "arguments": "{}"}
1362            }
1363        ]);
1364
1365        let messages = vec![
1366            ChatMessage {
1367                role: "user".to_string(),
1368                content: Some(MessageContent::Text("Question".to_string())),
1369                tool_calls: None,
1370                ..Default::default()
1371            },
1372            ChatMessage {
1373                role: "assistant".to_string(),
1374                content: Some(MessageContent::Text("Answer".to_string())),
1375                tool_calls: Some(tool_call),
1376                ..Default::default()
1377            },
1378            ChatMessage {
1379                role: "user".to_string(),
1380                content: Some(MessageContent::Text("Follow-up".to_string())),
1381                tool_calls: None,
1382                ..Default::default()
1383            },
1384        ];
1385
1386        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1387        let result = transform_messages_for_capabilities(messages, caps);
1388
1389        // Should remain 3 separate messages
1390        assert_eq!(result.len(), 3);
1391        assert_eq!(result[0].role, "user");
1392        assert_eq!(result[1].role, "assistant");
1393        assert_eq!(result[2].role, "user");
1394    }
1395
1396    /// Verify that array-form `content` (`OpenAI` multipart spec) deserializes
1397    /// correctly into `MessageContent::Parts` and is preserved on serialization.
1398    #[test]
1399    fn test_array_content_deserializes_to_parts() {
1400        let json = serde_json::json!({
1401            "role": "user",
1402            "content": [
1403                {"type": "text", "text": "What is in this image?"},
1404                {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}
1405            ]
1406        });
1407        let msg: ChatMessage = serde_json::from_value(json).unwrap();
1408        assert!(matches!(msg.content, Some(MessageContent::Parts(_))));
1409        // Round-trip: serialises back to array form
1410        let re_serialised = serde_json::to_value(&msg).unwrap();
1411        assert!(re_serialised["content"].is_array());
1412    }
1413
1414    /// Verify that two user messages where one has array content and one has
1415    /// text content are merged correctly into a Parts result.
1416    #[test]
1417    fn test_merge_array_content_with_text_content() {
1418        let messages = vec![
1419            ChatMessage {
1420                role: "user".to_string(),
1421                content: Some(MessageContent::Parts(vec![
1422                    serde_json::json!({"type": "text", "text": "Look at this:"}),
1423                    serde_json::json!({"type": "image_url", "image_url": {"url": "https://example.com/img.png"}}),
1424                ])),
1425                tool_calls: None,
1426                ..Default::default()
1427            },
1428            ChatMessage {
1429                role: "user".to_string(),
1430                content: Some(MessageContent::Text("What do you see?".to_string())),
1431                tool_calls: None,
1432                ..Default::default()
1433            },
1434        ];
1435        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1436        let result = transform_messages_for_capabilities(messages, caps);
1437
1438        assert_eq!(result.len(), 1);
1439        // Parts + Text → Parts with a trailing text block
1440        assert!(matches!(&result[0].content, Some(MessageContent::Parts(p)) if p.len() == 3));
1441    }
1442
1443    /// Verify that tool-result messages with array content survive the
1444    /// coalescing path unchanged (they are not mergeable roles).
1445    #[test]
1446    fn test_tool_message_with_array_content_passes_through() {
1447        let messages = vec![
1448            ChatMessage {
1449                role: "user".to_string(),
1450                content: Some(MessageContent::Text("Run the tool".to_string())),
1451                tool_calls: None,
1452                ..Default::default()
1453            },
1454            ChatMessage {
1455                role: "tool".to_string(),
1456                content: Some(MessageContent::Parts(vec![
1457                    serde_json::json!({"type": "text", "text": "tool result here"}),
1458                ])),
1459                tool_calls: None,
1460                ..Default::default()
1461            },
1462            ChatMessage {
1463                role: "user".to_string(),
1464                content: Some(MessageContent::Text("Thanks".to_string())),
1465                tool_calls: None,
1466                ..Default::default()
1467            },
1468        ];
1469        let caps = ModelCapabilities::REQUIRES_STRICT_TURNS;
1470        let result = transform_messages_for_capabilities(messages, caps);
1471
1472        // tool message is not merged; user messages are consecutive after tool
1473        // so the two user messages are separated by the tool message
1474        assert_eq!(result.len(), 3);
1475        assert_eq!(result[1].role, "tool");
1476        assert!(matches!(&result[1].content, Some(MessageContent::Parts(_))));
1477    }
1478}
1479
1480/// The JSON round-trip callers actually perform.
1481///
1482/// [`transform_messages_for_capabilities`] is reached by deserializing a
1483/// client's `messages` array into [`ChatMessage`] and re-serializing the
1484/// result, so anything the struct fails to model is deleted in transit. These
1485/// tests pin the round-trip itself rather than the transform, because that is
1486/// where fields go missing.
1487#[cfg(test)]
1488mod round_trip_tests {
1489    use super::*;
1490    use serde_json::{Value, json};
1491
1492    /// Deserialize → transform → serialize, exactly as the proxy and the
1493    /// adapter do.
1494    fn round_trip(messages: Value, capabilities: ModelCapabilities) -> Vec<Value> {
1495        let parsed: Vec<ChatMessage> = serde_json::from_value(messages).expect("valid messages");
1496        let transformed = transform_messages_for_capabilities(parsed, capabilities);
1497        serde_json::to_value(&transformed)
1498            .expect("serializable")
1499            .as_array()
1500            .expect("array")
1501            .clone()
1502    }
1503
1504    /// The regression this catch-all exists for.
1505    ///
1506    /// Mistral-family templates require `tool_call_id` on tool results and are
1507    /// exactly the models that set `REQUIRES_STRICT_TURNS`, so a transform that
1508    /// dropped it would break tool calling on the only models that need the
1509    /// transform at all.
1510    #[test]
1511    fn tool_call_id_survives_strict_turn_coalescing() {
1512        let out = round_trip(
1513            json!([
1514                {"role": "user", "content": "run it"},
1515                {"role": "assistant", "tool_calls": [
1516                    {"id": "call_1", "type": "function",
1517                     "function": {"name": "f", "arguments": "{}"}}
1518                ]},
1519                {"role": "tool", "tool_call_id": "call_1", "content": "result"},
1520            ]),
1521            ModelCapabilities::REQUIRES_STRICT_TURNS,
1522        );
1523
1524        let tool = out
1525            .iter()
1526            .find(|m| m["role"] == "tool")
1527            .expect("tool message");
1528        assert_eq!(
1529            tool["tool_call_id"], "call_1",
1530            "tool_call_id must survive the transform"
1531        );
1532    }
1533
1534    /// Every unmodelled key, not just the one that motivated the field.
1535    #[test]
1536    fn arbitrary_unknown_message_keys_are_preserved() {
1537        let out = round_trip(
1538            json!([{
1539                "role": "user",
1540                "content": "hi",
1541                "name": "alice",
1542                "x_vendor_extension": {"nested": [1, 2]},
1543            }]),
1544            ModelCapabilities::REQUIRES_STRICT_TURNS,
1545        );
1546
1547        assert_eq!(out[0]["name"], "alice");
1548        assert_eq!(out[0]["x_vendor_extension"], json!({"nested": [1, 2]}));
1549    }
1550
1551    /// `#[serde(flatten)]` routes the whole struct through serde's buffered
1552    /// content representation, and `MessageContent` is an untagged enum — the
1553    /// combination that most often silently degrades. Both content forms must
1554    /// still deserialize to the right variant and serialize back unchanged.
1555    #[test]
1556    fn flatten_does_not_disturb_untagged_content_forms() {
1557        let out = round_trip(
1558            json!([
1559                {"role": "user", "content": "plain string"},
1560                {"role": "assistant", "content": [{"type": "text", "text": "part"}]},
1561            ]),
1562            ModelCapabilities::empty(),
1563        );
1564
1565        assert_eq!(out[0]["content"], json!("plain string"));
1566        assert_eq!(out[1]["content"], json!([{"type": "text", "text": "part"}]));
1567    }
1568
1569    /// A merged message keeps its own extras and adopts only what it lacks.
1570    #[test]
1571    fn merging_prefers_the_surviving_message_extras() {
1572        let out = round_trip(
1573            json!([
1574                {"role": "user", "content": "one", "name": "first"},
1575                {"role": "user", "content": "two", "name": "second", "only_on_later": true},
1576            ]),
1577            ModelCapabilities::REQUIRES_STRICT_TURNS,
1578        );
1579
1580        assert_eq!(out.len(), 1, "consecutive user messages merge");
1581        assert_eq!(out[0]["name"], "first", "the surviving message's key wins");
1582        assert_eq!(out[0]["only_on_later"], true, "absent keys are adopted");
1583    }
1584
1585    /// The catch-all must not invent keys for messages that had none.
1586    #[test]
1587    fn messages_without_extras_serialize_unchanged() {
1588        let out = round_trip(
1589            json!([{"role": "user", "content": "hi"}]),
1590            ModelCapabilities::empty(),
1591        );
1592        assert_eq!(out[0], json!({"role": "user", "content": "hi"}));
1593    }
1594}