star

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

@@ -0,0 +1,19 @@
+from mlflow.entities.model_registry.model_version import ModelVersion
+from mlflow.entities.model_registry.model_version_search import ModelVersionSearch
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.entities.model_registry.prompt import Prompt
+from mlflow.entities.model_registry.registered_model import RegisteredModel
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+from mlflow.entities.model_registry.registered_model_search import RegisteredModelSearch
+from mlflow.entities.model_registry.registered_model_tag import RegisteredModelTag
+
+__all__ = [
+    "Prompt",
+    "RegisteredModel",
+    "ModelVersion",
+    "RegisteredModelAlias",
+    "RegisteredModelTag",
+    "ModelVersionTag",
+    "RegisteredModelSearch",
+    "ModelVersionSearch",
+]

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/_model_registry_entity.py

@@ -0,0 +1,13 @@
+from abc import abstractmethod
+
+from mlflow.entities._mlflow_object import _MlflowObject
+
+
+class _ModelRegistryEntity(_MlflowObject):
+    @classmethod
+    @abstractmethod
+    def from_proto(cls, proto):
+        pass
+
+    def __eq__(self, other):
+        return dict(self) == dict(other)

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version.py

@@ -0,0 +1,199 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.entities.model_registry.model_version_status import ModelVersionStatus
+from mlflow.entities.model_registry.model_version_tag import ModelVersionTag
+from mlflow.protos.model_registry_pb2 import ModelVersion as ProtoModelVersion
+from mlflow.protos.model_registry_pb2 import ModelVersionTag as ProtoModelVersionTag
+
+
+class ModelVersion(_ModelRegistryEntity):
+    """
+    MLflow entity for Model Version.
+    """
+
+    def __init__(
+        self,
+        name,
+        version,
+        creation_timestamp,
+        last_updated_timestamp=None,
+        description=None,
+        user_id=None,
+        current_stage=None,
+        source=None,
+        run_id=None,
+        status=ModelVersionStatus.to_string(ModelVersionStatus.READY),
+        status_message=None,
+        tags=None,
+        run_link=None,
+        aliases=None,
+    ):
+        super().__init__()
+        self._name = name
+        self._version = version
+        self._creation_time = creation_timestamp
+        self._last_updated_timestamp = last_updated_timestamp
+        self._description = description
+        self._user_id = user_id
+        self._current_stage = current_stage
+        self._source = source
+        self._run_id = run_id
+        self._run_link = run_link
+        self._status = status
+        self._status_message = status_message
+        self._tags = {tag.key: tag.value for tag in (tags or [])}
+        self._aliases = aliases or []
+
+    @property
+    def name(self):
+        """String. Unique name within Model Registry."""
+        return self._name
+
+    @name.setter
+    def name(self, new_name):
+        self._name = new_name
+
+    @property
+    def version(self):
+        """version"""
+        return self._version
+
+    @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 user_id(self):
+        """String. User ID that created this model version."""
+        return self._user_id
+
+    @property
+    def current_stage(self):
+        """String. Current stage of this model version."""
+        return self._current_stage
+
+    @current_stage.setter
+    def current_stage(self, stage):
+        self._current_stage = stage
+
+    @property
+    def source(self):
+        """String. Source path for the model."""
+        return self._source
+
+    @property
+    def run_id(self):
+        """String. MLflow run ID that generated this model."""
+        return self._run_id
+
+    @property
+    def run_link(self):
+        """String. MLflow run link referring to the exact run that generated this model version."""
+        return self._run_link
+
+    @property
+    def status(self):
+        """String. Current Model Registry status for this model."""
+        return self._status
+
+    @property
+    def status_message(self):
+        """String. Descriptive message for error status conditions."""
+        return self._status_message
+
+    @property
+    def tags(self):
+        """Dictionary of tag key (string) -> tag value for the current model version."""
+        return self._tags
+
+    @property
+    def aliases(self):
+        """List of aliases (string) for the current model version."""
+        return self._aliases
+
+    @aliases.setter
+    def aliases(self, aliases):
+        self._aliases = 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
+
+    # proto mappers
+    @classmethod
+    def from_proto(cls, proto):
+        # input: mlflow.protos.model_registry_pb2.ModelVersion
+        # returns: ModelVersion entity
+        model_version = cls(
+            proto.name,
+            proto.version,
+            proto.creation_timestamp,
+            proto.last_updated_timestamp,
+            proto.description if proto.HasField("description") else None,
+            proto.user_id,
+            proto.current_stage,
+            proto.source,
+            proto.run_id if proto.HasField("run_id") else None,
+            ModelVersionStatus.to_string(proto.status),
+            proto.status_message if proto.HasField("status_message") else None,
+            run_link=proto.run_link,
+            aliases=proto.aliases,
+        )
+        for tag in proto.tags:
+            model_version._add_tag(ModelVersionTag.from_proto(tag))
+        return model_version
+
+    def to_proto(self):
+        # input: ModelVersion entity
+        # returns mlflow.protos.model_registry_pb2.ModelVersion
+        model_version = ProtoModelVersion()
+        model_version.name = self.name
+        model_version.version = str(self.version)
+        model_version.creation_timestamp = self.creation_timestamp
+        if self.last_updated_timestamp is not None:
+            model_version.last_updated_timestamp = self.last_updated_timestamp
+        if self.description is not None:
+            model_version.description = self.description
+        if self.user_id is not None:
+            model_version.user_id = self.user_id
+        if self.current_stage is not None:
+            model_version.current_stage = self.current_stage
+        if self.source is not None:
+            model_version.source = str(self.source)
+        if self.run_id is not None:
+            model_version.run_id = str(self.run_id)
+        if self.run_link is not None:
+            model_version.run_link = str(self.run_link)
+        if self.status is not None:
+            model_version.status = ModelVersionStatus.from_string(self.status)
+        if self.status_message:
+            model_version.status_message = self.status_message
+        model_version.tags.extend(
+            [ProtoModelVersionTag(key=key, value=value) for key, value in self._tags.items()]
+        )
+        model_version.aliases.extend(self.aliases)
+        return model_version

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_search.py

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

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_stages.py

