mlrun 1.10.0rc12__tar.gz → 1.10.0rc14__tar.gz

{mlrun-1.10.0rc12/mlrun.egg-info → mlrun-1.10.0rc14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlrun
-Version: 1.10.0rc12
+Version: 1.10.0rc14
 Summary: Tracking and config of machine learning runs
 Home-page: https://github.com/mlrun/mlrun
 Author: Yaron Haviv
@@ -44,7 +44,7 @@ Requires-Dist: semver~=3.0
 Requires-Dist: dependency-injector~=4.41
 Requires-Dist: fsspec<2024.7,>=2023.9.2
 Requires-Dist: v3iofs~=0.1.17
-Requires-Dist: storey~=1.10.7
+Requires-Dist: storey~=1.10.8
 Requires-Dist: inflection~=0.5.0
 Requires-Dist: python-dotenv~=1.0
 Requires-Dist: setuptools>=75.2

mlrun-1.10.0rc14/mlrun/artifacts/llm_prompt.py

@@ -0,0 +1,255 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.schemas
+from mlrun.artifacts import Artifact, ArtifactMetadata, ArtifactSpec
+from mlrun.utils import StorePrefix, logger
+
+MAX_PROMPT_LENGTH = 1024
+
+
+class LLMPromptArtifactSpec(ArtifactSpec):
+    _dict_fields = ArtifactSpec._dict_fields + [
+        "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_template: Optional[list[dict]] = None,
+        prompt_path: Optional[str] = None,
+        prompt_legend: Optional[dict] = None,
+        model_configuration: Optional[dict] = None,
+        description: Optional[str] = None,
+        target_path: Optional[str] = None,
+        **kwargs,
+    ):
+        if prompt_template and prompt_path:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "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,
+            **kwargs,
+        )
+
+        self.prompt_template = prompt_template
+        self.prompt_legend = prompt_legend
+        self.model_configuration = model_configuration
+        self.description = description
+        self._model_artifact = (
+            model_artifact
+            if isinstance(model_artifact, model_art.ModelArtifact)
+            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):
+    """
+    LLM Prompt Artifact
+
+    This artifact is used to store and manage LLM prompts.
+    Stores the prompt string/path and a link to the related model artifact.
+    """
+
+    kind = mlrun.common.schemas.ArtifactCategories.llm_prompt
+    _store_prefix = StorePrefix.LLMPrompt
+
+    def __init__(
+        self,
+        key: Optional[str] = None,
+        project: Optional[str] = None,
+        model_artifact: Union[
+            model_art.ModelArtifact, str
+        ] = None,  # TODO support partial model uri
+        prompt_template: Optional[list[dict]] = None,
+        prompt_path: Optional[str] = None,
+        prompt_legend: Optional[dict] = None,
+        model_configuration: Optional[dict] = None,
+        description: Optional[str] = None,
+        target_path=None,
+        **kwargs,
+    ):
+        llm_prompt_spec = LLMPromptArtifactSpec(
+            prompt_template=prompt_template,
+            prompt_path=prompt_path,
+            prompt_legend=prompt_legend,
+            model_artifact=model_artifact,
+            model_configuration=model_configuration,
+            target_path=target_path,
+            description=description,
+        )
+
+        llm_metadata = ArtifactMetadata(
+            key=key,
+            project=project or "",
+        )
+
+        super().__init__(spec=llm_prompt_spec, metadata=llm_metadata, **kwargs)
+
+    @property
+    def spec(self) -> LLMPromptArtifactSpec:
+        return self._spec
+
+    @spec.setter
+    def spec(self, spec: LLMPromptArtifactSpec):
+        self._spec = self._verify_dict(spec, "spec", LLMPromptArtifactSpec)
+
+    @property
+    def model_artifact(self) -> Optional[model_art.ModelArtifact]:
+        """
+        Get the model artifact linked to this prompt artifact.
+        """
+        if self.spec._model_artifact:
+            return self.spec._model_artifact
+        if self.spec.model_uri:
+            self.spec._model_artifact, _ = (
+                mlrun.datastore.store_manager.get_store_artifact(self.spec.model_uri)
+            )
+            return self.spec._model_artifact
+        return None
+
+    def read_prompt(self) -> Optional[Union[str, list[dict]]]:
+        """
+        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_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:
+                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_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=".json"
+            ) as temp_file:
+                temp_file.write(json.dumps(self.spec.prompt_template))
+            self.spec.src_path = temp_file.name
+            self.spec.prompt_template = None
+            self._src_is_temp = True
+        super().before_log()

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/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-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/common/schemas/__init__.py

@@ -214,7 +214,7 @@ from .secret import (
     SecretsData,
     UserSecretCreationRequest,
 )
-from .serving import ModelRunnerStepData, MonitoringData
+from .serving import ModelRunnerStepData, ModelsData, MonitoringData
 from .tag import Tag, TagObjects
 from .workflow import (
     GetWorkflowResponse,

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -336,8 +336,8 @@ class ModelEndpointMonitoringMetricNoData(_ModelEndpointMonitoringMetricValuesBa
 
 class ApplicationBaseRecord(BaseModel):
     type: Literal["metric", "result"]
-    time: datetime
     value: float
+    time: Optional[datetime] = None
 
 
 class ApplicationResultRecord(ApplicationBaseRecord):

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/common/schemas/serving.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
+
 from pydantic.v1 import BaseModel
 
 from mlrun.common.types import StrEnum
@@ -40,3 +42,8 @@ class MonitoringData(StrEnum):
     MODEL_PATH = "model_path"
     MODEL_ENDPOINT_UID = "model_endpoint_uid"
     MODEL_CLASS = "model_class"
+
+
+class ModelsData(enum.Enum):
+    MODEL_CLASS = 0
+    MODEL_PARAMETERS = 1

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/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-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/config.py

@@ -125,6 +125,8 @@ default_config = {
                 "interval": "30",
                 # runs limit to fetch for retrying
                 "fetch_runs_limit": 1000,
+                # minutes until a run is considered stale and will be aborted
+                "staleness_threshold": 60 * 24 * 3,
             },
         },
         "projects": {

mlrun-1.10.0rc14/mlrun/datastore/model_provider/model_provider.py

@@ -0,0 +1,160 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.errors
+from mlrun.datastore.remote_client import (
+    BaseRemoteClient,
+)
+
+T = TypeVar("T")
+
+
+class ModelProvider(BaseRemoteClient):
+    """
+    The ModelProvider class is an abstract base for integrating with external
+    model providers, primarily generative AI (GenAI) services.
+
+    Designed to be subclassed, it defines a consistent interface and shared
+    functionality for tasks such as text generation, embeddings, and invoking
+    fine-tuned models. Subclasses should implement provider-specific logic,
+    including SDK client initialization, model invocation, and custom operations.
+
+    Key Features:
+    - Establishes a consistent, reusable client management for model provider integrations.
+    - Simplifies GenAI service integration by abstracting common operations.
+    - Reduces duplication through shared components for common tasks.
+    - Holds default invocation parameters (e.g., temperature, max_tokens) to avoid boilerplate
+    code and promote consistency.
+    """
+
+    support_async = False
+
+    def __init__(
+        self,
+        parent,
+        kind,
+        name,
+        endpoint="",
+        secrets: Optional[dict] = None,
+        default_invoke_kwargs: Optional[dict] = None,
+    ):
+        super().__init__(
+            parent=parent, name=name, kind=kind, endpoint=endpoint, secrets=secrets
+        )
+        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 load_client(self) -> None:
+        """
+        Initializes the SDK client for the model provider with the given keyword arguments
+        and assigns it to an instance attribute (e.g., self._client).
+
+        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")
+
+    def invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Optional[Union[str, T]]:
+        """
+        Invokes a generative AI model with the provided messages and additional parameters.
+        This method is designed to be a flexible interface for interacting with various
+        generative AI backends (e.g., OpenAI, Hugging Face, etc.). It allows users to send
+        a list of messages (following a standardized format) and receive a response. The
+        response can be returned as plain text or in its full structured format, depending
+        on the `as_str` parameter.
+
+        :param messages:    A list of dictionaries representing the conversation history or input messages.
+                            Each dictionary should follow the format::
+                            {"role": "system"| "user" | "assistant" ..., "content": "Message content as a string"}
+                            Example:
+
+                            .. code-block:: json
+
+                                [
+                                    {"role": "system", "content": "You are a helpful assistant."},
+                                    {"role": "user", "content": "What is the capital of France?"}
+                                ]
+
+                            This format is consistent across all backends. Defaults to None if no messages
+                            are provided.
+
+        :param as_str:      A boolean flag indicating whether to return the response as a plain string.
+                            - If True, the function extracts and returns the main content of the first
+                            response.
+                            - If False, the function returns the full response object,
+                            which may include additional metadata or multiple response options.
+                            Defaults to False.
+
+        :param invoke_kwargs:
+                            Additional keyword arguments to be passed to the underlying model API call.
+                            These can include parameters such as temperature, max tokens, etc.,
+                            depending on the capabilities of the specific backend being used.
+
+        :return:
+                            - If `as_str` is True: Returns the main content of the first response as a string.
+                            - If `as_str` is False: Returns the full response object.
+
+        """
+        raise NotImplementedError("invoke method is not implemented")
+
+    def customized_invoke(
+        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
+    ) -> Optional[T]:
+        raise NotImplementedError("customized_invoke method is not implemented")
+
+    @property
+    def client(self):
+        return self._client
+
+    @property
+    def model(self):
+        return None
+
+    def get_invoke_kwargs(self, invoke_kwargs):
+        kwargs = self.default_invoke_kwargs.copy()
+        kwargs.update(invoke_kwargs)
+        return kwargs
+
+    @property
+    def async_client(self):
+        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_invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Awaitable[str]:
+        raise NotImplementedError("async_invoke is not implemented")

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/datastore/model_provider/openai_provider.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Optional, TypeVar
+from typing import Callable, Optional, TypeVar, Union
 
 import mlrun
 from mlrun.datastore.model_provider.model_provider import ModelProvider
