feat(ph-ai): use artifacts for visualizations (#41947)

## Problem

We want to migrate to smarter artifact management. This PR introduces an artifact manager that can handle visualization objects, either created by the assistant, or retrieved from the user's ones.

## Changes

- Created a new `ArtifactManager` class to handle creation, retrieval, and enrichment of agent artifacts
- Updated stream processing to handle artifact messages and convert them to appropriate visualization formats
- Refactored frontend component to use artifact visualizations.

## How did you test this code?

- Locally, tests to be fixed

Author: kappa90

ee/hogai/api/serializers.py

@@ -8,6 +8,7 @@
 from posthog.api.shared import UserBasicSerializer
 from posthog.exceptions_capture import capture_exception
 
+from ee.hogai.artifacts.manager import ArtifactManager
 from ee.hogai.chat_agent import AssistantGraph
 from ee.hogai.utils.helpers import should_output_assistant_message
 from ee.hogai.utils.types import AssistantState
@@ -75,8 +76,11 @@ async def _aget_messages_with_flag(self, conversation: Conversation) -> tuple[li
                 {"configurable": {"thread_id": str(conversation.id), "checkpoint_ns": ""}}
             )
             state = state_class.model_validate(snapshot.values)
-            messages = [message.model_dump() for message in state.messages if should_output_assistant_message(message)]
-            return messages, False
+            messages = list(state.messages)
+            artifact_manager = ArtifactManager(team, user)
+            enriched_messages = await artifact_manager.aenrich_messages(messages)
+            return [m.model_dump() for m in enriched_messages if should_output_assistant_message(m)], False
+
         except pydantic.ValidationError as e:
             capture_exception(
                 e,

ee/hogai/api/test/test_serializers.py

@@ -1,12 +1,21 @@
+from uuid import uuid4
+
 from posthog.test.base import APIBaseTest
 from unittest.mock import AsyncMock, patch
 
-from posthog.schema import AssistantMessage, AssistantToolCallMessage, ContextMessage
+from posthog.schema import (
+    ArtifactContentType,
+    ArtifactSource,
+    AssistantMessage,
+    AssistantToolCallMessage,
+    ContextMessage,
+)
 
 from ee.hogai.api.serializers import ConversationSerializer
 from ee.hogai.chat_agent import AssistantGraph
 from ee.hogai.utils.types import AssistantState
-from ee.models.assistant import Conversation
+from ee.hogai.utils.types.base import ArtifactRefMessage
+from ee.models.assistant import AgentArtifact, Conversation
 
 
 class TestConversationSerializers(APIBaseTest):
@@ -67,8 +76,8 @@ class MockSnapshot:
             # Third message should be the AssistantToolCallMessage without UI payload
             self.assertEqual(filtered_messages[2]["ui_payload"], None)
 
-    def test_get_messages_handles_validation_errors(self):
-        """Gracefully fall back to an empty list when the stored state fails validation."""
+    def test_get_messages_handles_validation_errors_and_sets_unsupported_content(self):
+        """Gracefully fall back to an empty list when the stored state fails validation, and set has_unsupported_content."""
         conversation = Conversation.objects.create(
             user=self.user, team=self.team, title="Conversation with invalid state", type=Conversation.Type.ASSISTANT
         )
@@ -88,40 +97,74 @@ def test_get_messages_handles_validation_errors(self):
             ).data
 
         self.assertEqual(data["messages"], [])
+        self.assertTrue(data["has_unsupported_content"])
 
-    def test_has_unsupported_content_on_validation_error(self):
-        """When validation fails, has_unsupported_content should be True."""
+    def test_caching_prevents_duplicate_operations(self):
+        """This is to test that the caching works correctly as to not incurring in unnecessary operations (We would do a DRF call per field call)."""
         conversation = Conversation.objects.create(
-            user=self.user,
-            team=self.team,
-            title="Conversation with schema mismatch",
-            type=Conversation.Type.ASSISTANT,
+            user=self.user, team=self.team, title="Cached conversation", type=Conversation.Type.ASSISTANT
         )
 
-        invalid_snapshot = type("Snapshot", (), {"values": {"messages": [{"invalid": "schema"}]}})()
+        state = AssistantState(messages=[AssistantMessage(content="Cached message", type="ai")])
 
         with patch("langgraph.graph.state.CompiledStateGraph.aget_state", new_callable=AsyncMock) as mock_get_state:
-            mock_get_state.return_value = invalid_snapshot
 
-            data = ConversationSerializer(
+            class MockSnapshot:
+                values = state.model_dump()
+
+            mock_get_state.return_value = MockSnapshot()
+
+            serializer = ConversationSerializer(
                 conversation,
                 context={
                     "team": self.team,
                     "user": self.user,
                 },
-            ).data
+            )
 
-        self.assertEqual(data["messages"], [])
-        self.assertTrue(data["has_unsupported_content"])
+            # Explicitly access both fields multiple times
+            _ = serializer.data["messages"]
+            _ = serializer.data["has_unsupported_content"]
+            _ = serializer.data["messages"]
+            _ = serializer.data["has_unsupported_content"]
 
-    def test_has_unsupported_content_on_other_errors(self):
-        """On non-validation errors, has_unsupported_content should be False."""
+        # aget_state should only be called once though
+        self.assertEqual(mock_get_state.call_count, 1)
+
+
+class TestConversationSerializerArtifactEnrichment(APIBaseTest):
+    """Test artifact enrichment functionality in the serializer."""
+
+    def test_artifact_ref_message_enriched_in_response(self):
+        """Test that ArtifactRefMessage is enriched with content from database artifact."""
         conversation = Conversation.objects.create(
-            user=self.user, team=self.team, title="Conversation with graph error", type=Conversation.Type.ASSISTANT
+            user=self.user, team=self.team, title="Artifact test conversation", type=Conversation.Type.ASSISTANT
         )
 
+        # Create an artifact in the database
+        artifact = AgentArtifact.objects.create(
+            name="Test Artifact",
+            type=AgentArtifact.Type.VISUALIZATION,
+            data={"query": {"kind": "TrendsQuery", "series": []}, "name": "Chart Name"},
+            conversation=conversation,
+            team=self.team,
+        )
+
+        # Create state with an ArtifactRefMessage
+        artifact_message = ArtifactRefMessage(
+            id=str(uuid4()),
+            content_type=ArtifactContentType.VISUALIZATION,
+            artifact_id=artifact.short_id,
+            source=ArtifactSource.ARTIFACT,
+        )
+        state = AssistantState(messages=[artifact_message])
+
         with patch("langgraph.graph.state.CompiledStateGraph.aget_state", new_callable=AsyncMock) as mock_get_state:
-            mock_get_state.side_effect = RuntimeError("Graph compilation failed")
+
+            class MockSnapshot:
+                values = state.model_dump()
+
+            mock_get_state.return_value = MockSnapshot()
 
             data = ConversationSerializer(
                 conversation,
@@ -131,16 +174,27 @@ def test_has_unsupported_content_on_other_errors(self):
                 },
             ).data
 
