From 9afc4697258340100244e1759165790c792ebc47 Mon Sep 17 00:00:00 2001
From: Chesars <cesarponce19544@gmail.com>
Date: Tue, 17 Mar 2026 23:04:51 -0300
Subject: [PATCH] fix(mistral): preserve diarization segments in transcription
 response
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes #23890 — Mistral's Voxtral transcription with `diarize=true` returns
`segments` (with speaker_id, timestamps) and `language`, but these fields
were dropped when mapping the response to TranscriptionResponse.
---
 .../audio_transcription/transformation.py     |  7 +++
 ...tral_audio_transcription_transformation.py | 44 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/litellm/llms/mistral/audio_transcription/transformation.py b/litellm/llms/mistral/audio_transcription/transformation.py
index 4d294063499..8c6d604acb4 100644
--- a/litellm/llms/mistral/audio_transcription/transformation.py
+++ b/litellm/llms/mistral/audio_transcription/transformation.py
@@ -148,5 +148,12 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
 
         text = response_json.get("text") or ""
         response = TranscriptionResponse(text=text)
+
+        # Preserve Mistral-specific fields (e.g. diarization segments)
+        if "segments" in response_json:
+            response["segments"] = response_json["segments"]
+        if "language" in response_json:
+            response["language"] = response_json["language"]
+
         response._hidden_params = response_json
         return response
diff --git a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py
index 7ef50dede0c..4ca3e8ae0c7 100644
--- a/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py
+++ b/tests/test_litellm/llms/mistral/audio_transcription/test_mistral_audio_transcription_transformation.py
@@ -158,6 +158,50 @@ def test_mistral_audio_transcription_response_transform():
     assert response.text == "Four score and seven years ago..."
 
 
+def test_mistral_audio_transcription_response_transform_diarized():
+    """Test that diarized responses preserve segments and language."""
+    config = MistralAudioTranscriptionConfig()
+
+    mock_response = MagicMock(spec=httpx.Response)
+    mock_response.json.return_value = {
+        "model": "voxtral-mini-latest",
+        "text": "Hello, how are you? I am fine.",
+        "language": None,
+        "segments": [
+            {
+                "text": "Hello, how are you?",
+                "start": 0.3,
+                "end": 2.1,
+                "speaker_id": "speaker_1",
+                "type": "transcription_segment",
+            },
+            {
+                "text": "I am fine.",
+                "start": 2.5,
+                "end": 3.8,
+                "speaker_id": "speaker_2",
+                "type": "transcription_segment",
+            },
+        ],
+        "usage": {
+            "prompt_audio_seconds": 4,
+            "prompt_tokens": 5,
+            "total_tokens": 50,
+            "completion_tokens": 20,
+        },
+    }
+
+    response = config.transform_audio_transcription_response(mock_response)
+
+    assert isinstance(response, TranscriptionResponse)
+    assert response.text == "Hello, how are you? I am fine."
+    assert response["segments"] is not None
+    assert len(response["segments"]) == 2
+    assert response["segments"][0]["speaker_id"] == "speaker_1"
+    assert response["segments"][1]["speaker_id"] == "speaker_2"
+    assert response["language"] is None
+
+
 def test_mistral_audio_transcription_response_transform_empty():
     config = MistralAudioTranscriptionConfig()