feat(max): add keyword search across taxonomy properties

ee/hogai/chat_agent/query_planner/nodes.py

@@ -50,6 +50,7 @@
     retrieve_action_property_values,
     retrieve_event_properties,
     retrieve_event_property_values,
+    search_properties_tool,
 )
 
 
@@ -163,6 +164,7 @@ def _get_model(self, state: AssistantState):
                 retrieve_event_property_values,
                 retrieve_action_property_values,
                 dynamic_retrieve_entity_property_values,
+                search_properties_tool,
                 ask_user_for_help,
                 final_answer,
             ],
@@ -318,6 +320,8 @@ def _handle_tool(self, input: TaxonomyAgentTool, toolkit: TaxonomyAgentToolkit)
             output = toolkit.retrieve_entity_properties(input.arguments.entity)  # type: ignore
         elif input.name == "retrieve_entity_property_values":
             output = toolkit.retrieve_entity_property_values(input.arguments.entity, input.arguments.property_name)  # type: ignore
+        elif input.name == "search_properties":
+            output = toolkit.search_properties(input.arguments.term)  # type: ignore
         else:
             output = toolkit.handle_incorrect_response(input)
         return output

ee/hogai/chat_agent/query_planner/prompts.py

@@ -181,7 +181,7 @@
 """.strip()
 
 ITERATION_LIMIT_PROMPT = """
-The tool has reached the maximum number of iterations, a security measure to prevent infinite loops. To create this insight, you must request additional information from the user, such as specific events, properties, or property values.
+The tool has reached the maximum number of iterations, a security measure to prevent infinite loops. Before giving up, make sure you tried `search_properties` to look for the concept by keyword across persons, sessions, groups, and events - a property not found by name may still exist under different wording. If it's still not found, tell the user this specifically (rather than a generic "be more specific") and request additional information, such as the exact property/event name, a related value to filter on instead, or confirmation that the concept isn't tracked yet.
 """.strip()
 
 ACTIONS_EXPLANATION_PROMPT = """

ee/hogai/chat_agent/query_planner/test/test_toolkit.py