@@ -0,0 +1,25 @@
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+STAGE_NONE = "None"
+STAGE_STAGING = "Staging"
+STAGE_PRODUCTION = "Production"
+STAGE_ARCHIVED = "Archived"
+
+STAGE_DELETED_INTERNAL = "Deleted_Internal"
+
+ALL_STAGES = [STAGE_NONE, STAGE_STAGING, STAGE_PRODUCTION, STAGE_ARCHIVED]
+DEFAULT_STAGES_FOR_GET_LATEST_VERSIONS = [STAGE_STAGING, STAGE_PRODUCTION]
+_CANONICAL_MAPPING = {stage.lower(): stage for stage in ALL_STAGES}
+
+
+def get_canonical_stage(stage):
+    key = stage.lower()
+    if key not in _CANONICAL_MAPPING:
+        raise MlflowException(
+            "Invalid Model Version stage: {}. Value must be one of {}.".format(
+                stage, ", ".join(ALL_STAGES)
+            ),
+            INVALID_PARAMETER_VALUE,
+        )
+    return _CANONICAL_MAPPING[key]

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_status.py

@@ -0,0 +1,35 @@
+from mlflow.protos.model_registry_pb2 import ModelVersionStatus as ProtoModelVersionStatus
+
+
+class ModelVersionStatus:
+    """Enum for status of an :py:class:`mlflow.entities.model_registry.ModelVersion`."""
+
+    PENDING_REGISTRATION = ProtoModelVersionStatus.Value("PENDING_REGISTRATION")
+    FAILED_REGISTRATION = ProtoModelVersionStatus.Value("FAILED_REGISTRATION")
+    READY = ProtoModelVersionStatus.Value("READY")
+    _STRING_TO_STATUS = {
+        k: ProtoModelVersionStatus.Value(k) for k in ProtoModelVersionStatus.keys()
+    }
+    _STATUS_TO_STRING = {value: key for key, value in _STRING_TO_STATUS.items()}
+
+    @staticmethod
+    def from_string(status_str):
+        if status_str not in ModelVersionStatus._STRING_TO_STATUS:
+            raise Exception(
+                f"Could not get model version status corresponding to string {status_str}. "
+                f"Valid status strings: {list(ModelVersionStatus._STRING_TO_STATUS.keys())}"
+            )
+        return ModelVersionStatus._STRING_TO_STATUS[status_str]
+
+    @staticmethod
+    def to_string(status):
+        if status not in ModelVersionStatus._STATUS_TO_STRING:
+            raise Exception(
+                f"Could not get string corresponding to model version status {status}. "
+                f"Valid statuses: {list(ModelVersionStatus._STATUS_TO_STRING.keys())}"
+            )
+        return ModelVersionStatus._STATUS_TO_STRING[status]
+
+    @staticmethod
+    def all_status():
+        return list(ModelVersionStatus._STATUS_TO_STRING.keys())

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/model_version_tag.py

@@ -0,0 +1,35 @@
+from mlflow.entities.model_registry._model_registry_entity import _ModelRegistryEntity
+from mlflow.protos.model_registry_pb2 import ModelVersionTag as ProtoModelVersionTag
+
+
+class ModelVersionTag(_ModelRegistryEntity):
+    """Tag object associated with a model version."""
+
+    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 = ProtoModelVersionTag()
+        tag.key = self.key
+        tag.value = self.value
+        return tag

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/prompt.py

