genesis-flow 1.0.0__py3-none-any.whl

mlflow/ag2/ag2_logger.py

@@ -0,0 +1,294 @@
+import functools
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Optional, Union
+
+from autogen import Agent, ConversableAgent
+from autogen.logger.base_logger import BaseLogger
+from openai.types.chat import ChatCompletion
+
+from mlflow.entities.span import NoOpSpan, Span, SpanType
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatus, SpanStatusCode
+from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
+from mlflow.tracing.fluent import start_span_no_context
+from mlflow.tracing.utils import capture_function_input_args
+from mlflow.utils.autologging_utils import autologging_is_disabled
+from mlflow.utils.autologging_utils.safety import safe_patch
+
+# For GroupChat, a single "received_message" events are passed around multiple
+# internal layers and thus too verbose if we show them all. Therefore we ignore
+# some of the message senders listed below.
+_EXCLUDED_MESSAGE_SENDERS = ["chat_manager", "checking_agent"]
+
+_logger = logging.getLogger(__name__)
+
+
+FLAVOR_NAME = "ag2"
+
+
+@dataclass
+class ChatState:
+    """
+    Represents the state of a chat session.
+    """
+
+    # The root span object that scopes the entire single chat session. All spans
+    # such as LLM, function calls, in the chat session should be children of this span.
+    session_span: Optional[Span] = None
+    # The last message object in the chat session.
+    last_message: Optional[Any] = None
+    # The timestamp (ns) of the last message in the chat session.
+    last_message_timestamp: int = 0
+    # LLM/Tool Spans created after the last message in the chat session.
+    # We consider them as operations for generating the next message and
+    # re-locate them under the corresponding message span.
+    pending_spans: list[Span] = field(default_factory=list)
+
+    def clear(self):
+        self.session_span = None
+        self.last_message = None
+        self.last_message_timestamp = 0
+        self.pending_spans = []
+
+
+def _catch_exception(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            _logger.error(f"Error occurred during AutoGen tracing: {e}")
+
+    return wrapper
+
+
+class MlflowAg2Logger(BaseLogger):
+    def __init__(self):
+        self._chat_state = ChatState()
+
+    def start(self) -> str:
+        return "session_id"
+
+    @_catch_exception
+    def log_new_agent(self, agent: ConversableAgent, init_args: dict[str, Any]) -> None:
+        """
+        This handler is called whenever a new agent instance is created.
+        Here we patch the agent's methods to start and end a trace around its chat session.
+        """
+        # TODO: Patch generate_reply() method as well
+        if hasattr(agent, "initiate_chat"):
+            safe_patch(
+                FLAVOR_NAME,
+                agent.__class__,
+                "initiate_chat",
+                # Setting root_only = True because sometimes compounded agent calls initiate_chat()
+                # method of its sub-agents, which should not start a new trace.
+                self._get_patch_function(root_only=True),
+            )
+        if hasattr(agent, "register_function"):
+
+            def patched(original, _self, function_map, **kwargs):
+                original(_self, function_map, **kwargs)
+                # Wrap the newly registered tools to start and end a span around its invocation.
+                for name, f in function_map.items():
+                    if f is not None:
+                        _self._function_map[name] = functools.partial(
+                            self._get_patch_function(span_type=SpanType.TOOL), f
+                        )
+
+            safe_patch(FLAVOR_NAME, agent.__class__, "register_function", patched)
+
+    def _get_patch_function(self, span_type: str = SpanType.UNKNOWN, root_only: bool = False):
+        """
+        Patch a function to start and end a span around its invocation.
+
+        Args:
+            f: The function to patch.
+            span_name: The name of the span. If None, the function name is used.
+            span_type: The type of the span. Default is SpanType.UNKNOWN.
+            root_only: If True, only create a span if it is the root of the chat session.
+                When there is an existing root span for the chat session, the function will
+                not create a new span.
+        """
+
+        def _wrapper(original, *args, **kwargs):
+            # If autologging is disabled, just run the original function. This is a safety net to
+            # prevent patching side effects from being effective after autologging is disabled.
+            if autologging_is_disabled(FLAVOR_NAME):
+                return original(*args, **kwargs)
+
+            if self._chat_state.session_span is None:
+                # Create the trace per chat session
+                span = start_span_no_context(
+                    name=original.__name__,
+                    span_type=span_type,
+                    inputs=capture_function_input_args(original, args, kwargs),
+                )
+                self._chat_state.session_span = span
+                try:
+                    result = original(*args, **kwargs)
+                except Exception as e:
+                    result = None
+                    self._record_exception(span, e)
+                    raise e
+                finally:
+                    span.end(outputs=result)
+                    # Clear the state to start a new chat session
+                    self._chat_state.clear()
+            elif not root_only:
+                span = self._start_span_in_session(
+                    name=original.__name__,
+                    span_type=span_type,
+                    inputs=capture_function_input_args(original, args, kwargs),
+                )
+                try:
+                    result = original(*args, **kwargs)
+                except Exception as e:
+                    result = None
+                    self._record_exception(span, e)
+                    raise e
+                finally:
+                    span.end(outputs=result)
+                    self._chat_state.pending_spans.append(span)
+            else:
+                result = original(*args, **kwargs)
+            return result
+
+        return _wrapper
+
+    def _record_exception(self, span: Span, e: Exception):
+        try:
+            span.set_status(SpanStatus(SpanStatusCode.ERROR, str(e)))
+            span.add_event(SpanEvent.from_exception(e))
+        except Exception as e:
+            _logger.warning(
+                "Failed to record exception in span.", exc_info=_logger.isEnabledFor(logging.DEBUG)
+            )
+
+    def _start_span_in_session(
+        self,
+        name: str,
+        span_type: str,
+        inputs: dict[str, Any],
+        attributes: Optional[dict[str, Any]] = None,
+        start_time_ns: Optional[int] = None,
+    ) -> Span:
+        """
+        Start a span in the current chat session.
+        """
+        if self._chat_state.session_span is None:
+            _logger.warning("Failed to start span. No active chat session.")
+            return NoOpSpan()
+
+        return start_span_no_context(
+            # Tentatively set the parent ID to the session root span, because we
+            # cannot create a span without a parent span (otherwise it will start
+            # a new trace). The actual parent will be determined once the chat
+            # message is received.
+            parent_span=self._chat_state.session_span,
+            name=name,
+            span_type=span_type,
+            inputs=inputs,
+            attributes=attributes,
+            start_time_ns=start_time_ns,
+        )
+
+    @_catch_exception
+    def log_event(self, source: Union[str, Agent], name: str, **kwargs: dict[str, Any]):
+        event_end_time = time.time_ns()
+        if name == "received_message":
+            if (self._chat_state.last_message is not None) and (
+                kwargs.get("sender") not in _EXCLUDED_MESSAGE_SENDERS
+            ):
+                span = self._start_span_in_session(
+                    name=kwargs["sender"],
+                    # Last message is recorded as the input of the next message
+                    inputs=self._chat_state.last_message,
+                    span_type=SpanType.AGENT,
+                    start_time_ns=self._chat_state.last_message_timestamp,
+                )
+                span.end(outputs=kwargs, end_time_ns=event_end_time)
+
+                # Re-locate the pended spans under this message span
+                for child_span in self._chat_state.pending_spans:
+                    child_span._span._parent = span._span.context
+                self._chat_state.pending_spans = []
+
+            self._chat_state.last_message = kwargs
+            self._chat_state.last_message_timestamp = event_end_time
+
+    @_catch_exception
+    def log_chat_completion(
+        self,
+        invocation_id: uuid.UUID,
+        client_id: int,
+        wrapper_id: int,
+        source: Union[str, Agent],
+        request: dict[str, Union[float, str, list[dict[str, str]]]],
+        response: Union[str, ChatCompletion],
+        is_cached: int,
+        cost: float,
+        start_time: str,
+    ) -> None:
+        # The start_time passed from AutoGen is in UTC timezone.
+        start_dt = datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S.%f")
+        start_dt = start_dt.replace(tzinfo=timezone.utc)
+        start_time_ns = int(start_dt.timestamp() * 1e9)
+        span = self._start_span_in_session(
+            name="chat_completion",
+            span_type=SpanType.LLM,
+            inputs=request,
+            attributes={
+                "source": source,
+                "client_id": client_id,
+                "invocation_id": invocation_id,
+                "wrapper_id": wrapper_id,
+                "cost": cost,
+                "is_cached": is_cached,
+            },
+            start_time_ns=start_time_ns,
+        )
+        if usage := self._parse_usage(response):
+            span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage)
+
+        span.end(outputs=response, end_time_ns=time.time_ns())
+        self._chat_state.pending_spans.append(span)
+
+    def _parse_usage(self, output: Any) -> Optional[dict[str, int]]:
+        usage = getattr(output, "usage", None)
+        if usage is None:
+            return None
+        input_tokens = usage.prompt_tokens
+        output_tokens = usage.completion_tokens
+        total_tokens = usage.total_tokens
+        if total_tokens is None and None not in (input_tokens, output_tokens):
+            total_tokens = input_tokens + output_tokens
+        return {
+            TokenUsageKey.INPUT_TOKENS: input_tokens,
+            TokenUsageKey.OUTPUT_TOKENS: output_tokens,
+            TokenUsageKey.TOTAL_TOKENS: total_tokens,
+        }
+
+    # The following methods are not used but are required to implement the BaseLogger interface.
+    @_catch_exception
+    def log_function_use(self, *args: Any, **kwargs: Any):
+        pass
+
+    @_catch_exception
+    def log_new_wrapper(self, wrapper, init_args):
+        pass
+
+    @_catch_exception
+    def log_new_client(self, client, wrapper, init_args):
+        pass
+
+    @_catch_exception
+    def stop(self) -> None:
+        pass
+
+    @_catch_exception
+    def get_connection(self):
+        pass

