Skip to content

Commit f5b979d

Browse files
author
kappa90
committed
feat(ph-ai): agent artifacts schema
1 parent 23c326b commit f5b979d

File tree

5 files changed

+432
-0
lines changed

5 files changed

+432
-0
lines changed

‎ee/hogai/core/test/test_artifacts.py‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
import pytest
2+
from posthog.test.base import BaseTest
3+
4+
from parameterized import parameterized
5+
from pydantic import ValidationError
6+
7+
from posthog.schema import AssistantTrendsQuery
8+
9+
from ee.hogai.artifacts.schemas import (
10+
DocumentArtifactContent,
11+
MarkdownBlock,
12+
SessionReplayBlock,
13+
VisualizationArtifactContent,
14+
VisualizationBlock,
15+
)
16+
17+
18+
class TestDocumentBlocks(BaseTest):
19+
@parameterized.expand(
20+
[
21+
("simple_markdown", "# Hello World", "# Hello World"),
22+
("empty_content", "", ""),
23+
("multiline", "Line 1\nLine 2\n**Bold**", "Line 1\nLine 2\n**Bold**"),
24+
]
25+
)
26+
def test_markdown_block_valid(self, _name: str, content: str, expected: str):
27+
block = MarkdownBlock(content=content)
28+
self.assertEqual(block.type, "markdown")
29+
self.assertEqual(block.content, expected)
30+
31+
def test_visualization_block_valid(self):
32+
block = VisualizationBlock(artifact_id="abc123")
33+
self.assertEqual(block.type, "visualization")
34+
self.assertEqual(block.artifact_id, "abc123")
35+
36+
@parameterized.expand(
37+
[
38+
("with_title", "session_123", 5000, "Event at 00:05"),
39+
("without_title", "session_456", 0, None),
40+
("large_timestamp", "session_789", 3600000, None),
41+
]
42+
)
43+
def test_session_replay_block_valid(self, _name: str, session_id: str, timestamp_ms: int, title: str | None):
44+
block = SessionReplayBlock(session_id=session_id, timestamp_ms=timestamp_ms, title=title)
45+
self.assertEqual(block.type, "session_replay")
46+
self.assertEqual(block.session_id, session_id)
47+
self.assertEqual(block.timestamp_ms, timestamp_ms)
48+
self.assertEqual(block.title, title)
49+
50+
def test_session_replay_block_zero_timestamp_valid(self):
51+
# Timestamp validation happens at the application level, not schema level
52+
# since TypeScript schemas don't support numeric constraints
53+
block = SessionReplayBlock(session_id="session_123", timestamp_ms=0)
54+
self.assertEqual(block.timestamp_ms, 0)
55+
56+
57+
class TestDocumentArtifactContent(BaseTest):
58+
def test_empty_blocks(self):
59+
content = DocumentArtifactContent(blocks=[])
60+
self.assertEqual(content.blocks, [])
61+
62+
def test_mixed_blocks(self):
63+
blocks = [
64+
{"type": "markdown", "content": "# Introduction"},
65+
{"type": "visualization", "artifact_id": "vis123"},
66+
{"type": "session_replay", "session_id": "sess456", "timestamp_ms": 1000, "title": "Example"},
67+
{"type": "markdown", "content": "## Summary"},
68+
]
69+
content = DocumentArtifactContent(blocks=blocks)
70+
71+
self.assertEqual(len(content.blocks), 4)
72+
self.assertIsInstance(content.blocks[0], MarkdownBlock)
73+
self.assertIsInstance(content.blocks[1], VisualizationBlock)
74+
self.assertIsInstance(content.blocks[2], SessionReplayBlock)
75+
self.assertIsInstance(content.blocks[3], MarkdownBlock)
76+
77+
def test_invalid_block_type(self):
78+
with pytest.raises(ValidationError):
79+
DocumentArtifactContent(blocks=[{"type": "invalid", "content": "test"}])
80+
81+
def test_serialization_round_trip(self):
82+
original = DocumentArtifactContent(
83+
blocks=[
84+
MarkdownBlock(content="# Title"),
85+
VisualizationBlock(artifact_id="abc123"),
86+
SessionReplayBlock(session_id="sess", timestamp_ms=5000, title="Test"),
87+
]
88+
)
89+
serialized = original.model_dump()
90+
deserialized = DocumentArtifactContent.model_validate(serialized)
91+
92+
self.assertEqual(len(deserialized.blocks), 3)
93+
block0 = deserialized.blocks[0]
94+
block1 = deserialized.blocks[1]
95+
block2 = deserialized.blocks[2]
96+
assert isinstance(block0, MarkdownBlock)
97+
assert isinstance(block1, VisualizationBlock)
98+
assert isinstance(block2, SessionReplayBlock)
99+
self.assertEqual(block0.content, "# Title")
100+
self.assertEqual(block1.artifact_id, "abc123")
101+
self.assertEqual(block2.session_id, "sess")
102+
103+
104+
class TestVisualizationArtifactContent(BaseTest):
105+
def test_trends_query(self):
106+
trends = AssistantTrendsQuery(series=[])
107+
content = VisualizationArtifactContent(query=trends, name="Test Trends", description="Shows trend data")
108+
109+
self.assertEqual(content.query, trends)
110+
self.assertEqual(content.name, "Test Trends")
111+
self.assertEqual(content.description, "Shows trend data")
112+
113+
def test_minimal_content(self):
114+
trends = AssistantTrendsQuery(series=[])
115+
content = VisualizationArtifactContent(query=trends)
116+
117+
self.assertEqual(content.query, trends)
118+
self.assertIsNone(content.name)
119+
self.assertIsNone(content.description)
120+
121+
def test_serialization_round_trip(self):
122+
original = VisualizationArtifactContent(
123+
query=AssistantTrendsQuery(series=[]),
124+
name="My Chart",
125+
description="Chart description",
126+
)
127+
serialized = original.model_dump()
128+
deserialized = VisualizationArtifactContent.model_validate(serialized)
129+
130+
self.assertEqual(deserialized.name, "My Chart")
131+
self.assertEqual(deserialized.description, "Chart description")
132+
self.assertIsNotNone(deserialized.query)

