mlrun 1.10.0rc13__py3-none-any.whl → 1.10.0rc15__py3-none-any.whl

mlrun/artifacts/base.py

@@ -16,7 +16,6 @@ import os
 import pathlib
 import tempfile
 import typing
-import warnings
 import zipfile
 
 import yaml
@@ -876,36 +875,6 @@ def generate_target_path(item: Artifact, artifact_path, producer):
     return f"{artifact_path}{item.key}{suffix}"
 
 
-# TODO: Remove once data migration v5 is obsolete
-def convert_legacy_artifact_to_new_format(
-    legacy_artifact: dict,
-) -> Artifact:
-    """Converts a legacy artifact to a new format.
-    :param legacy_artifact: The legacy artifact to convert.
-    :return: The converted artifact.
-    """
-    artifact_key = legacy_artifact.get("key", "")
-    artifact_tag = legacy_artifact.get("tag", "")
-    if artifact_tag:
-        artifact_key = f"{artifact_key}:{artifact_tag}"
-    # TODO: Remove once data migration v5 is obsolete
-    warnings.warn(
-        f"Converting legacy artifact '{artifact_key}' to new format. This will not be supported in MLRun 1.10.0. "
-        f"Make sure to save the artifact/project in the new format.",
-        FutureWarning,
-    )
-
-    artifact = mlrun.artifacts.artifact_types.get(
-        legacy_artifact.get("kind", "artifact"), mlrun.artifacts.Artifact
-    )()
-
-    artifact.metadata = artifact.metadata.from_dict(legacy_artifact)
-    artifact.spec = artifact.spec.from_dict(legacy_artifact)
-    artifact.status = artifact.status.from_dict(legacy_artifact)
-
-    return artifact
-
-
 def fill_artifact_object_hash(object_dict, iteration=None, producer_id=None):
     # remove artifact related fields before calculating hash
     object_dict.setdefault("metadata", {})

mlrun/artifacts/llm_prompt.py

@@ -11,12 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
 import tempfile
 from typing import Optional, Union
 
 import mlrun
 import mlrun.artifacts.model as model_art
-import mlrun.common
+import mlrun.common.schemas
 from mlrun.artifacts import Artifact, ArtifactMetadata, ArtifactSpec
 from mlrun.utils import StorePrefix, logger
 
@@ -25,16 +26,18 @@ MAX_PROMPT_LENGTH = 1024
 
 class LLMPromptArtifactSpec(ArtifactSpec):
     _dict_fields = ArtifactSpec._dict_fields + [
-        "prompt_string",
+        "prompt_template",
         "prompt_legend",
         "model_configuration",
         "description",
     ]
+    PROMPT_TEMPLATE_KEYS = ("content", "role")
+    PROMPT_LEGENDS_KEYS = ("field", "description")
 
     def __init__(
         self,
         model_artifact: Union[model_art.ModelArtifact, str] = None,
-        prompt_string: Optional[str] = None,
+        prompt_template: Optional[list[dict]] = None,
         prompt_path: Optional[str] = None,
         prompt_legend: Optional[dict] = None,
         model_configuration: Optional[dict] = None,
@@ -42,22 +45,26 @@ class LLMPromptArtifactSpec(ArtifactSpec):
         target_path: Optional[str] = None,
         **kwargs,
     ):
-        if prompt_string and prompt_path:
+        if prompt_template and prompt_path:
             raise mlrun.errors.MLRunInvalidArgumentError(
-                "Cannot specify both 'prompt_string' and 'prompt_path'"
+                "Cannot specify both 'prompt_template' and 'prompt_path'"
             )
-
+        if prompt_legend:
+            self._verify_prompt_legend(prompt_legend)
+        if prompt_path:
+            self._verify_prompt_path(prompt_path)
+        if prompt_template:
+            self._verify_prompt_template(prompt_template)
         super().__init__(
             src_path=prompt_path,
             target_path=target_path,
             parent_uri=model_artifact.uri
             if isinstance(model_artifact, model_art.ModelArtifact)
             else model_artifact,
-            body=prompt_string,
             **kwargs,
         )
 
-        self.prompt_string = prompt_string
+        self.prompt_template = prompt_template
         self.prompt_legend = prompt_legend
         self.model_configuration = model_configuration
         self.description = description
@@ -67,10 +74,78 @@ class LLMPromptArtifactSpec(ArtifactSpec):
             else None
         )
 
