Literal support in output parser

.gitignore

@@ -12,4 +12,5 @@ bin
 .DS_Store
 *.sql
 *.log
-*.yaml
+*.yaml
+examples/local/*

examples/python/output_parser.py

@@ -34,9 +34,14 @@
             'name': 'middle_name',
         },
         {
-            'type': 'str',
-            'description': 'The last name of the person',
+            'type': 'literal',
+            'description': 'The last name of the person, the value can be either of Vishnu or Satis',
             'name': 'last_name',
+            'values': [
+                {'value': 'Vishnu', 'description': 'If the first_name starts with K'},
+                {'value': 'Satis', 'description': 'If the first_name starts with M'},
+            ],
+            'default_value_prompt': 'If none of the above value is suited, please use value other than the above in snake-case',
         },
     ],
 }

flo_ai/parsers/flo_json_parser.py

@@ -1,6 +1,8 @@
 import json
+import csv
+from io import StringIO
 from flo_ai.parsers.flo_parser import FloParser
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, Literal
 from pydantic import BaseModel, Field, create_model
 from flo_ai.error.flo_exception import FloException
 from langchain_core.output_parsers import PydanticOutputParser
@@ -18,15 +20,64 @@ def __init__(self, parse_contract: ParseContract):
         self.contract = parse_contract
         super().__init__()
 
+    def __dict_list_to_csv_string(self, data):
+        if not data or len(data) == 0:
+            return '```No data provided```'
+        headers = data[0].keys()
+        output = StringIO()
+
+        writer = csv.DictWriter(output, fieldnames=headers)
+        writer.writeheader()
+        writer.writerows(data)
+
+        csv_string = output.getvalue()
+        output.close()
+
+        return f'```\n{csv_string}```'
+
     def __create_contract_from_json(self) -> BaseModel:
-        type_mapping = {'str': str, 'int': int, 'bool': bool, 'float': float}
-        pydantic_fields = {
-            field['name']: (
-                type_mapping[field['type']],
-                Field(..., description=field['description']),
-            )
-            for field in self.contract.fields
+        type_mapping = {
+            'str': str,
+            'int': int,
+            'bool': bool,
+            'float': float,
+            'literal': Literal,
         }
+        pydantic_fields = {}
+        for field in self.contract.fields:
+            field_type = field['type']
+            if field_type == 'literal':
+                literal_values = field.get('values', [])
+                if not literal_values:
+                    raise ValueError(
+                        f"Field '{field['name']}' of type 'literal' must specify 'values'."
+                    )
+                literals = [literal_value['value'] for literal_value in literal_values]
+                field_type_annotation = Literal[tuple(literals)]
+                default_prompt = (
+                    field['default_value_prompt']
+                    if 'default_value_prompt' in field
+                    else ''
+                )
+                field_description = f"""
+                {field['description']}
+                Following are the list of possibles values and its correponding description:
+                {self.__dict_list_to_csv_string(literal_values)}
+
+                This should be one of the values in the `value` column in the above csv.
+                {default_prompt}
+                """
+            else:
+                field_type_annotation = type_mapping.get(field_type)
+                if field_type_annotation is None:
+                    raise ValueError(f'Unsupported type: {field_type}')
+                field_description = field['description']
+
+            pydantic_fields[field['name']] = (
+                field_type_annotation,
+                Field(..., description=field_description),
+            )
+
         DynamicModel = create_model(self.contract.name, **pydantic_fields)
         return DynamicModel
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flo-ai"
-version = "0.0.5-dev2"
+version = "0.0.5-dev3"
 description = "A easy way to create structured AI agents"
 authors = ["vizsatiz <vishnu@rootfor.xyz>"]
 license = "MIT"

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='flo-ai',
-    version='0.0.5-dev2',
+    version='0.0.5-dev3',
     author='Rootflo',
     description='Create composable AI agents',
     long_description=long_description,