kaxil commented on code in PR #71317: URL: https://github.com/apache/airflow/pull/71317#discussion_r3738772050
########## 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: Good catch. Reworded in 4362b0b to say `truncated` means you are not seeing the whole result, either because more rows matched or because it was too large, with `truncated_by` naming which. ########## 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: Fixed in 4362b0b. Stale wording: the loop keeps a contiguous prefix and stops at the first row that does not fit, so one wide row early in the result ends it rather than being skipped. Same text was wrong in `datafusion.py` and in the .rst. ########## 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: Fixed in 4362b0b along with the same wording in `sql.py` and the .rst. ########## 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: Fixed in 4362b0b. Now reads "not read out of a DBAPI cursor; what the driver has already transferred is its own call", which matches the qualification further down the page. The line 240 instance was the "dropped from the end" wording and is corrected too. -- 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]
