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

opik/integrations/bedrock/invoke_agent_decorator.py

@@ -1,11 +1,13 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing_extensions import override
 
-from opik import dict_utils
+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__)
 
@@ -22,12 +24,13 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
     overrides _generators_handler() method to work correctly with bedrock's streams
     """
 
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
         track_options: arguments_helpers.TrackOptions,
-        args: Optional[Tuple],
-        kwargs: Optional[Dict[str, Any]],
+        args: Tuple,
+        kwargs: Dict[str, Any],
     ) -> arguments_helpers.StartSpanParameters:
         assert (
             kwargs is not None
@@ -51,6 +54,7 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -67,13 +71,14 @@ class BedrockInvokeAgentDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(  # type: ignore
         self,
         output: Any,
         capture_output: bool,
         generations_aggregator: Optional[Callable[[List[Any]], Any]],
     ) -> Union[
-        helpers.ConverseStreamOutput,
+        types.ConverseStreamOutput,
         None,
     ]:
         DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR = (
@@ -88,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,
@@ -99,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]

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

@@ -0,0 +1,108 @@
+"""Meta Llama 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 LlamaAggregator(ChunkAggregator):
+    """
+    Aggregator for Meta Llama streaming format.
+
+    Returns Llama's native output format with Bedrock usage.
+
+    Llama chunk structure:
+    - generation: Contains generated text
+    - prompt_token_count: Input tokens (only in first chunk)
+    - generation_token_count: Cumulative output tokens
+    - stop_reason: Stop reason in final chunk
+    - amazon-bedrock-invocationMetrics: Final metrics
+    """
+
+    def aggregate(self, items: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Aggregate Llama chunks into native Llama output with Bedrock usage."""
+        LOGGER.debug("Llama aggregator processing %d items", len(items))
+
+        generation_text = ""
+        stop_reason = None
+        prompt_token_count = 0
+        generation_token_count = 0
+
+        for item in items:
+            if "chunk" not in item:
+                continue
+
+            try:
+                chunk_data = json.loads(item["chunk"]["bytes"])
+
+                # Extract generated text
+                if "generation" in chunk_data and chunk_data["generation"]:
+                    generation_text += chunk_data["generation"]
+
+                # Extract prompt token count from first chunk
+                if (
+                    "prompt_token_count" in chunk_data
+                    and chunk_data["prompt_token_count"]
+                ):
+                    prompt_token_count = chunk_data["prompt_token_count"]
+                    LOGGER.debug("Llama prompt_token_count: %d", prompt_token_count)
+
+                # Extract generation token count (cumulative)
+                if (
+                    "generation_token_count" in chunk_data
+                    and chunk_data["generation_token_count"]
+                ):
+                    generation_token_count = chunk_data["generation_token_count"]
+
+                # Extract stop reason
+                if "stop_reason" in chunk_data and chunk_data["stop_reason"]:
+                    stop_reason = chunk_data["stop_reason"]
+                    LOGGER.debug("Llama stop_reason: %s", stop_reason)
+
+                # Use bedrock metrics as authoritative source
+                metrics = chunk_data.get("amazon-bedrock-invocationMetrics", {})
+                if metrics:
+                    prompt_token_count = metrics.get(
+                        "inputTokenCount", prompt_token_count
+                    )
+                    generation_token_count = metrics.get(
+                        "outputTokenCount", generation_token_count
+                    )
+                    LOGGER.debug(
+                        "Llama bedrock metrics: input=%d, output=%d",
+                        prompt_token_count,
+                        generation_token_count,
+                    )
+
+            except (json.JSONDecodeError, KeyError, TypeError) as e:
+                LOGGER.debug("Llama aggregator error processing chunk: %s", e)
+                continue
+
+        LOGGER.debug(
+            "Llama aggregated: %d chars, prompt_tokens=%d, generation_tokens=%d",
+            len(generation_text),
+            prompt_token_count,
+            generation_token_count,
+        )
+
+        # Convert to Bedrock usage format using shared converter
+        bedrock_usage = usage_converters.llama_to_bedrock_usage(
+            {
+                "prompt_token_count": prompt_token_count,
+                "generation_token_count": generation_token_count,
+            }
+        )
+
+        # Return Llama's native output format with Bedrock usage
+        return {
+            "generation": generation_text,
+            "prompt_token_count": prompt_token_count,
+            "generation_token_count": generation_token_count,
+            "stop_reason": stop_reason,
+            "usage": bedrock_usage,
+        }

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

@@ -0,0 +1,118 @@
+"""Mistral (Pixtral) 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 MistralAggregator(ChunkAggregator):
+    """
+    Aggregator for Mistral (Pixtral) streaming format.
+
+    Returns Mistral's native OpenAI-like output format with Bedrock usage.
+
+    Mistral chunk structure (OpenAI-compatible):
+    - choices[0].message.content: Generated text
+    - choices[0].stop_reason: Stop reason
+    - usage: Token usage in last chunk (prompt_tokens, completion_tokens, total_tokens)
+    - amazon-bedrock-invocationMetrics: Bedrock metrics
+    """
+
+    def aggregate(self, items: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Aggregate Mistral chunks into native Mistral output with Bedrock usage."""
+        LOGGER.debug("Mistral aggregator processing %d items", len(items))
+
+        content = ""
+        stop_reason = None
+        input_tokens = 0
+        output_tokens = 0
+        model_id = None
+        completion_id = None
+
+        for item in items:
+            if "chunk" not in item:
+                continue
+
+            try:
+                chunk_data = json.loads(item["chunk"]["bytes"])
+
+                # Extract model ID and completion ID from first chunk
+                if model_id is None and "model" in chunk_data:
+                    model_id = chunk_data["model"]
+                if completion_id is None and "id" in chunk_data:
+                    completion_id = chunk_data["id"]
+
+                # Extract content from choices
+                if "choices" in chunk_data and chunk_data["choices"]:
+                    choice = chunk_data["choices"][0]
+
+                    # Extract message content
+                    if "message" in choice and choice["message"]:
+                        message_content = choice["message"].get("content")
+                        if message_content:
+                            content += message_content
+
+                    # Extract stop reason
+                    if "stop_reason" in choice and choice["stop_reason"]:
+                        stop_reason = choice["stop_reason"]
+                        LOGGER.debug("Mistral stop_reason: %s", stop_reason)
+
+                # Extract usage from last chunk
+                if "usage" in chunk_data and chunk_data["usage"]:
+                    usage = chunk_data["usage"]
+                    if "prompt_tokens" in usage:
+                        input_tokens = usage["prompt_tokens"]
+                    if "completion_tokens" in usage:
+                        output_tokens = usage["completion_tokens"]
+                    LOGGER.debug(
+                        "Mistral usage: prompt=%d, completion=%d",
+                        input_tokens,
+                        output_tokens,
+                    )
+
+                # Use bedrock metrics as authoritative source
+                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(
+                        "Mistral bedrock metrics: input=%d, output=%d",
+                        input_tokens,
+                        output_tokens,
+                    )
+
+            except (json.JSONDecodeError, KeyError, TypeError) as e:
+                LOGGER.debug("Mistral aggregator error processing chunk: %s", e)
+                continue
+
+        LOGGER.debug(
+            "Mistral aggregated: %d chars, input_tokens=%d, output_tokens=%d",
+            len(content),
+            input_tokens,
+            output_tokens,
+        )
+
+        # Convert to Bedrock usage format using shared converter
+        bedrock_usage = usage_converters.openai_to_bedrock_usage(
+            {"prompt_tokens": input_tokens, "completion_tokens": output_tokens}
+        )
+
+        # Return Mistral's native OpenAI-like format with Bedrock usage only
+        return {
+            "id": completion_id,
+            "object": "chat.completion",
+            "model": model_id,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "stop_reason": stop_reason,
+                }
+            ],
+            "usage": bedrock_usage,
+        }