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.
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, List
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from langchain.output_parsers.format_instructions import STRUCTURED_FORMAT_INSTRUCTIONS
|
|
from langchain.schema import BaseOutputParser, OutputParserException
|
|
|
|
line_template = '\t"{name}": {type} // {description}'
|
|
|
|
|
|
class ResponseSchema(BaseModel):
|
|
name: str
|
|
description: str
|
|
|
|
|
|
def _get_sub_string(schema: ResponseSchema) -> str:
|
|
return line_template.format(
|
|
name=schema.name, description=schema.description, type="string"
|
|
)
|
|
|
|
|
|
class StructuredOutputParser(BaseOutputParser):
|
|
response_schemas: List[ResponseSchema]
|
|
|
|
@classmethod
|
|
def from_response_schemas(
|
|
cls, response_schemas: List[ResponseSchema]
|
|
) -> StructuredOutputParser:
|
|
return cls(response_schemas=response_schemas)
|
|
|
|
def get_format_instructions(self) -> str:
|
|
schema_str = "\n".join(
|
|
[_get_sub_string(schema) for schema in self.response_schemas]
|
|
)
|
|
return STRUCTURED_FORMAT_INSTRUCTIONS.format(format=schema_str)
|
|
|
|
def parse(self, text: str) -> Any:
|
|
if "```json" not in text:
|
|
raise OutputParserException(
|
|
f"Got invalid return object. Expected markdown code snippet with JSON "
|
|
f"object, but got:\n{text}"
|
|
)
|
|
|
|
json_string = text.split("```json")[1].strip().strip("```").strip()
|
|
try:
|
|
json_obj = json.loads(json_string)
|
|
except json.JSONDecodeError as e:
|
|
raise OutputParserException(f"Got invalid JSON object. Error: {e}")
|
|
for schema in self.response_schemas:
|
|
if schema.name not in json_obj:
|
|
raise OutputParserException(
|
|
f"Got invalid return object. Expected key `{schema.name}` "
|
|
f"to be present, but got {json_obj}"
|
|
)
|
|
return json_obj
|
|
|
|
@property
|
|
def _type(self) -> str:
|
|
return "structured"
|