mlrun 1.10.0rc13__py3-none-any.whl → 1.10.0rc42__py3-none-any.whl

mlrun/datastore/model_provider/openai_provider.py

@@ -11,13 +11,21 @@
 # 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 Callable, Optional, TypeVar, Union
+import inspect
+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
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from openai._models import BaseModel  # noqa
+    from openai.types.chat.chat_completion import ChatCompletion
 
 
 class OpenAIProvider(ModelProvider):
@@ -33,6 +41,9 @@ class OpenAIProvider(ModelProvider):
     operations tailored to the OpenAI API.
     """
 
+    support_async = True
+    response_class = None
+
     def __init__(
         self,
         parent,
@@ -56,7 +67,31 @@ class OpenAIProvider(ModelProvider):
             default_invoke_kwargs=default_invoke_kwargs,
         )
         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 text content of the first choice from an OpenAI ChatCompletion response.
+        Only supports responses with a single choice. Raises an error if multiple choices exist.
+
+        :param response: The ChatCompletion response from OpenAI.
+        :return: The text content of the first message in the response.
+        :raises MLRunInvalidArgumentError: If the response contains more than one choice.
+        """
+        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):
@@ -67,32 +102,58 @@ class OpenAIProvider(ModelProvider):
         return endpoint, subpath
 
     @property
-    def model(self):
-        return self.endpoint
-
-    def load_client(self) -> None:
+    def client(self) -> Any:
         """
-        Initializes the OpenAI SDK client using the provided options.
+        Lazily return the synchronous OpenAI client.
 
-        This method imports the `OpenAI` class from the `openai` package, instantiates
-        a client with the given keyword arguments (`self.options`), and assigns it to
-        `self._client`.
+        If the client has not been initialized yet, it will be created
+        by calling `load_client`.
+        """
+        self.load_client()
+        return self._client
 
-        It also sets the default operation to `self.client.chat.completions.create`, which is
-        typically used for invoking chat-based model completions.
+    def load_client(self) -> None:
+        """
+        Lazily initialize the synchronous OpenAI client.
 
-        Raises:
-            ImportError: If the `openai` package is not installed.
+        The client is created only if it does not already exist.
+        Raises ImportError if the openai package is not installed.
         """
+        if self._client:
+            return
         try:
             from openai import OpenAI  # noqa
 
             self._client = OpenAI(**self.options)
-            self._default_operation = self.client.chat.completions.create
         except ImportError as exc:
             raise ImportError("openai package is not installed") from exc
 
-    def get_client_options(self):
+    def load_async_client(self) -> None:
+        """
+        Lazily initialize the asynchronous OpenAI client.
+
+        The client is created only if it does not already exist.
+        Raises ImportError if the openai package is not installed.
+        """
+        if not self._async_client:
+            try:
+                from openai import AsyncOpenAI  # noqa
+
+                self._async_client = AsyncOpenAI(**self.options)
+            except ImportError as exc:
+                raise ImportError("openai package is not installed") from exc
+
+    @property
+    def async_client(self) -> Any:
+        """
+        Return the asynchronous OpenAI client, creating it on first access.
+
+        The client is lazily initialized via `load_async_client`.
+        """
+        self.load_async_client()
+        return self._async_client
+
+    def get_client_options(self) -> dict:
         res = dict(
             api_key=self._get_secret_or_env("OPENAI_API_KEY"),
             organization=self._get_secret_or_env("OPENAI_ORG_ID"),
@@ -103,42 +164,234 @@ class OpenAIProvider(ModelProvider):
         )
         return self._sanitize_options(res)
 
