debug_rerank_binding.py

import os
import sqlalchemy
from sqlalchemy import text
from google.cloud.alloydb.connector import Connector, IPTypes

PROJECT_ID = "alloydbtest-374215"
REGION = "us-central1"
CLUSTER = "hybridsearch"
INSTANCE = "hybridsearch-primary"
DATABASE = "ecom"
DB_USER = "postgres"
DB_PASS = os.environ.get("DB_PASSWORD", "Qbertq!1")

def debug_binding():
    connector = Connector()
    pool = sqlalchemy.create_engine(
        "postgresql+pg8000://",
        creator=lambda: connector.connect(
            f"projects/{PROJECT_ID}/locations/{REGION}/clusters/{CLUSTER}/instances/{INSTANCE}",
            "pg8000",
            user=DB_USER,
            password=DB_PASS,
            db=DATABASE,
            ip_type=IPTypes.PUBLIC
        )
    )
    
    try:
        with pool.connect() as conn:
            print("Testing binding list to text[]...")
            query = text("""
                SELECT * FROM ai.rank(
                    model_id => 'semantic-ranker-512@latest',
                    search_string => :query,
                    documents => :docs
                )
            """)
            
            docs = ["red dress", "blue shoes"]
            q = "dress"
            
            # SQLAlchemy might need type hint or specific cast
            try:
                result = conn.execute(query, {"query": q, "docs": docs}).fetchall()
                print("Success!")
                for row in result:
                    print(row)
            except Exception as e:
                print(f"Binding failed: {e}")
                
            print("\nTesting CAST(:docs AS text[])...")
            query_cast = text("""
                 SELECT * FROM ai.rank(
                    model_id => 'semantic-ranker-512@latest',
                    search_string => :query,
                    documents => CAST(:docs AS text[])
                )
            """)
            try:
                result = conn.execute(query_cast, {"query": q, "docs": docs}).fetchall()
                print("Success with CAST!")
                for row in result:
                     print(row)
            except Exception as e:
                print(f"CAST also failed: {e}")

    except Exception as e:
        print(f"Connection Error: {e}")

if __name__ == "__main__":
    debug_binding()