feat: support descriptions of class labels in semantic.classify (#102)

## Summary

This PR enhances the semantic.classify API to support optional
descriptions for classification classes, improving LLM performance by
providing better context. The enhancement maintains backward
compatibility with string list input while removing enum support.

## Motivation

Previously, the classify API only accepted simple string labels or
enums, which limited the LLM's understanding of nuanced classification
categories. By allowing optional descriptions, we can provide richer
context that helps the model make more accurate classification
decisions.

## Changes

### New Features
- **ClassDefinition model**: New Pydantic model with `label` and
`description` fields for enhanced class definitions
- **Enhanced prompt generation**: System messages now include class
descriptions when provided, giving the LLM better context
- **Flexible API**: Support for both simple string lists (backward
compatible) and detailed ClassDefinition objects

### 🔧 Technical Changes
- Replaced `List[str] | type[Enum]` API with `Union[List[str],
List[ClassDefinition]]` for explicit type handling
- Added internal `ResolvedClassDefinition` dataclass for normalized
representation
- Updated validation logic to work with label sets instead of enums

## Usage Examples

### Basic usage (backward compatible)
```python
semantic.classify("support_ticket", ["technical_issue", "billing_inquiry", "general_support"])
```

### Enhanced usage with descriptions
```python
semantic.classify("support_ticket", [
    ClassDefinition(
        label="technical_issue",
        description="Problems with product functionality, bugs, or technical difficulties"
    ),
    ClassDefinition(
        label="billing_inquiry",
        description="Questions about charges, payments, subscriptions, or account billing"
    ),
    ClassDefinition(
        label="general_support",
        description="General questions, feature requests, or non-technical assistance"
    )
])
```

Author: Rohit Rastogi

src/fenic/__init__.py

@@ -52,6 +52,7 @@
 from fenic.core import (
     ArrayType,
     BooleanType,
+    ClassDefinition,
     ClassifyExample,
     ClassifyExampleCollection,
     ColumnField,
@@ -129,6 +130,7 @@
     "PredicateExample",
     "PredicateExampleCollection",
     "Schema",
+    "ClassDefinition",
     "ClassifyExample",
     "ClassifyExampleCollection",
     "JoinExample",

src/fenic/_backends/local/semantic_operators/analyze_sentiment.py

@@ -1,6 +1,5 @@
 import json
 import logging
-from enum import Enum
 from typing import List, Optional
 
 import polars as pl
@@ -110,17 +109,7 @@
     )
 )
 
