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

ibm_watsonx_gov/providers/llmevalkit/llm/providers/wxo_ai_gateway/wxo_ai_gateway.py

@@ -0,0 +1,134 @@
+# ----------------------------------------------------------------------------------------------------
+# IBM Confidential
+# Licensed Materials - Property of IBM
+# 5737-H76, 5900-A3Q
+# © Copyright IBM Corp. 2025  All Rights Reserved.
+# US Government Users Restricted Rights - Use, duplication or disclosure restricted by
+# GSA ADPSchedule Contract with IBM Corp.
+# ----------------------------------------------------------------------------------------------------
+import os
+from typing import Any, Dict, List, Optional, Type, Union
+
+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 .wxo_ai_gateway_inference import WxoAIGatewayInference
+
+SchemaType = Union[Dict[str, Any], Type["BaseModel"], Type]
+
+
+@register_llm("wxo_ai_gateway.output_val")
+class WxoAIGatewayClientOutputVal(ValidatingLLMClient):
+
+    def __init__(self, api_key: Optional[str] = None, url: Optional[str] = None, hooks: Optional[List[Hook]] = None, **kwargs):
+        provider_kwargs = {"api_key": api_key, "url": url}
+        super().__init__(
+            client=None, client_needs_init=True, hooks=hooks, **provider_kwargs
+        )
+
+    @classmethod
+    def provider_class(cls) -> Type:
+        """
+        Underlying client class: WxoAIGatewayInference.
+        """
+        return WxoAIGatewayInference
+
+    def _register_methods(self) -> None:
+        """
+        Register how to call wxo ai gateway methods for validation:
+          - 'chat'       -> ModelInference.chat
+          - 'chat_async' -> ModelInference.achat
+        """
+        self.set_method_config(GenerationMode.CHAT.value, "chat", "messages")
+        self.set_method_config(
+            GenerationMode.CHAT_ASYNC.value, "achat", "messages")
+
+    def _parse_llm_response(self, raw: Any) -> str:
+        """
+        Extract the assistant-generated text from a wxo ai gateway 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(raw.get("message", "Invalid response format"))
+
+    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}]
+
+        return super().generate(
+            **{
+                "prompt": prompt,
+                "schema": schema,
+                "retries": retries,
+                "mode": mode
+            }
+        )
+
+    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}]
+
+        return await super().generate_async(
+            **{
+                "prompt": prompt,
+                "schema": schema,
+                "retries": retries,
+                "mode": mode
+            }
+        )

ibm_watsonx_gov/providers/llmevalkit/llm/providers/wxo_ai_gateway/wxo_ai_gateway_inference.py

@@ -0,0 +1,214 @@
+# ----------------------------------------------------------------------------------------------------
+# IBM Confidential
+# Licensed Materials - Property of IBM
+# 5737-H76, 5900-A3Q
+# © Copyright IBM Corp. 2025  All Rights Reserved.
+# US Government Users Restricted Rights - Use, duplication or disclosure restricted by
+# GSA ADPSchedule Contract with IBM Corp.
+# ----------------------------------------------------------------------------------------------------
+
+import asyncio
+import json
+import logging
+import time
+import weakref
+from typing import Dict
+
+import aiohttp
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+class WxoAIGatewayInference():
+    RETRY_AFTER_STATUS_CODES = [502, 503, 504]
+    RETRY_COUNT = 3
+    BACK_OFF_FACTOR = 1
+
+    def __init__(self, api_key: str, url: str) -> None:
+        self.api_key = api_key
+        self.url = url
+        self.session: aiohttp.ClientSession = None
+        self._loop: asyncio.AbstractEventLoop = None
+        weakref.finalize(self, self._finalize_cleanup, self)
+
+    async def get_session(self) -> aiohttp.ClientSession:
+        """Get or create an aiohttp session for the current event loop."""
+        try:
+            current_loop = asyncio.get_running_loop()
+        except Exception:
+            raise logger.debug("get_session must be called from within an async context")
+
+        # Check if session exists and is valid for current loop
+        if (self.session is None or
+                self.session.closed or
+                self._loop is not current_loop):
+
+            # Close old session if it exists and is from a different loop
+            if self.session and not self.session.closed:
+                try:
+                    await self.session.close()
+                except Exception as e:
+                    logger.debug(f"Error closing old session: {e}")
+
+            # Create new session for current loop
+            self.session = aiohttp.ClientSession()
+            self._loop = current_loop
+            logger.debug(f"Created new session for event loop {id(current_loop)}")
+
+        return self.session
+
+    @staticmethod
+    def _finalize_cleanup(self):
+        """
+        Called automatically when the object is garbage-collected.
+        This function is synchronous, so it schedules async cleanup properly.
+        """
+        if self.session and not self.session.closed:
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    loop.create_task(self.session.close())
+                else:
+                    loop.run_until_complete(self.session.close())
+            except Exception as e:
+                logger.info(f"Error during cleanup: {e}")
+
+    def chat(self, messages: list[Dict], **kwargs) -> dict:
+        """
+        Sync chat method.
+        Args:
+            messages (list[Dict]): List of messages to send.
+
+        Returns:
+            dict: Response from the API.
+        """
+        payload_data = {
+            "messages": messages
+        }
+        response = self._post(payload_data)
+        return response
+
+    async def achat(self, messages: list[Dict], **kwargs) -> dict:
+        """
+        async implementation of the chat method.
+
+        Args:
+            messages (list[Dict]): List of messages to send.
+
+        Returns:
+            dict: Response from the API.
+        """
+        payload_data = {
+            "messages": messages
+        }
+        response = await self._apost(payload_data)
+        return response
+
+    def _get_headers(self) -> Dict[str, str]:
+        return {
+            "Content-Type": "application/json",
+            "accept": "application/json",
+            "IAM-API_KEY": self.api_key
+        }
+
+    def _post(self, payload_data: dict) -> dict:
+        """
+        Make a synchronous POST request to the Wxo AI Gateway API.
+
+        Args:
+            payload_data: Dictionary containing the request payload
+
+        Returns:
+            dict: JSON response from the API
+
+        Raises:
+            requests.HTTPError: If the request fails after all retries
+            json.JSONDecodeError: If response cannot be parsed as JSON
+        """
+        headers = self._get_headers()
+
+        for attempt in range(self.RETRY_COUNT):
+            try:
+                response = requests.post(
+                    url=self.url,
+                    json=payload_data,
+                    headers=headers,
+                    timeout=30.0,
+                    verify=False
+                )
+
+                response_status = response.status_code
+
+                if response_status == 200:
+                    return response.json()
+
+                if response_status in self.RETRY_AFTER_STATUS_CODES and attempt < self.RETRY_COUNT - 1:
+                    backoff_time = self.BACK_OFF_FACTOR * (2 ** attempt)
+                    logger.info(
+                        f"Received status {response_status}, retrying in {backoff_time}s (attempt {attempt + 1}/{self.RETRY_COUNT})...")
+                    time.sleep(backoff_time)
+                    continue
+
+                response.raise_for_status()
+                return response.json()
+
+            except requests.RequestException as e:
+                if attempt < self.RETRY_COUNT - 1:
+                    backoff_time = self.BACK_OFF_FACTOR * (2 ** attempt)
+                    logger.info(
+                        f"Request failed with {type(e).__name__}, retrying in {backoff_time}s (attempt {attempt + 1}/{self.RETRY_COUNT})...")
+                    time.sleep(backoff_time)
+                    continue
+                else:
+                    raise
+
+    async def _apost(self, payload_data: dict) -> dict:
+        """
+        Make an asynchronous POST request to the Wxo AI Gateway API.
+
+        Args:
+            payload_data: Dictionary containing the request payload
+
+        Returns:
+            dict: JSON response from the API
+
+        Raises:
+            aiohttp.ClientError: If the request fails after all retries
+            json.JSONDecodeError: If response cannot be parsed as JSON
+        """
+        headers = self._get_headers()
+        session = await self.get_session()
+        for attempt in range(self.RETRY_COUNT):
+            try:
+                async with session.post(
+                    url=self.url,
+                    json=payload_data,
+                    headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=30.0),
+                    ssl=False
+                ) as response:
+                    response_status = response.status
+
+                    if response_status == 200:
+                        return await response.json()
+
+                    if response_status in self.RETRY_AFTER_STATUS_CODES and attempt < self.RETRY_COUNT - 1:
+                        backoff_time = self.BACK_OFF_FACTOR * \
+                            (2 ** attempt)
+                        logger.info(
+                            f"Received status {response_status}, retrying in {backoff_time}s (attempt {attempt + 1}/{self.RETRY_COUNT})...")
+                        await asyncio.sleep(backoff_time)
+                        continue
+
+                    return await response.json()
+
+            except aiohttp.ClientError as e:
+                if attempt < self.RETRY_COUNT - 1:
+                    backoff_time = self.BACK_OFF_FACTOR * (2 ** attempt)
+                    logger.info(
+                        f"Request failed with {type(e).__name__}, retrying in {backoff_time}s (attempt {attempt + 1}/{self.RETRY_COUNT})...")
+                    await asyncio.sleep(backoff_time)
+                    continue
+                else:
+                    raise

ibm_watsonx_gov/providers/llmevalkit/llm/types.py

@@ -0,0 +1,136 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+from dataclasses import dataclass
+
+
+class GenerationMode(Enum):
+    """
+    Enum for different generation modes across LLM providers.
+    """
+
+    TEXT = "text"
+    CHAT = "chat"
+    TEXT_ASYNC = "text_async"
+    CHAT_ASYNC = "chat_async"
+
+
+class LLMResponse:
+    """Response object that can contain both content and tool calls"""
+
+    def __init__(self, content: str, tool_calls: Optional[List[Dict[str, Any]]] = None):
+        self.content = content
+        self.tool_calls = tool_calls or []
+
+    def __str__(self) -> str:
+        """Return the content of the response as a string."""
+        return self.content
+
+    def __repr__(self) -> str:
+        """Return a string representation of the LLMResponse object."""
+        return f"LLMResponse(content='{self.content}', tool_calls={self.tool_calls})"
+
+
+@dataclass
+class GenerationArgs:
+    """
+    Provider-agnostic generation arguments.
+
+    These arguments represent common parameters across LLM providers.
+    Each provider should implement mapping from these generic arguments
+    to their specific parameter names.
+    """
+
+    # Core generation parameters
+    max_tokens: Optional[int] = None
+    min_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    top_k: Optional[int] = None
+
+    # Penalties and biases
+    frequency_penalty: Optional[float] = None
+    presence_penalty: Optional[float] = None
+    repetition_penalty: Optional[float] = None
+
+    # Stop conditions
+    stop_sequences: Optional[List[str]] = None
+
+    # Randomness control
+    seed: Optional[int] = None
+
+    # Generation control
+    decoding_method: Optional[str] = None  # "greedy" or "sample"
+
+    # Output control
+    stream: Optional[bool] = None
+    echo: Optional[bool] = None
+    logprobs: Optional[bool] = None
+    top_logprobs: Optional[int] = None
+
+    # Other common parameters
+    timeout: Optional[float] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary, excluding None values."""
+        return {k: v for k, v in self.__dict__.items() if v is not None}
+
+
+class ParameterMapper:
+    """
+    Abstract base class for mapping generic generation arguments to provider-specific parameters.
+    """
+
+    def __init__(self):
+        self._text_mappings: Dict[str, str] = {}
+        self._chat_mappings: Dict[str, str] = {}
+        self._custom_transforms: Dict[str, callable] = {}
+
+    def set_text_mapping(self, generic_param: str, provider_param: str) -> None:
+        """Set parameter mapping for text generation mode."""
+        self._text_mappings[generic_param] = provider_param
+
+    def set_chat_mapping(self, generic_param: str, provider_param: str) -> None:
+        """Set parameter mapping for chat generation mode."""
+        self._chat_mappings[generic_param] = provider_param
+
+    def set_custom_transform(
+        self, generic_param: str, transform_func: callable
+    ) -> None:
+        """Set a custom transformation function for a parameter."""
+        self._custom_transforms[generic_param] = transform_func
+
+    def map_args(self, args: GenerationArgs, mode: str) -> Dict[str, Any]:
+        """
+        Map generic arguments to provider-specific parameters.
+
+        Args:
+            args: Generic generation arguments
+            mode: Generation mode ('text', 'chat', 'text_async', 'chat_async')
+
+        Returns:
+            Dictionary of provider-specific parameters
+        """
+        # Determine which mapping to use based on mode
+        is_chat_mode = mode in ["chat", "chat_async"]
+        mappings = self._chat_mappings if is_chat_mode else self._text_mappings
+
+        provider_args = {}
+        args_dict = args.to_dict()
+
+        for generic_param, value in args_dict.items():
+            # Check for custom transform first
+            if generic_param in self._custom_transforms:
+                transformed = self._custom_transforms[generic_param](value, mode)
+                if isinstance(transformed, dict):
+                    provider_args.update(transformed)
+                else:
+                    # If transform returns a single value, use the generic param name
+                    provider_args[generic_param] = transformed
+            # Use direct mapping if available
+            elif generic_param in mappings:
+                provider_args[mappings[generic_param]] = value
+            # Fall back to generic parameter name
+            else:
+                provider_args[generic_param] = value
+
+        return provider_args

ibm_watsonx_gov/providers/llmevalkit/metrics/__init__.py

@@ -0,0 +1,4 @@
+from .prompt import RelevancePrompt
+from .metric import StandardMetric, Metric
+from .prompt import MetricPrompt
+from .metrics_runner import MetricRunner, MetricRunResult