comet-ml--opik/sdks/python/tests/library_integration/bedrock/test_converse.py

import boto3
import pytest

import opik
from opik.integrations.bedrock import track_bedrock
from ...testlib import (
    ANY_BUT_NONE,
    ANY_DICT,
    ANY_STRING,
    SpanModel,
    TraceModel,
    assert_equal,
)
from .constants import (
    BEDROCK_MODEL_FOR_TESTS,
    EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT,
    MISTRAL_PIXTRAL_MODEL_FOR_TESTS,
    MISTRAL_PIXTRAL_REGION_FOR_TESTS,
)

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


@pytest.mark.parametrize(
    "project_name, expected_project_name",
    [
        (None, "Default Project"),
        ("bedrock-integration-test", "bedrock-integration-test"),
    ],
)
def test_bedrock_converse__happyflow(fake_backend, project_name, expected_project_name):
    """Test basic converse functionality with Bedrock client."""
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client, project_name=project_name)

    messages = [{"role": "user", "content": [{"text": "Hello, how are you?"}]}]

    system_prompt = [
        {
            "text": "You are a helpful AI assistant. Provide concise and accurate responses."
        }
    ]

    _ = tracked_client.converse(
        modelId=BEDROCK_MODEL_FOR_TESTS,
        messages=messages,
        system=system_prompt,
        inferenceConfig={
            "maxTokens": 50,
            "temperature": 0.1,
        },
    )

    opik.flush_tracker()

    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse",
        input={"messages": messages, "system": system_prompt},
        output={"output": ANY_DICT},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        project_name=expected_project_name,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse",
                type="llm",
                input={"messages": messages, "system": system_prompt},
                output={"output": ANY_DICT},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model=BEDROCK_MODEL_FOR_TESTS,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )
    assert len(fake_backend.trace_trees) == 1

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


def test_bedrock_converse__create_raises_an_error__span_and_trace_finished_gracefully__error_info_is_logged(
    fake_backend,
):
    """Test that errors are properly logged as error spans."""
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client)

    messages = [{"role": "user", "content": [{"text": "Test message"}]}]

    # Use an invalid model to trigger an error
    with pytest.raises(Exception):
        tracked_client.converse(
            modelId="invalid-model-id",
            messages=messages,
            inferenceConfig={"maxTokens": 50},
        )

    opik.flush_tracker()

    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse",
        input={"messages": messages},
        output=None,
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        error_info=ANY_DICT.containing(
            {
                "exception_type": ANY_STRING,
                "message": ANY_STRING,
                "traceback": ANY_STRING,
            }
        ),
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse",
                type="llm",
                input={"messages": messages},
                output=None,
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model="invalid-model-id",
                provider="bedrock",
                error_info=ANY_DICT.containing(
                    {
                        "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, trace_tree)


def test_bedrock_converse__converse_call_made_in_another_tracked_function__bedrock_span_attached_to_existing_trace(
    fake_backend,
):
    """Test that converse calls within tracked functions create proper nesting."""
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client)

    @opik.track()
    def ask_bedrock_question(question: str) -> str:
        messages = [{"role": "user", "content": [{"text": question}]}]

        response = tracked_client.converse(
            modelId=BEDROCK_MODEL_FOR_TESTS,
            messages=messages,
            inferenceConfig={"maxTokens": 50},
        )

        return response["output"]["message"]["content"][0]["text"]

    result = ask_bedrock_question("What is 2+2?")
    opik.flush_tracker()

    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="ask_bedrock_question",
        input={"question": "What is 2+2?"},
        output={"output": result},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="ask_bedrock_question",
                input={"question": "What is 2+2?"},
                output={"output": result},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                last_updated_at=ANY_BUT_NONE,
                spans=[
                    SpanModel(
                        id=ANY_BUT_NONE,
                        name="bedrock_converse",
                        type="llm",
                        input={
                            "messages": [
                                {"role": "user", "content": [{"text": "What is 2+2?"}]}
                            ]
                        },
                        output={"output": ANY_DICT},
                        start_time=ANY_BUT_NONE,
                        end_time=ANY_BUT_NONE,
                        tags=["bedrock"],
                        metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                        last_updated_at=ANY_BUT_NONE,
                        model=BEDROCK_MODEL_FOR_TESTS,
                        usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                        provider="bedrock",
                        spans=[],
                        source="sdk",
                    )
                ],
                source="sdk",
            )
        ],
        source="sdk",
    )
    assert len(fake_backend.trace_trees) == 1

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


