comet-ml--opik/sdks/python/tests/library_integration/litellm/test_litellm_streaming.py

import pytest

import litellm
import litellm.litellm_core_utils.streaming_handler

import opik
from opik.integrations.litellm import track_completion
from ... import llm_constants
from ...testlib import (
    ANY_BUT_NONE,
    ANY_DICT,
    ANY_LIST,
    ANY_STRING,
    SpanModel,
    TraceModel,
    assert_equal,
)

from . import constants

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

MODEL_FOR_TESTS = constants.MODEL_FOR_TESTS


@pytest.mark.parametrize(
    "model,expected_provider,extra_call_kwargs", constants.TEST_MODELS_PARAMETRIZE
)
def test_litellm_completion_streaming__happyflow(
    fake_backend, model, expected_provider, extra_call_kwargs
):
    """Test basic LiteLLM streaming completion tracking."""
    tracked_completion = track_completion()(litellm.completion)

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Say hello in one word"},
    ]

    stream = tracked_completion(
        model=model,
        messages=messages,
        max_tokens=64,
        stream=True,
        stream_options={"include_usage": True},
        **extra_call_kwargs,
    )

    # Consume the stream
    chunks_count = 0
    full_text = ""
    for chunk in stream:
        chunks_count += 1
        if chunk.choices and chunk.choices[0].delta.content:
            full_text += chunk.choices[0].delta.content

    opik.flush_tracker()

    # Verify we got chunks
    assert chunks_count > 0, "Should have received streaming chunks"
    assert len(full_text) > 0, "Should have received text content"

    # Verify the trace structure
    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="completion",
        input={"messages": messages},
        output={"choices": ANY_LIST},  # Aggregated output
        tags=["litellm"],
        metadata=ANY_DICT.containing(
            {
                "created_from": "litellm",
                "max_tokens": 64,
            }
        ),
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="completion",
                input={"messages": messages},
                output={"choices": ANY_LIST},  # Aggregated output
                tags=["litellm"],
                metadata=ANY_DICT.containing(
                    {
                        "created_from": "litellm",
                        "max_tokens": 64,
                    }
                ),
                usage=constants.EXPECTED_LITELLM_USAGE_LOGGED_FORMAT,  # Usage info must be present
                total_cost=ANY_BUT_NONE,  # Cost calculated by LiteLLM
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                spans=[],
                model=ANY_STRING,
                provider=expected_provider,
                source="sdk",
            )
        ],
        source="sdk",
    )

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


@pytest.mark.asyncio
async def test_litellm_acompletion_streaming__happyflow(fake_backend):
    """Test async LiteLLM streaming completion tracking."""
    tracked_acompletion = track_completion()(litellm.acompletion)

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Say hello in one word"},
    ]

    stream = await tracked_acompletion(
        model=MODEL_FOR_TESTS,
        messages=messages,
        max_tokens=64,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        stream=True,
        stream_options={"include_usage": True},
    )

    # Consume the stream
    chunks_count = 0
    full_text = ""
    async for chunk in stream:
        chunks_count += 1
        if chunk.choices and chunk.choices[0].delta.content:
            full_text += chunk.choices[0].delta.content

    opik.flush_tracker()

    # Verify we got chunks
    assert chunks_count > 0, "Should have received streaming chunks"
    assert len(full_text) > 0, "Should have received text content"

    # Verify the trace structure
    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="acompletion",
        input={"messages": messages},
        output={"choices": ANY_LIST},  # Aggregated output
        tags=["litellm"],
        metadata=ANY_DICT.containing(
            {
                "created_from": "litellm",
                "max_tokens": 64,
            }
        ),
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="acompletion",
                input={"messages": messages},
                output={"choices": ANY_LIST},  # Aggregated output
                tags=["litellm"],
                metadata=ANY_DICT.containing(
                    {
                        "created_from": "litellm",
                        "max_tokens": 64,
                    }
                ),
                usage=constants.EXPECTED_LITELLM_USAGE_LOGGED_FORMAT,  # Usage info must be present
                total_cost=ANY_BUT_NONE,  # Cost calculated by LiteLLM
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                spans=[],
                model=ANY_STRING,
                provider="openai",
                source="sdk",
            )
        ],
        source="sdk",
    )

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


def test_litellm_completion_streaming_with_opik_args__happyflow(fake_backend):
    """Test LiteLLM streaming with custom opik_args."""
    tracked_completion = track_completion()(litellm.completion)

    messages = [
        {"role": "user", "content": "Hello"},
    ]

    args_dict = {
        "span": {
            "tags": ["streaming-span"],
            "metadata": {"stream_key": "stream_value"},
        },
        "trace": {
            "thread_id": "stream-thread-1",
            "tags": ["streaming-trace"],
            "metadata": {"trace_key": "trace_value"},
        },
    }

    stream = tracked_completion(
        model=MODEL_FOR_TESTS,
        messages=messages,
        max_tokens=10,
        reasoning_effort=llm_constants.OPENAI_REASONING_EFFORT,
        stream=True,
        stream_options={"include_usage": True},
        opik_args=args_dict,
    )

    # Consume the stream
    for _ in stream:
        pass

    opik.flush_tracker()

    EXPECTED_TRACE_TREE = TraceModel(
        id=ANY_BUT_NONE,
        name="completion",
        input={"messages": messages},
        output={"choices": ANY_LIST},
        tags=["litellm", "streaming-span", "streaming-trace"],
        metadata=ANY_DICT.containing(
            {"created_from": "litellm", "max_tokens": 10, "trace_key": "trace_value"}
        ),
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        thread_id="stream-thread-1",
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                type="llm",
                name="completion",
                input={"messages": messages},
                output={"choices": ANY_LIST},
                tags=["litellm", "streaming-span"],
                metadata=ANY_DICT.containing(
                    {
                        "created_from": "litellm",
                        "max_tokens": 10,
                        "stream_key": "stream_value",
                    }
                ),
                usage=constants.EXPECTED_LITELLM_USAGE_LOGGED_FORMAT,
                total_cost=ANY_BUT_NONE,  # Cost calculated by LiteLLM
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                spans=[],
                model=ANY_STRING,
                provider="openai",
                source="sdk",
            )
        ],
        source="sdk",
    )

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