Merge 4ad2a69ae1 into 17db64ef78

litellm/llms/mistral/chat/transformation.py

@@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works
 Docs - https://docs.mistral.ai/api/
 """
 
-from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
+from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
 from itertools import accumulate, count
 from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
 
@@ -593,22 +593,38 @@ class MistralConfig(OpenAIGPTConfig):
                     # Only process if content is a list
                     if isinstance(content, list):
                         thinking_content = ""
-                        text_content = ""
+                        text_segments: list[str] = []  # mutable-ok: accumulated across this choice's blocks
+                        reference_blocks: list[dict[str, object]] = []  # mutable-ok: citation chunks for this choice
 
                         # Process each content block
                         for block in content:
-                            if block.get("type") == "thinking":
+                            block_type = block.get("type")  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # provider content blocks are untyped JSON
+                            if block_type == "thinking":
                                 thinking_blocks = block.get("thinking", [])
                                 thinking_texts = []
                                 for thinking_block in thinking_blocks:
                                     if thinking_block.get("type") == "text":
                                         thinking_texts.append(thinking_block.get("text", ""))
                                 thinking_content = "\n".join(thinking_texts)
-                            elif block.get("type") == "text":
-                                text_content = block.get("text", "")
+                            elif block_type == "text":
+                                text_segments.append(block.get("text", ""))  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType]  # provider content blocks are untyped JSON
+                            elif block_type == "reference":
+                                reference_blocks.append(block)  # pyright: ignore[reportUnknownArgumentType]  # provider content blocks are untyped JSON
 
                         # Set the extracted content
-                        choice["message"]["content"] = text_content
+                        choice["message"]["content"] = "".join(text_segments)
+                        if reference_blocks:
+                            # The flattened content string cannot carry Mistral
+                            # citation chunks, so expose the citation chunks
+                            # themselves instead of dropping them silently.
+                            #
+                            # Only the citation chunks are exposed, not a copy of
+                            # the whole block list: the text lives in `content`,
+                            # which the output PII-masking hook sanitises, while
+                            # provider-specific fields are not sanitised -- a
+                            # second copy of the text would reach the client
+                            # unmasked.
+                            choice["message"]["mistral_reference_blocks"] = reference_blocks
                         if thinking_content:
                             choice["message"]["reasoning_content"] = thinking_content
 
@@ -707,6 +723,7 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler):
                         normalized_text,
                         thinking_blocks,
                         reasoning_content,
+                        reference_blocks,
                     ) = self._normalize_content_blocks(content)
                     delta["content"] = normalized_text
                     if thinking_blocks:
@@ -715,6 +732,15 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler):
                     else:
                         delta.pop("thinking_blocks", None)
                         delta.pop("reasoning_content", None)
+                    if reference_blocks:
+                        # citation chunks carry no text; expose them on
+                        # provider_specific_fields instead of dropping them
+                        existing_fields = delta.get("provider_specific_fields")  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # provider_specific_fields comes from the provider payload, which is untyped JSON
+                        if isinstance(existing_fields, dict):
+                            # keep whatever the provider already put there
+                            existing_fields["mistral_reference_blocks"] = reference_blocks
+                        else:
+                            delta["provider_specific_fields"] = {"mistral_reference_blocks": reference_blocks}
         except Exception:
             # Fall back to default parsing if custom handling fails
             return super().chunk_parser(chunk)
@@ -724,13 +750,20 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler):
     @staticmethod
     def _normalize_content_blocks(
         content_blocks: list[dict],
-    ) -> tuple[str | None, list[dict], str | None]:
+    ) -> tuple[str | None, list[dict], str | None, Sequence[Mapping[str, object]]]:
         """
         Convert Mistral magistral content blocks into OpenAI-compatible content + thinking_blocks.
