star

venv/lib/python3.9/site-packages/mlflow/gemini/__init__.py

@@ -0,0 +1,81 @@
+"""
+The ``mlflow.gemini`` module provides an API for tracing the interaction with Gemini models.
+"""
+
+from mlflow.gemini.autolog import (
+    patched_class_call,
+    patched_module_call,
+)
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+
+FLAVOR_NAME = "gemini"
+
+
+@experimental
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    log_traces: bool = True,
+    disable: bool = False,
+    silent: bool = False,
+):
+    """
+    Enables (or disables) and configures autologging from Gemini to MLflow.
+    Currently, both legacy SDK google-generativeai and new SDK google-genai are supported.
+    Only synchronous calls are supported. Asynchnorous APIs and streaming are not recorded.
+
+    Args:
+        log_traces: If ``True``, traces are logged for Gemini models.
+            If ``False``, no traces are collected during inference. Default to ``True``.
+        disable: If ``True``, disables the Gemini autologging. Default to ``False``.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Gemini
+            autologging. If ``False``, show all events and warnings.
+    """
+    try:
+        from google import generativeai
+
+        for method in ["generate_content", "count_tokens"]:
+            safe_patch(
+                FLAVOR_NAME,
+                generativeai.GenerativeModel,
+                method,
+                patched_class_call,
+            )
+
+        safe_patch(
+            FLAVOR_NAME,
+            generativeai.ChatSession,
+            "send_message",
+            patched_class_call,
+        )
+
+        safe_patch(
+            FLAVOR_NAME,
+            generativeai,
+            "embed_content",
+            patched_module_call,
+        )
+    except ImportError:
+        pass
+
+    try:
+        from google import genai
+
+        # Since the genai SDK calls "_generate_content" iteratively within "generate_content",
+        # we need to patch both "generate_content" and "_generate_content".
+        for method in ["generate_content", "_generate_content", "count_tokens", "embed_content"]:
+            safe_patch(
+                FLAVOR_NAME,
+                genai.models.Models,
+                method,
+                patched_class_call,
+            )
+
+        safe_patch(
+            FLAVOR_NAME,
+            genai.chats.Chat,
+            "send_message",
+            patched_class_call,
+        )
+    except ImportError:
+        pass

venv/lib/python3.9/site-packages/mlflow/gemini/autolog.py

