diff --git a/src/sentry/ai_monitoring/conversation_query.py b/src/sentry/ai_monitoring/conversation_query.py new file mode 100644 index 000000000000..b5f14833bf8b --- /dev/null +++ b/src/sentry/ai_monitoring/conversation_query.py @@ -0,0 +1,150 @@ +from collections.abc import Iterator +from math import isfinite + +from parsimonious.exceptions import ParseError +from parsimonious.nodes import Node + +from sentry.ai_monitoring.constants import AI_CONVERSATIONS_FIELDS +from sentry.api.event_search import AggregateFilter, SearchFilter, event_search_grammar +from sentry.exceptions import InvalidSearchQuery +from sentry.search.eap.resolver import SearchResolver + + +def _nodes(node: Node, names: set[str], depth: int = 0) -> Iterator[Node]: + """Walk the search syntax tree and yield nodes with the requested grammar names. + + Names come from event_search_grammar, not span fields: + - filter: a complete field condition, e.g. toolCalls:>0. + - free_text: text without a field name, e.g. "hello world". + - boolean_operator: AND or OR between conditions. + - paren_group: conditions in parentheses, e.g. (toolCalls:>0 OR errors:0). + - search_key / text_key: the field name inside a filter, e.g. toolCalls. + - aggregate_key: a function expression, e.g. sum(span.duration). + - operator: a comparison inside a filter, e.g. >. + + Callers choose the level they need. For example, {"filter"} selects toolCalls:>0, + while {"search_key"} descends into that filter to select just toolCalls. + """ + # paren_group means a parenthesized search group: (toolCalls:>0) adds one level. + # Function parentheses, such as sum(span.duration), are not search groups. + if node.expr_name == "paren_group": + depth += 1 + if depth > 20: + raise InvalidSearchQuery("Conversation query has too many nested groups.") + if node.expr_name in names: + # Return the whole match and skip its contents. Selecting "filter" returns + # toolCalls:>0 once, without also walking its key and value. + yield node + else: + # Walk children in source order until we find a requested node. For example, + # (toolCalls:>0 errors:0) yields two filters, both at the same group depth. + for child in node.children: + yield from _nodes(child, names, depth) + + +def _is_summary_field(name: str) -> bool: + """Identify numeric summaries using the shared alias definitions. + + For example, totalCost and total_cost are summaries; conversationId and age are not. + """ + field = AI_CONVERSATIONS_FIELDS.get(name.strip('"')) + return field is not None and field[0] not in {"gen_ai.conversation.id", "max(timestamp)"} + + +def _compile_summary_filter(condition: Node, key: Node, resolver: SearchResolver) -> str: + """Expand a summary alias such as totalCost into an aggregate comparison. + + Let EAP interpret values, units, and nulls. For example, totalCost:0 may exclude + conversations with no recorded costs, even though the response displays zero. + """ + expression, _ = AI_CONVERSATIONS_FIELDS[key.text.strip('"')] + query = condition.text.replace(key.text, expression, 1) + terms = resolver.parse_search_query(query) + if len(terms) != 1 or not isinstance(terms[0], AggregateFilter): + raise InvalidSearchQuery(f"Invalid conversation aggregate filter: {condition.text}") + term = terms[0] + value = term.value.raw_value + if not isinstance(value, (int, float)) or not isfinite(value): + raise InvalidSearchQuery(f"Expected a finite numeric aggregate value: {condition.text}") + return term.to_query_string() + + +def _compile_condition(condition: Node, resolver: SearchResolver) -> str: + """Translate one user condition into a filter on a whole conversation. + + Expand summary aliases and reject user-written aggregates. For span fields and + free text, require at least one matching span, or no matching spans for exclusions. + For example, gen_ai.tool.name:search requires a matching span, while + !gen_ai.tool.name:search requires none. Keep source text to preserve those meanings. + """ + keys = list(_nodes(condition, {"aggregate_key", "search_key", "text_key"})) + if keys: + key = keys[0] + if key.expr_name == "aggregate_key": + raise InvalidSearchQuery( + "Explicit aggregates are not supported in conversation filters. " + "Use summary aliases or span fields instead." + ) + if _is_summary_field(key.text): + return _compile_summary_filter(condition, key, resolver) + if key.text == "has" and any(_is_summary_field(k.text) for k in keys[1:]): + raise InvalidSearchQuery("Use a numeric comparison for conversation summary fields.") + + # Preserve source text: !span.duration:>2s differs from span.duration:<=2s. + positive = condition.text + if "`" in positive: + raise InvalidSearchQuery("Literal backticks are not supported in conversation filters.") + excluded = False + if condition.expr_name == "filter": + excluded = positive.startswith("!") + positive = positive.removeprefix("!") + if any(operator.text == "!=" for operator in _nodes(condition, {"operator"})): + positive = positive.replace("!=", "", 1) + excluded = not excluded + if keys and keys[0].text.strip('"') == "conversationId": + positive = positive.replace(keys[0].text, "gen_ai.conversation.id", 1) + return f"count_if(`{positive}`,span.duration):{'=0' if excluded else '>0'}" + + +def compile_conversation_query(query: str, resolver: SearchResolver) -> str: + """Build the EAP query that selects matching conversation IDs before hydration. + + For general searches, restrict spans to those with a conversation ID and an AI + operation. Apply rewritten conditions to groups, preserving the user's Boolean + structure: toolCalls:>0 OR errors:>0 still means either condition can match. + Return exact ID lookups such as gen_ai.conversation.id:abc unchanged. + """ + # Keep source offsets: (toolCalls:>0 OR errors:0) yields two nodes, leaving OR untouched. + try: + conditions = list(_nodes(event_search_grammar.parse(query), {"filter", "free_text"})) + except (ParseError, RecursionError) as error: + # Parsing happens before EAP: turn malformed input such as "(errors:0" into a 400. + raise InvalidSearchQuery( + "Invalid conversation query. Check parentheses and quoting." + ) from error + # EAP's nested protobuf filters cannot handle long chains, e.g. 51 span predicates. + if len(conditions) > 50: + raise InvalidSearchQuery("Conversation queries may contain at most 50 conditions.") + parts: list[str] = [] + offset = 0 + for condition in conditions: + parts.extend((query[offset : condition.start], _compile_condition(condition, resolver))) + offset = condition.end + parts.append(query[offset:]) + group_query = "".join(parts).strip() + + # Conversation IDs imply AI operation types, so gen_ai.conversation.id:abc needs no scope. + if len(conditions) == 1: + terms = resolver.parse_search_query(query) + if ( + len(terms) == 1 + and isinstance(term := terms[0], SearchFilter) + and term.key.name == "gen_ai.conversation.id" + and term.operator == "=" + and term.value.raw_value + and not term.value.is_wildcard() + ): + return query + + conversation_scope = "has:gen_ai.conversation.id has:gen_ai.operation.type" + return f"{conversation_scope} AND ({group_query})" if group_query else conversation_scope diff --git a/src/sentry/ai_monitoring/endpoints/organization_ai_conversations.py b/src/sentry/ai_monitoring/endpoints/organization_ai_conversations.py index 3ac3f5ea4674..eddb9b5ba425 100644 --- a/src/sentry/ai_monitoring/endpoints/organization_ai_conversations.py +++ b/src/sentry/ai_monitoring/endpoints/organization_ai_conversations.py @@ -11,6 +11,7 @@ from sentry import features from sentry.ai_monitoring.constants import AI_CONVERSATIONS_FIELDS +from sentry.ai_monitoring.conversation_query import compile_conversation_query from sentry.ai_monitoring.conversation_titles import fetch_conversation_titles from sentry.ai_monitoring.serializers import OrganizationAIConversationsSerializer from sentry.ai_monitoring.utils import ( @@ -254,19 +255,29 @@ def get( return Response(as_validation_errors(serializer), status=400) validated_data = serializer.validated_data + user_query = validated_data.get("query", "") + query_string = _build_conversation_query( + "has:gen_ai.conversation.id has:gen_ai.operation.type", user_query + ) def data_fn(offset: int, limit: int) -> list[AIConversationResponse]: return self._get_conversations( snuba_params=snuba_params, offset=offset, limit=limit, - user_query=validated_data.get("query", ""), - sampling_mode=validated_data.get("samplingMode", "NORMAL"), + query_string=query_string, + sampling_mode=validated_data["samplingMode"], sorts=validated_data["sort"] if querying_enhancements_enabled else None, use_single_query=querying_enhancements_enabled, ) with handle_query_errors(): + if querying_enhancements_enabled: + resolver = Spans.get_resolver( + snuba_params, + SearchResolverConfig(auto_fields=True, disable_aggregate_extrapolation=True), + ) + query_string = compile_conversation_query(user_query, resolver) response = self.paginate( request=request, paginator=GenericOffsetPaginator(data_fn=data_fn), @@ -291,14 +302,11 @@ def _get_conversations( snuba_params: SnubaParams, offset: int, limit: int, - user_query: str, + query_string: str, sampling_mode: SAMPLING_MODES = "NORMAL", sorts: Sequence[str] | None = None, use_single_query: bool = False, ) -> list[AIConversationResponse]: - base_filter = "has:gen_ai.conversation.id has:gen_ai.operation.type" - query_string = _build_conversation_query(base_filter, user_query) - conversation_ids_results = self._fetch_conversation_ids( snuba_params, query_string, offset, limit, sampling_mode, sorts ) @@ -343,7 +351,6 @@ def _fetch_conversation_ids( if not any(column.removeprefix("-") == "gen_ai.conversation.id" for column in orderby): orderby.append("gen_ai.conversation.id") - # TODO (vgrozdanic): Sort on whole conversations instead of only matching spans. return Spans.run_table_query( params=snuba_params, query_string=query_string, diff --git a/src/sentry/api/event_search.py b/src/sentry/api/event_search.py index 1a1d7f403cf6..900a3c4b59f4 100644 --- a/src/sentry/api/event_search.py +++ b/src/sentry/api/event_search.py @@ -158,7 +158,8 @@ aggregate_key = key open_paren spaces function_args? spaces closed_paren function_args = aggregate_param (spaces comma spaces !comma aggregate_param?)* -aggregate_param = explicit_tag_key_aggregate_param / quoted_aggregate_param / raw_aggregate_param +aggregate_param = explicit_tag_key_aggregate_param / query_aggregate_param / quoted_aggregate_param / raw_aggregate_param +query_aggregate_param = "`" (quoted_value / ~r"[^`\"]+")* "`" raw_aggregate_param = ~r"[^()\t\n, \"]+" quoted_aggregate_param = '"' ('\\"' / ~r'[^\t\n\"]')* '"' explicit_tag_key_aggregate_param = explicit_tag_key / explicit_number_tag_key / explicit_string_tag_key / explicit_boolean_tag_key @@ -1642,6 +1643,9 @@ def visit_function_args( def visit_aggregate_param(self, node: Node, children: tuple[str]) -> str: return children[0] + def visit_query_aggregate_param(self, node: Node, children: object) -> str: + return node.text + def visit_raw_aggregate_param(self, node: Node, children: object) -> str: return node.text diff --git a/tests/sentry/api/endpoints/test_organization_ai_conversations.py b/tests/sentry/api/endpoints/test_organization_ai_conversations.py index ff29b81e6454..d165d7a611ac 100644 --- a/tests/sentry/api/endpoints/test_organization_ai_conversations.py +++ b/tests/sentry/api/endpoints/test_organization_ai_conversations.py @@ -7,6 +7,7 @@ import pytest from django.urls import reverse +from sentry.ai_monitoring.conversation_query import compile_conversation_query from sentry.ai_monitoring.endpoints.organization_ai_conversations import ( OrganizationAIConversationsEndpoint, ) @@ -17,7 +18,10 @@ from sentry.ai_monitoring.utils import ( get_last_output as _get_last_output, ) +from sentry.exceptions import InvalidSearchQuery +from sentry.search.eap.types import SearchResolverConfig from sentry.search.events.types import SnubaParams +from sentry.snuba.spans_rpc import Spans from sentry.testutils.helpers import parse_link_header from sentry.testutils.helpers.datetime import before_now @@ -240,6 +244,79 @@ def test_single_query_hydration_uses_one_aggregate_query(run_table_query: MagicM assert query["limit"] == 1 +@pytest.mark.parametrize( + "predicate", + [ + 'span.description:"hello, world"', + 'span.description:"sum(span.duration):>0"', + r'span.description:"say \"hello\""', + 'gen_ai.tool.name:["search docs",calculator]', + 'span.description:"literal\\*"', + "tags[custom,number]:>1.5", + "tags[custom,array][*]:value", + "timestamp:2023-06-01", + ], +) +def test_group_filter_preserves_source_predicate(predicate: str) -> None: + resolver = Spans.get_resolver(SnubaParams(), SearchResolverConfig()) + compiled = compile_conversation_query(predicate, resolver) + assert f"count_if(`{predicate}`,span.duration):>0" in compiled + where, having, _ = resolver.resolve_query(compiled) + assert having is not None + assert where is not None + assert "gen_ai.conversation.id" in str(where) + assert "gen_ai.operation.type" in str(where) + assert "custom" not in str(where) + assert "raw_description" not in str(where) + assert 'name: "sentry.project_id"' in str(having) + + +@pytest.mark.parametrize( + "query", + [ + "sum(gen_ai.cost.total_tokens):>10", + "count():>0", + "avg(span.duration):>2s", + "min(span.duration):>2s", + "max(timestamp):>2023-06-01", + "failure_count():0", + "!sum(span.duration):>2s", + "totalCost:>10 OR (sum(span.duration):>2s)", + "count_if(`gen_ai.tool.name:search`,span.duration):>0", + ], +) +def test_group_filter_rejects_explicit_aggregates(query: str) -> None: + resolver = Spans.get_resolver(SnubaParams(), SearchResolverConfig()) + with pytest.raises(InvalidSearchQuery, match="Explicit aggregates are not supported"): + compile_conversation_query(query, resolver) + + +@pytest.mark.parametrize("operator", ["=", "!=", ">", ">=", "<", "<="]) +def test_summary_filter_preserves_eap_null_semantics(operator: str) -> None: + resolver = Spans.get_resolver(SnubaParams(), SearchResolverConfig()) + compiled = compile_conversation_query(f"totalCost:{operator}0", resolver) + _, having, _ = resolver.resolve_query(compiled) + _, expected, _ = resolver.resolve_query( + f"sum_if(gen_ai.cost.total_tokens,gen_ai.operation.type,equals,ai_client):{operator}0" + ) + assert having == expected + + +def test_group_filter_accepts_long_text() -> None: + resolver = Spans.get_resolver(SnubaParams(), SearchResolverConfig()) + compiled = compile_conversation_query("x" * 4097, resolver) + _, having, _ = resolver.resolve_query(compiled) + assert having is not None + + +def test_group_filter_pushes_down_exact_id() -> None: + resolver = Spans.get_resolver(SnubaParams(), SearchResolverConfig()) + assert ( + compile_conversation_query('gen_ai.conversation.id:"session:123"', resolver) + == 'gen_ai.conversation.id:"session:123"' + ) + + class OrganizationAIConversationsEndpointTest(BaseAIConversationsTestCase): view = "sentry-api-0-organization-ai-conversations" @@ -318,6 +395,21 @@ def test_sorting_default_candidate_query(self, run_table_query: MagicMock) -> No assert query["selected_columns"] == ["gen_ai.conversation.id", "max(timestamp)"] assert query["orderby"] == ["-max(timestamp)", "gen_ai.conversation.id"] assert query["config"].disable_aggregate_extrapolation is True + assert query["sampling_mode"] == "HIGHEST_ACCURACY" + + @patch( + "sentry.ai_monitoring.endpoints.organization_ai_conversations.Spans.run_table_query", + return_value={"data": []}, + ) + def test_querying_preserves_requested_sampling_mode(self, run_table_query: MagicMock) -> None: + for sampling_mode in ["NORMAL", "HIGHEST_ACCURACY", "HIGHEST_ACCURACY_FLEX_TIME"]: + with self.feature("organizations:gen-ai-conversations-querying-enhancements"): + response = self.do_request( + {"project": [self.project.id], "samplingMode": sampling_mode} + ) + + assert response.status_code == 200, response.data + assert run_table_query.call_args.kwargs["sampling_mode"] == sampling_mode @patch( "sentry.ai_monitoring.endpoints.organization_ai_conversations.Spans.run_table_query", @@ -342,7 +434,8 @@ def test_sorting_cost_candidate_query(self, run_table_query: MagicMock) -> None: ] assert query["orderby"] == ["-total_cost", "gen_ai.conversation.id"] assert query["query_string"] == ( - "has:gen_ai.conversation.id has:gen_ai.operation.type gen_ai.tool.name:search" + "has:gen_ai.conversation.id has:gen_ai.operation.type " + "AND (count_if(`gen_ai.tool.name:search`,span.duration):>0)" ) @patch( @@ -468,7 +561,7 @@ def test_sorting_cost_pagination(self) -> None: assert response.status_code == 200, response.data assert [row["conversationId"] for row in response.data] == ["conversation-b"] - def test_sorting_uses_matching_spans_but_hydrates_whole_conversation(self) -> None: + def test_sorting_uses_whole_conversations_after_filtering(self) -> None: now = before_now(days=10).replace(microsecond=0) for conversation_id, tool_name, cost in [ ("conversation-a", "search", 1), @@ -494,12 +587,167 @@ def test_sorting_uses_matching_spans_but_hydrates_whole_conversation(self) -> No ) assert response.status_code == 200, response.data - # Temporary behavior: ascending matching-span cost is 1, 5, not hydrated cost 11, 5. assert [row["conversationId"] for row in response.data] == [ - "conversation-a", "conversation-b", + "conversation-a", ] - assert [row["totalCost"] for row in response.data] == [11, 5] + assert [row["totalCost"] for row in response.data] == [5, 11] + + def test_conversation_group_filters(self) -> None: + cases = [ + ("gen_ai.agent.name:researcher gen_ai.tool.name:search", ["a"]), + ("gen_ai.tool.name:search span.status:internal_error", ["a"]), + ("gen_ai.tool.name:search gen_ai.tool.name:calculator", ["a"]), + ("gen_ai.tool.name:[search,calculator]", ["a", "b"]), + ("!gen_ai.tool.name:search", ["c"]), + ("gen_ai.tool.name:!=search", ["c"]), + ("!gen_ai.tool.name:[search,calculator]", ["c"]), + ("has:gen_ai.tool.name", ["a", "b"]), + ("!has:gen_ai.tool.name", ["c"]), + ("totalCost:>10", ["a"]), + ('"totalCost":>10', ["a"]), + ("total_cost:>10 gen_ai.tool.name:search", ["a"]), + ( + "sum_if_gen_ai_cost_total_tokens_gen_ai_operation_type_equals_ai_client:>10", + ["a"], + ), + ("totalCost:0", []), + ("totalCost:<=0", []), + ("totalCost:>=0", ["a", "b"]), + ("!totalCost:>10", ["b"]), + ("toolCalls:>1", ["a"]), + ("errors:0", ["b", "c"]), + ("duration:>5s", ["a"]), + ("duration:>=0", ["a", "b", "c"]), + ("generationDuration:0", []), + ("span.duration:>2s", ["a"]), + ("span.duration:<=2s", ["a", "b", "c"]), + ("!span.duration:>2s", ["b", "c"]), + ("gen_ai.usage.total_tokens:>15", ["a"]), + ("gen_ai.conversation.id:b OR totalCost:>10", ["a", "b"]), + ("!gen_ai.tool.name:search OR totalCost:>10", ["a", "c"]), + ( + "gen_ai.agent.name:researcher AND (gen_ai.tool.name:search OR toolCalls:0)", + ["a", "c"], + ), + ('span.description:"hello world"', ["a"]), + ('"hello world"', ["a"]), + ("span.description:hello*", ["a"]), + ("span.description:non-ai-only", []), + ("conversationId:a", ["a"]), + ] + now = before_now(days=1).replace(microsecond=0) + self.store_ai_span( + conversation_id="a", + timestamp=now, + operation_type="ai_client", + agent_name="researcher", + cost=6, + tokens=10, + description="hello world", + ) + self.store_ai_span( + conversation_id="a", + timestamp=now, + operation_type="ai_client", + cost=6, + tokens=20, + ) + self.store_spans( + [ + self.create_span( + { + "ai_conversation_id": "a", + "data": {"gen_ai.operation.type": "tool", "gen_ai.tool.name": "search"}, + }, + start_ts=now, + duration=3000, + ) + ] + ) + self.store_ai_span( + conversation_id="a", + timestamp=now, + operation_type="tool", + tool_name="calculator", + status="internal_error", + ) + self.store_ai_span( + conversation_id="b", + timestamp=now, + operation_type="ai_client", + cost=2, + ) + self.store_ai_span( + conversation_id="b", + timestamp=now, + operation_type="tool", + tool_name="search", + ) + self.store_ai_span( + conversation_id="c", + timestamp=now, + operation_type="agent", + agent_name="researcher", + ) + # Ignore spans without an AI operation even when they belong to an AI conversation. + self.store_ai_span( + conversation_id="c", + timestamp=now, + description="non-ai-only", + tool_name="search", + status="internal_error", + ) + # Neither an out-of-window span nor a group without AI operations can match. + self.store_ai_span( + conversation_id="a", + timestamp=now - timedelta(days=2), + operation_type="ai_client", + cost=100, + ) + self.store_ai_span(conversation_id="not-ai", timestamp=now, tool_name="search", cost=100) + for search, expected_ids in cases: + with self.feature("organizations:gen-ai-conversations-querying-enhancements"): + response = self.do_request( + { + "project": [self.project.id], + "start": (now - timedelta(hours=1)).isoformat(), + "end": (now + timedelta(hours=1)).isoformat(), + "query": search, + "sort": "conversationId", + } + ) + assert response.status_code == 200, (search, response.data) + assert [row["conversationId"] for row in response.data] == expected_ids, search + assert [row["totalCost"] for row in response.data] == [ + {"a": 12, "b": 2, "c": 0}[conv_id] for conv_id in expected_ids + ], search + + def test_invalid_group_filter(self) -> None: + for search in [ + "toolCalls:banana", + "totalCost:NaN", + "duration:>oops", + "totalCost:>10 OR", + "AND toolCalls:1", + "toolCalls:1 OR OR errors:0", + "(gen_ai.tool.name:search", + "collect_unique(trace):>0", + "sum(gen_ai.cost.total_tokens):>10", + "count_if(`gen_ai.tool.name:search`,span.duration):>0", + 'span.description:"literal ` backtick"', + "has:totalCost", + 'has:"total_cost"', + "!has:totalCost", + "count_if(`gen_ai.tool.name:search sum(span.duration):>0`,span.duration):>0", + "count_if(`(broken`,span.duration):>0", + "gen_ai.tool.name:search " * 51, + "(" * 21 + "gen_ai.tool.name:search" + ")" * 21, + "count_if(`" + "(" * 21 + "gen_ai.tool.name:search" + ")" * 21 + "`,span.duration):>0", + ]: + with self.feature("organizations:gen-ai-conversations-querying-enhancements"): + response = self.do_request({"project": [self.project.id], "query": search}) + assert response.status_code == 400, (search, response.data) def test_single_conversation_single_trace(self) -> None: """Test a conversation with all spans in a single trace""" diff --git a/tests/sentry/api/test_event_search.py b/tests/sentry/api/test_event_search.py index afe449755524..bff671b3c071 100644 --- a/tests/sentry/api/test_event_search.py +++ b/tests/sentry/api/test_event_search.py @@ -349,6 +349,21 @@ def test_paren_expression(self) -> None: SearchFilter(key=SearchKey(name="z"), operator="=", value=SearchValue(raw_value="1")), ] + def test_conditional_aggregate_query_argument(self) -> None: + for predicate in [ + "x:1 AND (y:2 OR z:3)", + 'span.description:"hello world"', + 'gen_ai.tool.name:["search docs",calculator]', + r'span.description:"say \"hello\""', + ]: + assert parse_search_query(f"count_if(`{predicate}`,span.duration):>0") == [ + AggregateFilter( + key=AggregateKey(f"count_if(`{predicate}`, span.duration)"), + operator=">", + value=SearchValue(0.0), + ) + ] + def test_paren_expression_of_empty_string(self) -> None: assert parse_search_query('("")') == parse_search_query('""') == []