mlrun 1.10.0rc19__tar.gz → 1.10.0rc20__tar.gz

{mlrun-1.10.0rc19 → mlrun-1.10.0rc20}/Makefile

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # THIS BLOCK IS FOR VARIABLES USER MAY OVERRIDE
+DOCKER_DEFAULT_PLATFORM ?= linux/amd64
 MLRUN_VERSION ?= unstable
 # pip requires the python version to be according to some regex (so "unstable" is not valid for example) this regex only
 # allows us to have free text (like unstable) after the "+". on the contrary in a docker tag "+" is not a valid
@@ -235,11 +236,20 @@ update-version-file: ## Update the version file
 	python ./automation/version/version_file.py ensure --mlrun-version $(MLRUN_VERSION)
 
 .PHONY: generate-dockerignore
-generate-dockerignore: ## Copies the root .dockerignore and removes the tests pattern from it
+generate-dockerignore: ## Copies the root .dockerignore and removes test exclusions for test-system
 	$(eval TARGET := dockerfiles/${DEST}/Dockerfile.dockerignore)
-	@if [ -f "$(TARGET)" ]; then \
-		temp_file=$$(mktemp) && \
-		sed '/\*\*\/tests/d' .dockerignore > $$temp_file && \
+	@if [ -z "${DEST}" ]; then \
+		echo "Error: DEST variable must be set"; \
+		exit 1; \
+	fi; \
+	echo "Generating $(TARGET)..."; \
+	temp_file=$$(mktemp); \
+	if [ "$(DEST)" = "test-system" ]; then \
+		grep -vE '(\*\*/tests|\*\*/env\.yml|\*\*/test-[^/]*\.yml|\*\*/model_monitoring/assets)' .dockerignore > $$temp_file; \
+	else \
+		sed '/\*\*\/tests/d' .dockerignore > $$temp_file; \
+	fi; \
+	if [ -f "$(TARGET)" ]; then \
 		if cmp -s $$temp_file "$(TARGET)"; then \
 			echo "File $(TARGET) already exists and content is identical"; \
 			rm $$temp_file; \
@@ -249,7 +259,7 @@ generate-dockerignore: ## Copies the root .dockerignore and removes the tests pa
 			mv $$temp_file "$(TARGET)"; \
 		fi; \
 	else \
-		sed '/\*\*\/tests/d' .dockerignore > "$(TARGET)"; \
+		mv $$temp_file "$(TARGET)"; \
 	fi
 
 
@@ -292,7 +302,7 @@ MLRUN_CACHE_IMAGE_PUSH_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG),$(ML
 DEFAULT_IMAGES += $(MLRUN_IMAGE_NAME_TAGGED)
 
 .PHONY: mlrun
-mlrun: update-version-file ## Build mlrun docker image
+mlrun: common-image update-version-file ## Build mlrun docker image
 	$(MLRUN_CACHE_IMAGE_PULL_COMMAND)
 	docker build \
 		--file dockerfiles/mlrun/Dockerfile \
@@ -300,6 +310,7 @@ mlrun: update-version-file ## Build mlrun docker image
 		--build-arg MLRUN_PYTHON_VERSION=$(MLRUN_PYTHON_VERSION) \
 		--build-arg MLRUN_PIP_VERSION=$(MLRUN_PIP_VERSION) \
 		--build-arg MLRUN_UV_IMAGE=$(MLRUN_UV_IMAGE) \
+		--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
 		$(MLRUN_IMAGE_DOCKER_CACHE_FROM_FLAG) \
 		$(MLRUN_DOCKER_NO_CACHE_FLAG) \
 		--tag $(MLRUN_IMAGE_NAME_TAGGED) .
@@ -325,13 +336,14 @@ MLRUN_KFP_CACHE_IMAGE_PUSH_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG),
 DEFAULT_IMAGES += $(MLRUN_KFP_IMAGE_NAME_TAGGED)
 
 .PHONY: mlrun-kfp
