opik 1.8.39__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/integrations/bedrock/converse/chunks_aggregator.py

@@ -0,0 +1,188 @@
+import logging
+from typing import Any, Dict, List
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _handle_message_start(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract role from messageStart event."""
+    message_start = event.get("messageStart")
+    if isinstance(message_start, dict):
+        role = message_start.get("role")
+        if role:
+            result["output"]["message"]["role"] = role
+
+
+def _handle_content_block_delta(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """
+    Extract content from contentBlockDelta event.
+
+    Handles multiple delta types:
+    - delta.text: Regular text streaming
+    - delta.toolUse: Structured output / tool calls (Issue #3829)
+    """
+    content_block_delta = event.get("contentBlockDelta")
+    if not isinstance(content_block_delta, dict):
+        return
+
+    delta = content_block_delta.get("delta")
+    if not isinstance(delta, dict):
+        return
+
+    content = result["output"]["message"]["content"][0]
+
+    # Handle regular text streaming
+    if "text" in delta:
+        content["text"] += delta["text"]
+        return
+
+    # Handle structured output / tool use (Issue #3829)
+    # Ref: https://github.com/comet-ml/opik/issues/3829
+    if "toolUse" in delta:
+        if "toolUse" not in content:
+            content["toolUse"] = {}
+        content["toolUse"].update(delta["toolUse"])
+        return
+
+    # Log other delta types for future compatibility
+    LOGGER.debug("Unknown delta type in contentBlockDelta: %s", list(delta.keys()))
+
+
+def _handle_message_stop(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract stopReason from messageStop event."""
+    message_stop = event.get("messageStop")
+    if isinstance(message_stop, dict):
+        stop_reason = message_stop.get("stopReason")
+        if stop_reason:
+            result["stopReason"] = stop_reason
+
+
+def _handle_metadata(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract usage and metrics from metadata event."""
+    metadata = event.get("metadata")
+    if not isinstance(metadata, dict):
+        return
+
+    # Extract usage information
+    if "usage" in metadata:
+        result["usage"] = metadata["usage"]
+
+    # Extract metrics information
+    if "metrics" in metadata:
+        metrics = metadata["metrics"]
+        if isinstance(metrics, dict) and "latencyMs" in metrics:
+            result["metrics"] = {"latencyMs": metrics["latencyMs"]}
+
+
+def aggregate_converse_stream_chunks(items: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Aggregate streaming chunks from AWS Bedrock converse_stream API into a single response.
+
+    This function handles various event structures from different Bedrock model providers:
+    - Anthropic (Claude): Standard messageStart, contentBlockDelta, messageStop events
+    - Amazon (Nova): Amazon's proprietary event format with potential variations
+    - Meta (Llama): Open-source model events with different tokenization patterns
+    - Mistral (Pixtral): Multimodal model events that may include additional content types
+    - DeepSeek (R1): Reasoning model events with extended thought processes (OPIK-2910 fix)
+
+    Event Structure Variations by Provider:
+    ========================================
+
+    Standard Converse Stream Events (from AWS documentation):
+    - messageStart: Contains role and initial metadata
+    - contentBlockStart: Marks beginning of content block
+    - contentBlockDelta: Contains incremental text in delta.text
+    - contentBlockStop: Marks end of content block
+    - messageStop: Contains stopReason (e.g., "end_turn", "stop_sequence")
+    - metadata: Contains usage stats and performance metrics
+
+    Known Variations:
+    - DeepSeek R1: May have different delta structure or additional reasoning fields (OPIK-2910)
+    - Tool Use/Structured Output: delta.toolUse instead of delta.text (Issue #3829)
+    - Multimodal models: May include non-text content blocks (images, documents)
+    - Different models: Varying field names, nesting levels, or optional fields
+
+    References:
+    - AWS Bedrock Converse API: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
+    - Streaming Events: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStreamOutput.html
+    - 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
+    - DeepSeek R1 Issue: https://comet-ml.atlassian.net/browse/OPIK-2910
+    - Tool Use Issue: https://github.com/comet-ml/opik/issues/3829
+
+    Args:
+        items: List of streaming event dictionaries from Bedrock converse_stream
+
+    Returns:
+        Aggregated response dictionary with structure:
+        {
+            "output": {
+                "message": {
+                    "role": "assistant",
+                    "content": [{"text": "aggregated text"}]
+                }
+            },
+            "stopReason": "end_turn",  # Optional
+            "usage": {...},  # Optional
+            "metrics": {"latencyMs": ...}  # Optional
+        }
+    """
+
+    result: Dict[str, Any] = {
+        "output": {"message": {"role": "assistant", "content": [{"text": ""}]}}
+    }
+
+    for event in items:
+        if not isinstance(event, dict):
+            LOGGER.debug("Skipping non-dictionary event: %s", type(event))
+            continue
+
+        try:
+            if "messageStart" in event:
+                _handle_message_start(event, result)
+
+            if "contentBlockDelta" in event:
+                _handle_content_block_delta(event, result)
+
+            if "messageStop" in event:
+                _handle_message_stop(event, result)
+
+            if "metadata" in event:
+                _handle_metadata(event, result)
+
+        except Exception as e:
+            LOGGER.warning(
+                "Unexpected error processing event: %s. Event: %s",
+                str(e),
+                event,
+                exc_info=True,
+            )
+
+    return result
+
+
+def aggregate_invoke_agent_chunks(items: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Aggregate streaming chunks from AWS Bedrock invoke_agent API.
+
+    Note: The implementation uses a simplified approach as the completion payload
+    only contains chunks without additional metadata (as of implementation date).
+
+    Reference:
+    - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/invoke_agent.html
+
+    Args:
+        items: List of chunk event dictionaries from invoke_agent stream
+
+    Returns:
+        Aggregated response dictionary with decoded text output
+    """
+    merged_chunks = b""
+
+    for item in items:
+        if isinstance(item, dict) and "chunk" in item:
+            chunk = item["chunk"]
+            if isinstance(chunk, dict) and "bytes" in chunk:
+                merged_chunks += chunk["bytes"]
+
+    return {"output": merged_chunks.decode("utf-8")}

opik/integrations/bedrock/{converse_decorator.py → converse/converse_decorator.py}

@@ -8,7 +8,8 @@ import opik.llm_usage as llm_usage
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 
-from . import helpers, stream_wrappers
+from . import stream_wrappers
+from .. import types
 
 LOGGER = logging.getLogger(__name__)
 
@@ -91,7 +92,7 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
         capture_output: bool,
         generations_aggregator: Optional[Callable[[List[Any]], Any]],
     ) -> Union[
-        helpers.ConverseStreamOutput,
+        types.ConverseStreamOutput,
         None,
     ]:
         DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR = (
@@ -117,7 +118,7 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
             )
 
             output["stream"] = wrapped_stream
-            return cast(helpers.ConverseStreamOutput, output)
+            return cast(types.ConverseStreamOutput, output)
 
         STREAM_NOT_FOUND = None
 

opik/integrations/bedrock/invoke_agent_decorator.py

@@ -6,7 +6,8 @@ import opik.dict_utils as dict_utils
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 
-from . import helpers, stream_wrappers
+from . import types
+from .converse import stream_wrappers as converse_stream_wrappers
 
 LOGGER = logging.getLogger(__name__)
 
@@ -77,7 +78,7 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
         capture_output: bool,
         generations_aggregator: Optional[Callable[[List[Any]], Any]],
     ) -> Union[
-        helpers.ConverseStreamOutput,
+        types.ConverseStreamOutput,
         None,
     ]:
         DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR = (
@@ -92,7 +93,7 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
         if isinstance(output, dict) and "completion" in output:
             span_to_end, trace_to_end = base_track_decorator.pop_end_candidates()
 
-            wrapped_stream = stream_wrappers.wrap_stream(
+            wrapped_stream = converse_stream_wrappers.wrap_stream(
                 stream=output["completion"],
                 capture_output=capture_output,
                 span_to_end=span_to_end,
@@ -103,7 +104,7 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
             )
 
             output["completion"] = wrapped_stream
-            return cast(helpers.ConverseStreamOutput, output)
+            return cast(types.ConverseStreamOutput, output)
 
         STREAM_NOT_FOUND = None
 

opik/integrations/bedrock/invoke_model/chunks_aggregator/__init__.py

@@ -0,0 +1,78 @@
+"""
+Bedrock invoke_model_with_response_stream chunk aggregator.
+
+DISCLAIMER: This package was generated with AI assistance.
+
+This package provides a modular, extensible architecture for aggregating
+streaming chunks from different Bedrock model formats (Claude, Nova, etc.).
+
+## Public API
+
+The main public function is `aggregate_chunks_to_dataclass()` which returns a
+structured `BedrockAggregatedResponse` dataclass for type safety.
+
+```python
+from opik.integrations.bedrock.invoke_model.chunks_aggregator import aggregate_chunks_to_dataclass
+
+# Create typed aggregated response
+response = aggregate_chunks_to_dataclass(chunks)
+
+# Access native format and usage
+native_output = response.native_response  # Provider-specific format
+usage = response.usage  # Bedrock format: {inputTokens, outputTokens, totalTokens}
+
+# Use in span logging
+output = response.to_output_format()  # {"body": native_response}
+metadata = response.to_metadata_format()  # {"created_from": "bedrock", ...}
+```
+
+Each format aggregator returns the response in its native structure,
+with usage standardized to Bedrock format (camelCase: inputTokens, outputTokens, totalTokens).
+
+## Architecture
+
+- `base`: Base types and protocols
+- `claude`: Claude/Anthropic format aggregator
+- `llama`: Meta Llama format aggregator
+- `mistral`: Mistral/Pixtral format aggregator
+- `nova`: Amazon Nova format aggregator
+- `format_detector`: Format detection and aggregator registry
+- `response_types`: Dataclass definitions for structured responses
+
+## Adding New Formats
+
+To add support for a new model format:
+
+1. Create a new aggregator module (e.g., `titan.py`):
+   ```python
+   class TitanAggregator:
+       def aggregate(self, items) -> Dict[str, Any]:
+           # Return Titan's native structure with Bedrock usage
+           return {
+               "outputText": "...",
+               "usage": {"inputTokens": 10, "outputTokens": 20, "totalTokens": 30}
+           }
+   ```
+
+2. Add detection and registration in `format_detector.py`:
+   ```python
+   from . import titan
+
+   def _is_titan_format(chunk_data):
+       return "titan_specific_field" in chunk_data
+
+   # Add to _DETECTORS registry
+   _DETECTORS["titan"] = _is_titan_format
+
+   # Add to _AGGREGATORS registry
+   _AGGREGATORS["titan"] = titan.TitanAggregator()
+   ```
+
+The new format will be automatically detected and used.
+"""
+
+from .api import aggregate_chunks_to_dataclass
+
+__all__ = [
+    "aggregate_chunks_to_dataclass",
+]

opik/integrations/bedrock/invoke_model/chunks_aggregator/api.py

@@ -0,0 +1,45 @@
+"""
+API functions for Bedrock chunk aggregation.
+
+This module contains the main public functions for aggregating streaming chunks
+from different Bedrock model formats into structured responses.
+"""
+
+from typing import Any, Dict, List
+
+from . import format_detector
+from .. import response_types
+
+
+def aggregate_chunks_to_dataclass(
+    items: List[Dict[str, Any]],
+) -> response_types.BedrockAggregatedResponse:
+    """
+    Aggregate chunks directly to structured dataclass (used by stream wrapper).
+
+    This function is called by the stream wrapper and returns a structured
+    dataclass that the decorator can handle directly.
+
+    Args:
+        items: List of chunk items from the event stream
+
+    Returns:
+        BedrockAggregatedResponse with structured data
+    """
+    # Detect format and get appropriate aggregator
+    format_name = format_detector.detect_format(items)
+    aggregator = format_detector.get_aggregator(format_name)
+
+    # Get aggregated data
+    aggregated_data = aggregator.aggregate(items)
+
+    # Extract components directly
+    usage = aggregated_data.get("usage", {})
+    native_response = {
+        key: value for key, value in aggregated_data.items() if key != "usage"
+    }
+
+    # Create dataclass directly (response_metadata will be set by stream wrapper)
+    return response_types.BedrockAggregatedResponse(
+        native_response=native_response, usage=usage, response_metadata={}
+    )

opik/integrations/bedrock/invoke_model/chunks_aggregator/base.py

@@ -0,0 +1,23 @@
+"""Base types and protocols for chunk aggregators."""
+
+from typing import Any, Dict, List, Protocol
+
+
+class ChunkAggregator(Protocol):
+    """Protocol for chunk aggregators."""
+
+    def aggregate(self, items: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Aggregate streaming chunks in the format native to the provider.
+
+        The returned dictionary contains:
+        - Provider-specific response structure (e.g., Claude's content blocks, Nova's output)
+        - "usage" field with Bedrock format: {inputTokens, outputTokens, totalTokens}
+
+        Args:
+            items: List of chunk items from the event stream
+
+        Returns:
+            Dict with provider-native structure and Bedrock-standardized usage
+        """
+        ...

opik/integrations/bedrock/invoke_model/chunks_aggregator/claude.py

@@ -0,0 +1,121 @@
+"""Claude/Anthropic chunk aggregator."""
+
+import json
+import logging
+from typing import Any, Dict, List
+
+from .. import usage_converters
+from .base import ChunkAggregator
+
+LOGGER = logging.getLogger(__name__)
+
+
+class ClaudeAggregator(ChunkAggregator):
+    """
+    Aggregator for Claude/Anthropic streaming format.
+
+    Returns Claude's native message format with content blocks,
+    with usage standardized to Bedrock format (camelCase).
+
+    Claude chunk structure uses snake_case:
+    - message_start: Contains role and initial usage
+    - content_block_start: Content block metadata
+    - content_block_delta: Text chunks in delta.text
+    - content_block_stop: End of content block
+    - message_delta: Stop reason and updated usage
+    - message_stop: amazon-bedrock-invocationMetrics
+    """
+
+    def aggregate(self, items: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Aggregate Claude chunks into native Claude message structure with Bedrock usage."""
+        LOGGER.debug("Claude aggregator processing %d items", len(items))
+
+        content_blocks = []
+        current_text = ""
+        role = "assistant"
+        stop_reason = None
+        input_tokens = 0
+        output_tokens = 0
+
+        for item in items:
+            if "chunk" not in item:
+                continue
+
+            try:
+                chunk_data = json.loads(item["chunk"]["bytes"])
+                chunk_type = chunk_data.get("type", "")
+
+                if chunk_type == "message_start":
+                    message = chunk_data.get("message", {})
+                    role = message.get("role", "assistant")
+                    usage = message.get("usage", {})
+                    input_tokens = usage.get("input_tokens", 0)
+                    output_tokens = usage.get("output_tokens", 0)
+                    LOGGER.debug(
+                        "Claude message_start: input_tokens=%d, output_tokens=%d",
+                        input_tokens,
+                        output_tokens,
+                    )
+
+                elif chunk_type == "content_block_delta":
+                    delta = chunk_data.get("delta", {})
+                    if "text" in delta:
+                        current_text += delta["text"]
+
+                elif chunk_type == "content_block_stop":
+                    if current_text:
+                        content_blocks.append({"type": "text", "text": current_text})
+                        current_text = ""
+
+                elif chunk_type == "message_delta":
+                    delta = chunk_data.get("delta", {})
+                    if "stop_reason" in delta:
+                        stop_reason = delta["stop_reason"]
+                    usage = chunk_data.get("usage", {})
+                    if "output_tokens" in usage:
+                        output_tokens = usage["output_tokens"]
+                        LOGGER.debug(
+                            "Claude message_delta: output_tokens=%d", output_tokens
+                        )
+
+                elif chunk_type == "message_stop":
+                    metrics = chunk_data.get("amazon-bedrock-invocationMetrics", {})
+                    if metrics:
+                        input_tokens = metrics.get("inputTokenCount", input_tokens)
+                        output_tokens = metrics.get("outputTokenCount", output_tokens)
+                        LOGGER.debug(
+                            "Claude bedrock metrics: input=%d, output=%d",
+                            input_tokens,
+                            output_tokens,
+                        )
+
+            except (json.JSONDecodeError, KeyError, TypeError) as e:
+                LOGGER.debug("Claude aggregator error processing chunk: %s", e)
+                continue
+
+        # Add any remaining text
+        if current_text:
+            content_blocks.append({"type": "text", "text": current_text})
+
+        total_text = "".join(
+            block["text"] for block in content_blocks if block["type"] == "text"
+        )
+        LOGGER.debug(
+            "Claude aggregated: %d chars, input_tokens=%d, output_tokens=%d",
+            len(total_text),
+            input_tokens,
+            output_tokens,
+        )
+
+        # Convert to Bedrock usage format using shared converter
+        bedrock_usage = usage_converters.anthropic_to_bedrock_usage(
+            {"input_tokens": input_tokens, "output_tokens": output_tokens}
+        )
+
+        # Return Claude's native format with Bedrock usage
+        return {
+            "role": role,
+            "content": content_blocks,
+            "stop_reason": stop_reason,
+            "usage": bedrock_usage,
+        }

opik/integrations/bedrock/invoke_model/chunks_aggregator/format_detector.py

@@ -0,0 +1,107 @@
+"""Format detection and aggregator registry."""
+
+import json
+from typing import Any, Callable, Dict, List
+
+from .base import ChunkAggregator
+from . import claude
+from . import llama
+from . import mistral
+from . import nova
+
+
+# Format detection functions
+FormatDetector = Callable[[Dict[str, Any]], bool]
+
+
+def _is_nova_format(chunk_data: Dict[str, Any]) -> bool:
+    """Check if chunk is Nova format (camelCase fields)."""
+    return "contentBlockDelta" in chunk_data or "messageStart" in chunk_data
+
+
+def _is_claude_format(chunk_data: Dict[str, Any]) -> bool:
+    """Check if chunk is Claude format (snake_case fields with type)."""
+    return "type" in chunk_data
+
+
+def _is_llama_format(chunk_data: Dict[str, Any]) -> bool:
+    """Check if chunk is Llama format (generation field)."""
+    return "generation" in chunk_data
+
+
+def _is_mistral_format(chunk_data: Dict[str, Any]) -> bool:
+    """Check if chunk is Mistral/Pixtral format (OpenAI-like with choices and object)."""
+    return (
+        "object" in chunk_data
+        and chunk_data["object"] == "chat.completion.chunk"
+        and "choices" in chunk_data
+        and chunk_data["choices"]
+        and "message" in chunk_data["choices"][0]
+    )
+
+
+# Format detectors registry (ordered by specificity - most specific first)
+_DETECTORS: Dict[str, FormatDetector] = {
+    "mistral": _is_mistral_format,  # Specific (has object field)
+    "llama": _is_llama_format,  # Specific (has generation field)
+    "nova": _is_nova_format,  # Specific (has contentBlockDelta)
+    "claude": _is_claude_format,  # Generic (has type field)
+}
+
+# Aggregators registry
+_AGGREGATORS: Dict[str, ChunkAggregator] = {
+    "claude": claude.ClaudeAggregator(),
+    "llama": llama.LlamaAggregator(),
+    "mistral": mistral.MistralAggregator(),
+    "nova": nova.NovaAggregator(),
+}
+
+
+def detect_format(items: List[Dict[str, Any]]) -> str:
+    """
+    Detect streaming format from the first chunk.
+
+    Args:
+        items: List of chunk items from the event stream
+
+    Returns:
+        Format name (e.g., "claude", "nova") or "claude" as default
+    """
+    for item in items:
+        if "chunk" not in item:
+            continue
+
+        try:
+            chunk_data = json.loads(item["chunk"]["bytes"])
+
+            # Try each registered detector
+            for format_name, detector in _DETECTORS.items():
+                if detector(chunk_data):
+                    return format_name
+
+        except (json.JSONDecodeError, KeyError, TypeError):
+            continue
+
+    # Default to Claude format
+    return "claude"
+
+
+def get_aggregator(format_name: str) -> ChunkAggregator:
+    """
+    Get aggregator for the specified format.
+
+    Args:
+        format_name: Name of the format
+
+    Returns:
+        ChunkAggregator instance
+
+    Raises:
+        ValueError: If format is not registered
+    """
+    if format_name not in _AGGREGATORS:
+        raise ValueError(
+            f"Unknown format: {format_name}. Registered formats: {list(_AGGREGATORS.keys())}"
+        )
+
+    return _AGGREGATORS[format_name]