+    def _verify_prompt_template(self, prompt_template):
+        if not (
+            isinstance(prompt_template, list)
+            and all(isinstance(item, dict) for item in prompt_template)
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Expected prompt_template to be a list of dicts"
+            )
+        keys_to_pop = []
+        for message in prompt_template:
+            for key in message.keys():
+                if isinstance(key, str):
+                    if key.lower() not in self.PROMPT_TEMPLATE_KEYS:
+                        raise mlrun.errors.MLRunInvalidArgumentError(
+                            f"Expected prompt_template to contain dict that "
+                            f"only has keys from {self.PROMPT_TEMPLATE_KEYS}"
+                        )
+                    else:
+                        if not key.islower():
+                            message[key.lower()] = message[key]
+                            keys_to_pop.append(key)
+                else:
+                    raise mlrun.errors.MLRunInvalidArgumentError(
+                        f"Expected prompt_template to contain dict that only"
+                        f" has str keys got {key} of type {type(key)}"
+                    )
+            for key_to_pop in keys_to_pop:
+                message.pop(key_to_pop)
+
     @property
     def model_uri(self):
         return self.parent_uri
 
+    @staticmethod
+    def _verify_prompt_legend(prompt_legend: dict):
+        if prompt_legend is None:
+            return True
+        for place_holder, body_map in prompt_legend.items():
+            if isinstance(body_map, dict):
+                if body_map.get("field") is None:
+                    body_map["field"] = place_holder
+                body_map["description"] = body_map.get("description")
+                if diff := set(body_map.keys()) - set(
+                    LLMPromptArtifactSpec.PROMPT_LEGENDS_KEYS
+                ):
+                    raise mlrun.errors.MLRunInvalidArgumentError(
+                        "prompt_legend values must contain only 'field' and "
+                        f"'description' keys, got extra fields: {diff}"
+                    )
+            else:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Wrong prompt_legend format, {place_holder} is not mapped to dict"
+                )
+
+    @staticmethod
+    def _verify_prompt_path(prompt_path: str):
+        with mlrun.datastore.store_manager.object(prompt_path).open(mode="r") as p_file:
+            try:
+                json.load(p_file)
+            except json.JSONDecodeError:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Failed on decoding str in path "
+                    f"{prompt_path} expected file to contain a "
+                    f"json format."
+                )
+
+    def get_body(self):
+        if self.prompt_template:
+            return json.dumps(self.prompt_template)
+        else:
+            return None
+
 
 class LLMPromptArtifact(Artifact):
     """
@@ -90,7 +165,7 @@ class LLMPromptArtifact(Artifact):
         model_artifact: Union[
             model_art.ModelArtifact, str
         ] = None,  # TODO support partial model uri
-        prompt_string: Optional[str] = None,
+        prompt_template: Optional[list[dict]] = None,
         prompt_path: Optional[str] = None,
         prompt_legend: Optional[dict] = None,
         model_configuration: Optional[dict] = None,
@@ -99,7 +174,7 @@ class LLMPromptArtifact(Artifact):
         **kwargs,
     ):
         llm_prompt_spec = LLMPromptArtifactSpec(
-            prompt_string=prompt_string,
+            prompt_template=prompt_template,
             prompt_path=prompt_path,
             prompt_legend=prompt_legend,
             model_artifact=model_artifact,
@@ -137,33 +212,44 @@ class LLMPromptArtifact(Artifact):
             return self.spec._model_artifact
         return None
 
-    def read_prompt(self) -> Optional[str]:
+    def read_prompt(self) -> Optional[Union[str, list[dict]]]:
         """
