feat: output property description

ee/hogai/trends/test/test_toolkit.py

@@ -71,7 +71,7 @@ def test_retrieve_entity_properties(self):
         )
         self.assertEqual(
             toolkit.retrieve_entity_properties("person"),
-            "<properties><String><name>test</name><br /></String></properties>",
+            "<properties><String><prop><name>test</name></prop></String></properties>",
         )
 
         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"),
-            "<properties><Numeric><name>test</name><br /></Numeric></properties>",
+            "<properties><Numeric><prop><name>test</name></prop></Numeric></properties>",
         )
 
         self.assertNotEqual(
@@ -203,19 +203,19 @@ def test_retrieve_event_properties(self):
         prompt = toolkit.retrieve_event_properties("event1")
 
         self.assertIn(
-            "<Numeric><name>id</name><br /></Numeric>",
+            "<Numeric><prop><name>id</name></prop></Numeric>",
             prompt,
         )
         self.assertIn(
-            "<String><name>$browser</name><br /></String>",
+            "<String><prop><name>$browser</name><description>Name of the browser the user has used.</description></prop></String>",
             prompt,
         )
         self.assertIn(
-            "<DateTime><name>date</name><br /></DateTime>",
+            "<DateTime><prop><name>date</name></prop></DateTime>",
             prompt,
         )
         self.assertIn(
-            "<Boolean><name>bool</name><br /></Boolean>",
+            "<Boolean><prop><name>bool</name></prop></Boolean>",
             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)

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