Copilot commented on code in PR #71317: URL: https://github.com/apache/airflow/pull/71317#discussion_r3738682738
########## providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py: ########## @@ -0,0 +1,155 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Bounded, columnar payloads for the ``query`` tool of the SQL toolsets. + +A tool result stays in the model's message history for the rest of the run, so its +cost is re-paid on every subsequent request. Two things here keep that bounded: + +* **Columnar shape.** ``{"columns": [...], "rows": [[...], ...]}`` names each column + once instead of repeating it in a dict per row. On a table with thousands of + columns the repeated names, not the values, are the bulk of the payload. +* **A byte budget.** ``max_rows`` caps rows, which says nothing about size -- a + single row of a 3000-column table dwarfs a thousand rows of a narrow one. The + budget here is what actually bounds context, and when it bites the payload says + so, in terms the agent can act on (narrow the projection). +""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from typing import Any + +# A policy default, not a limit imposed by any storage, protocol, or model layer: +# roughly 16k tokens at 4 characters per token. Large enough that ordinary queries are +# unaffected, small enough that no single tool result can dominate the context window. +# Deployments that keep many results in history should lower it. +DEFAULT_MAX_RESULT_BYTES = 65_536 + +# Tool results are machine-read, so no whitespace. ensure_ascii=False matters as much as +# the separators: escaping one CJK character to \uXXXX costs six bytes instead of three, +# so an ASCII-escaped result is charged several times over against the budget and +# truncated that much earlier than an equivalent English one. +_DUMP_KWARGS: dict[str, Any] = {"default": str, "separators": (",", ":"), "ensure_ascii": False} + +#: Description for the ``query`` tool. States the columnar shape, since the model has +#: to align each row's values to ``columns`` positionally, and the truncation contract, +#: so a short result is not read as an empty table. +QUERY_TOOL_DESCRIPTION = ( + "Execute a SQL query. Returns JSON of the form " + '{"columns": [name, ...], "rows": [[value, ...], ...]}, where each row holds its ' + "values in column order. A `truncated` key means the query matched more than was " + "returned and `truncated_by` names the limit that was hit; narrow the projection or " + "aggregate in SQL rather than paging through the result." +) Review Comment: The `QUERY_TOOL_DESCRIPTION` says that a `truncated` key means “the query matched more than was returned”, but `build_query_result()` also sets `truncated` when the header (column names) alone exceeds `max_result_bytes`. In that case, truncation is due to the size budget, not necessarily because more rows matched. Updating the description avoids misleading agents/system prompts about what `truncated` implies. ########## providers/common/ai/docs/toolsets.rst: ########## @@ -203,6 +204,59 @@ Parameters Default ``False`` -- only SELECT-family and read-only metadata (``DESCRIBE``/``SHOW``) statements are permitted. - ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``. + Rows beyond it are never fetched from the cursor. +- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64 KiB. Review Comment: In the parameters list, this says rows beyond `max_rows` are “never fetched from the cursor”, but later in the same doc you note that some hooks fall back to a full fetch (non-DBAPI cursors) and that some drivers buffer client-side. Tweaking this line to qualify the claim (e.g. “not read from a DBAPI cursor”) will avoid a direct contradiction within the docs. This issue also appears on line 240 of the same file. ########## providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py: ########## @@ -143,7 +218,19 @@ class SQLToolset(AbstractToolset[Any]): :param allow_writes: Allow data-modifying SQL (INSERT, UPDATE, DELETE, etc.). Default ``False`` — only SELECT-family statements are permitted. :param max_rows: Maximum number of rows returned from the ``query`` tool. - Default ``50``. + Default ``50``. Rows beyond it are not pulled out of the cursor. How much that + saves is the driver's call, not this toolset's: a client-buffering driver + (psycopg2's default cursor, MySQLdb) has already received the whole result by + the time the first row is read, so only the per-row Python conversion is + skipped. Treat this as a bound on what the agent is shown, not as a guarantee + that ``SELECT * FROM huge_table`` is cheap. + :param max_result_bytes: Budget for the serialized ``query`` result, in bytes. + Default 64 KiB. ``max_rows`` bounds rows, which says nothing about size: one + row of a 3000-column table is larger than a thousand rows of a narrow one, and + a tool result stays in the model's message history for the rest of the run, so + its cost is re-paid on every subsequent request. Rows are dropped from the end + until the payload fits, and the result reports which limit it hit so the agent + can narrow its projection rather than page through the table. Review Comment: The `max_result_bytes` docstring says “Rows are dropped from the end until the payload fits”, but `build_query_result()` keeps a contiguous prefix and *stops at the first row that doesn’t fit* (it does not “pack” later rows, and it may stop early if an oversized row appears mid-result). Aligning the docstring with the actual truncation semantics will prevent confusion when debugging partial results. ########## providers/common/ai/src/airflow/providers/common/ai/toolsets/datafusion.py: ########## @@ -97,6 +102,13 @@ class DataFusionToolset(AbstractToolset[Any]): are permitted. :param max_rows: Maximum number of rows returned from the ``query`` tool. Default ``50``. + :param max_result_bytes: Budget for the serialized ``query`` result, in bytes. + Default 64 KiB. ``max_rows`` bounds rows, which says nothing about size: one + row of a 3000-column table is larger than a thousand rows of a narrow one, and + a tool result stays in the model's message history for the rest of the run, so + its cost is re-paid on every subsequent request. Rows are dropped from the end + until the payload fits, and the result reports which limit it hit so the agent + can narrow its projection rather than page through the table. Review Comment: The `max_result_bytes` docstring says “Rows are dropped from the end until the payload fits”, but `build_query_result()` returns a contiguous prefix and stops at the first row that doesn’t fit the remaining budget (it does not skip a wide row and continue with later rows). Updating the wording here to match the implementation will make the truncation behavior clearer to users. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
