mlrun 1.10.0rc18__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/datastore/model_provider/openai_provider.py

@@ -12,14 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import inspect
-from collections.abc import Awaitable
-from typing import Callable, Optional, TypeVar, Union
+from collections.abc import Awaitable, Callable
+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,
+)
 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):
@@ -36,6 +42,7 @@ class OpenAIProvider(ModelProvider):
     """
 
     support_async = True
+    response_class = None
 
     def __init__(
         self,
@@ -60,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):
@@ -70,25 +101,58 @@ class OpenAIProvider(ModelProvider):
             subpath = ""
         return endpoint, subpath
 
-    def load_client(self) -> None:
+    @property
+    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` and `self._async_client`.
+        If the client has not been initialized yet, it will be created
+        by calling `load_client`.
+        """
+        self.load_client()
+        return self._client
+
+    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, AsyncOpenAI  # noqa
+            from openai import OpenAI  # noqa
 
             self._client = OpenAI(**self.options)
-            self._async_client = AsyncOpenAI(**self.options)
         except ImportError as exc:
             raise ImportError("openai package is not installed") from exc
 
+    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"),
@@ -101,28 +165,40 @@ class OpenAIProvider(ModelProvider):
         return self._sanitize_options(res)
 
     def custom_invoke(
-        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
-    ) -> Optional[T]:
+        self, operation: Optional[Callable] = None, **invoke_kwargs
+    ) -> Union["ChatCompletion", "BaseModel"]:
         """
-        OpenAI-specific implementation of `ModelProvider.custom_invoke`.
+        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.
 
-        Invokes an OpenAI model operation using the sync client. For full details, see
-        `ModelProvider.custom_invoke`.
+        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.invoke(
+            result = openai_model_provider.custom_invoke(
                 openai_model_provider.client.images.generate,
                 prompt="A futuristic cityscape at sunset",
                 n=1,
                 size="1024x1024",
             )
             ```
-        :param      operation:      Same as ModelProvider.custom_invoke.
-        :param      invoke_kwargs:  Same as ModelProvider.custom_invoke.
-        :return:                    Same as ModelProvider.custom_invoke.
 
+        :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}
 
@@ -139,28 +215,39 @@ class OpenAIProvider(ModelProvider):
 
     async def async_custom_invoke(
         self,
-        operation: Optional[Callable[..., Awaitable[T]]] = None,
+        operation: Optional[Callable[..., Awaitable[Any]]] = None,
         **invoke_kwargs,
-    ) -> Optional[T]:
+    ) -> Union["ChatCompletion", "BaseModel"]:
         """
-        OpenAI-specific implementation of `ModelProvider.async_custom_invoke`.
+        Asynchronously invokes a model operation from the OpenAI client with the given keyword arguments.
 