@@ -21,6 +21,18 @@ T = TypeVar("T")
 
 
 class OpenAIProvider(ModelProvider):
+    """
+    OpenAIProvider is a wrapper around the OpenAI SDK that provides an interface
+    for interacting with OpenAI's generative AI services.
+
+    It supports both synchronous and asynchronous operations, allowing flexible
+    integration into various workflows.
+
+    This class extends the ModelProvider base class and implements OpenAI-specific
+    functionality, including client initialization, model invocation, and custom
+    operations tailored to the OpenAI API.
+    """
+
     def __init__(
         self,
         parent,
@@ -59,6 +71,19 @@ class OpenAIProvider(ModelProvider):
         return self.endpoint
 
     def load_client(self) -> None:
+        """
+        Initializes the OpenAI SDK client using the provided options.
+
+        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.
+
+        Raises:
+            ImportError: If the `openai` package is not installed.
+        """
         try:
             from openai import OpenAI  # noqa
 
@@ -87,34 +112,33 @@ class OpenAIProvider(ModelProvider):
         else:
             return self._default_operation(**invoke_kwargs, model=self.model)
 
-    def _get_messages_parameter(
-        self, prompt: Optional[str] = None, **invoke_kwargs
-    ) -> (str, dict):
-        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
-        messages = invoke_kwargs.get("messages")
-        if messages:
-            if prompt:
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    "can not provide 'messages' and 'prompt' to invoke"
-                )
-        elif prompt:
-            messages = [
-                {
-                    "role": "user",
-                    "content": prompt,
-                },
-            ]
-        else:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "must provide 'messages' or 'prompt' to invoke"
-            )
-        return messages, invoke_kwargs
+    def invoke(
+        self,
+        messages: Optional[list[dict]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Optional[Union[str, T]]:
+        """
+        OpenAI-specific implementation of `ModelProvider.invoke`.
+        Invokes an OpenAI model operation using the sync client.
+        For full details, see `ModelProvider.invoke`.
 
-    def invoke(self, prompt: Optional[str] = None, **invoke_kwargs) -> str:
-        messages, invoke_kwargs = self._get_messages_parameter(
-            prompt=prompt, **invoke_kwargs
-        )
+        :param messages:    Same as ModelProvider.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 response object, whose type depends on
+                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+
+        :param invoke_kwargs:
+                            Same as ModelProvider.invoke.
+
+        """
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
         response = self._default_operation(
             model=self.endpoint, messages=messages, **invoke_kwargs
         )
-        return response.choices[0].message.content
+        if as_str:
+            return response.choices[0].message.content
+        return response

{mlrun-1.10.0rc12 → mlrun-1.10.0rc14}/mlrun/datastore/remote_client.py

@@ -18,6 +18,17 @@ import mlrun
 
 
 class BaseRemoteClient:
+    """
+    The BaseRemoteClient class serves as a foundational component for managing
+    secrets and configurations.
+    It is designed to be extended by subclasses that interact with external services,
+     such as file systems (e.g., Datastore) or model providers (e.g., ModelProvider).
+
+    This class is intended to provide shared functionality and should not be
+    used directly. Instead, create a subclass to implement logic specific to
+    your use case, such as interactions with S3 storage or invoking model providers like OpenAI.
+    """
+
     def __init__(self, parent, kind, name, endpoint="", secrets: Optional[dict] = None):
         self._parent = parent
         self.kind = kind