From 37b4895cdf3ebf68e5b1dd836a01fa37b0336a85 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Thu, 13 Mar 2025 17:13:01 +0530 Subject: [PATCH] Adding support for get_format to parser --- flo_ai/flo_ai/parsers/flo_json_parser.py | 3 +++ flo_ai/flo_ai/parsers/flo_parser.py | 5 +++++ flo_ai/flo_ai/parsers/flo_pydantic_parser.py | 3 +++ 3 files changed, 11 insertions(+) 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()