-SENTIMENT_ANALYSIS_MODEL = create_classification_pydantic_model(
-    Enum(
-        "Sentiment",
-        [
-            ("POSITIVE", "positive"),
-            ("NEGATIVE", "negative"),
-            ("NEUTRAL", "neutral"),
-        ],
-    )
-)
-
+SENTIMENT_ANALYSIS_MODEL = create_classification_pydantic_model(["positive", "negative", "neutral"])
 
 class AnalyzeSentiment(BaseSingleColumnInputOperator[str, str]):
     SYSTEM_PROMPT = """You are a sentiment analysis expert.
@@ -166,6 +155,11 @@ def postprocess(self, responses: List[Optional[str]]) -> List[Optional[str]]:
                 try:
                     data = json.loads(response)["output"]
                     predictions.append(data)
+                    if data not in ["positive", "negative", "neutral"]:
+                        logger.warning(
+                            f"Model returned invalid label '{data}'. Valid labels: positive, negative, neutral"
+                        )
+                        predictions.append(None)
                 except Exception as e:
                     logger.warning(
                         f"Invalid model output: {response} for semantic.analyze_sentiment: {e}"

src/fenic/_backends/local/semantic_operators/classify.py

@@ -1,6 +1,5 @@
 import json
 import logging
-from enum import Enum
 from typing import List, Optional, Type
 
 import polars as pl
@@ -12,14 +11,14 @@
 )
 from fenic._backends.local.semantic_operators.utils import (
     create_classification_pydantic_model,
-    stringify_enum_type,
 )
 from fenic._constants import (
     MAX_TOKENS_DETERMINISTIC_OUTPUT_SIZE,
     TOKEN_OVERHEAD_JSON,
     TOKEN_OVERHEAD_MISC,
 )
 from fenic._inference.language_model import InferenceConfiguration, LanguageModel
+from fenic.core._logical_plan.expressions import ResolvedClassDefinition
 from fenic.core.types import ClassifyExample, ClassifyExampleCollection
 
 logger = logging.getLogger(__name__)
@@ -28,20 +27,24 @@
 class Classify(BaseSingleColumnInputOperator[str, str]):
     SYSTEM_PROMPT = (
         "You are a text classification expert. "
-        "Classify the following document into one of the following labels: {labels}. "
+        "Classify the following document into one of the following labels:"
+        "\n{classes}\n"
         "Respond with *only* the predicted label."
     )
 
     def __init__(
         self,
         input: pl.Series,
-        labels: Type[Enum],
+        classes: List[ResolvedClassDefinition],
         model: LanguageModel,
         temperature: float,
         examples: Optional[ClassifyExampleCollection] = None,
     ):
+        self.classes = classes
+        self.valid_labels = {class_def.label for class_def in classes}
+        # Create output model from class labels
+        labels = [class_def.label for class_def in classes]
         self.output_model = create_classification_pydantic_model(labels)
-        self.labels = labels
         super().__init__(
             input,
             CompletionOnlyRequestSender(
@@ -57,7 +60,17 @@ def __init__(
         )
 
     def build_system_message(self) -> str:
-        return self.SYSTEM_PROMPT.format(labels=stringify_enum_type(self.labels))
+        """Build system message with class descriptions."""
+        class_descriptions = []
+        for class_def in self.classes:
+            if class_def.description:
+                class_descriptions.append(f"- {class_def.label}: {class_def.description}")
+            else:
+                class_descriptions.append(f"- {class_def.label}")
+
+        classes_text = "\n".join(class_descriptions)
+
+        return self.SYSTEM_PROMPT.format(classes=classes_text)
 
     def postprocess(self, responses: List[Optional[str]]) -> List[Optional[str]]:
         predictions = []
@@ -67,7 +80,14 @@ def postprocess(self, responses: List[Optional[str]]) -> List[Optional[str]]:
             else:
                 try:
                     data = json.loads(response)["output"]
-                    predictions.append(data)
+                    # Validate the response is one of the valid labels
+                    if data not in self.valid_labels:
+                        logger.warning(
+                            f"Model returned invalid label '{data}'. Valid labels: {self.valid_labels}"
+                        )
+                        predictions.append(None)
+                    else:
+                        predictions.append(data)
                 except Exception as e:
                     logger.warning(
                         f"Invalid model output: {response} for semantic.classify: {e}"
@@ -82,5 +102,5 @@ def convert_example_to_assistant_message(self, example: ClassifyExample) -> str:
         return self.output_model(output=example.output).model_dump_json()
 
     def get_max_tokens(self) -> int:
-        max_label_length = max(len(str(label.value)) for label in self.labels)
+        max_label_length = max(len(class_def.label) for class_def in self.classes)
         return max(MAX_TOKENS_DETERMINISTIC_OUTPUT_SIZE, TOKEN_OVERHEAD_JSON + TOKEN_OVERHEAD_MISC + max_label_length)

src/fenic/_backends/local/semantic_operators/utils.py

@@ -1,6 +1,6 @@
 import re
 from enum import Enum
-from typing import Any, Dict, Type
+from typing import Any, Dict, List, Type
 
 import polars as pl
 from pydantic import BaseModel, create_model
@@ -37,15 +37,19 @@ def stringify_enum_type(enum_type: Type[Enum]) -> str:
     return ", ".join(f"{label.value}" for label in enum_type)
 
 
-def create_classification_pydantic_model(enum_cls: Type[Enum]) -> Type[BaseModel]:
-    """Creates a Pydantic model from an Enum class.
+def create_classification_pydantic_model(allowed_values: List[str]) -> type[BaseModel]:
+    """Creates a Pydantic model from a list of allowed string values using a dynamic Enum.
 
     Args:
-        enum_cls (Type[Enum]): The Enum class to convert.
+        allowed_values (List[str]): The list of allowed string values.
 
     Returns:
         Type[BaseModel]: A Pydantic model class with a field for the Enum values.
     """
