mlrun 1.10.0rc12__py3-none-any.whl → 1.10.0rc14__py3-none-any.whl

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,31 +45,107 @@ 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
-        self._model_artifact = None
+        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):
     """
@@ -86,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,
@@ -95,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,
@@ -133,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/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

@@ -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/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/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/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

@@ -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/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
+from typing import Callable, Optional, TypeVar, Union
 
 import mlrun.errors
 from mlrun.datastore.remote_client import (
@@ -23,6 +23,23 @@ 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__(
@@ -44,9 +61,65 @@ class ModelProvider(BaseRemoteClient):
         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, prompt: Optional[str] = None, **invoke_kwargs) -> str:
+    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(
@@ -78,5 +151,10 @@ class ModelProvider(BaseRemoteClient):
     async def async_customized_invoke(self, **kwargs):
         raise NotImplementedError("async_customized_invoke is not implemented")
 
-    async def async_invoke(self, prompt: str, **invoke_kwargs) -> Awaitable[str]:
+    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/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/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

mlrun/execution.py

@@ -94,6 +94,7 @@ class MLClientCtx:
         self._state_thresholds = {}
         self._retry_spec = {}
         self._retry_count = None
+        self._retries = []
 
         self._labels = {}
         self._annotations = {}
@@ -468,6 +469,7 @@ class MLClientCtx:
             for key, uri in status.get("artifact_uris", {}).items():
                 self._artifacts_manager.artifact_uris[key] = uri
             self._retry_count = status.get("retry_count", self._retry_count)
+            self._retries = status.get("retries", self._retries)
             # if run is a retry, the state needs to move to running
             if include_status:
                 self._state = status.get("state", self._state)
@@ -911,7 +913,7 @@ class MLClientCtx:
     def log_llm_prompt(
         self,
         key,
-        prompt_string: Optional[str] = None,
+        prompt_template: Optional[list[dict]] = None,
         prompt_path: Optional[str] = None,
         prompt_legend: Optional[dict] = None,
         model_artifact: Union[ModelArtifact, str] = None,
@@ -935,7 +937,7 @@ class MLClientCtx:
             # Log an inline prompt
             context.log_llm_prompt(
                 key="qa-prompt",
-                prompt_string="Q: {question}",
+                prompt_template=[{"role: "user", "content": "question with {place_holder}"}],
                 model_artifact=model,
                 prompt_legend={"question": "user_input"},
                 model_configuration={"temperature": 0.7, "max_tokens": 128},
@@ -943,10 +945,16 @@ class MLClientCtx:
             )
 
         :param key: Unique name of the artifact.
-        :param prompt_string: Raw prompt text as a string. Cannot be used with `prompt_path`.
+        :param prompt_template: Raw prompt list of dicts -
+         [{"role": "system", "content": "You are a {profession} advisor"},
+         "role": "user", "content": "I need your help with {profession}"]. only "role" and "content" keys allow in any
+         str format (upper/lower case), keys will be modified to lower case.
+         Cannot be used with `prompt_path`.
         :param prompt_path: Path to a file containing the prompt content. Cannot be used with `prompt_string`.
         :param prompt_legend: A dictionary where each key is a placeholder in the prompt (e.g., ``{user_name}``)
-               and the value is a description or explanation of what that placeholder represents.
+               and the value is a dictionary holding two keys, "field", "description". "field" points to the field in
+               the event where the value of the place-holder inside the event, if None or not exist will be replaced
+               with the place-holder name. "description" will point to explanation of what that placeholder represents.
                Useful for documenting and clarifying dynamic parts of the prompt.
         :param model_artifact: Reference to the parent model (either `ModelArtifact` or model URI string).
         :param model_configuration: Dictionary of generation parameters (e.g., temperature, max_tokens).
@@ -961,10 +969,15 @@ class MLClientCtx:
         :returns: The logged `LLMPromptArtifact` object.
         """
 
+        if not prompt_template and not prompt_path:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Either 'prompt_template' or 'prompt_path' must be provided"
+            )
+
         llm_prompt = LLMPromptArtifact(
             key=key,
             project=self.project or "",
-            prompt_string=prompt_string,
+            prompt_template=prompt_template,
             prompt_path=prompt_path,
             prompt_legend=prompt_legend,
             model_artifact=model_artifact,
@@ -1262,6 +1275,7 @@ class MLClientCtx:
                 "start_time": to_date_str(self._start_time),
                 "last_update": to_date_str(self._last_update),
                 "retry_count": self._retry_count,
+                "retries": self._retries,
             },
         }
 

mlrun/model.py

@@ -1375,6 +1375,7 @@ class RunStatus(ModelObj):
         notifications: Optional[dict[str, Notification]] = None,
         artifact_uris: Optional[dict[str, str]] = None,
         retry_count: Optional[int] = None,
+        retries: Optional[list[dict]] = None,
     ):
         self.state = state or "created"
         self.status_text = status_text
@@ -1393,6 +1394,7 @@ class RunStatus(ModelObj):
         # Artifact key -> URI mapping, since the full artifacts are not stored in the runs DB table
         self._artifact_uris = artifact_uris or {}
         self._retry_count = retry_count or None
+        self._retries = retries or []
 
     @classmethod
     def from_dict(
@@ -1461,6 +1463,19 @@ class RunStatus(ModelObj):
         """
         self._retry_count = retry_count
 
+    @property
+    def retries(self) -> list[dict]:
+        """List of metadata for each retry attempt."""
+        return self._retries
+
+    @retries.setter
+    def retries(self, retries: list[dict]):
+        """
+        Set the list of retry attempt metadata.
+        :param retries: A list of dictionaries, each representing a retry attempt.
+        """
+        self._retries = retries
+
     def is_failed(self) -> Optional[bool]:
         """
         This method returns whether a run has failed.