From 604dac0aff1c2a168c20769f1fdb7bb89aeb0e2e Mon Sep 17 00:00:00 2001 From: "posthog[bot]" <206114724+posthog[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 15:59:44 +0000 Subject: [PATCH] feat(max): add keyword search across taxonomy properties Add a `search_properties` query kind to the taxonomy tools so Max can search property names and descriptions by keyword across persons, sessions, groups, and events in one call, instead of guessing an entity and dumping its whole property list. Also fixes an inconsistency where the person branch of `retrieve_entity_properties` had no `max_properties` slice (unlike the group branch), and neither branch told the model when its list was truncated - so a property past the cutoff looked like it didn't exist. Both branches now truncate consistently and note it when they do. Updates the give-up message shown after the iteration limit is hit to distinguish "I looked and didn't find it" from a generic "be more specific", and suggests concrete next steps (exact name, related filter value, existing cohort, or instrumentation). Generated-By: PostHog Code Task-Id: 72f373b9-e2f0-4dc8-8ed3-9a888f47f915 --- ee/hogai/chat_agent/query_planner/nodes.py | 4 + ee/hogai/chat_agent/query_planner/prompts.py | 2 +- .../query_planner/test/test_toolkit.py | 91 ++++++++++++ ee/hogai/chat_agent/query_planner/toolkit.py | 131 ++++++++++++++++-- ee/hogai/chat_agent/taxonomy/prompts.py | 11 +- .../chat_agent/taxonomy/test/test_toolkit.py | 2 + ee/hogai/chat_agent/taxonomy/toolkit.py | 110 +++++++++++++++ ee/hogai/chat_agent/taxonomy/tools.py | 12 ++ ee/hogai/tools/read_taxonomy/core.py | 12 ++ .../tools/read_taxonomy/test/test_tool.py | 13 ++ ee/hogai/tools/read_taxonomy/tool.py | 3 + 11 files changed, 376 insertions(+), 15 deletions(-) diff --git a/ee/hogai/chat_agent/query_planner/nodes.py b/ee/hogai/chat_agent/query_planner/nodes.py index 7e2cc3340943..4bf945b38cd5 100644 --- a/ee/hogai/chat_agent/query_planner/nodes.py +++ b/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 diff --git a/ee/hogai/chat_agent/query_planner/prompts.py b/ee/hogai/chat_agent/query_planner/prompts.py index 449f88895a65..055877131882 100644 --- a/ee/hogai/chat_agent/query_planner/prompts.py +++ b/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 = """ diff --git a/ee/hogai/chat_agent/query_planner/test/test_toolkit.py b/ee/hogai/chat_agent/query_planner/test/test_toolkit.py index d596da8a2cc3..25ca3c891ce9 100644 --- a/ee/hogai/chat_agent/query_planner/test/test_toolkit.py +++ b/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): diff --git a/ee/hogai/chat_agent/query_planner/toolkit.py b/ee/hogai/chat_agent/query_planner/toolkit.py index a7cd99c52b37..aa20569b3448 100644 --- a/ee/hogai/chat_agent/query_planner/toolkit.py +++ b/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) diff --git a/ee/hogai/chat_agent/taxonomy/prompts.py b/ee/hogai/chat_agent/taxonomy/prompts.py index 0c5025956a45..794f8d7f5ad5 100644 --- a/ee/hogai/chat_agent/taxonomy/prompts.py +++ b/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""" diff --git a/ee/hogai/chat_agent/taxonomy/test/test_toolkit.py b/ee/hogai/chat_agent/taxonomy/test/test_toolkit.py index f2c6e9216679..7949a689ac28 100644 --- a/ee/hogai/chat_agent/taxonomy/test/test_toolkit.py +++ b/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 diff --git a/ee/hogai/chat_agent/taxonomy/toolkit.py b/ee/hogai/chat_agent/taxonomy/toolkit.py index 0c67f0d583eb..c8b514c0a14d 100644 --- a/ee/hogai/chat_agent/taxonomy/toolkit.py +++ b/ee/hogai/chat_agent/taxonomy/toolkit.py @@ -4,6 +4,8 @@ from typing import Optional, Union, cast from uuid import uuid4 +from django.db.models import Q + from langchain_core.agents import AgentAction from langchain_core.runnables import RunnableConfig from pydantic import BaseModel @@ -64,6 +66,7 @@ get_dynamic_entity_tools, retrieve_event_properties, retrieve_event_property_values, + search_properties, ) @@ -342,6 +345,100 @@ def handle_incorrect_response(self, response: BaseModel) -> str: """ return response.model_dump_json() + @database_sync_to_async(thread_sensitive=False) + def _search_property_definitions( + self, + property_type: PropertyDefinition.Type, + term: str, + exclude: set[str], + group_type_index: int | None = None, + ) -> list[str]: + """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)) + + return [name for name in [*name_matches, *description_matches] if name not in exclude] + + async 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]] = [] + + restricted_person = await self._restricted_property_names(PropertyDefinition.Type.PERSON) + person_names = await self._search_property_definitions(PropertyDefinition.Type.PERSON, term, restricted_person) + if person_names: + descriptions = await self._get_stored_property_descriptions(PropertyDefinition.Type.PERSON, person_names) + matches += [("person", name, descriptions.get(name)) for name in person_names] + + 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 = await self._restricted_property_names(PropertyDefinition.Type.GROUP) + groups = await self._get_groups() + for group in groups: + group_type_index = group["group_type_index"] + group_names = await self._search_property_definitions( + PropertyDefinition.Type.GROUP, term, restricted_group, group_type_index + ) + if group_names: + descriptions = await self._get_stored_property_descriptions( + PropertyDefinition.Type.GROUP, group_names, group_type_index + ) + matches += [(group["group_type"], name, descriptions.get(name)) for name in group_names] + + restricted_event = await self._restricted_property_names(PropertyDefinition.Type.EVENT) + event_names = await self._search_property_definitions(PropertyDefinition.Type.EVENT, term, restricted_event) + if event_names: + descriptions = await self._get_stored_property_descriptions(PropertyDefinition.Type.EVENT, event_names) + matches += [("event", name, descriptions.get(name)) for name in event_names] + + 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 get_tools(self) -> list: """Get all tools (default + custom). Override in subclasses to add custom tools.""" try: @@ -361,6 +458,7 @@ def _get_default_tools(self) -> list: dynamic_retrieve_entity_properties, retrieve_event_property_values, dynamic_retrieve_entity_property_values, + search_properties, ask_user_for_help, ] @@ -916,6 +1014,7 @@ def _collect_tools(self, tool_metadata: dict[str, list[tuple[TaxonomyTool, str]] "entity_mapping": {}, # entity -> [tool_call_id] "event_prop_mapping": {}, # (event, property) -> [tool_call_id] "event_mapping": {}, # event -> [tool_call_id] + "search_properties": [], # [(term, tool_call_id)] } for tool_name, tool_inputs in tool_metadata.items(): @@ -945,6 +1044,10 @@ def _collect_tools(self, tool_metadata: dict[str, list[tuple[TaxonomyTool, str]] event_name = tool_input.arguments.event_name # type: ignore result["event_properties"].append(event_name) result["event_mapping"].setdefault(event_name, []).append(tool_call_id) + + elif tool_name == "search_properties": + term = tool_input.arguments.term # type: ignore + result["search_properties"].append((term, tool_call_id)) else: raise TaxonomyToolNotFoundError(f"Tool {tool_name} not found in taxonomy toolkit.") @@ -996,6 +1099,13 @@ async def _execute_tools(self, collected_tools: dict) -> dict[str, str]: for tool_call_id in collected_tools["event_mapping"].get(event_name, []): results[tool_call_id] = result + if collected_tools["search_properties"]: + search_results = await asyncio.gather( + *[self.search_properties(term) for term, _ in collected_tools["search_properties"]] + ) + for (_, tool_call_id), result in zip(collected_tools["search_properties"], search_results): + results[tool_call_id] = result + return results async def handle_tools(self, tool_metadata: dict[str, list[tuple[TaxonomyTool, str]]]) -> dict[str, str]: diff --git a/ee/hogai/chat_agent/taxonomy/tools.py b/ee/hogai/chat_agent/taxonomy/tools.py index 3f693e7ab804..ce77cb729f71 100644 --- a/ee/hogai/chat_agent/taxonomy/tools.py +++ b/ee/hogai/chat_agent/taxonomy/tools.py @@ -69,6 +69,17 @@ class retrieve_entity_property_values(BaseModel): property_name: str = Field(..., description="The name of the property that you want to retrieve values for.") +class search_properties(BaseModel): + """ + Use this tool to search property names and descriptions for a keyword across persons, sessions, groups, and + events in one call. 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: str = Field(..., description="Keyword or short phrase to search for, e.g. 'internal user' or 'plan'.") + + class ask_user_for_help(BaseModel): """ Use this tool to ask a question to the user. Your question must be concise and clear. @@ -132,6 +143,7 @@ class base_final_answer(BaseModel, Generic[OutputType]): retrieve_entity_properties, retrieve_entity_property_values, retrieve_event_property_values, + search_properties, ask_user_for_help, ] diff --git a/ee/hogai/tools/read_taxonomy/core.py b/ee/hogai/tools/read_taxonomy/core.py index 0a441505def2..cd23f57f3a69 100644 --- a/ee/hogai/tools/read_taxonomy/core.py +++ b/ee/hogai/tools/read_taxonomy/core.py @@ -63,6 +63,15 @@ class ReadActionSamplePropertyValues(BaseModel): property_name: str = Field(description="Verified property name of an action.") +class SearchTaxonomyProperties(BaseModel): + """Searches property names and descriptions for a keyword across persons, sessions, groups, and events in one call. + 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.""" + + kind: Literal["search_properties"] = "search_properties" + term: str = Field(description="Keyword or short phrase to search for, e.g. 'internal user' or 'plan'.") + + ReadTaxonomyQuery = Union[ ReadEvents, ReadEventProperties, @@ -71,6 +80,7 @@ class ReadActionSamplePropertyValues(BaseModel): ReadEntitySamplePropertyValues, ReadActionProperties, ReadActionSamplePropertyValues, + SearchTaxonomyProperties, ] @@ -122,5 +132,7 @@ def execute_taxonomy_query(query: ReadTaxonomyQuery, toolkit: TaxonomyAgentToolk return result case ReadEntitySamplePropertyValues(): return toolkit.retrieve_entity_property_values(query.entity, query.property_name) + case SearchTaxonomyProperties(): + return toolkit.search_properties(query.term) case _: raise ValueError(f"Invalid query type: The query structure '{type(query).__name__}' is not recognized.") diff --git a/ee/hogai/tools/read_taxonomy/test/test_tool.py b/ee/hogai/tools/read_taxonomy/test/test_tool.py index c4d3265833fb..ad2c6effa17b 100644 --- a/ee/hogai/tools/read_taxonomy/test/test_tool.py +++ b/ee/hogai/tools/read_taxonomy/test/test_tool.py @@ -14,6 +14,7 @@ ReadEntityProperties, ReadEventProperties, ReadEvents, + SearchTaxonomyProperties, execute_taxonomy_query, ) from ee.hogai.tools.read_taxonomy.tool import ReadTaxonomyTool @@ -133,6 +134,18 @@ def test_non_person_entity_properties_exclude_dynamic_hint(self, mock_toolkit_cl self.assertNotIn(DYNAMIC_PERSON_PROPERTIES_HINT, result) + @patch("ee.hogai.tools.read_taxonomy.core.TaxonomyAgentToolkit") + def test_search_properties_dispatches_to_toolkit(self, mock_toolkit_class): + # Guards the match/case wiring: forgetting to add a branch for the new query kind would + # fall through to the `_` case and raise ValueError instead of running the search. + mock_toolkit = mock_toolkit_class.return_value + mock_toolkit.search_properties.return_value = 'Properties matching "internal":\n- [person] is_internal_user' + + result = execute_taxonomy_query(SearchTaxonomyProperties(term="internal"), mock_toolkit, self.team, self.user) + + mock_toolkit.search_properties.assert_called_once_with("internal") + self.assertIn("is_internal_user", result) + @patch("ee.hogai.tools.read_taxonomy.core.TaxonomyAgentToolkit") def test_event_properties_include_dynamic_hint(self, mock_toolkit_class): mock_toolkit = mock_toolkit_class.return_value diff --git a/ee/hogai/tools/read_taxonomy/tool.py b/ee/hogai/tools/read_taxonomy/tool.py index 89a11b01d082..f81534aef5a6 100644 --- a/ee/hogai/tools/read_taxonomy/tool.py +++ b/ee/hogai/tools/read_taxonomy/tool.py @@ -20,6 +20,7 @@ ReadEvents, ReadEventSamplePropertyValues, ReadTaxonomyToolArgs, + SearchTaxonomyProperties, execute_taxonomy_query, ) @@ -41,6 +42,7 @@ - kind: "entity_property_values" — sample values for an entity property. Required: `entity`, `property_name`. - kind: "action_properties" — properties for an action. Required: `action_id`. - kind: "action_property_values" — sample values for an action property. Required: `action_id`, `property_name`. +- kind: "search_properties" — search property names and descriptions for a keyword across persons, sessions, groups, and events in one call. Required: `term`. Use this when you know the concept (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. # Examples of when to use the read_taxonomy tool @@ -141,6 +143,7 @@ async def create_tool_class( ReadEntitySamplePropertyValuesWithGroups, # type: ignore[valid-type] ReadActionProperties, ReadActionSamplePropertyValues, + SearchTaxonomyProperties, ] class ReadTaxonomyToolArgsWithGroups(BaseModel):