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

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,
+        }

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

@@ -0,0 +1,99 @@
+"""Amazon Nova 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 NovaAggregator(ChunkAggregator):
+    """
+    Aggregator for Amazon Nova streaming format.
+
+    Returns Nova's native output format with Bedrock usage.
+
+    Nova chunk structure uses camelCase:
+    - messageStart: Contains role information
+    - contentBlockDelta: Contains text in delta.text
+    - contentBlockStop: End of content block
+    - messageStop: Contains stopReason
+    - metadata: Contains usage information with inputTokens/outputTokens
+    """
+
+    def aggregate(self, items: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Aggregate Nova chunks into native Nova output structure with Bedrock usage."""
+        LOGGER.debug("Nova aggregator processing %d items", len(items))
+
+        output_text = ""
+        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"])
+
+                if "contentBlockDelta" in chunk_data:
+                    delta = chunk_data["contentBlockDelta"].get("delta", {})
+                    if "text" in delta:
+                        output_text += delta["text"]
+
+                elif "messageStop" in chunk_data:
+                    stop_data = chunk_data["messageStop"]
+                    if "stopReason" in stop_data:
+                        stop_reason = stop_data["stopReason"]
+                        LOGGER.debug("Nova stop_reason: %s", stop_reason)
+
+                elif "metadata" in chunk_data:
+                    if "usage" in chunk_data["metadata"]:
+                        metadata_usage = chunk_data["metadata"]["usage"]
+                        input_tokens = metadata_usage.get("inputTokens", 0)
+                        output_tokens = metadata_usage.get("outputTokens", 0)
+                        LOGGER.debug(
+                            "Nova metadata usage: input=%d, output=%d",
+                            input_tokens,
+                            output_tokens,
+                        )
+
+                    # Use bedrock invocation 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(
+                            "Nova bedrock metrics: input=%d, output=%d",
+                            input_tokens,
+                            output_tokens,
+                        )
+
+            except (json.JSONDecodeError, KeyError, TypeError) as e:
+                LOGGER.debug("Nova aggregator error processing chunk: %s", e)
+                continue
+
+        LOGGER.debug(
+            "Nova aggregated: %d chars, input_tokens=%d, output_tokens=%d",
+            len(output_text),
+            input_tokens,
+            output_tokens,
+        )
+
+        # Convert to Bedrock usage format using shared converter
+        bedrock_usage = usage_converters.nova_to_bedrock_usage(
+            {"inputTokens": input_tokens, "outputTokens": output_tokens}
+        )
+
+        # Return Nova's native output format with Bedrock usage
+        return {
+            "output": {
+                "message": {"role": "assistant", "content": [{"text": output_text}]}
+            },
+            "stopReason": stop_reason,
+            "usage": bedrock_usage,
+        }

opik/integrations/bedrock/invoke_model/invoke_model_decorator.py

