Trying to encourage LLM to do small exploratory queries with results returned directly before doing the 'final' query with results not returned directly (because they overflow the context limit)

Author: kaleko

lib/idp_common_pkg/idp_common/agents/analytics/agent.py

@@ -68,7 +68,7 @@ def create_analytics_agent(
     When generating SQL:
     - ALWAYS put ALL column names in double quotes when including ANYHWERE inside of a query.
     - Use standard SQL syntax compatible with Amazon Athena, for example use standard date arithmetic that's compatible with Athena.
-    - Do not guess at table or column names. Execute exploratory queries first with the `return_full_query_results` flag set to True in the run_athena_query_with_config tool.
+    - Do not guess at table or column names. Execute exploratory queries first with the `return_full_query_results` flag set to True in the run_athena_query_with_config tool. Your final query should use `return_full_query_results` set to False. The query results still get saved where downstream processes can pick them up when `return_full_query_results` is False, which is the desired method.
     - Use a "SHOW TABLES" query to list all dynamic tables available to you.
     - Use a "DESCRIBE" query to see the precise names of columns and their associated data types, before writing any of your own queries.
     - Include appropriate table joins when needed
@@ -113,6 +113,7 @@ def run_athena_query_with_config(
             return_full_query_results: If True, includes the full query results as CSV string in the response.
                 WARNING: This can return very large strings and should only be used for small exploratory
                 queries like DESCRIBE, SHOW TABLES, or queries with LIMIT clauses. Default is False.
+                Use False whenever possible.
 
         Returns:
             Dict containing either query results or error information

lib/idp_common_pkg/idp_common/agents/analytics/tools/athena_tool.py

@@ -14,6 +14,9 @@
 
 logger = logging.getLogger(__name__)
 
+# Maximum number of rows that can be returned directly when return_full_query_results=True
+MAX_ROWS_TO_RETURN_DIRECTLY = 100
+
 
 @tool
 def run_athena_query(
@@ -24,7 +27,7 @@ def run_athena_query(
 
     Uses boto3 to execute the query on Athena. Query results are stored in s3.
     Successful execution will return a dict with result_column_metadata,
-        result_csv_s3_uri, rows_returned, and original_query.
+        result_csv_s3_uri, number of rows_returned, and original_query.
 
     Args:
         query: SQL query string to execute
@@ -101,6 +104,24 @@ def run_athena_query(
             # For queries with headers (like SELECT), Athena typically includes headers in the first row
             total_rows = len(results["ResultSet"]["Rows"])
 
+            # Check if return_full_query_results is True and we have too many rows
+            if return_full_query_results and total_rows > MAX_ROWS_TO_RETURN_DIRECTLY:
+                logger.warning(
+                    f"Query returned {total_rows} rows, which exceeds the limit of {MAX_ROWS_TO_RETURN_DIRECTLY} "
+                    f"for return_full_query_results=True"
+                )
+                return {
+                    "success": False,
+                    "error": (
+                        f"More than {MAX_ROWS_TO_RETURN_DIRECTLY} rows were retrieved when the tool was called with "
+                        "`return_full_query_results` set to True. This flag should only be used for small queries "
+                        "returning a few rows. Please try again with `return_full_query_results` set to False, "
+                        "in which case the query results will be saved rather than returned directly."
+                    ),
+                    "query": query,
+                    "rows_returned": total_rows,
+                }
+
             result_dict = {
                 "success": True,
                 "result_column_metadata": column_metadata,

lib/idp_common_pkg/idp_common/agents/common/monitoring.py

@@ -180,7 +180,7 @@ def on_before_tool_invocation(self, event) -> None:
         tool_name = event.tool_use.get("name", "unknown")
         tool_input = event.tool_use.get("input", {})
 
-        self.monitor_logger.info(f"🔧 Invoking tool: {tool_name}")
+        self.monitor_logger.info(f"🔧 Invoking tool: {tool_name} --> {event}")
         if self.enable_detailed_logging:
             self.monitor_logger.debug(f"Tool input: {json.dumps(tool_input, indent=2)}")