-    def customized_invoke(
-        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
-    ) -> Optional[T]:
+    def custom_invoke(
+        self, operation: Optional[Callable] = None, **invoke_kwargs
+    ) -> Union["ChatCompletion", "BaseModel"]:
+        """
+        Invokes a model operation from the OpenAI client with the given keyword arguments.
+
+        This method provides flexibility to either:
+        - Call a specific OpenAI client operation (e.g., `client.images.generate`).
+        - Default to `chat.completions.create` when no operation is provided.
+
+        The operation must be a callable that accepts keyword arguments. If the callable
+        does not accept a `model` parameter, it will be omitted from the call.
+
+        Example:
+            ```python
+            result = openai_model_provider.custom_invoke(
+                openai_model_provider.client.images.generate,
+                prompt="A futuristic cityscape at sunset",
+                n=1,
+                size="1024x1024",
+            )
+            ```
+
+        :param operation:       A callable representing the OpenAI operation to invoke.
+                                If not provided, defaults to `client.chat.completions.create`.
+
+        :param invoke_kwargs:   Additional keyword arguments to pass to the operation.
+                                These are merged with `default_invoke_kwargs` and may
+                                include parameters such as `temperature`, `max_tokens`,
+                                or `messages`.
+
+        :return:                The full response returned by the operation, typically
+                                an OpenAI `ChatCompletion` or other OpenAI SDK model.
+        """
+
+        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        model_kwargs = {"model": invoke_kwargs.pop("model", None) or self.model}
+
+        if operation:
+            if not callable(operation):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "OpenAI custom_invoke operation must be a callable"
+                )
+            if not accepts_param(operation, "model"):
+                model_kwargs = {}
+            return operation(**invoke_kwargs, **model_kwargs)
+        else:
+            return self.client.chat.completions.create(**invoke_kwargs, **model_kwargs)
+
+    async def async_custom_invoke(
+        self,
+        operation: Optional[Callable[..., Awaitable[Any]]] = None,
+        **invoke_kwargs,
+    ) -> Union["ChatCompletion", "BaseModel"]:
+        """
+        Asynchronously invokes a model operation from the OpenAI client with the given keyword arguments.
+
+        This method provides flexibility to either:
+        - Call a specific async OpenAI client operation (e.g., `async_client.images.generate`).
+        - Default to `chat.completions.create` when no operation is provided.
+
+        The operation must be an async callable that accepts keyword arguments.
+        If the callable does not accept a `model` parameter, it will be omitted from the call.
+
+        Example:
+            ```python
+            result = await openai_model_provider.async_custom_invoke(
+                openai_model_provider.async_client.images.generate,
+                prompt="A futuristic cityscape at sunset",
+                n=1,
+                size="1024x1024",
+            )
+            ```
+
+        :param operation:       An async callable representing the OpenAI operation to invoke.
+                                If not provided, defaults to `async_client.chat.completions.create`.
+
+        :param invoke_kwargs:   Additional keyword arguments to pass to the operation.
+                                These are merged with `default_invoke_kwargs` and may
+                                include parameters such as `temperature`, `max_tokens`,
+                                or `messages`.
+
+        :return:                The full response returned by the awaited operation,
+                                typically an OpenAI `ChatCompletion` or other OpenAI SDK model.
+
+        """
         invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
+        model_kwargs = {"model": invoke_kwargs.pop("model", None) or self.model}
         if operation:
-            return operation(**invoke_kwargs, model=self.model)
+            if not inspect.iscoroutinefunction(operation):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "OpenAI async_custom_invoke operation must be a coroutine function"
+                )
+            if not accepts_param(operation, "model"):
+                model_kwargs = {}
+            return await operation(**invoke_kwargs, **model_kwargs)
         else:
-            return self._default_operation(**invoke_kwargs, model=self.model)
+            return await self.async_client.chat.completions.create(
+                **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:
+                usage = response.to_dict()["usage"]
+                response = {
+                    UsageResponseKeys.ANSWER: str_response,
+                    UsageResponseKeys.USAGE: usage,
+                }
+        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,
-    ) -> Optional[Union[str, T]]:
+    ) -> Union[dict[str, Any], str, "ChatCompletion"]:
         """
         OpenAI-specific implementation of `ModelProvider.invoke`.
-        Invokes an OpenAI model operation using the sync client.
-        For full details, see `ModelProvider.invoke`.
+        Invokes an OpenAI model operation using the synchronous client.
+
+        :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?"}
+                ]
+
+            Defaults to None if no messages are provided.
+
+        :param invoke_response_format:
+            Specifies the format of the returned response. Options:
+
+            - "string": Returns only the generated text content, taken from a single response.
+            - "usage": Combines the generated text with metadata (e.g., token usage), returning a dictionary::
 
-        :param messages:    Same as ModelProvider.invoke.
+                .. code-block:: json
+                   {
+                       "answer": "<generated_text>",
+                       "usage": <ChatCompletion>.to_dict()["usage"]
+                   }
 
-        :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.).
+            - "full": Returns the full OpenAI `ChatCompletion` object.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the OpenAI client.
 
+        :return:
+            A string, dictionary, or `ChatCompletion` object, depending on `invoke_response_format`.
         """
