Update utils.py

utils.py

@@ -1,14 +1,25 @@
-from pydantic import BaseModel, Field
-from typing import Optional, List, Dict
-
-class SQLGenerationResult(BaseModel):
-    sql_query: str = Field(..., description="SQL query to execute")
-    reasoning: Optional[str] = Field(None, description="Optional explanation of the query")
-
-class AgentState(BaseModel):
-    question: str
-    sql_query: Optional[str] = None
-    reasoning: Optional[str] = None     
-    rows: Optional[List[Dict]] = None
-    answer: Optional[str] = None
-    error: Optional[str] = None
+from pydantic import BaseModel, Field
+from typing import Optional, List, Dict, Literal
+
+class SQLGenerationResult(BaseModel):
+    sql_query: str = Field(..., description="SQL query to execute")
+    reasoning: Optional[str] = Field(None, description="Optional explanation of the query")
+
+class AgentState(BaseModel):
+    question: str
+    sql_query: Optional[str] = None
+    reasoning: Optional[str] = None     
+    rows: Optional[List[Dict]] = None
+    answer: Optional[str] = None
+    error: Optional[str] = None
+    route: Optional[Literal["sql", "chat"]] = None
+
+class RouteDecision(BaseModel):
+    route: Literal["sql", "chat"] = Field(
+        ...,
+        description="Routing decision. Use 'sql' when the user's question requires querying the database. Use 'chat' for greetings, small talk, or general questions."
+    )
+    answer: Optional[str] = Field(
+        default=None,
+        description="Plain-text answer to return when route='chat'. Leave null when route='sql'."
+    )