mlrun 1.10.0rc19__py3-none-any.whl → 1.10.0rc21__py3-none-any.whl

mlrun/common/schemas/function.py

@@ -114,11 +114,21 @@ class StateThresholds(pydantic.v1.BaseModel):
     default: typing.Optional[dict[str, str]]
 
 
+class Backoff(pydantic.v1.BaseModel):
+    default_base_delay: typing.Optional[str]
+    min_base_delay: typing.Optional[str]
+
+
+class RetrySpec(pydantic.v1.BaseModel):
+    backoff: Backoff
+
+
 class FunctionSpec(pydantic.v1.BaseModel):
     image_pull_secret: typing.Optional[ImagePullSecret]
     security_context: typing.Optional[SecurityContext]
     service_account: typing.Optional[ServiceAccount]
     state_thresholds: typing.Optional[StateThresholds]
+    retry: typing.Optional[RetrySpec]
 
     class Config:
         extra = pydantic.v1.Extra.allow

mlrun/common/schemas/model_monitoring/constants.py

@@ -34,6 +34,7 @@ class ModelEndpointSchema(MonitoringStrEnum):
     UID = "uid"
     PROJECT = "project"
     ENDPOINT_TYPE = "endpoint_type"
+    MODE = "mode"
     NAME = "name"
     CREATED = "created"
     UPDATED = "updated"
@@ -326,18 +327,10 @@ class EndpointType(IntEnum):
     def top_level_list(cls):
         return [cls.NODE_EP, cls.ROUTER, cls.BATCH_EP]
 
-    @classmethod
-    def real_time_list(cls):
-        return [cls.NODE_EP, cls.ROUTER, cls.LEAF_EP]
-
-    @classmethod
-    def batch_list(cls):
-        return [cls.BATCH_EP]
 
-
-class EndpointMode(StrEnum):
-    REAL_TIME = "real_time"
-    BATCH = "batch"
+class EndpointMode(IntEnum):
+    REAL_TIME = 0
+    BATCH = 1
 
 
 class MonitoringFunctionNames(MonitoringStrEnum):

mlrun/common/schemas/model_monitoring/model_endpoints.py

@@ -28,6 +28,7 @@ from .constants import (
     FQN_REGEX,
     MODEL_ENDPOINT_ID_PATTERN,
     PROJECT_PATTERN,
+    EndpointMode,
     EndpointType,
     ModelEndpointMonitoringMetricType,
     ModelMonitoringMode,
@@ -118,6 +119,7 @@ class ModelEndpointMetadata(ObjectMetadata, ModelEndpointParser):
     project: constr(regex=PROJECT_PATTERN)
     endpoint_type: EndpointType = EndpointType.NODE_EP
     uid: Optional[constr(regex=MODEL_ENDPOINT_ID_PATTERN)]
+    mode: EndpointMode = EndpointMode.REAL_TIME
 
     @classmethod
     def mutable_fields(cls):

mlrun/datastore/model_provider/huggingface_provider.py

@@ -15,7 +15,11 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import mlrun
-from mlrun.datastore.model_provider.model_provider import ModelProvider
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
 
 if TYPE_CHECKING:
     from transformers.pipelines.base import Pipeline
@@ -61,15 +65,18 @@ class HuggingFaceProvider(ModelProvider):
         self.load_client()
 
     @staticmethod
-    def _extract_string_output(result) -> str:
+    def _extract_string_output(response: list[dict]) -> str:
         """
         Extracts the first generated string from Hugging Face pipeline output,
         regardless of whether it's plain text-generation or chat-style output.
         """
-        if not isinstance(result, list) or len(result) == 0:
+        if not isinstance(response, list) or len(response) == 0:
             raise ValueError("Empty or invalid pipeline output")
-
-        return result[0].get("generated_text")
+        if len(response) != 1:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "HuggingFaceProvider: extracting string from response is only supported for single-response outputs"
+            )
+        return response[0].get("generated_text")
 
     @classmethod
     def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
@@ -79,6 +86,68 @@ class HuggingFaceProvider(ModelProvider):
             subpath = ""
         return endpoint, subpath
 
