File size: 1,267 Bytes
f1e6b80 |
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 38 39 40 41 |
from enum import Enum
from typing import Dict, List, Type
from langchain_core.exceptions import OutputParserException
from langchain_core.output_parsers import BaseOutputParser
from langchain_core.utils import pre_init
class EnumOutputParser(BaseOutputParser[Enum]):
"""Parse an output that is one of a set of values."""
enum: Type[Enum]
"""The enum to parse. Its values must be strings."""
@pre_init
def raise_deprecation(cls, values: Dict) -> Dict:
enum = values["enum"]
if not all(isinstance(e.value, str) for e in enum):
raise ValueError("Enum values must be strings")
return values
@property
def _valid_values(self) -> List[str]:
return [e.value for e in self.enum]
def parse(self, response: str) -> Enum:
try:
return self.enum(response.strip())
except ValueError:
raise OutputParserException(
f"Response '{response}' is not one of the "
f"expected values: {self._valid_values}"
)
def get_format_instructions(self) -> str:
return f"Select one of the following options: {', '.join(self._valid_values)}"
@property
def OutputType(self) -> Type[Enum]:
return self.enum
|