-        Invokes an OpenAI model operation using the async client. For full details, see
-        `ModelProvider.async_custom_invoke`.
+        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 = openai_model_provider.invoke(
+            ```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 operation:       Same as ModelProvider.async_custom_invoke.
-        :param invoke_kwargs:   Same as ModelProvider.async_custom_invoke.
-        :return:                Same as ModelProvider.async_custom_invoke.
+        :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)
@@ -178,60 +265,133 @@ class OpenAIProvider(ModelProvider):
                 **invoke_kwargs, **model_kwargs
             )
 
+    def _response_handler(
+        self,
+        response: "ChatCompletion",
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **kwargs,
+    ) -> ["ChatCompletion", str, dict[str, Any]]:
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
+            str_response = self._extract_string_output(response)
+            if invoke_response_format == InvokeResponseFormat.STRING:
+                return str_response
+            if invoke_response_format == InvokeResponseFormat.USAGE:
+                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:
 
-        :param messages:    Same as ModelProvider.invoke.
+            - "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 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.).
+                .. code-block:: json
+                   {
+                       "answer": "<generated_text>",
+                       "usage": <ChatCompletion>.to_dict()["usage"]
+                   }
+
+            - "full": Returns the full OpenAI `ChatCompletion` object.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
-        :return:            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the OpenAI client.
 
+        :return:
+            A string, dictionary, or `ChatCompletion` object, depending on `invoke_response_format`.
         """
+
         response = self.custom_invoke(messages=messages, **invoke_kwargs)
-        if as_str:
-            return response.choices[0].message.content
-        return response
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )
 
     async def async_invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> str:
+    ) -> Union[str, "ChatCompletion", dict]:
         """
         OpenAI-specific implementation of `ModelProvider.async_invoke`.
-        Invokes an OpenAI model operation using the async client.
-        For full details, see `ModelProvider.async_invoke`.
+        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
 
-        :param messages:    Same as ModelProvider.async_invoke.
+                [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "What is the capital of France?"}
+                ]
 
-        :param as_str: bool
-                            If `True`, returns only the main content of the first response
-                            (`response.choices[0].message.content`).
-                            If `False`, returns the full awaited response object, whose type depends on
-                            the specific OpenAI SDK operation used (e.g., chat completion, completion, etc.).
+            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:
-                            Same as ModelProvider.async_invoke.
-        :returns            Same as ModelProvider.async_invoke.
+            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)
-        if as_str:
-            return response.choices[0].message.content
-        return response
+        return self._response_handler(
+            messages=messages,
+            invoke_response_format=invoke_response_format,
+            response=response,
+        )

mlrun/datastore/s3.py

@@ -13,21 +13,47 @@
 # limitations under the License.
 
 import time
+import warnings
 from typing import Optional
 from urllib.parse import urlparse
 
 import boto3
+import botocore.exceptions
 from boto3.s3.transfer import TransferConfig
 from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
 from .base import DataStore, FileStats, make_datastore_schema_sanitizer
+from .utils import parse_s3_bucket_and_key
+
+__all__ = ["parse_s3_bucket_and_key"]
 
 
 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 +67,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 +185,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")
@@ -203,9 +229,17 @@ class S3Store(DataStore):
     def get(self, key, size=None, offset=0):
         bucket, key = self.get_bucket_and_key(key)
         obj = self.s3.Object(bucket, key)
-        if size or offset:
-            return obj.get(Range=S3Store.get_range(size, offset))["Body"].read()
-        return obj.get()["Body"].read()
+        try:
+            if size or offset:
+                return obj.get(Range=S3Store.get_range(size, offset))["Body"].read()
+            return obj.get()["Body"].read()
+
+        except botocore.exceptions.ClientError as exc:
+            if exc.response["Error"]["Code"] == "NoSuchKey":
+                # "NoSuchKey" errors codes - equivalent to `FileNotFoundError`
+                raise FileNotFoundError(f"s3://{bucket}/{key}") from exc
+            # Other errors are raised as-is
+            raise
 
     def put(self, key, data, append=False):
         data, _ = self._prepare_put_data(data, append)
@@ -237,16 +271,3 @@ class S3Store(DataStore):
         #  In order to raise an error if there is connection error, ML-7056.
         self.filesystem.exists(path=path)
         self.filesystem.rm(path=path, recursive=recursive, maxdepth=maxdepth)
-
-
-def parse_s3_bucket_and_key(s3_path):
-    try:
-        path_parts = s3_path.replace("s3://", "").split("/")
-        bucket = path_parts.pop(0)
-        key = "/".join(path_parts)
-    except Exception as exc:
-        raise mlrun.errors.MLRunInvalidArgumentError(
-            "failed to parse s3 bucket and key"
-        ) from exc
-
-    return bucket, key

mlrun/datastore/sources.py

@@ -460,7 +460,7 @@ class ParquetSource(BaseSourceDriver):
             if not filter_tuple:
                 continue
             col_name, op, value = filter_tuple
-            if op.lower() in ("in", "not in") and isinstance(value, (list, tuple, set)):
+            if op.lower() in ("in", "not in") and isinstance(value, list | tuple | set):
                 none_exists = False
                 value = list(value)
                 for sub_value in value:

mlrun/datastore/store_resources.py

@@ -76,9 +76,9 @@ class ResourceCache:
             return self._tabels[uri]
 
         if uri.startswith("v3io://") or uri.startswith("v3ios://"):
-            endpoint, uri = parse_path(uri)
+            endpoint, path = parse_path(uri)
             self._tabels[uri] = Table(
-                uri,
+                path,
                 V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api),
                 flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
             )
@@ -87,10 +87,10 @@ class ResourceCache:
         if uri.startswith("redis://") or uri.startswith("rediss://"):
             from storey.redis_driver import RedisDriver
 
-            endpoint, uri = parse_path(uri)
+            endpoint, path = parse_path(uri)
             endpoint = endpoint or mlrun.mlconf.redis.url
             self._tabels[uri] = Table(
-                uri,
+                path,
                 RedisDriver(redis_url=endpoint, key_prefix="/"),
                 flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
             )

mlrun/datastore/storeytargets.py

@@ -18,12 +18,12 @@ from mergedeep import merge
 from storey import V3ioDriver
 
 import mlrun
-import mlrun.model_monitoring.helpers
+import mlrun.common.model_monitoring.helpers
 from mlrun.datastore.base import DataStore
 from mlrun.datastore.datastore_profile import (
-    DatastoreProfileKafkaSource,
+    DatastoreProfileKafkaStream,
     DatastoreProfileKafkaTarget,
-    DatastoreProfileTDEngine,
+    DatastoreProfilePostgreSQL,
     datastore_profile_read,
 )
 
@@ -49,17 +49,22 @@ def get_url_and_storage_options(path, external_storage_options=None):
     return url, DataStore._sanitize_options(storage_options)
 
 
-class TDEngineStoreyTarget(storey.TDEngineTarget):
+class TimescaleDBStoreyTarget(storey.TimescaleDBTarget):
     def __init__(self, *args, url: str, **kwargs):
         if url.startswith("ds://"):
             datastore_profile = datastore_profile_read(url)
-            if not isinstance(datastore_profile, DatastoreProfileTDEngine):
-                raise ValueError(
-                    f"Unexpected datastore profile type:{datastore_profile.type}."
-                    "Only DatastoreProfileTDEngine is supported"
+            if not isinstance(datastore_profile, DatastoreProfilePostgreSQL):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Unexpected datastore profile type: {type(datastore_profile)}. "
+                    "Only DatastoreProfilePostgreSQL is supported"
                 )
-            url = datastore_profile.dsn()
-        super().__init__(*args, url=url, **kwargs)
+            # Use the shared helper to determine the correct database name
+            # This ensures consistency with TimescaleDBConnector's database naming
+            database = mlrun.common.model_monitoring.helpers.get_tsdb_database_name(
+                datastore_profile.database
+            )
+            url = datastore_profile.dsn(database=database)
+        super().__init__(*args, dsn=url, **kwargs)
 
 
 class StoreyTargetUtils:
@@ -138,7 +143,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)}"
@@ -175,7 +180,7 @@ class RedisNoSqlStoreyTarget(storey.NoSqlTarget):
         endpoint, uri = mlrun.datastore.targets.RedisNoSqlTarget.get_server_endpoint(
             path
         )
-        kwargs["path"] = endpoint + "/" + uri
+        kwargs["path"] = f"{endpoint}/{uri}"
         super().__init__(*args, **kwargs)
 
 

mlrun/datastore/targets.py

@@ -532,7 +532,7 @@ class BaseStoreTarget(DataTargetBase):
             if (
                 file_system.protocol == "file"
                 # fsspec 2023.10.0 changed protocol from "file" to ("file", "local")
-                or isinstance(file_system.protocol, (tuple, list))
+                or isinstance(file_system.protocol, tuple | list)
                 and "file" in file_system.protocol
             ):
                 dir = os.path.dirname(target_path)