From c543b6cf482c7659e3f3e18ed3f76998c35cccde Mon Sep 17 00:00:00 2001 From: Adityavardhan Agrawal Date: Tue, 27 Jan 2026 10:37:58 -0800 Subject: [PATCH] fix set statement in probes --- core/vector_store/chunk_v2_store.py | 6 ++++-- core/vector_store/pgvector_store.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/core/vector_store/chunk_v2_store.py b/core/vector_store/chunk_v2_store.py index 2d8990cc..63366d3a 100644 --- a/core/vector_store/chunk_v2_store.py +++ b/core/vector_store/chunk_v2_store.py @@ -399,8 +399,10 @@ async def query_similar( ) async with self.get_session_with_retry() as session: - # PostgreSQL SET doesn't support parameterized values; safe since ivfflat_probes is a validated int - await session.execute(text(f"SET LOCAL ivfflat.probes = {self.ivfflat_probes}")) + # Use set_config() for parameterized config - SET doesn't support parameters + await session.execute( + text("SELECT set_config('ivfflat.probes', :probes, true)"), {"probes": str(self.ivfflat_probes)} + ) result = await session.execute(query) rows = result.all() diff --git a/core/vector_store/pgvector_store.py b/core/vector_store/pgvector_store.py index 2f86e4d2..66db9e09 100644 --- a/core/vector_store/pgvector_store.py +++ b/core/vector_store/pgvector_store.py @@ -458,8 +458,10 @@ async def query_similar( """ try: async with self.get_session_with_retry() as session: - # PostgreSQL SET doesn't support parameterized values; safe since ivfflat_probes is a validated int - await session.execute(text(f"SET LOCAL ivfflat.probes = {self.ivfflat_probes}")) + # Use set_config() for parameterized config - SET doesn't support parameters + await session.execute( + text("SELECT set_config('ivfflat.probes', :probes, true)"), {"probes": str(self.ivfflat_probes)} + ) # Build query with cosine distance calculation, which is normalized to [0, 2]. # A distance of 0 is perfect similarity. distance = VectorEmbedding.embedding.op("<=>")(query_embedding)