@@ -0,0 +1,178 @@
+import json
+import logging
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing_extensions import override
+
+import opik
+import opik.dict_utils as dict_utils
+from opik import llm_usage
+from opik.api_objects import span
+from opik.decorator import arguments_helpers, base_track_decorator
+
+from .. import types
+from . import stream_wrappers, usage_extraction, response_types
+
+import botocore.response
+import botocore.eventstream
+
+LOGGER = logging.getLogger(__name__)
+
+# Keys to extract from kwargs for input logging
+KWARGS_KEYS_TO_LOG_AS_INPUTS = ["body", "modelId"]
+# Keys to extract from response for output logging
+RESPONSE_KEYS_TO_LOG_AS_OUTPUTS = ["body"]
+
+
+class BedrockInvokeModelDecorator(base_track_decorator.BaseTrackDecorator):
+    """
+    An implementation of BaseTrackDecorator designed specifically for tracking
+    calls of AWS bedrock client `invoke_model` and `invoke_model_with_response_stream` functions.
+
+    Besides special processing for input arguments and response content, it
+    overrides _streams_handler() method to work correctly with bedrock's streams
+    """
+
+    @override
+    def _start_span_inputs_preprocessor(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> arguments_helpers.StartSpanParameters:
+        assert (
+            kwargs is not None
+        ), "Expected kwargs to be not None in BedrockRuntime.Client.invoke_model(**kwargs)"
+
+        name = track_options.name if track_options.name is not None else func.__name__
+        body_dict = json.loads(kwargs.get("body", "{}"))
+
+        kwargs_copy = kwargs.copy()
+        kwargs_copy["body"] = body_dict
+
+        input_data, metadata = dict_utils.split_dict_by_keys(
+            kwargs_copy, KWARGS_KEYS_TO_LOG_AS_INPUTS
+        )
+
+        metadata["created_from"] = "bedrock"
+        tags = ["bedrock", "invoke_model"]
+
+        result = arguments_helpers.StartSpanParameters(
+            name=name,
+            input=input_data,
+            type=track_options.type,
+            tags=tags,
+            metadata=metadata,
+            project_name=track_options.project_name,
+            model=kwargs.get("modelId", None),
+            provider=opik.LLMProvider.BEDROCK,
+        )
+
+        return result
+
+    @override
+    def _end_span_inputs_preprocessor(
+        self,
+        output: Any,
+        capture_output: bool,
+        current_span_data: span.SpanData,
+    ) -> arguments_helpers.EndSpanParameters:
+        # Check if this is a structured aggregated response dataclass
+        if isinstance(output, response_types.BedrockAggregatedResponse):
+            # This is a structured aggregated streaming response
+            opik_usage = llm_usage.build_opik_usage(
+                provider=opik.LLMProvider.BEDROCK, usage=output.usage
+            )
+
+            result = arguments_helpers.EndSpanParameters(
+                output=output.to_output_format(),  # Native format in body
+                provider=opik.LLMProvider.BEDROCK,
+                usage=opik_usage,
+                metadata=output.to_metadata_format(),
+            )
+        else:
+            # Regular non-streaming response (dict)
+            output, metadata = dict_utils.split_dict_by_keys(
+                output, RESPONSE_KEYS_TO_LOG_AS_OUTPUTS
+            )
+            subprovider = usage_extraction.extract_subprovider_from_model_id(
+                cast(str, current_span_data.model)
+            )
+            opik_usage = usage_extraction.try_extract_usage_from_bedrock_response(  # type: ignore
+                subprovider, output
+            )
+
+            result = arguments_helpers.EndSpanParameters(
+                output=output,
+                provider=opik.LLMProvider.BEDROCK,
+                usage=opik_usage,
+                metadata=metadata,
+            )
+
+        return result
+
+    @override
+    def _streams_handler(  # type: ignore
+        self,
+        output: Any,
+        capture_output: bool,
+        generations_aggregator: Optional[Callable[[List[Any]], Any]],
+    ) -> Union[
+        types.InvokeModelOutput,
+        None,
+    ]:
+        # Despite the name, StreamingBody is not a stream in traditional LLM provider sense (response chunks).
+        # It's an interface to a stream of bytes representing the response body.
+        streaming_body_detected = (
+            isinstance(output, dict)
+            and "body" in output
+            and isinstance(output["body"], botocore.response.StreamingBody)
+        )
+
+        if streaming_body_detected:
+            span_to_end, trace_to_end = base_track_decorator.pop_end_candidates()
+            return stream_wrappers.wrap_invoke_model_response(
+                output=output,
+                span_to_end=span_to_end,
+                trace_to_end=trace_to_end,
+                finally_callback=self._after_call,
+            )
+
+        DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR = (
+            generations_aggregator is None
+        )
+
+        if DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR:
+            return None
+
+        generations_aggregator = cast(
+            Callable[[List[Any]], Any], generations_aggregator
+        )
+
+        event_streaming_body_detected = (
+            isinstance(output, dict)
+            and "body" in output
+            and isinstance(output["body"], botocore.eventstream.EventStream)
+        )
+
+        if event_streaming_body_detected:
+            span_to_end, trace_to_end = base_track_decorator.pop_end_candidates()
+
+            wrapped_stream = (
+                stream_wrappers.wrap_invoke_model_with_response_stream_response(
+                    stream=output["body"],
+                    capture_output=capture_output,
+                    span_to_end=span_to_end,
+                    trace_to_end=trace_to_end,
+                    generations_aggregator=generations_aggregator,
+                    response_metadata=output["ResponseMetadata"],
+                    finally_callback=self._after_call,
+                )
+            )
+
+            output["body"] = wrapped_stream
+            return cast(types.InvokeModelWithResponseStreamOutput, output)
+
+        STREAM_NOT_FOUND = None
+
+        return STREAM_NOT_FOUND

opik/integrations/bedrock/invoke_model/response_types.py

@@ -0,0 +1,34 @@
+"""Response types for Bedrock invoke_model operations."""
+
+from dataclasses import dataclass
+from typing import Any, Dict
+
+
+@dataclass
+class BedrockAggregatedResponse:
+    """
+    Response from invoke_model_with_response_stream after chunk aggregation.
+
+    Contains the aggregated response in the provider's native format
+    with standardized Bedrock usage format.
+    """
+
+    # Provider-native response structure
+    native_response: Dict[str, Any]
+
+    # Standardized Bedrock usage format
+    usage: Dict[str, Any]  # {inputTokens, outputTokens, totalTokens}
+
+    # Response metadata from Bedrock
+    response_metadata: Dict[str, Any]
+
+    def to_output_format(self) -> Dict[str, Any]:
+        """Convert to output format for span logging."""
+        return {"body": self.native_response}
+
+    def to_metadata_format(self) -> Dict[str, Any]:
+        """Convert to metadata format for span logging."""
+        return {
+            "created_from": "bedrock",
+            **self.response_metadata,
+        }

opik/integrations/bedrock/invoke_model/stream_wrappers.py

@@ -0,0 +1,122 @@
+import logging
+import json
+from typing import Optional, Callable, List, Any, Dict, Generator
+
+import botocore.response
+import functools
+
+import opik.api_objects.span as span
+import opik.api_objects.trace as trace
+from opik.types import ErrorInfoDict
+from opik.decorator import generator_wrappers, error_info_collector
+from .. import types
+
+import botocore.eventstream
+
+LOGGER = logging.getLogger(__name__)
+
+
+__original_streaming_body_read = botocore.response.StreamingBody.read
+
+
+def wrap_invoke_model_response(
+    output: types.InvokeModelOutput,
+    span_to_end: span.SpanData,
+    trace_to_end: Optional[trace.TraceData],
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> types.InvokeModelOutput:
+    response_metadata = output["ResponseMetadata"]
+    streaming_body = output["body"]
+
+    @functools.wraps(__original_streaming_body_read)
+    def wrapped_read(self: botocore.response.StreamingBody, *args, **kwargs):  # type: ignore
+        error_info: Optional[ErrorInfoDict] = None
+        result = None
+        try:
+            result = __original_streaming_body_read(self, *args, **kwargs)
+            return result
+        except Exception as exception:
+            LOGGER.debug(
+                "Exception raised from botocore.response.StreamingBody: %s",
+                str(exception),
+                exc_info=True,
+            )
+            error_info = error_info_collector.collect(exception)
+            raise exception
+        finally:
+            if not hasattr(self, "opik_tracked_instance"):
+                return None
+
+            delattr(self, "opik_tracked_instance")
+
+            if error_info is None and result is not None:
+                try:
+                    parsed_body = json.loads(result)
+                    output = {
+                        "body": parsed_body,
+                        "ResponseMetadata": response_metadata,
+                    }
+                    LOGGER.debug(
+                        "Successfully parsed response body with keys: %s",
+                        list(parsed_body.keys()),
+                    )
+                except (json.JSONDecodeError, TypeError) as e:
+                    LOGGER.debug("Failed to parse response body as JSON: %s", e)
+                    output = {"body": {}, "ResponseMetadata": response_metadata}
+            else:
+                LOGGER.debug("Error occurred or result is None, using empty body")
+                output = {"body": {}, "ResponseMetadata": response_metadata}
+
+            finally_callback(
+                output=output,
+                error_info=error_info,
+                generators_span_to_end=span_to_end,
+                generators_trace_to_end=trace_to_end,
+                capture_output=True,
+            )
+
+    botocore.response.StreamingBody.read = wrapped_read
+    streaming_body.opik_tracked_instance = True
+
+    return output
+
+
+def wrap_invoke_model_with_response_stream_response(
+    stream: botocore.eventstream.EventStream,
+    capture_output: bool,
+    span_to_end: span.SpanData,
+    trace_to_end: Optional[trace.TraceData],
+    generations_aggregator: Callable[[List[Any]], Any],
+    response_metadata: Dict[str, Any],
+    finally_callback: generator_wrappers.FinishGeneratorCallback,
+) -> Generator[Any, None, None]:
+    items: List[Dict[str, Any]] = []
+    error_info: Optional[ErrorInfoDict] = None
+
+    try:
+        for item in stream:
+            items.append(item)
+
+            yield item
+    except Exception as exception:
+        LOGGER.debug(
+            "Exception raised from botocore.eventstream.EventStream: %s",
+            str(exception),
+            exc_info=True,
+        )
+        error_info = error_info_collector.collect(exception)
+        raise exception
+    finally:
+        if error_info is None:
+            output = generations_aggregator(items)
+            output.response_metadata = response_metadata
+        else:
+            output = None
+
+        finally_callback(
+            output=output,
+            error_info=error_info,
+            generators_span_to_end=span_to_end,
+            generators_trace_to_end=trace_to_end,
+            capture_output=capture_output,
+        )