-        self.assertEqual(data["messages"], [])
-        self.assertFalse(data["has_unsupported_content"])
+            # The message should be enriched as an ArtifactMessage
+            self.assertEqual(len(data["messages"]), 1)
+            enriched_msg = data["messages"][0]
+            self.assertEqual(enriched_msg["type"], "ai/artifact")
+            self.assertEqual(enriched_msg["artifact_id"], artifact.short_id)
+            self.assertEqual(enriched_msg["content"]["name"], "Chart Name")
 
-    def test_has_unsupported_content_on_success(self):
-        """On successful message fetch, has_unsupported_content should be False."""
+    def test_artifact_ref_message_filtered_when_not_found(self):
+        """Test that ArtifactRefMessage is filtered out when artifact not found in database."""
         conversation = Conversation.objects.create(
-            user=self.user, team=self.team, title="Valid conversation", type=Conversation.Type.ASSISTANT
+            user=self.user, team=self.team, title="Missing artifact conversation", type=Conversation.Type.ASSISTANT
         )
 
-        state = AssistantState(messages=[AssistantMessage(content="Test message", type="ai")])
+        # Create state with an ArtifactRefMessage pointing to non-existent artifact
+        artifact_message = ArtifactRefMessage(
+            id=str(uuid4()),
+            content_type=ArtifactContentType.VISUALIZATION,
+            artifact_id="nonexistent",
+            source=ArtifactSource.ARTIFACT,
+        )
+        state = AssistantState(messages=[artifact_message])
 
         with patch("langgraph.graph.state.CompiledStateGraph.aget_state", new_callable=AsyncMock) as mock_get_state:
 
