feat(ph-ai): agent artifacts schema

Author: kappa90

ee/hogai/core/test/test_artifacts.py

@@ -0,0 +1,131 @@
+import pytest
+from posthog.test.base import BaseTest
+
+from parameterized import parameterized
+from pydantic import ValidationError
+
+from posthog.schema import (
+    AssistantTrendsQuery,
+    DocumentArtifactContent,
+    MarkdownBlock,
+    SessionReplayBlock,
+    VisualizationArtifactContent,
+    VisualizationBlock,
+)
+
+
+class TestDocumentBlocks(BaseTest):
+    @parameterized.expand(
+        [
+            ("simple_markdown", "# Hello World", "# Hello World"),
+            ("empty_content", "", ""),
+            ("multiline", "Line 1\nLine 2\n**Bold**", "Line 1\nLine 2\n**Bold**"),
+        ]
+    )
+    def test_markdown_block_valid(self, _name: str, content: str, expected: str):
+        block = MarkdownBlock(content=content)
+        self.assertEqual(block.type, "markdown")
+        self.assertEqual(block.content, expected)
+
+    def test_visualization_block_valid(self):
+        block = VisualizationBlock(artifact_id="abc123")
+        self.assertEqual(block.type, "visualization")
+        self.assertEqual(block.artifact_id, "abc123")
+
+    @parameterized.expand(
+        [
+            ("with_title", "session_123", 5000, "Event at 00:05"),
+            ("without_title", "session_456", 0, None),
+            ("large_timestamp", "session_789", 3600000, None),
+        ]
+    )
+    def test_session_replay_block_valid(self, _name: str, session_id: str, timestamp_ms: int, title: str | None):
+        block = SessionReplayBlock(session_id=session_id, timestamp_ms=timestamp_ms, title=title)
+        self.assertEqual(block.type, "session_replay")
+        self.assertEqual(block.session_id, session_id)
+        self.assertEqual(block.timestamp_ms, timestamp_ms)
+        self.assertEqual(block.title, title)
+
+    def test_session_replay_block_zero_timestamp_valid(self):
+        # Timestamp validation happens at the application level, not schema level
+        # since TypeScript schemas don't support numeric constraints
+        block = SessionReplayBlock(session_id="session_123", timestamp_ms=0)
+        self.assertEqual(block.timestamp_ms, 0)
+
+
+class TestDocumentArtifactContent(BaseTest):
+    def test_empty_blocks(self):
+        content = DocumentArtifactContent(blocks=[])
+        self.assertEqual(content.blocks, [])
+
+    def test_mixed_blocks(self):
+        blocks = [
+            {"type": "markdown", "content": "# Introduction"},
+            {"type": "visualization", "artifact_id": "vis123"},
+            {"type": "session_replay", "session_id": "sess456", "timestamp_ms": 1000, "title": "Example"},
+            {"type": "markdown", "content": "## Summary"},
+        ]
+        content = DocumentArtifactContent(blocks=blocks)
+
+        self.assertEqual(len(content.blocks), 4)
+        self.assertIsInstance(content.blocks[0], MarkdownBlock)
+        self.assertIsInstance(content.blocks[1], VisualizationBlock)
+        self.assertIsInstance(content.blocks[2], SessionReplayBlock)
+        self.assertIsInstance(content.blocks[3], MarkdownBlock)
+
+    def test_invalid_block_type(self):
+        with pytest.raises(ValidationError):
+            DocumentArtifactContent(blocks=[{"type": "invalid", "content": "test"}])
+
+    def test_serialization_round_trip(self):
+        original = DocumentArtifactContent(
+            blocks=[
+                MarkdownBlock(content="# Title"),
+                VisualizationBlock(artifact_id="abc123"),
+                SessionReplayBlock(session_id="sess", timestamp_ms=5000, title="Test"),
+            ]
+        )
+        serialized = original.model_dump()
+        deserialized = DocumentArtifactContent.model_validate(serialized)
+
+        self.assertEqual(len(deserialized.blocks), 3)
+        block0 = deserialized.blocks[0]
+        block1 = deserialized.blocks[1]
+        block2 = deserialized.blocks[2]
+        assert isinstance(block0, MarkdownBlock)
+        assert isinstance(block1, VisualizationBlock)
+        assert isinstance(block2, SessionReplayBlock)
+        self.assertEqual(block0.content, "# Title")
+        self.assertEqual(block1.artifact_id, "abc123")
+        self.assertEqual(block2.session_id, "sess")
+
+
+class TestVisualizationArtifactContent(BaseTest):
+    def test_trends_query(self):
+        trends = AssistantTrendsQuery(series=[])
+        content = VisualizationArtifactContent(query=trends, name="Test Trends", description="Shows trend data")
+
+        self.assertEqual(content.query, trends)
+        self.assertEqual(content.name, "Test Trends")
+        self.assertEqual(content.description, "Shows trend data")
+
+    def test_minimal_content(self):
+        trends = AssistantTrendsQuery(series=[])
+        content = VisualizationArtifactContent(query=trends)
+
+        self.assertEqual(content.query, trends)
+        self.assertIsNone(content.name)
+        self.assertIsNone(content.description)
+
+    def test_serialization_round_trip(self):
+        original = VisualizationArtifactContent(
+            query=AssistantTrendsQuery(series=[]),
+            name="My Chart",
+            description="Chart description",
+        )
+        serialized = original.model_dump()
+        deserialized = VisualizationArtifactContent.model_validate(serialized)
+
+        self.assertEqual(deserialized.name, "My Chart")
+        self.assertEqual(deserialized.description, "Chart description")
+        self.assertIsNotNone(deserialized.query)