-        invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
-        response = self._default_operation(
-            model=self.endpoint, messages=messages, **invoke_kwargs
+
+        response = self.custom_invoke(messages=messages, **invoke_kwargs)
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )
+
+    async def async_invoke(
+        self,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
+        **invoke_kwargs,
+    ) -> Union[str, "ChatCompletion", dict]:
+        """
+        OpenAI-specific implementation of `ModelProvider.async_invoke`.
+        Invokes an OpenAI model operation using the asynchronous client.
+
+        :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?"}
+                ]
+
+            Defaults to None if no messages are provided.
+
+        :param invoke_response_format:
+            Specifies the format of the returned response. Options:
+
+            - "string": Returns only the generated text content, taken 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": <ChatCompletion>.to_dict()["usage"]
+                   }
+
+            - "full": Returns the full OpenAI `ChatCompletion` object.
+
+        :param invoke_kwargs:
+            Additional keyword arguments passed to the OpenAI client.
+
+        :return:
+            A string, dictionary, or `ChatCompletion` object, depending on `invoke_response_format`.
+        """
+        response = await self.async_custom_invoke(messages=messages, **invoke_kwargs)
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
         )
-        if as_str:
-            return response.choices[0].message.content
-        return response

mlrun/datastore/s3.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import time
+import warnings
 from typing import Optional
 from urllib.parse import urlparse
 
@@ -28,6 +29,27 @@ from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 class S3Store(DataStore):
     using_bucket = True
 
+    # TODO: Remove this in 1.12.0
+    def _get_endpoint_url_with_deprecation_warning(self):
+        """Get S3 endpoint URL with backward compatibility for deprecated S3_ENDPOINT_URL"""
+        # First try the new environment variable
+        endpoint_url = self._get_secret_or_env("AWS_ENDPOINT_URL_S3")
+        if endpoint_url:
+            return endpoint_url
+
+        # Check for deprecated environment variable
+        deprecated_endpoint_url = self._get_secret_or_env("S3_ENDPOINT_URL")
+        if deprecated_endpoint_url:
+            warnings.warn(
+                "S3_ENDPOINT_URL is deprecated in 1.10.0 and will be removed in 1.12.0, "
+                "use AWS_ENDPOINT_URL_S3 instead.",
+                # TODO: Remove this in 1.12.0
+                FutureWarning,
+            )
+            return deprecated_endpoint_url
+
+        return None
+
     def __init__(
         self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
     ):
@@ -41,7 +63,7 @@ class S3Store(DataStore):
         access_key_id = self._get_secret_or_env("AWS_ACCESS_KEY_ID")
         secret_key = self._get_secret_or_env("AWS_SECRET_ACCESS_KEY")
         token_file = self._get_secret_or_env("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE")
-        endpoint_url = self._get_secret_or_env("S3_ENDPOINT_URL")
+        endpoint_url = self._get_endpoint_url_with_deprecation_warning()
         force_non_anonymous = self._get_secret_or_env("S3_NON_ANONYMOUS")
         profile_name = self._get_secret_or_env("AWS_PROFILE")
         assume_role_arn = self._get_secret_or_env("MLRUN_AWS_ROLE_ARN")
@@ -159,7 +181,7 @@ class S3Store(DataStore):
     def get_storage_options(self):
         force_non_anonymous = self._get_secret_or_env("S3_NON_ANONYMOUS")
         profile = self._get_secret_or_env("AWS_PROFILE")
-        endpoint_url = self._get_secret_or_env("S3_ENDPOINT_URL")
+        endpoint_url = self._get_endpoint_url_with_deprecation_warning()
         access_key_id = self._get_secret_or_env("AWS_ACCESS_KEY_ID")
         secret = self._get_secret_or_env("AWS_SECRET_ACCESS_KEY")
         token_file = self._get_secret_or_env("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE")

mlrun/datastore/storeytargets.py

@@ -18,10 +18,9 @@ from mergedeep import merge
 from storey import V3ioDriver
 
 import mlrun
