genesis-flow 1.0.0__py3-none-any.whl

mlflow/azure/utils.py

@@ -0,0 +1,183 @@
+"""Utility functions for Azure authentication."""
+
+from __future__ import annotations
+
+import time
+import logging
+from datetime import datetime, timedelta
+from typing import Any, Optional
+from urllib.parse import parse_qs, urlparse
+
+logger = logging.getLogger(__name__)
+
+
+def is_token_expired(token_expiry: Optional[datetime], buffer_seconds: int = 300) -> bool:
+    """Check if a token is expired or will expire within the buffer time.
+
+    Args:
+        token_expiry: Token expiration datetime
+        buffer_seconds: Buffer time in seconds before actual expiry
+
+    Returns:
+        True if token is expired or will expire soon
+    """
+    if not token_expiry:
+        return True
+
+    buffer_time = timedelta(seconds=buffer_seconds)
+    return datetime.utcnow() + buffer_time >= token_expiry
+
+
+def parse_connection_string(connection_string: str) -> dict[str, Any]:
+    """Parse a PostgreSQL connection string and extract components.
+
+    Args:
+        connection_string: PostgreSQL connection string
+
+    Returns:
+        Dictionary with connection components
+    """
+    # Handle azure-postgres:// scheme by converting to postgresql:// for parsing
+    original_scheme = connection_string.split("://")[0] if "://" in connection_string else "unknown"
+    if connection_string.startswith("azure-postgres://"):
+        connection_string = connection_string.replace("azure-postgres://", "postgresql://", 1)
+
+    parsed = urlparse(connection_string)
+
+    # Parse query parameters
+    query_params = {}
+    if parsed.query:
+        query_params = {k: v[0] if len(v) == 1 else v for k, v in parse_qs(parsed.query).items()}
+
+    return {
+        "scheme": original_scheme,  # Use original scheme (azure-postgres or postgresql)
+        "username": parsed.username,
+        "password": parsed.password,
+        "hostname": parsed.hostname,
+        "port": parsed.port or 5432,
+        "database": parsed.path.lstrip("/") if parsed.path else "",
+        "params": query_params,
+    }
+
+
+def build_connection_string(
+    username: str,
+    password: str,
+    hostname: str,
+    port: int = 5432,
+    database: str = "postgres",
+    **params: Any,
+) -> str:
+    """Build a PostgreSQL connection string.
+
+    Args:
+        username: Database username
+        password: Database password
+        hostname: Database hostname
+        port: Database port
+        database: Database name
+        **params: Additional connection parameters
+
+    Returns:
+        PostgreSQL connection string
+    """
+    param_string = ""
+    if params:
+        param_string = "?" + "&".join(f"{k}={v}" for k, v in params.items())
+
+    return f"postgresql://{username}:{password}@{hostname}:{port}/{database}{param_string}"
+
+
+def sanitize_connection_string_for_logging(connection_string: str) -> str:
+    """Sanitize connection string for safe logging by removing sensitive information.
+
+    Args:
+        connection_string: Original connection string
+
+    Returns:
+        Sanitized connection string with password masked
+    """
+    try:
+        parsed = parse_connection_string(connection_string)
+
+        # Mask password
+        masked_password = "***" if parsed.get("password") else None
+
+        # Rebuild without sensitive info
+        safe_parts = []
+        safe_parts.append(f"scheme={parsed.get('scheme', 'unknown')}")
+        safe_parts.append(f"host={parsed.get('hostname', 'unknown')}")
+        safe_parts.append(f"port={parsed.get('port', 5432)}")
+        safe_parts.append(f"database={parsed.get('database', 'unknown')}")
+        safe_parts.append(f"user={parsed.get('username', 'unknown')}")
+
+        if masked_password:
+            safe_parts.append("password=***")
+
+        return " | ".join(safe_parts)
+
+    except Exception as e:
+        logger.warning("Failed to sanitize connection string", extra={"error": str(e)})
+        return "connection_string=<unable_to_parse>"
+
+
+def exponential_backoff(attempt: int, base_delay: float = 1.0, max_delay: float = 60.0) -> float:
+    """Calculate exponential backoff delay.
+
+    Args:
+        attempt: Current attempt number (0-indexed)
+        base_delay: Base delay in seconds
+        max_delay: Maximum delay in seconds
+
+    Returns:
+        Delay in seconds
+    """
+    delay = base_delay * (2**attempt)
+    return min(delay, max_delay)
+
+
+def retry_on_exception(
+    func: callable,
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    exceptions: tuple = (Exception,),
+) -> Any:
+    """Retry a function with exponential backoff.
+
+    Args:
+        func: Function to retry
+        max_retries: Maximum number of retries
+        base_delay: Base delay between retries
+        exceptions: Exception types to catch and retry on
+
+    Returns:
+        Function result
+
+    Raises:
+        Last exception if all retries failed
+    """
+    last_exception = None
+
+    for attempt in range(max_retries + 1):
+        try:
+            return func()
+        except exceptions as e:
+            last_exception = e
+
+            if attempt == max_retries:
+                logger.error(
+                    "Function failed after all retries",
+                    extra={"attempt": attempt, "max_retries": max_retries, "error": str(e)},
+                )
+                raise
+
+            delay = exponential_backoff(attempt, base_delay)
+            logger.warning(
+                "Function failed, retrying",
+                extra={"attempt": attempt, "delay": delay, "error": str(e)},
+            )
+
+            time.sleep(delay)
+
+    # This should never be reached
+    raise last_exception or Exception("Unexpected error in retry logic")

