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.
32 lines
706 B
Python
32 lines
706 B
Python
from enum import Enum
|
|
|
|
from langchain.output_parsers.enum import EnumOutputParser
|
|
from langchain.schema import OutputParserException
|
|
|
|
|
|
class Colors(Enum):
|
|
RED = "red"
|
|
GREEN = "green"
|
|
BLUE = "blue"
|
|
|
|
|
|
def test_enum_output_parser_parse() -> None:
|
|
parser = EnumOutputParser(enum=Colors)
|
|
|
|
# Test valid inputs
|
|
result = parser.parse("red")
|
|
assert result == Colors.RED
|
|
|
|
result = parser.parse("green")
|
|
assert result == Colors.GREEN
|
|
|
|
result = parser.parse("blue")
|
|
assert result == Colors.BLUE
|
|
|
|
# Test invalid input
|
|
try:
|
|
parser.parse("INVALID")
|
|
assert False, "Should have raised OutputParserException"
|
|
except OutputParserException:
|
|
pass
|