genesis-flow 1.0.0__py3-none-any.whl

mlflow/entities/model_registry/prompt_version.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import re
+from typing import Optional, Union
+
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.exceptions import MlflowException
+from mlflow.prompt.constants import (
+    IS_PROMPT_TAG_KEY,
+    PROMPT_TEMPLATE_VARIABLE_PATTERN,
+    PROMPT_TEXT_DISPLAY_LIMIT,
+    PROMPT_TEXT_TAG_KEY,
+)
+
+# Alias type
+PromptVersionTag = ModelVersionTag
+
+
+def _is_reserved_tag(key: str) -> bool:
+    return key in {IS_PROMPT_TAG_KEY, PROMPT_TEXT_TAG_KEY}
+
+
+class PromptVersion(_ModelRegistryEntity):
+    """
+    An entity representing a specific version of a prompt with its template content.
+
+    Args:
+        name: The name of the prompt.
+        version: The version number of the prompt.
+        template: The template text of the prompt. It can contain variables enclosed in
+            double curly braces, e.g. {{variable}}, which will be replaced with actual values
+            by the `format` method. MLflow use the same variable naming rules same as Jinja2
+            https://jinja.palletsprojects.com/en/stable/api/#notes-on-identifiers
+        commit_message: The commit message for the prompt version. Optional.
+        creation_timestamp: Timestamp of the prompt creation. Optional.
+        tags: A dictionary of tags associated with the **prompt version**.
+            This is useful for storing version-specific information, such as the author of
+            the changes. Optional.
+        aliases: List of aliases for this prompt version. Optional.
+        last_updated_timestamp: Timestamp of last update. Optional.
+        user_id: User ID that created this prompt version. Optional.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        version: int,
+        template: str,
+        commit_message: Optional[str] = None,
+        creation_timestamp: Optional[int] = None,
+        tags: Optional[dict[str, str]] = None,
+        aliases: Optional[list[str]] = None,
+        last_updated_timestamp: Optional[int] = None,
+        user_id: Optional[str] = None,
+    ):
+        super().__init__()
+
+        # Core PromptVersion attributes
+        self._name: str = name
+        self._version: str = str(version)  # Store as string internally
+        self._creation_time: int = creation_timestamp or 0
+
+        # Store template text as a tag
+        tags = tags or {}
+        tags[PROMPT_TEXT_TAG_KEY] = template
+        tags[IS_PROMPT_TAG_KEY] = "true"
+
+        # Store the tags dict
+        self._tags: dict[str, str] = tags
+
+        self._variables = set(PROMPT_TEMPLATE_VARIABLE_PATTERN.findall(template))
+        self._last_updated_timestamp: Optional[int] = last_updated_timestamp
+        self._description: Optional[str] = commit_message
+        self._user_id: Optional[str] = user_id
+        self._aliases: list[str] = aliases or []
+
+    def __repr__(self) -> str:
+        text = (
+            self.template[:PROMPT_TEXT_DISPLAY_LIMIT] + "..."
+            if len(self.template) > PROMPT_TEXT_DISPLAY_LIMIT
+            else self.template
+        )
+        return f"PromptVersion(name={self.name}, version={self.version}, template={text})"
+
+    # Core PromptVersion properties
+    @property
+    def template(self) -> str:
+        """
+        Return the template text of the prompt.
+        """
+        return self._tags[PROMPT_TEXT_TAG_KEY]
+
+    def to_single_brace_format(self) -> str:
+        """
+        Convert the template text to single brace format. This is useful for integrating with other
+        systems that use single curly braces for variable replacement, such as LangChain's prompt
+        template. Default is False.
+        """
+        t = self.template
+        for var in self.variables:
+            t = re.sub(r"\{\{\s*" + var + r"\s*\}\}", "{" + var + "}", t)
+        return t
+
+    @property
+    def variables(self) -> set[str]:
+        """
+        Return a list of variables in the template text.
+        The value must be enclosed in double curly braces, e.g. {{variable}}.
+        """
+        return self._variables
+
+    @property
+    def commit_message(self) -> Optional[str]:
+        """
+        Return the commit message of the prompt version.
+        """
+        return self.description
+
+    @property
+    def tags(self) -> dict[str, str]:
+        """
+        Return the version-level tags.
+        """
+        return {key: value for key, value in self._tags.items() if not _is_reserved_tag(key)}
+
+    @property
+    def uri(self) -> str:
+        """Return the URI of the prompt."""
+        return f"prompts:/{self.name}/{self.version}"
+
+    @property
+    def name(self) -> str:
+        """String. Unique name within Model Registry."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name: str):
+        self._name = new_name
+
+    @property
+    def version(self) -> int:
+        """Version"""
+        return int(self._version)
+
+    @property
+    def creation_timestamp(self) -> int:
+        """Integer. Prompt version creation timestamp (milliseconds since the Unix epoch)."""
+        return self._creation_time
+
+    @property
+    def last_updated_timestamp(self) -> Optional[int]:
+        """Integer. Timestamp of last update for this prompt version (milliseconds since the Unix
+        epoch).
+        """
+        return self._last_updated_timestamp
+
+    @last_updated_timestamp.setter
+    def last_updated_timestamp(self, updated_timestamp: int):
+        self._last_updated_timestamp = updated_timestamp
+
+    @property
+    def description(self) -> Optional[str]:
+        """String. Description"""
+        return self._description
+
+    @description.setter
+    def description(self, description: str):
+        self._description = description
+
+    @property
+    def user_id(self) -> Optional[str]:
+        """String. User ID that created this prompt version."""
+        return self._user_id
+
+    @property
+    def aliases(self) -> list[str]:
+        """List of aliases (string) for the current prompt version."""
+        return self._aliases
+
+    @aliases.setter
+    def aliases(self, aliases: list[str]):
+        self._aliases = aliases
+
+    # Methods
+    @classmethod
+    def _properties(cls) -> list[str]:
+        # aggregate with base class properties since cls.__dict__ does not do it automatically
+        return sorted(cls._get_properties_helper())
+
+    def _add_tag(self, tag: ModelVersionTag):
+        self._tags[tag.key] = tag.value
+
+    def format(self, allow_partial: bool = False, **kwargs) -> Union[PromptVersion, str]:
+        """
+        Format the template text with the given keyword arguments.
+        By default, it raises an error if there are missing variables. To format
+        the prompt text partially, set `allow_partial=True`.
+
+        Example:
+
+        .. code-block:: python
+
+            prompt = Prompt("my-prompt", 1, "Hello, {{title}} {{name}}!")
+            formatted = prompt.format(title="Ms", name="Alice")
+            print(formatted)
+            # Output: "Hello, Ms Alice!"
+
+            # Partial formatting
+            formatted = prompt.format(title="Ms", allow_partial=True)
+            print(formatted)
+            # Output: PromptVersion(name=my-prompt, version=1, template="Hello, Ms {{name}}!")
+
+
+        Args:
+            allow_partial: If True, allow partial formatting of the prompt text.
+                If False, raise an error if there are missing variables.
+            kwargs: Keyword arguments to replace the variables in the template.
+        """
+        input_keys = set(kwargs.keys())
+
+        template = self.template
+        for key, value in kwargs.items():
+            template = re.sub(r"\{\{\s*" + key + r"\s*\}\}", str(value), template)
+
+        if missing_keys := self.variables - input_keys:
+            if not allow_partial:
+                raise MlflowException.invalid_parameter_value(
+                    f"Missing variables: {missing_keys}. To partially format the prompt, "
+                    "set `allow_partial=True`."
+                )
+            else:
+                return PromptVersion(
+                    name=self.name,
+                    version=int(self.version),
+                    template=template,
+                    commit_message=self.commit_message,
+                    creation_timestamp=self.creation_timestamp,
+                    tags=self.tags,
+                    aliases=self.aliases,
+                    last_updated_timestamp=self.last_updated_timestamp,
+                    user_id=self.user_id,
+                )
+        return template

