comet-ml--opik/sdks/python/tests/library_integration/aisuite/test_aisuite.py

from typing import Any, Dict

import aisuite
import pytest

import opik
from opik.integrations.aisuite import track_aisuite
from ... import llm_constants
from ...testlib import (
    ANY_BUT_NONE,
    ANY_DICT,
    ANY_STRING,
    SpanModel,
    TraceModel,
    assert_dict_has_keys,
    assert_equal,
)

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

PROJECT_NAME = "aisuite-integration-test"
EXPECTED_OPENAI_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.completion_tokens_details.accepted_prediction_tokens": ANY_BUT_NONE,
    "original_usage.completion_tokens_details.audio_tokens": ANY_BUT_NONE,
    "original_usage.completion_tokens_details.reasoning_tokens": ANY_BUT_NONE,
    "original_usage.completion_tokens_details.rejected_prediction_tokens": ANY_BUT_NONE,
    "original_usage.prompt_tokens_details.audio_tokens": ANY_BUT_NONE,
    "original_usage.prompt_tokens_details.cached_tokens": ANY_BUT_NONE,
}


def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]):
    # max_tokens / max_completion_tokens is call-specific (OpenAI reasoning
    # models reject max_tokens; Anthropic takes it) so don't assert on it.
    REQUIRED_METADATA_KEYS = [
        "usage",
        "model",
        "created_from",
        "type",
        "id",
        "created",
        "object",
    ]
    assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS)


def test_aisuite__openai_provider__client_chat_completions_create__happyflow(
    fake_backend,
):
    client = aisuite.Client()
    wrapped_client = track_aisuite(
        aisuite_client=client,
        project_name=PROJECT_NAME,
    )
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell a fact"},
    ]

    _ = wrapped_client.chat.completions.create(
        model=llm_constants.AISUITE_OPENAI_GPT_NANO,
        messages=messages,
        max_completion_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
    )

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input={"messages": messages},
        output={"choices": ANY_BUT_NONE},
        tags=["aisuite"],
        metadata=ANY_DICT,
        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,
                type="llm",
                name="chat_completion_create",
                input={"messages": messages},
                output={"choices": ANY_BUT_NONE},
                tags=["aisuite"],
                metadata=ANY_DICT,
                usage=EXPECTED_OPENAI_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=PROJECT_NAME,
                spans=[],
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    trace_tree = fake_backend.trace_trees[0]

    assert_equal(EXPECTED_TRACE_TREE, trace_tree)

    llm_span_metadata = trace_tree.spans[0].metadata
    _assert_metadata_contains_required_keys(llm_span_metadata)


def test_aisuite__nonopenai_provider__client_chat_completions_create__happyflow(
    fake_backend,
):
    client = aisuite.Client()
    wrapped_client = track_aisuite(
        aisuite_client=client,
        project_name=PROJECT_NAME,
    )
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell a fact"},
    ]

    _ = wrapped_client.chat.completions.create(
        model=llm_constants.AISUITE_ANTHROPIC_CLAUDE_SONNET,
        messages=messages,
        max_tokens=10,
    )

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input={"messages": messages},
        output={"choices": ANY_BUT_NONE},
        tags=["aisuite"],
        metadata=ANY_DICT,
        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,
                type="llm",
                name="chat_completion_create",
                input={"messages": messages},
                output={"choices": ANY_BUT_NONE},
                tags=["aisuite"],
                metadata=ANY_DICT,
                usage=None,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=PROJECT_NAME,
                spans=[],
                model=ANY_STRING.starting_with(llm_constants.ANTHROPIC_CLAUDE_SONNET),
                provider="anthropic",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    trace_tree = fake_backend.trace_trees[0]

    assert_equal(EXPECTED_TRACE_TREE, trace_tree)


