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

mlrun/datastore/model_provider/model_provider.py

@@ -12,14 +12,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Awaitable
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Callable, Optional, Union
 
 import mlrun.errors
+from mlrun.common.types import StrEnum
 from mlrun.datastore.remote_client import (
     BaseRemoteClient,
 )
 
-T = TypeVar("T")
+
+class InvokeResponseFormat(StrEnum):
+    STRING = "string"
+    USAGE = "usage"
+    FULL = "full"
+
+    @classmethod
+    def is_str_response(cls, invoke_response_format: str) -> bool:
+        """
+        Returns True if the response key corresponds to a string-based response (not a full generation object).
+        """
+        return invoke_response_format in {
+            cls.USAGE,
+            cls.STRING,
+        }
+
+
+class UsageResponseKeys(StrEnum):
+    ANSWER = "answer"
+    USAGE = "usage"
+
+    @classmethod
+    def fields(cls) -> list[str]:
+        return [cls.ANSWER, cls.USAGE]
 
 
 class ModelProvider(BaseRemoteClient):
@@ -58,89 +82,61 @@ class ModelProvider(BaseRemoteClient):
         self._client = None
         self._async_client = None
 
-    def get_client_options(self) -> dict:
-        """
-        Returns a dictionary containing credentials and configuration
-        options required for client creation.
-
-        :return:           A dictionary with client-specific settings.
-        """
-        return {}
-
-    def load_client(self) -> None:
+    @staticmethod
+    def _extract_string_output(response: Any) -> str:
         """
-        Initializes the SDK client for the model provider with the given keyword arguments
-        and assigns it to an instance attribute (e.g., self._client).
-
-        Subclasses should override this method to:
-        - Create and configure the provider-specific client instance.
-        - Assign the client instance to self._client.
+        Extracts string response from response object
         """
+        pass
 
-        raise NotImplementedError("load_client method is not implemented")
-
-    def invoke(
+    def _response_handler(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
-        **invoke_kwargs,
-    ) -> Optional[Union[str, T]]:
+        response: Any,
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **kwargs,
+    ) -> Union[str, dict, 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.
+        Handles the model response according to the specified response format.
 
-        :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:
+        :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:
 
-                            .. code-block:: json
+                                       - 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>}
 
-                                [
-                                    {"role": "system", "content": "You are a helpful assistant."},
-                                    {"role": "user", "content": "What is the capital of France?"}
-                                ]
+                                       - FULL: Return the full raw response object.
 
-                            This format is consistent across all backends. Defaults to None if no messages
-                            are provided.
+        :param kwargs:                  Additional parameters that may be required by specific implementations.
 
-        :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.
+        :return:                        The processed response in the format specified by `invoke_response_format`.
+                                        Can be a string, dictionary, or the original response object.
+        """
+        return None
 
-        :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.
+    def get_client_options(self) -> dict:
+        """
+        Returns a dictionary containing credentials and configuration
+        options required for client creation.
 
-        :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:           A dictionary with client-specific settings.
+        """
+        return {}
 
+    def load_client(self) -> None:
         """
-        raise NotImplementedError("invoke method is not implemented")
+        Initialize the SDK client for the model provider and assign it to an instance attribute.
 
-    def custom_invoke(
-        self, operation: Optional[Callable[..., T]] = None, **invoke_kwargs
-    ) -> Optional[T]:
+        Subclasses should override this method to create and configure the provider-specific client.
         """
-        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.
+        raise NotImplementedError("load_client method is not implemented")
 
-        :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")
+    def load_async_client(self) -> Any:
+        raise NotImplementedError("load_async_client method is not implemented")
 
     @property
     def client(self) -> Any:
@@ -148,7 +144,12 @@ class ModelProvider(BaseRemoteClient):
 
     @property
     def model(self) -> Optional[str]:
-        return None
+        """
+        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()
@@ -163,9 +164,24 @@ class ModelProvider(BaseRemoteClient):
             )
         return self._async_client
 
+    def custom_invoke(
+        self, operation: Optional[Callable] = None, **invoke_kwargs
+    ) -> Any:
+        """
+        Invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.) with the given keyword arguments.
+
+        Useful for dynamically calling model methods like text generation, chat completions, or image generation.
+        The operation must be a callable that accepts keyword arguments.
+
+        :param operation:       A callable representing the model operation (e.g., a client method).
+        :param invoke_kwargs:   Keyword arguments to pass to the operation.
+        :return:                The full response returned by the operation.
+        """
+        raise NotImplementedError("custom_invoke method is not implemented")
+
     async def async_custom_invoke(
-        self, operation: Optional[Callable[..., Awaitable[T]]], **invoke_kwargs
-    ) -> Optional[T]:
+        self, operation: Optional[Callable[..., Awaitable[Any]]] = None, **invoke_kwargs
+    ) -> Any:
         """
         Asynchronously invokes a model operation from a provider (e.g., OpenAI, Hugging Face, etc.)
         with the given keyword arguments.
@@ -178,11 +194,132 @@ class ModelProvider(BaseRemoteClient):
         """
         raise NotImplementedError("async_custom_invoke is not implemented")
 
