Skip to content

Commit 514a053

Browse files
author
kappa90
committed
feat(ph-ai): agent artifacts API endpoint
1 parent 2b377bf commit 514a053

File tree

5 files changed

+135
-22
lines changed

5 files changed

+135
-22
lines changed

‎ee/hogai/chat_agent/usage/test/test_usage.py‎

Lines changed: 19 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -18,53 +18,53 @@
1818

1919

2020
class TestUsage(BaseTest):
21-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
22-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
21+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
22+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
2323
def test_get_ai_free_tier_credits_default(self, mock_region, mock_payload):
2424
"""Test that teams without custom limits get the default free tier."""
2525
mock_region.return_value = "EU"
2626
mock_payload.return_value = None
2727
credits = get_ai_free_tier_credits(team_id=999)
2828
self.assertEqual(credits, DEFAULT_FREE_TIER_CREDITS)
2929

30-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
31-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
30+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
31+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
3232
def test_get_ai_free_tier_credits_custom_eu(self, mock_region, mock_payload):
3333
"""Test that EU internal team gets custom free tier limit from feature flag."""
3434
mock_region.return_value = "EU"
3535
mock_payload.return_value = {"EU": {"1": 9999999}, "US": {"2": 9999999}}
3636
credits = get_ai_free_tier_credits(team_id=1)
3737
self.assertEqual(credits, 9999999)
3838

39-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
40-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
39+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
40+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
4141
def test_get_ai_free_tier_credits_custom_us(self, mock_region, mock_payload):
4242
"""Test that US internal team gets custom free tier limit from feature flag."""
4343
mock_region.return_value = "US"
4444
mock_payload.return_value = {"EU": {"1": 9999999}, "US": {"2": 9999999}}
4545
credits = get_ai_free_tier_credits(team_id=2)
4646
self.assertEqual(credits, 9999999)
4747

48-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
49-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
48+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
49+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
5050
def test_get_ai_free_tier_credits_fallback_when_region_unknown(self, mock_region, mock_payload):
5151
"""Test that unknown regions fall back to default."""
5252
mock_region.return_value = None
5353
mock_payload.return_value = {"EU": {"1": 9999999}}
5454
credits = get_ai_free_tier_credits(team_id=1)
5555
self.assertEqual(credits, DEFAULT_FREE_TIER_CREDITS)
5656

57-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
58-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
57+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
58+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
5959
def test_get_ai_free_tier_credits_team_not_in_payload(self, mock_region, mock_payload):
6060
"""Test that teams not in the payload get default credits."""
6161
mock_region.return_value = "EU"
6262
mock_payload.return_value = {"EU": {"1": 9999999}}
6363
credits = get_ai_free_tier_credits(team_id=999)
6464
self.assertEqual(credits, DEFAULT_FREE_TIER_CREDITS)
6565

66-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
67-
@patch("ee.hogai.chat_agent.queries.get_instance_region")
66+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
67+
@patch("ee.hogai.chat_agent.usage.queries.get_instance_region")
6868
def test_get_ai_free_tier_credits_invalid_payload(self, mock_region, mock_payload):
6969
"""Test that invalid payloads fall back to default."""
7070
mock_region.return_value = "EU"
@@ -90,21 +90,21 @@ def test_get_conversation_start_time_not_exists(self):
9090
start_time = get_conversation_start_time(uuid4())
9191
self.assertIsNone(start_time)
9292

93-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
93+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
9494
def test_get_ga_launch_date_from_payload(self, mock_payload):
9595
"""Test that GA launch date is fetched from feature flag payload."""
9696
mock_payload.return_value = {"ga_launch_date": "2025-12-01", "EU": {"1": 10000}}
9797
ga_date = get_ga_launch_date()
9898
self.assertEqual(ga_date, datetime(2025, 12, 1, tzinfo=UTC))
9999

100-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
100+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
101101
def test_get_ga_launch_date_fallback(self, mock_payload):
102102
"""Test that GA launch date falls back to default when not in payload."""
103103
mock_payload.return_value = None
104104
ga_date = get_ga_launch_date()
105105
self.assertEqual(ga_date, DEFAULT_GA_LAUNCH_DATE)
106106