+    enum_name = "LabelEnum"
+    enum_members = {value.upper(): value for value in allowed_values}
+    enum_cls = Enum(enum_name, enum_members)
+
     return create_model(
         "EnumModel",
         output=(enum_cls, ...),

src/fenic/_backends/local/transpiler/expr_converter.py

@@ -554,16 +554,12 @@ def sem_predicate_fn(batch: pl.Series) -> pl.Series:
     @_convert_expr.register(SemanticClassifyExpr)
     def _convert_semantic_classify_expr(self, logical: SemanticClassifyExpr) -> pl.Expr:
         def sem_classify_fn(batch: pl.Series) -> pl.Series:
-            labels_enum = (
-                SemanticClassifyExpr.transform_labels_list_into_enum(logical.labels)
-                if isinstance(logical.labels, list)
-                else logical.labels
-            )
             return SemanticClassify(
                 input=batch,
-                labels=labels_enum,
+                classes=logical.classes,
                 model=self.session_state.get_language_model(logical.model_alias),
                 temperature=logical.temperature,
+                examples=logical.examples,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(

src/fenic/api/functions/semantic.py

@@ -1,6 +1,5 @@
 """Semantic functions for Fenic DataFrames - LLM-based operations."""
 
-from enum import Enum
 from typing import List, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, validate_call
@@ -9,6 +8,7 @@
 from fenic.core._logical_plan.expressions import (
     AnalyzeSentimentExpr,
     EmbeddingsExpr,
+    ResolvedClassDefinition,
     SemanticClassifyExpr,
     SemanticExtractExpr,
     SemanticMapExpr,
@@ -22,6 +22,7 @@
 )
 from fenic.core.error import ValidationError
 from fenic.core.types import (
+    ClassDefinition,
     ClassifyExampleCollection,
     KeyPoints,
     MapExampleCollection,
@@ -259,60 +260,86 @@ def reduce(
 @validate_call(config=ConfigDict(strict=True, arbitrary_types_allowed=True))
 def classify(
         column: ColumnOrName,
-        labels: List[str] | type[Enum],
+        classes: Union[List[str], List[ClassDefinition]],
         examples: Optional[ClassifyExampleCollection] = None,
         model_alias: Optional[str] = None,
         temperature: float = 0,
 ) -> Column:
-    """Classifies a string column into one of the provided labels.
+    """Classifies a string column into one of the provided classes.
 
     This is useful for tagging incoming documents with predefined categories.
 
     Args:
         column: Column or column name containing text to classify.
-
-        labels: List of category strings or an Enum defining the categories to classify the text into.
-
+        classes: List of class labels or ClassDefinition objects defining the available classes. Use ClassDefinition objects to provide descriptions for the classes.
         examples: Optional collection of example classifications to guide the model.
             Examples should be created using ClassifyExampleCollection.create_example(),
             with instruction variables mapped to their expected classifications.
-
         model_alias: Optional alias for the language model to use for the mapping. If None, will use the language model configured as the default.
-
         temperature: Optional temperature parameter for the language model. If None, will use the default temperature (0.0).
 
     Returns:
         Column: Expression containing the classification results.
 
     Raises:
-        ValueError: If column is invalid or categories is not a list of strings.
+        ValueError: If column is invalid or classes is empty or has duplicate labels.
 
     Example: Categorizing incoming support requests
         ```python
         # Categorize incoming support requests
         semantic.classify("message", ["Account Access", "Billing Issue", "Technical Problem"])
         ```
 
-    Example: Categorizing incoming support requests with examples
+    Example: Categorizing incoming support requests using ClassDefinition objects
+        ```python
+        # Categorize incoming support requests
+        semantic.classify("message", [
+            ClassDefinition(label="Account Access", description="General questions, feature requests, or non-technical assistance"),
+            ClassDefinition(label="Billing Issue", description="Questions about charges, payments, subscriptions, or account billing"),
+            ClassDefinition(label="Technical Problem", description="Problems with product functionality, bugs, or technical difficulties")
+        ])
+        ```
+
+    Example: Categorizing incoming support requests with ClassDefinition objects and examples
         ```python
         examples = ClassifyExampleCollection()
+        class_definitions = [
+            ClassDefinition(label="Account Access", description="General questions, feature requests, or non-technical assistance"),
+            ClassDefinition(label="Billing Issue", description="Questions about charges, payments, subscriptions, or account billing"),
+            ClassDefinition(label="Technical Problem", description="Problems with product functionality, bugs, or technical difficulties")
+        ]
         examples.create_example(ClassifyExample(
             input="I can't reset my password or access my account.",
             output="Account Access"))
         examples.create_example(ClassifyExample(
             input="You charged me twice for the same month.",
             output="Billing Issue"))
-        semantic.classify("message", ["Account Access", "Billing Issue", "Technical Problem"], examples)
+        semantic.classify("message", class_definitions, examples)
         ```
     """
-    if isinstance(labels, List) and len(labels) == 0:
-        raise ValueError(
-            f"Must specify the categories for classification, found: {len(labels)} categories"
+    if len(classes) < 2:
+        raise ValidationError(
+            "The `classes` list must contain at least two ClassDefinition objects. "
+            "You provided only one. Classification requires at least two possible labels."
         )
+
+    # Validate unique labels
+    if isinstance(classes[0], ClassDefinition):
+        classes = [ResolvedClassDefinition(label=class_def.label, description=class_def.description) for class_def in classes]
+    else:
+        classes = [ResolvedClassDefinition(label=class_def, description=None) for class_def in classes]
+
+    labels = [class_def.label for class_def in classes]
+    duplicates = {label for label in labels if labels.count(label) > 1}
+    if duplicates:
+        raise ValidationError(
+            f"Class labels must be unique. The following duplicate label(s) were found: {sorted(duplicates)}"
+        )
+
     return Column._from_logical_expr(
         SemanticClassifyExpr(
             Column._from_col_or_name(column)._logical_expr,
-            labels,
+            classes,
             examples=examples,
             model_alias=model_alias,
             temperature=temperature,

src/fenic/core/__init__.py

@@ -10,6 +10,7 @@
     ArrayType,
     BooleanType,
     BranchSide,
+    ClassDefinition,
     ClassifyExample,
     ClassifyExampleCollection,
     ColumnField,
@@ -61,6 +62,7 @@
     "TranscriptType",
     "ColumnField",
     "Schema",
+    "ClassDefinition",
     "ClassifyExample",
     "ClassifyExampleCollection",
     "JoinExample",

src/fenic/core/_logical_plan/expressions/__init__.py

@@ -75,6 +75,9 @@
 from fenic.core._logical_plan.expressions.semantic import (
     EmbeddingsExpr as EmbeddingsExpr,
 )
+from fenic.core._logical_plan.expressions.semantic import (
+    ResolvedClassDefinition as ResolvedClassDefinition,
+)
 from fenic.core._logical_plan.expressions.semantic import (
     SemanticClassifyExpr as SemanticClassifyExpr,
 )