mlrun 1.10.0rc18__py3-none-any.whl → 1.10.0rc20__py3-none-any.whl

mlrun/__init__.py

@@ -31,6 +31,7 @@ from typing import Optional
 
 import dotenv
 
+from .common.constants import MLRUN_ACTIVE_PROJECT
 from .config import config as mlconf
 from .datastore import DataItem, ModelProvider, store_manager
 from .db import get_run_db
@@ -167,11 +168,29 @@ def set_environment(
 
 
 def get_current_project(silent: bool = False) -> Optional[MlrunProject]:
-    if not pipeline_context.project and not silent:
+    if pipeline_context.project:
+        return pipeline_context.project
+
+    project_name = environ.get(MLRUN_ACTIVE_PROJECT, None)
+    if not project_name:
+        if not silent:
+            raise MLRunInvalidArgumentError(
+                "No current project is initialized. Use new, get or load project functions first."
+            )
+        return None
+
+    project = load_project(
+        name=project_name,
+        url=project_name,
+        save=False,
+        sync_functions=False,
+    )
+
+    if not project and not silent:
         raise MLRunInvalidArgumentError(
             "No current project is initialized. Use new, get or load project functions first."
         )
-    return pipeline_context.project
+    return project
 
 
 def get_sample_path(subpath=""):

mlrun/common/constants.py

@@ -30,6 +30,7 @@ RESERVED_TAG_NAME_LATEST = "latest"
 JOB_TYPE_WORKFLOW_RUNNER = "workflow-runner"
 JOB_TYPE_PROJECT_LOADER = "project-loader"
 JOB_TYPE_RERUN_WORKFLOW_RUNNER = "rerun-workflow-runner"
+MLRUN_ACTIVE_PROJECT = "MLRUN_ACTIVE_PROJECT"
 
 
 class MLRunInternalLabels:

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/__init__.py

@@ -39,6 +39,7 @@ __all__ = [
 from urllib.parse import urlparse
 
 import fsspec
+import storey
 
 import mlrun.datastore.wasbfs
 from mlrun.datastore.datastore_profile import (
@@ -168,11 +169,12 @@ def get_stream_pusher(stream_path: str, **kwargs):
             raise ValueError(f"unsupported stream path {stream_path}")
 
 
-class _DummyStream:
+class _DummyStream(storey.MapClass):
     """stream emulator for tests and debug"""
 
     def __init__(self, event_list=None, **kwargs):
         self.event_list = event_list or []
+        super().__init__(**kwargs)
 
     def push(self, data, **kwargs):
         if not isinstance(data, list):
@@ -180,3 +182,9 @@ class _DummyStream:
         for item in data:
             logger.info(f"dummy stream got event: {item}, kwargs={kwargs}")
             self.event_list.append(item)
+
+    def do(self, event):
+        if not isinstance(event, list):
+            event = [event]
+        for item in event:
+            self.event_list.append(item)

mlrun/datastore/model_provider/huggingface_provider.py

@@ -12,16 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, TypeVar, Union
+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
-
-T = TypeVar("T")
-ChatType = list[dict[str, str]]  # according to transformers.pipelines.text_generation
+    from transformers.pipelines.text_generation import ChatType
 
 
 class HuggingFaceProvider(ModelProvider):
@@ -63,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):
@@ -81,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.
@@ -91,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:
@@ -117,7 +184,7 @@ class HuggingFaceProvider(ModelProvider):
 
     def custom_invoke(
         self, operation: Optional["Pipeline"] = None, **invoke_kwargs
-    ) -> Optional[T]:
+    ) -> Union[list, dict, Any]:
         """
         HuggingFace implementation of `ModelProvider.custom_invoke`.
         Use the default config in provider client/ user defined client:
@@ -150,34 +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,
-    ) -> Optional[Union[str, list, T]]:
+    ) -> 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 invoke_response_format: InvokeResponseFormat
+            Specifies the format of the returned response. Options:
 
-        :param as_str:
-                            If `True`, returns only the main content from a single response
-                            (intended for single-response use cases).
-                            If `False`, returns the full response object, whose type depends on
-                            the client (e.g., `pipeline`).
+            - "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

@@ -12,14 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Awaitable
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Callable, Optional, Union
 
 import mlrun.errors
+from mlrun.common.types import StrEnum
 from mlrun.datastore.remote_client import (
     BaseRemoteClient,
 )
 
-T = TypeVar("T")
+
+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):
@@ -58,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
@@ -79,69 +138,6 @@ class ModelProvider(BaseRemoteClient):
 
         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 custom_invoke(
-        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
-    ) -> Optional[T]:
-        """
-        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) -> Any:
         return self._client
@@ -168,9 +164,22 @@ class ModelProvider(BaseRemoteClient):
             )
         return self._async_client
 
+    def custom_invoke(self, operation: Optional[Callable], **invoke_kwargs) -> Any:
+        """
+        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")
+
     async def async_custom_invoke(
-        self, operation: Optional[Callable[..., Awaitable[T]]], **invoke_kwargs
-    ) -> Optional[T]:
+        self, operation: Optional[Callable[..., Awaitable[Any]]], **invoke_kwargs
+    ) -> Any:
         """
         Asynchronously invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.)
         with the given keyword arguments.
@@ -183,11 +192,76 @@ class ModelProvider(BaseRemoteClient):
         """
         raise NotImplementedError("async_custom_invoke is not implemented")
 
+    def invoke(
+        self,
+        messages: Union[list[dict], Any],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **invoke_kwargs,
+    ) -> 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.
+
+        :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.
+
+                                                Note: The usage dictionary may contain additional
+                                                keys depending on the model provider:
+
+                                    .. code-block:: json
+
+                                    {
+                                        "answer": "<generated_text>",
+                                        "usage": {
+                                        "prompt_tokens": <int>,
+                                        "completion_tokens": <int>,
+                                        "total_tokens": <int>
+                                        }
+
+                                    }
+
+                                    - 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.
+
+        :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,
-    ) -> Optional[str]:
+    ) -> Union[str, dict[str, Any], Any]:
         """Async version of `invoke`. See `invoke` for full documentation."""
         raise NotImplementedError("async_invoke is not implemented")