-        Read the prompt string from the artifact.
+        Read the prompt json from the artifact or if provided prompt template.
+        @:param as_str: True to return the prompt string or a list of dicts.
+        @:return prompt string or list of dicts
         """
-        if self.spec.prompt_string:
-            return self.spec.prompt_string
+        if self.spec.prompt_template:
+            return self.spec.prompt_template
         if self.spec.target_path:
             with mlrun.datastore.store_manager.object(url=self.spec.target_path).open(
                 mode="r"
             ) as p_file:
-                return p_file.read()
+                try:
+                    return json.load(p_file)
+                except json.JSONDecodeError:
+                    raise mlrun.errors.MLRunInvalidArgumentError(
+                        f"Failed on decoding str in path "
+                        f"{self.spec.target_path} expected file to contain a "
+                        f"json format."
+                    )
 
     def before_log(self):
         """
         Prepare the artifact before logging.
         This method is called before the artifact is logged.
         """
-        if self.spec.prompt_string and len(self.spec.prompt_string) > MAX_PROMPT_LENGTH:
+        if (
+            self.spec.prompt_template
+            and len(str(self.spec.prompt_template)) > MAX_PROMPT_LENGTH
+        ):
             logger.debug(
                 "Prompt string exceeds maximum length, saving to a temporary file."
             )
             with tempfile.NamedTemporaryFile(
-                delete=False, mode="w", suffix=".txt"
+                delete=False, mode="w", suffix=".json"
             ) as temp_file:
-                temp_file.write(self.spec.prompt_string)
+                temp_file.write(json.dumps(self.spec.prompt_template))
             self.spec.src_path = temp_file.name
-            self.spec.prompt_string = None
+            self.spec.prompt_template = None
             self._src_is_temp = True
-
         super().before_log()

mlrun/artifacts/manager.py

@@ -110,11 +110,6 @@ class ArtifactProducer:
 
 def dict_to_artifact(struct: dict) -> Artifact:
     kind = struct.get("kind", "")
-
-    # TODO: Remove once data migration v5 is obsolete
-    if mlrun.utils.is_legacy_artifact(struct):
-        return mlrun.artifacts.base.convert_legacy_artifact_to_new_format(struct)
-
     artifact_class = artifact_types[kind]
     return artifact_class.from_dict(struct)
 

mlrun/common/constants.py

@@ -81,7 +81,6 @@ class MLRunInternalLabels:
     kind = "kind"
     component = "component"
     mlrun_type = "mlrun__type"
-    rerun_of = "rerun-of"
     original_workflow_id = "original-workflow-id"
     workflow_id = "workflow-id"
 

mlrun/common/schemas/__init__.py

@@ -147,6 +147,7 @@ from .model_monitoring import (
     GrafanaTable,
     ModelEndpoint,
     ModelEndpointCreationStrategy,
+    ModelEndpointDriftValues,
     ModelEndpointList,
     ModelEndpointMetadata,
     ModelEndpointSchema,

mlrun/common/schemas/model_monitoring/__init__.py

@@ -59,6 +59,7 @@ from .model_endpoints import (
     Features,
     FeatureValues,
     ModelEndpoint,
+    ModelEndpointDriftValues,
     ModelEndpointList,
     ModelEndpointMetadata,
     ModelEndpointMonitoringMetric,

mlrun/common/schemas/model_monitoring/functions.py

@@ -64,5 +64,5 @@ class FunctionSummary(BaseModel):
             updated_time=func_dict["metadata"].get("updated"),
             status=func_dict["status"].get("state"),
             base_period=base_period,
-            stats=stats,
+            stats=stats or {},
         )

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -352,6 +352,16 @@ class ApplicationMetricRecord(ApplicationBaseRecord):
     type: Literal["metric"] = "metric"
 
 
+class _DriftBin(NamedTuple):
+    timestamp: datetime
+    count_suspected: int
+    count_detected: int
+
+
+class ModelEndpointDriftValues(BaseModel):
+    values: list[_DriftBin]
+
+
 def _mapping_attributes(
     model_class: type[Model],
     flattened_dictionary: dict,

mlrun/common/schemas/workflow.py

@@ -49,7 +49,6 @@ class WorkflowRequest(pydantic.v1.BaseModel):
 class RerunWorkflowRequest(pydantic.v1.BaseModel):
     run_name: typing.Optional[str] = None
     run_id: typing.Optional[str] = None
-    original_workflow_id: typing.Optional[str] = None
     notifications: typing.Optional[list[Notification]] = None
     workflow_runner_node_selector: typing.Optional[dict[str, str]] = None
 

mlrun/config.py

@@ -193,7 +193,7 @@ default_config = {
     },
     "v3io_framesd": "http://framesd:8080",
     "model_providers": {
-        "openai_default_model": "gpt-4",
+        "openai_default_model": "gpt-4o",
     },
     # default node selector to be applied to all functions - json string base64 encoded format
     "default_function_node_selector": "e30=",

mlrun/datastore/model_provider/model_provider.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Awaitable
-from typing import Callable, Optional, TypeVar, Union
+from typing import Any, Callable, Optional, TypeVar, Union
 
 import mlrun.errors
 from mlrun.datastore.remote_client import (
@@ -56,9 +56,16 @@ class ModelProvider(BaseRemoteClient):
         )
         self.default_invoke_kwargs = default_invoke_kwargs or {}
         self._client = None
-        self._default_operation = None
         self._async_client = None
-        self._default_async_operation = None
+
+    def get_client_options(self) -> dict:
+        """
+        Returns a dictionary containing credentials and configuration
+        options required for client creation.
+
+        :return:           A dictionary with client-specific settings.
+        """
+        return {}
 
     def load_client(self) -> None:
         """
