166 lines
4.3 KiB
Python
166 lines
4.3 KiB
Python
"""
|
|
Profile slow database queries to identify optimization opportunities.
|
|
|
|
Usage:
|
|
uv run python scripts/profile_queries.py
|
|
"""
|
|
import asyncio
|
|
import os
|
|
import asyncpg
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
|
|
async def profile_entry_points_query():
|
|
"""Profile the vector similarity entry points query."""
|
|
db_url = os.getenv("DATABASE_URL")
|
|
conn = await asyncpg.connect(db_url)
|
|
|
|
# Generate a dummy embedding vector (384 dimensions for bge-small-en-v1.5)
|
|
dummy_embedding = str([0.1] * 384)
|
|
|
|
print("=" * 80)
|
|
print("PROFILING: Entry Points Query (Vector Similarity)")
|
|
print("=" * 80)
|
|
|
|
# Run EXPLAIN ANALYZE
|
|
explain = await conn.fetch("""
|
|
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
|
|
SELECT id, text, context, event_date, access_count, embedding,
|
|
1 - (embedding <=> $1::vector) AS similarity
|
|
FROM memory_units
|
|
WHERE agent_id = $2
|
|
AND embedding IS NOT NULL
|
|
AND (1 - (embedding <=> $1::vector)) >= 0.5
|
|
ORDER BY embedding <=> $1::vector
|
|
LIMIT 3
|
|
""", dummy_embedding, "test_agent")
|
|
|
|
for row in explain:
|
|
print(row[0])
|
|
|
|
await conn.close()
|
|
|
|
|
|
async def profile_neighbors_query(sample_node_ids):
|
|
"""Profile the neighbors JOIN query."""
|
|
db_url = os.getenv("DATABASE_URL")
|
|
conn = await asyncpg.connect(db_url)
|
|
|
|
print("\n" + "=" * 80)
|
|
print("PROFILING: Neighbors Query (Graph Traversal)")
|
|
print(f"Sample size: {len(sample_node_ids)} nodes")
|
|
print("=" * 80)
|
|
|
|
# Run EXPLAIN ANALYZE
|
|
explain = await conn.fetch("""
|
|
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
|
|
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
|
|
mu.text, mu.context, mu.event_date, mu.access_count,
|
|
mu.id as neighbor_id
|
|
FROM memory_links ml
|
|
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
|
WHERE ml.from_unit_id = ANY($1::uuid[])
|
|
AND ml.weight >= 0.1
|
|
ORDER BY ml.from_unit_id, ml.weight DESC
|
|
""", sample_node_ids)
|
|
|
|
for row in explain:
|
|
print(row[0])
|
|
|
|
await conn.close()
|
|
|
|
|
|
async def profile_embeddings_query(sample_node_ids):
|
|
"""Profile the batch embeddings fetch query."""
|
|
db_url = os.getenv("DATABASE_URL")
|
|
conn = await asyncpg.connect(db_url)
|
|
|
|
print("\n" + "=" * 80)
|
|
print("PROFILING: Embeddings Query (Batch Fetch)")
|
|
print(f"Sample size: {len(sample_node_ids)} nodes")
|
|
print("=" * 80)
|
|
|
|
# Run EXPLAIN ANALYZE
|
|
explain = await conn.fetch("""
|
|
EXPLAIN (ANALYZE, BUFFERS, VERBOSE)
|
|
SELECT id, embedding
|
|
FROM memory_units
|
|
WHERE id = ANY($1::uuid[])
|
|
""", sample_node_ids)
|
|
|
|
for row in explain:
|
|
print(row[0])
|
|
|
|
await conn.close()
|
|
|
|
|
|
async def get_sample_node_ids(batch_size=50):
|
|
"""Get sample node IDs for profiling."""
|
|
db_url = os.getenv("DATABASE_URL")
|
|
conn = await asyncpg.connect(db_url)
|
|
|
|
rows = await conn.fetch(f"""
|
|
SELECT id FROM memory_units
|
|
LIMIT {batch_size}
|
|
""")
|
|
|
|
await conn.close()
|
|
return [row['id'] for row in rows]
|
|
|
|
|
|
async def check_indexes():
|
|
"""Check what indexes exist."""
|
|
db_url = os.getenv("DATABASE_URL")
|
|
conn = await asyncpg.connect(db_url)
|
|
|
|
print("\n" + "=" * 80)
|
|
print("CURRENT INDEXES")
|
|
print("=" * 80)
|
|
|
|
indexes = await conn.fetch("""
|
|
SELECT
|
|
tablename,
|
|
indexname,
|
|
indexdef
|
|
FROM pg_indexes
|
|
WHERE schemaname = 'public'
|
|
AND tablename IN ('memory_units', 'memory_links')
|
|
ORDER BY tablename, indexname
|
|
""")
|
|
|
|
for idx in indexes:
|
|
print(f"\nTable: {idx['tablename']}")
|
|
print(f"Index: {idx['indexname']}")
|
|
print(f"Definition: {idx['indexdef']}")
|
|
|
|
await conn.close()
|
|
|
|
|
|
async def main():
|
|
print("Starting Query Profiling...")
|
|
|
|
# Check indexes first
|
|
await check_indexes()
|
|
|
|
# Get sample node IDs
|
|
sample_ids = await get_sample_node_ids(50)
|
|
|
|
if sample_ids:
|
|
print(f"\nGot {len(sample_ids)} sample node IDs for profiling")
|
|
|
|
# Profile each query type
|
|
await profile_entry_points_query()
|
|
await profile_neighbors_query(sample_ids)
|
|
await profile_embeddings_query(sample_ids)
|
|
else:
|
|
print("\nNo data in database - run ingestion first")
|
|
|
|
print("\n" + "=" * 80)
|
|
print("PROFILING COMPLETE")
|
|
print("=" * 80)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|