mlrun 1.10.0rc17__tar.gz → 1.10.0rc19__tar.gz

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

@@ -752,22 +752,6 @@ run-api: api ## Run mlrun api (dockerized)
 		--env MLRUN_HTTPDB__REAL_PATH=$(MLRUN_HTTPDB__REAL_PATH) \
 		$(MLRUN_API_IMAGE_NAME_TAGGED)
 
-.PHONY: run-test-db
-run-test-db:
-	# clean up any previous test db container. Don't remove it after run to be able to debug failures
-	docker rm test-db --force || true
-	docker run \
-		--name=test-db \
-		--volume $(shell pwd):/mlrun \
-		--publish 3306:3306 \
-		--env MYSQL_ROOT_PASSWORD="" \
-		--env MYSQL_ALLOW_EMPTY_PASSWORD="true" \
-		--env MYSQL_ROOT_HOST=% \
-		--env MYSQL_DATABASE="mlrun" \
-		--detach \
-		gcr.io/iguazio/mlrun-mysql:8.0 \
-		--character-set-server=utf8 \
-		--collation-server=utf8_bin
 
 .PHONY: clean-html-docs
 clean-html-docs: ## Clean html docs

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlrun
-Version: 1.10.0rc17
+Version: 1.10.0rc19
 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.8
+Requires-Dist: storey~=1.10.9
 Requires-Dist: inflection~=0.5.0
 Requires-Dist: python-dotenv~=1.0
 Requires-Dist: setuptools>=75.2

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/dev-requirements.txt

@@ -12,7 +12,8 @@ pytest-alembic~=0.11.0
 pytest-httpserver~=1.0
 pytest-timeout~=2.0
 pytest-mock-resources[mysql]~=2.12
-testcontainers[k3s]~=4.10
+# pin to 4.10.0 to avoid issues with k3s compatibility
+testcontainers[k3s]~=4.10.0
 aioresponses~=0.7
 requests-mock~=1.8
 httpx~=0.27.0
@@ -47,3 +48,4 @@ openai~=1.88
 transformers~=4.53
 # for hugging face image classifier test:
 tf-keras~=2.18
+pillow~=11.3

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/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-1.10.0rc17 → mlrun-1.10.0rc19}/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-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/common/formatters/artifact.py

@@ -41,6 +41,7 @@ class ArtifactFormat(ObjectFormat, mlrun.common.types.StrEnum):
                     "spec.metrics",
                     "spec.target_path",
                     "spec.parent_uri",
+                    "spec.has_children",
                 ]
             ),
         }[_format]

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/common/schemas/model_monitoring/constants.py

@@ -487,25 +487,33 @@ class ModelMonitoringLabels:
 
 _RESERVED_FUNCTION_NAMES = MonitoringFunctionNames.list() + [SpecialApps.MLRUN_INFRA]
 
+_RESERVED_EVALUATE_FUNCTION_SUFFIX = "-batch"
+
 
 class ModelEndpointMonitoringMetricType(StrEnum):
     RESULT = "result"
     METRIC = "metric"
 
 
+# refer to `mlrun.utils.regex.project_name`
+_INNER_PROJECT_PATTERN = r"[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?"
+PROJECT_PATTERN = rf"^{_INNER_PROJECT_PATTERN}$"
+
+MODEL_ENDPOINT_ID_PATTERN = r"^[a-zA-Z0-9_-]+$"
+
 _FQN_PART_PATTERN = r"[a-zA-Z0-9_-]+"
+_RESULT_NAME_PATTERN = r"[a-zA-Z_][a-zA-Z0-9_]*"
+
 FQN_PATTERN = (
-    rf"^(?P<project>{_FQN_PART_PATTERN})\."
+    rf"^(?P<project>{_INNER_PROJECT_PATTERN})\."
     rf"(?P<app>{_FQN_PART_PATTERN})\."
     rf"(?P<type>{ModelEndpointMonitoringMetricType.RESULT}|{ModelEndpointMonitoringMetricType.METRIC})\."
-    rf"(?P<name>{_FQN_PART_PATTERN})$"
+    rf"(?P<name>{_RESULT_NAME_PATTERN})$"
 )
 FQN_REGEX = re.compile(FQN_PATTERN)
+APP_NAME_REGEX = re.compile(_FQN_PART_PATTERN)
+RESULT_NAME_REGEX = re.compile(_RESULT_NAME_PATTERN)
 
