diff --git a/.gitignore b/.gitignore index 25d51645..69cc4315 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,5 @@ bin .DS_Store *.sql *.log -*.yaml \ No newline at end of file +*.yaml +examples/local/* \ No newline at end of file diff --git a/examples/python/output_parser.py b/examples/python/output_parser.py index c07e7d0f..689717e9 100644 --- a/examples/python/output_parser.py +++ b/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', }, ], } diff --git a/flo_ai/parsers/flo_json_parser.py b/flo_ai/parsers/flo_json_parser.py index c1c11a7f..adf6db40 100644 --- a/flo_ai/parsers/flo_json_parser.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index c6edd8b5..4e3a6c5e 100644 --- a/pyproject.toml +++ b/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 "] license = "MIT" diff --git a/setup.py b/setup.py index f682b6f4..4f28b366 100644 --- a/setup.py +++ b/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,