Adding support for get_format to parser

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:

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

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()