feat: replay mode

Author: skoob13

ee/hogai/graph/agent_modes/presets/__init__.py

@@ -1,9 +1,12 @@
 from posthog.schema import AgentMode
 
+from ee.hogai.graph.agent_modes.presets.session_replay import session_replay_agent
+
 from .product_analytics import product_analytics_agent
 from .sql import sql_agent
 
 MODE_REGISTRY = {
     AgentMode.PRODUCT_ANALYTICS: product_analytics_agent,
     AgentMode.SQL: sql_agent,
+    AgentMode.SESSION_REPLAY: session_replay_agent,
 }

ee/hogai/graph/agent_modes/presets/product_analytics.py

@@ -4,7 +4,7 @@
 
 from posthog.schema import AgentMode
 
-from ee.hogai.tools import CreateAndQueryInsightTool, CreateDashboardTool, SessionSummarizationTool
+from ee.hogai.tools import CreateAndQueryInsightTool, CreateDashboardTool
 
 from ..factory import AgentDefinition
 from ..nodes import AgentToolkit
@@ -22,10 +22,6 @@ def custom_tools(self) -> list[type["MaxTool"]]:
         if not CreateAndQueryInsightTool.is_editing_mode(self._context_manager):
             tools.append(CreateAndQueryInsightTool)
 
-        # Add session summarization tool if enabled
-        if self._has_session_summarization_feature_flag():
-            tools.append(SessionSummarizationTool)
-
         # Add other lower-priority tools
         tools.append(CreateDashboardTool)
 

ee/hogai/graph/agent_modes/presets/session_replay.py

@@ -0,0 +1,44 @@
+from typing import TYPE_CHECKING
+
+import posthoganalytics
+
+from posthog.schema import AgentMode
+
+from ee.hogai.tools.replay.summarize_sessions import SummarizeSessionsTool
+
+from ..factory import AgentDefinition
+from ..nodes import AgentToolkit
+
+if TYPE_CHECKING:
+    from ee.hogai.tool import MaxTool
+
+
+class SessionReplayAgentToolkit(AgentToolkit):
+    @property
+    def custom_tools(self) -> list[type["MaxTool"]]:
+        tools: list[type[MaxTool]] = []
+
+        # Add session summarization tool if enabled
+        if self._has_session_summarization_feature_flag():
+            tools.append(SummarizeSessionsTool)
+
+        return tools
+
+    def _has_session_summarization_feature_flag(self) -> bool:
+        """
+        Check if the user has the session summarization feature flag enabled.
+        """
+        return posthoganalytics.feature_enabled(
+            "max-session-summarization",
+            str(self._user.distinct_id),
+            groups={"organization": str(self._team.organization_id)},
+            group_properties={"organization": {"id": str(self._team.organization_id)}},
+            send_feature_flag_events=False,
+        )
+
+
+session_replay_agent = AgentDefinition(
+    mode=AgentMode.SESSION_REPLAY,
+    mode_description="Specialized mode for analyzing session recordings and user behavior. This mode allows you to get summaries of session recordings and insights about them in natural language.",
+    toolkit_class=SessionReplayAgentToolkit,
+)

ee/hogai/mode_registry.py

@@ -1,9 +1,11 @@
 from posthog.schema import AgentMode
 
 from ee.hogai.graph.agent_modes.presets.product_analytics import product_analytics_agent
+from ee.hogai.graph.agent_modes.presets.session_replay import session_replay_agent
 from ee.hogai.graph.agent_modes.presets.sql import sql_agent
 
 MODE_REGISTRY = {
     AgentMode.PRODUCT_ANALYTICS: product_analytics_agent,
     AgentMode.SQL: sql_agent,
+    AgentMode.SESSION_REPLAY: session_replay_agent,
 }

ee/hogai/tools/create_and_query_insight.py

@@ -1,15 +1,18 @@
-from typing import Literal
+from typing import Literal, Self
 
+from langchain_core.runnables import RunnableConfig
 from pydantic import BaseModel, Field
 
 from posthog.schema import AssistantTool, AssistantToolCallMessage, VisualizationMessage
 
+from posthog.models import Team, User
+
 from ee.hogai.context.context import AssistantContextManager
 from ee.hogai.graph.insights_graph.graph import InsightsGraph
 from ee.hogai.graph.schema_generator.nodes import SchemaGenerationException
 from ee.hogai.tool import MaxTool, ToolMessagesArtifact
 from ee.hogai.utils.prompt import format_prompt_string
-from ee.hogai.utils.types.base import AssistantNodeName, AssistantState
+from ee.hogai.utils.types.base import AssistantNodeName, AssistantState, NodePath
 
 INSIGHT_TOOL_PROMPT = """
 Use this tool to generate an insight from a structured plan. It will return a visualization that the user will be able to analyze and textual representation for your analysis.
@@ -516,9 +519,25 @@ class CreateAndQueryInsightToolArgs(BaseModel):
 class CreateAndQueryInsightTool(MaxTool):
     name: Literal["create_and_query_insight"] = "create_and_query_insight"
     args_schema: type[BaseModel] = CreateAndQueryInsightToolArgs
-    description: str = INSIGHT_TOOL_PROMPT
     context_prompt_template: str = INSIGHT_TOOL_CONTEXT_PROMPT_TEMPLATE
-    thinking_message: str = "Coming up with an insight"
+
+    @classmethod
+    async def create_tool_class(
+        cls,
+        *,
+        team: Team,
+        user: User,
+        node_path: tuple[NodePath, ...] | None = None,
+        state: AssistantState | None = None,
+        config: RunnableConfig | None = None,
+        context_manager: AssistantContextManager | None = None,
+    ) -> Self:
+        context_manager = context_manager or AssistantContextManager(team, user, config)
+        prompt = format_prompt_string(
+            INSIGHT_TOOL_PROMPT,
+            groups=await context_manager.get_group_names(),
+        )
+        return cls(team=team, user=user, state=state, node_path=node_path, config=config, description=prompt)
 
     async def _arun_impl(
         self, query_description: str, insight_type: InsightType