-mlrun-kfp: update-version-file ## Build mlrun docker image with KFP
+mlrun-kfp: common-image-3.9 update-version-file ## Build mlrun docker image with KFP
 	$(MLRUN_KFP_CACHE_IMAGE_PULL_COMMAND)
 	docker build \
 		--file dockerfiles/mlrun-kfp/Dockerfile \
 		--build-arg MLRUN_DOCKER_REGISTRY=$(MLRUN_DOCKER_REGISTRY) \
 		--build-arg MLRUN_VERSION=$(MLRUN_VERSION) \
 		--build-arg MLRUN_PIP_VERSION=$(MLRUN_PIP_VERSION) \
+		--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
 		$(MLRUN_KFP_IMAGE_DOCKER_CACHE_FROM_FLAG) \
 		$(MLRUN_DOCKER_NO_CACHE_FLAG) \
 		--tag $(MLRUN_KFP_IMAGE_NAME):$(MLRUN_DOCKER_TAG)$(MLRUN_PYTHON_VERSION_SUFFIX) .
@@ -470,6 +482,49 @@ else
 	$(MAKE) -C server/go compile-schemas
 endif
 
+###############################################################################
+# Common base image
+###############################################################################
+
+COMMON_PLATFORM_TAG := $(subst /,_,$(DOCKER_DEFAULT_PLATFORM))
+COMMON_STAMP         ?= build/common-image.$(MLRUN_PYTHON_VERSION).$(COMMON_PLATFORM_TAG).stamp
+COMMON_DOCKERFILE     := dockerfiles/common/Dockerfile
+COMMON_IMAGE_NAME := mlrun_common_image:$(MLRUN_PYTHON_VERSION)
+
+common-image-3.11:
+	$(MAKE) common-image MLRUN_PYTHON_VERSION=3.11
+
+common-image-3.9:
+	$(MAKE) common-image MLRUN_PYTHON_VERSION=3.9
+
+# --- Build (cached) ----------------------------------------------------------
+ifeq ($(strip $(MLRUN_NO_CACHE)),)
+common-image: $(COMMON_STAMP)
+
+$(COMMON_STAMP): $(COMMON_DOCKERFILE)
+	docker build \
+	--build-arg MLRUN_PYTHON_VERSION=$(MLRUN_PYTHON_VERSION) \
+	--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
+	-f $(COMMON_DOCKERFILE) \
+	-t $(COMMON_IMAGE_NAME) . \
+	 && mkdir -p $(dir $@) && touch $@
+else  # when MLRUN_NO_CACHE is set
+.PHONY: common-image
+common-image:
+	docker build \
+	  --no-cache $(COMMON_DOCKER_ARGS) \
+ 	  --build-arg MLRUN_PYTHON_VERSION=$(MLRUN_PYTHON_VERSION) \
+ 	  --build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
+ 	  -f $(COMMON_DOCKERFILE) \
+	  -t $(COMMON_IMAGE_NAME) .
+endif
+
+.PHONY: clean-common-image
+clean-common-image:
+	docker rmi $(COMMON_IMAGE_TAG) || true
+	rm -f $(COMMON_STAMP)
+
+
 MLRUN_API_IMAGE_NAME := $(MLRUN_DOCKER_IMAGE_PREFIX)/mlrun-api
 MLRUN_API_CACHE_IMAGE_NAME := $(MLRUN_CACHE_DOCKER_IMAGE_PREFIX)/mlrun-api
 MLRUN_API_IMAGE_NAME_TAGGED := $(MLRUN_API_IMAGE_NAME):$(MLRUN_DOCKER_TAG)$(MLRUN_PYTHON_VERSION_SUFFIX)
