forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import json
|
|
import re
|
|
from typing import Type, TypeVar
|
|
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from langchain.output_parsers.format_instructions import PYDANTIC_FORMAT_INSTRUCTIONS
|
|
from langchain.schema import BaseOutputParser, OutputParserException
|
|
|
|
T = TypeVar("T", bound=BaseModel)
|
|
|
|
|
|
class PydanticOutputParser(BaseOutputParser[T]):
|
|
pydantic_object: Type[T]
|
|
|
|
def parse(self, text: str) -> T:
|
|
try:
|
|
# Greedy search for 1st json candidate.
|
|
match = re.search(
|
|
"\{.*\}", text.strip(), re.MULTILINE | re.IGNORECASE | re.DOTALL
|
|
)
|
|
json_str = ""
|
|
if match:
|
|
json_str = match.group()
|
|
json_object = json.loads(json_str)
|
|
return self.pydantic_object.parse_obj(json_object)
|
|
|
|
except (json.JSONDecodeError, ValidationError) as e:
|
|
name = self.pydantic_object.__name__
|
|
msg = f"Failed to parse {name} from completion {text}. Got: {e}"
|
|
raise OutputParserException(msg)
|
|
|
|
def get_format_instructions(self) -> str:
|
|
schema = self.pydantic_object.schema()
|
|
|
|
# Remove extraneous fields.
|
|
reduced_schema = schema
|
|
if "title" in reduced_schema:
|
|
del reduced_schema["title"]
|
|
if "type" in reduced_schema:
|
|
del reduced_schema["type"]
|
|
# Ensure json in context is well-formed with double quotes.
|
|
schema_str = json.dumps(reduced_schema)
|
|
|
|
return PYDANTIC_FORMAT_INSTRUCTIONS.format(schema=schema_str)
|