mlflow/entities/model_registry/registered_model.py

@@ -0,0 +1,175 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+from mlflow.entities.model_registry.registered_model_deployment_job_state import (
+    RegisteredModelDeploymentJobState,
+)
+from mlflow.entities.model_registry.registered_model_tag import RegisteredModelTag
+from mlflow.protos.model_registry_pb2 import RegisteredModel as ProtoRegisteredModel
+from mlflow.protos.model_registry_pb2 import RegisteredModelAlias as ProtoRegisteredModelAlias
+from mlflow.protos.model_registry_pb2 import RegisteredModelTag as ProtoRegisteredModelTag
+
+
+class RegisteredModel(_ModelRegistryEntity):
+    """
+    MLflow entity for Registered Model.
+    """
+
+    def __init__(
+        self,
+        name,
+        creation_timestamp=None,
+        last_updated_timestamp=None,
+        description=None,
+        latest_versions=None,
+        tags=None,
+        aliases=None,
+        deployment_job_id=None,
+        deployment_job_state=None,
+    ):
+        # Constructor is called only from within the system by various backend stores.
+        super().__init__()
+        self._name = name
+        self._creation_time = creation_timestamp
+        self._last_updated_timestamp = last_updated_timestamp
+        self._description = description
+        self._latest_version = latest_versions
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+        self._aliases = {alias.alias: alias.version for alias in (aliases or [])}
+        self._deployment_job_id = deployment_job_id
+        self._deployment_job_state = deployment_job_state
+
+    @property
+    def name(self):
+        """String. Registered model name."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name):
+        self._name = new_name
+
+    @property
+    def creation_timestamp(self):
+        """Integer. Model version creation timestamp (milliseconds since the Unix epoch)."""
+        return self._creation_time
+
+    @property
+    def last_updated_timestamp(self):
+        """Integer. Timestamp of last update for this model version (milliseconds since the Unix
+        epoch).
+        """
+        return self._last_updated_timestamp
+
+    @last_updated_timestamp.setter
+    def last_updated_timestamp(self, updated_timestamp):
+        self._last_updated_timestamp = updated_timestamp
+
+    @property
+    def description(self):
+        """String. Description"""
+        return self._description
+
+    @description.setter
+    def description(self, description):
+        self._description = description
+
+    @property
+    def latest_versions(self):
+        """List of the latest :py:class:`mlflow.entities.model_registry.ModelVersion` instances
+        for each stage.
+        """
+        return self._latest_version
+
+    @latest_versions.setter
+    def latest_versions(self, latest_versions):
+        self._latest_version = latest_versions
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current registered model."""
+        # Remove the is_prompt tag as it should not be user-facing
+        return {k: v for k, v in self._tags.items() if k != IS_PROMPT_TAG_KEY}
+
+    @property
+    def aliases(self):
+        """Dictionary of aliases (string) -> version for the current registered model."""
+        return self._aliases
+
+    @classmethod
+    def _properties(cls):
+        # aggregate with base class properties since cls.__dict__ does not do it automatically
+        return sorted(cls._get_properties_helper())
+
+    def _add_tag(self, tag):
+        self._tags[tag.key] = tag.value
+
+    def _add_alias(self, alias):
+        self._aliases[alias.alias] = alias.version
+
+    @property
+    def deployment_job_id(self):
+        """Deployment job ID for the current registered model."""
+        return self._deployment_job_id
+
+    @deployment_job_id.setter
+    def deployment_job_id(self, deployment_job_id):
+        self._deployment_job_id = deployment_job_id
+
+    @property
+    def deployment_job_state(self):
+        """Deployment job state for the current registered model."""
+        return self._deployment_job_state
+
+    # proto mappers
+    @classmethod
+    def from_proto(cls, proto):
+        # input: mlflow.protos.model_registry_pb2.RegisteredModel
+        # returns RegisteredModel entity
+        registered_model = cls(
+            proto.name,
+            proto.creation_timestamp,
+            proto.last_updated_timestamp,
+            proto.description,
+            [ModelVersion.from_proto(mvd) for mvd in proto.latest_versions],
+        )
+        for tag in proto.tags:
+            registered_model._add_tag(RegisteredModelTag.from_proto(tag))
+        for alias in proto.aliases:
+            registered_model._add_alias(RegisteredModelAlias.from_proto(alias))
+        registered_model._deployment_job_id = proto.deployment_job_id
+        registered_model._deployment_job_state = RegisteredModelDeploymentJobState.to_string(
+            proto.deployment_job_state
+        )
+        return registered_model
+
+    def to_proto(self):
+        # returns mlflow.protos.model_registry_pb2.RegisteredModel
+        rmd = ProtoRegisteredModel()
+        rmd.name = self.name
+        if self.creation_timestamp is not None:
+            rmd.creation_timestamp = self.creation_timestamp
+        if self.last_updated_timestamp:
+            rmd.last_updated_timestamp = self.last_updated_timestamp
+        if self.description:
+            rmd.description = self.description
+        if self.latest_versions is not None:
+            rmd.latest_versions.extend(
+                [model_version.to_proto() for model_version in self.latest_versions]
+            )
+        if self.deployment_job_id:
+            rmd.deployment_job_id = self.deployment_job_id
+        if self.deployment_job_state:
+            rmd.deployment_job_state = RegisteredModelDeploymentJobState.from_string(
+                self.deployment_job_state
+            )
+        rmd.tags.extend(
+            [ProtoRegisteredModelTag(key=key, value=value) for key, value in self._tags.items()]
+        )
+        rmd.aliases.extend(
+            [
+                ProtoRegisteredModelAlias(alias=alias, version=str(version))
+                for alias, version in self._aliases.items()
+            ]
+        )
+        return rmd

