refactor(ph-ai): clean up scaffolding

Author: kappa90

ee/hogai/README.md

@@ -234,8 +234,8 @@ class MaxToolTaxonomyOutput(BaseModel):
 
 ```python
 from pydantic import BaseModel, Field
-from ee.hogai.graph.taxonomy.toolkit import TaxonomyAgentToolkit
-from ee.hogai.graph.taxonomy.tools import base_final_answer
+from ee.hogai.chat_agent.taxonomy.toolkit import TaxonomyAgentToolkit
+from ee.hogai.chat_agent.taxonomy.tools import base_final_answer
 from posthog.models import Team
 
 
@@ -276,9 +276,9 @@ class YourToolkit(TaxonomyAgentToolkit):
 ```python
 from langchain_core.prompts import ChatPromptTemplate
 from posthog.models import Team, User
-from ee.hogai.graph.taxonomy.nodes import TaxonomyAgentNode, TaxonomyAgentToolsNode
-from ee.hogai.graph.taxonomy.agent import TaxonomyAgent
-from ee.hogai.graph.taxonomy.types import TaxonomyAgentState
+from ee.hogai.chat_agent.taxonomy.nodes import TaxonomyAgentNode, TaxonomyAgentToolsNode
+from ee.hogai.chat_agent.taxonomy.agent import TaxonomyAgent
+from ee.hogai.chat_agent.taxonomy.types import TaxonomyAgentState
 
 class LoopNode(TaxonomyAgentNode[TaxonomyAgentState, TaxonomyAgentState[MaxToolTaxonomyOutput]]):
     def __init__(self, team: Team, user: User, toolkit_class: type[YourToolkit]):

ee/hogai/api/serializers.py

@@ -8,9 +8,9 @@
 from posthog.api.shared import UserBasicSerializer
 from posthog.exceptions_capture import capture_exception
 
-from ee.hogai.chat_agent.graph import AssistantGraph
-from ee.hogai.graph.deep_research.graph import DeepResearchAssistantGraph
-from ee.hogai.graph.deep_research.types import DeepResearchState
+from ee.hogai.chat_agent import AssistantGraph
+from ee.hogai.research_agent.graph import DeepResearchAssistantGraph
+from ee.hogai.research_agent.types import DeepResearchState
 from ee.hogai.utils.helpers import should_output_assistant_message
 from ee.hogai.utils.types import AssistantState
 from ee.hogai.utils.types.composed import AssistantMaxGraphState

ee/hogai/api/test/test_serializers.py

@@ -4,7 +4,7 @@
 from posthog.schema import AssistantMessage, AssistantToolCallMessage, ContextMessage
 
 from ee.hogai.api.serializers import ConversationSerializer
-from ee.hogai.chat_agent.graph import AssistantGraph
+from ee.hogai.chat_agent import AssistantGraph
 from ee.hogai.utils.types import AssistantState
 from ee.models.assistant import Conversation
 

ee/hogai/chat_agent/__init__.py

@@ -0,0 +1,3 @@
+from .graph import AssistantGraph
+
+__all__ = ["AssistantGraph"]

ee/hogai/chat_agent/loop_graph/test/__init__.py renamed to ee/hogai/chat_agent/dashboards/__init__.py

@@ -20,13 +20,13 @@
 from posthog.sync import database_sync_to_async
 from posthog.utils import pluralize
 
