diff --git a/ee/hogai/trends/test/test_toolkit.py b/ee/hogai/trends/test/test_toolkit.py
index 12cd086b033c5..35c98d37b9556 100644
--- a/ee/hogai/trends/test/test_toolkit.py
+++ b/ee/hogai/trends/test/test_toolkit.py
@@ -71,7 +71,7 @@ def test_retrieve_entity_properties(self):
)
self.assertEqual(
toolkit.retrieve_entity_properties("person"),
- "test
",
+ "test",
)
GroupTypeMapping.objects.create(team=self.team, group_type_index=0, group_type="group")
@@ -80,7 +80,7 @@ def test_retrieve_entity_properties(self):
)
self.assertEqual(
toolkit.retrieve_entity_properties("group"),
- "test
",
+ "test",
)
self.assertNotEqual(
@@ -203,19 +203,19 @@ def test_retrieve_event_properties(self):
prompt = toolkit.retrieve_event_properties("event1")
self.assertIn(
- "id
",
+ "id",
prompt,
)
self.assertIn(
- "$browser
",
+ "$browserName of the browser the user has used.",
prompt,
)
self.assertIn(
- "date
",
+ "date",
prompt,
)
self.assertIn(
- "bool
",
+ "bool",
prompt,
)
@@ -233,3 +233,12 @@ def test_retrieve_event_property_values(self):
self.assertEqual(
toolkit.retrieve_event_property_values("event1", "date"), f'"{datetime(2024, 1, 1).isoformat()}"'
)
+
+ def test_enrich_props_with_descriptions(self):
+ toolkit = TrendsAgentToolkit(self.team)
+ res = toolkit._enrich_props_with_descriptions("event", [("$geoip_city_name", "String")])
+ self.assertEqual(len(res), 1)
+ prop, type, description = res[0]
+ self.assertEqual(prop, "$geoip_city_name")
+ self.assertEqual(type, "String")
+ self.assertIsNotNone(description)
diff --git a/ee/hogai/trends/toolkit.py b/ee/hogai/trends/toolkit.py
index a2f438756f3d1..da47e4df3880d 100644
--- a/ee/hogai/trends/toolkit.py
+++ b/ee/hogai/trends/toolkit.py
@@ -1,5 +1,6 @@
import json
import xml.etree.ElementTree as ET
+from collections.abc import Iterable
from functools import cached_property
from textwrap import dedent
from typing import Any, Literal, Optional, TypedDict, Union
@@ -238,23 +239,39 @@ def render_text_description(self) -> str:
descriptions.append(description)
return "\n".join(descriptions)
- def _generate_properties_xml(self, children: list[tuple[str, str | None]]):
+ def _generate_properties_xml(self, children: list[tuple[str, str | None, str | None]]):
root = ET.Element("properties")
- property_types = {property_type for _, property_type in children if property_type is not None}
- property_type_to_tag = {property_type: ET.SubElement(root, property_type) for property_type in property_types}
+ property_type_to_tag = {}
- for name, property_type in children:
+ for name, property_type, description in children:
# Do not include properties that are ambiguous.
if property_type is None:
continue
+ if property_type not in property_type_to_tag:
+ property_type_to_tag[property_type] = ET.SubElement(root, property_type)
type_tag = property_type_to_tag[property_type]
- ET.SubElement(type_tag, "name").text = name
- # Add a line break between names. Doubtful that it does anything.
- ET.SubElement(type_tag, "br")
+ prop = ET.SubElement(type_tag, "prop")
+ ET.SubElement(prop, "name").text = name
+ if description:
+ ET.SubElement(prop, "description").text = description
return ET.tostring(root, encoding="unicode")
+ def _enrich_props_with_descriptions(self, entity: str, props: Iterable[tuple[str, str | None]]):
+ enriched_props = []
+ mapping = {
+ "session": hardcoded_prop_defs["session_properties"],
+ "person": hardcoded_prop_defs["person_properties"],
+ "event": hardcoded_prop_defs["event_properties"],
+ }
+ for prop_name, prop_type in props:
+ description = None
+ if entity in mapping:
+ description = mapping[entity].get(prop_name, {}).get("description")
+ enriched_props.append((prop_name, prop_type, description))
+ return enriched_props
+
def retrieve_entity_properties(self, entity: str) -> str:
"""
Retrieve properties for an entitiy like person, session, or one of the groups.
@@ -266,14 +283,17 @@ def retrieve_entity_properties(self, entity: str) -> str:
qs = PropertyDefinition.objects.filter(team=self._team, type=PropertyDefinition.Type.PERSON).values_list(
"name", "property_type"
)
- props = list(qs)
+ props = self._enrich_props_with_descriptions("person", qs)
elif entity == "session":
# Session properties are not in the DB.
- props = [
- (prop_name, prop["type"])
- for prop_name, prop in hardcoded_prop_defs["session_properties"].items()
- if prop.get("type") is not None
- ]
+ props = self._enrich_props_with_descriptions(
+ "session",
+ [
+ (prop_name, prop["type"])
+ for prop_name, prop in hardcoded_prop_defs["session_properties"].items()
+ if prop.get("type") is not None
+ ],
+ )
else:
group_type_index = next(
(group.group_type_index for group in self.groups if group.group_type == entity), None
@@ -283,7 +303,7 @@ def retrieve_entity_properties(self, entity: str) -> str:
qs = PropertyDefinition.objects.filter(
team=self._team, type=PropertyDefinition.Type.GROUP, group_type_index=group_type_index
).values_list("name", "property_type")
- props = list(qs)
+ props = self._enrich_props_with_descriptions(entity, qs)
return self._generate_properties_xml(props)
@@ -305,15 +325,13 @@ def retrieve_event_properties(self, event_name: str) -> str:
team=self._team, type=PropertyDefinition.Type.EVENT, name__in=[item.property for item in response.results]
)
property_to_type = {property_definition.name: property_definition.property_type for property_definition in qs}
-
- return self._generate_properties_xml(
- [
- (item.property, property_to_type.get(item.property))
- for item in response.results
- # Exclude properties that exist in the taxonomy, but don't have a type.
- if item.property in property_to_type
- ]
- )
+ props = [
+ (item.property, property_to_type.get(item.property))
+ for item in response.results
+ # Exclude properties that exist in the taxonomy, but don't have a type.
+ if item.property in property_to_type
+ ]
+ return self._generate_properties_xml(self._enrich_props_with_descriptions("event", props))
def _format_property_values(
self, sample_values: list, sample_count: Optional[int] = 0, format_as_string: bool = False