mlflow/entities/model_registry/registered_model_alias.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import RegisteredModelAlias as ProtoRegisteredModelAlias
+
+
+class RegisteredModelAlias(_ModelRegistryEntity):
+    """Alias object associated with a registered model."""
+
+    def __init__(self, alias, version):
+        self._alias = alias
+        self._version = version
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def alias(self):
+        """String name of the alias."""
+        return self._alias
+
+    @property
+    def version(self):
+        """String model version number that the alias points to."""
+        return self._version
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.alias, proto.version)
+
+    def to_proto(self):
+        alias_proto = ProtoRegisteredModelAlias()
+        alias_proto.alias = self.alias
+        alias_proto.version = self.version
+        return alias_proto

mlflow/entities/model_registry/registered_model_deployment_job_state.py

@@ -0,0 +1,39 @@
+from mlflow.protos.databricks_uc_registry_messages_pb2 import DeploymentJobConnection
+
+
+class RegisteredModelDeploymentJobState:
+    """Enum for registered model deployment state of an
+    :py:class:`mlflow.entities.model_registry.RegisteredModel`.
+    """
+
+    NOT_SET_UP = DeploymentJobConnection.State.Value("NOT_SET_UP")
+    CONNECTED = DeploymentJobConnection.State.Value("CONNECTED")
+    NOT_FOUND = DeploymentJobConnection.State.Value("NOT_FOUND")
+    REQUIRED_PARAMETERS_CHANGED = DeploymentJobConnection.State.Value("REQUIRED_PARAMETERS_CHANGED")
+    _STRING_TO_STATE = {
+        k: DeploymentJobConnection.State.Value(k) for k in DeploymentJobConnection.State.keys()
+    }
+    _STATE_TO_STRING = {value: key for key, value in _STRING_TO_STATE.items()}
+
+    @staticmethod
+    def from_string(state_str):
+        if state_str not in RegisteredModelDeploymentJobState._STRING_TO_STATE:
+            raise Exception(
+                f"Could not get deployment job connection state corresponding to string "
+                f"{state_str}. "
+                f"Valid state strings: {RegisteredModelDeploymentJobState.all_states()}"
+            )
+        return RegisteredModelDeploymentJobState._STRING_TO_STATE[state_str]
+
+    @staticmethod
+    def to_string(state):
+        if state not in RegisteredModelDeploymentJobState._STATE_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to deployment job connection {state}. "
+                f"Valid states: {RegisteredModelDeploymentJobState.all_states()}"
+            )
+        return RegisteredModelDeploymentJobState._STATE_TO_STRING[state]
+
+    @staticmethod
+    def all_states():
+        return list(RegisteredModelDeploymentJobState._STATE_TO_STRING.keys())

