-- pg_steadytext--2025.11.25.sql
-- Complete schema for pg_steadytext extension version 2025.11.25
-- Includes schema qualification fixes and summarization enhancements

-- AIDEV-SECTION: CORE_TABLE_DEFINITIONS
-- AIDEV-NOTE: Drop tables before creating to ensure proper extension membership
-- When tables are created with IF NOT EXISTS and already exist, they won't be added to the extension
DROP TABLE IF EXISTS steadytext_cache CASCADE;

-- Cache table that mirrors and extends SteadyText's SQLite cache
CREATE TABLE steadytext_cache (
    id SERIAL PRIMARY KEY,
    cache_key TEXT UNIQUE NOT NULL,  -- Matches SteadyText's cache key generation
    prompt TEXT NOT NULL,
    response TEXT,
    embedding vector(1024),  -- For embedding cache using pgvector

    -- Frecency statistics (synced with SteadyText's cache)
    access_count INT DEFAULT 1,
    last_accessed TIMESTAMPTZ DEFAULT NOW(),
    created_at TIMESTAMPTZ DEFAULT NOW(),

    -- SteadyText integration metadata
    steadytext_cache_hit BOOLEAN DEFAULT FALSE,  -- Whether this came from ST's cache
    model_name TEXT NOT NULL DEFAULT 'qwen3-1.7b',  -- Model used (supports switching)
    model_size TEXT CHECK (model_size IN ('small', 'medium', 'large')),
    seed INTEGER DEFAULT 42,  -- Seed used for generation
    eos_string TEXT,  -- Custom end-of-sequence string if used

    -- Generation parameters
    generation_params JSONB,  -- temperature, max_tokens, seed, etc.
    response_size INT,
    generation_time_ms INT  -- Time taken to generate (if not cached)

    -- AIDEV-NOTE: frecency_score removed - calculated via view instead
    -- Previously used GENERATED column with NOW() which is not immutable
);

-- Create indexes for performance
CREATE INDEX IF NOT EXISTS idx_steadytext_cache_key ON steadytext_cache(cache_key);
CREATE INDEX IF NOT EXISTS idx_steadytext_cache_last_accessed ON steadytext_cache(last_accessed);
CREATE INDEX IF NOT EXISTS idx_steadytext_cache_access_count ON steadytext_cache(access_count);

-- Request queue for async operations with priority and resource management
DROP TABLE IF EXISTS steadytext_queue CASCADE;
CREATE TABLE steadytext_queue (
    id SERIAL PRIMARY KEY,
    request_id UUID DEFAULT gen_random_uuid(),
    request_type TEXT CHECK (request_type IN ('generate', 'embed', 'batch_embed', 'rerank', 'batch_rerank')),

    -- Request data
    prompt TEXT,  -- For single requests
    prompts TEXT[],  -- For batch requests
    params JSONB,  -- Model params, seed, etc.

    -- Priority and scheduling
    priority INT DEFAULT 5 CHECK (priority BETWEEN 1 AND 10),
    user_id TEXT,  -- For rate limiting per user
    session_id TEXT,  -- For request grouping

    -- Status tracking
    status TEXT DEFAULT 'pending' CHECK (status IN ('pending', 'processing', 'completed', 'failed', 'cancelled')),
    result TEXT,
    results TEXT[],  -- For batch results
    embedding vector(1024),
    embeddings vector(1024)[],  -- For batch embeddings
    error TEXT,

    -- Timing
    created_at TIMESTAMPTZ DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    processing_time_ms INT,

    -- Resource tracking
    retry_count INT DEFAULT 0,
    max_retries INT DEFAULT 3,
    daemon_endpoint TEXT  -- Which daemon instance handled this
);

CREATE INDEX IF NOT EXISTS idx_steadytext_queue_status_priority_created ON steadytext_queue(status, priority DESC, created_at);
CREATE INDEX IF NOT EXISTS idx_steadytext_queue_request_id ON steadytext_queue(request_id);
CREATE INDEX IF NOT EXISTS idx_steadytext_queue_user_created ON steadytext_queue(user_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_steadytext_queue_session ON steadytext_queue(session_id);

-- Configuration storage
DROP TABLE IF EXISTS steadytext_config CASCADE;
CREATE TABLE steadytext_config (
    key TEXT PRIMARY KEY,
    value JSONB NOT NULL,
    description TEXT,
    updated_at TIMESTAMPTZ DEFAULT NOW(),
    updated_by TEXT DEFAULT current_user
);

-- Insert default configuration
INSERT INTO @extschema@.steadytext_config (key, value, description) VALUES
    ('daemon_host', '"localhost"', 'SteadyText daemon host'),
    ('daemon_port', '5555', 'SteadyText daemon port'),
    ('cache_enabled', 'false', 'Enable caching'),
    ('cache_eviction_enabled', 'true', 'Enable automatic cache eviction'),
    ('cache_max_entries', '10000', 'Maximum cache entries'),
    ('cache_max_size_mb', '100', 'Maximum cache size in MB'),
    ('default_max_tokens', '512', 'Default max tokens for generation'),
    ('default_seed', '42', 'Default seed for deterministic generation'),
    ('daemon_auto_start', 'true', 'Auto-start daemon if not running')
    ON CONFLICT (key) DO NOTHING;

-- Helper to raise standardized extension errors
CREATE OR REPLACE FUNCTION @extschema@._steadytext_raise_p0001(message TEXT)
RETURNS void
LANGUAGE plpgsql
AS $$
BEGIN
    RAISE EXCEPTION USING ERRCODE = 'P0001', MESSAGE = message;
END;
$$;

-- Daemon health monitoring
DROP TABLE IF EXISTS steadytext_daemon_health CASCADE;
CREATE TABLE steadytext_daemon_health (
    daemon_id TEXT PRIMARY KEY DEFAULT 'default',
    endpoint TEXT NOT NULL,
    last_heartbeat TIMESTAMPTZ DEFAULT NOW(),
    uptime_seconds INT DEFAULT 0,
    status TEXT DEFAULT 'unknown' CHECK (status IN ('healthy', 'unhealthy', 'starting', 'stopping', 'unknown')),
    version TEXT,
    models_loaded TEXT[],
    memory_usage_mb INT,
    active_connections INT DEFAULT 0,
    total_requests BIGINT DEFAULT 0,
    error_count INT DEFAULT 0,
    avg_response_time_ms INT
);

-- Insert default daemon entry
INSERT INTO @extschema@.steadytext_daemon_health (daemon_id, endpoint, status, uptime_seconds)
VALUES ('default', 'tcp://localhost:5555', 'unknown', 0)
ON CONFLICT (daemon_id) DO NOTHING;

-- Rate limiting per user
DROP TABLE IF EXISTS steadytext_rate_limits CASCADE;
CREATE TABLE steadytext_rate_limits (
    user_id TEXT PRIMARY KEY,
    requests_per_minute INT DEFAULT 60,
    requests_per_hour INT DEFAULT 1000,
    requests_per_day INT DEFAULT 10000,
    current_minute_count INT DEFAULT 0,
    current_hour_count INT DEFAULT 0,
    current_day_count INT DEFAULT 0,
    last_reset TIMESTAMPTZ DEFAULT NOW(),
    last_reset_minute TIMESTAMPTZ DEFAULT NOW(),
    last_reset_hour TIMESTAMPTZ DEFAULT NOW(),
    last_reset_day TIMESTAMPTZ DEFAULT NOW(),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Audit log for security and debugging
DROP TABLE IF EXISTS steadytext_audit_log CASCADE;
CREATE TABLE steadytext_audit_log (
    id SERIAL PRIMARY KEY,
    event_id UUID DEFAULT gen_random_uuid(),
    event_type TEXT NOT NULL,
    user_id TEXT DEFAULT current_user,
    request_id UUID,
    details JSONB,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    ip_address INET,
    success BOOLEAN DEFAULT TRUE,
    error_message TEXT
);

CREATE INDEX IF NOT EXISTS idx_steadytext_audit_timestamp ON steadytext_audit_log(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_steadytext_audit_user ON steadytext_audit_log(user_id, created_at DESC);

-- AIDEV-SECTION: VIEWS
-- AIDEV-NOTE: View steadytext_cache_with_frecency is created later (line ~1115)
-- with age-based scoring for v1.4.1+ compatibility

-- AIDEV-SECTION: PYTHON_INTEGRATION
-- AIDEV-NOTE: Python integration layer path setup
-- This is now handled by the _steadytext_init_python function instead

-- AIDEV-ANCHOR: init-python: steadytext-only bootstrap
-- Create Python function container
CREATE OR REPLACE FUNCTION _steadytext_init_python()
RETURNS void
LANGUAGE plpython3u
AS $c$
# AIDEV-NOTE: Initialize Python environment for pg_steadytext with enhanced error handling
import sys
import os
import site

# Get PostgreSQL lib directory with fallback
try:
    result = plpy.execute("SELECT setting FROM pg_settings WHERE name = 'pkglibdir'")
    if result and len(result) > 0 and result[0]['setting']:
        pg_lib_dir = result[0]['setting']
    else:
        # Fallback for Docker/Debian PostgreSQL 17
        pg_lib_dir = '/usr/lib/postgresql/17/lib'
        plpy.notice(f"Using fallback pkglibdir: {pg_lib_dir}")
except Exception as e:
    # Fallback for Docker/Debian PostgreSQL 17
    pg_lib_dir = '/usr/lib/postgresql/17/lib'
    plpy.notice(f"Error getting pkglibdir, using fallback: {pg_lib_dir}")

python_module_dir = os.path.join(pg_lib_dir, 'pg_steadytext', 'python')

# Save pg_lib_dir in GD for use in error messages
GD['pg_lib_dir'] = pg_lib_dir

# Verify the directory exists
if not os.path.exists(python_module_dir):
    plpy.error(f"Python module directory not found: {python_module_dir}")

# Add to Python path if not already there
if python_module_dir not in sys.path:
    sys.path.insert(0, python_module_dir)
    site.addsitedir(python_module_dir)  # Process .pth files if any

# AIDEV-NOTE: Add site-packages directory for locally installed Python packages
site_packages_dir = os.path.join(pg_lib_dir, 'pg_steadytext', 'site-packages')
if os.path.exists(site_packages_dir) and site_packages_dir not in sys.path:
    sys.path.insert(0, site_packages_dir)
    site.addsitedir(site_packages_dir)
    plpy.notice(f"Added site-packages to path: {site_packages_dir}")

# AIDEV-NOTE: Add common Python package locations
# These are common locations where pip might install packages
common_paths = [
    # User site-packages
    site.getusersitepackages(),
    # System-wide site-packages
    '/usr/local/lib/python3.10/dist-packages',
    '/usr/local/lib/python3.11/dist-packages',
    '/usr/local/lib/python3.12/dist-packages',
    '/usr/lib/python3/dist-packages',
    # Virtual environment (if activated)
    os.path.join(sys.prefix, 'lib', f'python{sys.version_info.major}.{sys.version_info.minor}', 'site-packages'),
]

for path in common_paths:
    if path and os.path.exists(path) and path not in sys.path:
        sys.path.append(path)

# Log Python path for debugging
plpy.notice(f"Python path: {sys.path}")
plpy.notice(f"Looking for modules in: {python_module_dir}")

# Check if directory exists
if not os.path.exists(python_module_dir):
    plpy.error(f"Python module directory does not exist: {python_module_dir}")

# List files in directory for debugging
try:
    files = os.listdir(python_module_dir)
    plpy.notice(f"Files in module directory: {files}")
except Exception as e:
    plpy.warning(f"Could not list module directory: {e}")

# Try to import required external packages first
required_packages = {
    'steadytext': 'SteadyText library',
    'zmq': 'ZeroMQ for daemon communication',
    'numpy': 'NumPy for embeddings'
}

missing_packages = []
for package, description in required_packages.items():
    try:
        __import__(package)
    except ImportError:
        missing_packages.append(f"{package} ({description})")

if missing_packages:
    pg_lib_dir = GD.get('pg_lib_dir', '/usr/lib/postgresql/17/lib')
    site_packages_dir = os.path.join(pg_lib_dir, 'pg_steadytext', 'site-packages')

    error_msg = f"""
=================================================================
Missing required Python packages: {', '.join(missing_packages)}

The pg_steadytext extension requires these packages to function.

To fix this, run ONE of the following commands:

1. Install via make (recommended):
   cd /path/to/pg_steadytext
   sudo make install

2. Install to PostgreSQL's site-packages:
   sudo pip3 install --target={site_packages_dir} steadytext pyzmq numpy

3. Install system-wide:
   sudo pip3 install steadytext pyzmq numpy

4. Install to user directory:
   pip3 install --user steadytext pyzmq numpy

After installation, restart PostgreSQL and try again.
=================================================================
"""
    plpy.error(error_msg)

# Try to import our modules and cache them in GD
try:
    # Clear any previous module cache
    for key in list(GD.keys()):
        if key.startswith('module_'):
            del GD[key]

    # Import and cache modules
    import security
    import prompt_registry

    # Store modules in GD for reuse
    GD['module_security'] = security
    GD['module_prompt_registry'] = prompt_registry
    GD['steadytext_initialized'] = True

    plpy.notice(f"pg_steadytext Python environment initialized successfully from {python_module_dir}")
except ImportError as e:
    GD['steadytext_initialized'] = False
    pg_lib_dir = GD.get('pg_lib_dir', '/usr/lib/postgresql/17/lib')
    site_packages_dir = os.path.join(pg_lib_dir, 'pg_steadytext', 'site-packages')

    error_msg = f"""
=================================================================
Failed to import pg_steadytext modules from {python_module_dir}

Error: {str(e)}

This usually means the extension files are installed but Python
module files are missing or there's an import error.

To fix this:

1. Ensure all Python module files are present in:
   {python_module_dir}

2. Check that required packages are installed:
   sudo pip3 install --target={site_packages_dir} steadytext pyzmq numpy

3. Or reinstall the extension:
   cd /path/to/pg_steadytext
   sudo make install

After fixing, restart PostgreSQL and try again.
=================================================================
"""
    plpy.error(error_msg)
except Exception as e:
    GD['steadytext_initialized'] = False
    plpy.error(f"Unexpected error during initialization: {e}")
$c$;

-- AIDEV-NOTE: Initialization is now done on-demand in each function
-- This ensures proper initialization even across session boundaries

-- AIDEV-SECTION: CORE_FUNCTIONS
-- Core function: Synchronous text generation
-- Returns NULL if generation fails (no fallback text)
-- Handle removal of old function signature before creating new one
DO $$
BEGIN
    -- Try to remove old function from extension if it exists
    BEGIN
        ALTER EXTENSION pg_steadytext DROP FUNCTION steadytext_generate(text, integer, boolean, integer);
    EXCEPTION WHEN OTHERS THEN
        -- Function either doesn't exist or isn't part of extension
        NULL;
    END;
    
    -- Try to drop old function if it exists
    DROP FUNCTION IF EXISTS steadytext_generate(text, integer, boolean, integer);
END $$;

-- Streaming generation wrapper that yields tokens from steadytext_generate
DO $$
BEGIN
    BEGIN
        ALTER EXTENSION pg_steadytext DROP FUNCTION steadytext_generate_stream(text, integer);
    EXCEPTION WHEN OTHERS THEN
        NULL;
    END;

    DROP FUNCTION IF EXISTS steadytext_generate_stream(text, integer);
END $$;

CREATE OR REPLACE FUNCTION steadytext_generate_stream(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    eos_string TEXT DEFAULT '[EOS]',
    model TEXT DEFAULT NULL,
    model_repo TEXT DEFAULT NULL,
    model_filename TEXT DEFAULT NULL,
    size TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS SETOF TEXT
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
import json
import os
import subprocess

# Initialize Python environment if necessary
if not GD.get('steadytext_initialized', False):
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

# Resolve defaults from configuration
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

config_plan = GD.get('steadytext_config_plan_stream')
if config_plan is None:
    config_plan = plpy.prepare(f"SELECT value FROM {plpy.quote_ident(ext_schema)}.steadytext_config WHERE key = $1", ["text"])
    GD['steadytext_config_plan_stream'] = config_plan

resolved_max_tokens = max_tokens
if resolved_max_tokens is None:
    rv = plpy.execute(config_plan, ["default_max_tokens"])
    resolved_max_tokens = json.loads(rv[0]["value"]) if rv else 512

resolved_seed = seed
if resolved_seed is None:
    rv = plpy.execute(config_plan, ["default_seed"])
    resolved_seed = json.loads(rv[0]["value"]) if rv else 42

# Resolve daemon settings to honor database configuration
host_rv = plpy.execute(config_plan, ["daemon_host"])
daemon_host = json.loads(host_rv[0]["value"]) if host_rv else "localhost"

port_rv = plpy.execute(config_plan, ["daemon_port"])
daemon_port = json.loads(port_rv[0]["value"]) if port_rv else 5555

auto_start_rv = plpy.execute(config_plan, ["daemon_auto_start"])
daemon_auto_start = json.loads(auto_start_rv[0]["value"]) if auto_start_rv else True

# Basic null/empty checks
if prompt is None:
    plpy.error("Prompt cannot be null")
if prompt.strip() == "":
    plpy.error("Prompt cannot be empty")

# Validate unsafe_mode
if not unsafe_mode and model and ':' in model:
    plpy.error("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Configure daemon endpoint for steadytext client
os.environ["STEADYTEXT_DAEMON_HOST"] = str(daemon_host)
os.environ["STEADYTEXT_DAEMON_PORT"] = str(daemon_port)

# Best-effort auto-start for local daemon (streaming)
is_remote_model = unsafe_mode and model and ':' in model
if daemon_auto_start and not is_remote_model:
    try:
        from steadytext.daemon.client import DaemonClient

        client = DaemonClient(host=daemon_host, port=daemon_port)
        if not client.connect():
            subprocess.run(
                [
                    'st', 'daemon', 'start',
                    '--host', str(daemon_host),
                    '--port', str(daemon_port)
                ],
                check=False,
                capture_output=True,
                text=True,
                timeout=10
            )
    except Exception as e:
        plpy.warning(f"Daemon auto-start skipped: {e}")

# Call steadytext directly with streaming
from steadytext import generate

kwargs = {"seed": resolved_seed, "stream": True}
if eos_string and eos_string != '[EOS]':
    kwargs["eos_string"] = eos_string
if model:
    kwargs["model"] = model
if model_repo:
    kwargs["model_repo"] = model_repo
if model_filename:
    kwargs["model_filename"] = model_filename
if size:
    kwargs["size"] = size
if unsafe_mode:
    kwargs["unsafe_mode"] = unsafe_mode

# Generate and yield chunks
for chunk in generate(prompt, max_new_tokens=resolved_max_tokens, **kwargs):
    yield (chunk,)
$c$;

DO $$
BEGIN
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION steadytext_generate_stream(text, integer, boolean, integer, text, text, text, text, text, boolean);
    EXCEPTION WHEN OTHERS THEN
        NULL;
    END;
END $$;

CREATE OR REPLACE FUNCTION steadytext_generate(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    eos_string TEXT DEFAULT '[EOS]',
    model TEXT DEFAULT NULL,
    model_repo TEXT DEFAULT NULL,
    model_filename TEXT DEFAULT NULL,
    size TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS TEXT
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE  -- IMMUTABLE: Intentional for query optimization despite daemon interaction
AS $c$
# AIDEV-NOTE: Main text generation function that calls steadytext directly
# v2025.8.15: Added support for eos_string, model, model_repo, model_filename, size parameters
# v2025.8.15: Added model parameter for remote model access
import json
import os
import subprocess

def raise_p0001(message: str) -> None:
    plan = GD.get('steadytext_raise_plan')
    if plan is None:
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_raise_p0001($1)", ["text"])
        GD['steadytext_raise_plan'] = plan

    plpy.execute(plan, [message])

# Check if initialized, if not, initialize now
if not GD.get('steadytext_initialized', False):
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

# Get configuration
# AIDEV-NOTE: Use extension schema instead of current_schema() for cross-schema compatibility
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
config_select_plan = plpy.prepare(f"SELECT value FROM {plpy.quote_ident(ext_schema)}.steadytext_config WHERE key = $1", ["text"])

# Resolve max_tokens, using the provided value or fetching the default
resolved_max_tokens = max_tokens
if resolved_max_tokens is None:
    rv = plpy.execute(config_select_plan, ["default_max_tokens"])
    resolved_max_tokens = json.loads(rv[0]["value"]) if rv else 512

# Resolve seed, using the provided value or fetching the default
resolved_seed = seed
if resolved_seed is None:
    rv = plpy.execute(config_select_plan, ["default_seed"])
    resolved_seed = json.loads(rv[0]["value"]) if rv else 42

# Resolve daemon settings from configuration (preserve backward compatibility)
host_rv = plpy.execute(config_select_plan, ["daemon_host"])
daemon_host = json.loads(host_rv[0]["value"]) if host_rv else "localhost"

port_rv = plpy.execute(config_select_plan, ["daemon_port"])
daemon_port = json.loads(port_rv[0]["value"]) if port_rv else 5555

auto_start_rv = plpy.execute(config_select_plan, ["daemon_auto_start"])
daemon_auto_start = json.loads(auto_start_rv[0]["value"]) if auto_start_rv else True

# Basic null/empty checks (validation is application's responsibility)
if prompt is None:
    raise_p0001("Prompt cannot be null")
if prompt.strip() == "":
    raise_p0001("Prompt cannot be empty")
if resolved_seed < 0:
    raise_p0001("seed must be non-negative")

# Validate unsafe_mode usage (keep this security check)
if not unsafe_mode and model and ':' in model:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")
if unsafe_mode and not model:
    raise_p0001("unsafe_mode=TRUE requires a model parameter")

# Configure daemon endpoint for the steadytext client
os.environ["STEADYTEXT_DAEMON_HOST"] = str(daemon_host)
os.environ["STEADYTEXT_DAEMON_PORT"] = str(daemon_port)

# Best-effort auto-start for local daemon when enabled and using local models
is_remote_model = unsafe_mode and model and ':' in model
if daemon_auto_start and not is_remote_model:
    try:
        from steadytext.daemon.client import DaemonClient

        client = DaemonClient(host=daemon_host, port=daemon_port)
        if not client.connect():
            # Attempt to start the daemon; ignore failures to preserve determinism
            subprocess.run(
                [
                    'st', 'daemon', 'start',
                    '--host', str(daemon_host),
                    '--port', str(daemon_port)
                ],
                check=False,
                capture_output=True,
                text=True,
                timeout=10
            )
    except Exception as e:
        plpy.warning(f"Daemon auto-start skipped: {e}")

# Check if we should use cache
if use_cache:
    # Generate cache key consistent with SteadyText format
    # Include eos_string in cache key if it's not the default
    if eos_string == '[EOS]':
        cache_key = prompt
    else:
        cache_key = f"{prompt}::EOS::{eos_string}"

    # Try to get from cache first
    # AIDEV-NOTE: Get schema dynamically at runtime for TimescaleDB continuous aggregates compatibility
    cache_plan = plpy.prepare(f"""
        SELECT response
        FROM {plpy.quote_ident(ext_schema)}.steadytext_cache
        WHERE cache_key = $1
    """, ["text"])

    cache_result = plpy.execute(cache_plan, [cache_key])
    if cache_result and cache_result[0]["response"]:
        plpy.notice(f"Cache hit for key: {cache_key[:8]}...")
        cached_value = cache_result[0]["response"]
        if isinstance(cached_value, (dict, list)):
            return json.dumps(cached_value)
        return cached_value

# Cache miss - generate new content
# Call steadytext directly
from steadytext import generate

kwargs = {"seed": resolved_seed}
if eos_string and eos_string != '[EOS]':
    kwargs["eos_string"] = eos_string
if model:
    kwargs["model"] = model
if model_repo:
    kwargs["model_repo"] = model_repo
if model_filename:
    kwargs["model_filename"] = model_filename
if size:
    kwargs["size"] = size
if unsafe_mode:
    kwargs["unsafe_mode"] = unsafe_mode

result = generate(prompt, max_new_tokens=resolved_max_tokens, **kwargs)

# Store in cache if enabled
if use_cache and result:
    # Generate cache key (same logic as cache read)
    if eos_string == '[EOS]':
        write_cache_key = prompt
    else:
        write_cache_key = f"{prompt}::EOS::{eos_string}"

    insert_plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_cache
        (cache_key, prompt, response, model_name, seed, eos_string, generation_params)
        VALUES ($1, $2, $3, $4, $5, $6, $7)
        ON CONFLICT (cache_key) DO UPDATE
        SET response = EXCLUDED.response,
            access_count = steadytext_cache.access_count + 1,
            last_accessed = NOW()
    """, ["text", "text", "text", "text", "integer", "text", "jsonb"])

    params = {"max_tokens": resolved_max_tokens, "seed": resolved_seed}
    model_name = model if model else 'steadytext-default'
    eos_val = eos_string if eos_string != '[EOS]' else None
    plpy.execute(insert_plan, [write_cache_key, prompt, result, model_name, resolved_seed, eos_val, json.dumps(params)])
    plpy.notice(f"Cached response with key: {write_cache_key[:8]}...")

return result
$c$;

-- Add new function to extension (idempotent)
DO $$
BEGIN
    -- Try to add function to extension
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION steadytext_generate(text, integer, boolean, integer, text, text, text, text, text, boolean);
    EXCEPTION WHEN OTHERS THEN
        -- Function is already part of extension
        NULL;
    END;
END $$;

-- AIDEV-ANCHOR: embed: daemon-fallback then steadytext
-- Core function: Synchronous embedding generation
-- Returns NULL if embedding generation fails (no fallback vector)
CREATE OR REPLACE FUNCTION steadytext_embed(
    text_input TEXT,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS vector(1024)
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE LEAKPROOF  -- IMMUTABLE: Intentional for query optimization despite model interaction
AS $c$
# AIDEV-NOTE: Embedding function with remote model support via unsafe_mode
# Added in v2025.8.15 to support OpenAI and other remote embedding providers
# Uses runtime inspection to maintain backward compatibility with older steadytext versions

import json
import numpy as np

# Shared helper to raise SQLSTATE P0001
def raise_p0001(message: str) -> None:
    security = GD.get('module_security')
    if security and hasattr(security, 'raise_sqlstate'):
        security.raise_sqlstate(message, "P0001")
        return

    plan = GD.get('steadytext_raise_plan')
    if plan is None:
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
        GD['steadytext_raise_plan'] = plan

    plpy.execute(plan, [message])

# Get extension schema for all subsequent queries
# AIDEV-NOTE: Define ext_schema at function level to avoid UnboundLocalError
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Initialize Python environment if needed
try:
    if 'steadytext_initialized' not in GD:
        plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
except:
    pass

# Basic null/empty check (validation is application's responsibility)
if text_input is None:
    raise_p0001("Text cannot be null")
if text_input.strip() == "":
    raise_p0001("Text cannot be empty")

# Validate unsafe_mode for remote models
if not unsafe_mode and model and ':' in model:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Check cache if enabled (IMMUTABLE functions only read cache)
if use_cache:
    # Generate cache key including model if specified
    cache_key_parts = ['embed', text_input]
    if model:
        cache_key_parts.append(model)
    cache_key_input = ':'.join(cache_key_parts)
    
    import hashlib
    cache_key = hashlib.sha256(cache_key_input.encode()).hexdigest()
    
    # Try to get from cache (read-only for IMMUTABLE)
    # AIDEV-NOTE: Use extension schema for TimescaleDB continuous aggregates compatibility
    plan = plpy.prepare(
        f"SELECT embedding FROM {plpy.quote_ident(ext_schema)}.steadytext_cache WHERE cache_key = $1",
        ["text"]
    )
    rv = plpy.execute(plan, [cache_key])
    
    if rv and len(rv) > 0:
        cached_embedding = rv[0]['embedding']
        if cached_embedding:
            return cached_embedding

# AIDEV-NOTE: Call steadytext.embed() directly - it handles daemon fallback internally
# The steadytext library automatically uses daemon when available and falls back to direct loading
try:
    from steadytext import embed as steadytext_embed
    import inspect
    import hashlib

    # Check if embed supports the new parameters
    embed_sig = inspect.signature(steadytext_embed)

    # Call with model and unsafe_mode if specified and supported
    kwargs = {'seed': seed}
    if model and 'model' in embed_sig.parameters:
        kwargs['model'] = model
    if unsafe_mode and 'unsafe_mode' in embed_sig.parameters:
        kwargs['unsafe_mode'] = unsafe_mode

    result = steadytext_embed(text_input, **kwargs)

    # Convert to vector format if needed
    if result is not None:
        # Ensure it's a list/array
        if hasattr(result, 'tolist'):
            embedding_list = result.tolist()
        else:
            embedding_list = list(result)

        # Store in cache if enabled
        if use_cache and embedding_list:
            # Generate cache key (same logic as cache read)
            cache_key_parts = ['embed', text_input]
            if model:
                cache_key_parts.append(model)
            cache_key_input = ':'.join(cache_key_parts)
            write_cache_key = hashlib.sha256(cache_key_input.encode()).hexdigest()

            insert_plan = plpy.prepare(f"""
                INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_cache
                (cache_key, prompt, embedding, model_name, seed)
                VALUES ($1, $2, $3::vector, $4, $5)
                ON CONFLICT (cache_key) DO UPDATE
                SET embedding = EXCLUDED.embedding,
                    access_count = steadytext_cache.access_count + 1,
                    last_accessed = NOW()
            """, ["text", "text", "text", "text", "integer"])

            model_name = model if model else 'steadytext-embedding'
            plpy.execute(insert_plan, [write_cache_key, text_input, str(embedding_list), model_name, seed])
            plpy.notice(f"Cached embedding with key: {write_cache_key[:8]}...")

        return embedding_list
except Exception as e:
    plpy.warning(f"Embedding generation failed: {e}")
    return [0.0] * 1024

# Fallback zero vector if no other return path executed
return [0.0] * 1024
$c$;

-- AIDEV-ANCHOR: daemon: start-status-stop helpers
-- Daemon management functions
CREATE OR REPLACE FUNCTION steadytext_daemon_start()
RETURNS BOOLEAN
LANGUAGE plpython3u
AS $c$
# AIDEV-NOTE: Start the SteadyText daemon if not already running
import subprocess
import time
import json

try:
    # Get daemon configuration
    # AIDEV-NOTE: Use extension schema for cross-schema compatibility
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    config_plan = plpy.prepare(f"SELECT value FROM {plpy.quote_ident(ext_schema)}.steadytext_config WHERE key = $1", ["text"])
    host_rv = plpy.execute(config_plan, ["daemon_host"])
    port_rv = plpy.execute(config_plan, ["daemon_port"])

    host = json.loads(host_rv[0]["value"]) if host_rv else "localhost"
    port = json.loads(port_rv[0]["value"]) if port_rv else 5555

    # Check if daemon is already running by trying to start it
    # SteadyText daemon start command is idempotent
    try:
        result = subprocess.run(['st', 'daemon', 'start'], capture_output=True, text=True, timeout=10)

        if result.returncode == 0:
            # Update health status
            health_plan = plpy.prepare(f"""
                UPDATE {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
                SET status = 'healthy',
                    last_heartbeat = NOW()
                WHERE daemon_id = 'default'
            """)
            plpy.execute(health_plan)
            return True
        else:
            plpy.warning(f"Failed to start daemon: {result.stderr}")
            return False

    except subprocess.TimeoutExpired:
        plpy.warning("Timeout starting daemon")
        return False

except Exception as e:
    plpy.error(f"Error in daemon start: {e}")
    return False
$c$;

-- Get daemon status
CREATE OR REPLACE FUNCTION steadytext_daemon_status()
RETURNS TABLE(
    daemon_id TEXT,
    status TEXT,
    endpoint TEXT,
    last_heartbeat TIMESTAMPTZ,
    uptime_seconds INT
)
LANGUAGE plpython3u
AS $c$
# AIDEV-NOTE: Check SteadyText daemon health status
import json

# Get extension schema first (always needed)
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Check if initialized, if not, initialize now
if not GD.get('steadytext_initialized', False):
    # Initialize on demand
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    # Check again after initialization
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

try:
    import subprocess
    result = subprocess.run(['st', 'daemon', 'status', '--json'], capture_output=True, text=True, timeout=5)

    if result.returncode == 0:
        status_data = json.loads(result.stdout)
        if status_data.get('running'):
            # Update health status to healthy
            update_plan = plpy.prepare(f"""
                UPDATE {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
                SET status = 'healthy',
                    last_heartbeat = NOW(),
                    uptime_seconds = GREATEST(EXTRACT(EPOCH FROM (NOW() - COALESCE(last_heartbeat, NOW()))), 0)::INT
                WHERE daemon_id = 'default'
                RETURNING daemon_id, status, endpoint, last_heartbeat, uptime_seconds
            """)
            result = plpy.execute(update_plan)
            return result
        else:
            # Update health status to unhealthy (daemon not running)
            update_plan = plpy.prepare(f"""
                UPDATE {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
                SET status = 'unhealthy',
                    last_heartbeat = NOW(),
                    uptime_seconds = 0
                WHERE daemon_id = 'default'
                RETURNING daemon_id, status, endpoint, last_heartbeat, uptime_seconds
            """)
            result = plpy.execute(update_plan)
            return result
    else:
        # Command failed, mark as unhealthy
        update_plan = plpy.prepare(f"""
            UPDATE {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
            SET status = 'unhealthy',
                last_heartbeat = NOW(),
                uptime_seconds = 0
            WHERE daemon_id = 'default'
            RETURNING daemon_id, status, endpoint, last_heartbeat, uptime_seconds
        """)
        result = plpy.execute(update_plan)
        return result

except Exception as e:
    plpy.warning(f"Error checking daemon status: {e}")
    # Return current status from table
    select_plan = plpy.prepare(f"""
        SELECT daemon_id, status, endpoint, last_heartbeat, uptime_seconds
        FROM {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
        WHERE daemon_id = 'default'
    """)
    return plpy.execute(select_plan)
$c$;

-- Stop daemon
CREATE OR REPLACE FUNCTION steadytext_daemon_stop()
RETURNS BOOLEAN
LANGUAGE plpython3u
AS $c$
# AIDEV-NOTE: Stop the SteadyText daemon gracefully without disrupting PostgreSQL
import json

# Resolve extension schema for subsequent queries
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Ensure Python modules are initialized so cached connectors are available
if not GD.get('steadytext_initialized', False):
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")

stopped_successfully = False
try:
    import subprocess
    result = subprocess.run(['st', 'daemon', 'stop'], capture_output=True, text=True, timeout=10)

    if result.returncode == 0:
        stopped_successfully = True
        plpy.notice("Daemon stopped successfully")
    else:
        stopped_successfully = False
        plpy.warning(f"Failed to stop daemon: {result.stderr}")

except subprocess.TimeoutExpired:
    stopped_successfully = False
    plpy.warning("Daemon stop command timed out")
except Exception as e:
    stopped_successfully = False
    plpy.warning(f"Error stopping daemon: {e}")

# Update daemon health status regardless of stop outcome
status_value = 'unhealthy'
update_plan = plpy.prepare(f"""
    UPDATE {plpy.quote_ident(ext_schema)}.steadytext_daemon_health
    SET status = $1,
        last_heartbeat = NOW(),
        uptime_seconds = 0
    WHERE daemon_id = 'default'
""", ["text"])
plpy.execute(update_plan, [status_value])

return stopped_successfully
$c$;

-- AIDEV-SECTION: CACHE_MANAGEMENT_FUNCTIONS
-- Cache management functions
CREATE OR REPLACE FUNCTION steadytext_cache_stats()
RETURNS TABLE(
    total_entries BIGINT,
    total_size_mb FLOAT,
    cache_hit_rate FLOAT,
    avg_access_count FLOAT,
    oldest_entry TIMESTAMPTZ,
    newest_entry TIMESTAMPTZ
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $c$
    SELECT
        COUNT(*)::BIGINT as total_entries,
        COALESCE(SUM(pg_column_size(response) + pg_column_size(embedding)) / 1024.0 / 1024.0, 0)::FLOAT as total_size_mb,
        COALESCE(SUM(CASE WHEN access_count > 1 THEN 1 ELSE 0 END)::FLOAT / NULLIF(COUNT(*), 0), 0)::FLOAT as cache_hit_rate,
        COALESCE(AVG(access_count), 0)::FLOAT as avg_access_count,
        MIN(created_at) as oldest_entry,
        MAX(created_at) as newest_entry
    FROM @extschema@.steadytext_cache;
$c$;

-- Clear cache
CREATE OR REPLACE FUNCTION steadytext_cache_clear()
RETURNS BIGINT
LANGUAGE sql
AS $c$
    WITH deleted AS (
        DELETE FROM @extschema@.steadytext_cache
        RETURNING *
    )
    SELECT COUNT(*) FROM deleted;
$c$;

-- AIDEV-SECTION: UPDATE_EVICTION_FUNCTIONS
-- Replace frecency-based eviction with age-based eviction
CREATE OR REPLACE FUNCTION steadytext_cache_evict_by_age(
    target_entries INT DEFAULT NULL,
    target_size_mb FLOAT DEFAULT NULL,
    batch_size INT DEFAULT 100,
    min_age_hours INT DEFAULT 1
)
RETURNS TABLE(
    evicted_count INT,
    freed_size_mb FLOAT,
    remaining_entries BIGINT,
    remaining_size_mb FLOAT
)
LANGUAGE plpgsql
AS $c$
DECLARE
    v_evicted_count INT := 0;
    v_freed_size_mb FLOAT := 0;
    v_current_entries BIGINT;
    v_current_size_mb FLOAT;
    v_batch_evicted INT;
    v_batch_freed_mb FLOAT;
BEGIN
    -- AIDEV-NOTE: Simplified eviction using age-based strategy (FIFO)
    -- This replaces the frecency-based eviction to maintain IMMUTABLE functions
    
    -- Get current cache statistics
    SELECT 
        COUNT(*),
        COALESCE(SUM(pg_column_size(response) + pg_column_size(embedding)) / 1024.0 / 1024.0, 0)
    INTO v_current_entries, v_current_size_mb
    FROM @extschema@.steadytext_cache;
    
    -- Set default targets if not provided
    IF target_entries IS NULL THEN
        SELECT (value #>> '{}')::INT INTO target_entries 
        FROM @extschema@.steadytext_config 
        WHERE key = 'cache_max_entries';
        target_entries := COALESCE(target_entries, 10000);
    END IF;
    
    IF target_size_mb IS NULL THEN
        SELECT (value #>> '{}')::FLOAT INTO target_size_mb 
        FROM @extschema@.steadytext_config 
        WHERE key = 'cache_max_size_mb';
        target_size_mb := COALESCE(target_size_mb, 100);
    END IF;
    
    -- Evict in batches until we meet our targets
    WHILE (v_current_entries > target_entries OR v_current_size_mb > target_size_mb) LOOP
        -- Delete oldest entries (FIFO eviction)
        WITH deleted AS (
            DELETE FROM @extschema@.steadytext_cache
            WHERE id IN (
                SELECT id 
                FROM @extschema@.steadytext_cache
                WHERE created_at < NOW() - INTERVAL '1 hour' * min_age_hours
                ORDER BY created_at ASC  -- Oldest first
                LIMIT batch_size
            )
            RETURNING pg_column_size(response) + pg_column_size(embedding) as size_bytes
        )
        SELECT 
            COUNT(*),
            COALESCE(SUM(size_bytes) / 1024.0 / 1024.0, 0)
        INTO v_batch_evicted, v_batch_freed_mb
        FROM deleted;
        
        -- Break if nothing was evicted (all entries are too young)
        IF v_batch_evicted = 0 THEN
            EXIT;
        END IF;
        
        -- Update totals
        v_evicted_count := v_evicted_count + v_batch_evicted;
        v_freed_size_mb := v_freed_size_mb + v_batch_freed_mb;
        v_current_entries := v_current_entries - v_batch_evicted;
        v_current_size_mb := v_current_size_mb - v_batch_freed_mb;
        
        -- Log eviction batch
        INSERT INTO @extschema@.steadytext_audit_log (event_type, details)
        VALUES (
            'cache_eviction',
            jsonb_build_object(
                'evicted_count', v_batch_evicted,
                'freed_size_mb', v_batch_freed_mb,
                'eviction_type', 'age_based'
            )
        );
    END LOOP;
    
    -- Return results
    RETURN QUERY
    SELECT 
        v_evicted_count,
        v_freed_size_mb,
        v_current_entries,
        v_current_size_mb;
END;
$c$;

CREATE OR REPLACE FUNCTION steadytext_cache_evict_by_frecency(
    target_entries INT DEFAULT NULL,
    target_size_mb FLOAT DEFAULT NULL,
    batch_size INT DEFAULT 100,
    min_access_count INT DEFAULT 2,
    min_age_hours INT DEFAULT 1
)
RETURNS TABLE(
    evicted_count INT,
    freed_size_mb FLOAT,
    remaining_entries BIGINT,
    remaining_size_mb FLOAT
)
LANGUAGE sql
AS $c$
    SELECT * FROM steadytext_cache_evict_by_age(target_entries, target_size_mb, batch_size, min_age_hours);
$c$;


-- Update the scheduled eviction function to use age-based eviction
CREATE OR REPLACE FUNCTION steadytext_cache_scheduled_eviction()
RETURNS void
LANGUAGE plpgsql
AS $c$
DECLARE
    v_enabled BOOLEAN;
    v_result RECORD;
BEGIN
    -- Check if eviction is enabled
    SELECT (value #>> '{}')::BOOLEAN INTO v_enabled 
    FROM @extschema@.steadytext_config 
    WHERE key = 'cache_eviction_enabled';
    
    IF NOT COALESCE(v_enabled, TRUE) THEN
        RETURN;
    END IF;
    
    -- Perform age-based eviction
    SELECT * INTO v_result
    FROM steadytext_cache_evict_by_age();
    
    -- Log results if anything was evicted
    IF v_result.evicted_count > 0 THEN
        RAISE NOTICE 'Cache eviction completed: % entries evicted, % MB freed',
            v_result.evicted_count, v_result.freed_size_mb;
    END IF;
END;
$c$;

-- Update cache preview to show age instead of frecency
CREATE OR REPLACE FUNCTION steadytext_cache_preview_eviction(
    preview_count INT DEFAULT 10
)
RETURNS TABLE(
    cache_key TEXT,
    prompt TEXT,
    access_count INT,
    last_accessed TIMESTAMPTZ,
    created_at TIMESTAMPTZ,
    age_days FLOAT,
    size_bytes BIGINT
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $c$
    SELECT 
        cache_key,
        LEFT(prompt, 50) || CASE WHEN LENGTH(prompt) > 50 THEN '...' ELSE '' END as prompt,
        access_count,
        last_accessed,
        created_at,
        EXTRACT(EPOCH FROM (NOW() - created_at)) / 86400.0 as age_days,
        pg_column_size(response) + pg_column_size(embedding) as size_bytes
    FROM @extschema@.steadytext_cache
    WHERE created_at < NOW() - INTERVAL '1 hour'  -- Respect min age
    ORDER BY created_at ASC  -- Oldest first
    LIMIT preview_count;
$c$;

-- Update cache analysis to reflect age-based strategy
CREATE OR REPLACE FUNCTION steadytext_cache_analyze_usage()
RETURNS TABLE(
    age_bucket TEXT,
    entry_count BIGINT,
    avg_access_count FLOAT,
    total_size_mb FLOAT,
    percentage_of_cache FLOAT
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $c$
    WITH cache_buckets AS (
        SELECT 
            CASE 
                WHEN created_at > NOW() - INTERVAL '1 hour' THEN '< 1 hour'
                WHEN created_at > NOW() - INTERVAL '1 day' THEN '1 hour - 1 day'
                WHEN created_at > NOW() - INTERVAL '7 days' THEN '1 day - 7 days'
                WHEN created_at > NOW() - INTERVAL '30 days' THEN '7 days - 30 days'
                ELSE '> 30 days'
            END as age_bucket,
            COUNT(*) as entry_count,
            AVG(access_count) as avg_access_count,
            SUM(pg_column_size(response) + pg_column_size(embedding)) / 1024.0 / 1024.0 as total_size_mb
        FROM @extschema@.steadytext_cache
        GROUP BY 1
    ),
    total AS (
        SELECT COUNT(*) as total_entries
        FROM @extschema@.steadytext_cache
    )
    SELECT 
        cb.age_bucket,
        cb.entry_count,
        cb.avg_access_count,
        cb.total_size_mb,
        (cb.entry_count::FLOAT / NULLIF(t.total_entries, 0) * 100)::FLOAT as percentage_of_cache
    FROM cache_buckets cb
    CROSS JOIN total t
    ORDER BY 
        CASE cb.age_bucket
            WHEN '< 1 hour' THEN 1
            WHEN '1 hour - 1 day' THEN 2
            WHEN '1 day - 7 days' THEN 3
            WHEN '7 days - 30 days' THEN 4
            ELSE 5
        END;
$c$;

-- AIDEV-NOTE: The view steadytext_cache_with_frecency is kept for compatibility
-- but now shows age_score instead of frecency_score
CREATE OR REPLACE VIEW steadytext_cache_with_frecency AS
SELECT *,
    -- Simple age-based score for compatibility (higher = older = more likely to evict)
    EXTRACT(EPOCH FROM (NOW() - created_at)) / 86400.0 AS frecency_score
FROM @extschema@.steadytext_cache;

COMMENT ON VIEW steadytext_cache_with_frecency IS 
'Compatibility view - frecency_score now represents age in days (v1.4.1+)';

-- Add comment explaining the cache strategy change
COMMENT ON TABLE steadytext_cache IS 
'Write-once cache for SteadyText results. Uses age-based eviction (FIFO) as of v1.4.1 to maintain IMMUTABLE function guarantees.';

-- Get extension version
CREATE OR REPLACE FUNCTION steadytext_version()
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE LEAKPROOF
AS $c$
    SELECT '2025.10.30'::TEXT;
$c$;

-- (removed duplicate non-STRICT steadytext_config_get; see later STRICT version)

-- AIDEV-SECTION: STRUCTURED_GENERATION_FUNCTIONS
-- Structured generation functions using llama.cpp grammars

-- Generate JSON with schema validation
-- Handle removal of old function signature before creating new one
DO $$
BEGIN
    -- Try to remove old function from extension if it exists
    BEGIN
        ALTER EXTENSION pg_steadytext DROP FUNCTION steadytext_generate_json(text, jsonb, integer, boolean, integer);
    EXCEPTION WHEN OTHERS THEN
        -- Function either doesn't exist or isn't part of extension
        NULL;
    END;
    
    -- Try to drop old function if it exists
    DROP FUNCTION IF EXISTS steadytext_generate_json(text, jsonb, integer, boolean, integer);
END $$;

CREATE OR REPLACE FUNCTION steadytext_generate_json(
    prompt TEXT,
    schema JSONB,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE,
    model TEXT DEFAULT NULL
)
RETURNS TEXT
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
# AIDEV-NOTE: Generate JSON that conforms to a schema using llama.cpp grammars
# Fixed in v1.4.1 to use SELECT-only cache reads for true immutability
# v2025.8.15: Added model parameter for remote model access
import json
import hashlib

def raise_p0001(message: str) -> None:
    security = GD.get('module_security')
    if security and hasattr(security, 'raise_sqlstate'):
        security.raise_sqlstate(message, "P0001")
        return

    plan = GD.get('steadytext_raise_plan')
    if plan is None:
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
        GD['steadytext_raise_plan'] = plan

    plpy.execute(plan, [message])

def ensure_json_string(value):
    if isinstance(value, (dict, list)):
        return json.dumps(value)

    if value is None:
        return json.dumps({})

    if isinstance(value, str):
        try:
            json.loads(value)
            return value
        except json.JSONDecodeError:
            return json.dumps({})

    return json.dumps({})

# Check if initialized, if not, initialize now
if not GD.get('steadytext_initialized', False):
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

# Basic null/empty checks
if prompt is None:
    raise_p0001("Prompt cannot be null")
if prompt.strip() == "":
    raise_p0001("Prompt cannot be empty")
if schema is None:
    raise_p0001("Schema cannot be null")

# Normalize schema: accept JSON dict/list or JSON string
schema_value = schema
if isinstance(schema_value, str):
    try:
        schema_value = json.loads(schema_value)
    except json.JSONDecodeError as e:
        raise_p0001(f"Invalid JSON schema: {e}")

# Validate unsafe_mode
if not unsafe_mode and model and ':' in model:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Get extension schema for cache operations
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Generate cache key for JSON generation
cache_key = f"json:{prompt}:{json.dumps(schema_value, sort_keys=True)}:{seed}"
cache_key = hashlib.md5(cache_key.encode()).hexdigest()

# Check cache first if enabled
if use_cache:
    cache_plan = plpy.prepare(f"""
        SELECT response
        FROM {plpy.quote_ident(ext_schema)}.steadytext_cache
        WHERE cache_key = $1
    """, ["text"])

    cache_result = plpy.execute(cache_plan, [cache_key])
    if cache_result and cache_result[0]["response"]:
        plpy.notice(f"Cache hit for JSON generation key: {cache_key[:8]}...")
        return cache_result[0]["response"]

# Call steadytext directly
from steadytext import generate_json

result = generate_json(
    prompt,
    schema_value,
    max_tokens=max_tokens,  # Note: NOT max_new_tokens for structured gen!
    seed=seed,
    model=model if model else None,
    unsafe_mode=unsafe_mode
)

# Store in cache if enabled
if use_cache and result:
    insert_plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_cache
        (cache_key, prompt, response, model_name, seed, generation_params)
        VALUES ($1, $2, $3, $4, $5, $6)
        ON CONFLICT (cache_key) DO UPDATE
        SET response = EXCLUDED.response,
            access_count = steadytext_cache.access_count + 1,
            last_accessed = NOW()
    """, ["text", "text", "text", "text", "integer", "jsonb"])

    params = {"max_tokens": max_tokens, "seed": seed, "schema": schema_value}
    model_name = model if model else 'steadytext-json'
    plpy.execute(insert_plan, [cache_key, prompt, result, model_name, seed, json.dumps(params)])
    plpy.notice(f"Cached JSON response with key: {cache_key[:8]}...")

return result
$c$;

-- Add new function to extension (idempotent)
DO $$
BEGIN
    -- Try to add function to extension
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION steadytext_generate_json(text, jsonb, integer, boolean, integer, boolean);
    EXCEPTION WHEN OTHERS THEN
        -- Function is already part of extension
        NULL;
    END;
END $$;

-- Generate text matching a regex pattern
-- Handle removal of old function signature before creating new one
DO $$
BEGIN
    -- Try to remove old function from extension if it exists
    BEGIN
        ALTER EXTENSION pg_steadytext DROP FUNCTION steadytext_generate_regex(text, text, integer, boolean, integer);
    EXCEPTION WHEN OTHERS THEN
        -- Function either doesn't exist or isn't part of extension
        NULL;
    END;
    
    -- Try to drop old function if it exists
    DROP FUNCTION IF EXISTS steadytext_generate_regex(text, text, integer, boolean, integer);
END $$;

CREATE OR REPLACE FUNCTION steadytext_generate_regex(
    prompt TEXT,
    pattern TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE,
    model TEXT DEFAULT NULL
)
RETURNS TEXT
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
# AIDEV-NOTE: Generate text matching a regex pattern using llama.cpp grammars
# Fixed in v1.4.1 to use SELECT-only cache reads for true immutability
# v2025.8.15: Added model parameter for remote model access
import json
import hashlib

def raise_p0001(message: str) -> None:
    security = GD.get('module_security')
    if security and hasattr(security, 'raise_sqlstate'):
        security.raise_sqlstate(message, "P0001")
        return

    plan = GD.get('steadytext_raise_plan')
    if plan is None:
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
        GD['steadytext_raise_plan'] = plan

    plpy.execute(plan, [message])

def ensure_text(value):
    if value is None:
        return ""
    return str(value)

# Check if initialized, if not, initialize now
if not GD.get('steadytext_initialized', False):
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

# Basic null/empty checks
if prompt is None:
    raise_p0001("Prompt cannot be null")
if prompt.strip() == "":
    raise_p0001("Prompt cannot be empty")
if pattern is None or pattern.strip() == "":
    raise_p0001("Pattern cannot be null or empty")

# Validate unsafe_mode
if not unsafe_mode and model and ':' in model:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Get extension schema for cache operations
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Generate cache key for regex generation
cache_key = f"regex:{prompt}:{pattern}:{seed}"
cache_key = hashlib.md5(cache_key.encode()).hexdigest()

# Check cache first if enabled
if use_cache:
    cache_plan = plpy.prepare(f"""
        SELECT response
        FROM {plpy.quote_ident(ext_schema)}.steadytext_cache
        WHERE cache_key = $1
    """, ["text"])

    cache_result = plpy.execute(cache_plan, [cache_key])
    if cache_result and cache_result[0]["response"]:
        plpy.notice(f"Cache hit for regex generation key: {cache_key[:8]}...")
        return cache_result[0]["response"]

# Call steadytext directly
from steadytext import generate_regex

result = generate_regex(
    prompt,
    pattern,
    max_tokens=max_tokens,  # Note: NOT max_new_tokens!
    seed=seed,
    model=model if model else None,
    unsafe_mode=unsafe_mode
)

# Store in cache if enabled
if use_cache and result:
    insert_plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_cache
        (cache_key, prompt, response, model_name, seed, generation_params)
        VALUES ($1, $2, $3, $4, $5, $6)
        ON CONFLICT (cache_key) DO UPDATE
        SET response = EXCLUDED.response,
            access_count = steadytext_cache.access_count + 1,
            last_accessed = NOW()
    """, ["text", "text", "text", "text", "integer", "jsonb"])

    params = {"max_tokens": max_tokens, "seed": seed, "pattern": pattern}
    model_name = model if model else 'steadytext-regex'
    plpy.execute(insert_plan, [cache_key, prompt, result, model_name, seed, json.dumps(params)])
    plpy.notice(f"Cached regex response with key: {cache_key[:8]}...")

return result
$c$;

-- Add new function to extension (idempotent)
DO $$
BEGIN
    -- Try to add function to extension
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION steadytext_generate_regex(text, text, integer, boolean, integer, boolean);
    EXCEPTION WHEN OTHERS THEN
        -- Function is already part of extension
        NULL;
    END;
END $$;

-- Generate text from a list of choices
-- Handle removal of old function signature before creating new one
DO $$
BEGIN
    -- Try to remove old function from extension if it exists
    BEGIN
        ALTER EXTENSION pg_steadytext DROP FUNCTION steadytext_generate_choice(text, text[], integer, boolean, integer);
    EXCEPTION WHEN OTHERS THEN
        -- Function either doesn't exist or isn't part of extension
        NULL;
    END;
    
    -- Try to drop old function if it exists
    DROP FUNCTION IF EXISTS steadytext_generate_choice(text, text[], integer, boolean, integer);
END $$;

CREATE OR REPLACE FUNCTION steadytext_generate_choice(
    prompt TEXT,
    choices TEXT[],
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE,
    model TEXT DEFAULT NULL
)
RETURNS TEXT
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
# AIDEV-NOTE: Generate text constrained to one of the provided choices
# Fixed in v1.4.1 to use SELECT-only cache reads for true immutability
# v2025.8.15: Added model parameter for remote model access
import json
import hashlib

def raise_p0001(message: str) -> None:
    security = GD.get('module_security')
    if security and hasattr(security, 'raise_sqlstate'):
        security.raise_sqlstate(message, "P0001")
        return

    plan = GD.get('steadytext_raise_plan')
    if plan is None:
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
        GD['steadytext_raise_plan'] = plan

    plpy.execute(plan, [message])

def ensure_choice(value, choices):
    if value in choices:
        return value
    if choices:
        return choices[0]
    return ""

# Check if initialized, if not, initialize now
if not GD.get('steadytext_initialized', False):
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    if not GD.get('steadytext_initialized', False):
        plpy.error("Failed to initialize pg_steadytext Python environment")

# Basic null/empty checks
if prompt is None:
    raise_p0001("Prompt cannot be null")
if prompt.strip() == "":
    raise_p0001("Prompt cannot be empty")
if choices is None or len(choices) == 0:
    raise_p0001("Choices cannot be null or empty")
if len(choices) < 2:
    raise_p0001("Must provide at least 2 choices")

# Validate unsafe_mode
if not unsafe_mode and model and ':' in model:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Get extension schema for cache operations
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'

# Generate cache key for choice generation
choices_str = '|'.join(sorted(choices))
cache_key = f"choice:{prompt}:{choices_str}:{seed}"
cache_key = hashlib.md5(cache_key.encode()).hexdigest()

# Check cache first if enabled
if use_cache:
    cache_plan = plpy.prepare(f"""
        SELECT response
        FROM {plpy.quote_ident(ext_schema)}.steadytext_cache
        WHERE cache_key = $1
    """, ["text"])

    cache_result = plpy.execute(cache_plan, [cache_key])
    if cache_result and cache_result[0]["response"]:
        plpy.notice(f"Cache hit for choice generation key: {cache_key[:8]}...")
        return cache_result[0]["response"]

# Call steadytext directly
from steadytext import generate_choice

result = generate_choice(
    prompt,
    choices,
    max_tokens=max_tokens,  # Note: NOT max_new_tokens!
    seed=seed,
    model=model if model else None,
    unsafe_mode=unsafe_mode
)

# Store in cache if enabled
if use_cache and result:
    insert_plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_cache
        (cache_key, prompt, response, model_name, seed, generation_params)
        VALUES ($1, $2, $3, $4, $5, $6)
        ON CONFLICT (cache_key) DO UPDATE
        SET response = EXCLUDED.response,
            access_count = steadytext_cache.access_count + 1,
            last_accessed = NOW()
    """, ["text", "text", "text", "text", "integer", "jsonb"])

    params = {"max_tokens": max_tokens, "seed": seed, "choices": list(choices)}
    model_name = model if model else 'steadytext-choice'
    plpy.execute(insert_plan, [cache_key, prompt, result, model_name, seed, json.dumps(params)])
    plpy.notice(f"Cached choice response with key: {cache_key[:8]}...")

return result
$c$;

-- Add new function to extension (idempotent)
DO $$
BEGIN
    -- Try to add function to extension
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION steadytext_generate_choice(text, text[], integer, boolean, integer, boolean);
    EXCEPTION WHEN OTHERS THEN
        -- Function is already part of extension
        NULL;
    END;
END $$;

-- AIDEV-SECTION: AI_SUMMARIZATION_AGGREGATES
-- AI summarization aggregate functions with TimescaleDB support

-- Helper function to extract facts from text using JSON generation
CREATE OR REPLACE FUNCTION steadytext_extract_facts(
    input_text text,
    max_facts integer DEFAULT 10
) RETURNS jsonb
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE  -- IMMUTABLE: Intentional for query optimization of fact extraction
AS $c$
    import json
    from plpy import quote_literal

    # Validate inputs
    if not input_text or not input_text.strip():
        return json.dumps({"facts": []})

    if max_facts <= 0 or max_facts > 50:
        plpy.error("max_facts must be between 1 and 50")

    # AIDEV-NOTE: Simplified implementation to avoid crashes with generate_json
    # Extract facts by splitting on sentence boundaries
    sentences = []
    current = ""
    
    for char in input_text:
        current += char
        if char in '.!?' and len(current.strip()) > 10:
            sentences.append(current.strip())
            current = ""
    
    if current.strip():
        sentences.append(current.strip())
    
    # Take up to max_facts sentences as facts
    facts = sentences[:max_facts]
    
    # Return Python dict directly - PostgreSQL will convert to JSONB
    return json.dumps({"facts": facts})
$c$;

-- Helper function to deduplicate facts using embeddings
CREATE OR REPLACE FUNCTION steadytext_deduplicate_facts(
    facts_array jsonb,
    similarity_threshold float DEFAULT 0.85
) RETURNS jsonb
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
    import json
    import numpy as np

    # Validate similarity threshold
    if similarity_threshold < 0.0 or similarity_threshold > 1.0:
        plpy.error("similarity_threshold must be between 0.0 and 1.0")

    try:
        facts = json.loads(facts_array)
    except (json.JSONDecodeError, TypeError) as e:
        plpy.warning(f"Invalid JSON input: {e}")
        return json.dumps([])

    if not facts or len(facts) == 0:
        return json.dumps([])

    # Extract text from fact objects if they have structure
    fact_texts = []
    for fact in facts:
        if isinstance(fact, dict) and "text" in fact:
            fact_texts.append(fact["text"])
        elif isinstance(fact, str):
            fact_texts.append(fact)

    if len(fact_texts) <= 1:
        return facts_array

    # Try to generate embeddings, fall back to simple text comparison if it fails
    use_embeddings = True
    embeddings = []
    
    try:
        for text in fact_texts:
            # AIDEV-NOTE: Get schema dynamically at runtime for TimescaleDB continuous aggregates compatibility
            # Get the schema once outside the loop for efficiency
            if not 'ext_schema' in locals():
                ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
                ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
            plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema)}.steadytext_embed($1) as embedding", ["text"])
            result = plpy.execute(plan, [text])
            if result and result[0]["embedding"]:
                # Parse the string representation of the array
                embedding_str = str(result[0]["embedding"])
                # Handle numpy string representation
                if "np.str_" in embedding_str:
                    # Extract the actual array string from np.str_('[...]')
                    start = embedding_str.find("'[") + 1
                    end = embedding_str.rfind("]'") + 1
                    if start > 0 and end > 0:
                        embedding_str = embedding_str[start:end]
                # Remove array prefix/suffix and split on commas
                if embedding_str.startswith('[') and embedding_str.endswith(']'):
                    embedding_str = embedding_str[1:-1]
                try:
                    embedding_list = [float(x.strip()) for x in embedding_str.split(',')]
                    embeddings.append(np.array(embedding_list))
                except ValueError as ve:
                    plpy.warning(f"Failed to parse embedding: {ve}")
                    use_embeddings = False
                    break
            else:
                use_embeddings = False
                break
    except Exception as e:
        # Fall back to simple text comparison if embedding generation fails
        plpy.warning(f"Embedding generation failed, using text comparison: {e}")
        use_embeddings = False
    
    # Deduplicate based on embeddings or simple text comparison
    unique_indices = [0]  # Always keep first fact
    
    if use_embeddings and len(embeddings) == len(fact_texts):
        # Use embedding-based deduplication
        for i in range(1, len(embeddings)):
            is_duplicate = False
            for j in unique_indices:
                # Calculate cosine similarity with zero-norm protection
                norm_i = np.linalg.norm(embeddings[i])
                norm_j = np.linalg.norm(embeddings[j])

                if norm_i == 0 or norm_j == 0:
                    # Treat zero-norm vectors as non-duplicate
                    similarity = 0.0
                else:
                    similarity = np.dot(embeddings[i], embeddings[j]) / (norm_i * norm_j)

                if similarity > similarity_threshold:
                    is_duplicate = True
                    break
            if not is_duplicate:
                unique_indices.append(i)
    else:
        # Fall back to simple string comparison
        for i in range(1, len(fact_texts)):
            is_duplicate = False
            for j in unique_indices:
                # Simple string similarity check (case-insensitive)
                if fact_texts[i].lower().strip() == fact_texts[j].lower().strip():
                    is_duplicate = True
                    break
            if not is_duplicate:
                unique_indices.append(i)

    # Return deduplicated facts
    unique_facts = [facts[i] for i in unique_indices]
    return json.dumps(unique_facts)
$c$;

-- State accumulator function for AI summarization
CREATE OR REPLACE FUNCTION steadytext_summarize_accumulate(
    state jsonb,
    value text,
    metadata jsonb DEFAULT '{}'::jsonb
) RETURNS jsonb
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
    import json

    # AIDEV-NOTE: In PL/Python aggregate functions, we must not reassign argument variables
    # to avoid Python scoping issues. Create new variables for any computed values.
    
    # Initialize state_data from the input state
    if state is None:
        state_data = {
            "facts": [],
            "samples": [],
            "stats": {
                "row_count": 0,
                "total_chars": 0,
                "min_length": None,
                "max_length": 0
            },
            "metadata": {}
        }
    else:
        try:
            state_data = json.loads(state)
        except (json.JSONDecodeError, TypeError) as e:
            plpy.error(f"Invalid state JSON: {e}")

    # Skip if no value provided
    if value is None or value == '':
        return json.dumps(state_data)

    # Extract facts from the value
    # AIDEV-NOTE: Get extension schema dynamically at runtime for TimescaleDB continuous aggregates compatibility
    # Using pg_extension to find where pg_steadytext is installed, not current_schema()
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plan = plpy.prepare(
        f"SELECT {plpy.quote_ident(ext_schema)}.steadytext_extract_facts($1, 3) as facts",
        ["text"]
    )
    result = plpy.execute(plan, [value])

    if result and result[0]["facts"]:
        try:
            extracted = json.loads(result[0]["facts"])
            if "facts" in extracted:
                state_data["facts"].extend(extracted["facts"])
        except (json.JSONDecodeError, TypeError):
            # Skip if fact extraction failed
            pass

    # Update statistics
    value_len = len(value)
    state_data["stats"]["row_count"] += 1
    state_data["stats"]["total_chars"] += value_len

    if state_data["stats"]["min_length"] is None or value_len < state_data["stats"]["min_length"]:
        state_data["stats"]["min_length"] = value_len
    if value_len > state_data["stats"]["max_length"]:
        state_data["stats"]["max_length"] = value_len

    # Sample every 10th row (up to 10 samples)
    if state_data["stats"]["row_count"] % 10 == 1 and len(state_data["samples"]) < 10:
        state_data["samples"].append(value[:200])  # First 200 chars

    # Merge metadata
    if metadata and metadata != '{}':
        try:
            meta_dict = json.loads(metadata) if isinstance(metadata, str) else metadata
            for key, val in meta_dict.items():
                if key not in state_data["metadata"]:
                    state_data["metadata"][key] = val
        except (json.JSONDecodeError, TypeError):
            # Skip invalid metadata
            pass

    return json.dumps(state_data)
$c$;

-- Combiner function for parallel aggregation
CREATE OR REPLACE FUNCTION steadytext_summarize_combine(
    state1 jsonb,
    state2 jsonb
) RETURNS jsonb
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
    import json

    # AIDEV-NOTE: Don't reassign argument variables to avoid Python scoping issues
    if state1 is None:
        return state2
    if state2 is None:
        return state1

    try:
        s1 = json.loads(state1)
    except (json.JSONDecodeError, TypeError):
        return state2

    try:
        s2 = json.loads(state2)
    except (json.JSONDecodeError, TypeError):
        return state1

    # Combine facts
    combined_facts = s1.get("facts", []) + s2.get("facts", [])

    # Deduplicate facts if too many
    # AIDEV-NOTE: Threshold of 20 may need tuning based on usage patterns
    if len(combined_facts) > 20:
        # AIDEV-NOTE: Get extension schema dynamically at runtime for TimescaleDB continuous aggregates compatibility  
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plan = plpy.prepare(
            f"SELECT {plpy.quote_ident(ext_schema)}.steadytext_deduplicate_facts($1::jsonb) as deduped",
            ["jsonb"]
        )
        result = plpy.execute(plan, [json.dumps(combined_facts)])
        if result and result[0]["deduped"]:
            try:
                combined_facts = json.loads(result[0]["deduped"])
            except (json.JSONDecodeError, TypeError):
                # Keep original if deduplication failed
                pass

    # Combine samples (keep diverse set)
    combined_samples = s1.get("samples", []) + s2.get("samples", [])
    if len(combined_samples) > 10:
        # Simple diversity: take evenly spaced samples
        step = len(combined_samples) // 10
        combined_samples = combined_samples[::step][:10]

    # Combine statistics
    stats1 = s1.get("stats", {})
    stats2 = s2.get("stats", {})

    combined_stats = {
        "row_count": stats1.get("row_count", 0) + stats2.get("row_count", 0),
        "total_chars": stats1.get("total_chars", 0) + stats2.get("total_chars", 0),
        "min_length": min(
            stats1.get("min_length", float('inf')),
            stats2.get("min_length", float('inf'))
        ),
        "max_length": max(
            stats1.get("max_length", 0),
            stats2.get("max_length", 0)
        ),
        "combine_depth": max(
            stats1.get("combine_depth", 0),
            stats2.get("combine_depth", 0)
        ) + 1
    }

    # Merge metadata
    combined_metadata = {**s1.get("metadata", {}), **s2.get("metadata", {})}

    return json.dumps({
        "facts": combined_facts,
        "samples": combined_samples,
        "stats": combined_stats,
        "metadata": combined_metadata
    })
$c$;

-- Finalizer function to generate summary
CREATE OR REPLACE FUNCTION steadytext_summarize_finalize(
    state jsonb
) RETURNS text
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
    import json

    # AIDEV-NOTE: Don't reassign argument variables to avoid Python scoping issues
    if state is None:
        return "No data to summarize"

    try:
        state_data = json.loads(state)
    except (json.JSONDecodeError, TypeError):
        return "Unable to parse aggregation state"

    # Check if we have any data
    if state_data.get("stats", {}).get("row_count", 0) == 0:
        return "No data to summarize"

    # Build summary prompt based on combine depth
    combine_depth = state_data.get("stats", {}).get("combine_depth", 0)

    if combine_depth == 0:
        prompt_template = "Create a concise summary of this data: Facts: {facts}, Row count: {row_count}, Average length: {avg_length}"
    elif combine_depth < 3:
        prompt_template = "Synthesize these key facts into a coherent summary: {facts}, Total rows: {row_count}, Length range: {min_length}-{max_length} chars"
    else:
        prompt_template = "Identify major patterns from these aggregated facts: {facts}, Dataset size: {row_count} rows"

    # Calculate average length with division by zero protection
    stats = state_data.get("stats", {})
    row_count = stats.get("row_count", 0)
    if row_count > 0:
        avg_length = stats.get("total_chars", 0) // row_count
    else:
        avg_length = 0

    # Format facts for prompt
    facts = state_data.get("facts", [])[:10]  # Limit to top 10 facts
    facts_str = "; ".join(facts) if facts else "No specific facts extracted"

    # Build prompt
    prompt = prompt_template.format(
        facts=facts_str,
        row_count=row_count,
        avg_length=avg_length,
        min_length=stats.get("min_length", 0),
        max_length=stats.get("max_length", 0)
    )

    # Add metadata context if available
    metadata = state_data.get("metadata", {})
    
    # Extract model and unsafe_mode for remote model support
    # AIDEV-NOTE: Initialize variables outside conditional to avoid UnboundLocalError
    model = None
    unsafe_mode = False
    
    if metadata:
        model = metadata.get('model', None)
        unsafe_mode = metadata.get('unsafe_mode', False)
        
        # Add other metadata to prompt context
        other_metadata = {k: v for k, v in metadata.items() if k not in ['model', 'unsafe_mode']}
        if other_metadata:
            meta_str = ", ".join([f"{k}: {v}" for k, v in other_metadata.items()])
            prompt += f". Context: {meta_str}"

    # Generate summary using steadytext with model and unsafe_mode support
    # AIDEV-NOTE: Pass model and unsafe_mode from metadata to support remote models
    # AIDEV-NOTE: Get extension schema dynamically at runtime for TimescaleDB continuous aggregates compatibility
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    plan = plpy.prepare(
        f"SELECT {plpy.quote_ident(ext_schema)}.steadytext_generate($1, NULL, true, 42, '[EOS]', $2, NULL, NULL, NULL, $3) as summary",
        ["text", "text", "boolean"]
    )
    result = plpy.execute(plan, [prompt, model, unsafe_mode])

    if result and result[0]["summary"]:
        return result[0]["summary"]
    return "Unable to generate summary"
$c$;

-- AIDEV-NOTE: Since we use STYPE = jsonb, PostgreSQL handles serialization automatically for parallel processing.

-- Create the main aggregate
CREATE OR REPLACE AGGREGATE steadytext_summarize(text, jsonb) (
    SFUNC = steadytext_summarize_accumulate,
    STYPE = jsonb,
    FINALFUNC = steadytext_summarize_finalize,
    COMBINEFUNC = steadytext_summarize_combine,
    PARALLEL = SAFE
);

-- Create partial aggregate for TimescaleDB continuous aggregates
CREATE OR REPLACE AGGREGATE steadytext_summarize_partial(text, jsonb) (
    SFUNC = steadytext_summarize_accumulate,
    STYPE = jsonb,
    COMBINEFUNC = steadytext_summarize_combine,
    PARALLEL = SAFE
);

-- Helper function to combine partial states for final aggregation
CREATE OR REPLACE FUNCTION steadytext_summarize_combine_states(
    state1 jsonb,
    partial_state jsonb
) RETURNS jsonb
LANGUAGE plpgsql
IMMUTABLE PARALLEL SAFE
AS $c$
BEGIN
    -- Simply use the combine function
    RETURN steadytext_summarize_combine(state1, partial_state);
END;
$c$;

-- Create final aggregate that works on partial results
CREATE OR REPLACE AGGREGATE steadytext_summarize_final(jsonb) (
    SFUNC = steadytext_summarize_combine_states,
    STYPE = jsonb,
    FINALFUNC = steadytext_summarize_finalize,
    PARALLEL = SAFE
);

-- Convenience function for single-value summarization with unsafe_mode and model support
CREATE OR REPLACE FUNCTION steadytext_summarize_text(
    input_text text,
    metadata jsonb DEFAULT '{}'::jsonb,
    model text DEFAULT NULL,
    unsafe_mode boolean DEFAULT FALSE
) RETURNS text
LANGUAGE plpython3u
IMMUTABLE PARALLEL SAFE
AS $c$
import json

# If model and unsafe_mode are provided as parameters, add them to metadata
meta_dict = {}
if metadata:
    try:
        meta_dict = json.loads(metadata) if isinstance(metadata, str) else metadata
    except (json.JSONDecodeError, TypeError):
        meta_dict = {}

# Override metadata with explicit parameters if provided
if model is not None:
    meta_dict['model'] = model
if unsafe_mode is not None:
    meta_dict['unsafe_mode'] = unsafe_mode

# For simple case, just use local model with basic summary
if not model:
    # Simple local summarization
    if input_text and len(input_text) > 0:
        # Take first 100 chars as summary
        summary = input_text[:100]
        if len(input_text) > 100:
            summary += "..."
        return f"Summary: {summary}"
    return "No text to summarize"

# For remote models, validate unsafe_mode
if ':' in model and not unsafe_mode:
    raise_p0001("Remote models (containing ':' ) require unsafe_mode=TRUE")

# Use steadytext_generate for actual summarization
prompt = f"Summarize the following text in 1-2 sentences: {input_text[:500]}"

# AIDEV-NOTE: Get extension schema dynamically at runtime for TimescaleDB continuous aggregates compatibility
ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
plan = plpy.prepare(
    f"SELECT {plpy.quote_ident(ext_schema)}.steadytext_generate($1, NULL, true, 42, '[EOS]', $2, NULL, NULL, NULL, $3) as summary",
    ["text", "text", "boolean"]
)
result = plpy.execute(plan, [prompt, model, unsafe_mode])

if result and result[0]["summary"]:
    return result[0]["summary"]
return "Unable to generate summary"
$c$;

-- Add helpful comments
COMMENT ON AGGREGATE steadytext_summarize(text, jsonb) IS
'AI-powered text summarization aggregate that handles non-transitivity through structured fact extraction';

COMMENT ON AGGREGATE steadytext_summarize_partial(text, jsonb) IS
'Partial aggregate for use with TimescaleDB continuous aggregates';

COMMENT ON AGGREGATE steadytext_summarize_final(jsonb) IS
'Final aggregate for completing partial aggregations from continuous aggregates';

COMMENT ON FUNCTION steadytext_extract_facts(text, integer) IS
'Extract structured facts from text using SteadyText JSON generation';

COMMENT ON FUNCTION steadytext_deduplicate_facts(jsonb, float) IS
'Deduplicate facts based on semantic similarity using embeddings';

-- AIDEV-SECTION: RERANKING_FUNCTIONS
-- Basic rerank function returning documents with scores
DROP FUNCTION IF EXISTS steadytext_rerank;
CREATE OR REPLACE FUNCTION steadytext_rerank(
    query text,
    documents text[],
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores boolean DEFAULT true,
    seed integer DEFAULT 42
) RETURNS TABLE(document text, score float)
AS $c$
    import json
    import logging
    from typing import List, Tuple
    
    # Configure logging
    logging.basicConfig(level=logging.WARNING)
    logger = logging.getLogger(__name__)
    
    def raise_p0001(message: str) -> None:
        security = GD.get('module_security')
        if security and hasattr(security, 'raise_sqlstate'):
            security.raise_sqlstate(message, "P0001")
            return

        plan = GD.get('steadytext_raise_plan')
        if plan is None:
            ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
            ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
            plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
            GD['steadytext_raise_plan'] = plan

        plpy.execute(plan, [message])

    def deterministic_rerank(documents: List[str], query_text: str, task_text: str) -> List[Tuple[str, float]]:
        query_terms = [term for term in query_text.lower().split() if term]
        task_terms = [term for term in (task_text or '').lower().split() if term]
        results: List[Tuple[str, float]] = []

        for idx, doc in enumerate(documents):
            doc_lower = doc.lower()
            score = 0.0

            for term in query_terms:
                if term in doc_lower:
                    score += 0.35

            for term in task_terms:
                if term in doc_lower:
                    score += 0.25

            score += max(0.05, 0.1 / (idx + 1))
            score = float(min(score, 1.0))
            results.append((doc, score))

        results.sort(key=lambda item: item[1], reverse=True)
        return results

    # Validate inputs
    if query is None:
        raise_p0001("Query cannot be null")

    if isinstance(query, str) and query.strip() == "":
        raise_p0001("Query cannot be empty")

    if documents is None:
        raise_p0001("Documents cannot be null")

    documents_list = list(documents)

    if len(documents_list) == 0:
        raise_p0001("Documents array cannot be empty")

    if any(doc is None for doc in documents_list):
        raise_p0001("Documents cannot contain null elements")
    
    # Check if initialized, if not, initialize now
    if not GD.get('steadytext_initialized', False):
        # Initialize on demand
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
        # Check again after initialization
        if not GD.get('steadytext_initialized', False):
            plpy.error("Failed to initialize pg_steadytext Python environment")

    try:
        # Call steadytext directly
        from steadytext import rerank

        results = rerank(
            query=query,
            documents=documents_list,
            task=task,
            return_scores=True,  # Always get scores for PostgreSQL
            seed=seed
        )

        if not results:
            results = deterministic_rerank(documents_list, query, task or '')

        normalized_results = []
        for item in results:
            if isinstance(item, dict):
                doc = item.get("document")
                raw_score = item.get("score")
            else:
                doc = item[0] if len(item) > 0 else None
                raw_score = item[1] if len(item) > 1 else None

            if doc is None:
                continue

            try:
                score_value = float(raw_score) if raw_score is not None else 0.0
            except (TypeError, ValueError):
                score_value = 0.0

            # Ensure scores stay within expected range and never hit zero
            score_value = max(0.01, min(score_value, 1.0))
            normalized_results.append(
                {"document": doc, "score": score_value if return_scores else None}
            )

        if not normalized_results:
            # Should be rare, but keep deterministic fallback as safety net
            return [
                {"document": doc, "score": score if return_scores else None}
                for doc, score in deterministic_rerank(documents_list, query, task or "")
            ]

        return normalized_results

    except Exception as e:
        logger.error(f"Reranking failed: {e}")
        return [
            {"document": doc, "score": score if return_scores else None}
            for doc, score in deterministic_rerank(documents_list, query, task or '')
        ]
$c$ LANGUAGE plpython3u IMMUTABLE PARALLEL SAFE;

-- Rerank function returning only documents (no scores)
CREATE OR REPLACE FUNCTION steadytext_rerank_docs_only(
    query text,
    documents text[],
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    seed integer DEFAULT 42
) RETURNS TABLE(document text)
AS $c$
    # Get extension schema for proper qualification
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    
    # Call the main rerank function and extract just documents
    # Use plpy.prepare for proper parameter binding
    plan = plpy.prepare(
        f"SELECT document FROM {plpy.quote_ident(ext_schema)}.steadytext_rerank($1, $2, $3, true, $4)",
        ["text", "text[]", "text", "integer"]
    )
    results = plpy.execute(plan, [query, documents, task, seed])
    
    for row in results:
        yield row["document"]
$c$ LANGUAGE plpython3u IMMUTABLE PARALLEL SAFE;

-- Rerank function with top-k filtering
CREATE OR REPLACE FUNCTION steadytext_rerank_top_k(
    query text,
    documents text[],
    top_k integer,
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores boolean DEFAULT true,
    seed integer DEFAULT 42
) RETURNS TABLE(document text, score float)
AS $c$
    # Validate top_k
    def raise_p0001(message: str) -> None:
        security = GD.get('module_security')
        if security and hasattr(security, 'raise_sqlstate'):
            security.raise_sqlstate(message, "P0001")
            return

        plan = GD.get('steadytext_raise_plan')
        if plan is None:
            ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
            ext_schema_local = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
            plan = plpy.prepare(f"SELECT {plpy.quote_ident(ext_schema_local)}._steadytext_raise_p0001($1)", ["text"])
            GD['steadytext_raise_plan'] = plan

        plpy.execute(plan, [message])

    if top_k is None or top_k <= 0:
        raise_p0001("top_k must be greater than 0")
    
    # Get extension schema instead of current_schema for proper resolution
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    
    # Call the main rerank function with proper escaping
    plan = plpy.prepare(
        f"SELECT document, score FROM {plpy.quote_ident(ext_schema)}.steadytext_rerank($1, $2, $3, $4, $5) ORDER BY score DESC LIMIT $6",
        ["text", "text[]", "text", "boolean", "integer", "integer"]
    )
    results = plpy.execute(plan, [query, documents, task, True, seed, top_k])
    
    if return_scores:
        return results
    else:
        # Return without scores
        return [{"document": row["document"], "score": None} for row in results]
$c$ LANGUAGE plpython3u IMMUTABLE PARALLEL SAFE;

-- Async generate function
-- AIDEV-NOTE: Added in v2025.8.15 - was missing despite being referenced in aliases and TODO lists
CREATE OR REPLACE FUNCTION steadytext_generate_async(
    prompt TEXT,
    max_tokens INT DEFAULT 512
)
RETURNS UUID
LANGUAGE plpgsql
AS $$
DECLARE
    request_id UUID;
BEGIN
    -- AIDEV-NOTE: Queue a generation request for async processing
    -- AIDEV-NOTE: Uses same pattern as other async functions (JSON, regex, choice)
    
    -- Validate input
    IF prompt IS NULL OR trim(prompt) = '' THEN
        RAISE EXCEPTION 'Prompt cannot be empty';
    END IF;
    
    IF max_tokens < 1 OR max_tokens > 4096 THEN
        RAISE EXCEPTION 'max_tokens must be between 1 and 4096';
    END IF;
    
    -- Generate UUID for this request
    request_id := gen_random_uuid();
    
    -- Insert into queue
    INSERT INTO @extschema@.steadytext_queue (
        request_id,
        request_type,
        prompt,
        params,
        status
    ) VALUES (
        request_id,
        'generate',
        prompt,
        jsonb_build_object(
            'max_tokens', max_tokens,
            'use_cache', true,
            'seed', 42
        ),
        'pending'
    );
    
    -- Notify workers
    PERFORM pg_notify('steadytext_queue', request_id::text);
    
    RETURN request_id;
END;
$$;

DROP FUNCTION IF EXISTS steadytext_generate_batch_async(TEXT[], INT);
CREATE OR REPLACE FUNCTION steadytext_generate_batch_async(
    prompts TEXT[],
    max_tokens INT DEFAULT 512
)
RETURNS UUID[]
LANGUAGE plpgsql
AS $$
DECLARE
    prompt TEXT;
    request_ids UUID[] := ARRAY[]::UUID[];
BEGIN
    IF prompts IS NULL OR array_length(prompts, 1) = 0 THEN
        PERFORM _steadytext_raise_p0001('Batch cannot be empty');
    END IF;

    IF array_length(prompts, 1) > 100 THEN
        PERFORM _steadytext_raise_p0001('Batch size exceeds maximum limit');
    END IF;

    FOREACH prompt IN ARRAY prompts LOOP
        IF prompt IS NULL THEN
            PERFORM _steadytext_raise_p0001('Prompt cannot be null');
        ELSIF trim(prompt) = '' THEN
            PERFORM _steadytext_raise_p0001('Prompt cannot be empty');
        END IF;

        request_ids := request_ids || steadytext_generate_async(prompt, max_tokens);
    END LOOP;

    RETURN request_ids;
END;
$$;

-- Async rerank function
-- AIDEV-NOTE: Async reranking with proper UUID handling and schema qualification
-- AIDEV-FIX: Cast request_id to UUID explicitly to avoid type mismatch errors
CREATE OR REPLACE FUNCTION steadytext_rerank_async(
    query text,
    documents text[],
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores boolean DEFAULT true,
    seed integer DEFAULT 42
) RETURNS uuid
AS $c$
    import uuid
    import json
    import logging
    
    # Generate request ID
    request_id = str(uuid.uuid4())
    
    # Prepare parameters
    params = {
        'query': query,
        'documents': documents,
        'task': task,
        'return_scores': return_scores,
        'seed': seed
    }
    
    # Get extension schema for proper qualification
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    
    # Insert into queue - cast request_id to UUID
    plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_queue 
        (request_id, request_type, params, status, created_at, priority)
        VALUES ($1::uuid, 'rerank', $2::jsonb, 'pending', CURRENT_TIMESTAMP, 5)
    """, ["text", "text"])
    plpy.execute(plan, [request_id, json.dumps(params)])
    
    # Send notification to worker
    plpy.execute("NOTIFY steadytext_queue_notify")
    
    return request_id
$c$ LANGUAGE plpython3u VOLATILE;

-- Batch rerank function for multiple queries
CREATE OR REPLACE FUNCTION steadytext_rerank_batch(
    queries text[],
    documents text[],
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores boolean DEFAULT true,
    seed integer DEFAULT 42
) RETURNS TABLE(query_index integer, document text, score float)
AS $c$
    import json
    import logging
    from typing import List, Tuple
    
    # Configure logging
    logging.basicConfig(level=logging.WARNING)
    logger = logging.getLogger(__name__)
    
    # Validate inputs
    if not queries or len(queries) == 0:
        plpy.error("Queries cannot be empty")
        
    if not documents or len(documents) == 0:
        return []
    
    # Check if initialized, if not, initialize now
    if not GD.get('steadytext_initialized', False):
        # Initialize on demand
        ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
        # Check again after initialization
        if not GD.get('steadytext_initialized', False):
            plpy.error("Failed to initialize pg_steadytext Python environment")

    from steadytext import rerank

    all_results = []

    # Process each query
    for idx, query in enumerate(queries):
        try:
            # Call rerank for this query
            results = rerank(
                query=query,
                documents=list(documents),
                task=task,
                return_scores=True,
                seed=seed
            )
            
            # Normalize results (handle both tuple and dict formats)
            for item in results:
                if isinstance(item, dict):
                    doc = item.get("document")
                    raw_score = item.get("score")
                else:
                    doc = item[0] if len(item) > 0 else None
                    raw_score = item[1] if len(item) > 1 else None

                if doc is None:
                    continue

                try:
                    score_value = float(raw_score) if raw_score is not None else 0.0
                except (TypeError, ValueError):
                    score_value = 0.0

                all_results.append({
                    "query_index": idx,
                    "document": doc,
                    "score": max(0.01, min(score_value, 1.0))
                })

        except Exception as e:
            logger.error(f"Reranking failed for query {idx}: {e}")
            # Continue with next query
            continue
    
    return all_results
$c$ LANGUAGE plpython3u IMMUTABLE PARALLEL SAFE;

-- Batch async rerank function
-- AIDEV-NOTE: Batch async reranking with proper UUID array handling
-- AIDEV-FIX: Each UUID must be cast individually when inserting into queue
CREATE OR REPLACE FUNCTION steadytext_rerank_batch_async(
    queries text[],
    documents text[],
    task text DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores boolean DEFAULT true,
    seed integer DEFAULT 42
) RETURNS uuid[]
AS $c$
    import uuid
    import json
    
    request_ids = []
    
    # Get extension schema for proper qualification
    ext_schema_result = plpy.execute("SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'")
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    
    # Prepare the insert statement once - cast request_id to UUID
    plan = plpy.prepare(f"""
        INSERT INTO {plpy.quote_ident(ext_schema)}.steadytext_queue 
        (request_id, request_type, params, status, created_at, priority)
        VALUES ($1::uuid, 'batch_rerank', $2::jsonb, 'pending', CURRENT_TIMESTAMP, 5)
    """, ["text", "text"])
    
    # Create separate async request for each query
    for query in queries:
        request_id = str(uuid.uuid4())
        request_ids.append(request_id)
        
        params = {
            'query': query,
            'documents': documents,
            'task': task,
            'return_scores': return_scores,
            'seed': seed
        }
        
        plpy.execute(plan, [request_id, json.dumps(params)])
    
    # Send notification to worker
    plpy.execute("NOTIFY steadytext_queue_notify")
    
    return request_ids
$c$ LANGUAGE plpython3u VOLATILE;

COMMENT ON FUNCTION steadytext_rerank IS 'Rerank documents by relevance to a query using AI model';
COMMENT ON FUNCTION steadytext_rerank_docs_only IS 'Rerank documents returning only sorted documents without scores';
COMMENT ON FUNCTION steadytext_rerank_top_k IS 'Rerank documents and return only top K results';
COMMENT ON FUNCTION steadytext_rerank_async IS 'Asynchronously rerank documents (returns request UUID)';
COMMENT ON FUNCTION steadytext_rerank_batch IS 'Rerank documents for multiple queries in batch';
COMMENT ON FUNCTION steadytext_rerank_batch_async IS 'Asynchronously rerank documents for multiple queries';

-- AIDEV-SECTION: CONFIGURATION_MANAGEMENT
-- Functions for managing configuration

-- Function to set configuration values
DROP FUNCTION IF EXISTS steadytext_config_set(TEXT, TEXT);
CREATE OR REPLACE FUNCTION steadytext_config_set(
    config_key TEXT,
    config_value TEXT
)
RETURNS VOID
LANGUAGE plpgsql
VOLATILE PARALLEL SAFE STRICT
AS $$
DECLARE
    v_key_lower TEXT := lower(config_key);
    v_port_int INT;
BEGIN
    -- Validate inputs for security-sensitive configuration keys
    IF v_key_lower = 'daemon_host' THEN
        IF config_value IS NULL OR config_value !~ '^[A-Za-z0-9._-]+$' THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid host format');
        END IF;
    ELSIF v_key_lower = 'daemon_port' THEN
        BEGIN
            v_port_int := config_value::INT;
        EXCEPTION WHEN others THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid port format');
        END;

        IF v_port_int < 1 OR v_port_int > 65535 THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid port format');
        END IF;
    ELSIF v_key_lower = 'model_path' THEN
        IF config_value IS NULL OR config_value LIKE '/%' OR config_value ~ '[;|]' THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid file path');
        END IF;
    ELSIF v_key_lower = 'daemon_command' THEN
        IF config_value IS NULL OR config_value ~ '[;&|]' OR config_value ~ '\s' THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid command format');
        END IF;
    ELSIF v_key_lower LIKE '%table_name' THEN
        IF config_value IS NULL OR config_value !~ '^[A-Za-z_][A-Za-z0-9_]*$' THEN
            PERFORM @extschema@._steadytext_raise_p0001('Invalid table name');
        END IF;
    END IF;

    -- Update or insert configuration value
    INSERT INTO @extschema@.steadytext_config (key, value, description, updated_at, updated_by)
    VALUES (config_key, to_jsonb(config_value), NULL, NOW(), current_user)
    ON CONFLICT (key) DO UPDATE
    SET value = to_jsonb(config_value),
        updated_at = NOW(),
        updated_by = current_user;

    -- Audit the configuration change
    INSERT INTO @extschema@.steadytext_audit_log (event_type, details)
    VALUES (
        'config_change',
        jsonb_build_object(
            'key', config_key,
            'value', config_value
        )
    );
END;
$$;

-- Function to get configuration values
DROP FUNCTION IF EXISTS steadytext_config_get(TEXT);
CREATE OR REPLACE FUNCTION steadytext_config_get(
    config_key TEXT
)
RETURNS TEXT
LANGUAGE sql
STABLE PARALLEL SAFE STRICT
AS $$
    SELECT value #>> '{}' FROM @extschema@.steadytext_config WHERE key = config_key;
$$;

-- AIDEV-SECTION: ASYNC_RESULT_RETRIEVAL
-- Functions for checking and retrieving async operation results

-- Enhanced async check function that handles all result types
CREATE OR REPLACE FUNCTION steadytext_check_async(
    request_id UUID
)
RETURNS TABLE(
    status TEXT,
    result TEXT,
    results TEXT[],
    embedding vector(1024),
    embeddings vector(1024)[],
    error TEXT,
    created_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    processing_time_ms INT,
    request_type TEXT
)
LANGUAGE sql
STABLE PARALLEL SAFE LEAKPROOF
AS $$
    SELECT 
        status,
        result,
        results,
        embedding,
        embeddings,
        error,
        created_at,
        completed_at,
        processing_time_ms,
        request_type
    FROM @extschema@.steadytext_queue
    WHERE @extschema@.steadytext_queue.request_id = steadytext_check_async.request_id;
$$;

-- Convenience function to get result directly (blocks until ready or timeout)
CREATE OR REPLACE FUNCTION steadytext_get_async_result(
    request_id UUID,
    timeout_seconds INT DEFAULT 30
)
RETURNS TEXT
LANGUAGE plpgsql
AS $$
DECLARE
    start_time TIMESTAMPTZ := clock_timestamp();
    check_result RECORD;
BEGIN
    -- AIDEV-NOTE: Polls for async result with timeout

    IF timeout_seconds IS NULL THEN
        timeout_seconds := 30;
    ELSIF timeout_seconds < 0 THEN
        PERFORM @extschema@._steadytext_raise_p0001('Timeout cannot be negative');
    END IF;
    
    LOOP
        -- Check status
        SELECT * INTO check_result
        FROM steadytext_check_async(request_id);
        
        -- Check if request exists
        IF check_result IS NULL THEN
            RAISE EXCEPTION 'Request ID % not found', request_id;
        END IF;
        
        -- Return result if completed
        IF check_result.status = 'completed' THEN
            RETURN check_result.result;
        END IF;
        
        -- Return error if failed
        IF check_result.status = 'failed' THEN
            RAISE EXCEPTION 'Async request failed: %', check_result.error;
        END IF;
        
        -- Check timeout
        IF clock_timestamp() - start_time > interval '1 second' * timeout_seconds THEN
            RAISE EXCEPTION 'Timeout waiting for async result (% seconds)', timeout_seconds;
        END IF;
        
        -- Wait a bit before checking again
        PERFORM pg_sleep(0.1);
    END LOOP;
END;
$$;

-- Function to cancel an async request
CREATE OR REPLACE FUNCTION steadytext_cancel_async(
    request_id UUID
)
RETURNS BOOLEAN
LANGUAGE plpgsql
AS $$
DECLARE
    rows_updated INT;
BEGIN
    -- AIDEV-NOTE: Cancel a pending async request
    
    UPDATE @extschema@.steadytext_queue
    SET status = 'cancelled',
        completed_at = NOW()
    WHERE @extschema@.steadytext_queue.request_id = steadytext_cancel_async.request_id
    AND status IN ('pending', 'processing');
    
    GET DIAGNOSTICS rows_updated = ROW_COUNT;
    
    RETURN rows_updated > 0;
END;
$$;

-- Batch async functions for checking multiple requests
CREATE OR REPLACE FUNCTION steadytext_embed_batch_async(
    texts TEXT[]
)
RETURNS UUID[]
LANGUAGE plpgsql
AS $$
DECLARE
    request_ids UUID[] := '{}';
    request_id UUID;
    i INT;
BEGIN
    -- AIDEV-NOTE: Queue multiple embedding requests
    
    -- Validate input
    IF texts IS NULL OR array_length(texts, 1) IS NULL THEN
        RAISE EXCEPTION 'texts array cannot be null or empty';
    END IF;
    
    -- Create a request for each text
    FOR i IN 1..array_length(texts, 1) LOOP
        -- Insert into queue
        INSERT INTO @extschema@.steadytext_queue (
            request_type,
            prompt
        ) VALUES (
            'embed',
            texts[i]
        )
        RETURNING steadytext_queue.request_id INTO request_id;
        
        request_ids := array_append(request_ids, request_id);
        
        -- Notify workers
        PERFORM pg_notify('steadytext_queue', request_id::text);
    END LOOP;
    
    RETURN request_ids;
END;
$$;

-- Function to check multiple async requests at once
CREATE OR REPLACE FUNCTION steadytext_check_async_batch(
    request_ids UUID[]
)
RETURNS TABLE(
    request_id UUID,
    status TEXT,
    result TEXT,
    error TEXT,
    completed_at TIMESTAMPTZ
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $$
    SELECT 
        request_id,
        status,
        result,
        error,
        completed_at
    FROM @extschema@.steadytext_queue
    WHERE request_id = ANY(request_ids)
    ORDER BY array_position(request_ids, request_id);
$$;

-- AIDEV-SECTION: VOLATILE_WRAPPER_FUNCTIONS
-- These functions provide cache-writing capability for users who need it
-- They wrap the IMMUTABLE functions and handle cache population

CREATE OR REPLACE FUNCTION steadytext_generate_cached(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    seed INT DEFAULT 42
)
RETURNS TEXT
LANGUAGE plpgsql
VOLATILE
AS $c$
DECLARE
    v_result TEXT;
    v_cache_key TEXT;
    v_generation_params JSONB;
BEGIN
    -- AIDEV-NOTE: This VOLATILE wrapper allows cache writes
    -- Use this when you need automatic cache population
    
    -- Generate result using IMMUTABLE function
    v_result := steadytext_generate(prompt, max_tokens, true, seed);
    
    -- If result was generated (not from cache), store it
    IF v_result IS NOT NULL THEN
        -- Generate cache key
        v_cache_key := prompt;
        
        -- Check if already cached
        -- AIDEV-NOTE: Use schema qualification for TimescaleDB continuous aggregates
        PERFORM 1 FROM @extschema@.steadytext_cache WHERE cache_key = v_cache_key;
        
        IF NOT FOUND THEN
            -- Resolve max_tokens default
                            IF max_tokens IS NULL THEN
                SELECT (value #>> '{}')::INT INTO max_tokens 
                FROM @extschema@.steadytext_config 
                WHERE key = 'default_max_tokens';
                max_tokens := COALESCE(max_tokens, 512);
            END IF;
            
            v_generation_params := jsonb_build_object(
                'max_tokens', max_tokens,
                'seed', seed
            );
            
            -- Store in cache
            INSERT INTO @extschema@.steadytext_cache 
            (cache_key, prompt, response, model_name, seed, generation_params)
            VALUES (v_cache_key, prompt, v_result, 'steadytext-default', seed, v_generation_params)
            ON CONFLICT (cache_key) DO NOTHING;
        END IF;
    END IF;
    
    RETURN v_result;
END;
$c$;

CREATE OR REPLACE FUNCTION steadytext_embed_cached(
    text_input TEXT,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS vector(1024)
LANGUAGE plpgsql
VOLATILE
AS $c$
DECLARE
    v_result vector(1024);
    v_cache_key TEXT;
    v_cache_key_input TEXT;
BEGIN
    -- AIDEV-NOTE: This VOLATILE wrapper allows cache writes with remote model support
    
    -- Validate remote model requirements
    IF model IS NOT NULL AND position(':' IN model) > 0 AND NOT unsafe_mode THEN
        RAISE EXCEPTION USING 
            ERRCODE = 'P0001',
            MESSAGE = 'Remote models (containing '':'' ) require unsafe_mode=TRUE';
    END IF;
    
    -- Generate result using IMMUTABLE function
    v_result := steadytext_embed(text_input, true, seed, model, unsafe_mode);
    
    -- If result was generated, store it
    IF v_result IS NOT NULL THEN
        -- Generate cache key including model if specified
        IF model IS NOT NULL THEN
            v_cache_key_input := 'embed:' || text_input || ':' || model;
        ELSE
            v_cache_key_input := 'embed:' || text_input;
        END IF;
        v_cache_key := encode(digest(v_cache_key_input, 'sha256'), 'hex');
        
        -- Check if already cached
        -- AIDEV-NOTE: Use schema qualification for TimescaleDB continuous aggregates
        PERFORM 1 FROM @extschema@.steadytext_cache WHERE cache_key = v_cache_key;
        
        IF NOT FOUND THEN
            -- Store in cache
            INSERT INTO @extschema@.steadytext_cache 
            (cache_key, prompt, embedding, created_at)
            VALUES (v_cache_key, text_input, v_result, NOW())
            ON CONFLICT (cache_key) DO NOTHING;
        END IF;
    END IF;
    
    RETURN v_result;
END;
$c$;

-- Create async embedding function with model and unsafe_mode parameters
CREATE OR REPLACE FUNCTION steadytext_embed_async(
    text_input TEXT,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS UUID
LANGUAGE plpgsql
AS $$
DECLARE
    request_id UUID;
BEGIN
    -- AIDEV-NOTE: Queue an embedding request for async processing with remote model support
    
    -- Validate remote model requirements
    IF model IS NOT NULL AND position(':' IN model) > 0 AND NOT unsafe_mode THEN
        RAISE EXCEPTION USING 
            ERRCODE = 'P0001',
            MESSAGE = 'Remote models (containing '':'' ) require unsafe_mode=TRUE';
    END IF;
    
    request_id := gen_random_uuid();
    
    -- Build params JSON including new parameters
    DECLARE
        params_json JSONB;
    BEGIN
        params_json := jsonb_build_object(
            'text_input', text_input,
            'use_cache', use_cache,
            'seed', seed
        );
        
        IF model IS NOT NULL THEN
            params_json := params_json || jsonb_build_object('model', model);
        END IF;
        
        IF unsafe_mode IS NOT NULL THEN
            params_json := params_json || jsonb_build_object('unsafe_mode', unsafe_mode);
        END IF;
        
        INSERT INTO @extschema@.steadytext_queue (
            request_id,
            request_type,
            params,
            status
        ) VALUES (
            request_id,
            'embed',
            params_json,
            'pending'
        );
    END;
    
    -- Notify worker
    PERFORM pg_notify('steadytext_async', request_id::text);
    
    RETURN request_id;
END;
$$;

-- Add comments explaining the wrapper functions
COMMENT ON FUNCTION steadytext_generate_cached IS 
'VOLATILE wrapper for steadytext_generate that handles cache population. Use when automatic caching is needed.';

COMMENT ON FUNCTION steadytext_embed_cached IS 
'VOLATILE wrapper for steadytext_embed that handles cache population. Use when automatic caching is needed.';

COMMENT ON FUNCTION steadytext_embed_async IS
'Async embedding generation with remote model support. Returns UUID for checking status with steadytext_get_result.';

-- Add note about cache population strategies
COMMENT ON FUNCTION steadytext_generate(text, integer, boolean, integer, text, text, text, text, text, boolean) IS 
'IMMUTABLE function for text generation. Only reads from cache, never writes. For automatic cache population, use steadytext_generate_cached.';

COMMENT ON FUNCTION steadytext_embed IS 
'IMMUTABLE function for embeddings. Only reads from cache, never writes. For automatic cache population, use steadytext_embed_cached.';

-- Grant appropriate permissions
GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA public TO PUBLIC;
GRANT USAGE ON ALL SEQUENCES IN SCHEMA public TO PUBLIC;
GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO PUBLIC;

-- AIDEV-NOTE: This completes the base schema for pg_steadytext v1.4.2
--
-- AIDEV-SECTION: CHANGES_MADE_IN_REVIEW
-- The following issues were identified and fixed during review:
-- 1. Added missing columns: model_size, eos_string, response_size, daemon_endpoint
-- 2. Enhanced queue table with priority, user_id, session_id, batch support
-- 3. Added rate limiting and audit logging tables
-- 4. Fixed cache key generation to use SHA256 and match SteadyText format
-- 5. Fixed daemon integration to use proper SteadyText API methods
-- 6. Added proper indexes for performance
--
-- AIDEV-TODO: Future versions should add:
-- - Async processing functions (steadytext_generate_async, steadytext_get_result)
-- - Streaming generation function (steadytext_generate_stream)
-- - Batch operations (steadytext_embed_batch)
-- - FAISS index operations (steadytext_index_create, steadytext_index_search)
-- - Worker management functions
-- - Enhanced security and rate limiting functions
-- - Support for Pydantic models in structured generation (needs JSON serialization)
-- - Tests for structured generation functions

-- AIDEV-NOTE: Added in v1.0.1 (2025-07-07):
-- Marked all deterministic functions as IMMUTABLE, PARALLEL SAFE, and LEAKPROOF (where allowed):
-- - steadytext_generate(), steadytext_embed(), steadytext_generate_json(),
--   steadytext_generate_regex(), steadytext_generate_choice() are IMMUTABLE PARALLEL SAFE
-- - steadytext_version() is IMMUTABLE PARALLEL SAFE LEAKPROOF
-- - steadytext_cache_stats() and steadytext_config_get() are STABLE PARALLEL SAFE
-- - steadytext_config_get() is also LEAKPROOF since it's a simple SQL function
-- This enables use with TimescaleDB and in aggregates, and improves query optimization

-- AIDEV-NOTE: Added in v1.1.0 (2025-07-08):
-- AI summarization aggregate functions with the following features:
-- 1. Structured fact extraction to mitigate non-transitivity
-- 2. Semantic deduplication using embeddings
-- 3. Statistical tracking (row counts, character lengths)
-- 4. Sample preservation for context
-- 5. Combine depth tracking for adaptive prompts
-- 6. Full TimescaleDB continuous aggregate support
-- 7. Serialization for distributed aggregation

-- AIDEV-NOTE: Updated in v1.2.0 (2025-07-08):
-- Improved AI summarization aggregate functions with:
-- 1. Better error handling throughout all functions
-- 2. Input validation for all parameters
-- 3. Protection against division by zero in cosine similarity calculations
-- 4. Specific exception handling instead of bare except clauses
-- 5. Proper handling of invalid JSON inputs
-- 6. Zero-norm vector protection in similarity calculations
-- 7. Graceful fallback when parsing fails
-- 8. FIXED: Removed serialization functions to resolve "serialization functions may be
--    specified only when the aggregate transition data type is internal" error

-- AIDEV-NOTE: Added in v1.3.0 (2025-07-09):
-- Reranking functions for query-document relevance scoring:
-- 1. Basic rerank function with optional score return
-- 2. Document-only variant for simplified output
-- 3. Top-k filtering for returning best matches
-- 4. Async processing support via queue system
-- 5. Batch reranking for multiple queries
-- 6. Integration with SteadyText daemon for Qwen3-Reranker-4B model

-- AIDEV-NOTE: Added in v2025.8.15 (2025-07-31):
-- Short aliases for all steadytext functions (st_* for steadytext_*)
-- Manual creation is required to preserve default parameter values
-- Dynamic generation would create functions without DEFAULT clauses, breaking the API

-- st_generate alias
CREATE OR REPLACE FUNCTION st_generate(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    eos_string TEXT DEFAULT '[EOS]',
    model TEXT DEFAULT NULL,
    model_repo TEXT DEFAULT NULL,
    model_filename TEXT DEFAULT NULL,
    size TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);
$alias$;

-- st_generate_stream alias
CREATE OR REPLACE FUNCTION st_generate_stream(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    eos_string TEXT DEFAULT '[EOS]',
    model TEXT DEFAULT NULL,
    model_repo TEXT DEFAULT NULL,
    model_filename TEXT DEFAULT NULL,
    size TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS SETOF TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_generate_stream($1, $2, $3, $4, $5, $6, $7, $8, $9, $10);
$alias$;

DO $$
BEGIN
    BEGIN
        ALTER EXTENSION pg_steadytext ADD FUNCTION st_generate_stream(text, integer, boolean, integer, text, text, text, text, text, boolean);
    EXCEPTION WHEN OTHERS THEN
        NULL;
    END;
END $$;

-- st_embed alias
CREATE OR REPLACE FUNCTION st_embed(
    text_input TEXT,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS vector(1024)
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT steadytext_embed($1, $2, $3, $4, $5);
$alias$;

-- st_embed_cached alias
CREATE OR REPLACE FUNCTION st_embed_cached(
    text_input TEXT,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS vector(1024)
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_embed_cached($1, $2, $3, $4);
$alias$;

-- st_embed_async alias
CREATE OR REPLACE FUNCTION st_embed_async(
    text_input TEXT,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    model TEXT DEFAULT NULL,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS UUID
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_embed_async($1, $2, $3, $4, $5);
$alias$;

-- st_generate_cached alias
CREATE OR REPLACE FUNCTION st_generate_cached(
    prompt TEXT,
    max_tokens INT DEFAULT NULL,
    seed INT DEFAULT 42
)
RETURNS TEXT
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate_cached($1, $2, $3);
$alias$;

-- st_generate_json alias
CREATE OR REPLACE FUNCTION st_generate_json(
    prompt TEXT,
    schema JSONB,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate_json($1, $2, $3, $4, $5, $6);
$alias$;

-- st_generate_regex alias
CREATE OR REPLACE FUNCTION st_generate_regex(
    prompt TEXT,
    pattern TEXT,
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate_regex($1, $2, $3, $4, $5, $6);
$alias$;

-- st_generate_choice alias
CREATE OR REPLACE FUNCTION st_generate_choice(
    prompt TEXT,
    choices TEXT[],
    max_tokens INT DEFAULT NULL,
    use_cache BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42,
    unsafe_mode BOOLEAN DEFAULT FALSE
)
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate_choice($1, $2, $3, $4, $5, $6);
$alias$;

-- st_rerank alias
CREATE OR REPLACE FUNCTION st_rerank(
    query TEXT,
    documents TEXT[],
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42
)
RETURNS TABLE(document TEXT, score DOUBLE PRECISION)
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_rerank($1, $2, $3, $4, $5);
$alias$;

-- st_rerank_docs_only alias
CREATE OR REPLACE FUNCTION st_rerank_docs_only(
    query TEXT,
    documents TEXT[],
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    seed INT DEFAULT 42
)
RETURNS TABLE(document TEXT)
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_rerank_docs_only($1, $2, $3, $4);
$alias$;

-- st_rerank_top_k alias
CREATE OR REPLACE FUNCTION st_rerank_top_k(
    query TEXT,
    documents TEXT[],
    k INT DEFAULT 5,
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42
)
RETURNS TABLE(document TEXT, score DOUBLE PRECISION)
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_rerank_top_k($1, $2, $3, $4, $5, $6);
$alias$;

-- st_rerank_batch alias
DROP FUNCTION IF EXISTS st_rerank_batch(TEXT[], TEXT[], TEXT, BOOLEAN, INTEGER);
CREATE OR REPLACE FUNCTION st_rerank_batch(
    queries TEXT[],
    documents TEXT[],
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42
)
RETURNS TABLE(query_index INTEGER, document TEXT, score DOUBLE PRECISION)
LANGUAGE sql
IMMUTABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_rerank_batch($1, $2, $3, $4, $5);
$alias$;

-- Async function aliases
-- st_generate_async alias
CREATE OR REPLACE FUNCTION st_generate_async(
    prompt TEXT,
    max_tokens INT DEFAULT NULL
)
RETURNS UUID
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_generate_async($1, $2);
$alias$;

-- st_rerank_async alias
CREATE OR REPLACE FUNCTION st_rerank_async(
    query TEXT,
    documents TEXT[],
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42
)
RETURNS UUID
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_rerank_async($1, $2, $3, $4, $5);
$alias$;

-- st_rerank_batch_async alias
CREATE OR REPLACE FUNCTION st_rerank_batch_async(
    queries TEXT[],
    documents TEXT[],
    task TEXT DEFAULT 'Given a web search query, retrieve relevant passages that answer the query',
    return_scores BOOLEAN DEFAULT TRUE,
    seed INT DEFAULT 42
)
RETURNS UUID[]
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_rerank_batch_async($1, $2, $3, $4, $5);
$alias$;

-- st_check_async alias
CREATE OR REPLACE FUNCTION st_check_async(
    request_id UUID
)
RETURNS TABLE(
    status TEXT,
    result TEXT,
    results TEXT[],
    embedding vector(1024),
    embeddings vector(1024)[],
    error TEXT,
    created_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    processing_time_ms INT,
    request_type TEXT
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_check_async($1);
$alias$;

-- st_get_async_result alias
CREATE OR REPLACE FUNCTION st_get_async_result(
    request_id UUID,
    timeout_seconds INT DEFAULT 30
)
RETURNS TEXT
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_get_async_result($1, $2);
$alias$;

-- Cache management aliases
-- st_cache_stats alias
CREATE OR REPLACE FUNCTION st_cache_stats()
RETURNS TABLE(total_entries BIGINT, total_size_mb DOUBLE PRECISION, cache_hit_rate DOUBLE PRECISION, avg_access_count DOUBLE PRECISION, oldest_entry TIMESTAMP WITH TIME ZONE, newest_entry TIMESTAMP WITH TIME ZONE)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_cache_stats();
$alias$;

-- st_cache_clear alias
CREATE OR REPLACE FUNCTION st_cache_clear()
RETURNS BIGINT
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_cache_clear();
$alias$;

-- st_cache_evict_by_age alias
CREATE OR REPLACE FUNCTION st_cache_evict_by_age(
    target_entries INT DEFAULT NULL,
    target_size_mb DOUBLE PRECISION DEFAULT NULL,
    batch_size INT DEFAULT 100,
    min_age_hours INT DEFAULT 1
)
RETURNS TABLE(evicted_count INTEGER, freed_size_mb DOUBLE PRECISION, remaining_entries BIGINT, remaining_size_mb DOUBLE PRECISION)
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_cache_evict_by_age($1, $2, $3, $4);
$alias$;

-- st_cache_preview_eviction alias
CREATE OR REPLACE FUNCTION st_cache_preview_eviction(
    preview_count INT DEFAULT 10
)
RETURNS TABLE(cache_key TEXT, prompt TEXT, access_count INTEGER, last_accessed TIMESTAMP WITH TIME ZONE, created_at TIMESTAMP WITH TIME ZONE, age_days DOUBLE PRECISION, size_bytes BIGINT)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_cache_preview_eviction($1);
$alias$;

-- st_cache_analyze_usage alias
CREATE OR REPLACE FUNCTION st_cache_analyze_usage()
RETURNS TABLE(age_bucket TEXT, entry_count BIGINT, avg_access_count DOUBLE PRECISION, total_size_mb DOUBLE PRECISION, percentage_of_cache DOUBLE PRECISION)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM steadytext_cache_analyze_usage();
$alias$;

-- st_cache_scheduled_eviction alias
CREATE OR REPLACE FUNCTION st_cache_scheduled_eviction()
RETURNS VOID
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT steadytext_cache_scheduled_eviction();
$alias$;

-- Daemon management aliases
-- st_daemon_start alias
CREATE OR REPLACE FUNCTION st_daemon_start()
RETURNS BOOLEAN
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT @extschema@.steadytext_daemon_start();
$alias$;

-- st_daemon_status alias
CREATE OR REPLACE FUNCTION st_daemon_status()
RETURNS TABLE(daemon_id TEXT, status TEXT, endpoint TEXT, last_heartbeat TIMESTAMP WITH TIME ZONE, uptime_seconds INTEGER)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $alias$
    SELECT * FROM @extschema@.steadytext_daemon_status();
$alias$;

-- st_daemon_stop alias
CREATE OR REPLACE FUNCTION st_daemon_stop()
RETURNS BOOLEAN
LANGUAGE sql
VOLATILE PARALLEL SAFE
AS $alias$
    SELECT @extschema@.steadytext_daemon_stop();
$alias$;

-- Configuration aliases
-- st_config_get alias
CREATE OR REPLACE FUNCTION st_config_get(
    key TEXT
)
RETURNS TEXT
LANGUAGE sql
STABLE PARALLEL SAFE STRICT
AS $alias$
    SELECT steadytext_config_get($1);
$alias$;

-- st_config_set alias
CREATE OR REPLACE FUNCTION st_config_set(
    key TEXT,
    value TEXT
)
RETURNS VOID
LANGUAGE sql
VOLATILE PARALLEL SAFE STRICT
AS $alias$
    SELECT steadytext_config_set($1, $2);
$alias$;

-- st_version alias
CREATE OR REPLACE FUNCTION st_version()
RETURNS TEXT
LANGUAGE sql
IMMUTABLE PARALLEL SAFE LEAKPROOF
AS $alias$
    SELECT steadytext_version();
$alias$;

-- AIDEV-SECTION: SUMMARIZATION_ALIASES
-- Create st_ aliases for summarization functions
CREATE OR REPLACE FUNCTION st_summarize_text(
    input_text text,
    metadata jsonb DEFAULT '{}'::jsonb,
    model text DEFAULT NULL,
    unsafe_mode boolean DEFAULT FALSE
) RETURNS text
LANGUAGE sql IMMUTABLE PARALLEL SAFE
AS $$
    SELECT @extschema@.steadytext_summarize_text($1, $2, $3, $4);
$$;

CREATE OR REPLACE FUNCTION st_extract_facts(
    input_text text,
    max_facts integer DEFAULT 10
) RETURNS jsonb
LANGUAGE sql IMMUTABLE PARALLEL SAFE
AS $$
    SELECT @extschema@.steadytext_extract_facts($1, $2);
$$;

CREATE OR REPLACE FUNCTION st_deduplicate_facts(
    facts jsonb,
    similarity_threshold float DEFAULT 0.8
) RETURNS jsonb
LANGUAGE sql IMMUTABLE PARALLEL SAFE
AS $$
    SELECT @extschema@.steadytext_deduplicate_facts($1, $2);
$$;

-- Aggregate aliases
CREATE OR REPLACE AGGREGATE st_summarize(text, jsonb) (
    SFUNC = steadytext_summarize_accumulate,
    STYPE = jsonb,
    COMBINEFUNC = steadytext_summarize_combine,
    FINALFUNC = steadytext_summarize_finalize,
    PARALLEL = SAFE
);

-- ================================================================
-- PROMPT REGISTRY FEATURE (v2025.9.6)
-- ================================================================
-- AIDEV-NOTE: Lightweight Jinja2-based prompt template management with versioning

-- Drop tables if they exist (for clean installation)
DROP TABLE IF EXISTS steadytext_prompt_versions CASCADE;
DROP TABLE IF EXISTS steadytext_prompts CASCADE;

-- Create prompts table
CREATE TABLE steadytext_prompts (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    slug TEXT UNIQUE NOT NULL CHECK (slug ~ '^[a-z0-9]+(?:-[a-z0-9]+)*$' AND LENGTH(slug) BETWEEN 3 AND 100),
    description TEXT CHECK (LENGTH(description) <= 500),
    created_at TIMESTAMPTZ DEFAULT NOW(),
    created_by TEXT DEFAULT current_user,
    updated_at TIMESTAMPTZ DEFAULT NOW(),
    updated_by TEXT DEFAULT current_user
);

-- Create prompt versions table
CREATE TABLE steadytext_prompt_versions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    prompt_id UUID NOT NULL REFERENCES steadytext_prompts(id) ON DELETE CASCADE,
    version INTEGER NOT NULL CHECK (version > 0),
    template TEXT NOT NULL CHECK (LENGTH(template) > 0),
    required_variables TEXT[],
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    created_by TEXT DEFAULT current_user,
    is_active BOOLEAN DEFAULT FALSE,
    UNIQUE(prompt_id, version)
);

-- Create indexes for performance
CREATE INDEX idx_steadytext_prompts_slug ON steadytext_prompts(slug);
CREATE INDEX idx_steadytext_prompts_created_at ON steadytext_prompts(created_at DESC);
CREATE INDEX idx_steadytext_prompt_versions_prompt_id_active ON steadytext_prompt_versions(prompt_id, is_active) WHERE is_active = TRUE;
CREATE INDEX idx_steadytext_prompt_versions_created_at ON steadytext_prompt_versions(created_at DESC);

-- Helper function to get next version number
CREATE OR REPLACE FUNCTION _get_next_version(p_prompt_id UUID)
RETURNS INTEGER
LANGUAGE sql
IMMUTABLE PARALLEL SAFE  -- IMMUTABLE: Intentional, though reads from table (used in deterministic context)
AS $$
    SELECT COALESCE(MAX(version), 0) + 1
    FROM @extschema@.steadytext_prompt_versions
    WHERE prompt_id = p_prompt_id;
$$;

-- Python function to validate Jinja2 templates
CREATE OR REPLACE FUNCTION _validate_jinja2_template(template TEXT)
RETURNS TABLE(is_valid BOOLEAN, required_variables TEXT[], error_message TEXT)
LANGUAGE plpython3u
AS $$
    # AIDEV-NOTE: Validates Jinja2 template syntax and extracts required variables
    import re
    
    # Initialize Python environment if needed
    if not GD.get('steadytext_initialized', False):
        ext_schema_result = plpy.execute(
            "SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'"
        )
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    
    # Try to import Jinja2
    try:
        from jinja2 import Environment, meta, TemplateSyntaxError
    except ImportError:
        return [(False, None, "Jinja2 is not installed. Please install it using: pip install jinja2")]
    
    env = Environment()
    
    try:
        # Parse the template to check syntax
        ast = env.parse(template)
        
        # Extract all variables referenced in the template
        variables = meta.find_undeclared_variables(ast)
        
        # Convert to sorted list for consistency
        required_vars = sorted(list(variables))
        
        return [(True, required_vars, None)]
        
    except TemplateSyntaxError as e:
        return [(False, None, str(e))]
    except Exception as e:
        return [(False, None, f"Unexpected error: {str(e)}")]
$$;

-- Create prompt with first version
CREATE OR REPLACE FUNCTION steadytext_prompt_create(
    slug TEXT,
    template TEXT,
    description TEXT DEFAULT NULL,
    metadata JSONB DEFAULT '{}'
) RETURNS UUID
LANGUAGE plpgsql
AS $$
DECLARE
    v_prompt_id UUID;
    v_validation RECORD;
BEGIN
    -- Validate template
    SELECT * INTO v_validation 
    FROM @extschema@._validate_jinja2_template(template);
    
    IF NOT v_validation.is_valid THEN
        RAISE EXCEPTION 'Invalid Jinja2 template: %', v_validation.error_message
            USING ERRCODE = 'syntax_error';
    END IF;
    
    -- Insert prompt
    INSERT INTO @extschema@.steadytext_prompts (slug, description)
    VALUES (slug, description)
    RETURNING id INTO v_prompt_id;
    
    -- Insert first version
    INSERT INTO @extschema@.steadytext_prompt_versions (
        prompt_id, version, template, required_variables, metadata, is_active
    ) VALUES (
        v_prompt_id, 1, template, v_validation.required_variables, metadata, TRUE
    );
    
    RETURN v_prompt_id;
EXCEPTION
    WHEN unique_violation THEN
        RAISE EXCEPTION 'Prompt with slug "%" already exists', slug
            USING ERRCODE = 'unique_violation';
    WHEN check_violation THEN
        RAISE EXCEPTION 'Invalid slug format. Use lowercase letters, numbers, and hyphens only (3-100 chars)'
            USING ERRCODE = 'invalid_text_representation';
END;
$$;

-- Update prompt (create new version)
CREATE OR REPLACE FUNCTION steadytext_prompt_update(
    slug TEXT,
    template TEXT,
    metadata JSONB DEFAULT '{}'
) RETURNS UUID
LANGUAGE plpgsql
AS $$
DECLARE
    v_prompt_id UUID;
    v_next_version INTEGER;
    v_version_id UUID;
    v_validation RECORD;
BEGIN
    -- Get prompt ID
    SELECT id INTO v_prompt_id
    FROM @extschema@.steadytext_prompts
    WHERE steadytext_prompts.slug = steadytext_prompt_update.slug;
    
    IF v_prompt_id IS NULL THEN
        RAISE EXCEPTION 'Prompt with slug "%" not found', slug
            USING ERRCODE = 'no_data_found';
    END IF;
    
    -- Acquire advisory lock to prevent concurrent version updates
    PERFORM pg_advisory_xact_lock(hashtextextended(v_prompt_id::text, 0));
    
    -- Validate template
    SELECT * INTO v_validation 
    FROM @extschema@._validate_jinja2_template(template);
    
    IF NOT v_validation.is_valid THEN
        RAISE EXCEPTION 'Invalid Jinja2 template: %', v_validation.error_message
            USING ERRCODE = 'syntax_error';
    END IF;
    
    -- Get next version number
    SELECT @extschema@._get_next_version(v_prompt_id) INTO v_next_version;
    
    -- Deactivate current active version
    UPDATE @extschema@.steadytext_prompt_versions
    SET is_active = FALSE
    WHERE prompt_id = v_prompt_id AND is_active = TRUE;
    
    -- Insert new version
    INSERT INTO @extschema@.steadytext_prompt_versions (
        prompt_id, version, template, required_variables, metadata, is_active
    ) VALUES (
        v_prompt_id, v_next_version, template, v_validation.required_variables, metadata, TRUE
    ) RETURNING id INTO v_version_id;
    
    -- Update prompt's updated_at
    UPDATE @extschema@.steadytext_prompts
    SET updated_at = NOW(), updated_by = current_user
    WHERE id = v_prompt_id;
    
    RETURN v_version_id;
END;
$$;

-- Get prompt template
CREATE OR REPLACE FUNCTION steadytext_prompt_get(
    slug TEXT,
    version INTEGER DEFAULT NULL
) RETURNS TABLE(
    prompt_id UUID,
    version_num INTEGER,
    template TEXT,
    required_variables TEXT[],
    metadata JSONB,
    created_at TIMESTAMPTZ,
    created_by TEXT
)
LANGUAGE plpgsql
AS $$
DECLARE
    v_prompt_id UUID;
BEGIN
    -- Get prompt ID
    SELECT id INTO v_prompt_id
    FROM @extschema@.steadytext_prompts p
    WHERE p.slug = steadytext_prompt_get.slug;
    
    IF v_prompt_id IS NULL THEN
        RAISE EXCEPTION 'Prompt with slug "%" not found', slug
            USING ERRCODE = 'no_data_found';
    END IF;
    
    -- Return specific version or active version
    IF version IS NOT NULL THEN
        RETURN QUERY
        SELECT pv.prompt_id, pv.version AS version_num, pv.template, pv.required_variables, 
               pv.metadata, pv.created_at, pv.created_by
        FROM @extschema@.steadytext_prompt_versions pv
        WHERE pv.prompt_id = v_prompt_id AND pv.version = steadytext_prompt_get.version;
        
        IF NOT FOUND THEN
            RAISE EXCEPTION 'Version % not found for prompt "%"', version, slug
                USING ERRCODE = 'no_data_found';
        END IF;
    ELSE
        RETURN QUERY
        SELECT pv.prompt_id, pv.version AS version_num, pv.template, pv.required_variables, 
               pv.metadata, pv.created_at, pv.created_by
        FROM @extschema@.steadytext_prompt_versions pv
        WHERE pv.prompt_id = v_prompt_id AND pv.is_active = TRUE;
    END IF;
END;
$$;

-- Render prompt template
CREATE OR REPLACE FUNCTION steadytext_prompt_render(
    slug TEXT,
    variables JSONB,
    version INTEGER DEFAULT NULL,
    strict BOOLEAN DEFAULT TRUE
) RETURNS TEXT
LANGUAGE plpython3u
AS $$
    # AIDEV-NOTE: Renders Jinja2 template with provided variables
    
    # Initialize Python environment if needed
    if not GD.get('steadytext_initialized', False):
        ext_schema_result = plpy.execute(
            "SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'"
        )
        ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
        plpy.execute(f"SELECT {plpy.quote_ident(ext_schema)}._steadytext_init_python()")
    
    # Get extension schema
    ext_schema_result = plpy.execute(
        "SELECT nspname FROM pg_extension e JOIN pg_namespace n ON e.extnamespace = n.oid WHERE e.extname = 'pg_steadytext'"
    )
    ext_schema = ext_schema_result[0]['nspname'] if ext_schema_result else 'public'
    
    # Try to import Jinja2
    try:
        from jinja2 import Environment, TemplateSyntaxError, UndefinedError, StrictUndefined, Undefined
    except ImportError:
        plpy.error("Jinja2 is not installed. Please install it using: pip install jinja2")
    
    import json
    
    # Get the template
    if version is not None:
        query = f"""
            SELECT template, required_variables 
            FROM {plpy.quote_ident(ext_schema)}.steadytext_prompt_get($1, $2)
        """
        plan = plpy.prepare(query, ["text", "integer"])
        result = plpy.execute(plan, [slug, version])
    else:
        query = f"""
            SELECT template, required_variables 
            FROM {plpy.quote_ident(ext_schema)}.steadytext_prompt_get($1)
        """
        plan = plpy.prepare(query, ["text"])
        result = plpy.execute(plan, [slug])
    
    if not result:
        plpy.error(f"Prompt with slug '{slug}' not found")
    
    template_text = result[0]['template']
    required_vars = result[0]['required_variables'] or []
    
    # Parse variables
    # AIDEV-FIX: Ensure proper conversion of JSONB to dict
    if isinstance(variables, str):
        try:
            vars_dict = json.loads(variables)
        except json.JSONDecodeError as e:
            plpy.error(f"Invalid JSON in variables: {str(e)}")
    elif hasattr(variables, 'items'):
        # It's already dict-like, but ensure it's a plain dict
        vars_dict = dict(variables) if variables else {}
    else:
        vars_dict = {}
    
    # Check for missing required variables in strict mode
    if strict and required_vars:
        missing = [var for var in required_vars if var not in vars_dict]
        if missing:
            plpy.error(f"Missing required variables: {', '.join(missing)}")
    
    # Create Jinja2 environment
    if strict:
        env = Environment(undefined=StrictUndefined)
    else:
        env = Environment(undefined=Undefined)
    
    # Cache compiled templates in GD for performance
    # Include strict mode in cache key since it affects template behavior
    cache_key = f"jinja2_template:{slug}:{version}:{strict}:{hash(template_text)}"
    
    # AIDEV-FIX: Don't cache compiled templates - they're not serializable in GD
    # Instead, compile fresh each time (Jinja2 compilation is fast)
    try:
        compiled_template = env.from_string(template_text)
    except TemplateSyntaxError as e:
        plpy.error(f"Template syntax error: {str(e)}")
    
    # Render the template
    try:
        # AIDEV-NOTE: JSONB dict limitation - Cannot use 'items', 'keys', 'values' as JSON keys
        # These conflict with dict methods when JSONB is passed to Jinja2 templates
        rendered = compiled_template.render(**vars_dict)
        return rendered
    except UndefinedError as e:
        plpy.error(f"Undefined variable in template: {str(e)}")
    except Exception as e:
        plpy.error(f"Template rendering error: {str(e)}")
$$;

-- List all prompts
CREATE OR REPLACE FUNCTION steadytext_prompt_list()
RETURNS TABLE(
    prompt_id UUID,
    slug TEXT,
    description TEXT,
    latest_version_num INTEGER,
    total_versions INTEGER,
    created_at TIMESTAMPTZ,
    updated_at TIMESTAMPTZ
)
LANGUAGE sql
STABLE PARALLEL SAFE
AS $$
    SELECT 
        p.id,
        p.slug,
        p.description,
        MAX(pv.version) as latest_version_num,
        COUNT(pv.id)::INTEGER as total_versions,
        p.created_at,
        p.updated_at
    FROM @extschema@.steadytext_prompts p
    LEFT JOIN @extschema@.steadytext_prompt_versions pv ON p.id = pv.prompt_id
    GROUP BY p.id, p.slug, p.description, p.created_at, p.updated_at
    ORDER BY p.created_at DESC;
$$;

-- List versions of a prompt
CREATE OR REPLACE FUNCTION steadytext_prompt_versions(slug TEXT)
RETURNS TABLE(
    version_num INTEGER,
    template TEXT,
    required_variables TEXT[],
    metadata JSONB,
    created_at TIMESTAMPTZ,
    created_by TEXT,
    is_active BOOLEAN
)
LANGUAGE plpgsql
AS $$
DECLARE
    v_prompt_id UUID;
BEGIN
    -- Get prompt ID
    SELECT id INTO v_prompt_id
    FROM @extschema@.steadytext_prompts p
    WHERE p.slug = steadytext_prompt_versions.slug;
    
    IF v_prompt_id IS NULL THEN
        RAISE EXCEPTION 'Prompt with slug "%" not found', slug
            USING ERRCODE = 'no_data_found';
    END IF;
    
    RETURN QUERY
    SELECT pv.version AS version_num, pv.template, pv.required_variables, pv.metadata,
           pv.created_at, pv.created_by, pv.is_active
    FROM @extschema@.steadytext_prompt_versions pv
    WHERE pv.prompt_id = v_prompt_id
    ORDER BY pv.version DESC;
END;
$$;

-- Delete prompt and all versions
CREATE OR REPLACE FUNCTION steadytext_prompt_delete(slug TEXT)
RETURNS BOOLEAN
LANGUAGE plpgsql
AS $$
DECLARE
    v_prompt_id UUID;
BEGIN
    -- Get prompt ID
    SELECT id INTO v_prompt_id
    FROM @extschema@.steadytext_prompts p
    WHERE p.slug = steadytext_prompt_delete.slug;
    
    IF v_prompt_id IS NULL THEN
        RETURN FALSE;
    END IF;
    
    -- Delete prompt (cascade will delete versions)
    DELETE FROM @extschema@.steadytext_prompts
    WHERE id = v_prompt_id;
    
    RETURN TRUE;
END;
$$;

-- Create short aliases for all prompt functions
CREATE OR REPLACE FUNCTION st_prompt_create(
    slug TEXT,
    template TEXT,
    description TEXT DEFAULT NULL,
    metadata JSONB DEFAULT '{}'
) RETURNS UUID
LANGUAGE sql
AS $$
    SELECT @extschema@.steadytext_prompt_create($1, $2, $3, $4);
$$;

CREATE OR REPLACE FUNCTION st_prompt_update(
    slug TEXT,
    template TEXT,
    metadata JSONB DEFAULT '{}'
) RETURNS UUID
LANGUAGE sql
AS $$
    SELECT @extschema@.steadytext_prompt_update($1, $2, $3);
$$;

CREATE OR REPLACE FUNCTION st_prompt_get(
    slug TEXT,
    version INTEGER DEFAULT NULL
) RETURNS TABLE(
    prompt_id UUID,
    version_num INTEGER,
    template TEXT,
    required_variables TEXT[],
    metadata JSONB,
    created_at TIMESTAMPTZ,
    created_by TEXT
)
LANGUAGE sql
AS $$
    SELECT * FROM @extschema@.steadytext_prompt_get($1, $2);
$$;

CREATE OR REPLACE FUNCTION st_prompt_render(
    slug TEXT,
    variables JSONB,
    version INTEGER DEFAULT NULL,
    strict BOOLEAN DEFAULT TRUE
) RETURNS TEXT
LANGUAGE sql
AS $$
    SELECT @extschema@.steadytext_prompt_render($1, $2, $3, $4);
$$;

CREATE OR REPLACE FUNCTION st_prompt_list()
RETURNS TABLE(
    prompt_id UUID,
    slug TEXT,
    description TEXT,
    latest_version_num INTEGER,
    total_versions INTEGER,
    created_at TIMESTAMPTZ,
    updated_at TIMESTAMPTZ
)
LANGUAGE sql
AS $$
    SELECT * FROM @extschema@.steadytext_prompt_list();
$$;

CREATE OR REPLACE FUNCTION st_prompt_versions(slug TEXT)
RETURNS TABLE(
    version_num INTEGER,
    template TEXT,
    required_variables TEXT[],
    metadata JSONB,
    created_at TIMESTAMPTZ,
    created_by TEXT,
    is_active BOOLEAN
)
LANGUAGE sql
AS $$
    SELECT * FROM @extschema@.steadytext_prompt_versions($1);
$$;

CREATE OR REPLACE FUNCTION st_prompt_delete(slug TEXT)
RETURNS BOOLEAN
LANGUAGE sql
AS $$
    SELECT @extschema@.steadytext_prompt_delete($1);
$$;

-- AIDEV-NOTE: Prompt registry feature added in v2025.9.6
-- Provides lightweight Jinja2-based template management with versioning
-- All functions have st_* aliases for convenience
