Merge branch 'feature/optimize-agent-analytics-metadata' into 'develop'

Feature/optimize agent analytics metadata

See merge request genaiic-reusable-assets/engagement-artifacts/genaiic-idp-accelerator!322

Author: rstrahan

CHANGELOG.md

@@ -6,6 +6,15 @@ SPDX-License-Identifier: MIT-0
 ## [Unreleased]
 
 ### Added
+- **Analytics Agent 2-Phase Schema Optimization for Improved Performance**
+  - Implemented progressive schema disclosure system with efficient 2-phase approach
+  - Phase 1: `get_database_overview()` provides fast table listing and guidance (~500 tokens vs 3000+ tokens - 6x faster)
+  - Phase 2: `get_table_info(['specific_tables'])` loads detailed schemas only for tables actually needed by the query
+  - Enhanced SQL guidance with comprehensive Athena/Trino function reference and PostgreSQL operator warnings to prevent common query failures
+
+### Fixed
+- Fix missing data in Glue tables when using a document class that contains a dash (-). 
+
 
 ## [0.3.16]
 

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

@@ -15,7 +15,12 @@
 from ..common.config import load_result_format_description
 from ..common.strands_bedrock_model import create_strands_bedrock_model
 from .config import load_python_plot_generation_examples
-from .tools import CodeInterpreterTools, get_database_info, run_athena_query
+from .tools import (
+    CodeInterpreterTools,
+    get_database_overview,
+    get_table_info,
+    run_athena_query,
+)
 from .utils import register_code_interpreter_tools
 
 logger = logging.getLogger(__name__)