‎frontend/src/queries/schema.json‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -343,6 +343,70 @@
343343
"required": ["columns", "hogql", "limit", "offset", "results"],
344344
"type": "object"
345345
},
346+
"AgentArtifact": {
347+
"additionalProperties": false,
348+
"properties": {
349+
"content": {
350+
"$ref": "#/definitions/AgentArtifactContent"
351+
},
352+
"conversation": {
353+
"type": "string"
354+
},
355+
"created_at": {
356+
"type": "string"
357+
},
358+
"id": {
359+
"type": "string"
360+
},
361+
"name": {
362+
"type": "string"
363+
},
364+
"short_id": {
365+
"type": "string"
366+
},
367+
"team": {
368+
"type": "number"
369+
},
370+
"type": {
371+
"$ref": "#/definitions/AgentArtifactType"
372+
}
373+
},
374+
"required": ["id", "short_id", "name", "type", "content", "conversation", "team", "created_at"],
375+
"type": "object"
376+
},
377+
"AgentArtifactContent": {
378+
"anyOf": [
379+
{
380+
"$ref": "#/definitions/DocumentArtifactContent"
381+
},
382+
{
383+
"$ref": "#/definitions/VisualizationArtifactContent"
384+
}
385+
]
386+
},
387+
"AgentArtifactMinimal": {
388+
"additionalProperties": false,
389+
"properties": {
390+
"created_at": {
391+
"type": "string"
392+
},
393+
"name": {
394+
"type": "string"
395+
},
396+
"short_id": {
397+
"type": "string"
398+
},
399+
"type": {
400+
"$ref": "#/definitions/AgentArtifactType"
401+
}
402+
},
403+
"required": ["short_id", "name", "type", "created_at"],
404+
"type": "object"
405+
},
406+
"AgentArtifactType": {
407+
"enum": ["visualization", "document"],
408+
"type": "string"
409+
},
346410
"AgentMode": {
347411
"enum": ["product_analytics", "sql", "session_replay"],
348412
"type": "string"
@@ -10731,6 +10795,32 @@
1073110795
],
1073210796
"type": "string"
1073310797
},
10798+
"DocumentArtifactContent": {
10799+
"additionalProperties": false,
10800+
"properties": {
10801+
"blocks": {
10802+
"items": {
10803+
"$ref": "#/definitions/DocumentBlock"
10804+
},
10805+
"type": "array"
10806+
}
10807+
},
10808+
"required": ["blocks"],
10809+
"type": "object"
10810+
},
10811+
"DocumentBlock": {
10812+
"anyOf": [
10813+
{
10814+
"$ref": "#/definitions/MarkdownBlock"
10815+
},
10816+
{
10817+
"$ref": "#/definitions/VisualizationBlock"
10818+
},
10819+
{
10820+
"$ref": "#/definitions/SessionReplayBlock"
10821+
}
10822+
]
10823+
},
1073410824
"DocumentSimilarityQuery": {
1073510825
"additionalProperties": false,
1073610826
"properties": {
@@ -17493,6 +17583,20 @@
1749317583
"required": ["results"],
1749417584
"type": "object"
1749517585
},
17586+
"MarkdownBlock": {
17587+
"additionalProperties": false,
17588+
"properties": {
17589+
"content": {
17590+
"type": "string"
17591+
},
17592+
"type": {
17593+
"const": "markdown",
17594+
"type": "string"
17595+
}
17596+
},
17597+
"required": ["type", "content"],
17598+
"type": "object"
17599+
},
1749617600
"MarketingAnalyticsAggregatedQuery": {
1749717601
"additionalProperties": false,
1749817602
"properties": {
@@ -26214,6 +26318,26 @@
2621426318
"required": ["id", "viewed", "viewers", "recording_duration", "start_time", "end_time", "snapshot_source"],
2621526319
"type": "object"
2621626320
},
26321+
"SessionReplayBlock": {
26322+
"additionalProperties": false,
26323+
"properties": {
26324+
"session_id": {
26325+
"type": "string"
26326+
},
26327+
"timestamp_ms": {
26328+
"type": "number"
26329+
},
26330+
"title": {
26331+
"type": ["string", "null"]
26332+
},
26333+
"type": {
26334+
"const": "session_replay",
26335+
"type": "string"
26336+
}
26337+
},
26338+
"required": ["type", "session_id", "timestamp_ms"],
26339+
"type": "object"
26340+
},
2621726341
"SessionsQuery": {
2621826342
"additionalProperties": false,
2621926343
"properties": {
@@ -28747,6 +28871,49 @@
2874728871
"required": ["id", "distance"],
2874828872
"type": "object"
2874928873
},
28874+
"VisualizationArtifactContent": {
28875+
"additionalProperties": false,
28876+
"properties": {
28877+
"description": {
28878+
"type": ["string", "null"]
28879+
},
28880+
"name": {
28881+
"type": ["string", "null"]
28882+
},
28883+
"query": {
28884+
"anyOf": [
28885+
{
28886+
"$ref": "#/definitions/AssistantTrendsQuery"
28887+
},
28888+
{
28889+
"$ref": "#/definitions/AssistantFunnelsQuery"
28890+
},
28891+
{
28892+
"$ref": "#/definitions/AssistantRetentionQuery"
28893+
},
28894+
{
28895+
"$ref": "#/definitions/AssistantHogQLQuery"
28896+
}
28897+
]
28898+
}
28899+
},
28900+
"required": ["query"],
28901+
"type": "object"
28902+
},
28903+
"VisualizationBlock": {
28904+
"additionalProperties": false,
28905+
"properties": {
28906+
"artifact_id": {
28907+
"type": "string"
28908+
},
28909+
"type": {
28910+
"const": "visualization",
28911+
"type": "string"
28912+
}
28913+
},
28914+
"required": ["type", "artifact_id"],
28915+
"type": "object"
28916+
},
2875028917
"VisualizationItem": {
2875128918
"additionalProperties": false,
2875228919
"properties": {

‎frontend/src/queries/schema/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
// (even though our actual app's esbuild setup compiles perfectly well.)
66

77
// sort-imports-ignore
8+
export * from './schema-assistant-artifacts'
89
export * from './schema-assistant-messages'
910
export * from './schema-assistant-queries'
1011
export * from './schema-assistant-replay'

0 commit comments

Comments
 (0)