mlflow/bedrock/__init__.py

@@ -0,0 +1,45 @@
+import logging
+
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+_logger = logging.getLogger(__name__)
+
+FLAVOR_NAME = "bedrock"
+
+
+@experimental(version="2.20.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Amazon Bedrock to MLflow.
+    Only synchronous calls are supported. Asynchronous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Bedrock models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Bedrock autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Bedrock
+            autologging. If ``False``, show all events and warnings.
+    """
+    from botocore.client import ClientCreator
+
+    from mlflow.bedrock._autolog import patched_create_client
+
+    # NB: In boto3, the client class for each service is dynamically created at
+    # runtime via the ClientCreator factory class. Therefore, we cannot patch
+    # the service client directly, and instead patch the factory to return
+    # a patched client class.
+    safe_patch(FLAVOR_NAME, ClientCreator, "create_client", patched_create_client)
+
+    # Since we patch the ClientCreator factory, it only takes effect for new client instances.
+    if log_traces:
+        _logger.info(
+            "Enabled auto-tracing for Bedrock. Note that MLflow can only trace boto3 "
+            "service clients that are created after this call. If you have already "
+            "created one, please recreate the client by calling `boto3.client`."
+        )

mlflow/bedrock/_autolog.py

@@ -0,0 +1,202 @@
+import io
+import json
+import logging
+from typing import Any, Optional, Union
+
+from botocore.client import BaseClient
+from botocore.response import StreamingBody
+
+import mlflow
+from mlflow.bedrock import FLAVOR_NAME
+from mlflow.bedrock.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.bedrock.stream import ConverseStreamWrapper, InvokeModelStreamWrapper
+from mlflow.bedrock.utils import skip_if_trace_disabled
+from mlflow.entities import SpanType
+from mlflow.tracing.fluent import start_span_no_context
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.utils.autologging_utils import safe_patch
+
+_BEDROCK_RUNTIME_SERVICE_NAME = "bedrock-runtime"
+_BEDROCK_SPAN_PREFIX = "BedrockRuntime."
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_create_client(original, self, *args, **kwargs):
+    """
+    Patched version of the boto3 ClientCreator.create_client method that returns
+    a patched client class.
+    """
+    if kwargs.get("service_name") != _BEDROCK_RUNTIME_SERVICE_NAME:
+        return original(self, *args, **kwargs)
+
+    client = original(self, *args, **kwargs)
+    patch_bedrock_runtime_client(client.__class__)
+
+    return client
+
+
+def patch_bedrock_runtime_client(client_class: type[BaseClient]):
+    """
+    Patch the BedrockRuntime client to log traces and models.
+    """
+    # The most basic model invocation API
+    safe_patch(FLAVOR_NAME, client_class, "invoke_model", _patched_invoke_model)
+    safe_patch(
+        FLAVOR_NAME,
+        client_class,
+        "invoke_model_with_response_stream",
+        _patched_invoke_model_with_response_stream,
+    )
+
+    if hasattr(client_class, "converse"):
+        # The new "converse" API was introduced in boto3 1.35 to access all models
+        # with the consistent chat format.
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html
+        safe_patch(FLAVOR_NAME, client_class, "converse", _patched_converse)
+
+    if hasattr(client_class, "converse_stream"):
+        safe_patch(FLAVOR_NAME, client_class, "converse_stream", _patched_converse_stream)
+
+
+@skip_if_trace_disabled
+def _patched_invoke_model(original, self, *args, **kwargs):
+    with mlflow.start_span(name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}") as span:
+        # NB: Bedrock client doesn't accept any positional arguments
+        span.set_inputs(kwargs)
+
+        result = original(self, *args, **kwargs)
+
+        result["body"] = _buffer_stream(result["body"])
+        parsed_response_body = _parse_invoke_model_response_body(result["body"])
+
+        # Determine the span type based on the key in the response body.
+        # As of 2024 Dec 9th, all supported embedding models in Bedrock returns the response body
+        # with the key "embedding". This might change in the future.
+        span_type = SpanType.EMBEDDING if "embedding" in parsed_response_body else SpanType.LLM
+        span.set_span_type(span_type)
+        span.set_outputs({**result, "body": parsed_response_body})
+
+        return result
+
+
+@skip_if_trace_disabled
+def _patched_invoke_model_with_response_stream(original, self, *args, **kwargs):
+    span = start_span_no_context(
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        # NB: Since we don't inspect the response body for this method, the span type is unknown.
+        # We assume it is LLM as using streaming for embedding is not common.
+        span_type=SpanType.LLM,
+        inputs=kwargs,
+    )
+
+    result = original(self, *args, **kwargs)
+
+    # To avoid consuming the stream during serialization, set dummy outputs for the span.
+    span.set_outputs({**result, "body": "EventStream"})
+
+    result["body"] = InvokeModelStreamWrapper(stream=result["body"], span=span)
+    return result
+
+
+def _buffer_stream(raw_stream: StreamingBody) -> StreamingBody:
+    """
+    Create a buffered stream from the raw byte stream.
+
+    The boto3's invoke_model() API returns the LLM response as a byte stream.
+    We need to read the stream data to set the span outputs, however, the stream
+    can only be read once and not seekable (https://github.com/boto/boto3/issues/564).
+    To work around this, we create a buffered stream that can be read multiple times.
+    """
+    buffered_response = io.BytesIO(raw_stream.read())
+    buffered_response.seek(0)
+    return StreamingBody(buffered_response, raw_stream._content_length)
+
+
+def _parse_invoke_model_response_body(response_body: StreamingBody) -> Union[dict[str, Any], str]:
+    content = response_body.read()
+    try:
+        return json.loads(content)
+    except Exception:
+        # When failed to parse the response body as JSON, return the raw response
+        return content
+    finally:
+        # Reset the stream position to the beginning
+        response_body._raw_stream.seek(0)
+        # Boto3 uses this attribute to validate the amount of data read from the stream matches
+        # the content length, so we need to reset it as well.
+        # https://github.com/boto/botocore/blob/f88e981cb1a6cd0c64bc89da262ab76f9bfa9b7d/botocore/response.py#L164C17-L164C32
+        response_body._amount_read = 0
+
+
+@skip_if_trace_disabled
+def _patched_converse(original, self, *args, **kwargs):
+    with mlflow.start_span(
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        span_type=SpanType.CHAT_MODEL,
+    ) as span:
+        # NB: Bedrock client doesn't accept any positional arguments
+        span.set_inputs(kwargs)
+        _set_tool_attributes(span, kwargs)
+
+        result = None
+        try:
+            result = original(self, *args, **kwargs)
+            span.set_outputs(result)
+        finally:
+            _set_chat_messages_attributes(span, kwargs.get("messages", []), result)
+        return result
+
+
+@skip_if_trace_disabled
+def _patched_converse_stream(original, self, *args, **kwargs):
+    # NB: Do not use fluent API to create a span for streaming response. If we do so,
+    # the span context will remain active until the stream is fully exhausted, which
+    # can lead to super hard-to-debug issues.
+    span = start_span_no_context(
+        name=f"{_BEDROCK_SPAN_PREFIX}{original.__name__}",
+        span_type=SpanType.CHAT_MODEL,
+        inputs=kwargs,
+    )
+    _set_tool_attributes(span, kwargs)
+
+    result = original(self, *args, **kwargs)
+
+    if span:
+        result["stream"] = ConverseStreamWrapper(
+            stream=result["stream"],
+            span=span,
+            inputs=kwargs,
+        )
+
+    return result
+
+
+def _set_chat_messages_attributes(
+    span, messages: list[dict[str, Any]], response: Optional[dict[str, Any]]
+):
+    """
+    Extract standard chat span attributes for the Bedrock Converse API call.
+
+    NB: We only support standard attribute extraction for the Converse API, because
+    the InvokeModel API exposes the raw API spec from each LLM provider, hence
+    maintaining the compatibility for all providers is significantly cumbersome.
+    """
+    try:
+        messages = [*messages]  # shallow copy to avoid appending to the original list
+        if response:
+            messages.append(response["output"]["message"])
+        messages = [convert_message_to_mlflow_chat(msg) for msg in messages]
+        set_span_chat_messages(span, messages)
+    except Exception as e:
+        _logger.debug(f"Failed to set messages for {span}. Error: {e}")
+
+
+def _set_tool_attributes(span, kwargs):
+    """Extract tool attributes for the Bedrock Converse API call."""
+    if tool_config := kwargs.get("toolConfig"):
+        try:
+            tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tool_config["tools"]]
+            set_span_chat_tools(span, tools)
+        except Exception as e:
+            _logger.debug(f"Failed to set tools for {span}. Error: {e}")

mlflow/bedrock/chat.py

@@ -0,0 +1,122 @@
+import base64
+import json
+import logging
+from typing import Any, Optional, Union
+
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def convert_message_to_mlflow_chat(message: dict[str, Any]) -> ChatMessage:
+    """
+    Convert Bedrock Converse API's message object into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Message.html
+
+    Args:
+        message: Bedrock Converse API's message object.
+
+    Returns:
+        ChatMessage: MLflow's standard chat message object.
+    """
+    role = message["role"]
+    contents = []
+    tool_calls = []
+    tool_call_id = None
+    for content in message["content"]:
+        if tool_call := content.get("toolUse"):
+            input = tool_call.get("input")
+            tool_calls.append(
+                ToolCall(
+                    id=tool_call["toolUseId"],
+                    function=Function(
+                        name=tool_call["name"],
+                        arguments=input if isinstance(input, str) else json.dumps(input),
+                    ),
+                    type="function",
+                )
+            )
+        elif tool_result := content.get("toolResult"):
+            tool_call_id = tool_result["toolUseId"]
+            # "tool_result" content corresponds to the "tool" message in OpenAI.
+            # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolResultContentBlock.html
+            role = "tool"
+            for content in tool_result["content"]:
+                parsed_content = _parse_content(content)
+                if parsed_content:
+                    contents.append(parsed_content)
+        else:
+            parsed_content = _parse_content(content)
+            if parsed_content:
+                contents.append(parsed_content)
+
+    message = ChatMessage(role=role, content=contents)
+    if tool_calls:
+        message.tool_calls = tool_calls
+    if tool_call_id:
+        message.tool_call_id = tool_call_id
+    return message
+
+
+def _parse_content(content: dict[str, Any]) -> Optional[Union[TextContentPart, ImageContentPart]]:
+    """
+    Parse a single content block in the Bedrock message object.
+
+    Some content types like video and document are not supported by OpenAI's spec. This
+    function returns None for those content types.
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlock.html
+    """
+    if text := content.get("text"):
+        return TextContentPart(text=text, type="text")
+    elif json_content := content.get("json"):
+        return TextContentPart(text=json.dumps(json_content), type="text")
+    elif image := content.get("image"):
+        # Bedrock support passing images in both raw bytes and base64 encoded strings.
+        # OpenAI spec only supports base64 encoded images, so we encode the raw bytes to base64.
+        # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html
+        image_bytes = image["source"]["bytes"]
+        if isinstance(image_bytes, bytes):
+            data = base64.b64encode(image_bytes).decode("utf-8")
+        else:
+            data = image_bytes
+        format = "image/" + image["format"]
+        image_url = ImageUrl(url=f"data:{format};base64,{data}", detail="auto")
+        return ImageContentPart(type="image_url", image_url=image_url)
+    # NB: Video and Document content type are not supported by OpenAI's spec, so recording as text.
+    else:
+        _logger.debug(f"Received an unsupported content type: {list(content.keys())[0]}")
+        return None
+
+
+def convert_tool_to_mlflow_chat_tool(tool: dict[str, Any]) -> ChatTool:
+    """
+    Convert Bedrock tool definition into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Tool.html
+
+    Args:
+        tool: A dictionary represents a single tool definition in the input request.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    tool_spec = tool["toolSpec"]
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=tool_spec["name"],
+            description=tool_spec.get("description"),
+            parameters=tool_spec["inputSchema"].get("json"),
+        ),
+    )