@@ -0,0 +1,186 @@
+import inspect
+import logging
+
+import mlflow
+import mlflow.gemini
+from mlflow.entities import SpanType
+from mlflow.gemini.chat import (
+    convert_gemini_func_to_mlflow_chat_tool,
+    parse_gemini_content_to_mlflow_chat_messages,
+)
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.types.chat import ChatMessage
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+try:
+    # This is for supporting the previous Google GenAI SDK
+    # https://github.com/google-gemini/generative-ai-python
+    from google import generativeai
+
+    has_generativeai = True
+except ImportError:
+    has_generativeai = False
+
+try:
+    from google import genai
+
+    has_genai = True
+except ImportError:
+    has_genai = False
+
+_logger = logging.getLogger(__name__)
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    """
+    This method is used for patching class methods of gemini SDKs.
+    This patch creates a span and set input and output of the original method to the span.
+    """
+    config = AutoLoggingConfig.init(flavor_name=mlflow.gemini.FLAVOR_NAME)
+
+    if config.log_traces:
+        with mlflow.start_span(
+            name=f"{self.__class__.__name__}.{original.__name__}",
+            span_type=_get_span_type(original.__name__),
+        ) as span:
+            inputs = _construct_full_inputs(original, self, *args, **kwargs)
+            span.set_inputs(inputs)
+            if has_generativeai and isinstance(self, generativeai.GenerativeModel):
+                _log_generativeai_tool_definition(self, span)
+            if has_genai and isinstance(self, (genai.models.Models, genai.chats.Chat)):
+                _log_genai_tool_definition(self, inputs, span)
+
+            result = original(self, *args, **kwargs)
+
+            if (
+                has_generativeai and isinstance(result, generativeai.types.GenerateContentResponse)
+            ) or (has_genai and isinstance(result, genai.types.GenerateContentResponse)):
+                try:
+                    content = _get_keys(inputs, ["contents", "content", "message"])
+                    messages = parse_gemini_content_to_mlflow_chat_messages(content)
+                    messages += _parse_outputs(result)
+                    if messages:
+                        set_span_chat_messages(span=span, messages=messages)
+                except Exception as e:
+                    _logger.warning(
+                        f"An exception occurred on logging chat attributes for {span}. Error: {e}"
+                    )
+
+            # need to convert the response of generate_content for better visualization
+            outputs = result.to_dict() if hasattr(result, "to_dict") else result
+            span.set_outputs(outputs)
+
+            return result
+
+
+def patched_module_call(original, *args, **kwargs):
+    """
+    This method is used for patching standalone functions of the google.generativeai module.
+    This patch creates a span and set input and output of the original function to the span.
+    """
+    config = AutoLoggingConfig.init(flavor_name=mlflow.gemini.FLAVOR_NAME)
+
+    if config.log_traces:
+        with mlflow.start_span(
+            name=f"{original.__name__}",
+            span_type=_get_span_type(original.__name__),
+        ) as span:
+            inputs = _construct_full_inputs(original, *args, **kwargs)
+            span.set_inputs(inputs)
+            result = original(*args, **kwargs)
+            # need to convert the response of generate_content for better visualization
+            outputs = result.to_dict() if hasattr(result, "to_dict") else result
+            span.set_outputs(outputs)
+
+            return result
+
+
+def _get_keys(dic, keys):
+    for key in keys:
+        if key in dic:
+            return dic[key]
+
+    return None
+
+
+def _parse_outputs(outputs) -> list[ChatMessage]:
+    """
+    This method extract chat messages from genai.types.generation_types.GenerateContentResponse
+    """
+    # content always exist on output
+    # https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/generative_service.proto#L490
+    return sum(
+        [
+            parse_gemini_content_to_mlflow_chat_messages(candidate.content)
+            for candidate in outputs.candidates
+        ],
+        [],
+    )
+
+
+def _log_generativeai_tool_definition(model, span):
+    """
+    This method extract tool definition from generativeai tool type.
+    """
+    # when tools are not passed
+    if not getattr(model, "_tools", None):
+        return
+
+    try:
+        set_span_chat_tools(
+            span,
+            [
+                convert_gemini_func_to_mlflow_chat_tool(func)
+                for func in model._tools.to_proto()[0].function_declarations
+            ],
+        )
+    except Exception as e:
+        _logger.warning(f"Failed to set tool definitions for {span}. Error: {e}")
+
+
+def _log_genai_tool_definition(model, inputs, span):
+    """
+    This method extract tool definition from genai tool type.
+    """
+    config = inputs.get("config")
+    tools = getattr(config, "tools", None)
+    if not tools:
+        return
+    # Here, we use an internal function of gemini library to convert callable to Tool schema to
+    # avoid having the same logic on mlflow side and there is no public attribute for Tool schema.
+    # https://github.com/googleapis/python-genai/blob/01b15e32d3823a58d25534bb6eea93f30bf82219/google/genai/_transformers.py#L662
+    tools = genai._transformers.t_tools(model._api_client, tools)
+
+    try:
+        set_span_chat_tools(
+            span,
+            [
+                convert_gemini_func_to_mlflow_chat_tool(function_declaration)
+                for tool in tools
+                for function_declaration in tool.function_declarations
+            ],
+        )
+    except Exception as e:
+        _logger.warning(f"Failed to set tool definitions for {span}. Error: {e}")
+
+
+def _get_span_type(task_name: str) -> str:
+    span_type_mapping = {
+        "generate_content": SpanType.LLM,
+        "_generate_content": SpanType.LLM,
+        "send_message": SpanType.CHAT_MODEL,
+        "count_tokens": SpanType.LLM,
+        "embed_content": SpanType.EMBEDDING,
+    }
+    return span_type_mapping.get(task_name, SpanType.UNKNOWN)
+
+
+def _construct_full_inputs(func, *args, **kwargs):
+    signature = inspect.signature(func)
+    # this method does not create copy. So values should not be mutated directly
+    arguments = signature.bind_partial(*args, **kwargs).arguments
+
+    if "self" in arguments:
+        arguments.pop("self")
+
+    return arguments

venv/lib/python3.9/site-packages/mlflow/gemini/chat.py

