-- AI-AWARE QUERY PATTERNS
--
-- Idiomatic patterns for writing efficient and correct AI UDF queries against
-- AOCore's Trino plugin. Copy-paste any pattern as a starting point.
--
-- Requires: AOCore Trino plugin (Obj E); model catalog (Obj C) seeded with
--           the referenced models; Lance vector index (Obj D) for SEMANTIC_SEARCH.
--
-- Docs: docs/developer/ai-sql-functions.md
--       docs/developer/ai-sql-functions.md#performance-characteristics


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 1: CROSS JOIN UNNEST for ARRAY<ROW> results
-- ────────────────────────────────────────────────────────────────────────────
-- SEMANTIC_SEARCH and RERANK both return ARRAY<ROW(doc_id VARCHAR, score DOUBLE)>.
-- Expand the array to rows with CROSS JOIN UNNEST, then JOIN to the document table.

-- SEMANTIC_SEARCH example:
SELECT d.id, d.content, s.score
FROM documents d
CROSS JOIN UNNEST(SEMANTIC_SEARCH('kb_vectors', 'query text here', 10)) AS s(doc_id, score)
WHERE d.id = s.doc_id
ORDER BY s.score DESC;

-- RERANK example:
SELECT r.doc_id, r.score
FROM UNNEST(
    RERANK('cohere-rerank-3', 'query text here',
           ARRAY['candidate one', 'candidate two', 'candidate three'])
) AS r(doc_id, score)
ORDER BY r.score DESC;


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 2: NULL filtering for per-item failures
-- ────────────────────────────────────────────────────────────────────────────
-- Per-item errors (budget_exceeded, guardrail_blocked, eval_gate_blocked,
-- model_not_found, rate_limited, ...) return SQL NULL. Filter with IS NOT NULL.
--
-- ANTI-PATTERN: calling the UDF in both SELECT and WHERE evaluates it TWICE:
--   SELECT id FROM reviews
--   WHERE LLM_CLASSIFY('m', content, ARRAY['pos','neg']) = 'pos';   -- UDF called twice!
--
-- RECOMMENDED: use a CTE (materialized in one pass; UDF called once per row):

SELECT id, content, classification
FROM (
    SELECT id, content,
           LLM_CLASSIFY('claude-haiku-classify', content, ARRAY['positive', 'negative', 'neutral']) AS classification
    FROM reviews
    WHERE created_at > now() - INTERVAL '1' HOUR
)
WHERE classification IS NOT NULL;

-- Same pattern works for any UDF result:
--   WHERE summary IS NOT NULL        (LLM_GENERATE)
--   WHERE embedding IS NOT NULL      (EMBED)
--   WHERE extracted_json IS NOT NULL (LLM_EXTRACT)
--   WHERE rerank_results IS NOT NULL (RERANK — whole ARRAY is NULL on failure)


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 3: Materialized-table best practice (BigQuery / E-RESEARCH Domain 4)
-- ────────────────────────────────────────────────────────────────────────────
-- For workloads with >10K rows, separate the SCAN cost from the AI inference cost.
-- Write the filtered input subset to a temp table first; then run AI UDFs over
-- the materialized rows. This enables:
--   - Predictable batch sizes (known row count before the AI call begins)
--   - Separate retry budget (re-run only the AI step if it fails)
--   - Parallelism between SCAN and inference stages
--
-- ANTI-PATTERN (slow at scale — SCAN + AI inference serialized):
--   SELECT LLM_GENERATE('model', content) FROM big_table WHERE complex_filter;
--
-- RECOMMENDED:

-- Step 1: Materialize the filtered input
CREATE TABLE temp_ai_input AS
SELECT id, content
FROM big_table
WHERE complex_filter
  AND created_at > now() - INTERVAL '30' DAY;

-- Step 2: Run AI inference on the known-size input
SELECT id, LLM_GENERATE('claude-sonnet-4-7', 'Summarize: ' || content) AS summary
FROM temp_ai_input;

-- Step 3: Clean up
DROP TABLE temp_ai_input;

-- Why it's better: the SCAN of big_table happens once; AI calls fire against a
-- stable row count. Latency is predictable. The batch threshold (64 rows, 100ms
-- flush) operates against a defined input rather than an unbounded stream.


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 4: Hybrid retrieval — SEMANTIC_SEARCH + RERANK in one query
-- ────────────────────────────────────────────────────────────────────────────
-- Retrieve a broad top-K candidate set via SEMANTIC_SEARCH (recall-optimized),
-- then re-rank with RERANK for higher precision. Both UDFs route through the
-- gateway; the audit chain captures the full pipeline.
--
-- This is the standard two-stage retrieval pattern (retrieve → rerank).