@@ -68,8 +75,6 @@ class ModelProvider(BaseRemoteClient):
         Subclasses should override this method to:
         - Create and configure the provider-specific client instance.
         - Assign the client instance to self._client.
-        - Define a default operation callable (e.g., a method to invoke model completions)
-        and assign it to self._default_operation.
         """
 
         raise NotImplementedError("load_client method is not implemented")
@@ -122,39 +127,62 @@ class ModelProvider(BaseRemoteClient):
         """
         raise NotImplementedError("invoke method is not implemented")
 
-    def customized_invoke(
+    def custom_invoke(
         self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
     ) -> Optional[T]:
-        raise NotImplementedError("customized_invoke method is not implemented")
+        """
+        Invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.) with the given keyword arguments.
+
+        Useful for dynamically calling model methods like text generation, chat completions, or image generation.
+        The operation must be a callable that accepts keyword arguments.
+
+        :param operation:       A callable representing the model operation (e.g., a client method).
+        :param invoke_kwargs:   Keyword arguments to pass to the operation.
+        :return:                The full response returned by the operation.
+        """
+        raise NotImplementedError("custom_invoke method is not implemented")
 
     @property
-    def client(self):
+    def client(self) -> Any:
         return self._client
 
     @property
-    def model(self):
+    def model(self) -> Optional[str]:
         return None
 
-    def get_invoke_kwargs(self, invoke_kwargs):
+    def get_invoke_kwargs(self, invoke_kwargs) -> dict:
         kwargs = self.default_invoke_kwargs.copy()
         kwargs.update(invoke_kwargs)
         return kwargs
 
     @property
-    def async_client(self):
+    def async_client(self) -> Any:
         if not self.support_async:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"{self.__class__.__name__} does not support async operations"
             )
         return self._async_client
 
-    async def async_customized_invoke(self, **kwargs):
-        raise NotImplementedError("async_customized_invoke is not implemented")
+    async def async_custom_invoke(
+        self, operation: Optional[Callable[..., Awaitable[T]]], **invoke_kwargs
+    ) -> Optional[T]:
+        """
+        Asynchronously invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.)
+        with the given keyword arguments.
+
+        The operation must be an async callable (e.g., a method from an async client) that accepts keyword arguments.
+
+        :param operation:       An async callable representing the model operation (e.g., an async_client method).
+        :param invoke_kwargs:   Keyword arguments to pass to the operation.
+        :return:                The full response returned by the awaited operation.
+        """
+        raise NotImplementedError("async_custom_invoke is not implemented")
 
     async def async_invoke(
         self,
         messages: Optional[list[dict]] = None,
         as_str: bool = False,
         **invoke_kwargs,
-    ) -> Awaitable[str]:
+    ) -> Optional[str]:
+        """Async version of `invoke`. See `invoke` for full documentation."""
         raise NotImplementedError("async_invoke is not implemented")

mlrun/datastore/model_provider/openai_provider.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from collections.abc import Awaitable
 from typing import Callable, Optional, TypeVar, Union
 
 import mlrun
@@ -33,6 +33,8 @@ class OpenAIProvider(ModelProvider):
     operations tailored to the OpenAI API.
     """
 
+    support_async = True
+
     def __init__(
         self,
         parent,
@@ -67,7 +69,7 @@ class OpenAIProvider(ModelProvider):
         return endpoint, subpath
 
     @property
-    def model(self):
+    def model(self) -> Optional[str]:
         return self.endpoint
 
     def load_client(self) -> None:
@@ -76,23 +78,20 @@ class OpenAIProvider(ModelProvider):
 
         This method imports the `OpenAI` class from the `openai` package, instantiates
         a client with the given keyword arguments (`self.options`), and assigns it to
-        `self._client`.
-
-        It also sets the default operation to `self.client.chat.completions.create`, which is
-        typically used for invoking chat-based model completions.
+        `self._client` and `self._async_client`.
 
         Raises:
             ImportError: If the `openai` package is not installed.
         """
         try:
-            from openai import OpenAI  # noqa
+            from openai import OpenAI, AsyncOpenAI  # noqa
 
             self._client = OpenAI(**self.options)
-            self._default_operation = self.client.chat.completions.create
+            self._async_client = AsyncOpenAI(**self.options)
         except ImportError as exc:
             raise ImportError("openai package is not installed") from exc
 
