comet-ml--opik/sdks/python/tests/library_integration/mistral/test_mistral.py

import asyncio
import json
import os

import mistralai
import pydantic
import pytest

import opik
from opik.config import OPIK_PROJECT_DEFAULT_NAME
from opik.integrations.mistral import track_mistral
from opik.types import LLMProvider
from ... import llm_constants
from ...testlib import (
    ANY_BUT_NONE,
    ANY_DICT,
    SpanModel,
    TraceModel,
    assert_equal,
)

pytestmark = pytest.mark.usefixtures("ensure_mistral_configured")


class _Person(pydantic.BaseModel):
    name: str
    age: int


PARSE_MESSAGES = [{"role": "user", "content": "Extract this person: John is 30."}]

MODEL_FOR_TESTS = llm_constants.MISTRAL_SMALL
EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT = {
    "prompt_tokens": ANY_BUT_NONE,
    "completion_tokens": ANY_BUT_NONE,
    "total_tokens": ANY_BUT_NONE,
    "original_usage.prompt_tokens": ANY_BUT_NONE,
    "original_usage.completion_tokens": ANY_BUT_NONE,
    "original_usage.total_tokens": ANY_BUT_NONE,
    "original_usage.prompt_tokens_details.cached_tokens": ANY_BUT_NONE,
}

MESSAGES = [{"role": "user", "content": "Tell a short fact"}]


@pytest.mark.parametrize(
    "project_name, expected_project_name",
    [
        (None, OPIK_PROJECT_DEFAULT_NAME),
        ("mistral-integration-test", "mistral-integration-test"),
    ],
)
def test_mistral_chat_complete__happyflow(
    fake_backend, project_name, expected_project_name
):
    client = track_mistral(
        mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]),
        project_name=project_name,
    )

    _ = client.chat.complete(model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10)

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=expected_project_name,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=expected_project_name,
                spans=[],
                model=ANY_BUT_NONE,
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_mistral_chat_complete_async__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    async def async_call():
        return await client.chat.complete_async(
            model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10
        )

    _ = asyncio.run(async_call())

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model=ANY_BUT_NONE,
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_mistral_chat_stream__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    for _ in client.chat.stream(
        model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10
    ):
        pass

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_stream",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_stream",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model=ANY_BUT_NONE,
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_mistral_chat_stream_async__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    async def async_call():
        async for _ in await client.chat.stream_async(
            model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10
        ):
            pass

    asyncio.run(async_call())

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_stream",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_stream",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model=ANY_BUT_NONE,
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_mistral_chat_complete__custom_provider__provider_logged_but_usage_still_parsed(
    fake_backend,
):
    client = track_mistral(
        mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]),
        provider="my-mistral-host",
    )

    _ = client.chat.complete(model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10)

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model=ANY_BUT_NONE,
                provider="my-mistral-host",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_mistral_provider_enum__accepted__provider_logged(fake_backend):
    client = track_mistral(
        mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]),
        provider=LLMProvider.MISTRALAI,
    )

    _ = client.chat.complete(model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10)

    opik.flush_tracker()

    assert len(fake_backend.trace_trees) == 1
    assert fake_backend.trace_trees[0].spans[0].provider == "mistral"


def test_mistral_chat_complete__error_raised__span_and_trace_finished__error_info_logged(
    fake_backend,
):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    with pytest.raises(mistralai.models.SDKError):
        client.chat.complete(model="does-not-exist-xyz", messages=MESSAGES)

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input=ANY_DICT.containing({"messages": MESSAGES}),
        output=None,
        tags=["mistral"],
        metadata=ANY_DICT,
        error_info={
            "exception_type": "SDKError",
            "message": ANY_BUT_NONE,
            "traceback": ANY_BUT_NONE,
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input=ANY_DICT.containing({"messages": MESSAGES}),
                output=None,
                tags=["mistral"],
                metadata=ANY_DICT,
                error_info={
                    "exception_type": "SDKError",
                    "message": ANY_BUT_NONE,
                    "traceback": ANY_BUT_NONE,
                },
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model="does-not-exist-xyz",
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def _expected_parse_trace(span_name: str) -> TraceModel:
    # A single llm span (no nested primitive span) proves parse doesn't
    # double-log through the complete/stream method it calls internally. The
    # span is named after the primitive (create/stream), not "parse", since
    # parse delegates to it.
    return TraceModel(
        id=ANY_BUT_NONE,
        name=span_name,
        input=ANY_DICT.containing({"messages": PARSE_MESSAGES}),
        output={"choices": ANY_BUT_NONE},
        tags=["mistral"],
        metadata=ANY_DICT,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=OPIK_PROJECT_DEFAULT_NAME,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name=span_name,
                input=ANY_DICT.containing({"messages": PARSE_MESSAGES}),
                output={"choices": ANY_BUT_NONE},
                tags=["mistral"],
                metadata=ANY_DICT,
                usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=OPIK_PROJECT_DEFAULT_NAME,
                spans=[],
                model=ANY_BUT_NONE,
                provider="mistral",
                source="sdk",
            )
        ],
        source="sdk",
    )


