wehub-resource-sync
217 lines
7.7 KiB
Python
217 lines
7.7 KiB
Python
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from botocore.eventstream import EventStream
|
|
|
|
from mlflow.bedrock.utils import (
|
|
capture_exception,
|
|
parse_complete_token_usage_from_response,
|
|
parse_partial_token_usage_from_response,
|
|
)
|
|
from mlflow.entities.span import LiveSpan
|
|
from mlflow.entities.span_event import SpanEvent
|
|
from mlflow.tracing.constant import SpanAttributeKey
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseEventStreamWrapper:
|
|
"""
|
|
A wrapper class for a event stream to record events and accumulated response
|
|
in an MLflow span if possible.
|
|
|
|
A span should be ended when the stream is exhausted rather than when it is created.
|
|
|
|
Args:
|
|
stream: The original event stream to wrap.
|
|
span: The span to record events and response in.
|
|
inputs: The inputs to the converse API.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
stream: EventStream,
|
|
span: LiveSpan,
|
|
inputs: dict[str, Any] | None = None,
|
|
):
|
|
self._stream = stream
|
|
self._span = span
|
|
self._inputs = inputs
|
|
|
|
def __iter__(self):
|
|
for event in self._stream:
|
|
self._handle_event(self._span, event)
|
|
yield event
|
|
|
|
# End the span when the stream is exhausted
|
|
self._close()
|
|
|
|
def __getattr__(self, attr):
|
|
"""Delegate all other attributes to the original stream."""
|
|
return getattr(self._stream, attr)
|
|
|
|
def _handle_event(self, span, event):
|
|
"""Process a single event from the stream."""
|
|
raise NotImplementedError
|
|
|
|
def _close(self):
|
|
"""End the span and run any finalization logic."""
|
|
raise NotImplementedError
|
|
|
|
@capture_exception("Failed to handle event for the stream")
|
|
def _end_span(self):
|
|
"""End the span."""
|
|
self._span.end()
|
|
|
|
|
|
def _extract_token_usage_from_chunk(chunk: dict[str, Any]) -> dict[str, int] | None:
|
|
"""Extract partial token usage from streaming chunk.
|
|
|
|
Args:
|
|
chunk: A single streaming chunk from Bedrock API.
|
|
|
|
Returns:
|
|
Token usage dictionary with standardized keys, or None if no usage found.
|
|
"""
|
|
try:
|
|
usage = (
|
|
chunk.get("message", {}).get("usage")
|
|
if chunk.get("type") == "message_start"
|
|
else chunk.get("usage")
|
|
)
|
|
if isinstance(usage, dict):
|
|
return parse_partial_token_usage_from_response(usage)
|
|
return None
|
|
except (KeyError, TypeError, AttributeError) as e:
|
|
_logger.debug(f"Failed to extract token usage from chunk: {e}")
|
|
return None
|
|
|
|
|
|
class InvokeModelStreamWrapper(BaseEventStreamWrapper):
|
|
"""A wrapper class for a event stream returned by the InvokeModelWithResponseStream API.
|
|
|
|
This wrapper intercepts streaming events from Bedrock's invoke_model_with_response_stream
|
|
API and accumulates token usage information across multiple chunks. It buffers partial
|
|
token usage data as it arrives and sets the final aggregated usage on the span when
|
|
the stream is exhausted.
|
|
|
|
Attributes:
|
|
_usage_buffer (dict): Internal buffer to accumulate token usage data from
|
|
streaming chunks. Uses TokenUsageKey constants as keys.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._usage_buffer = {}
|
|
|
|
def _buffer_token_usage_from_chunk(self, chunk: dict[str, Any]):
|
|
"""Buffer token usage from streaming chunk."""
|
|
if usage_data := _extract_token_usage_from_chunk(chunk):
|
|
for token_key, token_value in usage_data.items():
|
|
self._usage_buffer[token_key] = token_value
|
|
|
|
@capture_exception("Failed to handle event for the stream")
|
|
def _handle_event(self, span, event):
|
|
"""Process streaming event and buffer token usage."""
|
|
chunk = json.loads(event["chunk"]["bytes"])
|
|
self._span.add_event(SpanEvent(name=chunk["type"], attributes={"json": json.dumps(chunk)}))
|
|
|
|
# Buffer usage information from streaming chunks
|
|
self._buffer_token_usage_from_chunk(chunk)
|
|
|
|
def _close(self):
|
|
"""Set accumulated token usage on span and end it."""
|
|
# Build a standardized usage dict from buffered data using the utility function
|
|
if usage_data := parse_complete_token_usage_from_response(self._usage_buffer):
|
|
self._span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage_data)
|
|
|
|
self._end_span()
|
|
|
|
|
|
class ConverseStreamWrapper(BaseEventStreamWrapper):
|
|
"""A wrapper class for a event stream returned by the ConverseStream API."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self._response_builder = _ConverseMessageBuilder()
|
|
|
|
def __getattr__(self, attr):
|
|
"""Delegate all other attributes to the original stream."""
|
|
return getattr(self._stream, attr)
|
|
|
|
@capture_exception("Failed to handle event for the stream")
|
|
def _handle_event(self, span, event):
|
|
"""
|
|
Process a single event from the stream.
|
|
|
|
Refer to the following documentation for the event format:
|
|
https://boto3.amazonaws.com/v1/documentation/api/1.35.8/reference/services/bedrock-runtime/client/converse_stream.html
|
|
"""
|
|
event_name = list(event.keys())[0]
|
|
self._response_builder.process_event(event_name, event[event_name])
|
|
# Record raw event as a span event
|
|
self._span.add_event(
|
|
SpanEvent(name=event_name, attributes={"json": json.dumps(event[event_name])})
|
|
)
|
|
|
|
@capture_exception("Failed to record the accumulated response in the span")
|
|
def _close(self):
|
|
"""Set final response and token usage on span and end it."""
|
|
# Build a standardized usage dict and set it on the span if valid
|
|
converse_response = self._response_builder.build()
|
|
self._span.set_outputs(converse_response)
|
|
|
|
raw_usage_data = converse_response.get("usage")
|
|
if isinstance(raw_usage_data, dict):
|
|
if usage_data := parse_complete_token_usage_from_response(raw_usage_data):
|
|
self._span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage_data)
|
|
|
|
self._end_span()
|
|
|
|
|
|
class _ConverseMessageBuilder:
|
|
"""A helper class to accumulate the chunks of a streaming Converse API response."""
|
|
|
|
def __init__(self):
|
|
self._role = "assistant"
|
|
self._text_content_buffer = ""
|
|
self._tool_use = {}
|
|
self._response = {}
|
|
|
|
def process_event(self, event_name: str, event_attr: dict[str, Any]):
|
|
if event_name == "messageStart":
|
|
self._role = event_attr["role"]
|
|
elif event_name == "contentBlockStart":
|
|
# ContentBlockStart event is only used for tool usage. It carries the tool id
|
|
# and the name, but not the input arguments.
|
|
self._tool_use = {
|
|
# In streaming, input is always string
|
|
"input": "",
|
|
**event_attr["start"]["toolUse"],
|
|
}
|
|
elif event_name == "contentBlockDelta":
|
|
delta = event_attr["delta"]
|
|
if text := delta.get("text"):
|
|
self._text_content_buffer += text
|
|
if tool_use := delta.get("toolUse"):
|
|
self._tool_use["input"] += tool_use["input"]
|
|
elif event_name == "contentBlockStop":
|
|
pass
|
|
elif event_name in {"messageStop", "metadata"}:
|
|
self._response.update(event_attr)
|
|
else:
|
|
_logger.debug(f"Unknown event, skipping: {event_name}")
|
|
|
|
def build(self) -> dict[str, Any]:
|
|
message = {
|
|
"role": self._role,
|
|
"content": [{"text": self._text_content_buffer}],
|
|
}
|
|
if self._tool_use:
|
|
message["content"].append({"toolUse": self._tool_use})
|
|
|
|
self._response.update({"output": {"message": message}})
|
|
|
|
return self._response
|