-from ee.hogai.graph.base import AssistantNode
-from ee.hogai.graph.parallel_task_execution.mixins import (
+from ee.hogai.chat_agent.parallel_task_execution.mixins import (
     WithInsightCreationTaskExecution,
     WithInsightSearchTaskExecution,
 )
-from ee.hogai.graph.parallel_task_execution.nodes import BaseTaskExecutorNode, TaskExecutionInputTuple
-from ee.hogai.graph.shared_prompts import HYPERLINK_USAGE_INSTRUCTIONS
+from ee.hogai.chat_agent.parallel_task_execution.nodes import BaseTaskExecutorNode, TaskExecutionInputTuple
+from ee.hogai.core.node import AssistantNode
+from ee.hogai.core.shared_prompts import HYPERLINK_USAGE_INSTRUCTIONS
 from ee.hogai.llm import MaxChatOpenAI
 from ee.hogai.utils.helpers import build_dashboard_url, build_insight_url, cast_assistant_query
 from ee.hogai.utils.types import AssistantState, PartialAssistantState

ee/hogai/graph/dashboards/nodes.py renamed to ee/hogai/chat_agent/dashboards/nodes.py

@@ -8,7 +8,7 @@
 
 from posthog.models import Dashboard, Insight, Team, User
 
-from ee.hogai.graph.dashboards.nodes import DashboardCreationExecutorNode, DashboardCreationNode, QueryMetadata
+from ee.hogai.chat_agent.dashboards.nodes import DashboardCreationExecutorNode, DashboardCreationNode, QueryMetadata
 from ee.hogai.utils.helpers import build_dashboard_url, build_insight_url
 from ee.hogai.utils.types import AssistantState, PartialAssistantState
 from ee.hogai.utils.types.base import (
@@ -163,7 +163,7 @@ def test_build_dashboard_url(self):
         expected_url = f"/project/{self.mock_team.id}/dashboard/{dashboard_id}"
         assert url == expected_url
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     async def test_arun_missing_search_insights_queries(self, mock_executor_node_class):
         """Test arun returns error when search_insights_queries is missing."""
         mock_executor_node = MagicMock()
@@ -183,7 +183,7 @@ async def test_arun_missing_search_insights_queries(self, mock_executor_node_cla
         assert isinstance(result.messages[0], AssistantToolCallMessage)
         assert "Search insights queries are required" in result.messages[0].content
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     @patch.object(DashboardCreationNode, "_search_insights")
     @patch.object(DashboardCreationNode, "_create_insights")
     @patch.object(DashboardCreationNode, "_create_dashboard_with_insights")
@@ -238,7 +238,7 @@ async def test_arun_successful_flow(
         assert "Dashboard Created" in result.messages[0].content
         assert "Test Dashboard" in result.messages[0].content
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     @patch.object(DashboardCreationNode, "_search_insights")
     @patch.object(DashboardCreationNode, "_create_insights")
     async def test_arun_no_insights_found_or_created(
@@ -275,9 +275,9 @@ async def test_arun_no_insights_found_or_created(
         assert isinstance(result.messages[0], AssistantToolCallMessage)
         assert "No existing insights matched" in result.messages[0].content
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     @patch.object(DashboardCreationNode, "_search_insights")
-    @patch("ee.hogai.graph.dashboards.nodes.logger")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.logger")
     async def test_arun_exception_handling(self, mock_logger, mock_search_insights, mock_executor_node_class):
         """Test arun handles exceptions properly."""
         mock_executor_node = MagicMock()
@@ -338,7 +338,7 @@ def test_create_no_insights_response(self):
 
     def test_create_error_response(self):
         """Test _create_error_response creates correct error message."""
-        with patch("ee.hogai.graph.dashboards.nodes.capture_exception") as mock_capture:
+        with patch("ee.hogai.chat_agent.dashboards.nodes.capture_exception") as mock_capture:
             result = self.node._create_error_response("Test error", "test_tool_call_id")
 
             assert isinstance(result, PartialAssistantState)
@@ -363,7 +363,7 @@ def setUp(self):
             node_path=(NodePath(name="test_node", message_id="test-id", tool_call_id="test_tool_call_id"),),
         )
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     async def test_create_insights(self, mock_executor_node_class):
         """Test _create_insights method."""
         # Setup mocks
@@ -411,7 +411,7 @@ async def mock_process_insight_creation_results(tool_calls, task_results, query_
             assert len(result["task_1"].created_insight_ids) == 2
             assert len(result["task_1"].created_insight_messages) == 1
 
-    @patch("ee.hogai.graph.dashboards.nodes.DashboardCreationExecutorNode")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.DashboardCreationExecutorNode")
     async def test_search_insights(self, mock_executor_node_class):
         """Test _search_insights method."""
         # Setup mocks
@@ -451,7 +451,7 @@ async def test_search_insights(self, mock_executor_node_class):
         assert len(result["task_1"].found_insight_ids) == 2
         assert len(result["task_1"].found_insight_messages) == 2
 
-    @patch("ee.hogai.graph.dashboards.nodes.database_sync_to_async")
+    @patch("ee.hogai.chat_agent.dashboards.nodes.database_sync_to_async")
     async def test_create_dashboard_with_insights(self, mock_db_sync):
         """Test _create_dashboard_with_insights method."""