fix: run linter and format

Author: Restodecoca

llama-index-integrations/vector_stores/llama-index-vector-store-paradedb/llama_index/vector_stores/paradedb/base.py

@@ -5,7 +5,6 @@
 
 import sqlalchemy
 from llama_index.core.bridge.pydantic import BaseModel, Field
-from llama_index.core.vector_stores.types import VectorStoreQuery
 from sqlalchemy.sql.selectable import Select
 
 from llama_index.vector_stores.postgres.base import (
@@ -36,7 +35,17 @@ def get_bm25_data_model(
     from pgvector.sqlalchemy import Vector, HALFVEC
     from sqlalchemy import Column
     from sqlalchemy.dialects.postgresql import BIGINT, JSON, JSONB, VARCHAR
-    from sqlalchemy import cast, column, String, Integer, Numeric, Float, Boolean, Date, DateTime
+    from sqlalchemy import (
+        cast,
+        column,
+        String,
+        Integer,
+        Numeric,
+        Float,
+        Boolean,
+        Date,
+        DateTime,
+    )
     from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, UUID
     from sqlalchemy.schema import Index
 
@@ -54,7 +63,7 @@ def get_bm25_data_model(
     }
 
     indexed_metadata_keys = indexed_metadata_keys or set()
-    
+
     for key, pg_type in indexed_metadata_keys:
         if pg_type not in pg_type_map:
             raise ValueError(
@@ -67,7 +76,9 @@ def get_bm25_data_model(
     indexname = f"{index_name}_idx"
 
     metadata_dtype = JSONB if use_jsonb else JSON
-    embedding_col = Column(HALFVEC(embed_dim)) if use_halfvec else Column(Vector(embed_dim))
+    embedding_col = (
+        Column(HALFVEC(embed_dim)) if use_halfvec else Column(Vector(embed_dim))
+    )
 
     metadata_indices = [
         Index(
@@ -107,7 +118,7 @@ class BM25AbstractData(base):
 class ParadeDBVectorStore(PGVectorStore, BaseModel):
     """
     ParadeDB Vector Store with BM25 search support.
-    
+
     Inherits from PGVectorStore and adds BM25 full-text search capabilities
     using ParadeDB's pg_search extension.
 
@@ -130,16 +141,19 @@ class ParadeDBVectorStore(PGVectorStore, BaseModel):
             use_halfvec=True
         )
         ```
+
     """
 
     connection_string: Optional[Union[str, sqlalchemy.engine.URL]] = Field(default=None)
-    async_connection_string: Optional[Union[str, sqlalchemy.engine.URL]] = Field(default=None)
+    async_connection_string: Optional[Union[str, sqlalchemy.engine.URL]] = Field(
+        default=None
+    )
     table_name: Optional[str] = Field(default=None)
     schema_name: Optional[str] = Field(default="paradedb")
     hybrid_search: bool = Field(default=False)
     text_search_config: str = Field(default="english")
     embed_dim: int = Field(default=1536)
-    cache_ok: bool = Field(default=False) 
+    cache_ok: bool = Field(default=False)
     perform_setup: bool = Field(default=True)
     debug: bool = Field(default=False)
     use_jsonb: bool = Field(default=False)
@@ -154,7 +168,7 @@ def __init__(
         table_name: Optional[str] = None,
         schema_name: Optional[str] = None,
         hybrid_search: bool = False,
-        text_search_config: str = "english", 
+        text_search_config: str = "english",
         embed_dim: int = 1536,
         cache_ok: bool = False,
         perform_setup: bool = True,
@@ -176,7 +190,7 @@ def __init__(
             self,
             connection_string=connection_string,
             async_connection_string=async_connection_string,
-            table_name=table_name, 
+            table_name=table_name,
             schema_name=schema_name or "paradedb",
             hybrid_search=hybrid_search,
             text_search_config=text_search_config,
@@ -187,14 +201,16 @@ def __init__(
             use_jsonb=use_jsonb,
             hnsw_kwargs=hnsw_kwargs,
             create_engine_kwargs=create_engine_kwargs,
-            use_bm25=use_bm25
+            use_bm25=use_bm25,
         )
-        
+
         # Call parent constructor
         PGVectorStore.__init__(
             self,
             connection_string=str(connection_string) if connection_string else None,
-            async_connection_string=str(async_connection_string) if async_connection_string else None,
+            async_connection_string=str(async_connection_string)
+            if async_connection_string
+            else None,
             table_name=table_name,
             schema_name=self.schema_name,
             hybrid_search=hybrid_search,
@@ -213,10 +229,11 @@ def __init__(
             indexed_metadata_keys=indexed_metadata_keys,
             customize_query_fn=customize_query_fn,
         )
-        
+
         # Override table model if using BM25
         if self.use_bm25:
             from sqlalchemy.orm import declarative_base
+
             self._base = declarative_base()
             self._table_class = get_bm25_data_model(
                 self._base,
@@ -270,6 +287,7 @@ def from_params(
 
         Returns:
             ParadeDBVectorStore: Instance of ParadeDBVectorStore.
+
         """
         conn_str = (
             connection_string
@@ -301,7 +319,7 @@ def from_params(
     def _create_extension(self) -> None:
         """Override to add pg_search extension for BM25."""
         super()._create_extension()
-        
+
         if self.use_bm25:
             with self._session() as session, session.begin():
                 try:
@@ -337,7 +355,7 @@ def _initialize(self) -> None:
         """Override to add BM25 index creation."""
         if not self._is_initialized:
             super()._initialize()
-            
+
             if self.use_bm25 and self.perform_setup:
                 try:
                     self._create_bm25_index()
@@ -355,10 +373,12 @@ def _build_sparse_query(
     ) -> Any:
         """Override to use BM25 if enabled, otherwise use parent's ts_vector."""
         if not self.use_bm25:
-            return super()._build_sparse_query(query_str, limit, metadata_filters, **kwargs)
-        
+            return super()._build_sparse_query(
+                query_str, limit, metadata_filters, **kwargs
+            )
+
         from sqlalchemy import text
-        
+
         if query_str is None:
             raise ValueError("query_str must be specified for a sparse vector query.")
 
@@ -373,14 +393,12 @@ def _build_sparse_query(
         if metadata_filters:
             _logger.warning("Metadata filters not fully implemented for BM25 raw SQL")
 
-        stmt = text(f"""
+        return text(f"""
             {base_query}
             ORDER BY rank DESC
             LIMIT :limit
         """).bindparams(query=query_str_clean, limit=limit)
 
-        return stmt
-
     def _sparse_query_with_rank(
         self,
         query_str: Optional[str] = None,
@@ -390,7 +408,7 @@ def _sparse_query_with_rank(
         """Override to handle BM25 results properly."""
         if not self.use_bm25:
             return super()._sparse_query_with_rank(query_str, limit, metadata_filters)
-        
+
         stmt = self._build_sparse_query(query_str, limit, metadata_filters)
         with self._session() as session, session.begin():
             res = session.execute(stmt)
@@ -417,8 +435,10 @@ async def _async_sparse_query_with_rank(
     ) -> List[DBEmbeddingRow]:
         """Override to handle async BM25 results properly."""
         if not self.use_bm25:
-            return await super()._async_sparse_query_with_rank(query_str, limit, metadata_filters)
-        
+            return await super()._async_sparse_query_with_rank(
+                query_str, limit, metadata_filters
+            )
+
         stmt = self._build_sparse_query(query_str, limit, metadata_filters)
         async with self._async_session() as session, session.begin():
             res = await session.execute(stmt)
@@ -435,4 +455,4 @@ async def _async_sparse_query_with_rank(
                     similarity=item.rank,
                 )
                 for item in res.all()
-            ]
+            ]

llama-index-integrations/vector_stores/llama-index-vector-store-paradedb/tests/test_paradedb.py

@@ -49,6 +49,7 @@ def _get_sample_vector(num: float) -> List[float]:
 @pytest.fixture(scope="session")
 def conn() -> Any:
     import psycopg2
+
     return psycopg2.connect(**PARAMS)  # type: ignore
 
 
@@ -434,24 +435,28 @@ async def test_bm25_extensions_created(db: None) -> None:
         hybrid_search=True,
         embed_dim=TEST_EMBED_DIM,
     )
-    
+
     # Force initialization
-    pg.add([
-        TextNode(
-            text="test",
-            id_="test",
-            embedding=_get_sample_vector(1.0),
-        )
-    ])
-    
+    pg.add(
+        [
+            TextNode(
+                text="test",
+                id_="test",
+                embedding=_get_sample_vector(1.0),
+            )
+        ]
+    )
+
     # Check that both extensions exist
     with psycopg2.connect(**PARAMS, database=TEST_DB) as conn:
         with conn.cursor() as c:
-            c.execute("SELECT COUNT(*) FROM pg_extension WHERE extname IN ('vector', 'pg_search');")
+            c.execute(
+                "SELECT COUNT(*) FROM pg_extension WHERE extname IN ('vector', 'pg_search');"
+            )
             ext_count = c.fetchone()[0]
-    
+
     assert ext_count == 2, "Both 'vector' and 'pg_search' extensions should exist"
-    
+
     await pg.close()
 
 
@@ -464,39 +469,37 @@ async def test_paradedb_inherits_pgvector_functionality(
     """Test that ParadeDBVectorStore inherits all PGVectorStore functionality."""
     # Add nodes
     pg_bm25.add(hybrid_node_embeddings)
-    
+
     # Test vector-only query (inherited from PGVectorStore)
     q = VectorStoreQuery(
         query_embedding=_get_sample_vector(0.1),
         similarity_top_k=2,
         mode=VectorStoreQueryMode.DEFAULT,
     )
-    
+
     res = pg_bm25.query(q)
     assert res.nodes
     assert len(res.nodes) == 2
-    
+
     # Test delete (inherited)
     pg_bm25.delete_nodes(["aaa"])
-    
+
     res = pg_bm25.query(q)
     assert "aaa" not in res.ids
-    
+
     # Test clear (inherited)
     await pg_bm25.aclear()
-    
+
     res = pg_bm25.query(q)
     assert len(res.nodes) == 0
 
 
 @pytest.mark.skipif(postgres_not_available, reason="postgres db is not available")
 @pytest.mark.asyncio
 async def test_bm25_vs_tsvector_different_results(
-    db: None,
-    hybrid_node_embeddings: List[TextNode]
-    ) -> None:
+    db: None, hybrid_node_embeddings: List[TextNode]
+) -> None:
     """Test that BM25 and ts_vector can produce different ranking results."""
-    
     # Create both stores
     pg_tsvector = PGVectorStore.from_params(
         **PARAMS,  # type: ignore
@@ -506,7 +509,7 @@ async def test_bm25_vs_tsvector_different_results(
         hybrid_search=True,
         embed_dim=TEST_EMBED_DIM,
     )
-    
+
     pg_bm25 = ParadeDBVectorStore.from_params(
         **PARAMS,  # type: ignore
         database=TEST_DB,
@@ -518,14 +521,14 @@ async def test_bm25_vs_tsvector_different_results(
     )
     pg_tsvector.add(hybrid_node_embeddings)
     pg_bm25.add(hybrid_node_embeddings)
-    
+
     q = VectorStoreQuery(
         query_str="fox",
         sparse_top_k=2,
         mode=VectorStoreQueryMode.SPARSE,
         query_embedding=_get_sample_vector(5.0),
     )
-    
+
     res_tsvector = pg_tsvector.query(q)
     res_bm25 = pg_bm25.query(q)
 
@@ -538,11 +541,11 @@ async def test_bm25_vs_tsvector_different_results(
     # Both should return results
     assert len(res_tsvector.nodes) == 2
     assert len(res_bm25.nodes) == 2
-    
+
     # BM25 uses BM25 ranking, ts_vector uses ts_rank
     # The implementation difference is verified
     assert pg_bm25.use_bm25 is True
     assert not hasattr(pg_tsvector, "use_bm25") or pg_tsvector.use_bm25 is False
-    
+
     await pg_tsvector.close()
-    await pg_bm25.close()
+    await pg_bm25.close()