114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
"""
|
|
Quick test script to debug vector query issue.
|
|
Run with: docker exec code-forge-app python3 /app/test_vector_query.py
|
|
"""
|
|
|
|
import os
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
|
|
def main():
|
|
print("=" * 60)
|
|
print("PostgreSQL Vector Query Test")
|
|
print("=" * 60)
|
|
|
|
# Connect to PostgreSQL
|
|
conn = psycopg2.connect(
|
|
host=os.getenv("POSTGRES_HOST", "host.docker.internal"),
|
|
port=int(os.getenv("POSTGRES_PORT", "5432")),
|
|
database=os.getenv("POSTGRES_DB", "codeforge"),
|
|
user=os.getenv("POSTGRES_USER", os.getenv("USER", "postgres")),
|
|
)
|
|
|
|
try:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
# Get a sample embedding from the database
|
|
print("\n📊 Getting sample embedding from database...")
|
|
cur.execute("SELECT embedding FROM documents WHERE kb_id = 2 LIMIT 1")
|
|
result = cur.fetchone()
|
|
if not result:
|
|
print("❌ No documents found in kb_id=2")
|
|
return
|
|
|
|
sample_embedding = result['embedding']
|
|
print(f"✅ Got sample embedding")
|
|
print(f" Type: {type(sample_embedding)}")
|
|
print(f" Length: {len(sample_embedding) if hasattr(sample_embedding, '__len__') else 'N/A'}")
|
|
|
|
# Test 1: Query using the sample embedding directly (as returned from DB)
|
|
print("\n🧪 Test 1: Query using embedding AS-IS from database")
|
|
try:
|
|
cur.execute(
|
|
"""
|
|
SELECT content, 1 - (embedding <=> %s) as similarity
|
|
FROM documents
|
|
WHERE kb_id = 2
|
|
ORDER BY embedding <=> %s
|
|
LIMIT 5
|
|
""",
|
|
(sample_embedding, sample_embedding)
|
|
)
|
|
results = cur.fetchall()
|
|
print(f" ✅ Results: {len(results)}")
|
|
if results:
|
|
print(f" Top similarity: {results[0]['similarity']:.4f}")
|
|
print(f" Content preview: {results[0]['content'][:80]}...")
|
|
except Exception as e:
|
|
print(f" ❌ Error: {e}")
|
|
|
|
# Test 2: Convert to list and back to string
|
|
print("\n🧪 Test 2: Query using string-formatted embedding with ::vector cast")
|
|
try:
|
|
embedding_list = list(sample_embedding)
|
|
embedding_str = "[" + ",".join(map(str, embedding_list)) + "]"
|
|
print(f" Embedding string (first 100 chars): {embedding_str[:100]}...")
|
|
|
|
cur.execute(
|
|
"""
|
|
SELECT content, 1 - (embedding <=> %s::vector) as similarity
|
|
FROM documents
|
|
WHERE kb_id = 2
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT 5
|
|
""",
|
|
(embedding_str, embedding_str)
|
|
)
|
|
results = cur.fetchall()
|
|
print(f" Results: {len(results)}")
|
|
if results:
|
|
print(f" ✅ Top similarity: {results[0]['similarity']:.4f}")
|
|
else:
|
|
print(" ❌ NO RESULTS - This is the problem!")
|
|
except Exception as e:
|
|
print(f" ❌ Error: {e}")
|
|
|
|
# Test 3: Try without ::vector cast but with string
|
|
print("\n🧪 Test 3: Query using string WITHOUT ::vector cast")
|
|
try:
|
|
cur.execute(
|
|
"""
|
|
SELECT content, 1 - (embedding <=> %s) as similarity
|
|
FROM documents
|
|
WHERE kb_id = 2
|
|
ORDER BY embedding <=> %s
|
|
LIMIT 5
|
|
""",
|
|
(embedding_str, embedding_str)
|
|
)
|
|
results = cur.fetchall()
|
|
print(f" Results: {len(results)}")
|
|
if results:
|
|
print(f" ✅ Top similarity: {results[0]['similarity']:.4f}")
|
|
except Exception as e:
|
|
print(f" ❌ Error: {e}")
|
|
|
|
finally:
|
|
conn.close()
|
|
|
|
print("\n" + "=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|