WITH
-- Stage 1: Broad semantic retrieval (top 50 candidates)
semantic_candidates AS (
    SELECT s.doc_id, s.score AS semantic_score
    FROM UNNEST(SEMANTIC_SEARCH('kb_vectors', 'how to configure rate limiting', 50)) AS s(doc_id, score)
),
-- Stage 2: Fetch content for candidate documents
candidate_content AS (
    SELECT c.doc_id, c.semantic_score, d.content
    FROM semantic_candidates c
    JOIN documents d ON c.doc_id = d.id
),
-- Stage 3: Re-rank by relevance (top 10 from 50 candidates)
reranked AS (
    SELECT r.doc_id, r.score AS rerank_score
    FROM UNNEST(
        RERANK(
            'cohere-rerank-3',
            'how to configure rate limiting',
            ARRAY(SELECT content FROM candidate_content)
        )
    ) AS r(doc_id, score)
)
SELECT cc.doc_id, cc.content, cc.semantic_score, rr.rerank_score
FROM candidate_content cc
JOIN reranked rr ON cc.doc_id = rr.doc_id
ORDER BY rr.rerank_score DESC
LIMIT 10;

-- Note: RERANK's doc_id is the candidate content string (not an id). The JOIN
-- matches on content = doc_id, so content values must be unique in the candidate
-- set for the join to work correctly. If content is not unique, include the
-- document id as a prefix:
--   ARRAY(SELECT id || '::' || content FROM candidate_content)
-- then strip the prefix after the JOIN.


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 5: Cross-tenant rejection awareness
-- ────────────────────────────────────────────────────────────────────────────
-- AI UDFs enforce tenant isolation at the /v1/batch/udf gateway boundary.
-- If your Trino principal (tenant-A:principal-1) calls a UDF referencing a
-- model or vector index owned by tenant-B, the per-item error 'tenant_mismatch'
-- returns SQL NULL + a WARNING log on the Trino worker.
--
-- This is NOT a bug in your query — it is an explicit security enforcement.
-- The wrong-tenant gate is checked after X-Trino-User identity resolution,
-- before any provider call is made.
--
-- To avoid cross-tenant surprises:
--   1. Ensure your Trino principal's tenant owns (or has access to) the model
--      asset and vector index asset you reference.
--   2. Check which models your tenant can access via the catalog browser
--      (AOCore admin UI) or via management API asset queries.
--   3. If you see NULL results with no other explanation, check Trino worker
--      WARNING logs for 'tenant_mismatch'.
--
-- Example: diagnosing tenant_mismatch from Trino worker log:
--   [WARN] UDFDispatch: per-item error for item id=<id> function=llm_generate
--          model=some-model: tenant_mismatch
--
-- The tenant_mismatch error code is also visible via the Trino query audit row:
--   SELECT updated_values->'ai_udfs_invoked' FROM audit_logs
--   WHERE action LIKE 'trino.query.%'
--     AND object_id LIKE 'trino-query/<query_id>/%';


-- ────────────────────────────────────────────────────────────────────────────
-- PATTERN 6: Eval-gate awareness (fallback warning)
-- ────────────────────────────────────────────────────────────────────────────
-- If a model version fails its eval gate and a previous approved version is
-- used as a fallback, the UDF output is returned with a JSON 'warning' field:
--   {"text": "...", "warning": "using_fallback_version:v1.0"}  (LLM_GENERATE)
--   {"label": "...", "warning": "..."}                        (LLM_CLASSIFY)
--
-- To detect fallback usage in results:
WITH results AS (
    SELECT id,
           LLM_GENERATE('gated-model', 'Summarize: ' || content) AS raw_output
    FROM articles
    WHERE published_at > now() - INTERVAL '1' DAY
)
SELECT
    id,
    json_extract_scalar(raw_output, '$.text')    AS summary,
    json_extract_scalar(raw_output, '$.warning') AS eval_gate_warning
FROM results
WHERE raw_output IS NOT NULL;
-- eval_gate_warning is NULL on normal (non-fallback) calls.
-- A non-null eval_gate_warning value indicates the response came from a
-- fallback model version, not the requested one. Consult the
-- udf.eval_gate.fallback audit chain entry for details.