mlflow/entities/model_registry/registered_model_search.py

@@ -0,0 +1,25 @@
+from mlflow.entities.model_registry import RegisteredModel
+
+
+class RegisteredModelSearch(RegisteredModel):
+    def __init__(self, *args, **kwargs):
+        kwargs["tags"] = []
+        kwargs["aliases"] = []
+        super().__init__(*args, **kwargs)
+
+    def tags(self):
+        raise Exception(
+            "UC Registered Models gathered through search_registered_models do not have tags. "
+            "Please use get_registered_model to obtain an individual model's tags."
+        )
+
+    def aliases(self):
+        raise Exception(
+            "UC Registered Models gathered through search_registered_models do not have aliases. "
+            "Please use get_registered_model to obtain an individual model's aliases."
+        )
+
+    def __eq__(self, other):
+        if type(other) in {type(self), RegisteredModel}:
+            return self.__dict__ == other.__dict__
+        return False

mlflow/entities/model_registry/registered_model_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import RegisteredModelTag as ProtoRegisteredModelTag
+
+
+class RegisteredModelTag(_ModelRegistryEntity):
+    """Tag object associated with a registered model."""
+
+    def __init__(self, key, value):
+        self._key = key
+        self._value = value
+
+    def __eq__(self, other):
+        if type(other) is type(self):
+            return self.__dict__ == other.__dict__
+        return False
+
+    @property
+    def key(self):
+        """String name of the tag."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the tag."""
+        return self._value
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def to_proto(self):
+        tag = ProtoRegisteredModelTag()
+        tag.key = self.key
+        tag.value = self.value
+        return tag

