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.messages import (12    AIMessage,13    AIMessageChunk,14    BaseMessage,15    ChatMessage,16    HumanMessage,17    InvalidToolCall,18    SystemMessage,19    ToolCall,20    ToolMessage,21)22from pydantic import SecretStr2324if TYPE_CHECKING:25    from langchain_core.messages import content as types2627from langchain_mistralai._compat import _convert_to_v1_from_mistral28from langchain_mistralai.chat_models import (  # type: ignore[import]29    ChatMistralAI,30    _convert_chunk_to_message_chunk,31    _convert_message_to_mistral_chat_message,32    _convert_mistral_chat_message_to_message,33    _convert_tool_call_id_to_mistral_compatible,34    _format_message_content,35    _is_valid_mistral_tool_call_id,36    _sanitize_chat_completions_content,37)3839os.environ["MISTRAL_API_KEY"] = "foo"404142def test_sanitize_chat_completions_text_blocks_strips_id() -> None:43    """LangChain auto-generated `id` on text blocks must not reach the wire.4445    Mistral's chat completions endpoint returns 422 with `extra_forbidden`46    on `messages[*].tool.content.list[...].text.id` if not stripped.47    """48    message = ToolMessage(49        content=[{"type": "text", "text": "foo", "id": "lc_abc123"}],50        tool_call_id="abc12345",51    )52    result = _convert_message_to_mistral_chat_message(message)53    assert result["content"] == [{"type": "text", "text": "foo"}]545556def test_sanitize_chat_completions_content_passthrough_string() -> None:57    assert _sanitize_chat_completions_content("hello") == "hello"585960def test_ai_message_reference_metadata_does_not_reach_wire() -> None:61    message = AIMessage(62        content=[63            {"type": "text", "text": "The answer is "},64            {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},65            {"type": "text", "text": "."},66        ],67        response_metadata={"model_provider": "mistralai"},68    )6970    result = _convert_message_to_mistral_chat_message(message)71    assert result["content"] == [72        {"type": "text", "text": "The answer is "},73        {"type": "text", "text": "42"},74        {"type": "text", "text": "."},75    ]767778def test_v1_ai_message_reference_metadata_does_not_reach_wire() -> None:79    message = AIMessage(80        content=[81            {"type": "text", "text": "The answer is "},82            {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},83            {"type": "text", "text": "."},84        ],85        response_metadata={"model_provider": "mistralai", "output_version": "v1"},86    )8788    result = _convert_message_to_mistral_chat_message(message)89    assert result["content"] == [90        {"type": "text", "text": "The answer is "},91        {"type": "text", "text": "42"},92        {"type": "text", "text": "."},93    ]949596def test_mistralai_model_param() -> None:97    llm = ChatMistralAI(model="foo")  # type: ignore[call-arg]98    assert llm.model == "foo"99100101def test_mistralai_initialization() -> None:102    """Test ChatMistralAI initialization."""103    # Verify that ChatMistralAI can be initialized using a secret key provided104    # as a parameter rather than an environment variable.105    for model in [106        ChatMistralAI(model="test", mistral_api_key="test"),  # type: ignore[call-arg, call-arg]107        ChatMistralAI(model="test", api_key="test"),  # type: ignore[call-arg, arg-type]108    ]:109        assert cast("SecretStr", model.mistral_api_key).get_secret_value() == "test"110111112@pytest.mark.parametrize(113    ("model", "expected_url"),114    [115        (ChatMistralAI(model="test"), "https://api.mistral.ai/v1"),  # type: ignore[call-arg, arg-type]116        (ChatMistralAI(model="test", endpoint="baz"), "baz"),  # type: ignore[call-arg, arg-type]117    ],118)119def test_mistralai_initialization_baseurl(120    model: ChatMistralAI, expected_url: str121) -> None:122    """Test ChatMistralAI initialization."""123    # Verify that ChatMistralAI can be initialized providing endpoint, but also124    # with default125126    assert model.endpoint == expected_url127128129@pytest.mark.parametrize(130    "env_var_name",131    [132        ("MISTRAL_BASE_URL"),133    ],134)135def test_mistralai_initialization_baseurl_env(136    env_var_name: str, monkeypatch: pytest.MonkeyPatch137) -> None:138    """Test ChatMistralAI initialization."""139    # Verify that ChatMistralAI can be initialized using env variable140    monkeypatch.setenv(env_var_name, "boo")141    model = ChatMistralAI(model="test")  # type: ignore[call-arg]142    assert model.endpoint == "boo"143144145@pytest.mark.parametrize(146    ("message", "expected"),147    [148        (149            SystemMessage(content="Hello"),150            {"role": "system", "content": "Hello"},151        ),152        (153            HumanMessage(content="Hello"),154            {"role": "user", "content": "Hello"},155        ),156        (157            AIMessage(content="Hello"),158            {"role": "assistant", "content": "Hello"},159        ),160        (161            AIMessage(content="{", additional_kwargs={"prefix": True}),162            {"role": "assistant", "content": "{", "prefix": True},163        ),164        (165            ChatMessage(role="assistant", content="Hello"),166            {"role": "assistant", "content": "Hello"},167        ),168    ],169)170def test_convert_message_to_mistral_chat_message(171    message: BaseMessage, expected: dict172) -> None:173    result = _convert_message_to_mistral_chat_message(message)174    assert result == expected175176177@pytest.mark.parametrize(178    ("content", "expected"),179    [180        ("hello", "hello"),181        ("", ""),182        (None, None),183        ([], []),184    ],185)186def test_format_message_content_passthrough_non_list(187    content: Any, expected: Any188) -> None:189    """Strings, None, and empty lists pass through `_format_message_content`."""190    assert _format_message_content(content) == expected191192193@pytest.mark.parametrize(194    ("block", "expected"),195    [196        (197            {"type": "image", "url": "https://example.com/img.png"},198            {199                "type": "image_url",200                "image_url": {"url": "https://example.com/img.png"},201            },202        ),203        (204            {"type": "image", "base64": "abc123", "mime_type": "image/jpeg"},205            {206                "type": "image_url",207                "image_url": {"url": "data:image/jpeg;base64,abc123"},208            },209        ),210        (211            {212                "type": "image",213                "source_type": "url",214                "url": "https://example.com/v0.png",215            },216            {217                "type": "image_url",218                "image_url": {"url": "https://example.com/v0.png"},219            },220        ),221        (222            {223                "type": "image",224                "source_type": "base64",225                "data": "v0data",226                "mime_type": "image/png",227            },228            {229                "type": "image_url",230                "image_url": {"url": "data:image/png;base64,v0data"},231            },232        ),233    ],234)235def test_format_message_content_translates_image_blocks(236    block: dict, expected: dict237) -> None:238    """v0 and v1 canonical image blocks translate to Mistral's `image_url` shape."""239    assert _format_message_content([block]) == [expected]240241242@pytest.mark.parametrize(243    "block",244    [245        {"type": "text", "text": "hello"},246        {"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},247        {"type": "image_url", "image_url": "https://example.com/img.png"},248    ],249)250def test_format_message_content_passthrough_known_blocks(block: dict) -> None:251    """Already-translated wire blocks and text blocks pass through unchanged."""252    assert _format_message_content([block]) == [block]253254255@pytest.mark.parametrize(256    "block_type",257    ["tool_use", "thinking", "reasoning_content", "document_url", "input_audio"],258)259def test_format_message_content_passes_unknown_blocks_through(block_type: str) -> None:260    """Non-canonical blocks pass through; the Mistral API validates them."""261    blocks = [262        {"type": "text", "text": "kept"},263        {"type": block_type, "data": "anything"},264    ]265    assert _format_message_content(blocks) == blocks266267268def test_format_message_content_preserves_order_for_mixed_blocks() -> None:269    """Multiple text + image blocks retain their order — vision prompts depend on it."""270    blocks: list[Any] = [271        {"type": "text", "text": "first"},272        {"type": "image", "url": "https://example.com/a.png"},273        {"type": "text", "text": "between"},274        {"type": "image", "base64": "xyz", "mime_type": "image/png"},275        "trailing string",276    ]277    expected = [278        {"type": "text", "text": "first"},279        {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}},280        {"type": "text", "text": "between"},281        {"type": "image_url", "image_url": {"url": "data:image/png;base64,xyz"}},282        "trailing string",283    ]284    assert _format_message_content(blocks) == expected285286287def test_format_message_content_image_missing_mime_type_raises() -> None:288    """Base64 image without `mime_type` raises via the core translator."""289    with pytest.raises(ValueError, match="mime_type"):290        _format_message_content([{"type": "image", "base64": "abc"}])291292293@pytest.mark.parametrize(294    ("message", "expected"),295    [296        (297            HumanMessage(298                content=[299                    {"type": "text", "text": "What is in this image?"},300                    {"type": "image", "url": "https://example.com/img.png"},301                ]302            ),303            {304                "role": "user",305                "content": [306                    {"type": "text", "text": "What is in this image?"},307                    {308                        "type": "image_url",309                        "image_url": {"url": "https://example.com/img.png"},310                    },311                ],312            },313        ),314        (315            HumanMessage(316                content=[317                    {"type": "text", "text": "Describe this image."},318                    {319                        "type": "image",320                        "base64": "abc123",321                        "mime_type": "image/png",322                    },323                ]324            ),325            {326                "role": "user",327                "content": [328                    {"type": "text", "text": "Describe this image."},329                    {330                        "type": "image_url",331                        "image_url": {"url": "data:image/png;base64,abc123"},332                    },333                ],334            },335        ),336    ],337)338def test_convert_human_message_with_images(339    message: BaseMessage, expected: dict340) -> None:341    result = _convert_message_to_mistral_chat_message(message)342    assert result == expected343344345def test_convert_human_message_with_string_content_unchanged() -> None:346    """Plain string `HumanMessage` content is not wrapped or modified."""347    result = _convert_message_to_mistral_chat_message(HumanMessage(content="hi"))348    assert result == {"role": "user", "content": "hi"}349350351def _make_completion_response_from_token(token: str) -> dict:352    return {353        "id": "abc123",354        "model": "fake_model",355        "choices": [356            {357                "index": 0,358                "delta": {"content": token},359                "finish_reason": None,360            }361        ],362    }363364365def mock_chat_stream(*args: Any, **kwargs: Any) -> Generator:366    def it() -> Generator:367        for token in ["Hello", " how", " can", " I", " help", "?"]:368            yield _make_completion_response_from_token(token)369370    return it()371372373async def mock_chat_astream(*args: Any, **kwargs: Any) -> AsyncGenerator:374    async def it() -> AsyncGenerator:375        for token in ["Hello", " how", " can", " I", " help", "?"]:376            yield _make_completion_response_from_token(token)377378    return it()379380381class MyCustomHandler(BaseCallbackHandler):382    last_token: str = ""383384    def on_llm_new_token(385        self, token: str | list[str | dict[str, Any]], **kwargs: Any386    ) -> None:387        if isinstance(token, str):388            self.last_token = token389390391@patch(392    "langchain_mistralai.chat_models.ChatMistralAI.completion_with_retry",393    new=mock_chat_stream,394)395def test_stream_with_callback() -> None:396    callback = MyCustomHandler()397    chat = ChatMistralAI(callbacks=[callback])398    for token in chat.stream("Hello"):399        assert callback.last_token == token.content400401402@patch("langchain_mistralai.chat_models.acompletion_with_retry", new=mock_chat_astream)403async def test_astream_with_callback() -> None:404    callback = MyCustomHandler()405    chat = ChatMistralAI(callbacks=[callback])406    async for token in chat.astream("Hello"):407        assert callback.last_token == token.content408409410def test__convert_dict_to_message_tool_call() -> None:411    raw_tool_call = {412        "id": "ssAbar4Dr",413        "function": {414            "arguments": '{"name": "Sally", "hair_color": "green"}',415            "name": "GenerateUsername",416        },417    }418    message = {"role": "assistant", "content": "", "tool_calls": [raw_tool_call]}419    result = _convert_mistral_chat_message_to_message(message)420    expected_output = AIMessage(421        content="",422        additional_kwargs={"tool_calls": [raw_tool_call]},423        tool_calls=[424            ToolCall(425                name="GenerateUsername",426                args={"name": "Sally", "hair_color": "green"},427                id="ssAbar4Dr",428                type="tool_call",429            )430        ],431        response_metadata={"model_provider": "mistralai"},432    )433    assert result == expected_output434    assert _convert_message_to_mistral_chat_message(expected_output) == message435436    # Test malformed tool call437    raw_tool_calls = [438        {439            "id": "pL5rEGzxe",440            "function": {441                "arguments": '{"name": "Sally", "hair_color": "green"}',442                "name": "GenerateUsername",443            },444        },445        {446            "id": "ssAbar4Dr",447            "function": {448                "arguments": "oops",449                "name": "GenerateUsername",450            },451        },452    ]453    message = {"role": "assistant", "content": "", "tool_calls": raw_tool_calls}454    result = _convert_mistral_chat_message_to_message(message)455    expected_output = AIMessage(456        content="",457        additional_kwargs={"tool_calls": raw_tool_calls},458        invalid_tool_calls=[459            InvalidToolCall(460                name="GenerateUsername",461                args="oops",462                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: E501463                id="ssAbar4Dr",464                type="invalid_tool_call",465            ),466        ],467        tool_calls=[468            ToolCall(469                name="GenerateUsername",470                args={"name": "Sally", "hair_color": "green"},471                id="pL5rEGzxe",472                type="tool_call",473            ),474        ],475        response_metadata={"model_provider": "mistralai"},476    )477    assert result == expected_output478    assert _convert_message_to_mistral_chat_message(expected_output) == message479480481def test__convert_dict_to_message_tool_call_with_null_content() -> None:482    raw_tool_call = {483        "id": "ssAbar4Dr",484        "function": {485            "arguments": '{"name": "Sally", "hair_color": "green"}',486            "name": "GenerateUsername",487        },488    }489    message = {"role": "assistant", "content": None, "tool_calls": [raw_tool_call]}490    result = _convert_mistral_chat_message_to_message(message)491    expected_output = AIMessage(492        content="",493        additional_kwargs={"tool_calls": [raw_tool_call]},494        tool_calls=[495            ToolCall(496                name="GenerateUsername",497                args={"name": "Sally", "hair_color": "green"},498                id="ssAbar4Dr",499                type="tool_call",500            )501        ],502        response_metadata={"model_provider": "mistralai"},503    )504    assert result == expected_output505506507def test__convert_dict_to_message_with_missing_content() -> None:508    raw_tool_call = {509        "id": "ssAbar4Dr",510        "function": {511            "arguments": '{"query": "test search"}',512            "name": "search",513        },514    }515    message = {"role": "assistant", "tool_calls": [raw_tool_call]}516    result = _convert_mistral_chat_message_to_message(message)517    expected_output = AIMessage(518        content="",519        additional_kwargs={"tool_calls": [raw_tool_call]},520        tool_calls=[521            ToolCall(522                name="search",523                args={"query": "test search"},524                id="ssAbar4Dr",525                type="tool_call",526            )527        ],528        response_metadata={"model_provider": "mistralai"},529    )530    assert result == expected_output531532533def test__convert_dict_to_message_with_citations() -> None:534    """Reference blocks normalized to text blocks with reference metadata."""535    cited_text = "the temperature is 20 degrees C"536    raw_content: list[str | dict] = [537        {"type": "text", "text": "According to the document, "},538        {"type": "reference", "reference_ids": [0], "text": cited_text},539        {"type": "text", "text": " on average."},540    ]541    message = {"role": "assistant", "content": raw_content}542    result = _convert_mistral_chat_message_to_message(message)543544    assert isinstance(result.content, list)545    content = result.content546    # The reference block is normalized to type="text" so .text includes it547    assert content[0] == {"type": "text", "text": "According to the document, "}548    assert isinstance(content[1], dict)549    block_1 = content[1]550    assert block_1["type"] == "text"551    assert block_1["text"] == cited_text552    assert block_1["reference"] == {"reference_ids": [0]}553    assert content[2] == {"type": "text", "text": " on average."}554    assert result.response_metadata["model_provider"] == "mistralai"555    assert "citations" not in result.response_metadata556557558def test__convert_dict_to_message_citations_text_accessor() -> None:559    """message.text includes cited spans from normalized reference blocks."""560    cited_text = "the temperature is 20 degrees C"561    raw_content: list[str | dict] = [562        {"type": "text", "text": "According to the document, "},563        {"type": "reference", "reference_ids": [0], "text": cited_text},564        {"type": "text", "text": " on average."},565    ]566    message = {"role": "assistant", "content": raw_content}567    result = _convert_mistral_chat_message_to_message(message)568569    # .text should include all visible text, including the cited span570    assert str(result.text) == (571        "According to the document, the temperature is 20 degrees C on average."572    )573574575def test__convert_dict_to_message_citations_to_content_blocks() -> None:576    """content_blocks translates reference metadata to TextContentBlock."""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    assert isinstance(result, AIMessage)587    blocks = _convert_to_v1_from_mistral(result)588    assert len(blocks) == 3589590    # First block: plain text591    assert blocks[0]["type"] == "text"592    assert blocks[0]["text"] == "According to the document, "593594    # Second block: text with citation annotation595    block_1 = cast("types.TextContentBlock", blocks[1])596    assert block_1["type"] == "text"597    assert block_1["text"] == cited_text598    annotations = block_1["annotations"]599    assert len(annotations) == 1600    assert annotations[0]["type"] == "citation"601    assert "cited_text" not in annotations[0]602    assert annotations[0]["extras"]["reference_ids"] == [0]603604    # Third block: plain text605    assert blocks[2]["type"] == "text"606    assert blocks[2]["text"] == " on average."607608609def test_create_chat_result_with_citations() -> None:610    """Citations are normalized to text blocks with reference metadata in .content."""611    chat = ChatMistralAI()612    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}613    raw_content: list[str | dict] = [614        {"type": "text", "text": "The answer is "},615        raw_citation,616        {"type": "text", "text": "."},617    ]618    response = {619        "choices": [620            {621                "message": {622                    "role": "assistant",623                    "content": raw_content,624                },625                "finish_reason": "stop",626            }627        ]628    }629630    result = chat._create_chat_result(response)631    message = result.generations[0].message632633    assert isinstance(message.content, list)634    content = message.content635    # The reference block is normalized; .text includes the cited span636    assert isinstance(content[1], dict)637    block_1 = content[1]638    assert block_1["type"] == "text"639    assert block_1["text"] == "42"640    assert block_1["reference"] == {"reference_ids": [0]}641    assert str(message.text) == "The answer is 42."642    assert "citations" not in message.response_metadata643644645def test__convert_chunk_to_message_chunk_with_citations() -> None:646    """Streaming reference blocks are normalized to text blocks in chunk .content."""647    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}648    text_chunk = {649        "choices": [650            {651                "delta": {"role": "assistant", "content": "The answer is "},652                "finish_reason": None,653            }654        ],655    }656    reference_chunk = {657        "choices": [658            {659                "delta": {660                    "role": "assistant",661                    "content": [662                        dict(raw_citation),663                    ],664                },665                "finish_reason": "stop",666            }667        ],668        "model": "mistral-small-latest",669    }670671    result_1, index, index_type = _convert_chunk_to_message_chunk(672        text_chunk, AIMessageChunk, -1, "", None673    )674    result_2, _, _ = _convert_chunk_to_message_chunk(675        reference_chunk, AIMessageChunk, index, index_type, None676    )677678    assert isinstance(result_2, AIMessageChunk)679    # Reference block is normalized to type="text" with reference metadata680    assert result_2.content == [681        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0},682    ]683    assert "citations" not in result_2.response_metadata684685    full = result_1 + result_2686    assert isinstance(full, AIMessageChunk)687    assert "citations" not in full.response_metadata688    assert full.response_metadata["finish_reason"] == "stop"689    # .text includes the cited span690    assert str(full.text) == "The answer is 42"691692693def test_citation_round_trip() -> None:694    """Round-trip through v1 preserves text and reference metadata."""695    from langchain_mistralai._compat import (696        _convert_from_v1_to_mistral,697        _convert_to_v1_from_mistral,698    )699700    # Start with normalized content (as produced by _convert_mistral_chat_message)701    original_content: list[str | dict] = [702        {"type": "text", "text": "The answer is "},703        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}},704        {"type": "text", "text": "."},705    ]706    message = AIMessage(content=original_content)707    v1_blocks = _convert_to_v1_from_mistral(message)708    round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai")709710    # Should have exactly 3 blocks, no duplication of cited text711    assert len(round_tripped) == 3712    assert round_tripped[0] == {"type": "text", "text": "The answer is "}713    assert isinstance(round_tripped[1], dict)714    block_1 = round_tripped[1]715    assert block_1["type"] == "text"716    assert block_1["text"] == "42"717    assert block_1["reference"] == {"reference_ids": [0]}718    assert round_tripped[2] == {"type": "text", "text": "."}719720721def test_citation_round_trip_preserves_extra_fields() -> None:722    """Extra provider fields on reference metadata survive the round-trip."""723    from langchain_mistralai._compat import (724        _convert_from_v1_to_mistral,725        _convert_to_v1_from_mistral,726    )727728    original_content: list[str | dict] = [729        {"type": "text", "text": "cited span", "reference": {"reference_ids": [1, 2]}},730    ]731    message = AIMessage(content=original_content)732    v1_blocks = _convert_to_v1_from_mistral(message)733    round_tripped = _convert_from_v1_to_mistral(v1_blocks, "mistralai")734735    assert len(round_tripped) == 1736    assert isinstance(round_tripped[0], dict)737    block_0 = round_tripped[0]738    assert block_0["type"] == "text"739    assert block_0["text"] == "cited span"740    assert block_0["reference"] == {"reference_ids": [1, 2]}741742743def test_citation_round_trip_preserves_annotated_response_text() -> None:744    """Serializing citations preserves block text, not citation source excerpts."""745    from langchain_mistralai._compat import _convert_from_v1_to_mistral746747    content: list[types.ContentBlock] = [748        {749            "type": "text",750            "text": "The answer is 42.",751            "annotations": [752                {753                    "type": "citation",754                    "cited_text": "source excerpt mentioning 42",755                    "extras": {"reference_ids": [0]},756                }757            ],758        }759    ]760    round_tripped = _convert_from_v1_to_mistral(content, "mistralai")761762    assert len(round_tripped) == 1763    assert isinstance(round_tripped[0], dict)764    block = round_tripped[0]765    assert block["type"] == "text"766    assert block["text"] == "The answer is 42."767    assert block["reference"]["reference_ids"] == [0]768    assert block["reference"]["cited_text"] == "source excerpt mentioning 42"769770771def test_citation_streaming_v1_reference_gets_separate_index() -> None:772    """Reference chunks do not merge into surrounding v1 text block indexes."""773    text_chunk = {774        "choices": [775            {776                "delta": {"role": "assistant", "content": "The answer is "},777                "finish_reason": None,778            }779        ],780    }781    reference_chunk = {782        "choices": [783            {784                "delta": {785                    "role": "assistant",786                    "content": [787                        {"type": "reference", "reference_ids": [0], "text": "42"},788                    ],789                },790                "finish_reason": "stop",791            }792        ],793        "model": "mistral-small-latest",794    }795796    result_1, index, index_type = _convert_chunk_to_message_chunk(797        text_chunk, AIMessageChunk, -1, "", "v1"798    )799    result_2, _, _ = _convert_chunk_to_message_chunk(800        reference_chunk, AIMessageChunk, index, index_type, "v1"801    )802803    assert result_1.content == [{"type": "text", "text": "The answer is ", "index": 0}]804    assert result_2.content == [805        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 1},806    ]807808809def test_citation_streaming_accumulated_content() -> None:810    """Streaming chunks accumulate normalized text blocks in full.content."""811    raw_citation = {"type": "reference", "reference_ids": [0], "text": "42"}812    text_chunk = {813        "choices": [814            {815                "delta": {"role": "assistant", "content": "The answer is "},816                "finish_reason": None,817            }818        ],819    }820    reference_chunk = {821        "choices": [822            {823                "delta": {824                    "role": "assistant",825                    "content": [dict(raw_citation)],826                },827                "finish_reason": "stop",828            }829        ],830        "model": "mistral-small-latest",831    }832833    result_1, index, index_type = _convert_chunk_to_message_chunk(834        text_chunk, AIMessageChunk, -1, "", None835    )836    result_2, _, _ = _convert_chunk_to_message_chunk(837        reference_chunk, AIMessageChunk, index, index_type, None838    )839840    full = result_1 + result_2841    # full.content should contain both the text and the normalized reference block842    assert isinstance(full.content, list)843    assert any(844        isinstance(b, dict)845        and b.get("type") == "text"846        and b.get("text") == "42"847        and isinstance(ref := b.get("reference"), dict)848        and ref.get("reference_ids") == [0]849        for b in full.content850    )851852853def test_citation_index_not_in_extras() -> None:854    """Streaming index should not leak into citation extras."""855    from langchain_mistralai._compat import _convert_to_v1_from_mistral856857    content: list[str | dict] = [858        {"type": "text", "text": "42", "reference": {"reference_ids": [0]}, "index": 0},859    ]860    message = AIMessageChunk(content=content)861    blocks = _convert_to_v1_from_mistral(message)862    assert len(blocks) == 1863    block_0 = cast("types.TextContentBlock", blocks[0])864    annotation = block_0["annotations"][0]865    extras = annotation.get("extras", {})866    assert isinstance(extras, dict)867    assert "index" not in extras868869870def test_citation_no_text_in_reference() -> None:871    """A reference block with no text still converts without error."""872    from langchain_mistralai._compat import _convert_to_v1_from_mistral873874    content: list[str | dict] = [875        {"type": "text", "text": "", "reference": {"reference_ids": [0]}},876    ]877    message = AIMessage(content=content)878    blocks = _convert_to_v1_from_mistral(message)879    assert len(blocks) == 1880    assert blocks[0]["type"] == "text"881    assert blocks[0]["text"] == ""882    block_0 = cast("types.TextContentBlock", blocks[0])883    assert "cited_text" not in block_0["annotations"][0]884885886def test_citation_empty_reference_metadata_still_adds_annotation() -> None:887    """Presence of reference metadata is the signal, even if the metadata is empty."""888    from langchain_mistralai._compat import _convert_to_v1_from_mistral889890    message = AIMessage(content=[{"type": "text", "text": "42", "reference": {}}])891    blocks = _convert_to_v1_from_mistral(message)892893    block_0 = cast("types.TextContentBlock", blocks[0])894    assert block_0["annotations"] == [{"type": "citation"}]895896897def test_malformed_annotation_does_not_crash() -> None:898    """Malformed annotations are skipped, not raised."""899    from langchain_mistralai._compat import _convert_from_v1_to_mistral900901    content: list = [902        {903            "type": "text",904            "text": "hello",905            "annotations": [906                None,  # not a dict907                {"type": "unknown"},  # unrecognized type908                {"type": "citation", "cited_text": "cited"},  # valid909            ],910        }911    ]912    result = _convert_from_v1_to_mistral(content, "mistralai")913    # The valid citation produces a text block with reference metadata;914    # the text block is not appended because a reference was emitted.915    assert len(result) == 1916    assert isinstance(result[0], dict)917    block_0 = result[0]918    assert block_0["type"] == "text"919    assert block_0["text"] == "hello"920    assert "reference" in block_0921922923def test_custom_token_counting() -> None:924    def token_encoder(text: str) -> list[int]:925        return [1, 2, 3]926927    llm = ChatMistralAI(custom_get_token_ids=token_encoder)928    assert llm.get_token_ids("foo") == [1, 2, 3]929930931def test_tool_id_conversion() -> None:932    assert _is_valid_mistral_tool_call_id("ssAbar4Dr")933    assert not _is_valid_mistral_tool_call_id("abc123")934    assert not _is_valid_mistral_tool_call_id("call_JIIjI55tTipFFzpcP8re3BpM")935936    result_map = {937        "ssAbar4Dr": "ssAbar4Dr",938        "abc123": "pL5rEGzxe",939        "call_JIIjI55tTipFFzpcP8re3BpM": "8kxAQvoED",940    }941    for input_id, expected_output in result_map.items():942        assert _convert_tool_call_id_to_mistral_compatible(input_id) == expected_output943        assert _is_valid_mistral_tool_call_id(expected_output)944945946def test_extra_kwargs() -> None:947    # Check that foo is saved in extra_kwargs.948    with pytest.warns(UserWarning, match="foo is not default parameter"):949        llm = ChatMistralAI(model="my-model", foo=3, max_tokens=10)  # type: ignore[call-arg]950    assert llm.max_tokens == 10951    assert llm.model_kwargs == {"foo": 3}952953    # Test that if extra_kwargs are provided, they are added to it.954    with pytest.warns(UserWarning, match="foo is not default parameter"):955        llm = ChatMistralAI(model="my-model", foo=3, model_kwargs={"bar": 2})  # type: ignore[call-arg]956    assert llm.model_kwargs == {"foo": 3, "bar": 2}957958    # Test that if provided twice it errors959    with pytest.raises(ValueError):960        ChatMistralAI(model="my-model", foo=3, model_kwargs={"foo": 2})  # type: ignore[call-arg]961962963def test_stop_stored_as_field() -> None:964    """`stop` is a first-class field, not routed into `model_kwargs`."""965    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]966    assert llm.stop == ["END"]967    assert "stop" not in llm.model_kwargs968969970def test_create_message_dicts_sends_instance_stop() -> None:971    """Instance-level `stop` is forwarded to the request params."""972    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]973    _, params = llm._create_message_dicts([HumanMessage("hi")], None)974    assert params["stop"] == ["END"]975976977def test_create_message_dicts_per_call_stop_overrides_instance() -> None:978    """A per-call `stop` (including an empty list) overrides the instance value."""979    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]980    # A non-empty per-call value wins over the instance default.981    _, params = llm._create_message_dicts([HumanMessage("hi")], ["STOP"])982    assert params["stop"] == ["STOP"]983984    # An explicit empty list overrides the instance default and is treated as985    # "no stop sequences", so it is omitted from the request rather than sent986    # as an empty array (which the API would reject).987    _, params = llm._create_message_dicts([HumanMessage("hi")], [])988    assert "stop" not in params989990991def test_create_message_dicts_omits_stop_when_unset() -> None:992    """No `stop` field and no per-call value means `stop` is not sent."""993    llm = ChatMistralAI(model="my-model")  # type: ignore[call-arg]994    _, params = llm._create_message_dicts([HumanMessage("hi")], None)995    assert "stop" not in params996997998def test_get_ls_params_stop_precedence() -> None:999    """`_get_ls_params` records instance `stop` and lets a per-call value win."""1000    llm = ChatMistralAI(model="my-model", stop=["END"])  # type: ignore[call-arg]1001    assert llm._get_ls_params().get("ls_stop") == ["END"]1002    assert llm._get_ls_params(stop=["STOP"]).get("ls_stop") == ["STOP"]10031004    # Without an instance default and no per-call value, `ls_stop` is omitted.1005    llm_no_stop = ChatMistralAI(model="my-model")  # type: ignore[call-arg]1006    assert "ls_stop" not in llm_no_stop._get_ls_params()100710081009def test_retry_with_failure_then_success() -> None:1010    """Test retry mechanism works correctly when fiest request fails, second succeed."""1011    # Create a real ChatMistralAI instance1012    chat = ChatMistralAI(max_retries=3)10131014    # Set up the actual retry mechanism (not just mocking it)1015    # We'll track how many times the function is called1016    call_count = 010171018    def mock_post(*args: Any, **kwargs: Any) -> MagicMock:1019        nonlocal call_count1020        call_count += 110211022        if call_count == 1:1023            msg = "Connection error"1024            raise httpx.RequestError(msg, request=MagicMock())10251026        mock_response = MagicMock()1027        mock_response.status_code = 2001028        mock_response.json.return_value = {1029            "choices": [1030                {1031                    "message": {1032                        "role": "assistant",1033                        "content": "Hello!",1034                    },1035                    "finish_reason": "stop",1036                }1037            ],1038            "usage": {1039                "prompt_tokens": 1,1040                "completion_tokens": 1,1041                "total_tokens": 2,1042            },1043        }1044        return mock_response10451046    with patch.object(chat.client, "post", side_effect=mock_post):1047        result = chat.invoke("Hello")1048        assert result.content == "Hello!"1049        assert call_count == 2, f"Expected 2 calls, but got {call_count}"105010511052def test_no_duplicate_tool_calls_when_multiple_tools() -> None:1053    """1054    Tests whether the conversion of an AIMessage with more than one tool call1055    to a Mistral assistant message correctly returns each tool call exactly1056    once in the final payload.10571058    The current implementation uses a faulty for loop which produces N*N entries in the1059    final tool_calls array of the payload (and thus duplicates tool call ids).1060    """1061    msg = AIMessage(1062        content="",  # content should be blank when tool_calls are present1063        tool_calls=[1064            ToolCall(name="tool_a", args={"x": 1}, id="id_a", type="tool_call"),1065            ToolCall(name="tool_b", args={"y": 2}, id="id_b", type="tool_call"),1066        ],1067        response_metadata={"model_provider": "mistralai"},1068    )10691070    mistral_msg = _convert_message_to_mistral_chat_message(msg)10711072    assert mistral_msg["role"] == "assistant"1073    assert "tool_calls" in mistral_msg, "Expected tool_calls to be present."10741075    tool_calls = mistral_msg["tool_calls"]1076    # With the bug, this would be 4 (2x2); we expect exactly 2 entries.1077    assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}"10781079    # Ensure there are no duplicate ids1080    ids = [tc.get("id") for tc in tool_calls if isinstance(tc, dict)]1081    assert len(ids) == 21082    assert len(set(ids)) == 2, f"Duplicate tool call IDs found: {ids}"108310841085def test_profile() -> None:1086    model = ChatMistralAI(model="mistral-large-latest")  # type: ignore[call-arg]1087    assert model.profile108810891090def test_metadata_versions() -> None:1091    """Test that metadata reports the correct version info."""1092    llm = ChatMistralAI(model="foo")  # type: ignore[call-arg]1093    assert llm.metadata is not None1094    versions = llm.metadata["lc_versions"]1095    assert "langchain-core" in versions1096    assert "langchain-mistralai" in versions