File size: 1,442 Bytes
9ad58c8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
# 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)
|