Repository: SylphAI-Inc/AdalFlow

Issue:
[Bug] [DataClass] Lack support for Enum type
@zjffdu reported the bug on not supporting Enum

class Label(str, enum.Enum):
    """Enumeration for single-label text classification."""
    SPAM = "spam"
    NOT_SPAM = "not_spam"

@dataclass
class ClassificationOutput(DataClass):
    """
    Class for a single class label prediction.
    """
    label: Label = field(metadata={"desc": "Label of the category."})
    
 We will add the support for this. Allow field type Enum and even list of Enum.

Discussion:

Solution patch:
diff --git a/lightrag/README.md b/lightrag/README.md
index 493daf74..dd593a74 100644
--- a/lightrag/README.md
+++ b/lightrag/README.md
@@ -5,9 +5,9 @@
 <!-- [![GitHub star chart](https://img.shields.io/github/stars/SylphAI-Inc/LightRAG?style=flat-square)](https://star-history.com/#SylphAI-Inc/LightRAG) -->
 [![License](https://img.shields.io/github/license/SylphAI-Inc/LightRAG)](https://opensource.org/license/MIT)
 [![PyPI](https://img.shields.io/pypi/v/lightRAG?style=flat-square)](https://pypi.org/project/lightRAG/)
-[![PyPI - Downloads](https://img.shields.io/pypi/dm/lightRAG?style=flat-square)](https://pypistats.org/packages/lightRAG)
 [![Open Issues](https://img.shields.io/github/issues-raw/SylphAI-Inc/LightRAG?style=flat-square)](https://github.com/SylphAI-Inc/LightRAG/issues)
 [![](https://dcbadge.vercel.app/api/server/zt2mTPcu?compact=true&style=flat)](https://discord.gg/zt2mTPcu)
+<!-- [![PyPI - Downloads](https://img.shields.io/pypi/dm/lightRAG?style=flat-square)](https://pypistats.org/packages/lightRAG) -->
 
 
 ### ⚡ The Lightning Library for Large Language Model Applications ⚡
diff --git a/lightrag/lightrag/core/functional.py b/lightrag/lightrag/core/functional.py
index d1282aec..6b329fbf 100644
--- a/lightrag/lightrag/core/functional.py
+++ b/lightrag/lightrag/core/functional.py
@@ -19,6 +19,7 @@ from typing import (
 )
 import logging
 import numpy as np
+from enum import Enum
 import re
 import json
 import yaml
@@ -381,12 +382,30 @@ def get_type_schema(
         return "Tuple"
 
     elif is_dataclass(type_obj):
+        if issubclass(type_obj, Enum):
+            # Handle Enum dataclass types
+            enum_members = ", ".join([f"{e.name}={e.value}" for e in type_obj])
+            return f"Enum[{type_obj.__name__}({enum_members})]"
         # Recursively handle nested dataclasses
         output = str(get_dataclass_schema(type_obj, exclude, type_var_map))
         return output
+
+    elif isinstance(type_obj, type) and issubclass(type_obj, Enum):
+        # Handle Enum types
+        enum_members = ", ".join([f"{e.name}={e.value}" for e in type_obj])
+        return f"Enum[{type_obj.__name__}({enum_members})]"
+
     return type_obj.__name__ if hasattr(type_obj, "__name__") else str(type_obj)
 
 
+def get_enum_schema(enum_cls: Type[Enum]) -> Dict[str, object]:
+    return {
+        "type": "string",
+        "enum": [e.value for e in enum_cls],
+        "description": enum_cls.__doc__ if enum_cls.__doc__ else "",
+    }
+
+
 def get_dataclass_schema(
     cls,
     exclude: ExcludeType = None,
@@ -406,9 +425,21 @@ def get_dataclass_schema(
 
     type_var_map = type_var_map or {}
 
-    schema = {"type": cls.__name__, "properties": {}, "required": []}
+    # TODO: Add support for having a description in the dataclass
+    schema = {
+        "type": cls.__name__,
+        "properties": {},
+        "required": [],
+        # "description": cls.__doc__ if cls.__doc__ else "",
+    }
     # get the exclude list for the current class
     current_exclude = exclude.get(cls.__name__, []) if exclude else []
+
+    # handle Combination of Enum and dataclass
+    if issubclass(cls, Enum):
+        schema["type"] = get_type_schema(cls, exclude, type_var_map)
+        return schema
+
     for f in fields(cls):
         if f.name in current_exclude:
             continue


Test patch:
diff --git a/lightrag/tests/test_base_data_class.py b/lightrag/tests/test_base_data_class.py
index 4d0db965..8084f4b7 100644
--- a/lightrag/tests/test_base_data_class.py
+++ b/lightrag/tests/test_base_data_class.py
@@ -2,6 +2,8 @@ import unittest
 from lightrag.core import DataClass, required_field
 from dataclasses import field, MISSING, dataclass, asdict
 
+from lightrag.core.functional import get_type_schema
+
 
 from typing import List, Dict, Optional, Set
 
@@ -27,6 +29,26 @@ class Address:
     zipcode: str
 
 
+import enum
+
+
+# @dataclass
+class Label(str, enum.Enum):
+    """Enumeration for single-label text classification."""
+
+    SPAM = "spam"
+    NOT_SPAM = "not_spam"
+
+
+@dataclass
+class ClassificationOutput(DataClass):
+    """
+    Class for a single class label prediction.
+    """
+
+    label: Label = field(metadata={"desc": "Label of the category."})
+
+
 # Example instance of the nested dataclasses and complex data types as list, dict
 @dataclass
 class Person(DataClass):
@@ -215,10 +237,115 @@ class TestBaseDataClass(unittest.TestCase):
         }
         output = Person.to_dict_class(exclude=exclude)
         print(f"output 1: {output}")
-        # self.assertEqual(output, expected_result)
+        self.assertEqual(output, expected_result)
 
     def test_error_non_dataclass(self):
         """Test error handling when to_dict is called on a non-dataclass."""
         with self.assertRaises(AttributeError):
             non_dataclass = "Not a dataclass"
             non_dataclass.to_dict()
+
+
+class TestGetTypeSchema(unittest.TestCase):
+
+    def test_enum_schema(self):
+        result = get_type_schema(Label)
+        expected = "Enum[Label(SPAM=spam, NOT_SPAM=not_spam)]"
+        self.assertEqual(result, expected)
+
+    def test_enum_field_in_dataclass(self):
+        result = get_type_schema(ClassificationOutput)
+        expected = "{'type': 'ClassificationOutput', 'properties': {'label': {'type': 'Enum[Label(SPAM=spam, NOT_SPAM=not_spam)]', 'desc': 'Label of the category.'}}, 'required': ['label']}"
+        self.assertEqual(result, expected)
+
+    def test_optional_enum_field(self):
+        @dataclass
+        class OptionalClassificationOutput:
+            label: Optional[Label] = field(
+                default=None, metadata={"desc": "Label of the category."}
+            )
+
+        result = get_type_schema(OptionalClassificationOutput)
+        expected = "{'type': 'OptionalClassificationOutput', 'properties': {'label': {'type': 'Optional[Enum[Label(SPAM=spam, NOT_SPAM=not_spam)]]', 'desc': 'Label of the category.'}}, 'required': []}"
+        self.assertEqual(result, expected)
+
+    def test_enum_as_list_element(self):
+        @dataclass
+        class EnumListClassificationOutput:
+            labels: List[Label] = field(
+                default_factory=list, metadata={"desc": "List of labels."}
+            )
+
+        result = get_type_schema(EnumListClassificationOutput)
+        print(f"result: {result}")
+        expected = "{'type': 'EnumListClassificationOutput', 'properties': {'labels': {'type': 'List[Enum[Label(SPAM=spam, NOT_SPAM=not_spam)]]', 'desc': 'List of labels.'}}, 'required': []}"
+        self.assertEqual(result, expected)
+
+        # test instance with DataClass
+        @dataclass
+        class EnumListClassificationOutput(DataClass):
+            labels: List[Label] = field(
+                default_factory=list, metadata={"desc": "List of labels."}
+            )
+
+        instance = EnumListClassificationOutput(labels=[Label.SPAM, Label.NOT_SPAM])
+        result = instance.to_dict()
+        print(f"EnumListClassificationOutput instance: {result}")
+        expected = "{'labels': [<Label.SPAM: 'spam'>, <Label.NOT_SPAM: 'not_spam'>]}"
+        self.assertEqual(str(result), expected)
+        restored_instance = EnumListClassificationOutput.from_dict(result)
+        print(f"restored_instance: {restored_instance}")
+        print(f"instance: {instance}")
+        self.assertEqual(restored_instance, instance)
+
+        # from lightrag.core import Generator
+        # from lightrag.components.model_client.openai_client import OpenAIClient
+        # from lightrag.components.output_parsers import JsonOutputParser
+        # from lightrag.utils import setup_env
+
+        # setup_env()
+        # parser = JsonOutputParser(
+        #     data_class=EnumListClassificationOutput, return_data_class=True
+        # )
+        # generator = Generator(
+        #     model_client=OpenAIClient(),
+        #     prompt_kwargs={
+        #         "task_desc_str": "Classify the following text as spam or not spam.",
+        #         "output_format_str": parser.format_instructions(),
+        #     },
+        #     model_kwargs={"model": "gpt-3.5-turbo"},
+        #     output_processors=parser,
+        # )
+        # prompt_kwargs = {"input_str": "This is a spam message."}
+        # result = generator(prompt_kwargs)
+        # print(f"result: {result}")
+
+    def test_enum_as_data_class(self):
+        @dataclass
+        class LabelDataClass(DataClass, str, enum.Enum):
+            """Enumeration for single-label text classification."""
+
+            SPAM = "spam"
+            NOT_SPAM = "not_spam"
+
+        schema = LabelDataClass.to_schema()
+        print(f"schema: {schema}")
+        type_schema = get_type_schema(LabelDataClass)
+        print(f"type_schema: {type_schema}")
+        self.assertEqual(
+            schema,
+            {
+                "type": "Enum[LabelDataClass(SPAM=spam, NOT_SPAM=not_spam)]",
+                "properties": {},
+                "required": [],
+            },
+        )
+        self.assertEqual(
+            type_schema,
+            "Enum[LabelDataClass(SPAM=spam, NOT_SPAM=not_spam)]",
+        )
+
+
+if __name__ == "__main__":
+
+    unittest.main()