@@ -481,6 +481,97 @@ def test_retrieve_event_property_values_hides_restricted_property(self, mock_res
         result = toolkit.retrieve_event_or_action_property_values("event1", "$browser")
         self.assertIn("does not exist", result)
 
+    def test_retrieve_entity_properties_marks_truncation_for_person_and_group(self):
+        # Regression guard: the group branch sliced on `max_properties` and the person branch didn't,
+        # and neither told the model the list was cut off, so a property past the slice looked like it
+        # didn't exist rather than "there's more, go search for it".
+        PropertyDefinition.objects.create(
+            team=self.team, type=PropertyDefinition.Type.PERSON, name="prop_a", property_type="String"
+        )
+        PropertyDefinition.objects.create(
+            team=self.team, type=PropertyDefinition.Type.PERSON, name="prop_b", property_type="String"
+        )
+        create_group_type_mapping_without_created_at(
+            team=self.team, project_id=self.team.project_id, group_type_index=0, group_type="group"
+        )
+        invalidate_group_types_cache(self.team.project_id)
+        PropertyDefinition.objects.create(
+            team=self.team,
+            type=PropertyDefinition.Type.GROUP,
+            group_type_index=0,
+            name="prop_a",
+            property_type="String",
+        )
+        PropertyDefinition.objects.create(
+            team=self.team,
+            type=PropertyDefinition.Type.GROUP,
+            group_type_index=0,
+            name="prop_b",
+            property_type="String",
+        )
+        toolkit = DummyToolkit(self.team, self.user)
+
+        truncated_person = toolkit.retrieve_entity_properties("person", max_properties=1)
+        self.assertIn("search_properties", truncated_person)
+
+        truncated_group = toolkit.retrieve_entity_properties("group", max_properties=1)
+        self.assertIn("search_properties", truncated_group)
+
+        full_person = toolkit.retrieve_entity_properties("person", max_properties=500)
+        self.assertNotIn("search_properties", full_person)
+
+    def test_search_properties_matches_by_name_across_entities(self):
+        PropertyDefinition.objects.create(
+            team=self.team, type=PropertyDefinition.Type.PERSON, name="is_internal_user", property_type="Boolean"
+        )
+        create_group_type_mapping_without_created_at(
+            team=self.team, project_id=self.team.project_id, group_type_index=0, group_type="organization"
+        )
+        invalidate_group_types_cache(self.team.project_id)
+        PropertyDefinition.objects.create(
+            team=self.team,
+            type=PropertyDefinition.Type.GROUP,
+            group_type_index=0,
+            name="is_internal_org",
+            property_type="Boolean",
+        )
+        PropertyDefinition.objects.create(
+            team=self.team, type=PropertyDefinition.Type.EVENT, name="internal_flag", property_type="Boolean"
+        )
+        toolkit = DummyToolkit(self.team, self.user)
+
+        result = toolkit.search_properties("internal")
+
+        self.assertIn("[person] is_internal_user", result)
+        self.assertIn("[organization] is_internal_org", result)
+        self.assertIn("[event] internal_flag", result)
+
+    def test_search_properties_matches_by_description(self):
+        # A property named without the searched-for word is still findable if a human documented it,
+        # matching the read_taxonomy tool's ability to search descriptions, not just names.
+        from ee.models.property_definition import EnterprisePropertyDefinition
+
+        EnterprisePropertyDefinition.objects.create(
+            team=self.team,
+            type=PropertyDefinition.Type.PERSON,
+            name="staff_flag",
+            property_type="Boolean",
+            description="Marks internal PostHog employees",
+        )
+        toolkit = DummyToolkit(self.team, self.user)
+
+        result = toolkit.search_properties("internal")
+
+        self.assertIn("[person] staff_flag", result)
+
+    def test_search_properties_no_match_offers_next_steps(self):
+        toolkit = DummyToolkit(self.team, self.user)
+
+        result = toolkit.search_properties("no_such_concept_anywhere")
+
+        self.assertIn("No property names or descriptions matched", result)
+        self.assertIn("cohort", result)
+
 
 class TestFinalAnswerTool(BaseTest):
     def test_normalize_plan(self):

ee/hogai/chat_agent/query_planner/toolkit.py

@@ -3,6 +3,8 @@
 from functools import cached_property
 from typing import Literal, Optional, Union, cast
 
+from django.db.models import Q
+
 from pydantic import BaseModel, field_validator
 
 from posthog.schema import (
@@ -36,6 +38,7 @@
     retrieve_entity_property_values,
     retrieve_event_properties,
     retrieve_event_property_values,
+    search_properties as search_properties_tool,
 )
 from ee.hogai.chat_agent.taxonomy.virtual_properties import (
     PropertyDefinitionOrVirtual,
@@ -80,6 +83,7 @@ def normalize_plan(cls, plan: str) -> str:
     retrieve_event_property_values,
     retrieve_action_property_values,
     retrieve_entity_property_values,
+    search_properties_tool,
     ask_user_for_help,
     final_answer,
 ]
@@ -212,14 +216,13 @@ def retrieve_entity_properties(self, entity: str, max_properties: int = 500) ->
         if entity not in ("person", "session", *[g["group_type"] for g in self._groups]):
             return f"Entity {entity} does not exist in the taxonomy."
 
+        truncated = False
         if entity == "person":
             restricted = self._restricted_property_names(PropertyDefinition.Type.PERSON)
+            person_qs = PropertyDefinition.objects.filter(team=self._team, type=PropertyDefinition.Type.PERSON)
+            truncated = person_qs.count() > max_properties
             stored_props = [
-                p
-                for p in PropertyDefinition.objects.filter(
-                    team=self._team, type=PropertyDefinition.Type.PERSON
-                ).values_list("name", "property_type")
-                if p[0] not in restricted
+                p for p in person_qs.values_list("name", "property_type")[:max_properties] if p[0] not in restricted
             ]
             stored_props += list_virtual_properties(
                 "person_properties", exclude={name for name, _ in stored_props} | restricted
@@ -244,12 +247,12 @@ def retrieve_entity_properties(self, entity: str, max_properties: int = 500) ->
             if group_type_index is None:
                 return f"Group {entity} does not exist in the taxonomy."
             restricted = self._restricted_property_names(PropertyDefinition.Type.GROUP)
+            group_qs = PropertyDefinition.objects.filter(
+                team=self._team, type=PropertyDefinition.Type.GROUP, group_type_index=group_type_index
+            )
+            truncated = group_qs.count() > max_properties
             stored_props = [
-                p
-                for p in PropertyDefinition.objects.filter(
-                    team=self._team, type=PropertyDefinition.Type.GROUP, group_type_index=group_type_index
-                ).values_list("name", "property_type")[:max_properties]
-                if p[0] not in restricted
+                p for p in group_qs.values_list("name", "property_type")[:max_properties] if p[0] not in restricted
             ]
             stored_props += list_virtual_properties("groups", exclude={name for name, _ in stored_props} | restricted)
             stored_descriptions = self._get_stored_property_descriptions(
@@ -262,7 +265,113 @@ def retrieve_entity_properties(self, entity: str, max_properties: int = 500) ->
         if not props:
             return f"Properties do not exist in the taxonomy for the entity {entity}."
 
-        return format_prompt_string(PROPERTIES_EXAMPLE_PROMPT, result=self._generate_properties_output(props))
+        result = format_prompt_string(PROPERTIES_EXAMPLE_PROMPT, result=self._generate_properties_output(props))
+        if truncated:
+            result += (
+                f"\n\nNOTE: This entity has more than {max_properties} properties, and the list above was cut off "
+                "at that limit. A property not listed here may still exist — use the `search_properties` query kind "
+                "to look for it by name or description instead of concluding it's missing."
+            )
+        return result
+
+    def search_properties(self, term: str, max_results: int = 25) -> str:
+        """
+        Search property names and descriptions for a keyword across persons, sessions, groups, and events.
+
+        Use this when you know the concept you're looking for (e.g. "internal user", "subscription plan")
+        but don't know which entity or event it lives on, instead of guessing an entity and dumping its
+        whole property list.
+        """
+        term = term.strip()
+        if not term:
+            return "Provide a non-empty search term."
+
+        matches: list[tuple[str, str, str | None]] = []  # (scope, name, description)
+
+        restricted_person = self._restricted_property_names(PropertyDefinition.Type.PERSON)
+        matches += [
+            ("person", name, description)
+            for name, description in self._search_property_definitions(
+                PropertyDefinition.Type.PERSON, term, exclude=restricted_person
+            )
+        ]
+
+        for prop_name, prop in CORE_FILTER_DEFINITIONS_BY_GROUP["session_properties"].items():
+            haystack = f"{prop_name} {prop.get('label') or ''} {prop.get('description') or ''}".lower()
+            if term.lower() in haystack:
+                matches.append(("session", prop_name, prop.get("description")))
+
+        restricted_group = self._restricted_property_names(PropertyDefinition.Type.GROUP)
+        for group in self._groups:
+            group_type_index = group["group_type_index"]
+            matches += [
+                (group["group_type"], name, description)
+                for name, description in self._search_property_definitions(
+                    PropertyDefinition.Type.GROUP, term, exclude=restricted_group, group_type_index=group_type_index
+                )
+            ]
+
+        restricted_event = self._restricted_property_names(PropertyDefinition.Type.EVENT)
+        matches += [
+            ("event", name, description)
+            for name, description in self._search_property_definitions(
+                PropertyDefinition.Type.EVENT, term, exclude=restricted_event
+            )
+        ]
+
+        if not matches:
+            return (
+                f'No property names or descriptions matched "{term}" across persons, sessions, groups, or events. '
+                "This doesn't necessarily mean the concept isn't tracked: it may be captured under a differently "
+                "worded property, computed from other properties, or not instrumented at all. Consider asking the "
+                "user for the exact property name, filtering by a related value (e.g. an email domain), or checking "
+                "if it's covered by an existing cohort."
+            )
+
+        truncated = len(matches) > max_results
+        matches = matches[:max_results]
+
+        output_parts = [f'Properties matching "{term}":']
+        for scope, name, description in matches:
+            suffix = f" – {description.replace(chr(10), ' ')}" if description else ""
+            output_parts.append(f"- [{scope}] {name}{suffix}")
+        if truncated:
+            output_parts.append(f"\n...and more matches were cut off at {max_results} results. Narrow your term.")
+
+        return "\n".join(output_parts)
+
+    def _search_property_definitions(
+        self,
+        property_type: "PropertyDefinition.Type",
+        term: str,
+        exclude: set[str],
+        group_type_index: int | None = None,
+    ) -> list[tuple[str, str | None]]:
+        """Search stored property names/descriptions for a team by keyword, name matches first."""
+        qs = PropertyDefinition.objects.filter(team=self._team, type=property_type)
+        if group_type_index is not None:
+            qs = qs.filter(group_type_index=group_type_index)
+        name_matches = list(qs.filter(name__icontains=term).values_list("name", flat=True))
+
+        description_matches: list[str] = []
+        if EE_AVAILABLE:
+            from ee.models.property_definition import (
+                EnterprisePropertyDefinition,  # noqa: PLC0415 — EE-only model, keep off the OSS import path
+            )
+
+            desc_qs = EnterprisePropertyDefinition.objects.filter(team=self._team, type=property_type).filter(
+                Q(description__icontains=term)
+            )
+            if group_type_index is not None:
+                desc_qs = desc_qs.filter(group_type_index=group_type_index)
+            description_matches = list(desc_qs.exclude(name__in=name_matches).values_list("name", flat=True))
+
+        names = [name for name in [*name_matches, *description_matches] if name not in exclude]
+        if not names:
+            return []
+
+        descriptions = self._get_stored_property_descriptions(property_type, names, group_type_index)
+        return [(name, descriptions.get(name)) for name in names]
 
     def _retrieve_event_or_action_taxonomy(self, event_name_or_action_id: str | int):
         is_event = isinstance(event_name_or_action_id, str)

ee/hogai/chat_agent/taxonomy/prompts.py

@@ -70,6 +70,7 @@
 3. **Tool Workflow**:
    - **For ENTITY properties** (person, session, organization, groups): Use `retrieve_entity_properties` and `retrieve_entity_property_values`
    - **For EVENT properties** (properties of specific events like pageview, signup, etc.): Use `retrieve_event_properties` and `retrieve_event_property_values`
+   - **When you know the concept but not which entity or event it lives on** (e.g. "internal user", "subscription plan"): Use `search_properties` with a keyword instead of guessing an entity and dumping its whole property list.
    - Use `ask_user_for_help` when you need clarification
    - Use `final_answer` only when you have complete information
    - *CRITICAL*: NEVER use entity tools for event properties. NEVER use event tools for entity properties.
@@ -113,6 +114,10 @@
 You must fix the exception and try again.
 """.strip()
 
-ITERATION_LIMIT_PROMPT = """I've tried several approaches but haven't been able to find the right options. Could you please be more specific about what kind of properties you're looking for? For example:
-- What type of events or actions are you interested in?
-- Are you looking for specific values or ranges?"""
+ITERATION_LIMIT_PROMPT = """I looked through the events, entities, and properties available in your project (including a keyword search) but couldn't confidently match what you're asking for to something that's tracked. This could mean the concept isn't captured yet, or it's just named differently than I guessed.
+
+Here's what would help me move forward:
+- Give me the exact property or event name, if you know it
+- Describe a related value I could filter on instead (e.g. an email domain or account ID)
+- Point me to an existing cohort that already captures this group
+- If it's genuinely not tracked yet, it may need to be instrumented before it can be used in an insight"""

ee/hogai/chat_agent/taxonomy/test/test_toolkit.py

@@ -86,12 +86,14 @@ class TestModel(BaseModel):
             ("retrieve_entity_property_values", {"entity": "person", "property_name": "email"}, "mocked"),
             ("retrieve_event_properties", {"event_name": "test_event"}, "mocked"),
             ("retrieve_event_property_values", {"event_name": "test_event", "property_name": "$browser"}, "mocked"),
+            ("search_properties", {"term": "internal"}, "mocked"),
         ]
     )
     @patch.object(DummyToolkit, "retrieve_entity_properties_parallel", return_value={"person": "mocked"})
     @patch.object(DummyToolkit, "retrieve_entity_property_values", return_value={"person": ["mocked"]})
     @patch.object(DummyToolkit, "retrieve_event_or_action_properties_parallel", return_value={"test_event": "mocked"})
     @patch.object(DummyToolkit, "retrieve_event_or_action_property_values", return_value={"test_event": ["mocked"]})
+    @patch.object(DummyToolkit, "search_properties", return_value="mocked")
     async def test_handle_tools(self, tool_name, tool_args, expected_result, *mocks):
         class Arguments(BaseModel):
             pass