-import mlrun.model_monitoring.helpers
 from mlrun.datastore.base import DataStore
 from mlrun.datastore.datastore_profile import (
-    DatastoreProfileKafkaSource,
+    DatastoreProfileKafkaStream,
     DatastoreProfileKafkaTarget,
     DatastoreProfileTDEngine,
     datastore_profile_read,
@@ -138,7 +137,7 @@ class KafkaStoreyTarget(storey.KafkaTarget):
             datastore_profile = datastore_profile_read(path)
             if not isinstance(
                 datastore_profile,
-                (DatastoreProfileKafkaSource, DatastoreProfileKafkaTarget),
+                (DatastoreProfileKafkaStream, DatastoreProfileKafkaTarget),
             ):
                 raise mlrun.errors.MLRunInvalidArgumentError(
                     f"Unsupported datastore profile type: {type(datastore_profile)}"

mlrun/datastore/utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import inspect
 import math
 import tarfile
 import tempfile
@@ -319,7 +320,13 @@ def parse_url(url):
     parsed_url = urlparse(url)
     schema = parsed_url.scheme.lower()
     endpoint = parsed_url.hostname
-    if endpoint:
+
+    # Special handling for WASBS URLs to preserve container information
+    if schema in ["wasbs", "wasb"] and parsed_url.netloc and "@" in parsed_url.netloc:
+        # For wasbs://container@host format, preserve the full netloc as endpoint
+        # This allows the datastore to extract container later
+        endpoint = parsed_url.netloc
+    elif endpoint:
         # HACK - urlparse returns the hostname after in lower case - we want the original case:
         # the hostname is a substring of the netloc, in which it's the original case, so we find the indexes of the
         # hostname in the netloc and take it from there
@@ -330,6 +337,11 @@ def parse_url(url):
         endpoint = netloc[
             hostname_index_in_netloc : hostname_index_in_netloc + len(lower_hostname)
         ]
-    if parsed_url.port:
-        endpoint += f":{parsed_url.port}"
+        if parsed_url.port:
+            endpoint += f":{parsed_url.port}"
     return schema, endpoint, parsed_url
+
+
+def accepts_param(func: callable, param_name):
+    sig = inspect.signature(func)
+    return param_name in sig.parameters

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
@@ -55,6 +53,12 @@ class RunDBInterface(ABC):
     def update_run(self, updates: dict, uid, project="", iter=0):
         pass
 
+    @abstractmethod
+    def set_run_retrying_status(
+        self, project: str, name: str, run_id: str, retrying: bool
+    ):
+        pass
+
     @abstractmethod
     def abort_run(self, uid, project="", iter=0, timeout=45, status_text=""):
         pass
@@ -439,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,
@@ -638,6 +625,11 @@ class RunDBInterface(ABC):
     ):
         pass
 
+    def wait_for_background_task_to_reach_terminal_state(
+        self, name: str, project: str = ""
+    ) -> mlrun.common.schemas.BackgroundTask:
+        pass
+
     @abstractmethod
     def retry_pipeline(
         self,
@@ -730,6 +722,9 @@ class RunDBInterface(ABC):
         tsdb_metrics: bool = False,
         metric_list: Optional[list[str]] = None,
         top_level: bool = False,
+        modes: Optional[
+            Union[mm_constants.EndpointMode, list[mm_constants.EndpointMode]]
+        ] = None,
         uids: Optional[list[str]] = None,
         latest_only: bool = False,
     ) -> mlrun.common.schemas.ModelEndpointList:
@@ -781,6 +776,7 @@ class RunDBInterface(ABC):
         item_name: Optional[str] = None,
         tag: Optional[str] = None,
         version: Optional[str] = None,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 
@@ -799,6 +795,7 @@ class RunDBInterface(ABC):
         version: Optional[str] = None,
         tag: Optional[str] = None,
         force_refresh: bool = False,
+        object_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 
@@ -810,6 +807,19 @@ class RunDBInterface(ABC):
         version: Optional[str] = None,
         tag: str = "latest",
         force_refresh: bool = False,
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
+    ):
+        pass
+
+    @abstractmethod
+    def get_hub_asset(
+        self,
+        source_name: str,
+        item_name: str,
+        asset_name: str,
+        version: Optional[str] = None,
+        tag: str = "latest",
+        item_type: mlrun.common.schemas.hub.HubSourceType = mlrun.common.schemas.hub.HubSourceType.functions,
     ):
         pass
 
@@ -1118,6 +1128,15 @@ class RunDBInterface(ABC):
     ) -> None:
         pass
 
+    @abstractmethod
+    def delete_model_monitoring_metrics(
+        self,
+        project: str,
+        application_name: str,
+        endpoint_ids: Optional[list[str]] = None,
+    ) -> None:
+        pass
+
     @abstractmethod
     def get_monitoring_function_summaries(
         self,
@@ -1145,3 +1164,12 @@ class RunDBInterface(ABC):
     @abstractmethod
     def get_project_summary(self, project: str) -> mlrun.common.schemas.ProjectSummary:
         pass
+
+    @abstractmethod
+    def get_drift_over_time(
+        self,
+        project: str,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+    ) -> mlrun.common.schemas.model_monitoring.ModelEndpointDriftValues:
+        pass