frontend/src/queries/schema.json

@@ -343,6 +343,16 @@
             "required": ["columns", "hogql", "limit", "offset", "results"],
             "type": "object"
         },
+        "AgentArtifactContent": {
+            "anyOf": [
+                {
+                    "$ref": "#/definitions/DocumentArtifactContent"
+                },
+                {
+                    "$ref": "#/definitions/VisualizationArtifactContent"
+                }
+            ]
+        },
         "AgentMode": {
             "enum": ["product_analytics", "sql", "session_replay"],
             "type": "string"
@@ -10731,6 +10741,32 @@
             ],
             "type": "string"
         },
+        "DocumentArtifactContent": {
+            "additionalProperties": false,
+            "properties": {
+                "blocks": {
+                    "items": {
+                        "$ref": "#/definitions/DocumentBlock"
+                    },
+                    "type": "array"
+                }
+            },
+            "required": ["blocks"],
+            "type": "object"
+        },
+        "DocumentBlock": {
+            "anyOf": [
+                {
+                    "$ref": "#/definitions/MarkdownBlock"
+                },
+                {
+                    "$ref": "#/definitions/VisualizationBlock"
+                },
+                {
+                    "$ref": "#/definitions/SessionReplayBlock"
+                }
+            ]
+        },
         "DocumentSimilarityQuery": {
             "additionalProperties": false,
             "properties": {
@@ -17490,6 +17526,20 @@
             "required": ["results"],
             "type": "object"
         },
+        "MarkdownBlock": {
+            "additionalProperties": false,
+            "properties": {
+                "content": {
+                    "type": "string"
+                },
+                "type": {
+                    "const": "markdown",
+                    "type": "string"
+                }
+            },
+            "required": ["type", "content"],
+            "type": "object"
+        },
         "MarketingAnalyticsAggregatedQuery": {
             "additionalProperties": false,
             "properties": {
@@ -26265,6 +26315,26 @@
             "required": ["id", "viewed", "viewers", "recording_duration", "start_time", "end_time", "snapshot_source"],
             "type": "object"
         },
+        "SessionReplayBlock": {
+            "additionalProperties": false,
+            "properties": {
+                "session_id": {
+                    "type": "string"
+                },
+                "timestamp_ms": {
+                    "type": "number"
+                },
+                "title": {
+                    "type": ["string", "null"]
+                },
+                "type": {
+                    "const": "session_replay",
+                    "type": "string"
+                }
+            },
+            "required": ["type", "session_id", "timestamp_ms"],
+            "type": "object"
+        },
         "SessionsQuery": {
             "additionalProperties": false,
             "properties": {
@@ -28813,6 +28883,49 @@
             "required": ["id", "distance"],
             "type": "object"
         },
+        "VisualizationArtifactContent": {
+            "additionalProperties": false,
+            "properties": {
+                "description": {
+                    "type": ["string", "null"]
+                },
+                "name": {
+                    "type": ["string", "null"]
+                },
+                "query": {
+                    "anyOf": [
+                        {
+                            "$ref": "#/definitions/AssistantTrendsQuery"
+                        },
+                        {
+                            "$ref": "#/definitions/AssistantFunnelsQuery"
+                        },
+                        {
+                            "$ref": "#/definitions/AssistantRetentionQuery"
+                        },
+                        {
+                            "$ref": "#/definitions/AssistantHogQLQuery"
+                        }
+                    ]
+                }
+            },
+            "required": ["query"],
+            "type": "object"
+        },
+        "VisualizationBlock": {
+            "additionalProperties": false,
+            "properties": {
+                "artifact_id": {
+                    "type": "string"
+                },
+                "type": {
+                    "const": "visualization",
+                    "type": "string"
+                }
+            },
+            "required": ["type", "artifact_id"],
+            "type": "object"
+        },
         "VisualizationItem": {
             "additionalProperties": false,
             "properties": {

frontend/src/queries/schema/index.ts

@@ -5,6 +5,7 @@
 // (even though our actual app's esbuild setup compiles perfectly well.)
 
 // sort-imports-ignore
+export * from './schema-assistant-artifacts'
 export * from './schema-assistant-messages'
 export * from './schema-assistant-queries'
 export * from './schema-assistant-replay'

frontend/src/queries/schema/schema-assistant-artifacts.ts

@@ -0,0 +1,38 @@
+// Agent Artifact Types - these will be auto-generated to Python via pnpm schema:build
+import {
+    AssistantFunnelsQuery,
+    AssistantHogQLQuery,
+    AssistantRetentionQuery,
+    AssistantTrendsQuery,
+} from './schema-assistant-queries'
+
+export interface MarkdownBlock {
+    type: 'markdown'
+    content: string
+}
+
+export interface VisualizationBlock {
+    type: 'visualization'
+    artifact_id: string
+}
+
+export interface SessionReplayBlock {
+    type: 'session_replay'
+    session_id: string
+    timestamp_ms: number
+    title?: string | null
+}
+
+export type DocumentBlock = MarkdownBlock | VisualizationBlock | SessionReplayBlock
+
+export interface DocumentArtifactContent {
+    blocks: DocumentBlock[]
+}
+
+export interface VisualizationArtifactContent {
+    query: AssistantTrendsQuery | AssistantFunnelsQuery | AssistantRetentionQuery | AssistantHogQLQuery
+    name?: string | null
+    description?: string | null
+}
+
+export type AgentArtifactContent = DocumentArtifactContent | VisualizationArtifactContent

frontend/src/scenes/max/maxTypes.ts

@@ -1,3 +1,4 @@
+import { AgentArtifactContent } from '~/queries/schema/schema-assistant-artifacts'
 import { AgentMode } from '~/queries/schema/schema-assistant-messages'
 import { DashboardFilter, HogQLVariable, QuerySchema } from '~/queries/schema/schema-general'
 import { integer } from '~/queries/schema/type-utils'
@@ -135,3 +136,19 @@ export const createMaxContextHelpers = {
 export function isAgentMode(mode: unknown): mode is AgentMode {
     return typeof mode === 'string' && Object.values(AgentMode).includes(mode as AgentMode)
 }
+
+export enum AgentArtifactType {
+    VISUALIZATION = 'visualization',
+    DOCUMENT = 'document',
+}
+
+export interface AgentArtifact {
+    id: string
+    short_id: string
+    name: string
+    type: AgentArtifactType
+    content: AgentArtifactContent
+    conversation: string
+    team: number
+    created_at: string
+}