Merge pull request #23925 from Chesars/fix/mistral-diarize-segments-response

fix(mistral): preserve diarization segments in transcription response

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

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()