diff --git a/libs/partners/mistralai/langchain_mistralai/_compat.py b/libs/partners/mistralai/langchain_mistralai/_compat.py index 16f525791f7..aadcef84f74 100644 --- a/libs/partners/mistralai/langchain_mistralai/_compat.py +++ b/libs/partners/mistralai/langchain_mistralai/_compat.py @@ -16,7 +16,45 @@ def _convert_from_v1_to_mistral( new_content: list = [] for block in content: if block["type"] == "text": - new_content.append({"text": block.get("text", ""), "type": "text"}) + annotations = block.get("annotations") + if model_provider == "mistralai" and isinstance(annotations, list): + reference_meta: dict[str, Any] = {} + has_reference = False + for annotation in annotations: + if not isinstance(annotation, dict): + continue + ann_type = annotation.get("type") + if ann_type == "non_standard_annotation": + value = annotation.get("value") + if isinstance(value, dict) and value.get("type") == "reference": + reference_meta.update( + { + k: v + for k, v in value.items() + if k not in ("type", "text", "index") + } + ) + has_reference = True + elif ann_type == "citation": + extras = annotation.get("extras", {}) + if isinstance(extras, dict): + reference_meta.update(extras) + cited_text = annotation.get("cited_text") + if cited_text and cited_text != block.get("text", ""): + reference_meta["cited_text"] = cited_text + has_reference = True + if has_reference: + new_content.append( + { + "type": "text", + "text": block.get("text", ""), + "reference": reference_meta, + } + ) + else: + new_content.append({"type": "text", "text": block.get("text", "")}) + else: + new_content.append({"text": block.get("text", ""), "type": "text"}) elif ( block["type"] == "reasoning" @@ -66,6 +104,22 @@ def _convert_to_v1_from_mistral(message: AIMessage) -> list[types.ContentBlock]: } if "index" in block: text_block["index"] = block["index"] + # If the text block carries reference metadata (from a + # normalized Mistral citation chunk), attach it as a + # Citation annotation so downstream consumers can map + # answer fragments back to source documents. + if "reference" in block: + citation: types.Citation = {"type": "citation"} + ref_meta = block.get("reference") + if isinstance(ref_meta, dict): + if cited_text := ref_meta.get("cited_text"): + citation["cited_text"] = cited_text + extras = { + k: v for k, v in ref_meta.items() if k != "cited_text" + } + if extras: + citation["extras"] = extras + text_block["annotations"] = [citation] content_blocks.append(text_block) elif block.get("type") == "thinking" and isinstance( diff --git a/libs/partners/mistralai/langchain_mistralai/chat_models.py b/libs/partners/mistralai/langchain_mistralai/chat_models.py index be86042beb4..ef335674b18 100644 --- a/libs/partners/mistralai/langchain_mistralai/chat_models.py +++ b/libs/partners/mistralai/langchain_mistralai/chat_models.py @@ -146,6 +146,42 @@ def _convert_tool_call_id_to_mistral_compatible(tool_call_id: str) -> str: return base62_str.rjust(9, "0") +def _normalize_mistral_content(content: Any) -> str | list[str | dict]: + """Normalize Mistral content so reference blocks are visible to .text. + + Mistral citation responses return content as a list of typed chunks where + `reference` blocks carry visible answer text alongside citation metadata. + The core `.text` accessor only concatenates blocks whose type is + `"text"`, so preserving `reference` as-is would drop cited answer spans + from `message.text` and `ChatGeneration.text`. + + To keep the answer text visible while preserving citation metadata, rewrite + each `reference` block to `type: "text"` and move the original block + (including `reference_ids`) under a `"reference"` key. The `_compat.py` + translator reads that key to produce standard `Citation` annotations. + """ + if not isinstance(content, list): + return content or "" + has_reference = False + new_blocks: list[str | dict] = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "reference": + has_reference = True + new_block = { + "type": "text", + "text": block.get("text", ""), + "reference": { + k: v for k, v in block.items() if k not in ("type", "text") + }, + } + if "index" in block: + new_block["index"] = block["index"] + new_blocks.append(new_block) + else: + new_blocks.append(block) + return new_blocks if has_reference else content + + def _convert_mistral_chat_message_to_message( _message: dict, ) -> BaseMessage: @@ -153,8 +189,11 @@ def _convert_mistral_chat_message_to_message( if role != "assistant": msg = f"Expected role to be 'assistant', got {role}" raise ValueError(msg) - # Mistral returns None for tool invocations - content = _message.get("content", "") or "" + # Mistral returns None for tool invocations. When citations are enabled, + # content is a list of typed chunks (text and reference). Normalize + # reference blocks so their answer text is visible via .text while + # citation metadata is preserved for _compat.py to translate. + content = _normalize_mistral_content(_message.get("content", "")) additional_kwargs: dict = {} tool_calls = [] @@ -260,11 +299,13 @@ def _convert_chunk_to_message_chunk( content = _delta.get("content") or "" if output_version == "v1" and isinstance(content, str): content = [{"type": "text", "text": content}] + content = _normalize_mistral_content(content) if isinstance(content, list): for block in content: if isinstance(block, dict): - if "type" in block and block["type"] != index_type: - index_type = block["type"] + block_type = "reference" if "reference" in block else block.get("type") + if block_type is not None and block_type != index_type: + index_type = block_type index = index + 1 if "index" not in block: block["index"] = index @@ -278,7 +319,7 @@ def _convert_chunk_to_message_chunk( return HumanMessageChunk(content=content), index, index_type if role == "assistant" or default_class == AIMessageChunk: additional_kwargs: dict = {} - response_metadata = {} + response_metadata: dict[str, Any] = {} if raw_tool_calls := _delta.get("tool_calls"): additional_kwargs["tool_calls"] = raw_tool_calls try: @@ -360,7 +401,10 @@ def _format_invalid_tool_call_for_mistral(invalid_tool_call: InvalidToolCall) -> def _clean_block(block: dict) -> dict: - # Remove "index" key added for message aggregation in langchain-core + # Remove internal keys added by LangChain or by provider response normalization. + if block.get("type") == "text" and "text" in block: + return {"type": "text", "text": block["text"]} + new_block = {k: v for k, v in block.items() if k != "index"} if block.get("type") == "thinking" and isinstance(block.get("thinking"), list): new_block["thinking"] = [ @@ -480,9 +524,7 @@ def _convert_message_to_mistral_chat_message( elif isinstance(content, list): content = [ - _clean_block(block) - if isinstance(block, dict) and "index" in block - else block + _clean_block(block) if isinstance(block, dict) else block for block in content ] else: diff --git a/libs/partners/mistralai/tests/unit_tests/test_chat_models.py b/libs/partners/mistralai/tests/unit_tests/test_chat_models.py index 35757ffe442..ade2a4a4874 100644 --- a/libs/partners/mistralai/tests/unit_tests/test_chat_models.py +++ b/libs/partners/mistralai/tests/unit_tests/test_chat_models.py @@ -2,7 +2,7 @@ import os from collections.abc import AsyncGenerator, Generator -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from unittest.mock import MagicMock, patch import httpx @@ -10,6 +10,7 @@ import pytest from langchain_core.callbacks.base import BaseCallbackHandler from langchain_core.messages import ( AIMessage, + AIMessageChunk, BaseMessage, ChatMessage, HumanMessage, @@ -20,8 +21,13 @@ from langchain_core.messages import ( ) from pydantic import SecretStr +if TYPE_CHECKING: + from langchain_core.messages import content as types + +from langchain_mistralai._compat import _convert_to_v1_from_mistral from langchain_mistralai.chat_models import ( # type: ignore[import] ChatMistralAI, + _convert_chunk_to_message_chunk, _convert_message_to_mistral_chat_message, _convert_mistral_chat_message_to_message, _convert_tool_call_id_to_mistral_compatible, @@ -51,6 +57,42 @@ def test_sanitize_chat_completions_content_passthrough_string() -> None: assert _sanitize_chat_completions_content("hello") == "hello" +def test_ai_message_reference_metadata_does_not_reach_wire() -> None: + message = AIMessage( + content=[ + {"type": "text", "text": "The answer is "}, + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}}, + {"type": "text", "text": "."}, + ], + response_metadata={"model_provider": "mistralai"}, + ) + + result = _convert_message_to_mistral_chat_message(message) + assert result["content"] == [ + {"type": "text", "text": "The answer is "}, + {"type": "text", "text": "42"}, + {"type": "text", "text": "."}, + ] + + +def test_v1_ai_message_reference_metadata_does_not_reach_wire() -> None: + message = AIMessage( + content=[ + {"type": "text", "text": "The answer is "}, + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}}, + {"type": "text", "text": "."}, + ], + response_metadata={"model_provider": "mistralai", "output_version": "v1"}, + ) + + result = _convert_message_to_mistral_chat_message(message) + assert result["content"] == [ + {"type": "text", "text": "The answer is "}, + {"type": "text", "text": "42"}, + {"type": "text", "text": "."}, + ] + + def test_mistralai_model_param() -> None: llm = ChatMistralAI(model="foo") # type: ignore[call-arg] assert llm.model == "foo" @@ -488,6 +530,396 @@ def test__convert_dict_to_message_with_missing_content() -> None: assert result == expected_output +def test__convert_dict_to_message_with_citations() -> None: + """Reference blocks normalized to text blocks with reference metadata.""" + cited_text = "the temperature is 20 degrees C" + raw_content: list[str | dict] = [ + {"type": "text", "text": "According to the document, "}, + {"type": "reference", "reference_ids": [0], "text": cited_text}, + {"type": "text", "text": " on average."}, + ] + message = {"role": "assistant", "content": raw_content} + result = _convert_mistral_chat_message_to_message(message) + + assert isinstance(result.content, list) + content = result.content + # The reference block is normalized to type="text" so .text includes it + assert content[0] == {"type": "text", "text": "According to the document, "} + assert isinstance(content[1], dict) + block_1 = content[1] + assert block_1["type"] == "text" + assert block_1["text"] == cited_text + assert block_1["reference"] == {"reference_ids": [0]} + assert content[2] == {"type": "text", "text": " on average."} + assert result.response_metadata["model_provider"] == "mistralai" + assert "citations" not in result.response_metadata + + +def test__convert_dict_to_message_citations_text_accessor() -> None: + """message.text includes cited spans from normalized reference blocks.""" + cited_text = "the temperature is 20 degrees C" + raw_content: list[str | dict] = [ + {"type": "text", "text": "According to the document, "}, + {"type": "reference", "reference_ids": [0], "text": cited_text}, + {"type": "text", "text": " on average."}, + ] + message = {"role": "assistant", "content": raw_content} + result = _convert_mistral_chat_message_to_message(message) + + # .text should include all visible text, including the cited span + assert str(result.text) == ( + "According to the document, the temperature is 20 degrees C on average." + ) + + +def test__convert_dict_to_message_citations_to_content_blocks() -> None: + """content_blocks translates reference metadata to TextContentBlock.""" + cited_text = "the temperature is 20 degrees C" + raw_content: list[str | dict] = [ + {"type": "text", "text": "According to the document, "}, + {"type": "reference", "reference_ids": [0], "text": cited_text}, + {"type": "text", "text": " on average."}, + ] + message = {"role": "assistant", "content": raw_content} + result = _convert_mistral_chat_message_to_message(message) + + assert isinstance(result, AIMessage) + blocks = _convert_to_v1_from_mistral(result) + assert len(blocks) == 3 + + # First block: plain text + assert blocks[0]["type"] == "text" + assert blocks[0]["text"] == "According to the document, " + + # Second block: text with citation annotation + block_1 = cast("types.TextContentBlock", blocks[1]) + assert block_1["type"] == "text" + assert block_1["text"] == cited_text + annotations = block_1["annotations"] + assert len(annotations) == 1 + assert annotations[0]["type"] == "citation" + assert "cited_text" not in annotations[0] + assert annotations[0]["extras"]["reference_ids"] == [0] + + # Third block: plain text + assert blocks[2]["type"] == "text" + assert blocks[2]["text"] == " on average." + + +def test_create_chat_result_with_citations() -> None: + """Citations are normalized to text blocks with reference metadata in .content.""" + chat = ChatMistralAI() + raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"} + raw_content: list[str | dict] = [ + {"type": "text", "text": "The answer is "}, + raw_citation, + {"type": "text", "text": "."}, + ] + response = { + "choices": [ + { + "message": { + "role": "assistant", + "content": raw_content, + }, + "finish_reason": "stop", + } + ] + } + + result = chat._create_chat_result(response) + message = result.generations[0].message + + assert isinstance(message.content, list) + content = message.content + # The reference block is normalized; .text includes the cited span + assert isinstance(content[1], dict) + block_1 = content[1] + assert block_1["type"] == "text" + assert block_1["text"] == "42" + assert block_1["reference"] == {"reference_ids": [0]} + assert str(message.text) == "The answer is 42." + assert "citations" not in message.response_metadata + + +def test__convert_chunk_to_message_chunk_with_citations() -> None: + """Streaming reference blocks are normalized to text blocks in chunk .content.""" + raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"} + text_chunk = { + "choices": [ + { + "delta": {"role": "assistant", "content": "The answer is "}, + "finish_reason": None, + } + ], + } + reference_chunk = { + "choices": [ + { + "delta": { + "role": "assistant", + "content": [ + dict(raw_citation), + ], + }, + "finish_reason": "stop", + } + ], + "model": "mistral-small-latest", + } + + result_1, index, index_type = _convert_chunk_to_message_chunk( + text_chunk, AIMessageChunk, -1, "", None + ) + result_2, _, _ = _convert_chunk_to_message_chunk( + reference_chunk, AIMessageChunk, index, index_type, None + ) + + assert isinstance(result_2, AIMessageChunk) + # Reference block is normalized to type="text" with reference metadata + assert result_2.content == [ + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0}, + ] + assert "citations" not in result_2.response_metadata + + full = result_1 + result_2 + assert isinstance(full, AIMessageChunk) + assert "citations" not in full.response_metadata + assert full.response_metadata["finish_reason"] == "stop" + # .text includes the cited span + assert str(full.text) == "The answer is 42" + + +def test_citation_round_trip() -> None: + """Round-trip through v1 preserves text and reference metadata.""" + from langchain_mistralai._compat import ( + _convert_from_v1_to_mistral, + _convert_to_v1_from_mistral, + ) + + # Start with normalized content (as produced by _convert_mistral_chat_message) + original_content: list[str | dict] = [ + {"type": "text", "text": "The answer is "}, + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}}, + {"type": "text", "text": "."}, + ] + message = AIMessage(content=original_content) + v1_blocks = _convert_to_v1_from_mistral(message) + round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai") + + # Should have exactly 3 blocks, no duplication of cited text + assert len(round_tripped) == 3 + assert round_tripped[0] == {"type": "text", "text": "The answer is "} + assert isinstance(round_tripped[1], dict) + block_1 = round_tripped[1] + assert block_1["type"] == "text" + assert block_1["text"] == "42" + assert block_1["reference"] == {"reference_ids": [0]} + assert round_tripped[2] == {"type": "text", "text": "."} + + +def test_citation_round_trip_preserves_extra_fields() -> None: + """Extra provider fields on reference metadata survive the round-trip.""" + from langchain_mistralai._compat import ( + _convert_from_v1_to_mistral, + _convert_to_v1_from_mistral, + ) + + original_content: list[str | dict] = [ + {"type": "text", "text": "cited span", "reference": {"reference_ids": [1, 2]}}, + ] + message = AIMessage(content=original_content) + v1_blocks = _convert_to_v1_from_mistral(message) + round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai") + + assert len(round_tripped) == 1 + assert isinstance(round_tripped[0], dict) + block_0 = round_tripped[0] + assert block_0["type"] == "text" + assert block_0["text"] == "cited span" + assert block_0["reference"] == {"reference_ids": [1, 2]} + + +def test_citation_round_trip_preserves_annotated_response_text() -> None: + """Serializing citations preserves block text, not citation source excerpts.""" + from langchain_mistralai._compat import _convert_from_v1_to_mistral + + content: list[types.ContentBlock] = [ + { + "type": "text", + "text": "The answer is 42.", + "annotations": [ + { + "type": "citation", + "cited_text": "source excerpt mentioning 42", + "extras": {"reference_ids": [0]}, + } + ], + } + ] + round_tripped = _convert_from_v1_to_mistral(content, "mistralai") + + assert len(round_tripped) == 1 + assert isinstance(round_tripped[0], dict) + block = round_tripped[0] + assert block["type"] == "text" + assert block["text"] == "The answer is 42." + assert block["reference"]["reference_ids"] == [0] + assert block["reference"]["cited_text"] == "source excerpt mentioning 42" + + +def test_citation_streaming_v1_reference_gets_separate_index() -> None: + """Reference chunks do not merge into surrounding v1 text block indexes.""" + text_chunk = { + "choices": [ + { + "delta": {"role": "assistant", "content": "The answer is "}, + "finish_reason": None, + } + ], + } + reference_chunk = { + "choices": [ + { + "delta": { + "role": "assistant", + "content": [ + {"type": "reference", "reference_ids": [0], "text": "42"}, + ], + }, + "finish_reason": "stop", + } + ], + "model": "mistral-small-latest", + } + + result_1, index, index_type = _convert_chunk_to_message_chunk( + text_chunk, AIMessageChunk, -1, "", "v1" + ) + result_2, _, _ = _convert_chunk_to_message_chunk( + reference_chunk, AIMessageChunk, index, index_type, "v1" + ) + + assert result_1.content == [{"type": "text", "text": "The answer is ", "index": 0}] + assert result_2.content == [ + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 1}, + ] + + +def test_citation_streaming_accumulated_content() -> None: + """Streaming chunks accumulate normalized text blocks in full.content.""" + raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"} + text_chunk = { + "choices": [ + { + "delta": {"role": "assistant", "content": "The answer is "}, + "finish_reason": None, + } + ], + } + reference_chunk = { + "choices": [ + { + "delta": { + "role": "assistant", + "content": [dict(raw_citation)], + }, + "finish_reason": "stop", + } + ], + "model": "mistral-small-latest", + } + + result_1, index, index_type = _convert_chunk_to_message_chunk( + text_chunk, AIMessageChunk, -1, "", None + ) + result_2, _, _ = _convert_chunk_to_message_chunk( + reference_chunk, AIMessageChunk, index, index_type, None + ) + + full = result_1 + result_2 + # full.content should contain both the text and the normalized reference block + assert isinstance(full.content, list) + assert any( + isinstance(b, dict) + and b.get("type") == "text" + and b.get("text") == "42" + and isinstance(ref := b.get("reference"), dict) + and ref.get("reference_ids") == [0] + for b in full.content + ) + + +def test_citation_index_not_in_extras() -> None: + """Streaming index should not leak into citation extras.""" + from langchain_mistralai._compat import _convert_to_v1_from_mistral + + content: list[str | dict] = [ + {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0}, + ] + message = AIMessageChunk(content=content) + blocks = _convert_to_v1_from_mistral(message) + assert len(blocks) == 1 + block_0 = cast("types.TextContentBlock", blocks[0]) + annotation = block_0["annotations"][0] + extras = annotation.get("extras", {}) + assert isinstance(extras, dict) + assert "index" not in extras + + +def test_citation_no_text_in_reference() -> None: + """A reference block with no text still converts without error.""" + from langchain_mistralai._compat import _convert_to_v1_from_mistral + + content: list[str | dict] = [ + {"type": "text", "text": "", "reference": {"reference_ids": [0]}}, + ] + message = AIMessage(content=content) + blocks = _convert_to_v1_from_mistral(message) + assert len(blocks) == 1 + assert blocks[0]["type"] == "text" + assert blocks[0]["text"] == "" + block_0 = cast("types.TextContentBlock", blocks[0]) + assert "cited_text" not in block_0["annotations"][0] + + +def test_citation_empty_reference_metadata_still_adds_annotation() -> None: + """Presence of reference metadata is the signal, even if the metadata is empty.""" + from langchain_mistralai._compat import _convert_to_v1_from_mistral + + message = AIMessage(content=[{"type": "text", "text": "42", "reference": {}}]) + blocks = _convert_to_v1_from_mistral(message) + + block_0 = cast("types.TextContentBlock", blocks[0]) + assert block_0["annotations"] == [{"type": "citation"}] + + +def test_malformed_annotation_does_not_crash() -> None: + """Malformed annotations are skipped, not raised.""" + from langchain_mistralai._compat import _convert_from_v1_to_mistral + + content: list = [ + { + "type": "text", + "text": "hello", + "annotations": [ + None, # not a dict + {"type": "unknown"}, # unrecognized type + {"type": "citation", "cited_text": "cited"}, # valid + ], + } + ] + result = _convert_from_v1_to_mistral(content, "mistralai") + # The valid citation produces a text block with reference metadata; + # the text block is not appended because a reference was emitted. + assert len(result) == 1 + assert isinstance(result[0], dict) + block_0 = result[0] + assert block_0["type"] == "text" + assert block_0["text"] == "hello" + assert "reference" in block_0 + + def test_custom_token_counting() -> None: def token_encoder(text: str) -> list[int]: return [1, 2, 3] diff --git a/libs/partners/mistralai/uv.lock b/libs/partners/mistralai/uv.lock index 5ab51a0a47f..702e221ac5f 100644 --- a/libs/partners/mistralai/uv.lock +++ b/libs/partners/mistralai/uv.lock @@ -257,7 +257,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [