feat(ph-ai): agent modes (#41284)

Author: skoob13

ee/hogai/context/context.py

@@ -1,6 +1,5 @@
 import asyncio
 from collections.abc import Sequence
-from functools import lru_cache
 from typing import Any, Optional, cast
 from uuid import uuid4
 
@@ -9,6 +8,7 @@
 from posthoganalytics import capture_exception
 
 from posthog.schema import (
+    AgentMode,
     ContextMessage,
     FunnelsQuery,
     HogQLQuery,
@@ -36,10 +36,13 @@
 
 from ee.hogai.graph.mixins import AssistantContextMixin
 from ee.hogai.graph.query_executor.query_executor import AssistantQueryExecutor, SupportedQueryTypes
-from ee.hogai.utils.helpers import find_start_message, insert_messages_before_start
+from ee.hogai.utils.feature_flags import has_agent_modes_feature_flag
+from ee.hogai.utils.helpers import find_start_message, find_start_message_idx, insert_messages_before_start
+from ee.hogai.utils.prompt import format_prompt_string
 from ee.hogai.utils.types.base import AnyAssistantSupportedQuery, AssistantMessageUnion, BaseStateWithMessages
 
 from .prompts import (
+    CONTEXT_MODE_PROMPT,
     CONTEXTUAL_TOOLS_REMINDER_PROMPT,
     ROOT_DASHBOARD_CONTEXT_PROMPT,
     ROOT_DASHBOARDS_CONTEXT_PROMPT,
@@ -133,7 +136,6 @@ def get_groups(self):
         """
         return GroupTypeMapping.objects.filter(project_id=self._team.project_id).order_by("group_type_index")
 
-    @lru_cache(maxsize=1)
     async def get_group_names(self) -> list[str]:
         """
         Returns the names of the team's groups.
@@ -384,23 +386,33 @@ def _render_user_context_template(
         ).to_string()
 
     async def _get_context_prompts(self, state: BaseStateWithMessages) -> list[str]:
+        are_modes_enabled = has_agent_modes_feature_flag(self._team, self._user)
+
         prompts: list[str] = []
-        if contextual_tools := self._get_contextual_tools_prompt():
+        if (
+            are_modes_enabled
+            and find_start_message_idx(state.messages, state.start_id) == 0
+            and (mode_prompt := self._get_mode_prompt(state.agent_mode))
+        ):
+            prompts.append(mode_prompt)
+        if contextual_tools := await self._get_contextual_tools_prompt():
             prompts.append(contextual_tools)
         if ui_context := await self._format_ui_context(self.get_ui_context(state)):
             prompts.append(ui_context)
         return self._deduplicate_context_messages(state, prompts)
 
-    def _get_contextual_tools_prompt(self) -> str | None:
+    async def _get_contextual_tools_prompt(self) -> str | None:
         from ee.hogai.registry import get_contextual_tool_class
 
-        contextual_tools_prompt = [
-            f"<{tool_name}>\n"
-            f"{get_contextual_tool_class(tool_name)(team=self._team, user=self._user).format_context_prompt_injection(tool_context)}\n"  # type: ignore
-            f"</{tool_name}>"
-            for tool_name, tool_context in self.get_contextual_tools().items()
-            if get_contextual_tool_class(tool_name) is not None
-        ]
+        contextual_tools_prompt: list[str] = []
+        for tool_name, tool_context in self.get_contextual_tools().items():
+            tool_class = get_contextual_tool_class(tool_name)
+            if tool_class is None:
+                continue
+            tool = await tool_class.create_tool_class(team=self._team, user=self._user, context_manager=self)
+            tool_prompt = tool.format_context_prompt_injection(tool_context)
+            contextual_tools_prompt.append(f"<{tool_name}>\n" f"{tool_prompt}\n" f"</{tool_name}>")
+
         if contextual_tools_prompt:
             tools = "\n".join(contextual_tools_prompt)
             return CONTEXTUAL_TOOLS_REMINDER_PROMPT.format(tools=tools)
@@ -417,3 +429,6 @@ def _inject_context_messages(
         context_messages = [ContextMessage(content=prompt, id=str(uuid4())) for prompt in context_prompts]
         # Insert context messages right before the start message
         return insert_messages_before_start(state.messages, context_messages, start_id=state.start_id)
+
+    def _get_mode_prompt(self, mode: AgentMode | None) -> str:
+        return format_prompt_string(CONTEXT_MODE_PROMPT, mode=mode.value if mode else AgentMode.PRODUCT_ANALYTICS.value)

ee/hogai/context/prompts.py

@@ -63,3 +63,7 @@
 IMPORTANT: this context may or may not be relevant to your tasks. You should not respond to this context unless it is highly relevant to your task.
 </system_reminder>
 """.strip()
+
+CONTEXT_MODE_PROMPT = """
+<system_reminder>Your initial mode is {{{mode}}}.</system_reminder>
+""".strip()

ee/hogai/context/test/test_context.py

@@ -8,6 +8,7 @@
 from parameterized import parameterized
 
 from posthog.schema import (
+    AgentMode,
     AssistantMessage,
     ContextMessage,
     DashboardFilter,
@@ -498,33 +499,25 @@ def test_format_entity_context_empty(self):
         result = self.context_manager._format_entity_context([], "events", "Event")
         self.assertEqual(result, "")
 
-    @patch("ee.hogai.registry.get_contextual_tool_class")
-    def test_get_contextual_tools_prompt(self, mock_get_contextual_tool_class):
+    async def test_get_contextual_tools_prompt(self):
         """Test generation of contextual tools prompt"""
-        # Mock the tool class
-        mock_tool = MagicMock()
-        mock_tool.format_context_prompt_injection.return_value = "Tool system prompt"
-        mock_get_contextual_tool_class.return_value = lambda team, user: mock_tool
-
         config = RunnableConfig(
             configurable={"contextual_tools": {"search_session_recordings": {"current_filters": {}}}}
         )
         context_manager = AssistantContextManager(self.team, self.user, config)
 
-        result = context_manager._get_contextual_tools_prompt()
-
-        self.assertIsNotNone(result)
-        assert result is not None  # Type guard for mypy
+        result = await context_manager._get_contextual_tools_prompt()
+        assert result is not None
         self.assertIn("<search_session_recordings>", result)
-        self.assertIn("Tool system prompt", result)
+        self.assertIn("Current recordings filters are", result)
         self.assertIn("</search_session_recordings>", result)
 
-    def test_get_contextual_tools_prompt_no_tools(self):
+    async def test_get_contextual_tools_prompt_no_tools(self):
         """Test generation of contextual tools prompt returns None when no tools"""
         config = RunnableConfig(configurable={})
         context_manager = AssistantContextManager(self.team, self.user, config)
 
-        result = context_manager._get_contextual_tools_prompt()
+        result = await context_manager._get_contextual_tools_prompt()
 
         self.assertIsNone(result)
 
@@ -646,3 +639,20 @@ async def test_get_billing_context(self):
 
         context_manager = AssistantContextManager(self.team, self.user, RunnableConfig(configurable={}))
         self.assertIsNone(context_manager.get_billing_context())
+
+    async def test_get_context_prompts_with_agent_mode_at_start(self):
+        """Test that mode prompt is added when feature flag is enabled and message is at start"""
+        with patch("ee.hogai.context.context.has_agent_modes_feature_flag", return_value=True):
+            state = AssistantState(
+                messages=[HumanMessage(content="Test", id="1")],
+                start_id="1",
+                agent_mode=AgentMode.PRODUCT_ANALYTICS,
+            )
+
+            result = await self.context_manager.get_state_messages_with_context(state)
+
+            assert result is not None
+            self.assertEqual(len(result), 2)
+            assert isinstance(result[0], ContextMessage)
+            self.assertIn("Your initial mode is", result[0].content)
+            self.assertIsInstance(result[1], HumanMessage)

ee/hogai/eval/ci/conftest.py

@@ -3,6 +3,7 @@
 from collections.abc import Generator
 
 import pytest
+from unittest.mock import patch
 
 from django.test import override_settings
 
@@ -33,59 +34,62 @@
 
 @pytest.fixture
 def call_root_for_insight_generation(demo_org_team_user):
-    # This graph structure will first get a plan, then generate the SQL query.
-    graph = (
-        AssistantGraph(demo_org_team_user[1], demo_org_team_user[2])
-        .add_edge(AssistantNodeName.START, AssistantNodeName.ROOT)
-        .add_root()
-        # TRICKY: We need to set a checkpointer here because async tests create a new event loop.
-        .compile(checkpointer=DjangoCheckpointer())
-    )
-
-    async def callable(query_with_extra_context: str | tuple[str, str]) -> PlanAndQueryOutput:
-        # If query_with_extra_context is a tuple, the first element is the query, the second is the extra context
-        # in case there's an ask_user tool call.
-        query = query_with_extra_context[0] if isinstance(query_with_extra_context, tuple) else query_with_extra_context
-        # Initial state for the graph
-        initial_state = AssistantState(
-            messages=[HumanMessage(content=f"Answer this question: {query}")],
+    with patch("ee.hogai.utils.feature_flags.has_agent_modes_feature_flag", return_value=True):
+        # This graph structure will first get a plan, then generate the SQL query.
+        graph = (
+            AssistantGraph(demo_org_team_user[1], demo_org_team_user[2])
+            .add_edge(AssistantNodeName.START, AssistantNodeName.ROOT)
+            .add_root()
+            # TRICKY: We need to set a checkpointer here because async tests create a new event loop.
+            .compile(checkpointer=DjangoCheckpointer())
         )
-        conversation = await Conversation.objects.acreate(team=demo_org_team_user[1], user=demo_org_team_user[2])
 
-        # Invoke the graph. The state will be updated through planner and then generator.
-        final_state_raw = await graph.ainvoke(initial_state, {"configurable": {"thread_id": conversation.id}})
+        async def callable(query_with_extra_context: str | tuple[str, str]) -> PlanAndQueryOutput:
+            # If query_with_extra_context is a tuple, the first element is the query, the second is the extra context
+            # in case there's an ask_user tool call.
+            query = (
+                query_with_extra_context[0] if isinstance(query_with_extra_context, tuple) else query_with_extra_context
+            )
+            # Initial state for the graph
+            initial_state = AssistantState(
+                messages=[HumanMessage(content=f"Answer this question: {query}")],
+            )
+            conversation = await Conversation.objects.acreate(team=demo_org_team_user[1], user=demo_org_team_user[2])
 
-        final_state = AssistantState.model_validate(final_state_raw)
+            # Invoke the graph. The state will be updated through planner and then generator.
+            final_state_raw = await graph.ainvoke(initial_state, {"configurable": {"thread_id": conversation.id}})
 
-        # If we have extra context for the potential ask_user tool, and there's no message of type ai/failure
-        # or ai/visualization, we should answer with that extra context. We only do this once at most in an eval case.
-        if isinstance(query_with_extra_context, tuple) and not any(
-            isinstance(m, VisualizationMessage | FailureMessage) for m in final_state.messages
-        ):
-            final_state.messages = [*final_state.messages, HumanMessage(content=query_with_extra_context[1])]
-            final_state.graph_status = "resumed"
-            final_state_raw = await graph.ainvoke(final_state, {"configurable": {"thread_id": conversation.id}})
             final_state = AssistantState.model_validate(final_state_raw)
 
-        # The order is a viz message, tool call message, and assistant message.
-        if (
-            not final_state.messages
-            or not len(final_state.messages) >= 3
-            or not isinstance(final_state.messages[-3], VisualizationMessage)
-        ):
+            # If we have extra context for the potential ask_user tool, and there's no message of type ai/failure
+            # or ai/visualization, we should answer with that extra context. We only do this once at most in an eval case.
+            if isinstance(query_with_extra_context, tuple) and not any(
+                isinstance(m, VisualizationMessage | FailureMessage) for m in final_state.messages
+            ):
+                final_state.messages = [*final_state.messages, HumanMessage(content=query_with_extra_context[1])]
+                final_state.graph_status = "resumed"
+                final_state_raw = await graph.ainvoke(final_state, {"configurable": {"thread_id": conversation.id}})
+                final_state = AssistantState.model_validate(final_state_raw)
+
+            # The order is a viz message, tool call message, and assistant message.
+            if (
+                not final_state.messages
+                or not len(final_state.messages) >= 3
+                or not isinstance(final_state.messages[-3], VisualizationMessage)
+            ):
+                return {
+                    "plan": None,
+                    "query": None,
+                    "query_generation_retry_count": final_state.query_generation_retry_count,
+                }
+
             return {
-                "plan": None,
-                "query": None,
+                "plan": final_state.messages[-3].plan,
+                "query": final_state.messages[-3].answer,
                 "query_generation_retry_count": final_state.query_generation_retry_count,
             }
 
-        return {
-            "plan": final_state.messages[-3].plan,
-            "query": final_state.messages[-3].answer,
-            "query_generation_retry_count": final_state.query_generation_retry_count,
-        }
-
-    return callable
+        yield callable
 
 
 @pytest.fixture(scope="package")

ee/hogai/graph/agent_executor/nodes.py

@@ -1,39 +1,29 @@
 from langchain_core.runnables import RunnableConfig
 
-from posthog.schema import AgentMode
-
 from ee.hogai.graph.agent_modes.mode_manager import AgentModeManager
 from ee.hogai.graph.base import AssistantNode
 from ee.hogai.utils.types import AssistantState, PartialAssistantState
 
 
 class AgentGraphNode(AssistantNode):
     async def arun(self, state: AssistantState, config: RunnableConfig) -> PartialAssistantState | None:
-        manager = AgentModeManager(
-            team=self._team, user=self._user, node_path=self.node_path, mode=AgentMode.PRODUCT_ANALYTICS
-        )
+        manager = AgentModeManager(team=self._team, user=self._user, node_path=self.node_path, mode=state.agent_mode)
         new_state = await manager.node(state, config)
         return new_state
 
     def router(self, state: AssistantState):
-        manager = AgentModeManager(
-            team=self._team, user=self._user, node_path=self.node_path, mode=AgentMode.PRODUCT_ANALYTICS
-        )
+        manager = AgentModeManager(team=self._team, user=self._user, node_path=self.node_path, mode=state.agent_mode)
         next_node = manager.node.router(state)
         return next_node
 
 
 class AgentGraphToolsNode(AssistantNode):
     async def arun(self, state: AssistantState, config: RunnableConfig) -> PartialAssistantState | None:
-        manager = AgentModeManager(
-            team=self._team, user=self._user, node_path=self.node_path, mode=AgentMode.PRODUCT_ANALYTICS
-        )
+        manager = AgentModeManager(team=self._team, user=self._user, node_path=self.node_path, mode=state.agent_mode)
         new_state = await manager.tools_node(state, config)
         return new_state
 
     def router(self, state: AssistantState):
-        manager = AgentModeManager(
-            team=self._team, user=self._user, node_path=self.node_path, mode=AgentMode.PRODUCT_ANALYTICS
-        )
+        manager = AgentModeManager(team=self._team, user=self._user, node_path=self.node_path, mode=state.agent_mode)
         next_node = manager.tools_node.router(state)
         return next_node

ee/hogai/graph/agent_modes/__init__.py

@@ -1,6 +1,6 @@
 from .compaction_manager import AnthropicConversationCompactionManager, ConversationCompactionManager
 from .const import SLASH_COMMAND_INIT, SLASH_COMMAND_REMEMBER
-from .factory import AgentExample, AgentModeDefinition
+from .factory import AgentModeDefinition
 from .mode_manager import AgentModeManager
 from .nodes import AgentExecutable, AgentToolkit, AgentToolsExecutable
 
@@ -10,7 +10,6 @@
     "AgentToolkit",
     "AgentModeManager",
     "AgentModeDefinition",
-    "AgentExample",
     "SLASH_COMMAND_INIT",
     "SLASH_COMMAND_REMEMBER",
     "AnthropicConversationCompactionManager",

ee/hogai/graph/agent_modes/factory.py

@@ -1,31 +1,10 @@
-from dataclasses import dataclass, field
-
-from pydantic import BaseModel
+from dataclasses import dataclass
 
 from posthog.schema import AgentMode
 
 from .nodes import AgentExecutable, AgentToolkit, AgentToolsExecutable
 
 
-class AgentExample(BaseModel):
-    """
-    Custom agent example to correct the agent's behavior through few-shot prompting.
-    The example will be formatted as follows:
-    ```
-    <example>
-    {example}
-
-    <reasoning>
-    {reasoning}
-    </reasoning>
-    </example>
-    ```
-    """
-
-    example: str
-    reasoning: str | None = None
-
-
 @dataclass
 class AgentModeDefinition:
     mode: AgentMode
@@ -38,7 +17,3 @@ class AgentModeDefinition:
     """A custom node class to use for the agent."""
     tools_node_class: type[AgentToolsExecutable] = AgentToolsExecutable
     """A custom tools node class to use for the agent."""
-    positive_todo_examples: list[AgentExample] = field(default_factory=list)
-    """Positive examples that will be injected into the `todo_write` tool. Use this field to explain the agent how it should orchestrate complex tasks using provided tools."""
-    negative_todo_examples: list[AgentExample] = field(default_factory=list)
-    """Negative examples that will be injected into the `todo_write` tool. Use this field to explain the agent how it should **NOT** orchestrate tasks using provided tools."""

ee/hogai/graph/agent_modes/mode_manager.py

@@ -6,6 +6,7 @@
 
 from posthog.models import Team, User
 
+from ee.hogai.utils.feature_flags import has_agent_modes_feature_flag
 from ee.hogai.utils.types.base import NodePath
 
 if TYPE_CHECKING:
@@ -31,7 +32,10 @@ def __init__(self, *, team: Team, user: User, node_path: tuple[NodePath, ...], m
         self._team = team
         self._user = user
         self._node_path = node_path
-        self._mode = mode or AgentMode.PRODUCT_ANALYTICS
+        if has_agent_modes_feature_flag(team, user):
+            self._mode = mode or AgentMode.PRODUCT_ANALYTICS
+        else:
+            self._mode = AgentMode.PRODUCT_ANALYTICS
 
     @property
     def node(self) -> "AgentExecutable":