refactor(ph-ai): move chat agent specifics to dedicated classes

Author: kappa90

ee/api/conversation.py

@@ -30,8 +30,8 @@
 )
 from posthog.utils import get_instance_region
 
-from ee.hogai.agent.executor import AgentExecutor
 from ee.hogai.api.serializers import ConversationSerializer
+from ee.hogai.core.executor import AgentExecutor
 from ee.hogai.utils.aio import async_to_sync
 from ee.hogai.utils.sse import AssistantSSESerializer
 from ee.hogai.utils.types.base import AssistantMode

ee/api/max_tools.py

@@ -15,7 +15,7 @@
 from posthog.rate_limit import AIBurstRateThrottle, AISustainedRateThrottle
 from posthog.renderers import SafeJSONRenderer
 
-from ee.hogai.utils.types import AssistantMode, AssistantState
+from ee.hogai.utils.types import AssistantState
 from ee.models.assistant import Conversation
 
 
@@ -53,17 +53,16 @@ class MaxToolsViewSet(TeamAndOrgViewSetMixin, GenericViewSet):
         required_scopes=["insight:read", "query:read"],
     )
     def create_and_query_insight(self, request: Request, *args, **kwargs):
-        from ee.hogai.assistant import Assistant
+        from ee.hogai.insights_assistant import InsightsAssistant
 
         serializer = InsightsToolCallSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         conversation = self.get_queryset().create(user=request.user, team=self.team, type=Conversation.Type.TOOL_CALL)
-        assistant = Assistant.create(
+        assistant = InsightsAssistant(
             self.team,
             conversation,
             user=cast(User, request.user),
             is_new_conversation=False,  # we don't care about the conversation id being sent back to the client
-            mode=AssistantMode.INSIGHTS_TOOL,
             initial_state=serializer.validated_data["state"],
         )
 

ee/api/test/test_conversation.py

@@ -102,7 +102,7 @@ def test_create_conversation(self):
         conversation_id = str(uuid.uuid4())
 
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_start_workflow_and_stream:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -129,7 +129,7 @@ def test_create_conversation(self):
 
     def test_add_message_to_existing_conversation(self):
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_start_workflow_and_stream:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -183,7 +183,7 @@ def test_invalid_message_format(self):
     def test_rate_limit_burst(self):
         # Create multiple requests to trigger burst rate limit
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ):
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -215,7 +215,7 @@ def test_none_content_with_existing_conversation(self):
             user=self.user, team=self.team, status=Conversation.Status.IN_PROGRESS
         )
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_stream_conversation:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -251,7 +251,7 @@ def test_missing_trace_id(self):
 
     def test_nonexistent_conversation(self):
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ):
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -284,7 +284,7 @@ def test_unauthenticated_request(self):
         )
         self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
 
-    @patch("ee.hogai.agent.executor.AgentExecutor.cancel_workflow")
+    @patch("ee.hogai.core.executor.AgentExecutor.cancel_workflow")
     def test_cancel_conversation(self, mock_cancel):
         conversation = Conversation.objects.create(
             user=self.user,
@@ -327,7 +327,7 @@ def test_cancel_other_teams_conversation(self):
         )
         self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
 
-    @patch("ee.hogai.agent.executor.AgentExecutor.cancel_workflow")
+    @patch("ee.hogai.core.executor.AgentExecutor.cancel_workflow")
     def test_cancel_conversation_with_async_cleanup(self, mock_cancel):
         """Test that cancel endpoint properly handles async cleanup."""
         conversation = Conversation.objects.create(
@@ -347,7 +347,7 @@ def test_cancel_conversation_with_async_cleanup(self, mock_cancel):
 
         self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
 
-    @patch("ee.hogai.agent.executor.AgentExecutor.cancel_workflow")
+    @patch("ee.hogai.core.executor.AgentExecutor.cancel_workflow")
     def test_cancel_conversation_async_cleanup_failure(self, mock_cancel):
         """Test cancel endpoint behavior when async cleanup fails."""
         conversation = Conversation.objects.create(
@@ -415,7 +415,7 @@ def test_stream_from_in_progress_conversation(self):
             user=self.user, team=self.team, status=Conversation.Status.IN_PROGRESS
         )
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_stream_conversation:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -581,7 +581,7 @@ def test_billing_context_validation_valid_data(self):
         conversation = Conversation.objects.create(user=self.user, team=self.team)
 
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_start_workflow_and_stream:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):
@@ -605,7 +605,7 @@ def test_billing_context_validation_invalid_data(self):
         conversation = Conversation.objects.create(user=self.user, team=self.team)
 
         with patch(
-            "ee.hogai.agent.executor.AgentExecutor.astream",
+            "ee.hogai.core.executor.AgentExecutor.astream",
             return_value=_async_generator(),
         ) as mock_start_workflow_and_stream:
             with patch("ee.api.conversation.StreamingHttpResponse", side_effect=self._create_mock_streaming_response):

ee/api/test/test_max_tools.py

@@ -3,7 +3,7 @@
 
 from mistralai_azure import AssistantMessage
 
-from ee.hogai.assistant.insights_assistant import InsightsAssistant
+from ee.hogai.insights_assistant import InsightsAssistant
 
 
 class TestMaxToolsAPI(APIBaseTest):

ee/hogai/api/serializers.py

@@ -7,9 +7,9 @@
 
 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.graph.graph import AssistantGraph
 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.graph.graph import AssistantGraph
+from ee.hogai.chat_agent.graph import AssistantGraph
 from ee.hogai.utils.types import AssistantState
 from ee.models.assistant import Conversation
 

ee/hogai/assistant/__init__.py

@@ -1,12 +1,8 @@
 from collections.abc import Callable
 
+from ee.hogai.chat_agent.loop_graph.graph import ChatAgentLoopGraph
 from ee.hogai.django_checkpoint.checkpointer import DjangoCheckpointer
-from ee.hogai.graph.agent_executor import AgentExecutorGraph
-from ee.hogai.graph.title_generator.nodes import TitleGeneratorNode
-from ee.hogai.graph.usage import UsageNode
-from ee.hogai.utils.types.base import AssistantGraphName, AssistantNodeName, AssistantState
-
-from .memory.nodes import (
+from ee.hogai.graph.memory.nodes import (
     MemoryCollectorNode,
     MemoryCollectorToolsNode,
     MemoryInitializerInterruptNode,
@@ -16,9 +12,12 @@
     MemoryOnboardingFinalizeNode,
     MemoryOnboardingNode,
 )
+from ee.hogai.graph.title_generator.nodes import TitleGeneratorNode
+from ee.hogai.graph.usage.nodes import UsageNode
+from ee.hogai.utils.types.base import AssistantGraphName, AssistantNodeName, AssistantState
 
 
-class AssistantGraph(AgentExecutorGraph):
+class AssistantGraph(ChatAgentLoopGraph):
     @property
     def graph_name(self) -> AssistantGraphName:
         return AssistantGraphName.ASSISTANT