@@ -0,0 +1,228 @@
+from __future__ import annotations
+
+import re
+from typing import Optional, Union
+
+from mlflow.entities.model_registry.model_version import ModelVersion
+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_ASSOCIATED_RUN_IDS_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, PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY}
+
+
+# Prompt is implemented as a special type of ModelVersion. MLflow stores both prompts
+# and model versions in the model registry as ModelVersion DB records, but distinguishes
+# them using the special tag "mlflow.prompt.is_prompt".
+class Prompt(ModelVersion):
+    """
+    An entity representing a prompt (template) for GenAI applications.
+
+    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.
+        version_metadata: A dictionary of metadata associated with the **prompt version**.
+            This is useful for storing version-specific information, such as the author of
+            the changes. Optional.
+        prompt_tags: A dictionary of tags associated with the entire prompt. This is different
+            from the `version_metadata` as it is not tied to a specific version of the prompt.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        version: int,
+        template: str,
+        commit_message: Optional[str] = None,
+        creation_timestamp: Optional[int] = None,
+        version_metadata: Optional[dict[str, str]] = None,
+        prompt_tags: Optional[dict[str, str]] = None,
+        aliases: Optional[list[str]] = None,
+    ):
+        # Store template text as a tag
+        version_metadata = version_metadata or {}
+        version_metadata[PROMPT_TEXT_TAG_KEY] = template
+        version_metadata[IS_PROMPT_TAG_KEY] = "true"
+
+        super().__init__(
+            name=name,
+            version=version,
+            creation_timestamp=creation_timestamp,
+            description=commit_message,
+            # "version_metadata" is represented as ModelVersion tags.
+            tags=[ModelVersionTag(key=key, value=value) for key, value in version_metadata.items()],
+            aliases=aliases,
+        )
+
+        self._variables = set(PROMPT_TEMPLATE_VARIABLE_PATTERN.findall(self.template))
+
+        # Store the prompt-level tags (from RegisteredModel).
+        self._prompt_tags = prompt_tags 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"Prompt(name={self.name}, version={self.version}, template={text})"
+
+    @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  # inherited from ModelVersion
+
+    @property
+    def version_metadata(self) -> dict[str, str]:
+        """Return the tags of the prompt as a dictionary."""
+        # Remove the prompt text tag as it should not be user-facing
+        return {key: value for key, value in self._tags.items() if not _is_reserved_tag(key)}
+
+    @property
+    def tags(self) -> dict[str, str]:
+        """
+        Return the prompt-level tags (from RegisteredModel).
+        """
+        return {key: value for key, value in self._prompt_tags.items() if not _is_reserved_tag(key)}
+
+    @property
+    def run_ids(self) -> list[str]:
+        """Get the run IDs associated with the prompt."""
+        run_tag = self._tags.get(PROMPT_ASSOCIATED_RUN_IDS_TAG_KEY)
+        if not run_tag:
+            return []
+        return run_tag.split(",")
+
+    @property
+    def uri(self) -> str:
+        """Return the URI of the prompt."""
+        return f"prompts:/{self.name}/{self.version}"
+
+    def format(self, allow_partial: bool = False, **kwargs) -> Union[Prompt, 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: Prompt(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 Prompt(
+                    name=self.name,
+                    version=self.version,
+                    template=template,
+                    commit_message=self.commit_message,
+                    creation_timestamp=self.creation_timestamp,
+                    prompt_tags=self._prompt_tags,
+                    version_metadata=self.version_metadata,
+                    aliases=self.aliases,
+                )
+        return template
+
+    @classmethod
+    def from_model_version(
+        cls, model_version: ModelVersion, prompt_tags: Optional[dict[str, str]] = None
+    ) -> Prompt:
+        """
+        Create a Prompt object from a ModelVersion object.
+
+        Args:
+            model_version: The ModelVersion object to convert to a Prompt.
+            prompt_tags: The prompt-level tags (from RegisteredModel). Optional.
+        """
+        if IS_PROMPT_TAG_KEY not in model_version.tags:
+            raise MlflowException.invalid_parameter_value(
+                f"Name `{model_version.name}` is registered as a model, not a prompt. MLflow "
+                "does not allow registering a prompt with the same name as an existing model.",
+            )
+
+        if PROMPT_TEXT_TAG_KEY not in model_version.tags:
+            raise MlflowException.invalid_parameter_value(
+                f"Prompt `{model_version.name}` does not contain a prompt text"
+            )
+
+        return cls(
+            name=model_version.name,
+            version=model_version.version,
+            template=model_version.tags[PROMPT_TEXT_TAG_KEY],
+            commit_message=model_version.description,
+            creation_timestamp=model_version.creation_timestamp,
+            version_metadata=model_version.tags,
+            prompt_tags=prompt_tags,
+            aliases=model_version.aliases,
+        )

venv/lib/python3.9/site-packages/mlflow/entities/model_registry/registered_model.py

@@ -0,0 +1,144 @@
+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 import IS_PROMPT_TAG_KEY
+from mlflow.entities.model_registry.registered_model_alias import RegisteredModelAlias
+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,
+    ):
+        # 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 [])}
+
+    @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
+
+    # 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))
+        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]
+            )
+        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

venv/lib/python3.9/site-packages/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

venv/lib/python3.9/site-packages/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

venv/lib/python3.9/site-packages/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