107-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
107+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
108108
def test_get_ga_launch_date_invalid_format(self, mock_payload):
109109
"""Test that invalid date format falls back to default."""
110110
mock_payload.return_value = {"ga_launch_date": "invalid-date"}
@@ -114,20 +114,20 @@ def test_get_ga_launch_date_invalid_format(self, mock_payload):
114114
def test_get_past_month_start_normal(self):
115115
"""Test past month start when 30 days ago is after GA launch."""
116116
now = datetime(2025, 12, 20, tzinfo=UTC)
117-
with patch("ee.hogai.chat_agent.queries.datetime") as mock_datetime:
117+
with patch("ee.hogai.chat_agent.usage.queries.datetime") as mock_datetime:
118118
mock_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)
119119
mock_datetime.now.return_value = now
120120
past_month_start = get_past_month_start()
121121
expected = datetime(2025, 11, 20, tzinfo=UTC)
122122
self.assertEqual(past_month_start, expected)
123123

124-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
124+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
125125
def test_get_past_month_start_capped_at_ga_launch(self, mock_payload):
126126
"""Test that past month start is capped at GA launch date."""
127127
mock_payload.return_value = None
128128
# Set current time to shortly after GA launch
129129
now = DEFAULT_GA_LAUNCH_DATE + timedelta(days=10)
130-
with patch("ee.hogai.chat_agent.queries.datetime") as mock_datetime:
130+
with patch("ee.hogai.chat_agent.usage.queries.datetime") as mock_datetime:
131131
mock_datetime.side_effect = lambda *args, **kw: datetime(*args, **kw)
132132
mock_datetime.now.return_value = now
133133
past_month_start = get_past_month_start()
@@ -169,7 +169,7 @@ def test_format_usage_message_over_limit(self):
169169
self.assertIn("**Overage**: 500 credits over limit", message)
170170
self.assertIn("125% of free tier", message)
171171

172-
@patch("ee.hogai.chat_agent.queries.posthoganalytics.get_feature_flag_payload")
172+
@patch("ee.hogai.chat_agent.usage.queries.posthoganalytics.get_feature_flag_payload")
173173
def test_format_usage_message_with_ga_cap(self, mock_payload):
174174
"""Test formatting when GA cap is active."""
175175
mock_payload.return_value = None

‎ee/migrations/0031_agentartifact.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
# Generated by Django 4.2.26 on 2025-11-21 16:00
2+
3+
import django.db.models.deletion
4+
from django.conf import settings
5+
from django.db import migrations, models
6+
7+
import posthog.models.utils
8+
9+
import ee.models.assistant
10+
11+
12+
class Migration(migrations.Migration):
13+
dependencies = [
14+
("posthog", "0914_alter_userproductlist_reason"),
15+
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
16+
("ee", "0030_singlesessionsummary_distinct_id_and_more"),
17+
]
18+
19+
operations = [
20+
migrations.CreateModel(
21+
name="AgentArtifact",
22+
fields=[
23+
("created_at", models.DateTimeField(auto_now_add=True)),
24+
("updated_at", models.DateTimeField(auto_now=True, null=True)),
25+
("deleted", models.BooleanField(blank=True, default=False, null=True)),
26+
("deleted_at", models.DateTimeField(blank=True, null=True)),
27+
(
28+
"id",
29+
models.UUIDField(
30+
default=posthog.models.utils.uuid7, editable=False, primary_key=True, serialize=False
31+
),
32+
),
33+
("short_id", models.CharField(default=ee.models.assistant.generate_short_id, max_length=4)),
34+
("name", models.CharField(max_length=400)),
35+
(
36+
"type",
37+
models.CharField(
38+
choices=[("visualization", "Visualization"), ("notebook", "Notebook")], max_length=50
39+
),
40+
),
41+
("data", models.JSONField(help_text="Artifact content. Structure depends on artifact type.")),
42+
(
43+
"conversation",
44+
models.ForeignKey(
45+
on_delete=django.db.models.deletion.CASCADE, related_name="artifacts", to="ee.conversation"
46+
),
47+
),
48+
(
49+
"created_by",
50+
models.ForeignKey(
51+
blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to=settings.AUTH_USER_MODEL
52+
),
53+
),
54+
("team", models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to="posthog.team")),
55+
],
56+
options={
57+
"indexes": [
58+
models.Index(fields=["team", "short_id"], name="ee_agentart_team_id_d05587_idx"),
59+
models.Index(fields=["team", "conversation", "created_at"], name="ee_agentart_team_id_e402e1_idx"),
60+
],
61+
},
62+
),
63+
migrations.AddConstraint(
64+
model_name="agentartifact",
65+
constraint=models.UniqueConstraint(fields=("team", "short_id"), name="unique_team_short_id"),
66+
),
67+
]