def test_aisuite_client_chat_completions_create__create_raises_an_error__span_and_trace_finished_gracefully__error_info_is_logged(
    fake_backend,
):
    client = aisuite.Client()
    wrapped_client = track_aisuite(
        aisuite_client=client,
        project_name=PROJECT_NAME,
    )

    # aisuite 0.1.3 stopped wrapping upstream errors in LLMError for the
    # OpenAI provider — the raw openai.BadRequestError now bubbles up. We
    # only care that Opik finishes the span gracefully on any failure.
    with pytest.raises(Exception):
        _ = wrapped_client.chat.completions.create(
            messages=None,
            model=llm_constants.AISUITE_OPENAI_GPT_NANO,
        )

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input={"messages": None},
        output=None,
        tags=["aisuite"],
        metadata={
            "created_from": "aisuite",
            "type": "aisuite_chat",
            "model": llm_constants.AISUITE_OPENAI_GPT_NANO,
        },
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=PROJECT_NAME,
        error_info={
            "exception_type": ANY_STRING,
            "message": ANY_STRING,
            "traceback": ANY_STRING,
        },
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input={"messages": None},
                output=None,
                tags=["aisuite"],
                metadata={
                    "created_from": "aisuite",
                    "type": "aisuite_chat",
                    "model": llm_constants.AISUITE_OPENAI_GPT_NANO,
                },
                usage=None,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=PROJECT_NAME,
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                error_info={
                    "exception_type": ANY_STRING,
                    "message": ANY_STRING,
                    "traceback": ANY_STRING,
                },
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1

    trace_tree = fake_backend.trace_trees[0]
    assert_equal(EXPECTED_TRACE_TREE, trace_tree)


def test_aisuite_client_chat_completions_create__openai_call_made_in_another_tracked_function__openai_span_attached_to_existing_trace(
    fake_backend,
):
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell a fact"},
    ]

    @opik.track(project_name=PROJECT_NAME)
    def f():
        client = aisuite.Client()
        wrapped_client = track_aisuite(
            aisuite_client=client,
            # we are trying to log span into another project, but parent's project name will be used
            project_name=f"{PROJECT_NAME}-nested-level",
        )

        _ = wrapped_client.chat.completions.create(
            model=llm_constants.AISUITE_OPENAI_GPT_NANO,
            messages=messages,
            max_completion_tokens=10,
            reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        )

    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={"messages": messages},
                        output={"choices": ANY_BUT_NONE},
                        tags=["aisuite"],
                        metadata=ANY_DICT,
                        usage=EXPECTED_OPENAI_USAGE_LOGGED_FORMAT,
                        start_time=ANY_BUT_NONE,
                        end_time=ANY_BUT_NONE,
                        project_name=PROJECT_NAME,
                        spans=[],
                        model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                        provider="openai",
                        source="sdk",
                    )
                ],
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1

    trace_tree = fake_backend.trace_trees[0]

    assert_equal(EXPECTED_TRACE_TREE, trace_tree)

    llm_span_metadata = trace_tree.spans[0].spans[0].metadata
    _assert_metadata_contains_required_keys(llm_span_metadata)


def test_aisuite__openai_provider__client_chat_completions_create__opik_args__happyflow(
    fake_backend,
):
    client = aisuite.Client()
    wrapped_client = track_aisuite(
        aisuite_client=client,
        project_name=PROJECT_NAME,
    )
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell a fact"},
    ]

    args_dict = {
        "span": {"tags": ["span_tag"], "metadata": {"span_key": "span_value"}},
        "trace": {
            "thread_id": "conversation-2",
            "tags": ["trace_tag"],
            "metadata": {"trace_key": "trace_value"},
        },
    }

    _ = wrapped_client.chat.completions.create(
        model=llm_constants.AISUITE_OPENAI_GPT_NANO,
        messages=messages,
        max_completion_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        opik_args=args_dict,
    )

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="chat_completion_create",
        input={"messages": messages},
        output={"choices": ANY_BUT_NONE},
        tags=["aisuite", "span_tag", "trace_tag"],
        metadata=ANY_DICT.containing({"trace_key": "trace_value"}),
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        project_name=PROJECT_NAME,
        thread_id="conversation-2",
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="chat_completion_create",
                input={"messages": messages},
                output={"choices": ANY_BUT_NONE},
                tags=["aisuite", "span_tag"],
                metadata=ANY_DICT.containing({"span_key": "span_value"}),
                usage=EXPECTED_OPENAI_USAGE_LOGGED_FORMAT,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                project_name=PROJECT_NAME,
                spans=[],
                model=ANY_STRING.starting_with(llm_constants.OPENAI_GPT_NANO),
                provider="openai",
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert len(fake_backend.trace_trees) == 1
    trace_tree = fake_backend.trace_trees[0]

    assert_equal(EXPECTED_TRACE_TREE, trace_tree)

    llm_span_metadata = trace_tree.spans[0].metadata
    _assert_metadata_contains_required_keys(llm_span_metadata)