Update src/schemas/_pydantic_agent.py

src/schemas/_pydantic_agent.py

@@ -1,4 +1,5 @@
 from typing import Optional
+from enum import Enum
 from pydantic import BaseModel, Field, model_validator
 
 
@@ -117,6 +118,15 @@ class OrderConditions(BaseModel):
             "Use 'ASC' for: oldest, lowest, alphabetical, earliest."
         ),
     )
+    
+    
+class ChartsType(Enum):
+    """
+    Enumeration of supported chart types for visualization.
+    """
+    PIE: str = "PIE"
+    BAR: str = "BAR"
+    LINE: str = "LINE"
 
 
 class SQLQueryExtractor(BaseModel):
@@ -211,6 +221,18 @@ class SQLQueryExtractor(BaseModel):
             "This must be set True if provided tool return true else set False. "
         ),
     )
+    
+    best_suitable_chart: ChartsType = Field(
+        ...,
+        description=(
+            "The type of chart that best represents the data for visualization. "
+            "Choose based on the user query and the nature of the data: "
+            "- PIE: Use for showing proportions or percentages of a whole (e.g., market share, distribution). "
+            "- BAR: Use for comparing quantities across categories (e.g., sales by region, counts). "
+            "- LINE: Use for showing trends over time or continuous data (e.g., stock prices, temperature changes)."
+        ),
+    )
+
 
     @model_validator(mode="after")
     def validate_sql_query(self):