"""
Test Mistral instrumentation against local Latitude instance.

Required env vars:
- LATITUDE_API_KEY
- MISTRAL_API_KEY

Install: uv add mistralai
"""

import os

import mistralai
from mistralai import Mistral

from latitude_telemetry import Latitude, capture

# Initialize telemetry pointing to local instance
latitude = Latitude(
    api_key=os.environ["LATITUDE_API_KEY"],
    project=os.environ["LATITUDE_PROJECT_SLUG"],
    instrumentations={"mistralai": mistralai},
    disable_batch=True,
)


@capture("test-mistral-completion", {"tags": ["python", "test"], "session_id": "example"})
def test_mistral_completion():
    from mistralai.models import UserMessage

    client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])

    response = client.chat.complete(
        model="mistral-small-latest",
        messages=[UserMessage(role="user", content="Say 'Hello from Mistral!' in exactly 5 words.")],
        max_tokens=50,
    )

    return response.choices[0].message.content


if __name__ == "__main__":
    test_mistral_completion()
    latitude.flush()
