Feat/math validation agent (#6012)

Co-authored-by: James Brunton <[email protected]>
Co-authored-by: EthanHealy01 <[email protected]>

engine/src/stirling/contracts/common.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 from enum import StrEnum
-from typing import Literal
+from typing import Literal, assert_never
 
 from pydantic import Field, model_validator
 
-from stirling.models import OPERATIONS, ApiModel, OperationId, ParamToolModel
+from stirling.models import OPERATIONS, ApiModel, OperationId
+from stirling.models.agent_tool_models import AGENT_OPERATIONS, AgentToolId, AnyParamModel, AnyToolId
 
 
 class PdfContentType(StrEnum):
@@ -82,6 +83,7 @@ class SupportedCapability(StrEnum):
     AGENT_DRAFT = "agent_draft"
     AGENT_REVISE = "agent_revise"
     AGENT_NEXT_ACTION = "agent_next_action"
+    MATH_AUDITOR_AGENT = "math_auditor_agent"
 
 
 class ConversationMessage(ApiModel):
@@ -101,14 +103,19 @@ class ExtractedFileText(ApiModel):
 
 class ToolOperationStep(ApiModel):
     kind: Literal[StepKind.TOOL] = StepKind.TOOL
-    tool: OperationId
-    parameters: ParamToolModel
+    tool: AnyToolId
+    parameters: AnyParamModel
 
     @model_validator(mode="after")
     def validate_tool_parameter_pairing(self) -> ToolOperationStep:
-        expected_type = OPERATIONS[self.tool]
+        if isinstance(self.tool, AgentToolId):
+            expected_type = AGENT_OPERATIONS[self.tool]
+        elif isinstance(self.tool, OperationId):
+            expected_type = OPERATIONS[self.tool]
+        else:
+            assert_never(self.tool)
+
         if not isinstance(self.parameters, expected_type):
             actual_type = type(self.parameters).__name__
-            expected_type_name = expected_type.__name__
-            raise ValueError(f"Parameters for tool {self.tool.value} must be {expected_type_name}, got {actual_type}.")
+            raise ValueError(f"Parameters for tool {self.tool} must be {expected_type.__name__}, got {actual_type}.")
         return self