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

ibm_watsonx_gov/providers/llmevalkit/llm/providers/litellm/rits.py

@@ -0,0 +1,224 @@
+import os
+import os
+from typing import Any, List, Optional
+from .litellm import (
+    LiteLLMClient,
+    LiteLLMClientOutputVal,
+)
+from llmevalkit.llm.base import Hook, register_llm
+from ..consts import RITS_API_KEY, RITS_API_URL, XGRAMMAR
+
+import requests
+
+
+def get_rits_model_list():
+    url = "https://rits.fmaas.res.ibm.com/ritsapi/inferenceinfo"
+    response = requests.get(url, headers={RITS_API_KEY: os.getenv(RITS_API_KEY, "")})
+    if response.status_code == 200:
+        return {m["model_name"]: m["endpoint"].split("/")[-1] for m in response.json()}
+    else:
+        raise Exception(f"Failed getting RITS model list:\n\n{response.text}")
+
+
+@register_llm("litellm.rits")
+class RITSLiteLLMClient(LiteLLMClient):
+    """
+    Specialized LiteLLMClient for RITS-hosted models.
+
+    Automatically injects:
+      - model_name = "hosted_vllm/{model_name}"
+      - api_base URL = "{api_url}/{model_url}/v1"
+      - authentication headers with RITS_API_KEY
+      - guided_decoding_backend
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: Optional[str] = None,
+        model_url: Optional[str] = None,
+        api_url: Optional[str] = None,
+        guided_decoding_backend: Optional[str] = XGRAMMAR,
+        *,
+        hooks: Optional[List[Hook]] = None,
+        **lite_kwargs: Any,
+    ) -> None:
+        """
+        Initialize the RITS LiteLLM client.
+
+        Args:
+            model_name: Name of the hosted RITS model (e.g. "my-model").
+            api_key: RITS API key (falls back to env var RITS_API_KEY).
+            model_url: URL fragment for the model; derived from model_name if omitted.
+            api_url: Base RITS API URL (defaults to None).
+            guided_decoding_backend: Backend identifier for guided decoding (defaults to XGRAMMAR).
+            hooks: Optional observability hooks to receive events.
+            lite_kwargs: Additional parameters passed to the underlying LiteLLM constructor.
+
+        Raises:
+            ValueError: If model_url derivation fails.
+            EnvironmentError: If API key is missing.
+        """
+        # Derive model_url from model_name if not provided
+        if not model_url:
+            try:
+                model_url = get_rits_model_list().get(model_name)
+            except Exception as e:
+                raise ValueError(f"Unable to derive model_url from '{model_name}': {e}")
+
+        # Obtain API key from environment if still not provided
+        if not api_key:
+            api_key = os.getenv(RITS_API_KEY)
+            if not api_key:
+                raise EnvironmentError(
+                    f"Missing API key; please set the '{RITS_API_KEY}' environment variable."
+                )
+
+        # Ensure api_url is set, if not provided, it tries to obtain it from the environment variable, otherwise it raises an error.
+        if not api_url:
+            api_url = os.getenv(RITS_API_URL)
+            if not api_url:
+                raise EnvironmentError(
+                    f"Missing API URL; please set the '{RITS_API_URL}' environment variable."
+                )
+
+        # Construct the full API base endpoint
+        api_base = f"{api_url.rstrip('/')}/{model_url}/v1"
+
+        # Call parent constructor with all required lite parameters
+        super().__init__(
+            model_name=f"hosted_vllm/{model_name}",
+            hooks=hooks,
+            api_base=api_base,
+            api_key=api_key,
+            headers={RITS_API_KEY: api_key},
+            guided_decoding_backend=guided_decoding_backend,
+            **lite_kwargs,
+        )
+
+
+@register_llm("litellm.rits.output_val")
+class RITSLiteLLMClientOutputVal(LiteLLMClientOutputVal):
+    """
+    Specialized LiteLLMClientOutputVal for RITS-hosted models.
+
+    Automatically injects:
+      - model_path = "hosted_vllm/{model_name}"
+      - api_base URL = "{api_url}/{model_url}/v1"
+      - authentication headers with RITS_API_KEY
+      - guided_decoding_backend
+
+    Inherits full JSON / Pydantic / type-based output validation,
+    retry logic, batch & async support from LiteLLMClientOutputVal.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: Optional[str] = None,
+        model_url: Optional[str] = None,
+        api_url: Optional[str] = None,
+        guided_decoding_backend: Optional[str] = XGRAMMAR,
+        *,
+        hooks: Optional[List[Hook]] = None,
+        **lite_kwargs: Any,
+    ) -> None:
+        """
+        Initialize the RITS LiteLLM client with output validation.
+
+        Args:
+            model_name: Name of the hosted RITS model (e.g. "my-model").
+            api_key: RITS API key (falls back to env var RITS_API_KEY).
+            model_url: URL fragment for the model; derived from model_name if omitted.
+            api_url: Base RITS API URL (defaults to None).
+            guided_decoding_backend: Backend identifier for guided decoding (defaults to XGRAMMAR).
+            hooks: Optional observability hooks to receive events.
+            lite_kwargs: Additional parameters passed to the underlying LiteLLM constructor.
+
+        Raises:
+            ValueError: If model_url derivation fails.
+            EnvironmentError: If API key is missing.
+        """
+        # Derive model_url from model_name if not provided
+        if not model_url:
+            try:
+                model_url = get_rits_model_list().get(model_name)
+            except Exception as e:
+                raise ValueError(f"Unable to derive model_url from '{model_name}': {e}")
+
+        # Obtain API key from environment if still not provided
+        if not api_key:
+            api_key = os.getenv(RITS_API_KEY)
+            if not api_key:
+                raise EnvironmentError(
+                    f"Missing API key; please set the '{RITS_API_KEY}' environment variable."
+                )
+
+        # Ensure api_url is set, if not provided, it tries to obtain it from the environment variable, otherwise it raises an error.
+        if not api_url:
+            api_url = os.getenv(RITS_API_URL)
+            if not api_url:
+                raise EnvironmentError(
+                    f"Missing API URL; please set the '{RITS_API_URL}' environment variable."
+                )
+
+        # Construct the full API base endpoint
+        api_base = f"{api_url.rstrip('/')}/{model_url}/v1"
+
+        # Call parent constructor with all required lite parameters
+        super().__init__(
+            model_name=f"hosted_vllm/{model_name}",
+            hooks=hooks,
+            api_base=api_base,
+            api_key=api_key,
+            headers={RITS_API_KEY: api_key},
+            guided_decoding_backend=guided_decoding_backend,
+            **lite_kwargs,
+        )
+
+    def generate(
+        self,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Synchronous chat generation with validation + retries.
+        This method is a wrapper around the generate method of the parent class,
+        ensuring that the schema_field is set to None, as RITS has problems with litellm schema validation.
+        Therefore, we disable schema validation for RITS models, and use the default validation.
+
+        Args:
+            **kwargs: Additional keyword arguments passed to the generate method.
+        Returns:
+            Any: The generated response from the model.
+        """
+
+        # Delegate to ValidatingLLMClient.generate
+        return super().generate(
+            **{
+                "schema_field": None,
+                **kwargs,
+            }
+        )
+
+    async def generate_async(
+        self,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Asynchronous chat generation with validation + retries.
+        This method is a wrapper around the generate_async method of the parent class,
+        ensuring that the schema_field is set to None, as RITS has problems with litellm schema validation.
+        Therefore, we disable schema validation for RITS models, and use the default validation.
+
+        Args:
+            **kwargs: Additional keyword arguments passed to the generate_async method.
+        Returns:
+            Any: The generated response from the model.
+        """
+
+        return await super().generate_async(
+            **{
+                "schema_field": None,
+                **kwargs,
+            }
+        )

ibm_watsonx_gov/providers/llmevalkit/llm/providers/litellm/watsonx.py

@@ -0,0 +1,60 @@
+from typing import Any, List, Optional
+from .litellm import (
+    LiteLLMClient,
+    LiteLLMClientOutputVal,
+)
+from llmevalkit.llm.base import Hook, register_llm
+
+
+@register_llm("litellm.watsonx")
+class WatsonxLiteLLMClient(LiteLLMClient):
+    """
+    Specialized LiteLLM client for Watsox models.
+
+    Automatically prefixes the model path with "watsonx/".
+    """
+
+    def __init__(
+        self, model_name: str, hooks: Optional[List[Hook]] = None, **lite_kwargs: Any
+    ) -> None:
+        """
+        Initialize a Watsonx LiteLLM client.
+
+        Args:
+            model_name: Watsonx model identifier (e.g. "gpt-j-6b").
+            hooks: Optional observability hooks (callable(event, payload)).
+            lite_kwargs: Additional keyword args passed to the underlying LiteLLM constructor.
+        """
+        # Construct the model_path for Watsonx
+        model_path = f"watsonx/{model_name}"
+
+        # Delegate to the validating LiteLLMClient
+        super().__init__(model_name=model_path, hooks=hooks, **lite_kwargs)
+
+
+@register_llm("litellm.watsonx.output_val")
+class WatsonxLiteLLMClientOutputVal(LiteLLMClientOutputVal):
+    """
+    Validating LiteLLM client for IBM Watsonx models.
+
+    Inherits all JSON/Pydantic/type-based output validation, retry logic,
+    and batch/async support from LiteLLMClientOutputVal. Automatically
+    prefixes the model path with "watsonx/".
+    """
+
+    def __init__(
+        self, model_name: str, hooks: Optional[List[Hook]] = None, **lite_kwargs: Any
+    ) -> None:
+        """
+        Initialize a Watsonx LiteLLM client with output validation.
+
+        Args:
+            model_name: Watsonx model identifier (e.g. "gpt-j-6b").
+            hooks: Optional observability hooks (callable(event, payload)).
+            lite_kwargs: Additional keyword args passed to the underlying LiteLLM constructor.
+        """
+        # Construct the model_path for Watsonx
+        model_path = f"watsonx/{model_name}"
+
+        # Delegate to the validating LiteLLMClient
+        super().__init__(model_name=model_path, hooks=hooks, **lite_kwargs)

ibm_watsonx_gov/providers/llmevalkit/llm/providers/mock_llm_client.py

@@ -0,0 +1,75 @@
+from typing import Any, Dict, List, Optional, Union
+
+from llmevalkit.llm.output_parser import ValidatingLLMClient
+from ..base import (
+    GenerationMode,
+    GenerationArgs,
+    register_llm,
+)
+
+
+@register_llm("mock.output_val")
+class MockLLMClient(ValidatingLLMClient):
+    """
+    Mock implementation of LLMClient for testing purposes.
+    Assumes that the reflection is done in runtime (i.e., the generate output does not contain actionable recommendations).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._client = None  # No real client needed for mock
+
+    @classmethod
+    def provider_class(cls):
+        # Return a dummy class for type checking
+        class DummyProvider:
+            pass
+
+        return DummyProvider
+
+    def _register_methods(self) -> None:
+        # Register mock method configs
+        self.set_method_config("text", "mock_text", "prompt")
+        self.set_method_config("chat", "mock_chat", "prompt")
+        self.set_method_config("text_async", "mock_text_async", "prompt")
+        self.set_method_config("chat_async", "mock_chat_async", "prompt")
+
+    def _setup_parameter_mapper(self) -> None:
+        # No parameter mapping needed for mock
+        pass
+
+    def _parse_llm_response(self, raw: Any) -> str:
+        # Return the raw value as the mock response
+        return str(raw)
+
+    def generate(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        mode: Union[str, GenerationMode] = GenerationMode.CHAT,
+        generation_args: Optional[GenerationArgs] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        # Return a mock response resembling expected model output
+        return {
+            "evidence": f"Assistant message: 'Logged in as user 55; auth token acquired.'",
+            "explanation": "The user_id=55 parameter is properly grounded in the conversation history, as evidenced by the assistant's explicit statement 'Logged in as user 55.'",
+            "output": 3,
+            "confidence": 0.95,
+            "correction": {},
+        }
+
+    async def generate_async(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        mode: Union[str, GenerationMode] = GenerationMode.CHAT_ASYNC,
+        generation_args: Optional[GenerationArgs] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        # Return a mock async response resembling expected model output
+        return {
+            "evidence": f"Assistant message: 'Logged in as user 55; auth token acquired.'",
+            "explanation": "The user_id=55 parameter is properly grounded in the conversation history, as evidenced by the assistant's explicit statement 'Logged in as user 55.'",
+            "output": 3,
+            "confidence": 0.95,
+            "correction": {},
+        }