genesis-flow 1.0.0__py3-none-any.whl

mlflow/openai/utils/chat_schema.py

@@ -0,0 +1,367 @@
+import json
+import logging
+from collections.abc import Iterable
+from typing import Any, Optional, Union
+
+from pydantic import BaseModel
+
+from mlflow.entities.span import LiveSpan
+from mlflow.exceptions import MlflowException
+from mlflow.tracing import set_span_chat_messages, set_span_chat_tools
+from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    ContentType,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_RESPONSE_API_BUILT_IN_TOOLS = {
+    "file_search",
+    "computer_use_preview",
+    "web_search_preview",
+}
+
+
+def set_span_chat_attributes(span: LiveSpan, inputs: dict[str, Any], output: Any):
+    # NB: This function is also used for setting chat attributes for ResponsesAgent tracing spans
+    # (TODO: Add doc link). Therefore, the core logic should still run without openai package.
+    messages = _parse_inputs_output(inputs, output)
+    try:
+        set_span_chat_messages(span, messages)
+    except MlflowException:
+        _logger.debug(
+            "Failed to set chat messages on span",
+            exc_info=True,
+        )
+
+    if tools := _parse_tools(inputs):
+        try:
+            set_span_chat_tools(span, tools)
+        except MlflowException:
+            _logger.debug("Failed to set chat tools on span", exc_info=True)
+
+    # Extract and set usage information if available
+    if usage := _parse_usage(output):
+        span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage)
+
+
+def _parse_inputs_output(inputs: dict[str, Any], output: Any) -> list[ChatMessage]:
+    if _is_responses_output(output):
+        return _parse_responses_inputs_outputs(inputs, output)
+
+    messages = []
+    if "messages" in inputs:
+        messages.extend(inputs["messages"])
+
+    try:
+        from openai.types.chat import ChatCompletion
+
+        if isinstance(output, ChatCompletion):
+            messages.append(output.choices[0].message.to_dict(exclude_unset=True))
+            return messages
+    except ImportError:
+        pass
+
+    if isinstance(output, str):
+        messages.append({"role": "assistant", "content": output})
+
+    return messages
+
+
+def _is_responses_output(output: Any) -> bool:
+    """
+    Check whether the output is OpenAI Responses API output, or
+    a response from the MLflow ResponsesAgent instance.
+    """
+    try:
+        from openai.types.responses import Response
+
+        if isinstance(output, Response):
+            return True
+    except ImportError:
+        pass
+
+    try:
+        from mlflow.types.responses import ResponsesAgentResponse
+
+        if ResponsesAgentResponse.validate_compat(output):
+            return True
+    except Exception:
+        pass
+
+    return False
+
+
+def _parse_responses_inputs_outputs(
+    inputs: dict[str, Any], output: Union[BaseModel, dict[str, Any]]
+) -> list[ChatMessage]:
+    messages = []
+    if _input := inputs.get("input"):
+        if isinstance(_input, str):
+            messages.append(ChatMessage(role="user", content=_input))
+        elif isinstance(_input, list):
+            for item in _input:
+                messages.extend(_parse_response_item(item, messages))
+
+    output = output if isinstance(output, dict) else output.model_dump()
+    for output_item in output["output"]:
+        messages.extend(_parse_response_item(output_item, messages))
+
+    return messages
+
+
+def _parse_response_item(
+    item: Union[dict[str, Any], BaseModel],
+    past_messages: list[ChatMessage],
+) -> list[ChatMessage]:
+    """Parse Response API output into MLflow standard chat messages"""
+    if isinstance(item, BaseModel):
+        item = item.model_dump()
+
+    item_type = item.get("type", "message")
+    if item_type == "message":
+        content, refusal = _parse_message_content(item["content"], past_messages)
+        message = ChatMessage(role=item["role"])
+        if content:
+            message.content = content
+        if refusal:
+            message.refusal = refusal
+        return [message]
+
+    elif item_type == "function_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["call_id"], tool_name=item["name"], arguments=item["arguments"]
+            )
+        ]
+
+    elif item_type == "function_call_output":
+        return [
+            ChatMessage(
+                role="tool",
+                content=item["output"],
+                tool_call_id=item["call_id"],
+            )
+        ]
+
+    elif item_type == "file_search_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["id"],
+                tool_name=item["type"],
+                arguments=json.dumps({"queries": item["queries"]}),
+            ),
+            ChatMessage(role="tool", tool_call_id=item.get("id"), content=item_type),
+        ]
+
+    elif item_type == "web_search_call":
+        return [
+            _get_tool_call_message(tool_id=item["id"], tool_name=item["type"], arguments=""),
+            ChatMessage(role="tool", tool_call_id=item.get("id"), content=item_type),
+        ]
+
+    elif item_type == "computer_call":
+        return [
+            _get_tool_call_message(
+                tool_id=item["call_id"],
+                tool_name=item["type"],
+                arguments=json.dumps({"action": item["action"]}),
+            )
+        ]
+    elif item_type == "computer_call_output":
+        output = item["output"]
+        return [
+            ChatMessage(
+                role="tool",
+                content=[
+                    # Screenshot of the computer after taking the action
+                    ImageContentPart(
+                        image_url=ImageUrl(url=output.get("image_url")),
+                        type="image_url",
+                    ),
+                ],
+                tool_call_id=item["call_id"],
+            )
+        ]
+
+    elif item_type == "reasoning":
+        summary = item["summary"][0]["text"] if item["summary"] else None
+        return [ChatMessage(role="assistant", content=summary)]
+
+    raise MlflowException(f"Unknown output type: {type(item)}")
+
+
+def _parse_message_content(
+    content: Union[str, list[dict[str, Any]]], past_messages: Optional[list[ChatMessage]] = None
+) -> tuple[ContentType, Optional[str]]:
+    if isinstance(content, str):
+        return content, None
+
+    if not isinstance(content, list):
+        raise MlflowException(f"Invalid content type: {type(content)}")
+
+    parsed_contents = []
+    refusal = None
+    for item in content:
+        content_type = item.get("type")
+        if content_type == "input_text":
+            parsed_contents.append(TextContentPart(text=item.get("text"), type="text"))
+
+        elif content_type == "input_image":
+            parsed_contents.append(
+                ImageContentPart(
+                    image_url=ImageUrl(
+                        url=item["image_url"],
+                        detail=item.get("detail"),
+                    ),
+                    type="image_url",
+                )
+            )
+
+        elif content_type == "input_file":
+            # TODO: MLflow chat schema doesn't support file content yet. Even if it does,
+            #   including the full file data in an attribute is not a good idea.
+            parsed_contents.append(
+                TextContentPart(
+                    text=f"{item.get('file_id')}:{item.get('file_name')}",
+                    type="text",
+                )
+            )
+
+        elif content_type == "output_text":
+            if annotations := item.get("annotations"):
+                _populate_tool_result_message(annotations, past_messages)
+            parsed_contents.append(TextContentPart(text=item.get("text"), type="text"))
+
+        elif content_type == "refusal":
+            refusal = item.get("refusal")
+
+        else:
+            raise MlflowException(f"Unknown content type: {content_type}")
+
+    return parsed_contents, refusal
+
+
+def _get_tool_call_message(tool_id: str, tool_name: str, arguments: str) -> ChatMessage:
+    return ChatMessage(
+        role="assistant",
+        tool_calls=[
+            ToolCall(
+                id=tool_id,
+                type="function",
+                function=Function(name=tool_name, arguments=arguments),
+            )
+        ],
+    )
+
+
+def _populate_tool_result_message(
+    annotations: dict[str, Any], messages: list[ChatMessage]
+) -> ChatMessage:
+    """
+    Parses annotations from the Response API output into MLflow standard chat spec.
+
+    In OpenAI spec, annotations are used for populating information from file/web
+    search tools in the Responses API response. When converting to MLflow standard
+    chat spec, it should be presented as a tool result message.
+    """
+    file_citations = [a for a in annotations if a.get("type") == "file_citation"]
+    web_citations = [a for a in annotations if a.get("type") == "url_citation"]
+
+    if file_search_result_msg := _find_tool_message(messages, "file_search_call"):
+        file_search_result_msg.content = json.dumps(file_citations)
+
+    if web_search_result_msg := _find_tool_message(messages, "web_search_call"):
+        web_search_result_msg.content = json.dumps(web_citations)
+
+
+def _find_tool_message(messages, tool_type):
+    return next((msg for msg in messages if msg.role == "tool" and msg.content == tool_type), None)
+
+
+def _parse_tools(inputs: dict[str, Any]) -> list[ChatTool]:
+    tools = inputs.get("tools", [])
+
+    if tools is None or not isinstance(tools, Iterable):
+        return []
+
+    parsed_tools = []
+    for tool in tools:
+        tool_type = tool.get("type", "function")
+        if tool_type == "function":
+            if "function" in tool:
+                # ChatCompletion API style
+                parsed_tools.append(ChatTool(**tool))
+            else:
+                # Responses API style
+                definition = {k: v for k, v in tool.items() if k != "type"}
+                parsed_tools.append(
+                    ChatTool(
+                        type="function",
+                        function=FunctionToolDefinition(**definition),
+                    )
+                )
+        elif tool_type in _RESPONSE_API_BUILT_IN_TOOLS:
+            parsed_tools.append(
+                ChatTool(
+                    type="function",
+                    function=FunctionToolDefinition(
+                        name=tool_type,
+                    ),
+                )
+            )
+        else:
+            raise MlflowException(f"Unknown tool type: {tool_type}")
+
+    return parsed_tools
+
+
+def _parse_usage(output: Any) -> Optional[dict[str, Any]]:
+    """
+    Parse token usage information from OpenAI response objects.
+
+    Args:
+        output: The response object from OpenAI API calls
+
+    Returns:
+        A dictionary containing token usage information.
+    """
+    if output is None:
+        return None
+
+    # Handle OpenAI ChatCompletion API response
+    try:
+        from openai.types.chat import ChatCompletion
+
+        if isinstance(output, ChatCompletion) and (usage := output.usage):
+            return {
+                TokenUsageKey.INPUT_TOKENS: usage.prompt_tokens,
+                TokenUsageKey.OUTPUT_TOKENS: usage.completion_tokens,
+                TokenUsageKey.TOTAL_TOKENS: usage.total_tokens,
+            }
+    except ImportError:
+        pass
+
+    # Handle OpenAI Responses API response
+    try:
+        from openai.types.responses import Response
+
+        if isinstance(output, Response) and (usage := output.usage):
+            return {
+                TokenUsageKey.INPUT_TOKENS: usage.input_tokens,
+                TokenUsageKey.OUTPUT_TOKENS: usage.output_tokens,
+                TokenUsageKey.TOTAL_TOKENS: usage.total_tokens,
+            }
+    except ImportError:
+        pass
+
+    return None

mlflow/optuna/__init__.py

@@ -0,0 +1,3 @@
+from mlflow.optuna.storage import MlflowStorage
+
+__all__ = ["MlflowStorage"]