1"""Test MistralAI Chat API wrapper."""23import os4from collections.abc import AsyncGenerator, Generator5from typing import TYPE_CHECKING, Any, cast6from unittest.mock import MagicMock, patch78import httpx9import pytest10from langchain_core.callbacks.base import BaseCallbackHandler11from langchain_core.exceptions import (12    ModelAPIError,13    ModelAuthenticationError,14    ModelError,15    ModelInvalidRequestError,16    ModelNotFoundError,17    ModelPermissionDeniedError,18    ModelRateLimitError,19)20from langchain_core.messages import (21    AIMessage,22    AIMessageChunk,23    BaseMessage,24    ChatMessage,25    HumanMessage,26    InvalidToolCall,27    SystemMessage,28    ToolCall,29    ToolMessage,30)31from pydantic import SecretStr3233if TYPE_CHECKING:34    from langchain_core.messages import content as types3536from langchain_mistralai._compat import _convert_to_v1_from_mistral37from langchain_mistralai.chat_models import (  # type: ignore[import]38    ChatMistralAI,39    _araise_on_error,40    _convert_chunk_to_message_chunk,41    _convert_message_to_mistral_chat_message,42    _convert_mistral_chat_message_to_message,43    _convert_tool_call_id_to_mistral_compatible,44    _format_message_content,45    _is_valid_mistral_tool_call_id,46    _raise_on_error,47    _sanitize_chat_completions_content,48)4950os.environ["MISTRAL_API_KEY"] = "foo"515253def _error_response(status_code: int) -> httpx.Response:54    """Build a response with the given status for the error handlers."""55    request = httpx.Request("POST", "https://api.mistral.ai/v1/chat/completions")56    return httpx.Response(status_code, request=request, content=b'{"message": "boom"}')575859def test_sanitize_chat_completions_text_blocks_strips_id() -> None:60    """LangChain auto-generated `id` on text blocks must not reach the wire.6162    Mistral's chat completions endpoint returns 422 with `extra_forbidden`63    on `messages[*].tool.content.list[...].text.id` if not stripped.64    """65    message = ToolMessage(66        content=[{"type": "text", "text": "foo", "id": "lc_abc123"}],67        tool_call_id="abc12345",68    )69    result = _convert_message_to_mistral_chat_message(message)70    assert result["content"] == [{"type": "text", "text": "foo"}]717273def test_sanitize_chat_completions_content_passthrough_string() -> None:74    assert _sanitize_chat_completions_content("hello") == "hello"757677def test_ai_message_reference_metadata_does_not_reach_wire() -> None:78    message = AIMessage(79        content=[80            {"type": "text", "text": "The answer is "},81            {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},82            {"type": "text", "text": "."},83        ],84        response_metadata={"model_provider": "mistralai"},85    )8687    result = _convert_message_to_mistral_chat_message(message)88    assert result["content"] == [89        {"type": "text", "text": "The answer is "},90        {"type": "text", "text": "42"},91        {"type": "text", "text": "."},92    ]939495def test_v1_ai_message_reference_metadata_does_not_reach_wire() -> None:96    message = AIMessage(97        content=[98            {"type": "text", "text": "The answer is "},99            {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},100            {"type": "text", "text": "."},101        ],102        response_metadata={"model_provider": "mistralai", "output_version": "v1"},103    )104105    result = _convert_message_to_mistral_chat_message(message)106    assert result["content"] == [107        {"type": "text", "text": "The answer is "},108        {"type": "text", "text": "42"},109        {"type": "text", "text": "."},110    ]111112113def test_mistralai_model_param() -> None:114    llm = ChatMistralAI(model="foo")  # type: ignore[call-arg]115    assert llm.model == "foo"116117118def test_mistralai_initialization() -> None:119    """Test ChatMistralAI initialization."""120    # Verify that ChatMistralAI can be initialized using a secret key provided121    # as a parameter rather than an environment variable.122    for model in [123        ChatMistralAI(model="test", mistral_api_key="test"),  # type: ignore[call-arg, call-arg]124        ChatMistralAI(model="test", api_key="test"),  # type: ignore[call-arg, arg-type]125    ]:126        assert cast("SecretStr", model.mistral_api_key).get_secret_value() == "test"127128129@pytest.mark.parametrize(130    ("model", "expected_url"),131    [132        (ChatMistralAI(model="test"), "https://api.mistral.ai/v1"),  # type: ignore[call-arg, arg-type]133        (ChatMistralAI(model="test", endpoint="baz"), "baz"),  # type: ignore[call-arg, arg-type]134    ],135)136def test_mistralai_initialization_baseurl(137    model: ChatMistralAI, expected_url: str138) -> None:139    """Test ChatMistralAI initialization."""140    # Verify that ChatMistralAI can be initialized providing endpoint, but also141    # with default142143    assert model.endpoint == expected_url144145146@pytest.mark.parametrize(147    "env_var_name",148    [149        ("MISTRAL_BASE_URL"),150    ],151)152def test_mistralai_initialization_baseurl_env(153    env_var_name: str, monkeypatch: pytest.MonkeyPatch154) -> None:155    """Test ChatMistralAI initialization."""156    # Verify that ChatMistralAI can be initialized using env variable157    monkeypatch.setenv(env_var_name, "boo")158    model = ChatMistralAI(model="test")  # type: ignore[call-arg]159    assert model.endpoint == "boo"160161162@pytest.mark.parametrize(163    ("message", "expected"),164    [165        (166            SystemMessage(content="Hello"),167            {"role": "system", "content": "Hello"},168        ),169        (170            HumanMessage(content="Hello"),171            {"role": "user", "content": "Hello"},172        ),173        (174            AIMessage(content="Hello"),175            {"role": "assistant", "content": "Hello"},176        ),177        (178            AIMessage(content="{", additional_kwargs={"prefix": True}),179            {"role": "assistant", "content": "{", "prefix": True},180        ),181        (182            ChatMessage(role="assistant", content="Hello"),183            {"role": "assistant", "content": "Hello"},184        ),185    ],186)187def test_convert_message_to_mistral_chat_message(188    message: BaseMessage, expected: dict189) -> None:190    result = _convert_message_to_mistral_chat_message(message)191    assert result == expected192193194@pytest.mark.parametrize(195    ("content", "expected"),196    [197        ("hello", "hello"),198        ("", ""),199        (None, None),200        ([], []),201    ],202)203def test_format_message_content_passthrough_non_list(204    content: Any, expected: Any205) -> None:206    """Strings, None, and empty lists pass through `_format_message_content`."""207    assert _format_message_content(content) == expected208209210@pytest.mark.parametrize(211    ("block", "expected"),212    [213        (214            {"type": "image", "url": "https://example.com/img.png"},215            {216                "type": "image_url",217                "image_url": {"url": "https://example.com/img.png"},218            },219        ),220        (221            {"type": "image", "base64": "abc123", "mime_type": "image/jpeg"},222            {223                "type": "image_url",224                "image_url": {"url": "data:image/jpeg;base64,abc123"},225            },226        ),227        (228            {229                "type": "image",230                "source_type": "url",231                "url": "https://example.com/v0.png",232            },233            {234                "type": "image_url",235                "image_url": {"url": "https://example.com/v0.png"},236            },237        ),238        (239            {240                "type": "image",241                "source_type": "base64",242                "data": "v0data",243                "mime_type": "image/png",244            },245            {246                "type": "image_url",247                "image_url": {"url": "data:image/png;base64,v0data"},248            },249        ),250    ],251)252def test_format_message_content_translates_image_blocks(253    block: dict, expected: dict254) -> None:255    """v0 and v1 canonical image blocks translate to Mistral's `image_url` shape."""256    assert _format_message_content([block]) == [expected]257258259@pytest.mark.parametrize(260    "block",261    [262        {"type": "text", "text": "hello"},263        {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},264        {"type": "image_url", "image_url": "https://example.com/img.png"},265    ],266)267def test_format_message_content_passthrough_known_blocks(block: dict) -> None:268    """Already-translated wire blocks and text blocks pass through unchanged."""269    assert _format_message_content([block]) == [block]270271272@pytest.mark.parametrize(273    "block_type",274    ["tool_use", "thinking", "reasoning_content", "document_url", "input_audio"],275)276def test_format_message_content_passes_unknown_blocks_through(block_type: str) -> None:277    """Non-canonical blocks pass through; the Mistral API validates them."""278    blocks = [279        {"type": "text", "text": "kept"},280        {"type": block_type, "data": "anything"},281    ]282    assert _format_message_content(blocks) == blocks283284285def test_format_message_content_preserves_order_for_mixed_blocks() -> None:286    """Multiple text + image blocks retain their order — vision prompts depend on it."""287    blocks: list[Any] = [288        {"type": "text", "text": "first"},289        {"type": "image", "url": "https://example.com/a.png"},290        {"type": "text", "text": "between"},291        {"type": "image", "base64": "xyz", "mime_type": "image/png"},292        "trailing string",293    ]294    expected = [295        {"type": "text", "text": "first"},296        {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},297        {"type": "text", "text": "between"},298        {"type": "image_url", "image_url": {"url": "data:image/png;base64,xyz"}},299        "trailing string",300    ]301    assert _format_message_content(blocks) == expected302303304def test_format_message_content_image_missing_mime_type_raises() -> None:305    """Base64 image without `mime_type` raises via the core translator."""306    with pytest.raises(ValueError, match="mime_type"):307        _format_message_content([{"type": "image", "base64": "abc"}])308309310@pytest.mark.parametrize(311    ("message", "expected"),312    [313        (314            HumanMessage(315                content=[316                    {"type": "text", "text": "What is in this image?"},317                    {"type": "image", "url": "https://example.com/img.png"},318                ]319            ),320            {321                "role": "user",322                "content": [323                    {"type": "text", "text": "What is in this image?"},324                    {325                        "type": "image_url",326                        "image_url": {"url": "https://example.com/img.png"},327                    },328                ],329            },330        ),331        (332            HumanMessage(333                content=[334                    {"type": "text", "text": "Describe this image."},335                    {336                        "type": "image",337                        "base64": "abc123",338                        "mime_type": "image/png",339                    },340                ]341            ),342            {343                "role": "user",344                "content": [345                    {"type": "text", "text": "Describe this image."},346                    {347                        "type": "image_url",348                        "image_url": {"url": "data:image/png;base64,abc123"},349                    },350                ],351            },352        ),353    ],354)355def test_convert_human_message_with_images(356    message: BaseMessage, expected: dict357) -> None:358    result = _convert_message_to_mistral_chat_message(message)359    assert result == expected360361362def test_convert_human_message_with_string_content_unchanged() -> None:363    """Plain string `HumanMessage` content is not wrapped or modified."""364    result = _convert_message_to_mistral_chat_message(HumanMessage(content="hi"))365    assert result == {"role": "user", "content": "hi"}366367368def _make_completion_response_from_token(token: str) -> dict:369    return {370        "id": "abc123",371        "model": "fake_model",372        "choices": [373            {374                "index": 0,375                "delta": {"content": token},376                "finish_reason": None,377            }378        ],379    }380381382def mock_chat_stream(*args: Any, **kwargs: Any) -> Generator:383    def it() -> Generator:384        for token in ["Hello", " how", " can", " I", " help", "?"]:385            yield _make_completion_response_from_token(token)386387    return it()388389390async def mock_chat_astream(*args: Any, **kwargs: Any) -> AsyncGenerator:391    async def it() -> AsyncGenerator:392        for token in ["Hello", " how", " can", " I", " help", "?"]:393            yield _make_completion_response_from_token(token)394395    return it()396397398class MyCustomHandler(BaseCallbackHandler):399    last_token: str = ""400401    def on_llm_new_token(402        self, token: str | list[str | dict[str, Any]], **kwargs: Any403    ) -> None:404        if isinstance(token, str):405            self.last_token = token406407408@patch(409    "langchain_mistralai.chat_models.ChatMistralAI.completion_with_retry",410    new=mock_chat_stream,411)412def test_stream_with_callback() -> None:413    callback = MyCustomHandler()414    chat = ChatMistralAI(callbacks=[callback])415    for token in chat.stream("Hello"):416        assert callback.last_token == token.content417418419@patch("langchain_mistralai.chat_models.acompletion_with_retry", new=mock_chat_astream)420async def test_astream_with_callback() -> None:421    callback = MyCustomHandler()422    chat = ChatMistralAI(callbacks=[callback])423    async for token in chat.astream("Hello"):424        assert callback.last_token == token.content425426427def test__convert_dict_to_message_tool_call() -> None:428    raw_tool_call = {429        "id": "ssAbar4Dr",430        "function": {431            "arguments": '{"name": "Sally", "hair_color": "green"}',432            "name": "GenerateUsername",433        },434    }435    message = {"role": "assistant", "content": "", "tool_calls": [raw_tool_call]}436    result = _convert_mistral_chat_message_to_message(message)437    expected_output = AIMessage(438        content="",439        additional_kwargs={"tool_calls": [raw_tool_call]},440        tool_calls=[441            ToolCall(442                name="GenerateUsername",443                args={"name": "Sally", "hair_color": "green"},444                id="ssAbar4Dr",445                type="tool_call",446            )447        ],448        response_metadata={"model_provider": "mistralai"},449    )450    assert result == expected_output451    assert _convert_message_to_mistral_chat_message(expected_output) == message452453    # Test malformed tool call454    raw_tool_calls = [455        {456            "id": "pL5rEGzxe",457            "function": {458                "arguments": '{"name": "Sally", "hair_color": "green"}',459                "name": "GenerateUsername",460            },461        },462        {463            "id": "ssAbar4Dr",464            "function": {465                "arguments": "oops",466                "name": "GenerateUsername",467            },468        },469    ]470    message = {"role": "assistant", "content": "", "tool_calls": raw_tool_calls}471    result = _convert_mistral_chat_message_to_message(message)472    expected_output = AIMessage(473        content="",474        additional_kwargs={"tool_calls": raw_tool_calls},475        invalid_tool_calls=[476            InvalidToolCall(477                name="GenerateUsername",478                args="oops",479                error="Function GenerateUsername arguments:\n\noops\n\nare not valid JSON. Received JSONDecodeError Expecting value: line 1 column 1 (char 0)\nFor troubleshooting, visit: https://docs.langchain.com/oss/python/langchain/errors/OUTPUT_PARSING_FAILURE ",  # noqa: E501480                id="ssAbar4Dr",481                type="invalid_tool_call",482            ),483        ],484        tool_calls=[485            ToolCall(486                name="GenerateUsername",487                args={"name": "Sally", "hair_color": "green"},488                id="pL5rEGzxe",489                type="tool_call",490            ),491        ],492        response_metadata={"model_provider": "mistralai"},493    )494    assert result == expected_output495    assert _convert_message_to_mistral_chat_message(expected_output) == message496497498def test__convert_dict_to_message_tool_call_with_null_content() -> None:499    raw_tool_call = {500        "id": "ssAbar4Dr",501        "function": {502            "arguments": '{"name": "Sally", "hair_color": "green"}',503            "name": "GenerateUsername",504        },505    }506    message = {"role": "assistant", "content": None, "tool_calls": [raw_tool_call]}507    result = _convert_mistral_chat_message_to_message(message)508    expected_output = AIMessage(509        content="",510        additional_kwargs={"tool_calls": [raw_tool_call]},511        tool_calls=[512            ToolCall(513                name="GenerateUsername",514                args={"name": "Sally", "hair_color": "green"},515                id="ssAbar4Dr",516                type="tool_call",517            )518        ],519        response_metadata={"model_provider": "mistralai"},520    )521    assert result == expected_output522523524def test__convert_dict_to_message_with_missing_content() -> None:525    raw_tool_call = {526        "id": "ssAbar4Dr",527        "function": {528            "arguments": '{"query": "test search"}',529            "name": "search",530        },531    }532    message = {"role": "assistant", "tool_calls": [raw_tool_call]}533    result = _convert_mistral_chat_message_to_message(message)534    expected_output = AIMessage(535        content="",536        additional_kwargs={"tool_calls": [raw_tool_call]},537        tool_calls=[538            ToolCall(539                name="search",540                args={"query": "test search"},541                id="ssAbar4Dr",542                type="tool_call",543            )544        ],545        response_metadata={"model_provider": "mistralai"},546    )547    assert result == expected_output548549550def test__convert_dict_to_message_with_citations() -> None:551    """Reference blocks normalized to text blocks with reference metadata."""552    cited_text = "the temperature is 20 degrees C"553    raw_content: list[str | dict] = [554        {"type": "text", "text": "According to the document, "},555        {"type": "reference", "reference_ids": [0], "text": cited_text},556        {"type": "text", "text": " on average."},557    ]558    message = {"role": "assistant", "content": raw_content}559    result = _convert_mistral_chat_message_to_message(message)560561    assert isinstance(result.content, list)562    content = result.content563    # The reference block is normalized to type="text" so .text includes it564    assert content[0] == {"type": "text", "text": "According to the document, "}565    assert isinstance(content[1], dict)566    block_1 = content[1]567    assert block_1["type"] == "text"568    assert block_1["text"] == cited_text569    assert block_1["reference"] == {"reference_ids": [0]}570    assert content[2] == {"type": "text", "text": " on average."}571    assert result.response_metadata["model_provider"] == "mistralai"572    assert "citations" not in result.response_metadata573574575def test__convert_dict_to_message_citations_text_accessor() -> None:576    """message.text includes cited spans from normalized reference blocks."""577    cited_text = "the temperature is 20 degrees C"578    raw_content: list[str | dict] = [579        {"type": "text", "text": "According to the document, "},580        {"type": "reference", "reference_ids": [0], "text": cited_text},581        {"type": "text", "text": " on average."},582    ]583    message = {"role": "assistant", "content": raw_content}584    result = _convert_mistral_chat_message_to_message(message)585586    # .text should include all visible text, including the cited span587    assert str(result.text) == (588        "According to the document, the temperature is 20 degrees C on average."589    )590591592def test__convert_dict_to_message_citations_to_content_blocks() -> None:593    """content_blocks translates reference metadata to TextContentBlock."""594    cited_text = "the temperature is 20 degrees C"595    raw_content: list[str | dict] = [596        {"type": "text", "text": "According to the document, "},597        {"type": "reference", "reference_ids": [0], "text": cited_text},598        {"type": "text", "text": " on average."},599    ]600    message = {"role": "assistant", "content": raw_content}601    result = _convert_mistral_chat_message_to_message(message)602603    assert isinstance(result, AIMessage)604    blocks = _convert_to_v1_from_mistral(result)605    assert len(blocks) == 3606607    # First block: plain text608    assert blocks[0]["type"] == "text"609    assert blocks[0]["text"] == "According to the document, "610611    # Second block: text with citation annotation612    block_1 = cast("types.TextContentBlock", blocks[1])613    assert block_1["type"] == "text"614    assert block_1["text"] == cited_text615    annotations = block_1["annotations"]616    assert len(annotations) == 1617    assert annotations[0]["type"] == "citation"618    assert "cited_text" not in annotations[0]619    assert annotations[0]["extras"]["reference_ids"] == [0]620621    # Third block: plain text622    assert blocks[2]["type"] == "text"623    assert blocks[2]["text"] == " on average."624625626def test_create_chat_result_with_citations() -> None:627    """Citations are normalized to text blocks with reference metadata in .content."""628    chat = ChatMistralAI()629    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}630    raw_content: list[str | dict] = [631        {"type": "text", "text": "The answer is "},632        raw_citation,633        {"type": "text", "text": "."},634    ]635    response = {636        "choices": [637            {638                "message": {639                    "role": "assistant",640                    "content": raw_content,641                },642                "finish_reason": "stop",643            }644        ]645    }646647    result = chat._create_chat_result(response)648    message = result.generations[0].message649650    assert isinstance(message.content, list)651    content = message.content652    # The reference block is normalized; .text includes the cited span653    assert isinstance(content[1], dict)654    block_1 = content[1]655    assert block_1["type"] == "text"656    assert block_1["text"] == "42"657    assert block_1["reference"] == {"reference_ids": [0]}658    assert str(message.text) == "The answer is 42."659    assert "citations" not in message.response_metadata660661662def test__convert_chunk_to_message_chunk_with_citations() -> None:663    """Streaming reference blocks are normalized to text blocks in chunk .content."""664    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}665    text_chunk = {666        "choices": [667            {668                "delta": {"role": "assistant", "content": "The answer is "},669                "finish_reason": None,670            }671        ],672    }673    reference_chunk = {674        "choices": [675            {676                "delta": {677                    "role": "assistant",678                    "content": [679                        dict(raw_citation),680                    ],681                },682                "finish_reason": "stop",683            }684        ],685        "model": "mistral-small-latest",686    }687688    result_1, index, index_type = _convert_chunk_to_message_chunk(689        text_chunk, AIMessageChunk, -1, "", None690    )691    result_2, _, _ = _convert_chunk_to_message_chunk(692        reference_chunk, AIMessageChunk, index, index_type, None693    )694695    assert isinstance(result_2, AIMessageChunk)696    # Reference block is normalized to type="text" with reference metadata697    assert result_2.content == [698        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0},699    ]700    assert "citations" not in result_2.response_metadata701702    full = result_1 + result_2703    assert isinstance(full, AIMessageChunk)704    assert "citations" not in full.response_metadata705    assert full.response_metadata["finish_reason"] == "stop"706    # .text includes the cited span707    assert str(full.text) == "The answer is 42"708709710def test_citation_round_trip() -> None:711    """Round-trip through v1 preserves text and reference metadata."""712    from langchain_mistralai._compat import (713        _convert_from_v1_to_mistral,714        _convert_to_v1_from_mistral,715    )716717    # Start with normalized content (as produced by _convert_mistral_chat_message)718    original_content: list[str | dict] = [719        {"type": "text", "text": "The answer is "},720        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},721        {"type": "text", "text": "."},722    ]723    message = AIMessage(content=original_content)724    v1_blocks = _convert_to_v1_from_mistral(message)725    round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai")726727    # Should have exactly 3 blocks, no duplication of cited text728    assert len(round_tripped) == 3729    assert round_tripped[0] == {"type": "text", "text": "The answer is "}730    assert isinstance(round_tripped[1], dict)731    block_1 = round_tripped[1]732    assert block_1["type"] == "text"733    assert block_1["text"] == "42"734    assert block_1["reference"] == {"reference_ids": [0]}735    assert round_tripped[2] == {"type": "text", "text": "."}736737738def test_citation_round_trip_preserves_extra_fields() -> None:739    """Extra provider fields on reference metadata survive the round-trip."""740    from langchain_mistralai._compat import (741        _convert_from_v1_to_mistral,742        _convert_to_v1_from_mistral,743    )744745    original_content: list[str | dict] = [746        {"type": "text", "text": "cited span", "reference": {"reference_ids": [1, 2]}},747    ]748    message = AIMessage(content=original_content)749    v1_blocks = _convert_to_v1_from_mistral(message)750    round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai")751752    assert len(round_tripped) == 1753    assert isinstance(round_tripped[0], dict)754    block_0 = round_tripped[0]755    assert block_0["type"] == "text"756    assert block_0["text"] == "cited span"757    assert block_0["reference"] == {"reference_ids": [1, 2]}758759760def test_citation_round_trip_preserves_annotated_response_text() -> None:761    """Serializing citations preserves block text, not citation source excerpts."""762    from langchain_mistralai._compat import _convert_from_v1_to_mistral763764    content: list[types.ContentBlock] = [765        {766            "type": "text",767            "text": "The answer is 42.",768            "annotations": [769                {770                    "type": "citation",771                    "cited_text": "source excerpt mentioning 42",772                    "extras": {"reference_ids": [0]},773                }774            ],775        }776    ]777    round_tripped = _convert_from_v1_to_mistral(content, "mistralai")778779    assert len(round_tripped) == 1780    assert isinstance(round_tripped[0], dict)781    block = round_tripped[0]782    assert block["type"] == "text"783    assert block["text"] == "The answer is 42."784    assert block["reference"]["reference_ids"] == [0]785    assert block["reference"]["cited_text"] == "source excerpt mentioning 42"786787788def test_citation_streaming_v1_reference_gets_separate_index() -> None:789    """Reference chunks do not merge into surrounding v1 text block indexes."""790    text_chunk = {791        "choices": [792            {793                "delta": {"role": "assistant", "content": "The answer is "},794                "finish_reason": None,795            }796        ],797    }798    reference_chunk = {799        "choices": [800            {801                "delta": {802                    "role": "assistant",803                    "content": [804                        {"type": "reference", "reference_ids": [0], "text": "42"},805                    ],806                },807                "finish_reason": "stop",808            }809        ],810        "model": "mistral-small-latest",811    }812813    result_1, index, index_type = _convert_chunk_to_message_chunk(814        text_chunk, AIMessageChunk, -1, "", "v1"815    )816    result_2, _, _ = _convert_chunk_to_message_chunk(817        reference_chunk, AIMessageChunk, index, index_type, "v1"818    )819820    assert result_1.content == [{"type": "text", "text": "The answer is ", "index": 0}]821    assert result_2.content == [822        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 1},823    ]824825826def test_citation_streaming_accumulated_content() -> None:827    """Streaming chunks accumulate normalized text blocks in full.content."""828    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}829    text_chunk = {830        "choices": [831            {832                "delta": {"role": "assistant", "content": "The answer is "},833                "finish_reason": None,834            }835        ],836    }837    reference_chunk = {838        "choices": [839            {840                "delta": {841                    "role": "assistant",842                    "content": [dict(raw_citation)],843                },844                "finish_reason": "stop",845            }846        ],847        "model": "mistral-small-latest",848    }849850    result_1, index, index_type = _convert_chunk_to_message_chunk(851        text_chunk, AIMessageChunk, -1, "", None852    )853    result_2, _, _ = _convert_chunk_to_message_chunk(854        reference_chunk, AIMessageChunk, index, index_type, None855    )856857    full = result_1 + result_2858    # full.content should contain both the text and the normalized reference block859    assert isinstance(full.content, list)860    assert any(861        isinstance(b, dict)862        and b.get("type") == "text"863        and b.get("text") == "42"864        and isinstance(ref := b.get("reference"), dict)865        and ref.get("reference_ids") == [0]866        for b in full.content867    )868869870def test_citation_index_not_in_extras() -> None:871    """Streaming index should not leak into citation extras."""872    from langchain_mistralai._compat import _convert_to_v1_from_mistral873874    content: list[str | dict] = [875        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0},876    ]877    message = AIMessageChunk(content=content)878    blocks = _convert_to_v1_from_mistral(message)879    assert len(blocks) == 1880    block_0 = cast("types.TextContentBlock", blocks[0])881    annotation = block_0["annotations"][0]882    extras = annotation.get("extras", {})883    assert isinstance(extras, dict)884    assert "index" not in extras885886887def test_citation_no_text_in_reference() -> None:888    """A reference block with no text still converts without error."""889    from langchain_mistralai._compat import _convert_to_v1_from_mistral890891    content: list[str | dict] = [892        {"type": "text", "text": "", "reference": {"reference_ids": [0]}},893    ]894    message = AIMessage(content=content)895    blocks = _convert_to_v1_from_mistral(message)896    assert len(blocks) == 1897    assert blocks[0]["type"] == "text"898    assert blocks[0]["text"] == ""899    block_0 = cast("types.TextContentBlock", blocks[0])900    assert "cited_text" not in block_0["annotations"][0]901902903def test_citation_empty_reference_metadata_still_adds_annotation() -> None:904    """Presence of reference metadata is the signal, even if the metadata is empty."""905    from langchain_mistralai._compat import _convert_to_v1_from_mistral906907    message = AIMessage(content=[{"type": "text", "text": "42", "reference": {}}])908    blocks = _convert_to_v1_from_mistral(message)909910    block_0 = cast("types.TextContentBlock", blocks[0])911    assert block_0["annotations"] == [{"type": "citation"}]912913914def test_malformed_annotation_does_not_crash() -> None:915    """Malformed annotations are skipped, not raised."""916    from langchain_mistralai._compat import _convert_from_v1_to_mistral917918    content: list = [919        {920            "type": "text",921            "text": "hello",922            "annotations": [923                None,  # not a dict924                {"type": "unknown"},  # unrecognized type925                {"type": "citation", "cited_text": "cited"},  # valid926            ],927        }928    ]929    result = _convert_from_v1_to_mistral(content, "mistralai")930    # The valid citation produces a text block with reference metadata;931    # the text block is not appended because a reference was emitted.932    assert len(result) == 1933    assert isinstance(result[0], dict)934    block_0 = result[0]935    assert block_0["type"] == "text"936    assert block_0["text"] == "hello"937    assert "reference" in block_0938939940def test_custom_token_counting() -> None:941    def token_encoder(text: str) -> list[int]:942        return [1, 2, 3]943944    llm = ChatMistralAI(custom_get_token_ids=token_encoder)945    assert llm.get_token_ids("foo") == [1, 2, 3]946947948def test_tool_id_conversion() -> None:949    assert _is_valid_mistral_tool_call_id("ssAbar4Dr")950    assert not _is_valid_mistral_tool_call_id("abc123")951    assert not _is_valid_mistral_tool_call_id("call_JIIjI55tTipFFzpcP8re3BpM")952953    result_map = {954        "ssAbar4Dr": "ssAbar4Dr",955        "abc123": "pL5rEGzxe",956        "call_JIIjI55tTipFFzpcP8re3BpM": "8kxAQvoED",957    }958    for input_id, expected_output in result_map.items():959        assert _convert_tool_call_id_to_mistral_compatible(input_id) == expected_output960        assert _is_valid_mistral_tool_call_id(expected_output)961962963def test_extra_kwargs() -> None:964    # Check that foo is saved in extra_kwargs.965    with pytest.warns(UserWarning, match="foo is not default parameter"):966        llm = ChatMistralAI(model="my-model", foo=3, max_tokens=10)  # type: ignore[call-arg]967    assert llm.max_tokens == 10968    assert llm.model_kwargs == {"foo": 3}969970    # Test that if extra_kwargs are provided, they are added to it.971    with pytest.warns(UserWarning, match="foo is not default parameter"):972        llm = ChatMistralAI(model="my-model", foo=3, model_kwargs={"bar": 2})  # type: ignore[call-arg]973    assert llm.model_kwargs == {"foo": 3, "bar": 2}974975    # Test that if provided twice it errors976    with pytest.raises(ValueError):977        ChatMistralAI(model="my-model", foo=3, model_kwargs={"foo": 2})  # type: ignore[call-arg]978979980def test_stop_stored_as_field() -> None:981    """`stop` is a first-class field, not routed into `model_kwargs`."""982    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]983    assert llm.stop == ["END"]984    assert "stop" not in llm.model_kwargs985986987def test_create_message_dicts_sends_instance_stop() -> None:988    """Instance-level `stop` is forwarded to the request params."""989    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]990    _, params = llm._create_message_dicts([HumanMessage("hi")], None)991    assert params["stop"] == ["END"]992993994def test_create_message_dicts_per_call_stop_overrides_instance() -> None:995    """A per-call `stop` (including an empty list) overrides the instance value."""996    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]997    # A non-empty per-call value wins over the instance default.998    _, params = llm._create_message_dicts([HumanMessage("hi")], ["STOP"])999    assert params["stop"] == ["STOP"]10001001    # An explicit empty list overrides the instance default and is treated as1002    # "no stop sequences", so it is omitted from the request rather than sent1003    # as an empty array (which the API would reject).1004    _, params = llm._create_message_dicts([HumanMessage("hi")], [])1005    assert "stop" not in params100610071008def test_create_message_dicts_omits_stop_when_unset() -> None:1009    """No `stop` field and no per-call value means `stop` is not sent."""1010    llm = ChatMistralAI(model="my-model")  # type: ignore[call-arg]1011    _, params = llm._create_message_dicts([HumanMessage("hi")], None)1012    assert "stop" not in params101310141015def test_get_ls_params_stop_precedence() -> None:1016    """`_get_ls_params` records instance `stop` and lets a per-call value win."""1017    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]1018    assert llm._get_ls_params().get("ls_stop") == ["END"]1019    assert llm._get_ls_params(stop=["STOP"]).get("ls_stop") == ["STOP"]10201021    # Without an instance default and no per-call value, `ls_stop` is omitted.1022    llm_no_stop = ChatMistralAI(model="my-model")  # type: ignore[call-arg]1023    assert "ls_stop" not in llm_no_stop._get_ls_params()102410251026def test_retry_with_failure_then_success() -> None:1027    """Test retry mechanism works correctly when fiest request fails, second succeed."""1028    # Create a real ChatMistralAI instance1029    chat = ChatMistralAI(max_retries=3)10301031    # Set up the actual retry mechanism (not just mocking it)1032    # We'll track how many times the function is called1033    call_count = 010341035    def mock_post(*args: Any, **kwargs: Any) -> MagicMock:1036        nonlocal call_count1037        call_count += 110381039        if call_count == 1:1040            msg = "Connection error"1041            raise httpx.RequestError(msg, request=MagicMock())10421043        mock_response = MagicMock()1044        mock_response.status_code = 2001045        mock_response.json.return_value = {1046            "choices": [1047                {1048                    "message": {1049                        "role": "assistant",1050                        "content": "Hello!",1051                    },1052                    "finish_reason": "stop",1053                }1054            ],1055            "usage": {1056                "prompt_tokens": 1,1057                "completion_tokens": 1,1058                "total_tokens": 2,1059            },1060        }1061        return mock_response10621063    with patch.object(chat.client, "post", side_effect=mock_post):1064        result = chat.invoke("Hello")1065        assert result.content == "Hello!"1066        assert call_count == 2, f"Expected 2 calls, but got {call_count}"106710681069def test_no_duplicate_tool_calls_when_multiple_tools() -> None:1070    """1071    Tests whether the conversion of an AIMessage with more than one tool call1072    to a Mistral assistant message correctly returns each tool call exactly1073    once in the final payload.10741075    The current implementation uses a faulty for loop which produces N*N entries in the1076    final tool_calls array of the payload (and thus duplicates tool call ids).1077    """1078    msg = AIMessage(1079        content="",  # content should be blank when tool_calls are present1080        tool_calls=[1081            ToolCall(name="tool_a", args={"x": 1}, id="id_a", type="tool_call"),1082            ToolCall(name="tool_b", args={"y": 2}, id="id_b", type="tool_call"),1083        ],1084        response_metadata={"model_provider": "mistralai"},1085    )10861087    mistral_msg = _convert_message_to_mistral_chat_message(msg)10881089    assert mistral_msg["role"] == "assistant"1090    assert "tool_calls" in mistral_msg, "Expected tool_calls to be present."10911092    tool_calls = mistral_msg["tool_calls"]1093    # With the bug, this would be 4 (2x2); we expect exactly 2 entries.1094    assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}"10951096    # Ensure there are no duplicate ids1097    ids = [tc.get("id") for tc in tool_calls if isinstance(tc, dict)]1098    assert len(ids) == 21099    assert len(set(ids)) == 2, f"Duplicate tool call IDs found: {ids}"110011011102def test_profile() -> None:1103    model = ChatMistralAI(model="mistral-large-latest")  # type: ignore[call-arg]1104    assert model.profile110511061107def test_metadata_versions() -> None:1108    """Test that metadata reports the correct version info."""1109    llm = ChatMistralAI(model="foo")  # type: ignore[call-arg]1110    assert llm.metadata is not None1111    versions = llm.metadata["lc_versions"]1112    assert "langchain-core" in versions1113    assert "langchain-mistralai" in versions111411151116@pytest.mark.parametrize(1117    ("status_code", "model_error_type", "is_retryable"),1118    [1119        (400, ModelInvalidRequestError, False),1120        (401, ModelAuthenticationError, False),1121        (403, ModelPermissionDeniedError, False),1122        (404, ModelNotFoundError, False),1123        (422, ModelInvalidRequestError, False),1124        (429, ModelRateLimitError, True),1125        (500, ModelAPIError, True),1126        (503, ModelAPIError, True),1127    ],1128)1129def test_error_classification(1130    status_code: int,1131    model_error_type: type[ModelError],1132    *,1133    is_retryable: bool,1134) -> None:1135    """Errors are raised as both `httpx.HTTPStatusError` and the LangChain type."""1136    response = _error_response(status_code)11371138    with pytest.raises(httpx.HTTPStatusError) as exc_info:1139        _raise_on_error(response)11401141    assert isinstance(exc_info.value, model_error_type)1142    assert exc_info.value.is_retryable is is_retryable1143    assert exc_info.value.response.status_code == status_code114411451146async def test_error_classification_async() -> None:1147    """The async response handler classifies errors the same way."""1148    with pytest.raises(httpx.HTTPStatusError) as exc_info:1149        await _araise_on_error(_error_response(429))11501151    assert isinstance(exc_info.value, ModelRateLimitError)115211531154def test_unclassified_status_stays_a_plain_status_error() -> None:1155    """Status codes outside the taxonomy keep the previous behavior."""1156    with pytest.raises(httpx.HTTPStatusError) as exc_info:1157        _raise_on_error(_error_response(409))11581159    assert not isinstance(exc_info.value, ModelError)116011611162def test_success_response_does_not_raise() -> None:1163    """A non-error response is left alone."""1164    _raise_on_error(_error_response(200))