+    def _response_handler(
+        self,
+        response: Union[str, list],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        messages: Union[str, list[str], "ChatType", list["ChatType"]] = None,
+        **kwargs,
+    ) -> Union[str, list, dict[str, Any]]:
+        """
+        Same as `ModelProvider._response_handler`.
+
+        * Expected to receive the response with `return_full_text=False`.
+
+        :param messages:                Same as in `ModelProvider._response_handler`.
+        :param response:                Same as in `ModelProvider._response_handler`.
+        :param invoke_response_format:  Same as in `ModelProvider._response_handler`, in full and string modes.
+
+                                        For usage mode, generate 3 statistics:
+                                        prompt_tokens, completion_tokens and total_tokens.
+
+                                        NOTE: Token counts are estimated after answer generation and
+                                        may differ from the actual tokens generated by the model due to
+                                        internal decoding behavior and implementation details.
+
+        :param kwargs:                  Same as in `ModelProvider._response_handler`.
+
+        :return: The result formatted according to the `invoke_response_format`.
+
+        :raises MLRunInvalidArgumentError: If extracting the string response fails.
+        :raises MLRunRuntimeError: If applying the chat template to the model fails.
+        """
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
+            str_response = self._extract_string_output(response)
+            if invoke_response_format == InvokeResponseFormat.STRING:
+                return str_response
+            if invoke_response_format == InvokeResponseFormat.USAGE:
+                tokenizer = self.client.tokenizer
+                if not isinstance(messages, str):
+                    try:
+                        messages = tokenizer.apply_chat_template(
+                            messages, tokenize=False, add_generation_prompt=True
+                        )
+                    except Exception as e:
+                        raise mlrun.errors.MLRunRuntimeError(
+                            f"Failed to apply chat template using the tokenizer for model '{self.model}'. "
+                            "This may indicate that the tokenizer does not support chat formatting, "
+                            "or that the input format is invalid. "
+                            f"Original error: {e}"
+                        )
+                prompt_tokens = len(tokenizer.encode(messages))
+                completion_tokens = len(tokenizer.encode(str_response))
+                total_tokens = prompt_tokens + completion_tokens
+                usage = {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                }
+                response = {
+                    UsageResponseKeys.ANSWER: str_response,
+                    UsageResponseKeys.USAGE: usage,
+                }
+        return response
+
     def load_client(self) -> None:
         """
         Initializes the Hugging Face pipeline using the provided options.
@@ -89,7 +158,7 @@ class HuggingFaceProvider(ModelProvider):
 
         Note: Hugging Face pipelines are synchronous and do not support async invocation.
 
-        Raises:
+        :raises:
             ImportError: If the `transformers` package is not installed.
         """
         try:
@@ -148,35 +217,55 @@ class HuggingFaceProvider(ModelProvider):
 
     def invoke(
         self,
-        messages: Union[str, list[str], "ChatType", list["ChatType"]] = None,
-        as_str: bool = False,
+        messages: Union[str, list[str], "ChatType", list["ChatType"]],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Union[str, list]:
+    ) -> Union[str, list, dict[str, Any]]:
         """
         HuggingFace-specific implementation of `ModelProvider.invoke`.
         Invokes a HuggingFace model operation using the synchronous client.
-        For complete usage details, refer to `ModelProvider.invoke`.
+        For full details, see `ModelProvider.invoke`.
 
         :param messages:
-                            Same as ModelProvider.invoke.
+            Same as `ModelProvider.invoke`.
 
-        :param as_str:
-                            If `True`, return only the main content (e.g., generated text) from a
-                            **single-response output** — intended for use cases where you expect exactly one result.
+        :param invoke_response_format: InvokeResponseFormat
+            Specifies the format of the returned response. Options:
 
-                            If `False`, return the **full raw response object**, which is a list of dictionaries.
+            - "string": Returns only the generated text content, extracted from a single response.
+            - "usage":  Combines the generated text with metadata (e.g., token usage), returning a dictionary:
+
+            .. code-block:: json
+                {
+                    "answer": "<generated_text>",
+                    "usage": {
+                        "prompt_tokens": <int>,
+                        "completion_tokens": <int>,
+                        "total_tokens": <int>
+                    }
+                }
+
+            - "full":   Returns the raw response object from the HuggingFace model,
+                        typically a list of generated sequences (dictionaries).
+                        This format does not include token usage statistics.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
-        :return:            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the HuggingFace client. Same as in `ModelProvider.invoke`.
+
+        :return:
+            A string, dictionary, or list of model outputs, depending on `invoke_response_format`.
         """
+
         if self.client.task != "text-generation":
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "HuggingFaceProvider.invoke supports text-generation task only"
             )
-        if as_str:
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
             invoke_kwargs["return_full_text"] = False
         response = self.custom_invoke(text_inputs=messages, **invoke_kwargs)
-        if as_str:
-            return self._extract_string_output(response)
+        response = self._response_handler(
+            messages=messages,
+            response=response,
+            invoke_response_format=invoke_response_format,
+        )
         return response

mlrun/datastore/model_provider/model_provider.py