@@ -157,16 +211,32 @@ class MockSnapshot:
                 },
             ).data
 
-        self.assertEqual(len(data["messages"]), 1)
-        self.assertFalse(data["has_unsupported_content"])
+            # The message should be filtered out
+            self.assertEqual(len(data["messages"]), 0)
 
-    def test_caching_prevents_duplicate_operations(self):
-        """This is to test that the caching works correctly as to not incurring in unnecessary operations (We would do a DRF call per field call)."""
+    def test_mixed_messages_with_artifacts(self):
+        """Test serialization with mixed message types including artifacts."""
         conversation = Conversation.objects.create(
-            user=self.user, team=self.team, title="Cached conversation", type=Conversation.Type.ASSISTANT
+            user=self.user, team=self.team, title="Mixed messages conversation", type=Conversation.Type.ASSISTANT
         )
 
-        state = AssistantState(messages=[AssistantMessage(content="Cached message", type="ai")])
+        artifact = AgentArtifact.objects.create(
+            name="Mixed Artifact",
+            type=AgentArtifact.Type.VISUALIZATION,
+            data={"query": {"kind": "TrendsQuery", "series": []}, "name": "Mixed Chart"},
+            conversation=conversation,
+            team=self.team,
+        )
+
+        # Create state with mixed message types
+        assistant_message = AssistantMessage(content="Hello from assistant", type="ai")
+        artifact_message = ArtifactRefMessage(
+            id=str(uuid4()),
+            content_type=ArtifactContentType.VISUALIZATION,
+            artifact_id=artifact.short_id,
+            source=ArtifactSource.ARTIFACT,
+        )
+        state = AssistantState(messages=[assistant_message, artifact_message])
 
         with patch("langgraph.graph.state.CompiledStateGraph.aget_state", new_callable=AsyncMock) as mock_get_state:
 
@@ -175,19 +245,17 @@ class MockSnapshot:
 
             mock_get_state.return_value = MockSnapshot()
 
-            serializer = ConversationSerializer(
+            data = ConversationSerializer(
                 conversation,
                 context={
                     "team": self.team,
                     "user": self.user,
                 },
-            )
-
-            # Explicitly access both fields multiple times
-            _ = serializer.data["messages"]
-            _ = serializer.data["has_unsupported_content"]
-            _ = serializer.data["messages"]
-            _ = serializer.data["has_unsupported_content"]
+            ).data
 
-        # aget_state should only be called once though
-        self.assertEqual(mock_get_state.call_count, 1)
+            # Both messages should be included (AssistantMessage and enriched ArtifactMessage)
+            self.assertEqual(len(data["messages"]), 2)
+            self.assertEqual(data["messages"][0]["type"], "ai")
+            self.assertEqual(data["messages"][0]["content"], "Hello from assistant")
+            self.assertEqual(data["messages"][1]["type"], "ai/artifact")
+            self.assertEqual(data["messages"][1]["content"]["name"], "Mixed Chart")

ee/hogai/artifacts/__init__.py

@@ -0,0 +1,20 @@
+from posthog.schema import (
+    ArtifactSource,
+    DocumentArtifactContent,
+    MarkdownBlock,
+    SessionReplayBlock,
+    VisualizationArtifactContent,
+    VisualizationBlock,
+)
+
+from .manager import ArtifactManager
+
+__all__ = [
+    "ArtifactSource",
+    "ArtifactManager",
+    "DocumentArtifactContent",
+    "MarkdownBlock",
+    "SessionReplayBlock",
+    "VisualizationArtifactContent",
+    "VisualizationBlock",
+]