@@ -479,13 +534,18 @@ MLRUN_API_CACHE_IMAGE_PULL_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG),
 MLRUN_API_CACHE_IMAGE_PUSH_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG),$(MLRUN_PUSH_DOCKER_CACHE_IMAGE)),docker tag $(MLRUN_API_IMAGE_NAME_TAGGED) $(MLRUN_API_CACHE_IMAGE_NAME_TAGGED) && docker push $(MLRUN_API_CACHE_IMAGE_NAME_TAGGED),)
 DEFAULT_IMAGES += $(MLRUN_API_IMAGE_NAME_TAGGED)
 
+
+# The API (and the common image it inherits from) must *always* be built on
+# Python 3.11, regardless of what the rest of the matrix is doing.
+api: export MLRUN_PYTHON_VERSION = 3.11
 .PHONY: api
-api: compile-schemas update-version-file ## Build mlrun-api docker image
+api: common-image-3.11	 compile-schemas update-version-file ## Build mlrun-api docker image
 	$(MLRUN_API_CACHE_IMAGE_PULL_COMMAND)
 	docker build \
 		--file dockerfiles/mlrun-api/Dockerfile \
 		--build-arg MLRUN_PYTHON_VERSION=$(MLRUN_PYTHON_VERSION) \
 		--build-arg MLRUN_UV_IMAGE=$(MLRUN_UV_IMAGE) \
+		--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
 		$(MLRUN_API_IMAGE_DOCKER_CACHE_FROM_FLAG) \
 		$(MLRUN_DOCKER_NO_CACHE_FLAG) \
 		--tag $(MLRUN_API_IMAGE_NAME_TAGGED) .
@@ -508,7 +568,7 @@ MLRUN_TEST_CACHE_IMAGE_PULL_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG)
 MLRUN_TEST_CACHE_IMAGE_PUSH_COMMAND := $(if $(and $(MLRUN_DOCKER_CACHE_FROM_TAG),$(MLRUN_PUSH_DOCKER_CACHE_IMAGE)),docker tag $(MLRUN_TEST_IMAGE_NAME_TAGGED) $(MLRUN_TEST_CACHE_IMAGE_NAME_TAGGED) && docker push $(MLRUN_TEST_CACHE_IMAGE_NAME_TAGGED),)
 
 .PHONY: build-test
-build-test: compile-schemas update-version-file ## Build test docker image
+build-test: common-image compile-schemas update-version-file ## Build test docker image
 	$(MAKE) generate-dockerignore DEST=test
 	$(MLRUN_TEST_CACHE_IMAGE_PULL_COMMAND)
 	docker build \
@@ -517,6 +577,7 @@ build-test: compile-schemas update-version-file ## Build test docker image
 		--build-arg MLRUN_PIP_VERSION=$(MLRUN_PIP_VERSION) \
 		--build-arg MLRUN_PIPELINES_KFP_VERSION=$(MLRUN_PIPELINES_KFP_VERSION) \
 		--build-arg MLRUN_UV_VERSION=$(MLRUN_UV_VERSION) \
+		--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
 		$(MLRUN_TEST_IMAGE_DOCKER_CACHE_FROM_FLAG) \
 		$(MLRUN_DOCKER_NO_CACHE_FLAG) \
 		--tag $(MLRUN_TEST_IMAGE_NAME_TAGGED) .
@@ -529,13 +590,13 @@ push-test: build-test ## Push test docker image
 MLRUN_SYSTEM_TEST_IMAGE_NAME := $(MLRUN_DOCKER_IMAGE_PREFIX)/test-system:$(MLRUN_DOCKER_TAG)
 
 .PHONY: build-test-system
-build-test-system: compile-schemas update-version-file ## Build system tests docker image
+build-test-system: common-image compile-schemas update-version-file ## Build system tests docker image
 	$(MAKE) generate-dockerignore DEST=test-system
 	docker build \
 		--file dockerfiles/test-system/Dockerfile \
