How to Parse Output to Enum in Pydantic Class

Question

Hi everybody... I have 2 questions, I solved them... But I don't like the way, maybe you have other solutions

Imagine I have some Enum:

Plain Text

class MyEnum(Enum): ...

Does anybody know how to:

Parse output to Enum?
I did the following: I created the Pydantic class with this field and used PydanticOutputParser:

Plain Text

  class MyEnumType(BaseModel):
    my_enum: MyEnum
  
  program = LLMTextCompletionProgram.from_defaults(
    llm=llm,
    output_parser=PydanticOutputParser(output_cls=MyEnumType),
    prompt_template_str='......',
  )
  result = program(prompt=....)
  return query_type.my_enum

How can I pass Enum as a parameter to the agent function?

Plain Text

def my_agent(param: MyEnum):
  ....

As I understand it passes only basic Python types. So I just manually convert it back:

Plain Text

def my_agent(value: str):
  param = MyEnum(value)

Logan M · Answer

Pydantic class/json schema is really the only way to do this. Alternatively, you could just use llm.complete() or llm.chat() and manually prompt the llm, and parse the output yourself. If you are using openai (or some model that has a tools API), I would use the pydantic program instead.If your enum is just strings, define it like class MyEnumType(str, Enum):, then you can use it directly in a function definition, and FunctionTool will parse it just fine. (you'll probably still need to convert it back to an enum, although there's maybe some pydantic trickery that could work here, def my_agent(value: MyEnum = Field()) ?

Find answers from the community

How to Parse Output to Enum in Pydantic Class