mlflow/entities/multipart_upload.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from mlflow.protos.mlflow_artifacts_pb2 import (
+    CreateMultipartUpload as ProtoCreateMultipartUpload,
+)
+from mlflow.protos.mlflow_artifacts_pb2 import (
+    MultipartUploadCredential as ProtoMultipartUploadCredential,
+)
+
+
+@dataclass
+class MultipartUploadPart:
+    part_number: int
+    etag: str
+    url: Optional[str] = None
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(
+            proto.part_number,
+            proto.etag or None,
+            proto.url or None,
+        )
+
+    def to_dict(self):
+        return {
+            "part_number": self.part_number,
+            "etag": self.etag,
+            "url": self.url,
+        }
+
+
+@dataclass
+class MultipartUploadCredential:
+    url: str
+    part_number: int
+    headers: dict[str, Any]
+
+    def to_proto(self):
+        credential = ProtoMultipartUploadCredential()
+        credential.url = self.url
+        credential.part_number = self.part_number
+        credential.headers.update(self.headers)
+        return credential
+
+    @classmethod
+    def from_dict(cls, dict_):
+        return cls(
+            url=dict_["url"],
+            part_number=dict_["part_number"],
+            headers=dict_.get("headers", {}),
+        )
+
+
+@dataclass
+class CreateMultipartUploadResponse:
+    upload_id: Optional[str]
+    credentials: list[MultipartUploadCredential]
+
+    def to_proto(self):
+        response = ProtoCreateMultipartUpload.Response()
+        if self.upload_id:
+            response.upload_id = self.upload_id
+        response.credentials.extend([credential.to_proto() for credential in self.credentials])
+        return response
+
+    @classmethod
+    def from_dict(cls, dict_):
+        credentials = [MultipartUploadCredential.from_dict(cred) for cred in dict_["credentials"]]
+        return cls(
+            upload_id=dict_.get("upload_id"),
+            credentials=credentials,
+        )

mlflow/entities/param.py

@@ -0,0 +1,49 @@
+import sys
+
+from mlflow.entities._mlflow_object import _MlflowObject
+from mlflow.protos.service_pb2 import Param as ProtoParam
+
+
+class Param(_MlflowObject):
+    """
+    Parameter object.
+    """
+
+    def __init__(self, key, value):
+        if "pyspark.ml" in sys.modules:
+            import pyspark.ml.param
+
+            if isinstance(key, pyspark.ml.param.Param):
+                key = key.name
+                value = str(value)
+        self._key = key
+        self._value = value
+
+    @property
+    def key(self):
+        """String key corresponding to the parameter name."""
+        return self._key
+
+    @property
+    def value(self):
+        """String value of the parameter."""
+        return self._value
+
+    def to_proto(self):
+        param = ProtoParam()
+        param.key = self.key
+        param.value = self.value
+        return param
+
+    @classmethod
+    def from_proto(cls, proto):
+        return cls(proto.key, proto.value)
+
+    def __eq__(self, __o):
+        if isinstance(__o, self.__class__):
+            return self._key == __o._key
+
+        return False
+
+    def __hash__(self):
+        return hash(self._key)