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