-    def get_client_options(self):
+    def get_client_options(self) -> dict:
         res = dict(
             api_key=self._get_secret_or_env("OPENAI_API_KEY"),
             organization=self._get_secret_or_env("OPENAI_ORG_ID"),
@@ -103,14 +102,69 @@ class OpenAIProvider(ModelProvider):
         )
         return self._sanitize_options(res)
 
-    def customized_invoke(
+    def custom_invoke(
         self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
     ) -> Optional[T]:
+        """
+        OpenAI-specific implementation of `ModelProvider.custom_invoke`.
+
+        Invokes an OpenAI model operation using the sync client. For full details, see
+        `ModelProvider.custom_invoke`.
+
+        Example:
+            ```python
+            result = openai_model_provider.invoke(
+                openai_model_provider.client.images.generate,
+                prompt="A futuristic cityscape at sunset",
+                n=1,
+                size="1024x1024",
+            )
+            ```
+        :param      operation:      Same as ModelProvider.custom_invoke.
+        :param      invoke_kwargs:  Same as ModelProvider.custom_invoke.
+        :return:                    Same as ModelProvider.custom_invoke.
+
+        """
         invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
         if operation:
             return operation(**invoke_kwargs, model=self.model)
         else:
-            return self._default_operation(**invoke_kwargs, model=self.model)
+            return self.client.chat.completions.create(
+                **invoke_kwargs, model=self.model
+            )
+
+    async def async_custom_invoke(
+        self,
+        operation: Optional[Callable[..., Awaitable[T]]] = None,
+        **invoke_kwargs,
+    ) -> Optional[T]:
+        """
+        OpenAI-specific implementation of `ModelProvider.async_custom_invoke`.
+
+        Invokes an OpenAI model operation using the async client. For full details, see
+        `ModelProvider.async_custom_invoke`.
+
+        Example:
+            ```python
+            result = openai_model_provider.invoke(
+                openai_model_provider.async_client.images.generate,
+                prompt="A futuristic cityscape at sunset",
+                n=1,
+                size="1024x1024",
+            )
+            ```
+        :param operation:       Same as ModelProvider.async_custom_invoke.
+        :param invoke_kwargs:   Same as ModelProvider.async_custom_invoke.
+        :return:                Same as ModelProvider.async_custom_invoke.
+
+        """
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        if operation:
+            return await operation(**invoke_kwargs, model=self.model)
+        else:
+            return await self.async_client.chat.completions.create(
+                **invoke_kwargs, model=self.model
+            )
 
     def invoke(
         self,
@@ -133,12 +187,39 @@ class OpenAIProvider(ModelProvider):
 
         :param invoke_kwargs:
                             Same as ModelProvider.invoke.
+        :return:            Same as ModelProvider.invoke.
 
         """
-        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
-        response = self._default_operation(
-            model=self.endpoint, messages=messages, **invoke_kwargs
-        )
+        response = self.custom_invoke(messages=messages, **invoke_kwargs)
+        if as_str:
+            return response.choices[0].message.content
+        return response
+
+    async def async_invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> str:
+        """
+        OpenAI-specific implementation of `ModelProvider.async_invoke`.
+        Invokes an OpenAI model operation using the async client.
+        For full details, see `ModelProvider.async_invoke`.
+
+        :param messages:    Same as ModelProvider.async_invoke.
+
+        :param as_str: bool
+                            If `True`, returns only the main content of the first response
+                            (`response.choices[0].message.content`).
+                            If `False`, returns the full awaited response object, whose type depends on
+                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+
+        :param invoke_kwargs:
+                            Same as ModelProvider.async_invoke.
+        :returns            Same as ModelProvider.async_invoke.
+
+        """
+        response = await self.async_custom_invoke(messages=messages, **invoke_kwargs)
         if as_str:
             return response.choices[0].message.content
         return response

mlrun/db/base.py

@@ -638,6 +638,11 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    def wait_for_background_task_to_reach_terminal_state(
+        self, name: str, project: str = ""
+    ) -> mlrun.common.schemas.BackgroundTask:
+        pass
+
     @abstractmethod
     def retry_pipeline(
         self,
@@ -1145,3 +1150,12 @@ class RunDBInterface(ABC):
     @abstractmethod
     def get_project_summary(self, project: str) -> mlrun.common.schemas.ProjectSummary:
         pass
+
+    @abstractmethod
+    def get_drift_over_time(
+        self,
+        project: str,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+    ) -> mlrun.common.schemas.model_monitoring.ModelEndpointDriftValues:
+        pass