@@ -0,0 +1,261 @@
+import json
+import logging
+from typing import TYPE_CHECKING, Optional, Union
+
+from mlflow.types.chat import (
+    ChatMessage,
+    ChatTool,
+    Function,
+    FunctionParams,
+    FunctionToolDefinition,
+    ImageContentPart,
+    ImageUrl,
+    ParamProperty,
+    TextContentPart,
+    ToolCall,
+)
+
+if TYPE_CHECKING:
+    from google import genai
+
+_logger = logging.getLogger(__name__)
+
+
+def convert_gemini_func_to_mlflow_chat_tool(
+    function_def: "genai.types.FunctionDeclaration",
+) -> ChatTool:
+    """
+    Convert Gemini function definition into MLflow's standard format (OpenAI compatible).
+    Ref: https://ai.google.dev/gemini-api/docs/function-calling
+
+    Args:
+        function_def: A genai.types.FunctionDeclaration or genai.protos.FunctionDeclaration object
+                      representing a function definition.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=function_def.name,
+            description=function_def.description,
+            parameters=_convert_gemini_function_param_to_mlflow_function_param(
+                function_def.parameters
+            ),
+        ),
+    )
+
+
+def convert_gemini_func_call_to_mlflow_tool_call(
+    func_call: "genai.types.FunctionCall",
+) -> ToolCall:
+    """
+    Convert Gemini function call into MLflow's standard format (OpenAI compatible).
+    Ref: https://ai.google.dev/gemini-api/docs/function-calling
+
+    Args:
+        func_call: A genai.types.FunctionCall or genai.protos.FunctionCall object
+                   representing a single func call.
+
+    Returns:
+        ToolCall: MLflow's standard tool call object.
+    """
+    # original args object is not json serializable
+    args = func_call.args or {}
+
+    return ToolCall(
+        # Gemini does not have func call id
+        id=func_call.name,
+        type="function",
+        function=Function(name=func_call.name, arguments=json.dumps(dict(args))),
+    )
+
+
+def parse_gemini_content_to_mlflow_chat_messages(
+    content: "genai.types.ContentsType",
+) -> list[ChatMessage]:
+    """
+    Convert a gemini content to chat messages.
+
+    Args:
+        content: A genai.types.ContentsType object representing the model content.
+
+    Returns:
+        list[ChatMessage]: A list of MLflow's standard chat messages.
+    """
+    if isinstance(content, str):
+        # Assume str content is used only for user input
+        return [
+            ChatMessage(
+                role="user",
+                content=content,
+            )
+        ]
+    elif isinstance(content, list):
+        # either list of user inputs or multi-turn conversation
+        if not content:
+            return []
+        # when chat history is passed, parse content recursively
+        if hasattr(content[0], "parts"):
+            return sum(
+                [
+                    parse_gemini_content_to_mlflow_chat_messages(content_block)
+                    for content_block in content
+                ],
+                [],
+            )
+
+        # when multiple contents are passed by user
+        return [_construct_chat_message(content, "user")]
+    elif hasattr(content, "parts"):
+        # eigher genai.types.Content or ContentDict
+
+        # This could be unset for single turn conversation even if this is content proto
+        # https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L64
+        role = getattr(content, "role", "model") or "model"
+        # we normalize role and use assistant
+        if role == "model":
+            role = "assistant"
+
+        return [_construct_chat_message(content.parts, role)]
+    else:
+        _logger.debug(f"Received an unsupported content type: {content.__class__}")
+        return []
+
+
+def _construct_chat_message(parts: list["genai.types.PartType"], role: str) -> ChatMessage:
+    tool_calls = []
+    content_parts = []
+    for content_part in parts:
+        part = _parse_content_part(content_part)
+        if isinstance(part, (TextContentPart, ImageContentPart)):
+            content_parts.append(part)
+        elif isinstance(part, ToolCall):
+            tool_calls.append(part)
+    chat_message = ChatMessage(
+        role=role,
+        content=content_parts or None,
+    )
+
+    if tool_calls:
+        chat_message.tool_calls = tool_calls
+
+    return chat_message
+
+
+def _parse_content_part(part: "genai.types.PartType") -> Optional[Union[TextContentPart, ToolCall]]:
+    """
+    Convert Gemini part type into MLflow's standard format (OpenAI compatible).
+    Ref: https://ai.google.dev/gemini-api/docs/function-calling
+
+    Args:
+        part: A genai.types.PartType object representing a part of content.
+
+    Returns:
+        Optional[Union[TextContentPart, ToolCall]]: MLflow's standard content part.
+    """
+    # The schema of the Part proto is available at https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L76
+    if function_call := getattr(part, "function_call", None):
+        # FunctionCall part: https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L316
+        return convert_gemini_func_call_to_mlflow_tool_call(function_call)
+    elif function_response := getattr(part, "function_response", None):
+        # FunctionResponse part: https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L332
+        if hasattr(function_response, "json"):
+            # genai
+            return TextContentPart(text=function_response.json(), type="text")
+        # generativeai
+        return TextContentPart(
+            text=str(type(function_response).to_dict(function_response)), type="text"
+        )
+    elif blob := getattr(part, "inline_data", None):
+        # Blob part: https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L109C9-L109C13
+        return ImageContentPart(
+            image_url=ImageUrl(
+                url=f"data:{blob.mime_type};base64,{blob.data}",
+                detail="auto",
+            ),
+            type="image_url",
+        )
+    elif file := getattr(part, "file_data", None):
+        # FileData part: https://github.com/googleapis/googleapis/blob/9e966149c59f47f6305d66c98e2a9e7d9c26a2eb/google/ai/generativelanguage/v1beta/content.proto#L124
+        return ImageContentPart(
+            image_url=ImageUrl(
+                url=file.file_uri,
+                detail="auto",
+            ),
+            type="image_url",
+        )
+    elif hasattr(part, "mime_type"):
+        # Blob part or FileData part
+        url = (
+            part.file_uri
+            if hasattr(part, "file_uri")
+            else f"data:{part.mime_type};base64,{part.data}"
+        )
+        return ImageContentPart(
+            image_url=ImageUrl(url=url, detail="auto"),
+            type="image_url",
+        )
+    elif isinstance(part, dict):
+        if "mime_type" in part:
+            # genai.types.BlobDict
+            return ImageContentPart(
+                image_url=ImageUrl(
+                    url=f"data:{part['mime_type']};base64,{part['data']}", detail="auto"
+                ),
+                type="image_url",
+            )
+        elif "text" in part:
+            return TextContentPart(text=part["text"], type="text")
+    elif text := getattr(part, "text", None):
+        # Text part
+        return TextContentPart(text=text, type="text")
+    elif isinstance(part, str):
+        return TextContentPart(text=part, type="text")
+    # TODO: Gemini supports more types. Consider including unsupported types (e.g. PIL image)
+    _logger.debug(f"Received an unsupported content block type: {part.__class__}")
+
+
+def _convert_gemini_param_property_to_mlflow_param_property(param_property) -> ParamProperty:
+    """
+    Convert Gemini parameter property definition into MLflow's standard format (OpenAI compatible).
+    Ref: https://ai.google.dev/gemini-api/docs/function-calling
+
+    Args:
+        param_property: A genai.types.Schema or genai.protos.Schema object
+                        representing a parameter property.
+
+    Returns:
+        ParamProperty: MLflow's standard param property object.
+    """
+    type_name = param_property.type
+    type_name = type_name.name.lower() if hasattr(type_name, "name") else type_name.lower()
+    return ParamProperty(
+        description=param_property.description,
+        enum=param_property.enum,
+        type=type_name,
+    )
+
+
+def _convert_gemini_function_param_to_mlflow_function_param(
+    function_params: "genai.types.Schema",
+) -> FunctionParams:
+    """
+    Convert Gemini function parameter definition into MLflow's standard format (OpenAI compatible).
+    Ref: https://ai.google.dev/gemini-api/docs/function-calling
+
+    Args:
+        function_params: A genai.types.Schema or genai.protos.Schema object
+                         representing function parameters.
+
+    Returns:
+        FunctionParams: MLflow's standard function parameter object.
+    """
+    return FunctionParams(
+        properties={
+            k: _convert_gemini_param_property_to_mlflow_param_property(v)
+            for k, v in function_params.properties.items()
+        },
+        required=function_params.required,
+    )