ibm-watsonx-gov 1.3.3__cp313-cp313-macosx_11_0_arm64.whl

ibm_watsonx_gov/providers/llmevalkit/llm/providers/ibm_watsonx_ai/ibm_watsonx_ai.py

@@ -0,0 +1,656 @@
+import os
+from typing import Any, Dict, List, Optional, Type, TypeVar, Union
+
+try:
+    from ibm_watsonx_ai import Credentials
+    from ibm_watsonx_ai.foundation_models import ModelInference
+except ImportError as e:
+    raise ImportError(
+        "Please install the ibm-watsonx-ai package: pip install 'llmevalkit[ibm_watsonx_ai]'"
+    ) from e
+
+from llmevalkit.llm.base import Hook, LLMClient, register_llm
+from llmevalkit.llm.output_parser import ValidatingLLMClient
+from llmevalkit.llm.types import GenerationMode, LLMResponse, ParameterMapper
+from pydantic import BaseModel
+
+from ..consts import WX_API_KEY, WX_PROJECT_ID, WX_SPACE_ID, WX_URL
+
+T = TypeVar("T", bound="WatsonxLLMClient")
+SchemaType = Union[Dict[str, Any], Type["BaseModel"], Type]
+
+# -------------------------------------------------------------------
+# 1. Non-validating Watsonx wrapper
+# -------------------------------------------------------------------
+
+
+@register_llm("watsonx")
+class WatsonxLLMClient(LLMClient):
+    """
+    Adapter for IBM watsonx.ai Foundation Model (via ibm_watsonx_ai.foundation_models.ModelInference).
+
+    Supports:
+      - text:       sync generation (ModelInference.generate)
+      - chat:       sync chat     (ModelInference.chat)
+      - text_async: async generation (ModelInference.agenerate)
+      - chat_async: async chat       (ModelInference.achat)
+    """
+
+    def __init__(
+        self,
+        model_name: Optional[str] = None,
+        api_key: Optional[str] = None,
+        project_id: Optional[str] = None,
+        space_id: Optional[str] = None,
+        deployment_id: Optional[str] = None,
+        url: Optional[str] = None,
+        hooks: Optional[List[Hook]] = None,
+        model_id: Optional[str] = None,
+        **model_kwargs: Any,
+    ) -> None:
+        """
+        Initialize the Watsonx client.
+
+        Args:
+            model_name:   Identifier of the watsonx model (e.g., "meta-llama/llama-3-3-70b-instruct").
+            api_key:    (Optional) Your IBM Cloud API Key for watsonx.ai.
+            project_id: (Optional) watsonx project ID.
+            space_id:   (Optional) watsonx space ID.
+            deployment_id: (Optional) watsonx deployment ID.
+            url:        (Optional) Base URL for the watsonx endpoint (e.g., "https://us-south.ml.cloud.ibm.com").
+            hooks:      Optional observability hooks.
+            model_kwargs: Additional keyword args passed to ModelInference constructor.
+        """
+        self.model_name = model_name
+        self._model_kwargs = model_kwargs
+
+        if not url:
+            url = os.getenv(WX_URL)
+            if not url:
+                raise EnvironmentError(
+                    f"Missing API URL; please set the '{WX_URL}' environment variable."
+                )
+
+        if not api_key:
+            api_key = os.getenv(WX_API_KEY)
+            if not api_key:
+                raise EnvironmentError(
+                    f"Missing API key; please set the '{WX_API_KEY}' environment variable."
+                )
+
+        if not project_id:
+            project_id = os.getenv(WX_PROJECT_ID)
+            if not project_id:
+                if not space_id:
+                    space_id = os.getenv(WX_SPACE_ID)
+                raise EnvironmentError(
+                    f"Missing project ID; please set the '{WX_PROJECT_ID}' or '{WX_SPACE_ID}' environment variable."
+                )
+
+        creds = Credentials(api_key=api_key, url=url)
+
+        if model_id:
+            # If model_id is provided, use it as the model_name
+            model_name = model_id
+
+        self.model_name = model_name
+        self.model_id = model_name
+
+        if not model_name:
+            raise ValueError("model_name or model_id must be provided")
+
+        # Assemble provider_kwargs for LLMClient base class
+        provider_kwargs: Dict[str, Any] = {
+            "model_id": model_name,
+            "credentials": creds,
+        }
+        if project_id:
+            provider_kwargs["project_id"] = project_id
+        elif space_id:
+            provider_kwargs["space_id"] = space_id
+
+        if deployment_id:
+            provider_kwargs["deployment_id"] = deployment_id
+
+        # Pass through any additional ModelInference args (params, space_id, verify, validate, etc.)
+        provider_kwargs.update(model_kwargs)
+
+        # Initialize underlying ModelInference instance via LLMClient logic
+        super().__init__(
+            client=None, client_needs_init=True, hooks=hooks, **provider_kwargs
+        )
+
+    @classmethod
+    def provider_class(cls) -> Type:
+        """
+        Underlying SDK client class for watsonx.ai: ModelInference.
+        """
+        return ModelInference
+
+    def _register_methods(self) -> None:
+        """
+        Register how to call watsonx methods:
+
+          - 'text'       -> ModelInference.generate
+          - 'text_async' -> ModelInference.agenerate
+          - 'chat'       -> ModelInference.chat
+          - 'chat_async' -> ModelInference.achat
+        """
+        self.set_method_config(GenerationMode.TEXT.value, "generate", "prompt")
+        self.set_method_config(
+            GenerationMode.TEXT_ASYNC.value, "agenerate", "prompt")
+        self.set_method_config(GenerationMode.CHAT.value, "chat", "messages")
+        self.set_method_config(
+            GenerationMode.CHAT_ASYNC.value, "achat", "messages")
+
+    def _setup_parameter_mapper(self) -> None:
+        """Setup parameter mapping for IBM WatsonX provider."""
+        self._parameter_mapper = ParameterMapper()
+
+        # Text generation parameters (based on TextGenParameters)
+        self._parameter_mapper.set_text_mapping("temperature", "temperature")
+        self._parameter_mapper.set_text_mapping("top_p", "top_p")
+        self._parameter_mapper.set_text_mapping("top_k", "top_k")
+        self._parameter_mapper.set_text_mapping("max_tokens", "max_new_tokens")
+        self._parameter_mapper.set_text_mapping("min_tokens", "min_new_tokens")
+        self._parameter_mapper.set_text_mapping(
+            "repetition_penalty", "repetition_penalty"
+        )
+        self._parameter_mapper.set_text_mapping("seed", "random_seed")
+        self._parameter_mapper.set_text_mapping(
+            "stop_sequences", "stop_sequences")
+        self._parameter_mapper.set_text_mapping("timeout", "time_limit")
+        self._parameter_mapper.set_text_mapping(
+            "decoding_method", "decoding_method")
+
+        # Chat parameters (based on TextChatParameters)
+        self._parameter_mapper.set_chat_mapping("temperature", "temperature")
+        self._parameter_mapper.set_chat_mapping("top_p", "top_p")
+        self._parameter_mapper.set_chat_mapping("max_tokens", "max_tokens")
+        self._parameter_mapper.set_chat_mapping(
+            "frequency_penalty", "frequency_penalty"
+        )
+        self._parameter_mapper.set_chat_mapping(
+            "presence_penalty", "presence_penalty")
+        self._parameter_mapper.set_chat_mapping("seed", "seed")
+        self._parameter_mapper.set_chat_mapping("stop_sequences", "stop")
+        self._parameter_mapper.set_chat_mapping("timeout", "time_limit")
+        self._parameter_mapper.set_chat_mapping("logprobs", "logprobs")
+        self._parameter_mapper.set_chat_mapping("top_logprobs", "top_logprobs")
+
+        # Custom transforms for complex parameters
+        def transform_echo_text_mode(value, mode):
+            if mode in ["text", "text_async"]:
+                # Text mode can include input text in response
+                return (
+                    {"include_stop_sequence": value}
+                    if "stop" in str(value).lower()
+                    else {}
+                )
+            return {}
+
+        self._parameter_mapper.set_custom_transform(
+            "echo", transform_echo_text_mode)
+
+    def _parse_llm_response(self, raw: Any) -> Union[str, LLMResponse]:
+        """
+        Extract the generated text and tool calls from a watsonx response.
+
+        - For text generation: raw['results'][0]['generated_text']
+        - For chat:           raw['choices'][0]['message']['content']
+        """
+        content = ""
+        tool_calls = []
+
+        # Text‐generation style
+        if isinstance(raw, dict) and "results" in raw:
+            results = raw["results"]
+            if isinstance(results, list) and results:
+                first = results[0]
+                content = first.get("generated_text", "")
+
+        # Chat style
+        elif isinstance(raw, dict) and "choices" in raw:
+            choices = raw["choices"]
+            if isinstance(choices, list) and choices:
+                first = choices[0]
+                msg = first.get("message")
+                if isinstance(msg, dict):
+                    content = msg.get("content", "")
+                    # Extract tool calls if present
+                    if "tool_calls" in msg and msg["tool_calls"]:
+                        tool_calls = []
+                        for tool_call in msg["tool_calls"]:
+                            tool_call_dict = {
+                                "id": tool_call.get("id"),
+                                "type": tool_call.get("type", "function"),
+                                "function": {
+                                    "name": tool_call.get("function", {}).get("name"),
+                                    "arguments": tool_call.get("function", {}).get(
+                                        "arguments"
+                                    ),
+                                },
+                            }
+                            tool_calls.append(tool_call_dict)
+                elif "text" in first:
+                    content = first["text"]
+
+        if not content and not tool_calls:
+            raise ValueError(f"Unexpected watsonx response format: {raw!r}")
+
+        # Return LLMResponse if tool calls exist, otherwise just content
+        if tool_calls:
+            return LLMResponse(content=content, tool_calls=tool_calls)
+        return content
+
+    def generate(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        mode: Union[str, GenerationMode] = GenerationMode.CHAT,
+        generation_args: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Synchronous generation override for WatsonX.
+
+        - If mode is 'chat' and prompt is str, wrap into messages list.
+        - If mode is 'text', prompt must be str or list of strings.
+        - Handle WatsonX-specific params structure.
+        """
+        mode_str = mode.value if isinstance(mode, GenerationMode) else mode
+        mode_str = mode_str.lower()
+
+        if mode_str not in ("text", "chat"):
+            raise KeyError(
+                f"Unsupported mode '{mode_str}' for WatsonxLLMClient.generate"
+            )
+
+        # Normalize prompt format based on mode
+        if mode_str == GenerationMode.CHAT.value:
+            # Chat mode expects list of messages
+            if isinstance(prompt, str):
+                prompt = [{"role": "user", "content": prompt}]
+            elif isinstance(prompt, list):
+                prompt = prompt
+            else:
+                raise TypeError(
+                    "For chat mode, prompt must be a string or List[Dict[str,str]]"
+                )
+        elif mode_str == GenerationMode.TEXT.value:
+            # Text mode expects a string prompt
+            if isinstance(prompt, list):
+                # Convert messages to simple string
+                prompt = "\n".join(
+                    [msg.get("content", "")
+                     for msg in prompt if msg.get("content")]
+                )
+
+        # Handle WatsonX params structure
+        watsonx_kwargs = {}
+
+        # Extract any existing params from kwargs
+        existing_params = kwargs.pop("params", {})
+
+        # Map generation_args to WatsonX parameters if provided
+        if generation_args and self._parameter_mapper:
+            from llmevalkit.llm.types import GenerationArgs
+
+            if isinstance(generation_args, GenerationArgs):
+                mapped_args = self._parameter_mapper.map_args(
+                    generation_args, mode_str)
+                # Merge mapped args with existing params
+                existing_params.update(mapped_args)
+
+        # Set params if we have any
+        if existing_params:
+            watsonx_kwargs["params"] = existing_params
+
+        # Add any other kwargs that aren't generation parameters
+        watsonx_kwargs.update(kwargs)
+
+        return super().generate(prompt=prompt, mode=mode_str, **watsonx_kwargs)
+
+    async def generate_async(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        mode: Union[str, GenerationMode] = GenerationMode.CHAT_ASYNC,
+        generation_args: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Asynchronous generation override for WatsonX.
+
+        - If mode is 'chat_async', wrap prompt into messages.
+        - If mode is 'text_async', prompt must be str or list of strings.
+        - Handle WatsonX-specific params structure.
+        """
+        mode_str = mode.value if isinstance(mode, GenerationMode) else mode
+        mode_str = mode_str.lower()
+
+        if mode_str not in ("text_async", "chat_async"):
+            raise KeyError(
+                f"Unsupported mode '{mode_str}' for WatsonxLLMClient.generate_async"
+            )
+
+        if mode_str == GenerationMode.CHAT_ASYNC.value:
+            # Chat mode expects list of messages
+            if isinstance(prompt, str):
+                prompt = [{"role": "user", "content": prompt}]
+            elif isinstance(prompt, list):
+                prompt = prompt
+            else:
+                raise TypeError(
+                    "For chat_async mode, prompt must be a string or List[Dict[str,str]]"
+                )
+        elif mode_str == GenerationMode.TEXT_ASYNC.value:
+            # Text mode expects a string prompt
+            if isinstance(prompt, list):
+                # Convert messages to simple string
+                prompt = "\n".join(
+                    [msg.get("content", "")
+                     for msg in prompt if msg.get("content")]
+                )
+
+        # Handle WatsonX params structure
+        watsonx_kwargs = {}
+
+        # Extract any existing params from kwargs
+        existing_params = kwargs.pop("params", {})
+
+        # Map generation_args to WatsonX parameters if provided
+        if generation_args and self._parameter_mapper:
+            from llmevalkit.llm.types import GenerationArgs
+
+            if isinstance(generation_args, GenerationArgs):
+                mapped_args = self._parameter_mapper.map_args(
+                    generation_args, mode_str)
+                # Merge mapped args with existing params
+                existing_params.update(mapped_args)
+
+        # Set params if we have any
+        if existing_params:
+            watsonx_kwargs["params"] = existing_params
+
+        # Add any other kwargs that aren't generation parameters
+        watsonx_kwargs.update(kwargs)
+
+        return await super().generate_async(
+            prompt=prompt, mode=mode_str, **watsonx_kwargs
+        )
+
+
+# -------------------------------------------------------------------
+# 2. Validating Watsonx wrapper
+# -------------------------------------------------------------------
+
+
+@register_llm("watsonx.output_val")
+class WatsonxLLMClientOutputVal(ValidatingLLMClient):
+    """
+    Validating adapter for IBM watsonx.ai Foundation Model.
+
+    Extends ValidatingLLMClient to enforce output structure (via JSON Schema,
+    Pydantic models, or simple Python types) on all generate calls,
+    with retries and batch support (sync & async).
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        api_key: Optional[str] = None,
+        project_id: Optional[str] = None,
+        space_id: Optional[str] = None,
+        username: Optional[str] = None,
+        version: Optional[str] = None,
+        instance_id: Optional[str] = None,
+        password: Optional[str] = None,
+        url: Optional[str] = "https://us-south.ml.cloud.ibm.com",
+        hooks: Optional[List[Hook]] = None,
+        **model_kwargs: Any,
+    ) -> None:
+        """
+        Initialize a Watsonx client with output validation.
+
+        Args:
+            model_id:   Identifier of the watsonx model.
+            api_key:    Your IBM Cloud API Key.
+            project_id: (Optional) watsonx project ID.
+            url:        (Optional) Base URL for the watsonx endpoint.
+            hooks:      Optional observability hooks.
+            model_kwargs: Additional arguments passed to the ModelInference constructor.
+        """
+        self.model_id = model_id
+        self._model_kwargs = model_kwargs
+
+        creds_args = {"url": url}
+        if api_key is not None:
+            creds_args["api_key"] = api_key
+        if version is not None:
+            creds_args["version"] = version
+            creds_args["instance_id"] = instance_id
+            creds_args["username"] = username
+            if api_key is None and password is not None:
+                creds_args["password"] = password
+
+        creds = Credentials(**creds_args)
+        provider_kwargs: Dict[str, Any] = {
+            "model_id": model_id,
+            "credentials": creds,
+        }
+        if project_id is not None:
+            provider_kwargs["project_id"] = project_id
+        if space_id is not None:
+            provider_kwargs["space_id"] = space_id
+
+        provider_kwargs.update(model_kwargs)
+
+        super().__init__(
+            client=None, client_needs_init=True, hooks=hooks, **provider_kwargs
+        )
+
+    @classmethod
+    def provider_class(cls) -> Type:
+        """
+        Underlying SDK client class: ModelInference.
+        """
+        return ModelInference
+
+    def _register_methods(self) -> None:
+        """
+        Register how to call watsonx methods for validation:
+
+          - 'text'       -> ModelInference.generate
+          - 'text_async' -> ModelInference.agenerate
+          - 'chat'       -> ModelInference.chat
+          - 'chat_async' -> ModelInference.achat
+        """
+        self.set_method_config(GenerationMode.TEXT.value, "generate", "prompt")
+        self.set_method_config(
+            GenerationMode.TEXT_ASYNC.value, "agenerate", "prompt")
+        self.set_method_config(GenerationMode.CHAT.value, "chat", "messages")
+        self.set_method_config(
+            GenerationMode.CHAT_ASYNC.value, "achat", "messages")
+
+    def _setup_parameter_mapper(self) -> None:
+        """Setup parameter mapping for IBM WatsonX provider (same as regular WatsonX)."""
+        self._parameter_mapper = ParameterMapper()
+
+        # Text generation parameters (based on TextGenParameters)
+        self._parameter_mapper.set_text_mapping("temperature", "temperature")
+        self._parameter_mapper.set_text_mapping("top_p", "top_p")
+        self._parameter_mapper.set_text_mapping("top_k", "top_k")
+        self._parameter_mapper.set_text_mapping("max_tokens", "max_new_tokens")
+        self._parameter_mapper.set_text_mapping("min_tokens", "min_new_tokens")
+        self._parameter_mapper.set_text_mapping(
+            "repetition_penalty", "repetition_penalty"
+        )
+        self._parameter_mapper.set_text_mapping("seed", "random_seed")
+        self._parameter_mapper.set_text_mapping(
+            "stop_sequences", "stop_sequences")
+        self._parameter_mapper.set_text_mapping("timeout", "time_limit")
+
+        # Chat parameters (based on TextChatParameters)
+        self._parameter_mapper.set_chat_mapping("temperature", "temperature")
+        self._parameter_mapper.set_chat_mapping("top_p", "top_p")
+        self._parameter_mapper.set_chat_mapping("max_tokens", "max_tokens")
+        self._parameter_mapper.set_chat_mapping(
+            "frequency_penalty", "frequency_penalty"
+        )
+        self._parameter_mapper.set_chat_mapping(
+            "presence_penalty", "presence_penalty")
+        self._parameter_mapper.set_chat_mapping("seed", "seed")
+        self._parameter_mapper.set_chat_mapping("stop_sequences", "stop")
+        self._parameter_mapper.set_chat_mapping("timeout", "time_limit")
+        self._parameter_mapper.set_chat_mapping("logprobs", "logprobs")
+        self._parameter_mapper.set_chat_mapping("top_logprobs", "top_logprobs")
+
+        def transform_echo_text_mode(value, mode):
+            if mode in ["text", "text_async"]:
+                return (
+                    {"include_stop_sequence": value}
+                    if "stop" in str(value).lower()
+                    else {}
+                )
+            return {}
+
+        self._parameter_mapper.set_custom_transform(
+            "echo", transform_echo_text_mode)
+
+    def _parse_llm_response(self, raw: Any) -> str:
+        """
+        Extract the assistant-generated text from a watsonx response.
+
+        Same logic as non-validating client.
+        """
+        if isinstance(raw, dict) and "results" in raw:
+            results = raw["results"]
+            if isinstance(results, list) and results:
+                first = results[0]
+                return first.get("generated_text", "")
+        if isinstance(raw, dict) and "choices" in raw:
+            choices = raw["choices"]
+            if isinstance(choices, list) and choices:
+                first = choices[0]
+                msg = first.get("message")
+                if isinstance(msg, dict) and "content" in msg:
+                    return msg["content"]
+                if "text" in first:
+                    return first["text"]
+        raise ValueError(f"Unexpected watsonx response format: {raw!r}")
+
+    def generate(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        *,
+        schema: SchemaType,
+        retries: int = 3,
+        generation_args: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Synchronous chat generation with validation + retries.
+
+        Args:
+            prompt: Either a string or a list of chat messages.
+            schema: JSON Schema dict, Pydantic model class, or built-in Python type.
+            retries: Maximum attempts (including the first).
+            generation_args: GenerationArgs to map to provider parameters.
+            **kwargs: Passed to the underlying ModelInference call (e.g., temperature).
+        """
+        mode = "chat"
+
+        # Normalize prompt to chat-messages
+        if isinstance(prompt, str):
+            prompt = [{"role": "user", "content": prompt}]
+
+        # Handle WatsonX params structure
+        watsonx_kwargs = {}
+
+        # Extract any existing params from kwargs
+        existing_params = kwargs.pop("params", {})
+
+        # Map generation_args to WatsonX parameters if provided
+        if generation_args and self._parameter_mapper:
+            from llmevalkit.llm.types import GenerationArgs
+
+            if isinstance(generation_args, GenerationArgs):
+                mapped_args = self._parameter_mapper.map_args(
+                    generation_args, mode)
+                # Merge mapped args with existing params
+                existing_params.update(mapped_args)
+
+        # Set params if we have any
+        if existing_params:
+            watsonx_kwargs["params"] = existing_params
+
+        # Add any other kwargs that aren't generation parameters
+        watsonx_kwargs.update(kwargs)
+
+        return super().generate(
+            **{
+                "prompt": prompt,
+                "schema": schema,
+                "retries": retries,
+                "mode": mode,
+                **self._model_kwargs,
+                **watsonx_kwargs,
+            }
+        )
+
+    async def generate_async(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        *,
+        schema: SchemaType,
+        retries: int = 3,
+        generation_args: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Asynchronous chat generation with validation + retries.
+
+        Args:
+            prompt: Either a string or a list of chat messages.
+            schema: JSON Schema dict, Pydantic model class, or built-in Python type.
+            retries: Maximum attempts.
+            generation_args: GenerationArgs to map to provider parameters.
+            **kwargs: Passed to the underlying ModelInference call.
+        """
+        mode = "chat_async"
+
+        if isinstance(prompt, str):
+            prompt = [{"role": "user", "content": prompt}]
+
+        # Handle WatsonX params structure
+        watsonx_kwargs = {}
+
+        # Extract any existing params from kwargs
+        existing_params = kwargs.pop("params", {})
+
+        # Map generation_args to WatsonX parameters if provided
+        if generation_args and self._parameter_mapper:
+            from llmevalkit.llm.types import GenerationArgs
+
+            if isinstance(generation_args, GenerationArgs):
+                mapped_args = self._parameter_mapper.map_args(
+                    generation_args, mode)
+                # Merge mapped args with existing params
+                existing_params.update(mapped_args)
+
+        # Set params if we have any
+        if existing_params:
+            watsonx_kwargs["params"] = existing_params
+
+        # Add any other kwargs that aren't generation parameters
+        watsonx_kwargs.update(kwargs)
+
+        return await super().generate_async(
+            **{
+                "prompt": prompt,
+                "schema": schema,
+                "retries": retries,
+                "mode": mode,
+                **self._model_kwargs,
+                **watsonx_kwargs,
+            }
+        )