@@ -15,11 +15,37 @@ from collections.abc import Awaitable
 from typing import Any, Callable, Optional, Union
 
 import mlrun.errors
+from mlrun.common.types import StrEnum
 from mlrun.datastore.remote_client import (
     BaseRemoteClient,
 )
 
 
+class InvokeResponseFormat(StrEnum):
+    STRING = "string"
+    USAGE = "usage"
+    FULL = "full"
+
+    @classmethod
+    def is_str_response(cls, invoke_response_format: str) -> bool:
+        """
+        Returns True if the response key corresponds to a string-based response (not a full generation object).
+        """
+        return invoke_response_format in {
+            cls.USAGE,
+            cls.STRING,
+        }
+
+
+class UsageResponseKeys(StrEnum):
+    ANSWER = "answer"
+    USAGE = "usage"
+
+    @classmethod
+    def fields(cls) -> list[str]:
+        return [cls.ANSWER, cls.USAGE]
+
+
 class ModelProvider(BaseRemoteClient):
     """
     The ModelProvider class is an abstract base for integrating with external
@@ -56,6 +82,41 @@ class ModelProvider(BaseRemoteClient):
         self._client = None
         self._async_client = None
 
+    @staticmethod
+    def _extract_string_output(response: Any) -> str:
+        """
+        Extracts string response from response object
+        """
+        pass
+
+    def _response_handler(
+        self,
+        response: Any,
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **kwargs,
+    ) -> Union[str, dict, Any]:
+        """
+        Handles the model response according to the specified response format.
+
+        :param response: The raw response returned from the model invocation.
+        :param invoke_response_format: Determines how the response should be processed and returned.
+                                       Options include:
+
+                                       - STRING: Return only the main generated content as a string,
+                                                 typically for single-answer responses.
+                                       - USAGE: Return a dictionary combining the string response with
+                                                additional metadata or token usage statistics, in this format:
+                                                {"answer": <string>, "usage": <dict>}
+
+                                       - FULL: Return the full raw response object unmodified.
+
+        :param kwargs:                  Additional parameters that may be required by specific implementations.
+
+        :return:                        The processed response in the format specified by `invoke_response_format`.
+                                        Can be a string, dictionary, or the original response object.
+        """
+        return None
+
     def get_client_options(self) -> dict:
         """
         Returns a dictionary containing credentials and configuration
@@ -133,57 +194,74 @@ class ModelProvider(BaseRemoteClient):
 
     def invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: Union[list[dict], Any],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Union[str, Any]:
+    ) -> Union[str, dict[str, Any], Any]:
         """
         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.
+        a list of messages (following a standardized format) and receive a response.
+
+        :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 invoke_response_format:   Determines how the model response is returned:
+
+                                    - string:   Returns only the generated text content from the model output,
+                                                for single-answer responses only.
+
+                                    - usage:    Combines the STRING response with additional metadata (token usage),
+                                                and returns the result in a dictionary.
 
-        :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:
+                                                Note: The usage dictionary may contain additional
+                                                keys depending on the model provider:
 
-                            .. code-block:: json
+                                    .. code-block:: json
 
-                                [
-                                    {"role": "system", "content": "You are a helpful assistant."},
-                                    {"role": "user", "content": "What is the capital of France?"}
-                                ]
+                                    {
+                                        "answer": "<generated_text>",
+                                        "usage": {
+                                        "prompt_tokens": <int>,
+                                        "completion_tokens": <int>,
+                                        "total_tokens": <int>
+                                        }
 
-                            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.
+                                    - full:   Returns the full model output.
 
         :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.
+                                    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.
+        :return:                    The invoke result formatted according to the specified
+                                    invoke_response_format parameter.
 
         """
         raise NotImplementedError("invoke method is not implemented")
 
     async def async_invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Union[str, Any]:
+    ) -> Union[str, dict[str, Any], Any]:
         """Async version of `invoke`. See `invoke` for full documentation."""
         raise NotImplementedError("async_invoke is not implemented")

mlrun/datastore/model_provider/openai_provider.py

@@ -16,7 +16,11 @@ from collections.abc import Awaitable
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union
 
 import mlrun
-from mlrun.datastore.model_provider.model_provider import ModelProvider
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
 from mlrun.datastore.utils import accepts_param
 
 if TYPE_CHECKING:
