structured_output.py

# ruff: noqa: E501

import json

from pydantic import BaseModel

STRUCTURED_OUTPUT_FORMAT_INSTRUCTIONS = """The output should be formatted as a JSON instance that conforms to the JSON schema below.

As an example, for the schema {{"properties": {{"foo": {{"title": "Foo", "description": "a list of strings", "type": "array", "items": {{"type": "string"}}}}}}, "required": ["foo"]}}
the object {{"foo": ["bar", "baz"]}} is a well-formatted instance of the schema. The object {{"properties": {{"foo": ["bar", "baz"]}}}} is not well-formatted.

Here is the output schema:
```
{schema}
```"""


class StructuredOutputParser(BaseModel):
    pydantic_class: type[BaseModel]

    def get_schema(self, indent: int = 2) -> str:
        # Copy schema to avoid altering original Pydantic schema.
        schema = self.pydantic_class.model_json_schema()
        # Iterate over fields to remove from the schema
        for field_name in ["title", "type"]:
            if field_name in schema:
                schema.pop(field_name)
        return json.dumps(schema, indent=indent, ensure_ascii=False)

    def get_format_instructions(self) -> str:
        schema_str = self.get_schema()
        # Ensure json in context is well-formed with double quotes.
        return STRUCTURED_OUTPUT_FORMAT_INSTRUCTIONS.format(schema=schema_str)

    def parse(self, json_string: str) -> BaseModel:
        return self.pydantic_class.model_validate_json(json_string)