-# refer to `mlrun.utils.regex.project_name`
-PROJECT_PATTERN = r"^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$"
-MODEL_ENDPOINT_ID_PATTERN = r"^[a-zA-Z0-9_-]+$"
-RESULT_NAME_PATTERN = r"[a-zA-Z_][a-zA-Z0-9_]*"
 
 INTERSECT_DICT_KEYS = {
     ModelEndpointMonitoringMetricType.METRIC: "intersect_metrics",

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/config.py

@@ -194,6 +194,7 @@ default_config = {
     "v3io_framesd": "http://framesd:8080",
     "model_providers": {
         "openai_default_model": "gpt-4o",
+        "huggingface_default_model": "microsoft/Phi-3-mini-4k-instruct",
     },
     # default node selector to be applied to all functions - json string base64 encoded format
     "default_function_node_selector": "e30=",
@@ -1238,6 +1239,19 @@ class Config:
         """
         return self.is_running_on_iguazio()
 
+    @staticmethod
+    def get_run_retry_staleness_threshold_timedelta() -> timedelta:
+        """
+        Get the staleness threshold in timedelta for run retries.
+        This is used to determine if a run is stale and should be retried.
+
+        :return: The staleness threshold in timedelta.
+        """
+        staleness_threshold = int(
+            mlrun.mlconf.monitoring.runs.retry.staleness_threshold
+        )
+        return timedelta(minutes=staleness_threshold)
+
     def to_dict(self):
         return copy.deepcopy(self._cfg)
 

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/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-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/datastore/datastore.py

@@ -38,6 +38,7 @@ from ..utils import DB_SCHEMA, RunKeys
 from .base import DataItem, DataStore, HttpStore
 from .filestore import FileStore
 from .inmem import InMemoryStore
+from .model_provider.huggingface_provider import HuggingFaceProvider
 from .model_provider.openai_provider import OpenAIProvider
 from .store_resources import get_store_resource, is_store_uri
 from .v3io import V3ioStore
@@ -102,8 +103,7 @@ def schema_to_store(schema) -> DataStore.__subclasses__():
 def schema_to_model_provider(
     schema: str, raise_missing_schema_exception=True
 ) -> type[ModelProvider]:
-    #  TODO add hugging face and http
-    schema_dict = {"openai": OpenAIProvider}
+    schema_dict = {"openai": OpenAIProvider, "huggingface": HuggingFaceProvider}
     provider_class = schema_dict.get(schema, None)
     if not provider_class:
         if raise_missing_schema_exception:
@@ -247,7 +247,7 @@ class StoreManager:
 
         if schema == "ds":
             datastore_profile = datastore_profile_read(url, project_name, secrets)
-            secrets = merge(secrets or {}, datastore_profile.secrets() or {})
+            secrets = merge({}, secrets or {}, datastore_profile.secrets() or {})
             url = datastore_profile.url(subpath)
             schema, endpoint, parsed_url = parse_url(url)
             subpath = parsed_url.path
@@ -281,7 +281,7 @@ class StoreManager:
                 endpoint, subpath
             )
             remote_client = remote_client_class(
-                self, schema, cache_key, parsed_url.netloc, secrets=secrets, **kwargs
+                self, schema, cache_key, endpoint, secrets=secrets, **kwargs
             )
             if not secrets and not mlrun.config.is_running_as_api():
                 cache[cache_key] = remote_client

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/datastore/datastore_profile.py

@@ -486,6 +486,31 @@ class OpenAIProfile(DatastoreProfile):
         return f"{self.type}://{subpath.lstrip('/')}"
 
 
+class HuggingFaceProfile(DatastoreProfile):
+    type: str = pydantic.v1.Field("huggingface")
+    _private_attributes = ("token", "model_kwargs")
+    task: typing.Optional[str] = None
+    token: typing.Optional[str] = None
+    device: typing.Optional[typing.Union[int, str]] = None
+    device_map: typing.Union[str, dict[str, typing.Union[int, str]], None] = None
+    trust_remote_code: bool = None
+    model_kwargs: typing.Optional[dict[str, typing.Any]] = None
+
+    def secrets(self) -> dict:
+        keys = {
+            "HF_TASK": self.task,
+            "HF_TOKEN": self.token,
+            "HF_DEVICE": self.device,
+            "HF_DEVICE_MAP": self.device_map,
+            "HF_TRUST_REMOTE_CODE": self.trust_remote_code,
+            "HF_MODEL_KWARGS": self.model_kwargs,
+        }
+        return {k: v for k, v in keys.items() if v}
+
+    def url(self, subpath):
+        return f"{self.type}://{subpath.lstrip('/')}"
+
+
 _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "v3io": DatastoreProfileV3io,
     "s3": DatastoreProfileS3,
@@ -500,6 +525,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "taosws": DatastoreProfileTDEngine,
     "config": ConfigProfile,
     "openai": OpenAIProfile,
+    "huggingface": HuggingFaceProfile,
 }
 
 

mlrun-1.10.0rc19/mlrun/datastore/model_provider/huggingface_provider.py

@@ -0,0 +1,182 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import mlrun
+from mlrun.datastore.model_provider.model_provider import ModelProvider
+
+if TYPE_CHECKING:
+    from transformers.pipelines.base import Pipeline
+    from transformers.pipelines.text_generation import ChatType
+
+
+class HuggingFaceProvider(ModelProvider):
+    """
+    HuggingFaceProvider is a wrapper around the Hugging Face Transformers pipeline
+    that provides an interface for interacting with a wide range of Hugging Face models.
+
+    It supports synchronous operations, enabling flexible integration into various workflows.
+
+    This class extends the ModelProvider base class and implements Hugging Face-specific
+    functionality, including pipeline initialization, default text generation operations,
+    and custom operations tailored to the Hugging Face Transformers pipeline API.
+    """
+
+    def __init__(
+        self,
+        parent,
+        schema,
+        name,
+        endpoint="",
+        secrets: Optional[dict] = None,
+        default_invoke_kwargs: Optional[dict] = None,
+    ):
+        endpoint = endpoint or mlrun.mlconf.model_providers.huggingface_default_model
+        if schema != "huggingface":
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "HuggingFaceProvider supports only 'huggingface' as the provider kind."
+            )
+        super().__init__(
+            parent=parent,
+            kind=schema,
+            name=name,
+            endpoint=endpoint,
+            secrets=secrets,
+            default_invoke_kwargs=default_invoke_kwargs,
+        )
+        self.options = self.get_client_options()
+        self._expected_operation_type = None
+        self.load_client()
+
+    @staticmethod
+    def _extract_string_output(result) -> 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:
+            raise ValueError("Empty or invalid pipeline output")
+
+        return result[0].get("generated_text")
+
+    @classmethod
+    def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
+        if endpoint and subpath:
+            endpoint = endpoint + subpath
+            # In HuggingFace, "/" in a model name is part of the name — `subpath` is not used.
+            subpath = ""
+        return endpoint, subpath
+
+    def load_client(self) -> None:
+        """
+        Initializes the Hugging Face pipeline using the provided options.
+
+        This method imports the `pipeline` function from the `transformers` package,
+        creates a pipeline instance with the specified task and model (from `self.options`),
+        and assigns it to `self._client`.
+
+        Note: Hugging Face pipelines are synchronous and do not support async invocation.
+
+        Raises:
+            ImportError: If the `transformers` package is not installed.
+        """
+        try:
+            from transformers import pipeline, AutoModelForCausalLM  # noqa
+            from transformers import AutoTokenizer  # noqa
+            from transformers.pipelines.base import Pipeline  # noqa
+
+            self._client = pipeline(model=self.model, **self.options)
+            self._expected_operation_type = Pipeline
+        except ImportError as exc:
+            raise ImportError("transformers package is not installed") from exc
+
+    def get_client_options(self):
+        res = dict(
+            task=self._get_secret_or_env("HF_TASK") or "text-generation",
+            token=self._get_secret_or_env("HF_TOKEN"),
+            device=self._get_secret_or_env("HF_DEVICE"),
+            device_map=self._get_secret_or_env("HF_DEVICE_MAP"),
+            trust_remote_code=self._get_secret_or_env("HF_TRUST_REMOTE_CODE"),
+            model_kwargs=self._get_secret_or_env("HF_MODEL_KWARGS"),
+        )
+        return self._sanitize_options(res)
+
+    def custom_invoke(
+        self, operation: Optional["Pipeline"] = None, **invoke_kwargs
+    ) -> Union[list, dict, Any]:
+        """
+        HuggingFace implementation of `ModelProvider.custom_invoke`.
+        Use the default config in provider client/ user defined client:
+
+        Example:
+        ```python
+            image = Image.open(image_path)
+            pipeline_object =  pipeline("image-classification", model="microsoft/resnet-50")
+            result = hf_provider.custom_invoke(
+                pipeline_object,
+                inputs=image,
+            )
+        ```
+
+
+        :param operation:               A pipeline object
+        :param invoke_kwargs:           Keyword arguments to pass to the operation.
+        :return:                        The full response returned by the operation.
+
+        """
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        if operation:
+            if not isinstance(operation, self._expected_operation_type):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Huggingface operation must inherit" " from 'Pipeline' object"
+                )
+            return operation(**invoke_kwargs)
+        else:
+            return self.client(**invoke_kwargs)
+
+    def invoke(
+        self,
+        messages: Union[str, list[str], "ChatType", list["ChatType"]] = None,
+        as_str: bool = False,
+        **invoke_kwargs,
+    ) -> Union[str, list]:
+        """
+        HuggingFace-specific implementation of `ModelProvider.invoke`.
+        Invokes a HuggingFace model operation using the synchronous client.
+        For complete usage details, refer to `ModelProvider.invoke`.
+
+        :param messages:
+                            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.
+
+                            If `False`, return the **full raw response object**, which is a list of dictionaries.
+
+        :param invoke_kwargs:
+                            Same as ModelProvider.invoke.
+        :return:            Same as ModelProvider.invoke.
+        """
+        if self.client.task != "text-generation":
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "HuggingFaceProvider.invoke supports text-generation task only"
+            )
+        if as_str:
+            invoke_kwargs["return_full_text"] = False
+        response = self.custom_invoke(text_inputs=messages, **invoke_kwargs)
+        if as_str:
+            return self._extract_string_output(response)
+        return response

{mlrun-1.10.0rc17 → mlrun-1.10.0rc19}/mlrun/datastore/model_provider/model_provider.py

@@ -12,15 +12,13 @@
 # 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.datastore.remote_client import (
     BaseRemoteClient,
 )
 
-T = TypeVar("T")
-
 
 class ModelProvider(BaseRemoteClient):
     """
@@ -79,12 +77,66 @@ class ModelProvider(BaseRemoteClient):
 
         raise NotImplementedError("load_client method is not implemented")
 
+    @property
+    def client(self) -> Any:
+        return self._client
+
+    @property
+    def model(self) -> Optional[str]:
+        """
+        Returns the model identifier used by the underlying SDK.
+
+        :return: A string representing the model ID, or None if not set.
+        """
+        return self.endpoint
+
+    def get_invoke_kwargs(self, invoke_kwargs) -> dict:
+        kwargs = self.default_invoke_kwargs.copy()
+        kwargs.update(invoke_kwargs)
+        return kwargs
+
+    @property
+    def async_client(self) -> Any:
+        if not self.support_async:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"{self.__class__.__name__} does not support async operations"
+            )
+        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[Any]]], **invoke_kwargs
+    ) -> Any:
+        """
+        Asynchronously invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.)
+        with the given keyword arguments.
+
+        The operation must be an async callable (e.g., a method from an async client) that accepts keyword arguments.
+
+        :param operation:       An async callable representing the model operation (e.g., an async_client method).
+        :param invoke_kwargs:   Keyword arguments to pass to the operation.
+        :return:                The full response returned by the awaited operation.
+        """
+        raise NotImplementedError("async_custom_invoke is not implemented")
+
     def invoke(
         self,
         messages: Optional[list[dict]] = None,
         as_str: bool = False,
         **invoke_kwargs,
-    ) -> Optional[Union[str, T]]:
+    ) -> Union[str, 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
@@ -127,62 +179,11 @@ class ModelProvider(BaseRemoteClient):
         """
         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
-
-    @property
-    def model(self) -> Optional[str]:
-        return None
-
-    def get_invoke_kwargs(self, invoke_kwargs) -> dict:
-        kwargs = self.default_invoke_kwargs.copy()
-        kwargs.update(invoke_kwargs)
-        return kwargs
-
-    @property
-    def async_client(self) -> Any:
-        if not self.support_async:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"{self.__class__.__name__} does not support async operations"
-            )
-        return self._async_client
-
-    async def async_custom_invoke(
-        self, operation: Optional[Callable[..., Awaitable[T]]], **invoke_kwargs
-    ) -> Optional[T]:
-        """
-        Asynchronously invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.)
-        with the given keyword arguments.
-
-        The operation must be an async callable (e.g., a method from an async client) that accepts keyword arguments.
-
-        :param operation:       An async callable representing the model operation (e.g., an async_client method).
-        :param invoke_kwargs:   Keyword arguments to pass to the operation.
-        :return:                The full response returned by the awaited operation.
-        """
-        raise NotImplementedError("async_custom_invoke is not implemented")
-
     async def async_invoke(
         self,
         messages: Optional[list[dict]] = None,
         as_str: bool = False,
         **invoke_kwargs,
-    ) -> Optional[str]:
+    ) -> Union[str, Any]:
         """Async version of `invoke`. See `invoke` for full documentation."""
         raise NotImplementedError("async_invoke is not implemented")