feat(ph-ai): create form tool (#42085)

## Problem

PostHog AI needs a way to collect structured information from users through multi-question forms. This will be particularly helpful in plan/research mode.

![image.png](https://app.graphite.com/user-attachments/assets/b0ade642-331f-4acb-8bc2-67cf635cd171.png)

![image.png](https://app.graphite.com/user-attachments/assets/7a1edf23-97ce-4358-a86f-e05288fc0057.png)

![image.png](https://app.graphite.com/user-attachments/assets/85697675-61d7-4115-9bc6-15cc6577a655.png)

## Changes

- Added a new `create_form` tool that allows PostHog AI to create multi-question forms
- Implemented form rendering in the frontend with support for:
    - Multiple questions with predefined options
    - Optional custom answer input
- Form answers are sent as UI payload
- Added feature flag `phai-create-form-tool` to control availability
- Runner can now handle form responses and resume the conversation (answers are added to the `create_form`​ tool message response and sent back as `ui_payload`​, so they can populate the frontend components on reload)

## How did you test this code?

New stories, new tests

Author: kappa90

ee/hogai/chat_agent/mode_manager.py

@@ -28,7 +28,6 @@
 )
 from ee.hogai.context import AssistantContextManager
 from ee.hogai.core.agent_modes.factory import AgentModeDefinition
-from ee.hogai.core.agent_modes.feature_flags import has_agent_modes_feature_flag
 from ee.hogai.core.agent_modes.mode_manager import AgentModeManager
 from ee.hogai.core.agent_modes.presets.product_analytics import product_analytics_agent
 from ee.hogai.core.agent_modes.presets.session_replay import session_replay_agent
@@ -39,7 +38,8 @@
 from ee.hogai.core.shared_prompts import CORE_MEMORY_PROMPT
 from ee.hogai.registry import get_contextual_tool_class
 from ee.hogai.tool import MaxTool
-from ee.hogai.tools import ReadDataTool, ReadTaxonomyTool, SearchTool, SwitchModeTool, TodoWriteTool
+from ee.hogai.tools import CreateFormTool, ReadDataTool, ReadTaxonomyTool, SearchTool, SwitchModeTool, TodoWriteTool
+from ee.hogai.utils.feature_flags import has_agent_modes_feature_flag, has_create_form_tool_feature_flag
 from ee.hogai.utils.prompt import format_prompt_string
 from ee.hogai.utils.types.base import AssistantState, NodePath
 
@@ -63,7 +63,10 @@
 class ChatAgentToolkit(AgentToolkit):
     @property
     def tools(self) -> list[type["MaxTool"]]:
-        return DEFAULT_TOOLS if has_agent_modes_feature_flag(self._team, self._user) else LEGACY_DEFAULT_TOOLS
+        tools = list(DEFAULT_TOOLS if has_agent_modes_feature_flag(self._team, self._user) else LEGACY_DEFAULT_TOOLS)
+        if has_create_form_tool_feature_flag(self._team, self._user):
+            tools.append(CreateFormTool)
+        return tools
 
 
 class ChatAgentToolkitManager(AgentToolkitManager):

ee/hogai/core/agent_modes/executables.py

@@ -13,6 +13,7 @@
     ToolMessage as LangchainToolMessage,
 )
 from langchain_core.runnables import RunnableConfig
+from langgraph.errors import NodeInterrupt
 from langgraph.types import Send
 from posthoganalytics import capture_exception
 from pydantic import ValidationError
@@ -29,7 +30,6 @@
 from posthog.event_usage import groups
 from posthog.models import Team, User
 
-from ee.hogai.core.agent_modes.feature_flags import has_agent_modes_feature_flag
 from ee.hogai.core.agent_modes.prompt_builder import AgentPromptBuilder
 from ee.hogai.core.agent_modes.prompts import (
     ROOT_CONVERSATION_SUMMARY_PROMPT,
@@ -43,6 +43,7 @@
 from ee.hogai.tool_errors import MaxToolError
 from ee.hogai.utils.anthropic import add_cache_control, convert_to_anthropic_messages
 from ee.hogai.utils.conversation_summarizer import AnthropicConversationSummarizer
+from ee.hogai.utils.feature_flags import has_agent_modes_feature_flag
 from ee.hogai.utils.helpers import convert_tool_messages_to_dict, normalize_ai_message
 from ee.hogai.utils.types import (
     AssistantMessageUnion,
@@ -419,6 +420,9 @@ async def arun(self, state: AssistantState, config: RunnableConfig) -> PartialAs
                     )
                 ],
             )
+        except NodeInterrupt:
+            # Let NodeInterrupt propagate to the graph engine for tool interrupts
+            raise
         except Exception as e:
             logger.exception("Error calling tool", extra={"tool_name": tool_call.name, "error": str(e)})
             capture_exception(

ee/hogai/core/agent_modes/feature_flags.py

@@ -9,10 +9,9 @@
 from ee.hogai.core.agent_modes.prompt_builder import AgentPromptBuilder
 from ee.hogai.core.agent_modes.toolkit import AgentToolkit, AgentToolkitManager
 from ee.hogai.core.mixins import AssistantContextMixin
+from ee.hogai.utils.feature_flags import has_agent_modes_feature_flag
 from ee.hogai.utils.types.base import NodePath
 
-from .feature_flags import has_agent_modes_feature_flag
-
 if TYPE_CHECKING:
     from .executables import AgentExecutable, AgentToolsExecutable
     from .factory import AgentModeDefinition

ee/hogai/core/agent_modes/mode_manager.py

@@ -19,6 +19,7 @@
     AssistantEventType,
     AssistantGenerationStatusEvent,
     AssistantMessage,
+    AssistantToolCallMessage,
     AssistantUpdateEvent,
     FailureMessage,
     HumanMessage,
@@ -36,7 +37,7 @@
 from ee.hogai.core.stream_processor import AssistantStreamProcessorProtocol
 from ee.hogai.utils.exceptions import LLM_API_EXCEPTIONS, LLM_PROVIDER_ERROR_COUNTER, GenerationCanceled
 from ee.hogai.utils.feature_flags import is_privacy_mode_enabled
-from ee.hogai.utils.helpers import extract_stream_update
+from ee.hogai.utils.helpers import extract_stream_update, find_last_message_of_type
 from ee.hogai.utils.state import validate_state_update
 from ee.hogai.utils.types.base import (
     AssistantDispatcherEvent,
@@ -62,7 +63,7 @@ class BaseAgentRunner(ABC):
     _contextual_tools: dict[str, Any]
     _conversation: Conversation
     _session_id: Optional[str]
-    _latest_message: Optional[HumanMessage]
+    _latest_message: Optional[HumanMessage | AssistantToolCallMessage]
     _state: Optional[AssistantMaxGraphState]
     _callback_handlers: list[BaseCallbackHandler]
     _trace_id: Optional[str | UUID]
@@ -218,6 +219,8 @@ async def astream(
                     interrupt_messages = []
                     for task in state.tasks:
                         for interrupt in task.interrupts:
+                            if interrupt.value is None:
+                                continue  # Skip None interrupts (used by create_form)
                             interrupt_message = (
                                 AssistantMessage(content=interrupt.value, id=str(uuid4()))
                                 if isinstance(interrupt.value, str)
@@ -309,6 +312,10 @@ async def _init_or_update_state(self):
         saved_state = validate_state_update(snapshot.values, self._state_type)
         last_recorded_dt = saved_state.start_dt
 
+        # When resuming after a create_form interrupt, create the tool call response message
+        if form_response_message := self._get_form_response_message(saved_state):
+            self._latest_message = form_response_message
+
         # Add existing ids to streamed messages, so we don't send the messages again.
         for message in saved_state.messages:
             if message.id is not None:
@@ -406,3 +413,41 @@ def _capture_exception(self, e: Exception):
                 "$groups": event_usage.groups(team=self._team),
             },
         )
+
+    def _get_form_response_message(self, saved_state: AssistantMaxGraphState) -> AssistantToolCallMessage | None:
+        """
+        When resuming after a create_form tool call (which raises NodeInterrupt(None)),
+        create an AssistantToolCallMessage with the user's response content and parsed answers in ui_payload.
+        """
+        if not saved_state.messages or not self._latest_message:
+            return None
+
+        # Form responses must come from a HumanMessage
+        if not isinstance(self._latest_message, HumanMessage):
+            return None
+
+        # Check if we have form answers in the ui_context
+        if not self._latest_message.ui_context or not self._latest_message.ui_context.form_answers:
+            return None
+
+        # Find the last assistant message with tool calls
+        last_assistant_message = find_last_message_of_type(saved_state.messages, AssistantMessage)
+        if not last_assistant_message or not last_assistant_message.tool_calls:
+            return None
+
+        # Find the create_form tool call
+        create_form_tool_call = next(
+            (tc for tc in last_assistant_message.tool_calls if tc.name == "create_form"),
+            None,
+        )
+        if not create_form_tool_call:
+            return None
+
+        answers = self._latest_message.ui_context.form_answers
+
+        return AssistantToolCallMessage(
+            content=self._latest_message.content or "",
+            id=str(uuid4()),
+            tool_call_id=create_form_tool_call.id,
+            ui_payload={"create_form": {"answers": answers}},
+        )

ee/hogai/core/runner.py

@@ -0,0 +1,186 @@
+from posthog.test.base import BaseTest
+
+from posthog.schema import AssistantMessage, AssistantToolCall, AssistantToolCallMessage, HumanMessage, MaxUIContext
+
+from ee.hogai.core.runner import BaseAgentRunner
+from ee.hogai.utils.types import AssistantState
+
+
+class MockAgentRunner(BaseAgentRunner):
+    """Minimal concrete implementation for testing BaseAgentRunner methods."""
+
+    def __init__(self, team, user, latest_message=None):
+        self._team = team
+        self._user = user
+        self._latest_message = latest_message
+
+    def get_initial_state(self) -> AssistantState:
+        return AssistantState(messages=[])
+
+    def get_resumed_state(self):
+        return {}
+
+
+class TestGetFormResponseMessage(BaseTest):
+    def setUp(self):
+        super().setUp()
+
+    def _create_runner_with_message(self, latest_message):
+        return MockAgentRunner(team=self.team, user=self.user, latest_message=latest_message)
+
+    def test_returns_none_when_no_messages(self):
+        runner = self._create_runner_with_message(HumanMessage(content="Test"))
+        saved_state = AssistantState(messages=[])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)
+
+    def test_returns_none_when_no_latest_message(self):
+        runner = self._create_runner_with_message(None)
+        saved_state = AssistantState(messages=[AssistantToolCallMessage(content="test", tool_call_id="tc1")])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)
+
+    def test_returns_none_when_latest_message_is_not_human_message(self):
+        runner = self._create_runner_with_message(
+            AssistantToolCallMessage(content="Tool response", tool_call_id="some-id")
+        )
+        assistant_message = AssistantMessage(
+            content="Please answer:",
+            tool_calls=[
+                AssistantToolCall(
+                    id="create_form_tc_1",
+                    name="create_form",
+                    args={"questions": [{"id": "q1", "question": "Question 1"}]},
+                    type="tool_call",
+                )
+            ],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)
+
+    def test_returns_none_when_no_form_answers_in_ui_context(self):
+        runner = self._create_runner_with_message(HumanMessage(content="My answer"))
+        assistant_message = AssistantMessage(
+            content="test",
+            tool_calls=[AssistantToolCall(id="tc1", name="create_form", args={}, type="tool_call")],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)
+
+    def test_returns_none_when_no_create_form_tool_call(self):
+        runner = self._create_runner_with_message(
+            HumanMessage(content="My answer", ui_context=MaxUIContext(form_answers={"q1": "answer"}))
+        )
+        assistant_message = AssistantMessage(
+            content="test",
+            tool_calls=[AssistantToolCall(id="tc1", name="other_tool", args={}, type="tool_call")],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)
+
+    def test_creates_tool_call_message_with_form_answers(self):
+        user_content = "What is your name: John\nWhat is your role: Engineer"
+        message = HumanMessage(
+            content=user_content, ui_context=MaxUIContext(form_answers={"name": "John", "role": "Engineer"})
+        )
+        runner = self._create_runner_with_message(message)
+
+        tool_call_id = "create_form_tc_1"
+        assistant_message = AssistantMessage(
+            content="Please answer these questions:",
+            tool_calls=[
+                AssistantToolCall(
+                    id=tool_call_id,
+                    name="create_form",
+                    args={
+                        "questions": [
+                            {"id": "name", "question": "What is your name"},
+                            {"id": "role", "question": "What is your role"},
+                        ]
+                    },
+                    type="tool_call",
+                )
+            ],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNotNone(result)
+        self.assertIsInstance(result, AssistantToolCallMessage)
+        assert isinstance(result, AssistantToolCallMessage)
+        self.assertEqual(result.content, user_content)
+        self.assertEqual(result.tool_call_id, tool_call_id)
+        self.assertIsNotNone(result.ui_payload)
+        assert result.ui_payload is not None
+        self.assertIn("create_form", result.ui_payload)
+        self.assertEqual(
+            result.ui_payload["create_form"]["answers"],
+            {"name": "John", "role": "Engineer"},
+        )
+
+    def test_builds_correct_ui_payload_structure(self):
+        user_content = "Pick a color: Blue"
+        message = HumanMessage(content=user_content, ui_context=MaxUIContext(form_answers={"color": "Blue"}))
+        runner = self._create_runner_with_message(message)
+
+        tool_call_id = "form_tc_123"
+        assistant_message = AssistantMessage(
+            content="Choose your favorite color:",
+            tool_calls=[
+                AssistantToolCall(
+                    id=tool_call_id,
+                    name="create_form",
+                    args={
+                        "questions": [
+                            {"id": "color", "question": "Pick a color"},
+                        ]
+                    },
+                    type="tool_call",
+                )
+            ],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNotNone(result)
+        assert isinstance(result, AssistantToolCallMessage)
+        assert result.ui_payload is not None
+        self.assertEqual(result.ui_payload, {"create_form": {"answers": {"color": "Blue"}}})
+
+    def test_returns_none_with_empty_form_answers_dict(self):
+        user_content = "What is your name: John"
+        ui_context = MaxUIContext(form_answers={})  # Empty dict
+        runner = self._create_runner_with_message(HumanMessage(content=user_content, ui_context=ui_context))
+
+        tool_call_id = "create_form_tc_1"
+        assistant_message = AssistantMessage(
+            content="Please answer:",
+            tool_calls=[
+                AssistantToolCall(
+                    id=tool_call_id,
+                    name="create_form",
+                    args={"questions": [{"id": "name", "question": "What is your name"}]},
+                    type="tool_call",
+                )
+            ],
+        )
+        saved_state = AssistantState(messages=[assistant_message])
+
+        result = runner._get_form_response_message(saved_state)
+
+        self.assertIsNone(result)

ee/hogai/core/test/test_form_response.py

@@ -1,5 +1,6 @@
 from .create_and_query_insight import CreateAndQueryInsightTool, CreateAndQueryInsightToolArgs
 from .create_dashboard import CreateDashboardTool, CreateDashboardToolArgs
+from .create_form import CreateFormTool, CreateFormToolArgs
 from .create_insight import CreateInsightTool, CreateInsightToolArgs
 from .execute_sql.tool import ExecuteSQLTool, ExecuteSQLToolArgs
 from .read_data import ReadDataTool, ReadDataToolArgs
@@ -14,6 +15,8 @@
     "CreateAndQueryInsightToolArgs",
     "CreateDashboardTool",
     "CreateDashboardToolArgs",
+    "CreateFormTool",
+    "CreateFormToolArgs",
     "ReadDataTool",
     "ReadDataToolArgs",
     "ReadTaxonomyTool",