diff --git a/flo_ai/flo_ai/parsers/flo_json_parser.py b/flo_ai/flo_ai/parsers/flo_json_parser.py index 4d196e5f..a060c9fc 100644 --- a/flo_ai/flo_ai/parsers/flo_json_parser.py +++ b/flo_ai/flo_ai/parsers/flo_json_parser.py @@ -96,6 +96,9 @@ def __create_literal_type(self, field: Dict[str, Any]) -> Any: literals = [literal_value['value'] for literal_value in literal_values] return Literal[tuple(literals)] + def get_format(self) -> BaseModel: + return self.__create_contract_from_json() + def __create_contract_from_json(self) -> BaseModel: pydantic_fields = {} for field in self.contract.fields: diff --git a/flo_ai/flo_ai/parsers/flo_parser.py b/flo_ai/flo_ai/parsers/flo_parser.py index 65a06323..ad2b32e3 100644 --- a/flo_ai/flo_ai/parsers/flo_parser.py +++ b/flo_ai/flo_ai/parsers/flo_parser.py @@ -1,7 +1,12 @@ from abc import ABC, abstractmethod +from pydantic import BaseModel class FloParser(ABC): @abstractmethod def get_format_instructions(self): pass + + @abstractmethod + def get_format(self) -> BaseModel: + pass diff --git a/flo_ai/flo_ai/parsers/flo_pydantic_parser.py b/flo_ai/flo_ai/parsers/flo_pydantic_parser.py index a140b34c..b7bca7b9 100644 --- a/flo_ai/flo_ai/parsers/flo_pydantic_parser.py +++ b/flo_ai/flo_ai/parsers/flo_pydantic_parser.py @@ -13,6 +13,9 @@ def get_format_instructions(self): pydantic_object=self.model ).get_format_instructions() + def get_format(self): + return self.model + def create(output_model: BaseModel): return FloPydanticParser.Builder(output_model).build()