def test_mistral_chat_parse__happyflow__single_span(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    response = client.chat.parse(
        model=MODEL_FOR_TESTS,
        messages=PARSE_MESSAGES,
        response_format=_Person,
        max_tokens=50,
    )

    opik.flush_tracker()

    # parse() delegates to complete(); only the primitive is patched, so this
    # produces exactly one span (asserted via the single-span trace tree below),
    # named after that primitive (chat_completion_create).
    assert response.choices[0].message.parsed == _Person(name="John", age=30)
    assert len(fake_backend.trace_trees) == 1
    assert_equal(
        _expected_parse_trace("chat_completion_create"), fake_backend.trace_trees[0]
    )


def test_mistral_chat_parse_async__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    async def async_call():
        return await client.chat.parse_async(
            model=MODEL_FOR_TESTS,
            messages=PARSE_MESSAGES,
            response_format=_Person,
            max_tokens=50,
        )

    asyncio.run(async_call())

    opik.flush_tracker()

    assert len(fake_backend.trace_trees) == 1
    assert_equal(
        _expected_parse_trace("chat_completion_create"), fake_backend.trace_trees[0]
    )


def test_mistral_chat_parse_stream__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    for _ in client.chat.parse_stream(
        model=MODEL_FOR_TESTS,
        messages=PARSE_MESSAGES,
        response_format=_Person,
        max_tokens=50,
    ):
        pass

    opik.flush_tracker()

    assert len(fake_backend.trace_trees) == 1
    assert_equal(
        _expected_parse_trace("chat_completion_stream"),
        fake_backend.trace_trees[0],
    )


def test_mistral_chat_parse_stream_async__happyflow(fake_backend):
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    async def async_call():
        async for _ in await client.chat.parse_stream_async(
            model=MODEL_FOR_TESTS,
            messages=PARSE_MESSAGES,
            response_format=_Person,
            max_tokens=50,
        ):
            pass

    asyncio.run(async_call())

    opik.flush_tracker()

    assert len(fake_backend.trace_trees) == 1
    assert_equal(
        _expected_parse_trace("chat_completion_stream"),
        fake_backend.trace_trees[0],
    )


def test_mistral_chat_complete__called_in_tracked_function__span_nested_under_track(
    fake_backend,
):
    project_name = "mistral-integration-test"
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    @opik.track(project_name=project_name)
    def f():
        client.chat.complete(model=MODEL_FOR_TESTS, messages=MESSAGES, max_tokens=10)

    f()

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="f",
        input={},
        output=None,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=project_name,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="f",
                input={},
                output=None,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=project_name,
                model=None,
                provider=None,
                spans=[
                    SpanModel(
                        id=ANY_BUT_NONE,
                        type="llm",
                        name="chat_completion_create",
                        input=ANY_DICT.containing({"messages": MESSAGES}),
                        output={"choices": ANY_BUT_NONE},
                        tags=["mistral"],
                        metadata=ANY_DICT,
                        usage=EXPECTED_MISTRAL_USAGE_LOGGED_FORMAT,
                        start_time=ANY_BUT_NONE,
                        end_time=ANY_BUT_NONE,
                        project_name=project_name,
                        spans=[],
                        model=ANY_BUT_NONE,
                        provider="mistral",
                        source="sdk",
                    )
                ],
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])


def test_track_mistral__unsupported_old_version__raises(monkeypatch):
    from opik.integrations.mistral import opik_tracker

    monkeypatch.setattr(
        opik_tracker.importlib.metadata, "version", lambda _pkg: "1.2.0"
    )

    with pytest.raises(RuntimeError, match=r"mistralai>=1\.3\.0"):
        track_mistral(mistralai.Mistral(api_key="dummy-key"))


_WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the weather for a city",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    },
}


def test_mistral_chat_stream__tool_calls__aggregated_into_span_output(fake_backend):
    # Guards against losing streamed tool calls: the chunk aggregator must keep
    # the tool call(s) (with complete function name + arguments) rather than
    # overwriting them per chunk.
    client = track_mistral(mistralai.Mistral(api_key=os.environ["MISTRAL_API_KEY"]))

    for _ in client.chat.stream(
        model=MODEL_FOR_TESTS,
        messages=[{"role": "user", "content": "What is the weather in Paris?"}],
        tools=[_WEATHER_TOOL],
        tool_choice="any",
        max_tokens=100,
    ):
        pass

    opik.flush_tracker()

    assert len(fake_backend.trace_trees) == 1
    tool_calls = (
        fake_backend.trace_trees[0]
        .spans[0]
        .output["choices"][0]["message"]["tool_calls"]
    )
    assert tool_calls, "streamed tool call was lost during aggregation"
    first_call = tool_calls[0]["function"]
    assert first_call["name"] == "get_weather"
    # arguments must be complete/valid JSON, not a truncated fragment
    assert "city" in json.loads(first_call["arguments"])