Fix: Support both Pydantic models and dicts in step() and reward calculation

- Updated env.py step() to handle both dict and Action Pydantic models
- Updated reward.py calculate() to handle both dict and Pydantic models
- All 120 tests pass
- API endpoints validated and working correctly
- Supports flexible action format: dict or Action model

app/env.py

@@ -104,7 +104,7 @@ class CRMQueryEnv:
         Execute action in environment.
         
         Args:
-            action: Action dict with 'tool' and 'arguments' keys
+            action: Action dict with 'tool' and 'arguments' keys, or Action Pydantic model
         
         Returns:
             (observation, reward, done, info)
@@ -118,9 +118,14 @@ class CRMQueryEnv:
         self.step_count += 1
         task = get_task_by_id(self.current_task_id)
 
-        # Validate action
-        tool = action.get("tool", "")
-        arguments = action.get("arguments", {})
+        # Validate action - handle both dict and Pydantic models
+        if isinstance(action, dict):
+            tool = action.get("tool", "")
+            arguments = action.get("arguments", {})
+        else:
+            # Pydantic model
+            tool = action.tool
+            arguments = action.arguments
 
         # Execute action
         action_result: Dict[str, Any] = {}

app/reward.py

@@ -34,7 +34,7 @@ class RewardCalculator:
         Calculate dense reward with shape.
         
         Args:
-            action: The action taken
+            action: The action taken (dict or Pydantic model)
             action_result: Result of the action
             done: Whether episode ended
             task_ground_truth: Ground truth answer
@@ -49,8 +49,14 @@ class RewardCalculator:
         components: Dict[str, float] = {}
         base_reward = 0.0
 
-        tool = action.get("tool", "")
-        arguments = action.get("arguments", {})
+        # Handle both dict and Pydantic model
+        if isinstance(action, dict):
+            tool = action.get("tool", "")
+            arguments = action.get("arguments", {})
+        else:
+            # Pydantic model
+            tool = action.tool
+            arguments = action.arguments
 
         # Schema validation reward
         if tool in ["search_customers", "search_orders", "search_tickets", "submit_answer"]:
@@ -103,7 +109,12 @@ class RewardCalculator:
 
         # Submit answer reward
         if tool == "submit_answer":
-            submitted = action.get("arguments", {}).get("customer_ids", [])
+            # Handle both dict and Pydantic model
+            if isinstance(action, dict):
+                submitted = action.get("arguments", {}).get("customer_ids", [])
+            else:
+                submitted = action.arguments.get("customer_ids", [])
+            
             ground_truth = task_ground_truth.get("customer_ids", [])
 
             if isinstance(submitted, list) and isinstance(ground_truth, list):