@@ -38,6 +42,7 @@ class OpenAIProvider(ModelProvider):
     """
 
     support_async = True
+    response_class = None
 
     def __init__(
         self,
@@ -64,6 +69,27 @@ class OpenAIProvider(ModelProvider):
         self.options = self.get_client_options()
         self.load_client()
 
+    @classmethod
+    def _import_response_class(cls) -> None:
+        if not cls.response_class:
+            try:
+                from openai.types.chat.chat_completion import ChatCompletion
+            except ImportError as exc:
+                raise ImportError("openai package is not installed") from exc
+            cls.response_class = ChatCompletion
+
+    @staticmethod
+    def _extract_string_output(response: "ChatCompletion") -> str:
+        """
+        Extracts the first generated string from Hugging Face pipeline output,
+        regardless of whether it's plain text-generation or chat-style output.
+        """
+        if len(response.choices) != 1:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "OpenAIProvider: extracting string from response is only supported for single-response outputs"
+            )
+        return response.choices[0].message.content
+
     @classmethod
     def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
         if endpoint and subpath:
@@ -180,60 +206,90 @@ class OpenAIProvider(ModelProvider):
                 **invoke_kwargs, **model_kwargs
             )
 
+    def _response_handler(
+        self,
+        response: "ChatCompletion",
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **kwargs,
+    ) -> ["ChatCompletion", str, dict[str, Any]]:
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
+            str_response = self._extract_string_output(response)
+            if invoke_response_format == InvokeResponseFormat.STRING:
+                return str_response
+            if invoke_response_format == InvokeResponseFormat.USAGE:
+                stats = response.to_dict()["usage"]
+                response = {
+                    UsageResponseKeys.ANSWER: str_response,
+                    UsageResponseKeys.USAGE: stats,
+                }
+        return response
+
     def invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Union[str, "ChatCompletion"]:
+    ) -> Union[dict[str, Any], str, "ChatCompletion"]:
         """
         OpenAI-specific implementation of `ModelProvider.invoke`.
-        Invokes an OpenAI model operation using the sync client.
+        Invokes an OpenAI model operation using the synchronous client.
         For full details, see `ModelProvider.invoke`.
 
-        :param messages:    Same as ModelProvider.invoke.
+        :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_response_format: InvokeResponseFormat
+            Specifies the format of the returned response. Options:
+
+            - "string": Returns only the generated text content, taken from a single response.
+            - "stats": Combines the generated text with metadata (e.g., token usage), returning a dictionary:
+
+              .. code-block:: json
+                 {
+                     "answer": "<generated_text>",
+                     "stats": <ChatCompletion>.to_dict()["usage"]
+                 }
+
+            - "full": Returns the full OpenAI `ChatCompletion` object.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
-        :return:            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the OpenAI client. Same as in `ModelProvider.invoke`.
 
+        :return:
+            A string, dictionary, or `ChatCompletion` object, depending on `invoke_response_format`.
         """
+
         response = self.custom_invoke(messages=messages, **invoke_kwargs)
-        if as_str:
-            return response.choices[0].message.content
-        return response
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )
 
     async def async_invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Union[str, "ChatCompletion"]:
+    ) -> Union[str, "ChatCompletion", dict]:
         """
         OpenAI-specific implementation of `ModelProvider.async_invoke`.
         Invokes an OpenAI model operation using the async client.
-        For full details, see `ModelProvider.async_invoke`.
+        For full details, see `ModelProvider.async_invoke` and `OpenAIProvider.invoke`.
 
-        :param messages:    Same as ModelProvider.async_invoke.
+        :param messages:    Same as `OpenAIProvider.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_response_format: InvokeResponseFormat
+                            Same as `OpenAIProvider.invoke`.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.async_invoke.
-        :returns            Same as ModelProvider.async_invoke.
+                            Same as `OpenAIProvider.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
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )

mlrun/db/base.py

@@ -16,8 +16,6 @@ import datetime
 from abc import ABC, abstractmethod
 from typing import Literal, Optional, Union
 
-from deprecated import deprecated
-
 import mlrun.alerts
 import mlrun.common
 import mlrun.common.formatters
@@ -445,23 +443,6 @@ class RunDBInterface(ABC):
     ) -> dict:
         pass
 
-    # TODO: remove in 1.10.0
-    @deprecated(
-        version="1.7.0",
-        reason="'list_features' will be removed in 1.10.0, use 'list_features_v2' instead",
-        category=FutureWarning,
-    )
-    @abstractmethod
-    def list_features(
-        self,
-        project: str,
-        name: Optional[str] = None,
-        tag: Optional[str] = None,
-        entities: Optional[list[str]] = None,
-        labels: Optional[Union[str, dict[str, Optional[str]], list[str]]] = None,
-    ) -> mlrun.common.schemas.FeaturesOutput:
-        pass
-
     @abstractmethod
     def list_features_v2(
         self,