genesis-flow 1.0.0__py3-none-any.whl

mlflow/mistral/autolog.py

@@ -0,0 +1,71 @@
+import inspect
+import logging
+
+import mlflow
+import mlflow.mistral
+from mlflow.entities import SpanType
+from mlflow.mistral.chat import convert_message_to_mlflow_chat, convert_tool_to_mlflow_chat_tool
+from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
+from mlflow.tracing.utils import set_span_chat_messages, set_span_chat_tools
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+
+_logger = logging.getLogger(__name__)
+
+
+def _construct_full_inputs(func, *args, **kwargs):
+    signature = inspect.signature(func)
+    # this 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
+
+
+def patched_class_call(original, self, *args, **kwargs):
+    config = AutoLoggingConfig.init(flavor_name=mlflow.mistral.FLAVOR_NAME)
+
+    if config.log_traces:
+        with mlflow.start_span(
+            name=f"{self.__class__.__name__}.{original.__name__}",
+            span_type=SpanType.CHAT_MODEL,
+        ) as span:
+            inputs = _construct_full_inputs(original, self, *args, **kwargs)
+            span.set_inputs(inputs)
+
+            if (tools := inputs.get("tools")) is not None:
+                try:
+                    tools = [convert_tool_to_mlflow_chat_tool(tool) for tool in tools if tool]
+                    set_span_chat_tools(span, tools)
+                except Exception as e:
+                    _logger.debug(f"Failed to set tools for {span}. Error: {e}")
+
+            try:
+                messages = [convert_message_to_mlflow_chat(m) for m in inputs.get("messages", [])]
+            except Exception as e:
+                _logger.debug(f"Failed to convert chat messages for {span}. Error: {e}")
+
+            try:
+                outputs = original(self, *args, **kwargs)
+                span.set_outputs(outputs)
+                if usage := getattr(outputs, "usage", None):
+                    span.set_attribute(
+                        SpanAttributeKey.CHAT_USAGE,
+                        {
+                            TokenUsageKey.INPUT_TOKENS: usage.prompt_tokens,
+                            TokenUsageKey.OUTPUT_TOKENS: usage.completion_tokens,
+                            TokenUsageKey.TOTAL_TOKENS: usage.total_tokens,
+                        },
+                    )
+            finally:
+                # Set message attribute once at the end to avoid multiple JSON serialization
+                try:
+                    for choice in getattr(outputs, "choices", []):
+                        choice_message = getattr(choice, "message", {})
+                        messages.append(convert_message_to_mlflow_chat(choice_message))
+                    set_span_chat_messages(span, messages)
+                except Exception as e:
+                    _logger.debug(f"Failed to set chat messages for {span}. Error: {e}")
+
+            return outputs

mlflow/mistral/chat.py

@@ -0,0 +1,135 @@
+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 _to_dict(obj: BaseModel):
+    if IS_PYDANTIC_V2_OR_NEWER:
+        return obj.model_dump()
+    return obj.dict()
+
+
+def convert_message_to_mlflow_chat(message: Union[BaseModel, dict[str, Any]]) -> ChatMessage:
+    """
+    Convert Mistral AI message object into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.mistral.ai/api/#tag/chat/operation/chat_completion_v1_chat_completions_post
+
+    Args:
+        message: Mistral AI 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")
+        tool_calls = message.get("tool_calls")
+        tool_call_id = message.get("tool_call_id")
+    elif isinstance(message, BaseModel):
+        content = message.content
+        role = message.role
+        # tool_calls is available if message is an AssistantMessage object
+        tool_calls = getattr(message, "tool_calls", None)
+        if tool_calls:
+            tool_calls = [_to_dict(tool_call) for tool_call in tool_calls]
+        # tool_call_id is available if message is a ToolMessage object
+        tool_call_id = getattr(message, "tool_call_id", None)
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Message must be either a dict or a Message object, but got: {type(message)}."
+        )
+
+    if tool_calls:
+        tool_calls = [
+            ToolCall(
+                id=tool_call["id"],
+                function=Function(
+                    name=tool_call["function"]["name"],
+                    arguments=json.dumps(tool_call["function"]["arguments"]),
+                ),
+                type="function",
+            )
+            for tool_call in tool_calls
+        ]
+
+    if isinstance(content, str):
+        return ChatMessage(
+            role=role, content=content, tool_calls=tool_calls, tool_call_id=tool_call_id
+        )
+
+    elif isinstance(content, list):
+        contents = []
+        tool_calls = []
+        tool_call_id = None
+        for content_chunk in content:
+            if isinstance(content_chunk, BaseModel):
+                content_chunk = _to_dict(content_chunk)
+            contents.append(_parse_content(content_chunk))
+
+        return ChatMessage(
+            role=role, content=contents, tool_calls=tool_calls, tool_call_id=tool_call_id
+        )
+
+    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_url":
+        return ImageContentPart(
+            image_url=ImageUrl(url=content["image_url"], detail="auto"),
+            type="image_url",
+        )
+    else:
+        raise MlflowException.invalid_parameter_value(
+            f"Unknown content type: {content_type['type']}. Please make sure the message "
+            "is a valid Mistral AI 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 Mistral AI tool definition into MLflow's standard format (OpenAI compatible).
+
+    Ref: https://docs.mistral.ai/capabilities/function_calling/#tools
+
+    Args:
+        tool: A dictionary represents a single tool definition in the input request.
+
+    Returns:
+        ChatTool: MLflow's standard tool definition object.
+    """
+    function = tool["function"]
+    return ChatTool(
+        type="function",
+        function=FunctionToolDefinition(
+            name=function["name"],
+            description=function.get("description"),
+            parameters=function["parameters"],
+        ),
+    )