@pytest.mark.parametrize(
    "model_id",
    [
        # Standard Claude model - baseline for converse_stream event structure
        # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html
        BEDROCK_MODEL_FOR_TESTS,
        # DeepSeek R1 reasoning model - OPIK-2910: Different event structure for reasoning models
        # Reasoning models may have unique streaming patterns, including reasoning traces
        # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-deepseek.html
        "us.deepseek.r1-v1:0",
        # Amazon Nova - Tests Amazon's proprietary model streaming format
        # Nova models use different internal event structures
        # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-nova.html
        "us.amazon.nova-pro-v1:0",
        # Meta Llama - Tests open-source model integration
        # Llama models may have different tokenization and streaming patterns
        # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-meta.html
        "us.meta.llama3-1-8b-instruct-v1:0",
        # Mistral Pixtral - Tests multimodal model streaming (text focus in this test)
        # Multimodal models may include additional event types
        # Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mistral.html
        MISTRAL_PIXTRAL_MODEL_FOR_TESTS,
    ],
)
def test_bedrock_converse__stream_mode_is_on__generator_tracked_correctly(
    fake_backend, model_id
):
    region_name_by_model_id = {
        MISTRAL_PIXTRAL_MODEL_FOR_TESTS: MISTRAL_PIXTRAL_REGION_FOR_TESTS,
    }

    client = boto3.client(
        "bedrock-runtime",
        region_name=region_name_by_model_id.get(model_id, "us-east-1"),
    )
    tracked_client = track_bedrock(client)

    messages = [{"role": "user", "content": [{"text": "Hello, tell me a story"}]}]

    response = tracked_client.converse_stream(
        modelId=model_id,
        messages=messages,
        inferenceConfig={"maxTokens": 50},
    )

    for _ in response["stream"]:
        pass

    opik.flush_tracker()

    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse_stream",
        input={"messages": messages},
        output={"output": ANY_DICT},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse_stream",
                type="llm",
                input={"messages": messages},
                output={"output": ANY_DICT},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model=model_id,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )
    assert len(fake_backend.trace_trees) == 1

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


def test_bedrock_converse__stream_with_tool_use__structured_output_tracked_correctly(
    fake_backend,
):
    """
    Test converse_stream with tool use / structured output.

    This test verifies the fix for Issue #3829: KeyError 'text' when using streaming
    with structured output via toolConfig. When using tool use, the contentBlockDelta
    events contain delta.toolUse instead of delta.text.

    References:
    - Issue #3829: https://github.com/comet-ml/opik/issues/3829
    - ContentBlockDelta: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlockDelta.html
    - Tool Use Guide: https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html
    """
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client)

    messages = [{"role": "user", "content": [{"text": "What's the weather in Tokyo?"}]}]

    # Define a simple weather tool
    tool_config = {
        "tools": [
            {
                "toolSpec": {
                    "name": "get_weather",
                    "description": "Get the current weather for a location",
                    "inputSchema": {
                        "json": {
                            "type": "object",
                            "properties": {
                                "location": {
                                    "type": "string",
                                    "description": "City name, e.g., Tokyo",
                                }
                            },
                            "required": ["location"],
                        }
                    },
                }
            }
        ]
    }

    response = tracked_client.converse_stream(
        modelId=BEDROCK_MODEL_FOR_TESTS,
        messages=messages,
        toolConfig=tool_config,
        inferenceConfig={"maxTokens": 100},
    )

    # Consume the stream - should not raise KeyError
    for _ in response["stream"]:
        pass

    opik.flush_tracker()

    # Verify trace was created successfully with output
    assert len(fake_backend.trace_trees) == 1

    trace_tree = fake_backend.trace_trees[0]

    # Verify basic structure
    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse_stream",
        input={"messages": messages, "toolConfig": tool_config},
        output={"output": ANY_DICT},  # May contain text or toolUse
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse_stream",
                type="llm",
                input={"messages": messages, "toolConfig": tool_config},
                output={
                    "output": ANY_DICT
                },  # Simplified - just verify structure exists
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model=BEDROCK_MODEL_FOR_TESTS,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert_equal(expected_trace, trace_tree)