+
+        The fourth element carries the citation chunks (`reference` blocks) as
+        the provider sent them. They are returned separately from the text
+        because the text is flattened into `content`, where the output
+        PII-masking hook can sanitise it, while provider-specific fields are
+        not sanitised.
         """
         text_segments: Final[list[str]] = []
         thinking_blocks: Final[list[dict]] = []
         reasoning_segments: Final[list[str]] = []
+        reference_blocks: Final[list[dict[str, object]]] = []  # mutable-ok: citation chunks per response
 
         for block in content_blocks:
             block_type = block.get("type")
@@ -752,7 +785,9 @@ class MistralChatResponseIterator(OpenAIChatCompletionStreamingHandler):
                     )
             elif block_type == "text":
                 text_segments.append(block.get("text", ""))
+            elif block_type == "reference":
+                reference_blocks.append(block)  # pyright: ignore[reportUnknownArgumentType]  # provider content blocks are untyped JSON
 
         normalized_text: Final = "".join(text_segments) if text_segments else None
         reasoning_content: Final = "\n".join(reasoning_segments) if reasoning_segments else None
-        return normalized_text, thinking_blocks, reasoning_content
+        return normalized_text, thinking_blocks, reasoning_content, reference_blocks

tests/unit/llms/mistral/test_mistral_chat_transformation.py

@@ -1162,3 +1162,134 @@ def test_mistral_transform_request_drops_an_empty_tool_message_whose_call_id_is_
     )
 
     assert [m.get("role") for m in request["messages"]] == ["user", "assistant"]
+
+
+class TestMistralContentListConversion:
+    """
+    Mistral content-as-list must keep every text chunk and expose citation chunks.
+
+    Issue #45378: the non-streaming conversion overwrote the content per text
+    block (only the last chunk survived — an answer closed by a citation was
+    flattened to "."), and reference chunks were dropped on both the
+    non-streaming and streaming paths.
+    """
+
+    CITATION_CONTENT = [
+        {"type": "text", "text": "The sky is blue during the day"},
+        {"type": "reference", "reference_ids": ["doc1"]},
+        {"type": "text", "text": "."},
+    ]
+    REFERENCE_BLOCK = {"type": "reference", "reference_ids": ["doc1"]}
+
+    def test_non_streaming_concatenates_all_text_chunks(self):
+        response_data = {"choices": [{"message": {"role": "assistant", "content": self.CITATION_CONTENT}}]}
+
+        MistralConfig._handle_content_list_to_str_conversion(response_data)
+
+        message = response_data["choices"][0]["message"]
+        assert message["content"] == "The sky is blue during the day."
+
+    def test_non_streaming_exposes_reference_chunks_in_provider_specific_fields(self):
+        response_data = {"choices": [{"message": {"role": "assistant", "content": self.CITATION_CONTENT}}]}
+
+        MistralConfig._handle_content_list_to_str_conversion(response_data)
+
+        message = response_data["choices"][0]["message"]
+        assert message["mistral_reference_blocks"] == [self.REFERENCE_BLOCK]
+
+    def test_non_streaming_exposed_reference_blocks_carry_no_free_text(self):
+        # The output PII-masking hook sanitises `message.content` but not
+        # provider-specific fields, so the exposed copy must not repeat the
+        # text -- a second copy would reach the client unmasked.
+        content = [
+            {"type": "text", "text": "card 4111 1111 1111 1111"},
+            {"type": "reference", "reference_ids": ["doc1"]},
+        ]
+        response_data = {"choices": [{"message": {"role": "assistant", "content": content}}]}
+
+        MistralConfig._handle_content_list_to_str_conversion(response_data)
+
+        message = response_data["choices"][0]["message"]
+        exposed = message["mistral_reference_blocks"]
+        assert exposed == [self.REFERENCE_BLOCK]
+        assert not any("text" in block for block in exposed)
+
+    def test_non_streaming_adds_no_content_blocks_key_without_citations(self):
+        content = [{"type": "text", "text": "plain answer"}]
+        response_data = {"choices": [{"message": {"role": "assistant", "content": content}}]}
+
+        MistralConfig._handle_content_list_to_str_conversion(response_data)
+
+        message = response_data["choices"][0]["message"]
+        assert message["content"] == "plain answer"
+        assert "mistral_reference_blocks" not in message
+
+    def test_non_streaming_thinking_blocks_still_extracted(self):
+        response_data = {
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": [
+                            {"type": "thinking", "thinking": [{"type": "text", "text": "reasoning"}]},
+                            {"type": "text", "text": "answer part one"},
+                            {"type": "text", "text": " and two"},
+                        ],
+                    }
+                }
+            ]
+        }
+
+        MistralConfig._handle_content_list_to_str_conversion(response_data)
+
+        message = response_data["choices"][0]["message"]
+        assert message["content"] == "answer part one and two"
+        assert message["reasoning_content"] == "reasoning"
+
+    def test_non_streaming_reference_blocks_reach_model_response_provider_specific_fields(self):
+        # the exposed key must survive convert_to_model_response_object so
+        # callers actually receive the citation chunks
+        from litellm.utils import convert_to_model_response_object
+
+        response_data = {
+            "id": "chatcmpl-test",
+            "object": "chat.completion",
+            "created": 1,
+            "model": "mistral-medium-2604",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "stop",
+                    "message": {"role": "assistant", "content": self.CITATION_CONTENT},
+                }
+            ],
+        }
+
+        response_data = MistralConfig._handle_content_list_to_str_conversion(response_data)
+        model_response = convert_to_model_response_object(
+            response_object=response_data, model_response_object=ModelResponse()
+        )
+
+        provider_specific_fields = model_response.choices[0].message.provider_specific_fields
+        assert provider_specific_fields["mistral_reference_blocks"] == [self.REFERENCE_BLOCK]
+
+    def test_streaming_mixed_blocks_concatenate_and_expose_references(self):
+        iterator = MistralChatResponseIterator(streaming_response=iter([]), sync_stream=True)
+        chunk = {"id": "x", "choices": [{"delta": {"content": self.CITATION_CONTENT}}]}
+
+        result = iterator.chunk_parser(chunk)
+
+        delta = result.choices[0].delta
+        assert delta.content == "The sky is blue during the day."
+        assert delta.provider_specific_fields["mistral_reference_blocks"] == [self.REFERENCE_BLOCK]
+
+    def test_streaming_reference_only_delta_exposes_references(self):
+        iterator = MistralChatResponseIterator(streaming_response=iter([]), sync_stream=True)
+        content = [{"type": "reference", "reference_ids": ["doc1"]}]
+        chunk = {"id": "x", "choices": [{"delta": {"content": content}}]}
+
+        result = iterator.chunk_parser(chunk)
+
+        delta = result.choices[0].delta
+        assert not delta.content
+        assert delta.provider_specific_fields["mistral_reference_blocks"] == content