mlflow/ml_package_versions.py

@@ -0,0 +1,452 @@
+# This file was auto-generated by update_ml_package_versions.py.
+# Please do not edit it manually.
+
+_ML_PACKAGE_VERSIONS = {
+    "sklearn": {
+        "package_info": {
+            "pip_release": "scikit-learn"
+        },
+        "models": {
+            "minimum": "1.3.0",
+            "maximum": "1.7.0"
+        },
+        "autologging": {
+            "minimum": "1.3.0",
+            "maximum": "1.7.0"
+        }
+    },
+    "pytorch": {
+        "package_info": {
+            "pip_release": "torch",
+            "module_name": "torch"
+        },
+        "models": {
+            "minimum": "2.1.0",
+            "maximum": "2.7.1"
+        },
+        "autologging": {
+            "minimum": "2.1.0",
+            "maximum": "2.7.1"
+        }
+    },
+    "pytorch-lightning": {
+        "package_info": {
+            "pip_release": "pytorch-lightning",
+            "module_name": "lightning"
+        },
+        "autologging": {
+            "minimum": "2.0.4",
+            "maximum": "2.5.1.post0"
+        }
+    },
+    "keras": {
+        "package_info": {
+            "pip_release": "keras"
+        },
+        "models": {
+            "minimum": "3.0.2",
+            "maximum": "3.10.0"
+        },
+        "autologging": {
+            "minimum": "3.0.2",
+            "maximum": "3.10.0"
+        }
+    },
+    "tensorflow": {
+        "package_info": {
+            "pip_release": "tensorflow"
+        },
+        "models": {
+            "minimum": "2.12.1",
+            "maximum": "2.19.0"
+        },
+        "autologging": {
+            "minimum": "2.12.1",
+            "maximum": "2.19.0"
+        }
+    },
+    "xgboost": {
+        "package_info": {
+            "pip_release": "xgboost"
+        },
+        "models": {
+            "minimum": "1.7.6",
+            "maximum": "3.0.2"
+        },
+        "autologging": {
+            "minimum": "1.7.6",
+            "maximum": "3.0.2"
+        }
+    },
+    "lightgbm": {
+        "package_info": {
+            "pip_release": "lightgbm"
+        },
+        "models": {
+            "minimum": "4.0.0",
+            "maximum": "4.6.0"
+        },
+        "autologging": {
+            "minimum": "4.0.0",
+            "maximum": "4.6.0"
+        }
+    },
+    "catboost": {
+        "package_info": {
+            "pip_release": "catboost"
+        },
+        "models": {
+            "minimum": "1.2.1",
+            "maximum": "1.2.8"
+        }
+    },
+    "onnx": {
+        "package_info": {
+            "pip_release": "onnx"
+        },
+        "models": {
+            "minimum": "1.14.1",
+            "maximum": "1.18.0"
+        }
+    },
+    "spacy": {
+        "package_info": {
+            "pip_release": "spacy"
+        },
+        "models": {
+            "minimum": "3.5.4",
+            "maximum": "3.8.7"
+        }
+    },
+    "statsmodels": {
+        "package_info": {
+            "pip_release": "statsmodels"
+        },
+        "models": {
+            "minimum": "0.14.1",
+            "maximum": "0.14.4"
+        },
+        "autologging": {
+            "minimum": "0.14.1",
+            "maximum": "0.14.4"
+        }
+    },
+    "spark": {
+        "package_info": {
+            "pip_release": "pyspark",
+            "module_name": "pyspark"
+        },
+        "models": {
+            "minimum": "3.2.1",
+            "maximum": "4.0.0"
+        },
+        "autologging": {
+            "minimum": "3.2.1",
+            "maximum": "4.0.0"
+        }
+    },
+    "prophet": {
+        "package_info": {
+            "pip_release": "prophet"
+        },
+        "models": {
+            "minimum": "1.1.6",
+            "maximum": "1.1.7"
+        }
+    },
+    "pmdarima": {
+        "package_info": {
+            "pip_release": "pmdarima"
+        },
+        "models": {
+            "minimum": "2.0.4",
+            "maximum": "2.0.4"
+        }
+    },
+    "diviner": {
+        "package_info": {
+            "pip_release": "diviner"
+        },
+        "models": {
+            "minimum": "0.1.1",
+            "maximum": "0.1.1"
+        }
+    },
+    "h2o": {
+        "package_info": {
+            "pip_release": "h2o"
+        },
+        "models": {
+            "minimum": "3.42.0.1",
+            "maximum": "3.46.0.7"
+        }
+    },
+    "shap": {
+        "package_info": {
+            "pip_release": "shap"
+        },
+        "models": {
+            "minimum": "0.42.1",
+            "maximum": "0.47.2"
+        }
+    },
+    "paddle": {
+        "package_info": {
+            "pip_release": "paddlepaddle"
+        },
+        "models": {
+            "minimum": "2.5.2",
+            "maximum": "3.0.0"
+        },
+        "autologging": {
+            "minimum": "2.5.2",
+            "maximum": "3.0.0"
+        }
+    },
+    "transformers": {
+        "package_info": {
+            "pip_release": "transformers"
+        },
+        "models": {
+            "minimum": "4.35.2",
+            "maximum": "4.52.4"
+        },
+        "autologging": {
+            "minimum": "4.35.2",
+            "maximum": "4.52.4"
+        }
+    },
+    "openai": {
+        "package_info": {
+            "pip_release": "openai"
+        },
+        "models": {
+            "minimum": "1.0.1",
+            "maximum": "1.85.0"
+        },
+        "autologging": {
+            "minimum": "1.17.0",
+            "maximum": "1.85.0"
+        }
+    },
+    "dspy": {
+        "package_info": {
+            "pip_release": "dspy"
+        },
+        "models": {
+            "minimum": "2.5.17",
+            "maximum": "2.6.27"
+        },
+        "autologging": {
+            "minimum": "2.5.17",
+            "maximum": "2.6.27"
+        }
+    },
+    "langchain": {
+        "package_info": {
+            "pip_release": "langchain"
+        },
+        "models": {
+            "minimum": "0.0.354",
+            "maximum": "0.3.25"
+        },
+        "autologging": {
+            "minimum": "0.1.0",
+            "maximum": "0.3.25"
+        }
+    },
+    "langgraph": {
+        "package_info": {
+            "pip_release": "langgraph"
+        },
+        "models": {
+            "minimum": "0.2.0",
+            "maximum": "0.4.8"
+        },
+        "autologging": {
+            "minimum": "0.2.0",
+            "maximum": "0.4.8"
+        }
+    },
+    "llama_index": {
+        "package_info": {
+            "pip_release": "llama-index",
+            "module_name": "llama_index.core"
+        },
+        "models": {
+            "minimum": "0.10.44",
+            "maximum": "0.12.41"
+        },
+        "autologging": {
+            "minimum": "0.10.44",
+            "maximum": "0.12.41"
+        }
+    },
+    "ag2": {
+        "package_info": {
+            "pip_release": "ag2",
+            "module_name": "autogen"
+        },
+        "autologging": {
+            "minimum": "0.7.0",
+            "maximum": "0.9.2"
+        }
+    },
+    "autogen": {
+        "package_info": {
+            "pip_release": "autogen-agentchat",
+            "module_name": "autogen_agentchat"
+        },
+        "autologging": {
+            "minimum": "0.4.9",
+            "maximum": "0.6.1"
+        }
+    },
+    "gemini": {
+        "package_info": {
+            "pip_release": "google-genai",
+            "module_name": "google.genai"
+        },
+        "autologging": {
+            "minimum": "1.0.0",
+            "maximum": "1.19.0"
+        }
+    },
+    "anthropic": {
+        "package_info": {
+            "pip_release": "anthropic"
+        },
+        "autologging": {
+            "minimum": "0.30.0",
+            "maximum": "0.53.0"
+        }
+    },
+    "crewai": {
+        "package_info": {
+            "pip_release": "crewai",
+            "module_name": "crewai"
+        },
+        "autologging": {
+            "minimum": "0.80.0",
+            "maximum": "0.126.0"
+        }
+    },
+    "pydantic_ai": {
+        "package_info": {
+            "pip_release": "pydantic-ai",
+            "module_name": "pydantic_ai"
+        },
+        "autologging": {
+            "minimum": "0.1.9",
+            "maximum": "0.2.16"
+        }
+    },
+    "smolagents": {
+        "package_info": {
+            "pip_release": "smolagents",
+            "module_name": "smolagents"
+        },
+        "autologging": {
+            "minimum": "1.14.0",
+            "maximum": "1.17.0"
+        }
+    },
+    "mistral": {
+        "package_info": {
+            "pip_release": "mistralai",
+            "module_name": "mistralai"
+        },
+        "autologging": {
+            "minimum": "1.0.0",
+            "maximum": "1.8.1"
+        }
+    },
+    "sentence_transformers": {
+        "package_info": {
+            "pip_release": "sentence-transformers"
+        },
+        "models": {
+            "minimum": "2.3.1",
+            "maximum": "4.1.0"
+        }
+    },
+    "johnsnowlabs": {
+        "package_info": {
+            "pip_release": "johnsnowlabs"
+        },
+        "models": {
+            "minimum": "4.4.8",
+            "maximum": "6.0.2"
+        }
+    },
+    "promptflow": {
+        "package_info": {
+            "pip_release": "promptflow"
+        },
+        "models": {
+            "minimum": "1.3.0",
+            "maximum": "1.18.0"
+        }
+    },
+    "litellm": {
+        "package_info": {
+            "pip_release": "litellm"
+        },
+        "autologging": {
+            "minimum": "1.52.9",
+            "maximum": "1.58.4"
+        }
+    },
+    "groq": {
+        "package_info": {
+            "pip_release": "groq"
+        },
+        "autologging": {
+            "minimum": "0.13.0",
+            "maximum": "0.26.0"
+        }
+    },
+    "bedrock": {
+        "package_info": {
+            "pip_release": "boto3",
+            "module_name": "boto3"
+        },
+        "autologging": {
+            "minimum": "1.33.0",
+            "maximum": "1.38.33"
+        }
+    }
+}
+
+# A mapping of flavor name to the module name to be imported for autologging.
+# This is used for checking version compatibility in autologging.
+# DO NOT EDIT MANUALLY
+FLAVOR_TO_MODULE_NAME = {
+    "sklearn": "sklearn",
+    "pytorch": "torch",
+    "pytorch-lightning": "lightning",
+    "keras": "keras",
+    "tensorflow": "tensorflow",
+    "xgboost": "xgboost",
+    "lightgbm": "lightgbm",
+    "statsmodels": "statsmodels",
+    "spark": "pyspark",
+    "paddle": "paddle",
+    "transformers": "transformers",
+    "openai": "openai",
+    "dspy": "dspy",
+    "langchain": "langchain",
+    "llama_index": "llama_index.core",
+    "ag2": "autogen",
+    "autogen": "autogen_agentchat",
+    "gemini": "google.genai",
+    "anthropic": "anthropic",
+    "crewai": "crewai",
+    "pydantic_ai": "pydantic_ai",
+    "smolagents": "smolagents",
+    "mistral": "mistralai",
+    "litellm": "litellm",
+    "groq": "groq",
+    "bedrock": "boto3",
+    "pyspark.ml": "pyspark"
+}