@@ -50,29 +55,179 @@ def create_analytics_agent(
     # Task
     Your task is to:
     1. Understand the user's question
-    2. Use get_database_info tool to understand initial information about the database schema
-    3. Generate a valid Athena query that answers the question OR that will provide you information to write a second Athena query which answers the question (e.g. listing tables first, if not enough information was provided by the get_database_info tool)
-    4. Before executing the Athena query, re-read it and make sure _all_ column names mentioned _anywhere inside of the query_ are enclosed in double quotes.
-    5. Execute your revised query using the run_athena_query tool. If you receive an error message, correct your Athena query and try again a maximum of 5 times, then STOP. Do not ever make up fake data. For exploratory queries you can return the athena results directly. For larger or final queries, the results should need to be returned because downstream tools will download them separately.
-    6. Use the write_query_results_to_code_sandbox to convert the athena response into a file called "query_results.csv" in the same environment future python scripts will be executed.
-    7. If the query is best answered with a plot or a table, write python code to analyze the query results to create a plot or table. If the final response to the user's question is answerable with a human readable string, return it as described in the result format description section below.
-    8. To execute your plot generation code, use the execute_python tool and directly return its output without doing any more analysis.
+    2. **EFFICIENT APPROACH**: Use get_database_overview() to get a fast overview of available tables and their purposes
+    3. Apply the Question-to-Table mapping rules below to select the correct tables for your query
+    4. Use get_table_info(['table1', 'table2']) to get detailed schemas ONLY for the tables you need
+    5. Generate a valid Athena query based on the targeted schema information
+    6. **VALIDATE YOUR SQL**: Before executing, check for these common mistakes:
+       - All column names enclosed in double quotes: `"column_name"`
+       - No PostgreSQL operators: Replace `~` with `REGEXP_LIKE()`
+       - No invalid functions: Replace `CONTAINS()` with `LIKE`, `ILIKE` with `LOWER() + LIKE`
+       - Only valid Trino functions used
+       - Proper date formatting and casting
+    7. Execute your validated query using the run_athena_query tool. If you receive an error message, correct your Athena query and try again a maximum of 5 times, then STOP. Do not ever make up fake data. For exploratory queries you can return the athena results directly. For larger or final queries, the results should need to be returned because downstream tools will download them separately.
+    8. Use the write_query_results_to_code_sandbox to convert the athena response into a file called "query_results.csv" in the same environment future python scripts will be executed.
+    9. If the query is best answered with a plot or a table, write python code to analyze the query results to create a plot or table. If the final response to the user's question is answerable with a human readable string, return it as described in the result format description section below.
+    10. To execute your plot generation code, use the execute_python tool and directly return its output without doing any more analysis.
+
+    # CRITICAL: Two-Step Database Information Approach
+    **For optimal performance and accuracy:**
+    
+    ## Step 1: Overview (Fast)
+    - Always start with `get_database_overview()` to see available tables
+    - This gives you table names, purposes, and question-to-table mapping guidance
+    - **~500 tokens vs 3000+ tokens** - much faster for simple questions
+    
+    ## Step 2: Detailed Schemas (On-Demand) 
+    - Use `get_table_info(['table1', 'table2'])` for specific tables you need
+    - Only request detailed info for tables relevant to your query
+    - Get complete column listings, sample queries, and aggregation rules
+    
+    # CRITICAL: Question-to-Table Mapping Rules
+    **ALWAYS follow these rules to select the correct table:**
+    
+    ## For Classification/Document Type Questions:
+    - "How many X documents?" → Use `document_sections_x` table
+    - "Documents classified as Y" → Use `document_sections_y` table  
+    - "What document types processed?" → Query document_sections_* tables
+    - **NEVER use metering table for classification info - it only has usage/cost data**
+    
+    Examples:
+    ```sql
+    -- ✅ CORRECT: Count W2 documents
+    SELECT COUNT(DISTINCT "document_id") FROM document_sections_w2 WHERE "date" = CAST(CURRENT_DATE AS VARCHAR)
+    
+    -- ❌ WRONG: Don't use metering for classification
+    SELECT COUNT(*) FROM metering WHERE "service_api" LIKE '%w2%'
+    ```
+    
+    ## For Volume/Cost/Consumption Questions:
+    - "How much did processing cost?" → Use `metering` table
+    - "Token usage by model" → Use `metering` table  
+    - "Pages processed" → Use `metering` table (with proper MAX aggregation)
+    
+    ## For Accuracy Questions:
+    - "Document accuracy" → Use `evaluation` tables (may be empty)
+    - "Precision/recall metrics" → Use `evaluation` tables
+    
+    ## For Content/Extraction Questions:
+    - "What was extracted from documents?" → Use appropriate `document_sections_*` table
+    - "Show invoice amounts" → Use `document_sections_invoice` table
     
     DO NOT attempt to execute multiple tools in parallel. The input of some tools depend on the output of others. Only ever execute one tool at a time.
     
-    When generating Athena:
-    - ALWAYS put ALL column names in double quotes when including ANYHWERE inside of a query.
-    - Use standard Athena 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. 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
-    - Use column names exactly as they appear in the schema, ALWAYS in double quotes within your query.
-    - When querying strings, be aware that tables may contain ALL CAPS strings (or they may not). So, make your queries agnostic to case whenever possible.
-    - If you cannot get your query to work successfully, stop. DO NOT EVER generate fake or synthetic data. Instead, return a text response indicating that you were unable to answer the question based on the data available to you.
-    - The Athena query does not have to answer the question directly, it just needs to return the data required to answer the question. Python code will read the results and further analyze the data as necessary. If the Athena query is too complicated, you can simplify it to rely on post processing logic later.
-    - If your query returns 0 rows, it may be that the query needs to be changed and tried again. If you try a few variations and keep getting 0 rows, then perhaps that tells you the answer to the user's question and you can stop trying.
-    - If you get an error related to the column not existing or not having permissions to access the column, this is likely fixed by putting the column name in double quotes within your Athena query.
+    # CRITICAL: Athena SQL Function Reference (Trino-based)
+    **Athena engine version 3 uses Trino functions. DO NOT use PostgreSQL-style operators or invalid functions.**
+    
+    ## CRITICAL: Regular Expression Operators
+    **Athena does NOT support PostgreSQL-style regex operators:**
+    - ❌ NEVER use `~`, `~*`, `!~`, or `!~*` operators (these will cause query failures)
+    - ✅ ALWAYS use `REGEXP_LIKE(column, 'pattern')` for regex matching
+    - ✅ Use `NOT REGEXP_LIKE(column, 'pattern')` for negative matching
+
+    ### Common Regex Examples:
+    ```sql
+    -- ❌ WRONG: PostgreSQL-style (will fail with operator error)
+    WHERE "inference_result.wages" ~ '^[0-9.]+$'
+    WHERE "service_api" ~* 'classification'
+    WHERE "document_type" !~ 'invalid'
+    
+    -- ✅ CORRECT: Athena/Trino style
+    WHERE REGEXP_LIKE("inference_result.wages", '^[0-9.]+$')
+    WHERE REGEXP_LIKE(LOWER("service_api"), 'classification') 
+    WHERE NOT REGEXP_LIKE("document_type", 'invalid')
+    ```
+    
+    ## Valid String Functions (Trino-based):
+    - `LIKE '%pattern%'` - Pattern matching (NOT CONTAINS function)
+    - `REGEXP_LIKE(string, pattern)` - Regular expression matching (NOT ~ operator)
+    - `LOWER()`, `UPPER()` - Case conversion
+    - `POSITION(substring IN string)` - Find substring position (NOT STRPOS)
+    - `SUBSTRING(string, start, length)` - String extraction
+    - `CONCAT(string1, string2)` - String concatenation
+    - `LENGTH(string)` - String length
+    - `TRIM(string)` - Remove whitespace
+    
+    ## ❌ COMMON MISTAKES - Functions/Operators that DON'T exist in Athena:
+    - `CONTAINS(string, substring)` → Use `string LIKE '%substring%'`
+    - `ILIKE` operator → Use `LOWER(column) LIKE LOWER('pattern')`
+    - `STRPOS(string, substring)` → Use `POSITION(substring IN string)`
+    - `~` regex operator → Use `REGEXP_LIKE(column, 'pattern')`
+    
+    ## Valid Date/Time Functions:
+    - `CURRENT_DATE` - Current date
+    - `DATE_ADD(unit, value, date)` - Date arithmetic (e.g., `DATE_ADD('day', 1, CURRENT_DATE)`)
+    - `CAST(expression AS type)` - Type conversion
+    - `FORMAT_DATETIME(timestamp, format)` - Date formatting
+    
+    ## Critical Query Patterns:
+    ```sql
+    -- ✅ CORRECT: String matching
+    WHERE LOWER("service_api") LIKE '%classification%'
+    
+    -- ❌ WRONG: Invalid function
+    WHERE CONTAINS("service_api", 'classification')
+    
+    -- ✅ CORRECT: Numeric validation with regex
+    WHERE REGEXP_LIKE("inference_result.amount", '^[0-9]+\.?[0-9]*$')
+    
+    -- ❌ WRONG: PostgreSQL regex operator
+    WHERE "inference_result.amount" ~ '^[0-9.]+$'
+    
+    -- ✅ CORRECT: Case-insensitive pattern matching
+    WHERE LOWER("document_type") LIKE LOWER('%invoice%')
+    
+    -- ❌ WRONG: ILIKE operator
+    WHERE "document_type" ILIKE '%invoice%'
+    
+    -- ✅ CORRECT: Today's data
+    WHERE "date" = CAST(CURRENT_DATE AS VARCHAR)
+    
+    -- ✅ CORRECT: Date range  
+    WHERE "date" >= '2024-01-01' AND "date" <= '2024-12-31'
+    ```
+   
+    **TRUST THIS INFORMATION - Do not run discovery queries like SHOW TABLES or DESCRIBE unless genuinely needed.**
+
+    When generating Athena queries:
+    - **ALWAYS put ALL column names in double quotes** - this includes dot-notation columns like `"document_class.type"`
+    - **Use only valid Trino functions** listed above - Athena engine v3 is Trino-based
+    - **Leverage comprehensive schema first** - it contains complete table/column information
+    - **Follow aggregation patterns**: MAX for page counts per document (not SUM), SUM for costs
+    - **Use case-insensitive matching**: `WHERE LOWER("column") LIKE LOWER('%pattern%')`
+    - **Handle dot-notation carefully**: `"document_class.type"` is a SINGLE column name with dots
+    - **Prefer simple queries**: Complex logic can be handled in Python post-processing
+    
+    ## Error Recovery Patterns:
+    - **`~ operator not found`** → Replace with `REGEXP_LIKE(column, 'pattern')`
+    - **`ILIKE operator not found`** → Use `LOWER(column) LIKE LOWER('pattern')`
+    - **`Function CONTAINS not found`** → Use `column LIKE '%substring%'`
+    - **`Function STRPOS not found`** → Use `POSITION(substring IN column)`
+    - **Column not found** → Check double quotes: `"column_name"`
+    - **Function not found** → Use valid Trino functions only
+    - **0 rows returned** → Check table names, date filters, and case sensitivity  
+    - **Case sensitivity** → Use `LOWER()` for string comparisons
+    
+    ## Standard Query Templates:
+    ```sql
+    -- Document classification count
+    SELECT COUNT(DISTINCT "document_id") 
+    FROM document_sections_{type} 
+    WHERE "date" = CAST(CURRENT_DATE AS VARCHAR)
+    
+    -- Cost analysis
+    SELECT "context", SUM("estimated_cost") as total_cost
+    FROM metering 
+    WHERE "date" >= '2024-01-01'
+    GROUP BY "context"
+    
+    -- Joined analysis
+    SELECT ds."document_class.type", AVG(CAST(m."estimated_cost" AS DOUBLE)) as avg_cost
+    FROM document_sections_w2 ds
+    JOIN metering m ON ds."document_id" = m."document_id"
+    WHERE ds."date" = CAST(CURRENT_DATE AS VARCHAR)
+    GROUP BY ds."document_class.type"
+    ```
     
     When writing python:
     - Only write python code to generate plots or tables. Do not use python for any other purpose.
@@ -132,7 +287,8 @@ def run_athena_query_with_config(
         run_athena_query_with_config,
         code_interpreter_tools.write_query_results_to_code_sandbox,
         code_interpreter_tools.execute_python,
-        get_database_info,
+        get_database_overview,  # Fast, lightweight table overview
+        get_table_info,  # Detailed schema for specific tables
     ]
 
     # Get model ID from environment variable