Aquery ivfflat fix

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/llama_index/vector_stores/postgres/base.py

@@ -647,12 +647,16 @@ def _query_with_score(
         stmt = self._build_query(embedding, limit, metadata_filters)
         with self._session() as session, session.begin():
             from sqlalchemy import text
+            from psycopg2 import sql
 
             if kwargs.get("ivfflat_probes"):
                 ivfflat_probes = kwargs.get("ivfflat_probes")
                 session.execute(
-                    text(f"SET ivfflat.probes = :ivfflat_probes"),
-                    {"ivfflat_probes": ivfflat_probes},
+                    text(
+                        sql.SQL("SET ivfflat.probes = {}")
+                        .format(sql.Literal(ivfflat_probes))
+                        .as_string(context=self._engine.raw_connection().connection),
+                    )
                 )
             if self.hnsw_kwargs:
                 hnsw_ef_search = (
@@ -686,6 +690,7 @@ async def _aquery_with_score(
         stmt = self._build_query(embedding, limit, metadata_filters)
         async with self._async_session() as async_session, async_session.begin():
             from sqlalchemy import text
+            from psycopg2 import sql
 
             if self.hnsw_kwargs:
                 hnsw_ef_search = (
@@ -697,8 +702,11 @@ async def _aquery_with_score(
             if kwargs.get("ivfflat_probes"):
                 ivfflat_probes = kwargs.get("ivfflat_probes")
                 await async_session.execute(
-                    text(f"SET ivfflat.probes = :ivfflat_probes"),
-                    {"ivfflat_probes": ivfflat_probes},
+                    text(
+                        sql.SQL("SET ivfflat.probes = {}")
+                        .format(sql.Literal(ivfflat_probes))
+                        .as_string(context=self._engine.raw_connection().connection),
+                    )
                 )
 
             res = await async_session.execute(stmt)

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/pyproject.toml

@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-vector-stores-postgres"
 readme = "README.md"
-version = "0.4.2"
+version = "0.4.3"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"

llama-index-integrations/vector_stores/llama-index-vector-stores-postgres/tests/test_postgres.py

@@ -2,13 +2,16 @@
 from typing import Any, Dict, Generator, List, Union, Optional
 
 import pytest
+from llama_index.core import VectorStoreIndex
 from llama_index.core.schema import (
     BaseNode,
     IndexNode,
     NodeRelationship,
     RelatedNodeInfo,
     TextNode,
 )
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core.response_synthesizers import get_response_synthesizer
 from llama_index.core.vector_stores.types import (
     ExactMatchFilter,
     FilterOperator,
@@ -968,6 +971,34 @@ async def test_delete_nodes_metadata(
     assert "ccc" in res.ids
 
 
+@pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
+@pytest.mark.asyncio()
+@pytest.mark.parametrize("pg_fixture", ["pg", "pg_hybrid"], indirect=True)
+@pytest.mark.parametrize("use_async", [True, False])
+async def test_set_ivfflat(pg_fixture: PGVectorStore, use_async: bool) -> None:
+    vector_index = VectorStoreIndex.from_vector_store(vector_store=pg_fixture)
+
+    vector_retriever = vector_index.as_retriever(
+        vector_store_kwargs={"ivfflat_probes": 20}
+    )
+    vector_response_synthesizer = get_response_synthesizer(
+        use_async=use_async,
+    )
+
+    vector_query_engine = RetrieverQueryEngine(
+        retriever=vector_retriever,
+        response_synthesizer=vector_response_synthesizer,
+    )
+
+    query = "lorem ipsum"
+    if use_async:
+        response = await vector_query_engine.aquery(query)
+    else:
+        response = vector_query_engine.query(query)
+
+    assert response
+
+
 @pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
 @pytest.mark.asyncio()
 @pytest.mark.parametrize("use_async", [True, False])