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

ibm_watsonx_gov/providers/llmevalkit/llm/base.py

@@ -0,0 +1,426 @@
+import inspect
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union
+
+# Import from parent module to ensure singleton registry
+from llmevalkit.llm import _REGISTRY
+from .types import GenerationMode, GenerationArgs, ParameterMapper
+
+T = TypeVar("T", bound="LLMClient")
+Hook = Callable[[str, Dict[str, Any]], None]
+
+
+def register_llm(name: str) -> Callable[[Type[T]], Type[T]]:
+    """
+    Register an LLM client class with the global registry.
+
+    Args:
+        name: Unique identifier for the client
+
+    Returns:
+        Decorator function that registers the class
+    """
+
+    def deco(cls: Type["LLMClient"]) -> Type["LLMClient"]:
+        _REGISTRY[name] = cls
+        return cls
+
+    return deco
+
+
+def get_llm(name: str) -> Type["LLMClient"]:
+    """
+    Retrieve an LLM client class from the global registry.
+
+    Args:
+        name: Identifier for the client
+
+    Returns:
+        LLMClient class
+
+    Raises:
+        ValueError: If no client is registered under the given name
+    """
+    if not _REGISTRY:
+        # Import providers to populate registry
+        from llmevalkit.llm import _import_providers
+
+        _import_providers()
+
+    try:
+        return _REGISTRY[name]
+    except KeyError:
+        available = list(_REGISTRY.keys())
+        raise ValueError(
+            f"No LLMClient registered under '{name}'. "
+            f"Available clients: {available}"
+        )
+
+
+def list_available_llms() -> List[str]:
+    """
+    List all registered LLM client names.
+
+    Returns:
+        List of registered LLM client names
+    """
+    if not _REGISTRY:
+        # Import providers to populate registry
+        from llmevalkit.llm import _import_providers
+
+        _import_providers()
+
+    return list(_REGISTRY.keys())
+
+
+class MethodConfig:
+    """
+    Configuration for a provider method.
+
+    Attributes:
+        path: Dot-delimited attribute path on the client (e.g. "chat.completions.create").
+        prompt_arg: Name of the parameter used for the prompt/messages.
+    """
+
+    def __init__(self, path: str, prompt_arg: str) -> None:
+        self.path = path
+        self.prompt_arg = prompt_arg
+
+    def resolve(self, client: Any) -> Callable[..., Any]:
+        """
+        Traverse `path` on `client` to retrieve the bound callable.
+
+        Raises:
+            AttributeError: if any attribute in the path is missing.
+            TypeError: if the resolved attribute is not callable.
+        """
+        obj: Any = client
+        for attr in self.path.split("."):
+            obj = getattr(obj, attr, None)
+            if obj is None:
+                raise AttributeError(
+                    f"Could not resolve method path '{self.path}' on {client}"
+                )
+        if not callable(obj):
+            raise TypeError(f"Resolved '{self.path}' is not callable on {client}")
+        return obj
+
+
+class LLMClient(ABC):
+    """
+    Abstract base wrapper for any LLM provider.
+
+    Responsibilities:
+      - Accept an existing SDK client or construct one from kwargs.
+      - Register provider methods via MethodConfig.
+      - Provide sync/async/batch generate calls.
+      - Emit observability hooks.
+      - Parse raw responses into plain text.
+    """
+
+    def __init__(
+        self,
+        *,
+        client: Optional[Any] = None,
+        client_needs_init: bool = False,
+        hooks: Optional[List[Hook]] = None,
+        **provider_kwargs: Any,
+    ) -> None:
+        """
+        Initialize the wrapper.
+
+        Args:
+            client: Pre-initialized provider SDK instance.
+            client_needs_init: If True, client is not initialized and will be
+                initialized with provider_kwargs.
+            hooks: Callables(event_name, payload) for observability.
+            provider_kwargs: Passed to provider_class constructor if client is None.
+
+        Raises:
+            TypeError: if `client` is provided but is not instance of provider_class.
+            RuntimeError: if provider_class instantiation fails.
+        """
+        self._hooks: List[Hook] = hooks or []
+        self._method_configs: Dict[str, MethodConfig] = {}
+        self._parameter_mapper: Optional[ParameterMapper] = None
+
+        self._other_kwargs = provider_kwargs
+        if client is not None:
+            if not isinstance(client, self.provider_class()):
+                raise TypeError(
+                    f"Expected client of type {self.provider_class().__name__}, "
+                    f"got {type(client).__name__}"
+                )
+            self._client = client
+        else:
+            if client_needs_init:
+                sig = inspect.signature(self.provider_class().__init__)
+                init_kwargs = {
+                    k: v
+                    for k, v in provider_kwargs.items()
+                    if k in sig.parameters and k != "self"
+                }
+                self._other_kwargs = {
+                    k: v
+                    for k, v in provider_kwargs.items()
+                    if k not in sig.parameters and k != "self"
+                }
+                try:
+                    self._client = self.provider_class()(**init_kwargs)
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to initialize {self.provider_class().__name__}: {e}"
+                    ) from e
+            else:
+                self._client = self.provider_class()
+
+        self._register_methods()
+        self._setup_parameter_mapper()
+
+    @classmethod
+    @abstractmethod
+    def provider_class(cls) -> Type:
+        """
+        Underlying SDK client class, e.g. openai.OpenAI or litellm.LiteLLM.
+        """
+
+    @abstractmethod
+    def _register_methods(self) -> None:
+        """
+        Subclasses register MethodConfig entries by calling:
+            self.set_method_config(key, path, prompt_arg)
+        for keys: 'text', 'chat', 'text_async', 'chat_async'.
+        """
+    
+    def _setup_parameter_mapper(self) -> None:
+        """
+        Setup parameter mapping for the provider. Override in subclasses to configure
+        mapping from generic GenerationArgs to provider-specific parameters.
+        """
+        pass
+
+    def set_method_config(self, key: str, path: str, prompt_arg: str) -> None:
+        """
+        Register how to invoke a provider method.
+
+        Args:
+            key: Identifier ('text', 'chat', 'text_async', 'chat_async').
+            path: Dot-separated path on the SDK client.
+            prompt_arg: Name of the argument carrying the prompt/messages.
+        """
+        self._method_configs[key] = MethodConfig(path, prompt_arg)
+
+    def get_method_config(self, key: str) -> MethodConfig:
+        """
+        Retrieve a previously registered MethodConfig.
+
+        Raises:
+            KeyError: if no config exists for `key`.
+        """
+        try:
+            return self._method_configs[key]
+        except KeyError:
+            raise KeyError(f"No method config registered under '{key}'")
+
+    def get_client(self) -> Any:
+        """Return the raw underlying SDK client."""
+        return self._client
+
+    def _emit(self, event: str, payload: Dict[str, Any]) -> None:
+        """Invoke all observability hooks, swallowing errors."""
+        for hook in self._hooks:
+            try:
+                hook(event, payload)
+            except Exception:
+                pass
+            
+    def _filter_sensitive_args(self, args: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Filter out sensitive information from call arguments before logging.
+
+        Args:
+            args: Dictionary of call arguments that may contain sensitive data
+
+        Returns:
+            Filtered dictionary with sensitive keys masked or removed
+        """
+        # Define sensitive keys that should be filtered out
+        sensitive_keys = {
+            # RITS-specific sensitive keys
+            "api_key",
+            "headers",
+            # IBM WatsonX AI sensitive keys
+            "project_id",
+            "credentials",
+            # General sensitive keys
+            "authorization",
+            "secret",
+            "password",
+            "key",
+        }
+
+        filtered_args = {}
+        for key, value in args.items():
+            if key.lower() in sensitive_keys or any(
+                sensitive in key.lower()
+                for sensitive in ["key", "secret", "auth", "credential"]
+            ):
+                filtered_args[key] = "SECRET"
+            elif isinstance(value, dict) and key == "headers":
+                # Special handling for headers dict to filter out sensitive header values
+                filtered_headers = {}
+                for header_key, header_value in value.items():
+                    if any(
+                        sensitive in header_key.lower()
+                        for sensitive in ["key", "auth", "secret"]
+                    ):
+                        filtered_headers[header_key] = "SECRET"
+                    else:
+                        filtered_headers[header_key] = header_value
+                filtered_args[key] = filtered_headers
+            else:
+                filtered_args[key] = value
+
+        return filtered_args
+
+    @abstractmethod
+    def _parse_llm_response(self, raw: Any) -> str:
+        """
+        Extract the generated text from a single raw response.
+
+        Raises:
+            ValueError: if extraction fails.
+        """
+
+    def generate(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        mode: Union[str, GenerationMode] = GenerationMode.CHAT,
+        generation_args: Optional[GenerationArgs] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Synchronous generation.
+
+        Args:
+            prompt: Either a plain string or a list of chat messages dicts.
+            mode: Generation mode (text, chat, text_async, chat_async).
+            generation_args: Provider-agnostic generation arguments.
+            **kwargs: Additional provider-specific parameters.
+
+        Returns:
+            The generated text.
+
+        Raises:
+            KeyError: if no MethodConfig for `mode`.
+            Exception: if the underlying call or parsing fails.
+        """
+        # Convert enum to string if needed
+        mode_str = mode.value if isinstance(mode, GenerationMode) else mode
+
+        # Convert single string prompt to chat format if in chat mode
+        if mode is GenerationMode.CHAT and isinstance(prompt, str):
+            prompt = [{"role": "user", "content": prompt}]
+
+        cfg = self.get_method_config(mode_str)
+        fn = cfg.resolve(self._client)
+
+        call_args = {cfg.prompt_arg: prompt, **kwargs}
+
+        # Map generic arguments to provider-specific parameters
+        if generation_args and self._parameter_mapper:
+            mapped_args = self._parameter_mapper.map_args(generation_args, mode_str)
+            # Provider-specific kwargs take precedence over mapped args
+            for k, v in mapped_args.items():
+                if k not in call_args:
+                    call_args[k] = v
+
+        sig = inspect.signature(fn)
+
+        for k, v in self._other_kwargs.items():
+            if k not in call_args and k in sig.parameters:
+                call_args[k] = v
+
+        # Filter sensitive arguments before logging
+        filtered_args = self._filter_sensitive_args(call_args)
+        self._emit("before_generate", {"mode": mode_str, "args": filtered_args})
+        try:
+            raw = fn(**call_args)
+        except Exception as e:
+            self._emit("error", {"phase": "generate", "error": str(e)})
+            raise
+        text = self._parse_llm_response(raw)
+        self._emit("after_generate", {"mode": mode_str, "response": text})
+        return text
+
+    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,
+    ) -> str:
+        """
+        Asynchronous generation.
+
+        Uses provider async method if registered, otherwise falls back to thread.
+
+        Args:
+            prompt: string or messages list.
+            mode: Generation mode (text_async, chat_async).
+            generation_args: Provider-agnostic generation arguments.
+            **kwargs: Additional provider-specific parameters.
+
+        Returns:
+            The generated text.
+
+        Raises:
+            Exception: if generation or parsing fails.
+        """
+        # Convert enum to string if needed
+        mode_str = mode.value if isinstance(mode, GenerationMode) else mode
+
+        # Convert single string prompt to chat format if in chat mode
+        if mode is GenerationMode.CHAT_ASYNC and isinstance(prompt, str):
+            prompt = [{"role": "user", "content": prompt}]
+
+        if mode_str in self._method_configs:
+            cfg = self.get_method_config(mode_str)
+            fn = cfg.resolve(self._client)
+
+            call_args = {cfg.prompt_arg: prompt, **kwargs}
+
+            # Map generic arguments to provider-specific parameters
+            if generation_args and self._parameter_mapper:
+                mapped_args = self._parameter_mapper.map_args(generation_args, mode_str)
+                # Provider-specific kwargs take precedence over mapped args
+                for k, v in mapped_args.items():
+                    if k not in call_args:
+                        call_args[k] = v
+
+            sig = inspect.signature(fn)
+
+            for k, v in self._other_kwargs.items():
+                if k not in call_args and k in sig.parameters:
+                    call_args[k] = v
+
+            # Filter sensitive arguments before logging
+            filtered_args = self._filter_sensitive_args(call_args)
+            self._emit(
+                "before_generate_async", {"mode": mode_str, "args": filtered_args}
+            )
+            try:
+                raw = await fn(**call_args)
+            except Exception as e:
+                self._emit("error", {"phase": "generate_async", "error": str(e)})
+                raise
+            text = self._parse_llm_response(raw)
+            self._emit("after_generate_async", {"mode": mode_str, "response": text})
+            return text
+
+        kwargs["mode"] = mode_str.replace("_async", "")
+        kwargs["generation_args"] = generation_args
+
+        # fallback to sync generate in thread
+        return await asyncio.to_thread(self.generate, prompt, **kwargs)