mlflow/anthropic/__init__.py

@@ -0,0 +1,40 @@
+from mlflow.anthropic.autolog import async_patched_class_call, patched_class_call
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+FLAVOR_NAME = "anthropic"
+
+
+@experimental(version="2.19.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Anthropic to MLflow.
+    Only synchronous calls and asynchronous APIs are supported. Streaming is not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Anthropic models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Anthropic autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Anthropic
+            autologging. If ``False``, show all events and warnings.
+    """
+    from anthropic.resources import AsyncMessages, Messages
+
+    safe_patch(
+        FLAVOR_NAME,
+        Messages,
+        "create",
+        patched_class_call,
+    )
+
+    safe_patch(
+        FLAVOR_NAME,
+        AsyncMessages,
+        "create",
+        async_patched_class_call,
+    )

mlflow/anthropic/autolog.py

@@ -0,0 +1,129 @@
+import logging
+from typing import Any, Optional
+
+import mlflow
+import mlflow.anthropic
+from mlflow.anthropic.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.entities import SpanType
+from mlflow.entities.span import LiveSpan
+from mlflow.entities.span_event import SpanEvent
+from mlflow.entities.span_status import SpanStatusCode
+from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
+from mlflow.tracing.fluent import start_span_no_context
+from mlflow.tracing.utils import (
+    construct_full_inputs,
+    set_span_chat_messages,
+    set_span_chat_tools,
+)
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    with TracingSession(original, self, args, kwargs) as manager:
+        output = original(self, *args, **kwargs)
+        manager.output = output
+        return output
+
+
+async def async_patched_class_call(original, self, *args, **kwargs):
+    async with TracingSession(original, self, args, kwargs) as manager:
+        output = await original(self, *args, **kwargs)
+        manager.output = output
+        return output
+
+
+class TracingSession:
+    """Context manager for handling MLflow spans in both sync and async contexts."""
+
+    def __init__(self, original, instance, args, kwargs):
+        self.original = original
+        self.instance = instance
+        self.inputs = construct_full_inputs(original, instance, *args, **kwargs)
+
+        # These attributes are set outside the constructor.
+        self.span = None
+        self.output = None
+
+    def __enter__(self):
+        return self._enter_impl()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        return self._enter_impl()
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        config = AutoLoggingConfig.init(flavor_name=mlflow.anthropic.FLAVOR_NAME)
+
+        if config.log_traces:
+            attributes = {}
+            self.span = start_span_no_context(
+                name=f"{self.instance.__class__.__name__}.{self.original.__name__}",
+                span_type=_get_span_type(self.original.__name__),
+                inputs=self.inputs,
+                attributes=attributes,
+            )
+            _set_tool_attribute(self.span, self.inputs)
+
+        return self
+
+    def _exit_impl(self, exc_type, exc_val, exc_tb) -> None:
+        if self.span:
+            if exc_val:
+                self.span.add_event(SpanEvent.from_exception(exc_val))
+                status = SpanStatusCode.ERROR
+            else:
+                status = SpanStatusCode.OK
+
+            _set_chat_message_attribute(self.span, self.inputs, self.output)
+            self.span.end(status=status, outputs=self.output)
+
+
+def _get_span_type(task_name: str) -> str:
+    # Anthropic has a few APIs in beta, e.g., count_tokens.
+    # Once they are stable, we can add them to the mapping.
+    span_type_mapping = {
+        "create": SpanType.CHAT_MODEL,
+    }
+    return span_type_mapping.get(task_name, SpanType.UNKNOWN)
+
+
+def _set_tool_attribute(span: LiveSpan, inputs: dict[str, Any]):
+    if (tools := inputs.get("tools")) is not None:
+        try:
+            tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tools]
+            set_span_chat_tools(span, tools)
+        except Exception as e:
+            _logger.debug(f"Failed to set tools for {span}. Error: {e}")
+
+
+def _set_chat_message_attribute(span: LiveSpan, inputs: dict[str, Any], output: Any):
+    try:
+        messages = [convert_message_to_mlflow_chat(msg) for msg in inputs.get("messages", [])]
+        if output is not None:
+            messages.append(convert_message_to_mlflow_chat(output))
+        set_span_chat_messages(span, messages)
+        if usage := _parse_usage(output):
+            span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage)
+    except Exception as e:
+        _logger.debug(f"Failed to set chat messages for {span}. Error: {e}")
+
+
+def _parse_usage(output: Any) -> Optional[dict[str, int]]:
+    try:
+        usage = getattr(output, "usage", None)
+        if usage:
+            return {
+                TokenUsageKey.INPUT_TOKENS: usage.input_tokens,
+                TokenUsageKey.OUTPUT_TOKENS: usage.output_tokens,
+                TokenUsageKey.TOTAL_TOKENS: usage.input_tokens + usage.output_tokens,
+            }
+    except Exception as e:
+        _logger.debug(f"Failed to parse token usage from output: {e}")
+    return None

mlflow/anthropic/chat.py

@@ -0,0 +1,144 @@
+import json
+from typing import Any, Union
+
+from pydantic import BaseModel
+
+from mlflow.exceptions import MlflowException
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    TextContentPart,
+    ToolCall,
+)
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
+
+
+def convert_message_to_mlflow_chat(message: Union[BaseModel, dict[str, Any]]) -> ChatMessage:
+    """
+    Convert Anthropic message object into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.anthropic.com/en/api/messages#body-messages
+
+    Args:
+        message: Anthropic message object or a dictionary representing the message.
+
+    Returns:
+        ChatMessage: MLflow's standard chat message object.
+    """
+    if isinstance(message, dict):
+        content = message.get("content")
+        role = message.get("role")
+    elif isinstance(message, BaseModel):
+        content = message.content
+        role = message.role
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Message must be either a dict or a Message object, but got: {type(message)}."
+        )
+
+    if isinstance(content, str):
+        return ChatMessage(role=role, content=content)
+
+    elif isinstance(content, list):
+        contents = []
+        tool_calls = []
+        tool_call_id = None
+        for content_block in content:
+            if isinstance(content_block, BaseModel):
+                if IS_PYDANTIC_V2_OR_NEWER:
+                    content_block = content_block.model_dump()
+                else:
+                    content_block = content_block.dict()
+
+            content_type = content_block.get("type")
+            if content_type == "tool_use":
+                # Anthropic response contains tool calls in the content block
+                # Ref: https://docs.anthropic.com/en/docs/build-with-claude/tool-use#example-api-response-with-a-tool-use-content-block
+                tool_calls.append(
+                    ToolCall(
+                        id=content_block["id"],
+                        function=Function(
+                            name=content_block["name"], arguments=json.dumps(content_block["input"])
+                        ),
+                        type="function",
+                    )
+                )
+            elif content_type == "tool_result":
+                # In Anthropic, the result of tool execution is returned as a special content type
+                # "tool_result" with "user" role, which corresponds to the "tool" role in OpenAI.
+                role = "tool"
+                tool_call_id = content_block["tool_use_id"]
+                if result_content := content_block.get("content"):
+                    contents.append(_parse_content(result_content))
+                else:
+                    contents.append(TextContentPart(text="", type="text"))
+            else:
+                contents.append(_parse_content(content_block))
+
+        message = ChatMessage(role=role, content=contents)
+        # Only set tool_calls field when it is present
+        if tool_calls:
+            message.tool_calls = tool_calls
+        if tool_call_id:
+            message.tool_call_id = tool_call_id
+        return message
+
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Invalid content type. Must be either a string or a list, but got: {type(content)}."
+        )
+
+
+def _parse_content(content: Union[str, dict[str, Any]]) -> Union[TextContentPart, ImageContentPart]:
+    if isinstance(content, str):
+        return TextContentPart(text=content, type="text")
+
+    content_type = content.get("type")
+    if content_type == "text":
+        return TextContentPart(text=content["text"], type="text")
+    elif content_type == "image":
+        source = content["source"]
+        return ImageContentPart(
+            image_url=ImageUrl(
+                url=f"data:{source['media_type']};{source['type']},{source['data']}"
+            ),
+            type="image_url",
+        )
+    # Claude 3.7 added new "thinking" content block, which is essentially a text block as of now.
+    # TODO: We should consider adding a new ContentPart type if more providers support this.
+    # https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking
+    elif content_type == "thinking":
+        return TextContentPart(text=content["thinking"], type="text")
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Unknown content type: {content_type['type']}. Please make sure the message "
+            "is a valid Anthropic message object. If it is a valid type, contact to the "
+            "MLflow maintainer via https://github.com/mlflow/mlflow/issues/new/choose for "
+            "requesting support for a new message type."
+        )
+
+
+def convert_tool_to_mlflow_chat_tool(tool: dict[str, Any]) -> ChatTool:
+    """
+    Convert Anthropic tool definition into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.anthropic.com/en/docs/build-with-claude/tool-use
+
+    Args:
+        tool: A dictionary represents a single tool definition in the input request.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=tool.get("name"),
+            description=tool.get("description"),
+            parameters=tool.get("input_schema"),
+        ),
+    )