-		--build-arg MLRUN_PYTHON_VERSION=$(MLRUN_PYTHON_VERSION) \
 		--build-arg MLRUN_PIP_VERSION=$(MLRUN_PIP_VERSION) \
 		--build-arg MLRUN_UV_VERSION=$(MLRUN_UV_VERSION) \
+		--build-arg DOCKER_DEFAULT_PLATFORM=$(DOCKER_DEFAULT_PLATFORM) \
 		$(MLRUN_DOCKER_NO_CACHE_FLAG) \
 		--tag $(MLRUN_SYSTEM_TEST_IMAGE_NAME) .
 
@@ -619,7 +680,7 @@ test: clean ## Run mlrun tests
 
 
 .PHONY: test-integration-dockerized
-test-integration-dockerized: build-test ## Run mlrun integration tests in docker container
+test-integration-dockerized: build-test api ## Run mlrun integration tests in docker container, some tests require the api image to be built
 	COVERAGE_MOUNT_PATH="/tmp/coverage_reports/integration_tests" ;\
 	$(SETUP_COVERAGE_MOUNTING)  && \
 	docker run \

{mlrun-1.10.0rc19/mlrun.egg-info → mlrun-1.10.0rc20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlrun
-Version: 1.10.0rc19
+Version: 1.10.0rc20
 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.9
+Requires-Dist: storey~=1.10.10
 Requires-Dist: inflection~=0.5.0
 Requires-Dist: python-dotenv~=1.0
 Requires-Dist: setuptools>=75.2
@@ -112,7 +112,7 @@ Requires-Dist: apscheduler<4,>=3.11; extra == "api"
 Requires-Dist: objgraph~=3.6; extra == "api"
 Requires-Dist: igz-mgmt~=0.4.1; extra == "api"
 Requires-Dist: humanfriendly~=10.0; extra == "api"
-Requires-Dist: fastapi~=0.115.6; extra == "api"
+Requires-Dist: fastapi~=0.116.0; extra == "api"
 Requires-Dist: sqlalchemy~=2.0; extra == "api"
 Requires-Dist: sqlalchemy-utils~=0.41.2; extra == "api"
 Requires-Dist: pymysql~=1.1; extra == "api"
@@ -203,7 +203,7 @@ Requires-Dist: dask~=2024.12.1; python_version >= "3.11" and extra == "complete-
 Requires-Dist: databricks-sdk~=0.20.0; extra == "complete-api"
 Requires-Dist: distributed~=2023.12.1; python_version < "3.11" and extra == "complete-api"
 Requires-Dist: distributed~=2024.12.1; python_version >= "3.11" and extra == "complete-api"
-Requires-Dist: fastapi~=0.115.6; extra == "complete-api"
+Requires-Dist: fastapi~=0.116.0; extra == "complete-api"
 Requires-Dist: gcsfs<2024.7,>=2023.9.2; extra == "complete-api"
 Requires-Dist: google-cloud-bigquery-storage~=2.17; extra == "complete-api"
 Requires-Dist: google-cloud-bigquery[bqstorage,pandas]==3.14.1; extra == "complete-api"

{mlrun-1.10.0rc19 → mlrun-1.10.0rc20}/dockerfiles/mlrun-api/requirements.txt

@@ -6,7 +6,7 @@ apscheduler>=3.11,<4
 objgraph~=3.6
 igz-mgmt~=0.4.1
 humanfriendly~=10.0
-fastapi~=0.115.6
+fastapi~=0.116.0
 sqlalchemy~=2.0
 sqlalchemy-utils~=0.41.2
 pymysql~=1.1

{mlrun-1.10.0rc19 → mlrun-1.10.0rc20}/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-1.10.0rc19 → mlrun-1.10.0rc20}/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-1.10.0rc19 → mlrun-1.10.0rc20}/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-1.10.0rc19 → mlrun-1.10.0rc20}/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-1.10.0rc19 → mlrun-1.10.0rc20}/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")