feat: Add required input handling to agent definitions and update UI components

Signed-off-by: Andre Bossard <anbossar@microsoft.com>

Author: abossard

backend/agent_workbench/models.py

@@ -68,6 +68,14 @@ class AgentDefinition(SQLModel, table=True):
     name: str = SField(index=True, description="Human-readable agent name")
     description: str = SField(default="", description="Optional description")
     system_prompt: str = SField(description="System prompt sent to the LLM")
+    requires_input: bool = SField(
+        default=False,
+        description="When true, runs must include required_input_value",
+    )
+    required_input_description: str = SField(
+        default="",
+        description="Description shown to operators for required runtime input",
+    )
     # Stored as JSON arrays in TEXT columns
     tool_names_json: str = SField(
         default="[]",
@@ -115,6 +123,8 @@ def to_dict(self) -> dict[str, Any]:
             "name": self.name,
             "description": self.description,
             "system_prompt": self.system_prompt,
+            "requires_input": self.requires_input,
+            "required_input_description": self.required_input_description,
             "tool_names": self.tool_names,
             "success_criteria": [c.model_dump() for c in self.success_criteria],
             "created_at": self.created_at.isoformat(),
@@ -228,6 +238,8 @@ class AgentDefinitionCreate(BaseModel):
     name: str = Field(..., min_length=1, max_length=200)
     description: str = Field(default="")
     system_prompt: str = Field(..., min_length=1)
+    requires_input: bool = Field(default=False)
+    required_input_description: str = Field(default="")
     tool_names: list[str] = Field(default_factory=list)
     success_criteria: list[SuccessCriteria] = Field(default_factory=list)
 
@@ -236,12 +248,15 @@ class AgentDefinitionUpdate(BaseModel):
     name: Optional[str] = Field(default=None)
     description: Optional[str] = Field(default=None)
     system_prompt: Optional[str] = Field(default=None)
+    requires_input: Optional[bool] = Field(default=None)
+    required_input_description: Optional[str] = Field(default=None)
     tool_names: Optional[list[str]] = Field(default=None)
     success_criteria: Optional[list[SuccessCriteria]] = Field(default=None)
 
 
 class AgentRunCreate(BaseModel):
-    input_prompt: str = Field(..., min_length=1, max_length=10000)
+    input_prompt: str = Field(default="", max_length=10000)
+    required_input_value: Optional[str] = Field(default=None, max_length=2000)
 
 
 # ============================================================================
@@ -258,6 +273,18 @@ def build_engine(db_path: Path):
 
 def _run_migrations(engine) -> None:
     """Apply lightweight SQLite migrations for new columns."""
+    _ensure_column(
+        engine,
+        "workbench_agent_definitions",
+        "requires_input",
+        "BOOLEAN NOT NULL DEFAULT 0",
+    )
+    _ensure_column(
+        engine,
+        "workbench_agent_definitions",
+        "required_input_description",
+        "TEXT NOT NULL DEFAULT ''",
+    )
     _ensure_column(
         engine,
         "workbench_agent_runs",

backend/agent_workbench/service.py

@@ -158,11 +158,56 @@ def _build_agent_snapshot(self, agent: AgentDefinition) -> dict[str, Any]:
             "name": agent.name,
             "description": agent.description,
             "system_prompt": agent.system_prompt,
+            "requires_input": agent.requires_input,
+            "required_input_description": agent.required_input_description,
             "tool_names": list(agent.tool_names),
             "success_criteria": [criteria.model_dump() for criteria in agent.success_criteria],
             "captured_at": datetime.now().isoformat(),
         }
 
+    def _normalize_input_contract(
+        self,
+        requires_input: bool,
+        required_input_description: str,
+    ) -> tuple[bool, str]:
+        normalized_description = (required_input_description or "").strip()
+        if requires_input and not normalized_description:
+            raise ValueError(
+                "required_input_description must be provided when requires_input is true"
+            )
+        if not requires_input:
+            normalized_description = ""
+        return requires_input, normalized_description
+
+    def _build_run_user_message(
+        self,
+        agent_def: AgentDefinition,
+        run_request: AgentRunCreate,
+    ) -> tuple[str, str]:
+        run_prompt = (run_request.input_prompt or "").strip()
+        required_input_value = (run_request.required_input_value or "").strip()
+        message_parts: list[str] = []
+
+        if run_prompt:
+            message_parts.append(run_prompt)
+
+        if agent_def.requires_input:
+            if not required_input_value:
+                raise ValueError(
+                    "Missing required_input_value for this agent. "
+                    f"Expected: {agent_def.required_input_description}"
+                )
+            message_parts.append(
+                f"Required input ({agent_def.required_input_description}): {required_input_value}"
+            )
+        elif required_input_value:
+            message_parts.append(f"Additional input: {required_input_value}")
+
+        if not message_parts:
+            message_parts.append("Proceed with the configured system instructions and tools.")
+
+        return "\n\n".join(message_parts), required_input_value
+
     def _criteria_from_run_snapshot(self, run: AgentRun) -> list[SuccessCriteria]:
         snapshot = run.agent_snapshot
         raw = snapshot.get("success_criteria")
@@ -184,10 +229,16 @@ def _criteria_from_run_snapshot(self, run: AgentRun) -> list[SuccessCriteria]:
 
     def create_agent(self, data: AgentDefinitionCreate) -> AgentDefinition:
         validated_tool_names = self._validate_tool_names(data.tool_names)
+        requires_input, required_input_description = self._normalize_input_contract(
+            data.requires_input,
+            data.required_input_description,
+        )
         agent = AgentDefinition(
             name=data.name,
             description=data.description,
             system_prompt=data.system_prompt,
+            requires_input=requires_input,
+            required_input_description=required_input_description,
         )
         agent.tool_names = validated_tool_names
         agent.success_criteria = data.success_criteria
@@ -218,6 +269,19 @@ def update_agent(
                 agent.description = data.description
             if data.system_prompt is not None:
                 agent.system_prompt = data.system_prompt
+            next_requires_input = agent.requires_input if data.requires_input is None else data.requires_input
+            next_required_input_description = (
+                agent.required_input_description
+                if data.required_input_description is None
+                else data.required_input_description
+            )
+            (
+                agent.requires_input,
+                agent.required_input_description,
+            ) = self._normalize_input_contract(
+                next_requires_input,
+                next_required_input_description,
+            )
             if data.tool_names is not None:
                 agent.tool_names = self._validate_tool_names(data.tool_names)
             if data.success_criteria is not None:
@@ -279,11 +343,16 @@ async def run_agent(
 
         validated_tool_names = self._validate_tool_names(agent_def.tool_names)
         agent_snapshot = self._build_agent_snapshot(agent_def)
+        user_message, normalized_required_input = self._build_run_user_message(agent_def, run_request)
+        normalized_prompt = (run_request.input_prompt or "").strip()
+        agent_snapshot["input_prompt"] = normalized_prompt
+        agent_snapshot["required_input_value"] = normalized_required_input
+        agent_snapshot["composed_user_message"] = user_message
 
         # -- Persist a PENDING run --
         run = AgentRun(
             agent_id=agent_id,
-            input_prompt=run_request.input_prompt,
+            input_prompt=normalized_prompt,
             status=RunStatus.RUNNING.value,
         )
         run.agent_snapshot = agent_snapshot
@@ -300,7 +369,7 @@ async def run_agent(
             react = _build_react_agent(self.llm, tools, agent_def.system_prompt)
 
             result = await react.ainvoke(
-                {"messages": [("user", run_request.input_prompt)]},
+                {"messages": [("user", user_message)]},
                 config={"recursion_limit": self._recursion_limit},
             )
 

backend/tests/test_workbench_integration_e2e.py

@@ -30,6 +30,13 @@ def __init__(self, tools: list[object]) -> None:
 
     async def ainvoke(self, _payload: dict, config: dict | None = None) -> dict:
         _ = config
+        user_message = ""
+        messages = _payload.get("messages", [])
+        if messages:
+            first = messages[0]
+            if isinstance(first, (list, tuple)) and len(first) >= 2:
+                user_message = str(first[1])
+
         tool = next(
             (item for item in self._tools if getattr(item, "name", "") == "csv_ticket_stats"),
             None,
@@ -39,7 +46,7 @@ async def ainvoke(self, _payload: dict, config: dict | None = None) -> dict:
 
         stats = await tool.ainvoke({})
         total = stats.get("total", 0) if isinstance(stats, dict) else 0
-        output = f"Used csv_ticket_stats successfully. total={total}"
+        output = f"Used csv_ticket_stats successfully. total={total}. context={user_message}"
         return {
             "messages": [
                 _ToolCallMessage("csv_ticket_stats"),
@@ -141,6 +148,67 @@ async def test_create_run_and_evaluate_agent_with_csv_tool(self) -> None:
                 self.assertTrue(evaluate_data["overall_passed"])
                 self.assertEqual(evaluate_data["score"], 1.0)
 
+    async def test_required_input_agent_run_validation_and_context(self) -> None:
+        with patch("agent_workbench.service._build_react_agent", new=_fake_build_react_agent):
+            async with backend_app_module.app.test_app() as test_app:
+                client = test_app.test_client()
+
+                invalid_create_resp = await client.post(
+                    "/api/workbench/agents",
+                    json={
+                        "name": "Needs Input Invalid",
+                        "description": "",
+                        "system_prompt": "Use CSV tools.",
+                        "requires_input": True,
+                        "required_input_description": "",
+                        "tool_names": ["csv_ticket_stats"],
+                        "success_criteria": [],
+                    },
+                )
+                self.assertEqual(invalid_create_resp.status_code, 400)
+
+                create_resp = await client.post(
+                    "/api/workbench/agents",
+                    json={
+                        "name": "Needs Input",
+                        "description": "",
+                        "system_prompt": "Use CSV tools.",
+                        "requires_input": True,
+                        "required_input_description": "Ticket INC number",
+                        "tool_names": ["csv_ticket_stats"],
+                        "success_criteria": [],
+                    },
+                )
+                create_data = await create_resp.get_json()
+                self.assertEqual(create_resp.status_code, 201, create_data)
+                self.assertTrue(create_data["requires_input"])
+                self.assertEqual(create_data["required_input_description"], "Ticket INC number")
+                agent_id = create_data["id"]
+
+                missing_input_resp = await client.post(
+                    f"/api/workbench/agents/{agent_id}/runs",
+                    json={"input_prompt": ""},
+                )
+                self.assertEqual(missing_input_resp.status_code, 400)
+
+                run_resp = await client.post(
+                    f"/api/workbench/agents/{agent_id}/runs",
+                    json={"required_input_value": "INC-12345"},
+                )
+                run_data = await run_resp.get_json()
+                self.assertEqual(run_resp.status_code, 200, run_data)
+                self.assertEqual(run_data["status"], "completed")
+                self.assertEqual(run_data["input_prompt"], "")
+                self.assertIn("INC-12345", run_data["output"] or "")
+                self.assertEqual(
+                    run_data["agent_snapshot"].get("required_input_value"),
+                    "INC-12345",
+                )
+                self.assertIn(
+                    "Required input (Ticket INC number): INC-12345",
+                    run_data["agent_snapshot"].get("composed_user_message", ""),
+                )
+
 
 if __name__ == "__main__":
     unittest.main()