def test_bedrock_converse__stream_called_2_times__generator_tracked_correctly(
    fake_backend,
):
    """Test that multiple converse_stream calls create separate spans."""
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client)

    # Make first stream call
    response1 = tracked_client.converse_stream(
        modelId=BEDROCK_MODEL_FOR_TESTS,
        messages=[{"role": "user", "content": [{"text": "Hello"}]}],
        inferenceConfig={"maxTokens": 20},
    )

    # Consume the first stream
    for _ in response1["stream"]:
        pass

    # Make second stream call
    response2 = tracked_client.converse_stream(
        modelId=BEDROCK_MODEL_FOR_TESTS,
        messages=[{"role": "user", "content": [{"text": "Goodbye"}]}],
        inferenceConfig={"maxTokens": 20},
    )

    # Consume the second stream
    for _ in response2["stream"]:
        pass

    opik.flush_tracker()

    # Should have two separate trace trees
    assert len(fake_backend.trace_trees) == 2

    # Verify first trace
    messages1 = [{"role": "user", "content": [{"text": "Hello"}]}]
    expected_trace1 = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse_stream",
        input={"messages": messages1},
        output={"output": ANY_DICT},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse_stream",
                type="llm",
                input={"messages": messages1},
                output={"output": ANY_DICT},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model=BEDROCK_MODEL_FOR_TESTS,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )

    # Verify second trace
    messages2 = [{"role": "user", "content": [{"text": "Goodbye"}]}]
    expected_trace2 = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse_stream",
        input={"messages": messages2},
        output={"output": ANY_DICT},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        tags=["bedrock"],
        metadata=ANY_DICT,
        last_updated_at=ANY_BUT_NONE,
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse_stream",
                type="llm",
                input={"messages": messages2},
                output={"output": ANY_DICT},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock"],
                metadata=ANY_DICT.containing({"created_from": "bedrock"}),
                last_updated_at=ANY_BUT_NONE,
                model=BEDROCK_MODEL_FOR_TESTS,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )

    assert_equal(expected_trace1, fake_backend.trace_trees[0])
    assert_equal(expected_trace2, fake_backend.trace_trees[1])


@pytest.mark.parametrize(
    "project_name, expected_project_name",
    [
        (None, "Default Project"),
        ("bedrock-integration-test", "bedrock-integration-test"),
    ],
)
def test_bedrock_converse__opik_args__happyflow(
    fake_backend, project_name, expected_project_name
):
    """Test basic converse functionality with Bedrock client."""
    client = boto3.client("bedrock-runtime", region_name="us-east-1")
    tracked_client = track_bedrock(client, project_name=project_name)

    messages = [{"role": "user", "content": [{"text": "Hello, how are you?"}]}]

    system_prompt = [
        {
            "text": "You are a helpful AI assistant. Provide concise and accurate responses."
        }
    ]

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

    _ = tracked_client.converse(
        modelId=BEDROCK_MODEL_FOR_TESTS,
        messages=messages,
        system=system_prompt,
        inferenceConfig={
            "maxTokens": 50,
            "temperature": 0.1,
        },
        opik_args=args_dict,
    )

    opik.flush_tracker()

    expected_trace = TraceModel(
        id=ANY_BUT_NONE,
        name="bedrock_converse",
        input={"messages": messages, "system": system_prompt},
        output={"output": ANY_DICT},
        start_time=ANY_BUT_NONE,
        end_time=ANY_BUT_NONE,
        project_name=expected_project_name,
        tags=["bedrock", "span_tag", "trace_tag"],
        metadata=ANY_DICT.containing({"trace_key": "trace_value"}),
        last_updated_at=ANY_BUT_NONE,
        thread_id="conversation-2",
        spans=[
            SpanModel(
                id=ANY_BUT_NONE,
                name="bedrock_converse",
                type="llm",
                input={"messages": messages, "system": system_prompt},
                output={"output": ANY_DICT},
                start_time=ANY_BUT_NONE,
                end_time=ANY_BUT_NONE,
                tags=["bedrock", "span_tag"],
                metadata=ANY_DICT.containing(
                    {"created_from": "bedrock", "span_key": "span_value"}
                ),
                last_updated_at=ANY_BUT_NONE,
                model=BEDROCK_MODEL_FOR_TESTS,
                usage=ANY_DICT.containing(EXPECTED_BEDROCK_USAGE_LOGGED_FORMAT),
                provider="bedrock",
                spans=[],
                source="sdk",
            )
        ],
        source="sdk",
    )
    assert len(fake_backend.trace_trees) == 1

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