+    def invoke(
+        self,
+        messages: Union[list[dict], Any],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **invoke_kwargs,
+    ) -> Union[str, dict[str, Any], Any]:
+        """
+        Invokes a generative AI model with the provided messages and additional parameters.
+        This method is designed to be a flexible interface for interacting with various
+        generative AI backends (e.g., OpenAI, Hugging Face, etc.). It allows users to send
+        a list of messages (following a standardized format) and receive a response.
+
+        :param messages:            A list of dictionaries representing the conversation history or input messages.
+                                    Each dictionary should follow the format::
+                                    {"role": "system"| "user" | "assistant" ..., "content":
+                                    "Message content as a string"}
+
+                                    Example:
+
+                                    .. code-block:: json
+
+                                        [
+                                            {"role": "system", "content": "You are a helpful assistant."},
+                                            {"role": "user", "content": "What is the capital of France?"}
+                                        ]
+
+                                    This format is consistent across all backends. Defaults to None if no messages
+                                    are provided.
+
+        :param invoke_response_format:   Determines how the model response is returned:
+
+                                    - string:   Returns only the generated text content from the model output,
+                                                for single-answer responses only.
+
+                                    - usage:    Combines the STRING response with additional metadata (token usage),
+                                                and returns the result in a dictionary.
+
+                                                Note: The usage dictionary may contain additional
+                                                keys depending on the model provider:
+
+                                    .. code-block:: json
+
+                                    {
+                                        "answer": "<generated_text>",
+                                        "usage": {
+                                        "prompt_tokens": <int>,
+                                        "completion_tokens": <int>,
+                                        "total_tokens": <int>
+                                        }
+
+                                    }
+
+                                    - full:   Returns the full model output.
+
+        :param invoke_kwargs:
+                                    Additional keyword arguments to be passed to the underlying model API call.
+                                    These can include parameters such as temperature, max tokens, etc.,
+                                    depending on the capabilities of the specific backend being used.
+
+        :return:                    The invoke result formatted according to the specified
+                                    invoke_response_format parameter.
+
+        """
+        raise NotImplementedError("invoke method is not implemented")
+
     async def async_invoke(
         self,
-        messages: Optional[list[dict]] = None,
-        as_str: bool = False,
+        messages: list[dict],
+        invoke_response_format=InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Optional[str]:
-        """Async version of `invoke`. See `invoke` for full documentation."""
+    ) -> Union[str, dict[str, Any], Any]:
+        """
+        Asynchronously invokes a generative AI model with the provided messages and additional parameters.
+        This method is designed to be a flexible interface for interacting with various
+        generative AI backends (e.g., OpenAI, Hugging Face, etc.). It allows users to send
+        a list of messages (following a standardized format) and receive a response.
+
+        :param messages:            A list of dictionaries representing the conversation history or input messages.
+                                    Each dictionary should follow the format::
+                                    {"role": "system"| "user" | "assistant" ..., "content":
+                                    "Message content as a string"}
+
+                                    Example:
+
+                                    .. code-block:: json
+
+                                        [
+                                            {"role": "system", "content": "You are a helpful assistant."},
+                                            {"role": "user", "content": "What is the capital of France?"}
+                                        ]
+
+                                    This format is consistent across all backends. Defaults to None if no messages
+                                    are provided.
+
+        :param invoke_response_format:   Determines how the model response is returned:
+
+                                    - string:   Returns only the generated text content from the model output,
+                                                for single-answer responses only.
+
+                                    - usage:    Combines the STRING response with additional metadata (token usage),
+                                                and returns the result in a dictionary.
+
+                                                Note: The usage dictionary may contain additional
+                                                keys depending on the model provider:
+
+                                    .. code-block:: json
+
+                                    {
+                                        "answer": "<generated_text>",
+                                        "usage": {
+                                        "prompt_tokens": <int>,
+                                        "completion_tokens": <int>,
+                                        "total_tokens": <int>
+                                        }
+
+                                    }
+
+                                    - full:   Returns the full model output.
+
+        :param invoke_kwargs:
+                                    Additional keyword arguments to be passed to the underlying model API call.
+                                    These can include parameters such as temperature, max tokens, etc.,
+                                    depending on the capabilities of the specific backend being used.
+
+        :return:                    The invoke result formatted according to the specified
+                                    invoke_response_format parameter.
+
+        """
         raise NotImplementedError("async_invoke is not implemented")