‎ee/migrations/max_migration.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
0030_singlesessionsummary_distinct_id_and_more
1+
0031_agentartifact

‎ee/models/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from .assistant import (
2+
AgentArtifact,
23
Conversation,
34
ConversationCheckpoint,
45
ConversationCheckpointBlob,
@@ -16,6 +17,7 @@
1617

1718
__all__ = [
1819
"AccessControl",
20+
"AgentArtifact",
1921
"ConversationCheckpoint",
2022
"ConversationCheckpointBlob",
2123
"ConversationCheckpointWrite",

‎ee/models/assistant.py‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,22 @@
1+
import string
2+
import secrets
13
from datetime import timedelta
24

3-
from django.db import models
5+
from django.db import IntegrityError, models
46
from django.utils import timezone
57

68
from posthog.models.team.team import Team
79
from posthog.models.user import User
8-
from posthog.models.utils import UUIDTModel
10+
from posthog.models.utils import CreatedMetaFields, DeletedMetaFields, UpdatedMetaFields, UUIDModel, UUIDTModel
11+
12+
13+
def generate_short_id():
14+
"""Generate securely random 4 characters long alphanumeric ID.
15+
16+
With team-scoped uniqueness, 4 characters (62^4 = 14.7M combinations)
17+
is sufficient to avoid collisions within a single team.
18+
"""
19+
return "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(4))
920

1021

1122
class Conversation(UUIDTModel):
@@ -183,3 +194,36 @@ def answers_left(self) -> int:
183194
if self.initial_text.endswith("\nAnswer:"):
184195
answers_given -= 1
185196
return MAX_ONBOARDING_QUESTIONS - answers_given
197+
198+
199+
class AgentArtifact(UUIDModel, CreatedMetaFields, UpdatedMetaFields, DeletedMetaFields):
200+
class Type(models.TextChoices):
201+
VISUALIZATION = "visualization", "Visualization"
202+
NOTEBOOK = "notebook", "Notebook"
203+
204+
short_id = models.CharField(max_length=4, default=generate_short_id)
205+
name = models.CharField(max_length=400)
206+
type = models.CharField(max_length=50, choices=Type.choices)
207+
data = models.JSONField(help_text="Artifact content. Structure depends on artifact type.")
208+
conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="artifacts")
209+
team = models.ForeignKey(Team, on_delete=models.CASCADE)
210+
211+
class Meta:
212+
indexes = [
213+
models.Index(fields=["team", "short_id"]),
214+
models.Index(fields=["team", "conversation", "created_at"]),
215+
]
216+
constraints = [
217+
models.UniqueConstraint(fields=["team", "short_id"], name="unique_team_short_id"),
218+
]
219+
220+
def save(self, *args, **kwargs):
221+
max_retries = 5
222+
for attempt in range(max_retries):
223+
try:
224+
return super().save(*args, **kwargs)
225+
except IntegrityError as e:
226+
if "short_id" in str(e) and attempt < max_retries - 1:
227+
self.short_id = generate_short_id()
228+
else:
229+
raise

0 commit comments

Comments
 (0)