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

ibm_watsonx_gov/providers/llmevalkit/llm/output_parser.py

@@ -0,0 +1,364 @@
+import re
+import json
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+)
+
+try:
+    import jsonschema
+except ImportError:
+    jsonschema = None
+
+from pydantic import (
+    BaseModel,
+    create_model,
+    Field,
+    ValidationError as PydanticValidationError,
+)
+
+from .base import LLMClient
+
+T = TypeVar("T")
+
+
+def json_schema_to_pydantic_model(
+    schema: Dict[str, Any], model_name: str = "AutoModel"
+) -> Type[BaseModel]:
+    fields = {}
+    required_fields = set(schema.get("required", []))
+    properties = schema.get("properties", {})
+
+    # Validate that all required fields exist in properties
+    missing_properties = required_fields - set(properties.keys())
+    if missing_properties:
+        raise ValueError(
+            f"Required fields {missing_properties} are not defined in properties"
+        )
+
+    # Validate that all properties are covered in required (for Azure OpenAI compliance)
+    all_property_keys = set(properties.keys())
+    schema = dict(schema)  # Make a copy
+    schema["required"] = list(all_property_keys)
+
+    type_mapping = {
+        "string": str,
+        "integer": int,
+        "number": float,
+        "boolean": bool,
+        "array": list,
+        "object": dict,
+        "null": type(None),
+    }
+
+    def parse_type(type_def: Union[str, List[str]]) -> Type[T]:
+        if isinstance(type_def, list):
+            python_types = [type_mapping.get(t, Any) for t in type_def]
+            if type(None) in python_types:
+                python_types.remove(type(None))
+                if len(python_types) == 1:
+                    return Optional[python_types[0]]
+                else:
+                    return Optional[Union[tuple(python_types)]]
+            else:
+                return Union[tuple(python_types)]
+        else:
+            return type_mapping.get(type_def, Any)
+
+    for prop_name, prop_schema in properties.items():
+        field_type = parse_type(prop_schema.get("type"))
+        default = prop_schema.get("default", None)
+        description = prop_schema.get("description", None)
+        field_args = {"description": description} if description else {}
+        fields[prop_name] = (field_type, Field(default, **field_args))
+
+    # Create the model with custom configuration
+    model = create_model(model_name, **fields)
+
+    return model
+
+
+class OutputValidationError(Exception):
+    """Raised when LLM output cannot be validated against the provided schema."""
+
+
+class ValidatingLLMClient(LLMClient, ABC):
+    """
+    An LLMClient wrapper enforcing output structure via:
+      - JSON Schema (dict),
+      - Pydantic model (BaseModel subclass),
+      - or Python built-in types (int, float, str, bool, list, dict).
+
+    Features:
+      - Injects a system-level prompt describing the required format.
+      - Cleans raw responses (strips Markdown, extracts JSON).
+      - Validates and parses the response.
+      - Retries only invalid items (single or batch) up to `retries` times.
+      - Falls back to single-item loops if no batch method is configured.
+    """
+
+    @classmethod
+    @abstractmethod
+    def provider_class(cls) -> Type:
+        """Return the underlying SDK client class, e.g. openai.OpenAI."""
+
+    @abstractmethod
+    def _register_methods(self) -> None:
+        """
+        Register MethodConfig entries:
+          self.set_method_config("text", ...),
+          self.set_method_config("chat", ...),
+          self.set_method_config("text_async", ...),
+          self.set_method_config("chat_async", ...),
+        """
+
+    def _make_instruction(
+        self, schema: Union[Dict[str, Any], Type[BaseModel], Type]
+    ) -> str:
+        """Produce a clear instruction describing exactly the required output format."""
+        if isinstance(schema, dict):
+            schema_json = json.dumps(schema, indent=2)
+            return (
+                "Please output ONLY a JSON object conforming exactly to the following JSON Schema:\n"
+                f"{schema_json}"
+            )
+        if isinstance(schema, type) and issubclass(schema, BaseModel):
+            model_schema = schema.model_json_schema()
+            return (
+                "Please output ONLY a JSON object conforming exactly to this Pydantic model schema:\n"
+                f"{model_schema}"
+            )
+        if isinstance(schema, type) and schema in (int, float, str, bool, list, dict):
+            # For simple types, no JSON wrapper required
+            return f"Please output ONLY a value of type `{schema.__name__}`."
+        raise TypeError(f"Unsupported schema type: {schema!r}")
+
+    @staticmethod
+    def _extract_json(raw: str) -> str:
+        """
+        Extract JSON from markdown fences or inline braces.
+        Falls back to returning the entire raw string.
+        """
+        # Code fence (```json ... ```)
+        fence = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", raw)
+        if fence:
+            return fence.group(1)
+        # Inline {...}
+        inline = re.search(r"(\{[\s\S]*\})", raw)
+        if inline:
+            return inline.group(1)
+        return raw
+
+    def _clean_raw(self, raw: str) -> str:
+        """Strip extraneous markdown and whitespace."""
+        cleaned = self._extract_json(raw)
+        return cleaned.strip()
+
+    def _validate(
+        self, raw: str, schema: Union[Dict[str, Any], Type[BaseModel], Type]
+    ) -> Any:
+        """
+        Clean, parse, and validate raw text against the schema/type.
+        Returns the parsed object or Pydantic instance.
+        Raises OutputValidationError on any failure.
+        """
+
+        cleaned = self._clean_raw(raw)
+        try:
+            if isinstance(schema, str):
+                data = cleaned
+            else:
+                data = json.loads(cleaned)
+        except json.JSONDecodeError:
+            try:
+                data = json.loads(cleaned.encode("unicode_escape").decode("utf-8"))
+            except Exception as e:
+                data = cleaned
+
+        # JSON Schema validation
+        if isinstance(schema, dict):
+            if jsonschema is None:
+                raise ImportError(
+                    "jsonschema is required for JSON Schema validation. Install with: pip install jsonschema"
+                )
+            try:
+                jsonschema.validate(instance=data, schema=schema)
+            except jsonschema.ValidationError as e:
+                raise OutputValidationError(
+                    f"JSON Schema validation error: {e.message}"
+                )
+            return data
+
+        # Pydantic model validation
+        if isinstance(schema, type) and issubclass(schema, BaseModel):
+            try:
+                return schema.model_validate(data)
+            except PydanticValidationError as e:
+                raise OutputValidationError(f"Pydantic validation error: {e}")
+
+        # Built-in type enforcement
+        if isinstance(schema, type) and schema in (int, float, str, bool, list, dict):
+            if not isinstance(data, schema):
+                raise OutputValidationError(
+                    f"Type mismatch: expected {schema.__name__}, got {type(data).__name__}"
+                )
+            return data
+
+        raise TypeError(f"Unsupported schema type: {schema!r}")
+
+    def _inject_system(
+        self, prompt: Union[str, List[Dict[str, Any]]], instr: str
+    ) -> Union[str, List[Dict[str, Any]]]:
+        """
+        Combine instruction and user prompt:
+        - For text: prepend the instruction.
+        - For chat messages: if first role=system, append instr to it;
+          otherwise insert a new system message.
+        """
+        if isinstance(prompt, str):
+            return f"{instr}\n\n{prompt}"
+
+        msgs = prompt.copy()
+        if msgs and msgs[0].get("role") == "system":
+            msgs[0]["content"] = msgs[0]["content"].rstrip() + "\n\n" + instr
+        else:
+            msgs.insert(0, {"role": "system", "content": instr})
+        return msgs
+
+    def generate(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        *,
+        schema: Union[Dict[str, Any], Type[BaseModel], Type],
+        schema_field: Optional[str] = None,
+        retries: Optional[int] = 3,
+        include_schema_in_system_prompt: bool = False,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Synchronous single-item generation with validation + retries.
+        """
+        current = prompt
+        instr = None
+        if include_schema_in_system_prompt:
+            instr = self._make_instruction(schema)
+            current = self._inject_system(prompt, instr)
+        if schema_field:
+            kwargs[schema_field] = schema
+            if isinstance(schema, dict):
+                new_schema = json_schema_to_pydantic_model(schema)
+                kwargs[schema_field] = new_schema
+
+        last_error: Optional[str] = None
+        for _ in range(1, retries + 1):
+            # Filter out schema-related kwargs for the base class
+            filtered_kwargs = {
+                k: v
+                for k, v in kwargs.items()
+                if k
+                not in [
+                    "schema",
+                    "schema_field",
+                    "retries",
+                    "include_schema_in_system_prompt",
+                ]
+            }
+            raw = super().generate(**{"prompt": current, **filtered_kwargs})
+            try:
+                if isinstance(raw, str):
+                    return self._validate(raw, schema)
+                return raw
+            except OutputValidationError as e:
+                last_error = str(e)
+                correction = (
+                    f"The previous response did not conform: {last_error}\nPlease correct it."
+                    " And remember to output ONLY the requested schema, without any additional text."
+                )
+                if isinstance(current, str):
+                    if instr:
+                        current = (
+                            f"{instr}\n\nPrevious output:\n{raw}\n\n"
+                            f"{correction}\n\n{prompt}"
+                        )
+                    else:
+                        current = (
+                            f"Previous output:\n{raw}\n\n" f"{correction}\n\n{prompt}"
+                        )
+                else:
+                    current = current + [
+                        {"role": "assistant", "content": raw},
+                        {"role": "user", "content": correction},
+                    ]
+        raise OutputValidationError(f"Failed after {retries} attempts: {last_error}")
+
+    async def generate_async(
+        self,
+        prompt: Union[str, List[Dict[str, Any]]],
+        *,
+        schema: Union[Dict[str, Any], Type[BaseModel], Type],
+        schema_field: Optional[str] = None,
+        retries: Optional[int] = 3,
+        include_schema_in_system_prompt: bool = False,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Asynchronous single-item generation with validation + retries.
+        """
+        current = prompt
+        instr = None
+        if include_schema_in_system_prompt:
+            instr = self._make_instruction(schema)
+            current = self._inject_system(prompt, instr)
+        if schema_field:
+            kwargs[schema_field] = schema
+            if isinstance(schema, dict):
+                new_schema = json_schema_to_pydantic_model(schema)
+                kwargs[schema_field] = new_schema
+
+        last_error: Optional[str] = None
+        for _ in range(1, retries + 1):
+            # Filter out schema-related kwargs for the base class
+            filtered_kwargs = {
+                k: v
+                for k, v in kwargs.items()
+                if k
+                not in [
+                    "schema",
+                    "schema_field",
+                    "retries",
+                    "include_schema_in_system_prompt",
+                ]
+            }
+            raw = await super().generate_async(**{"prompt": current, **filtered_kwargs})
+            try:
+                if isinstance(raw, str):
+                    return self._validate(raw, schema)
+                return raw
+            except OutputValidationError as e:
+                last_error = str(e)
+                correction = (
+                    f"The previous response did not conform: {last_error}\nPlease correct it."
+                    " And remember to output ONLY the requested schema, without any additional text."
+                )
+                if isinstance(current, str):
+                    if instr:
+                        current = (
+                            f"{instr}\n\nPrevious output:\n{raw}\n\n"
+                            f"{correction}\n\n{prompt}"
+                        )
+                    else:
+                        current = (
+                            f"Previous output:\n{raw}\n\n" f"{correction}\n\n{prompt}"
+                        )
+                else:
+                    current = current + [
+                        {"role": "assistant", "content": raw},
+                        {"role": "user", "content": correction},
+                    ]
+        raise OutputValidationError(f"Failed after {retries} attempts: {last_error}")

ibm_watsonx_gov/providers/llmevalkit/llm/providers/consts.py

@@ -0,0 +1,7 @@
+RITS_API_KEY = "RITS_API_KEY"
+RITS_API_URL = "RITS_API_URL"
+WX_URL = "WX_URL"
+WX_API_KEY = "WX_API_KEY"
+WX_PROJECT_ID = "WX_PROJECT_ID"
+WX_SPACE_ID = "WX_SPACE_ID"
+XGRAMMAR = "xgrammar"