adaptive-harmony 0.1.23__py3-none-any.whl

adaptive_harmony/core/structured_output.py

@@ -0,0 +1,385 @@
+import json
+import re
+from enum import Enum
+from typing import Literal, Type, Union, get_args, get_origin
+
+from pydantic import BaseModel, ValidationError
+
+from adaptive_harmony import InferenceModel, StringThread
+from adaptive_harmony.core.reasoning import remove_reasoning
+
+FIX_OUTPUT_FORMAT = """Below, the COMPLETION did not satisfy the constraints given in the PROMPT. Please rewrite the completion to comply with constraints, nothing else.
+
+PROMPT
+The output should be a well-formatted JSON instance that conforms to the JSON schema below. All fields are required. Do not output anything else other than the JSON.
+
+As an example, for the schema
+{{
+    "foo": {{
+        "items":{{"type": "string"}},
+        "type": "array"
+    }},
+    "bar": {{"type": "integer"}}
+}}
+the object {{"foo": ["hey", "bye"], "bar": 1}} is a well-formatted instance of the schema. The object {{"properties": {{"foo": ["hey", "bye"], "bar":"1" }}}} is not well-formatted.
+
+Here is the output JSON schema:
+{json_schema}
+
+COMPLETION
+{completion}
+"""
+
+
+class JsonParseError(Exception):
+    def __init__(self, message: str, completion: str):
+        super().__init__(message)
+        self.completion = completion
+
+
+def get_pydantic_schema(base_model: Type[BaseModel]) -> str:
+    schema = base_model.model_json_schema()
+    for prop in schema.get("properties", {}).values():
+        prop.pop("title", None)
+    return json.dumps(schema, indent=2)
+
+
+class OutputParserException(Exception):
+    """Exception raised for parsing errors."""
+
+    def __init__(self, message: str, llm_output: str | None = None):
+        super().__init__(message)
+        self.llm_output = llm_output
+
+
+def pydantic_parse[T: BaseModel](text: str, pydantic_object: type[T]) -> T:
+    """Parse the output of an LLM call to a pydantic object.
+
+    Args:
+        text: The output of the LLM call.
+        pydantic_object: The pydantic model to parse into.
+
+    Returns:
+        The parsed pydantic object.
+    """
+    # Remove Qwen3 reasoning
+    text = remove_reasoning(text)
+
+    # Strip initial whitespace
+    text = text.strip()
+
+    def parse_json_with_completion(json_text):
+        """Parse JSON, handling partial JSON by completing missing brackets."""
+        # Strip whitespace and backticks
+        json_text = json_text.strip(" \n\r\t`")
+
+        # Handle action_input special case - escape special chars
+        if '"action_input"' in json_text:
+
+            def fix_action_input(match):
+                value = match.group(2)
+                value = re.sub(r"\n", r"\\n", value)
+                value = re.sub(r"\r", r"\\r", value)
+                value = re.sub(r"\t", r"\\t", value)
+                value = re.sub(r'(?<!\\)"', r"\"", value)
+                return match.group(1) + value + match.group(3)
+
+            json_text = re.sub(r'("action_input"\:\s*")(.*?)(")', fix_action_input, json_text, flags=re.DOTALL)
+
+        # NOTE Axel: gemma likes to escape the left bracket, patching for now
+        json_text = json_text.replace(r"\\[", "[")
+
+        # Try parsing as-is first
+        try:
+            return json.loads(json_text)
+        except json.JSONDecodeError:
+            pass
+
+        # Handle partial JSON - complete missing brackets and quotes
+        chars = list(json_text)
+        stack = []
+        in_string = False
+        escaped = False
+
+        for i, char in enumerate(chars):
+            if in_string:
+                if char == '"' and not escaped:
+                    in_string = False
+                elif char == "\n" and not escaped:
+                    chars[i] = "\\n"
+                escaped = char == "\\" and not escaped
+            elif char == '"':
+                in_string = True
+                escaped = False
+            elif char == "{":
+                stack.append("}")
+            elif char == "[":
+                stack.append("]")
+            elif char in {"}", "]"}:
+                if stack and stack[-1] == char:
+                    stack.pop()
+
+        # Close unterminated string
+        if in_string:
+            if escaped and chars:
+                chars.pop()
+            chars.append('"')
+
+        # Add missing closing brackets
+        chars.extend(reversed(stack))
+
+        # Try parsing with progressively fewer characters
+        while chars:
+            try:
+                return json.loads("".join(chars))
+            except json.JSONDecodeError:
+                chars.pop()
+
+        # If nothing worked, raise with original
+        raise json.JSONDecodeError("Invalid JSON", json_text, 0)
+
+    # Try parsing the original text first
+    try:
+        json_object = parse_json_with_completion(text)
+    except json.JSONDecodeError:
+        # Try extracting from markdown blocks
+        markdown_match = re.search(r"```(json)(.*?)```", text, re.DOTALL)
+        if not markdown_match:
+            markdown_match = re.search(r"```(json)?(.*)", text, re.DOTALL)
+        xml_match = re.search(r"<json>(.*?)</json>", text, re.DOTALL)
+        if not xml_match:
+            xml_match = re.search(r"<json>(.*)", text, re.DOTALL)
+
+        if markdown_match or xml_match:
+            try:
+                json_object = parse_json_with_completion(
+                    markdown_match.group(2) if markdown_match else (xml_match.group(1) if xml_match else "")
+                )
+            except json.JSONDecodeError:
+                msg = f"Invalid json output: {text}"
+                raise OutputParserException(msg, llm_output=text)
+        else:
+            msg = f"Invalid json output: {text}"
+            raise OutputParserException(msg, llm_output=text)
+
+    try:
+        return pydantic_object.model_validate(json_object)
+    except ValidationError as e:
+        json_string = json.dumps(json_object)
+        msg = f"Failed to parse {pydantic_object.__name__} from completion {json_string}. Got: {e}"
+        raise OutputParserException(msg, llm_output=json_string) from e
+
+
+async def generate_and_validate[T: BaseModel](
+    model: InferenceModel,
+    thread: StringThread,
+    pydantic_model: Type[T],
+    max_parsing_retries: int = 1,
+) -> tuple[str, T]:
+    """
+    Generates with InferenceModel, validates completion against Pydantic model and retries
+    if validation fails. It's recommended you use a StructuredJSONOutputBaseModel as
+    the pydantic_object to clean up the JSON schema for the LLM. Does not support RootModel.
+    """
+
+    json_schema = get_pydantic_schema(pydantic_model)
+
+    response_thread = await model.generate(thread)
+    completion = response_thread.last_content()
+
+    current_retries = 0
+    while current_retries <= max_parsing_retries:
+        try:
+            parsed = pydantic_parse(completion, pydantic_model)
+            return (completion, parsed)
+        except Exception:
+            if current_retries == max_parsing_retries:
+                break
+
+            # Create repair prompt
+            repair_thread = StringThread(
+                [("user", FIX_OUTPUT_FORMAT.format(json_schema=json_schema, completion=completion))]
+            )
+            response_thread = await model.generate(repair_thread)
+            completion = response_thread.last_content()
+            current_retries += 1
+
+    raise JsonParseError(f"Could not parse json output after {max_parsing_retries} retries", completion)
+
+
+def _get_simplified_type(field_type):
+    origin = get_origin(field_type)
+    args = get_args(field_type)
+
+    if origin is list:
+        if args:
+            return [_get_simplified_type(args[0])]
+        else:
+            return "array"
+    elif origin is dict:
+        if len(args) == 2:
+            key_type = _get_simplified_type(args[0])
+            value_type = _get_simplified_type(args[1])
+            return f"Dict[{key_type}, {value_type}]"
+        else:
+            return "dict"
+    elif origin is tuple:
+        if args:
+            if len(args) == 2 and args[1] is ...:
+                # Variable length tuple like Tuple[str, ...]
+                element_type = _get_simplified_type(args[0])
+                return f"Tuple[{element_type}, ...]"
+            else:
+                # Fixed length tuple like Tuple[str, int]
+                element_types = [_get_simplified_type(arg) for arg in args]
+                return f"Tuple[{', '.join(element_types)}]"
+        else:
+            return "tuple"
+    elif origin is set:
+        if args:
+            element_type = _get_simplified_type(args[0])
+            return f"Set[{element_type}]"
+        else:
+            return "set"
+    elif origin is type(None):
+        return "null"
+    elif origin is Literal:
+        # Handle Literal types by showing them as Literal["value1", "value2"]
+        # Use double quotes for strings to match JSON format and prevent LLM confusion
+        literal_values = []
+        for arg in args:
+            if isinstance(arg, str):
+                # Use double quotes for strings to match JSON format
+                literal_values.append(f'"{arg}"')
+            else:
+                # Use repr() for non-strings (numbers, booleans, etc.)
+                literal_values.append(repr(arg))
+        return f"Literal[{', '.join(literal_values)}]"
+    elif origin is Union:
+        # Handle Union types by showing all possible types
+        if len(args) == 2 and type(None) in args:
+            # This is Optional[T] which is Union[T, None]
+            non_none_type = [arg for arg in args if arg is not type(None)][0]
+            simplified_type = _get_simplified_type(non_none_type)
+            # Convert to string representation if needed
+            if isinstance(simplified_type, (list, dict)):
+                simplified_type = str(simplified_type).replace("'", '"')
+            return f"Optional[{simplified_type}]"
+        else:
+            # Regular Union with multiple types
+            union_types = []
+            for arg in args:
+                simplified = _get_simplified_type(arg)
+                # Convert to string representation if needed
+                if isinstance(simplified, (list, dict)):
+                    union_types.append(str(simplified).replace("'", '"'))
+                else:
+                    union_types.append(str(simplified))
+            return f"Union[{', '.join(union_types)}]"
+    elif origin is not None:
+        return str(origin.__name__) if origin else str(field_type.__name__)
+    elif hasattr(field_type, "__bases__") and issubclass(field_type, BaseModel):
+        return get_simple_pydantic_schema(field_type)
+    elif hasattr(field_type, "__bases__") and issubclass(field_type, Enum):
+        # Handle Enum types by showing possible values
+        enum_values = [f'"{value.value}"' for value in field_type]
+        return f"Enum[{', '.join(enum_values)}]"
+    elif field_type is str:
+        return "str"
+    elif field_type is int:
+        return "int"
+    elif field_type is float:
+        return "float"
+    elif field_type is bool:
+        return "bool"
+    else:
+        return str(field_type.__name__)
+
+
+def get_simple_pydantic_schema(model: type[BaseModel]):
+    representation = {}
+    for field_name, field in model.model_fields.items():
+        representation[field_name] = _get_simplified_type(field.annotation)
+    return representation
+
+
+def _format_schema_value(value, indent=0):
+    """Format a schema value for display, handling nested structures."""
+    indent_str = "  " * indent
+    if isinstance(value, dict):
+        if not value:
+            return "{}"
+        lines = ["{"]
+        for k, v in value.items():
+            formatted_value = _format_schema_value(v, indent + 1)
+            lines.append(f'  {indent_str}"{k}": {formatted_value},')
+        # Remove trailing comma from last item
+        if lines[-1].endswith(","):
+            lines[-1] = lines[-1][:-1]
+        lines.append(f"{indent_str}" + "}")
+        return "\n".join(lines)
+    elif isinstance(value, list):
+        if not value:
+            return "[]"
+        elif len(value) == 1:
+            formatted_item = _format_schema_value(value[0], indent)
+            return f"[{formatted_item}]"
+        else:
+            lines = ["["]
+            for item in value:
+                formatted_item = _format_schema_value(item, indent + 1)
+                lines.append(f"  {indent_str}{formatted_item},")
+            # Remove trailing comma from last item
+            if lines[-1].endswith(","):
+                lines[-1] = lines[-1][:-1]
+            lines.append(f"{indent_str}]")
+            return "\n".join(lines)
+    elif isinstance(value, str) and (
+        value.startswith(("Literal[", "Union[", "Optional[", "Dict[", "Tuple[", "Set[", "Enum["))
+        or value in ("str", "int", "float", "bool")
+    ):
+        # Don't add quotes around type annotations or basic type names
+        return value
+    else:
+        # Regular string values get quotes
+        return f'"{value}"'
+
+
+def render_schema(pydantic_model: type[BaseModel], with_field_descriptions: bool = True) -> str:
+    simplified_schema = get_simple_pydantic_schema(pydantic_model)
+    # Use custom formatting instead of json.dumps to handle Literal types properly
+    schema_str = _format_schema_value(simplified_schema)
+
+    if not with_field_descriptions:
+        return schema_str
+
+    descriptions = []
+    for field_name, field in pydantic_model.model_fields.items():
+        if not field.description:
+            raise ValueError(f"Field '{field_name}' in model '{pydantic_model.__name__}' is missing a description.")
+        descriptions.append(f"{field_name}: {field.description}")
+
+    for field_name, field in pydantic_model.model_fields.items():
+        if isinstance(field.annotation, type) and issubclass(field.annotation, BaseModel):
+            nested_model = field.annotation
+            for nested_field_name, nested_field in nested_model.model_fields.items():
+                if not nested_field.description:
+                    raise ValueError(
+                        f"Field '{nested_field_name}' in nested model '{nested_model.__name__}' is missing a description."
+                    )
+                descriptions.append(f"{field_name}.{nested_field_name}: {nested_field.description}")
+        elif get_origin(field.annotation) is list and get_args(field.annotation):
+            list_item_type = get_args(field.annotation)[0]
+            if isinstance(list_item_type, type) and issubclass(list_item_type, BaseModel):
+                for nested_field_name, nested_field in list_item_type.model_fields.items():
+                    if not nested_field.description:
+                        raise ValueError(
+                            f"Field '{nested_field_name}' in list item model '{list_item_type.__name__}' is missing a description."
+                        )
+                    descriptions.append(f"{field_name}[].{nested_field_name}: {nested_field.description}")
+
+    return f"{schema_str}\n\n{'\n'.join(descriptions)}"
+
+
+def render_pydantic_model(pydantic_model: BaseModel) -> str:
+    return pydantic_model.model_dump_json()