ragbits-core 0.16.0__py3-none-any.whl → 1.4.0.dev202512021005__py3-none-any.whl

ragbits/core/vector_stores/in_memory.py

@@ -1,10 +1,9 @@
+import math
 from itertools import islice
 from typing import cast
 from uuid import UUID
 
-import numpy as np
-
-from ragbits.core.audit import trace, traceable
+from ragbits.core.audit.traces import trace, traceable
 from ragbits.core.embeddings import Embedder, SparseVector
 from ragbits.core.vector_stores.base import (
     EmbeddingType,
@@ -90,6 +89,14 @@ class InMemoryVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
             results: list[VectorStoreResult] = []
 
             for entry_id, vector in self._embeddings.items():
+                entry = self._entries[entry_id]
+
+                # Apply metadata filtering
+                if merged_options.where and not all(
+                    entry.metadata.get(key) == value for key, value in merged_options.where.items()
+                ):
+                    continue
+
                 # Calculate score based on vector type
                 if isinstance(query_vector, SparseVector) and isinstance(vector, SparseVector):
                     # For sparse vectors, use dot product between query and document vectors
@@ -105,7 +112,9 @@ class InMemoryVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
                     # For dense vectors, use negative L2 distance
                     query_vector_dense = cast(list[float], query_vector)
                     vector_dense = cast(list[float], vector)
-                    score = float(np.linalg.norm(np.array(vector_dense) - np.array(query_vector_dense))) * -1
+                    score = -math.sqrt(
+                        sum((a - b) ** 2 for a, b in zip(vector_dense, query_vector_dense, strict=False))
+                    )
 
                 result = VectorStoreResult(entry=self._entries[entry_id], vector=vector, score=score)
                 if merged_options.score_threshold is None or result.score >= merged_options.score_threshold:

ragbits/core/vector_stores/pgvector.py

@@ -1,13 +1,13 @@
 import json
 import re
-from typing import Any, NamedTuple, cast
+from typing import Any, NamedTuple
 from uuid import UUID
 
 import asyncpg
 from pydantic.json import pydantic_encoder
 
-from ragbits.core.audit import trace
-from ragbits.core.embeddings.base import Embedder, SparseVector
+from ragbits.core.audit.traces import trace
+from ragbits.core.embeddings.base import Embedder, SparseVector, VectorSize
 from ragbits.core.embeddings.sparse.base import SparseEmbedder
 from ragbits.core.vector_stores.base import (
     EmbeddingType,
@@ -33,14 +33,16 @@ class DistanceOp(NamedTuple):
 DISTANCE_OPS = {
     "cosine": DistanceOp("vector_cosine_ops", "<=>", "1 - distance"),
     "l2": DistanceOp("vector_l2_ops", "<->", "distance * -1"),
+    "halfvec_l2": DistanceOp("halfvec_l2_ops", "<->", "distance * -1"),
     "l1": DistanceOp("vector_l1_ops", "<+>", "distance * -1"),
     "ip": DistanceOp("vector_ip_ops", "<#>", "distance * -1"),
     "bit_hamming": DistanceOp("bit_hamming_ops", "<~>", "distance * -1"),
     "bit_jaccard": DistanceOp("bit_jaccard_ops", "<%>", "distance * -1"),
     "sparsevec_l2": DistanceOp("sparsevec_l2_ops", "<->", "distance * -1"),
-    "halfvec_l2": DistanceOp("halfvec_l2_ops", "<->", "distance * -1"),
 }
 
+MAX_VECTOR_SIZE = 2000
+
 
 class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
     """
@@ -53,11 +55,12 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         self,
         client: asyncpg.Pool,
         table_name: str,
-        vector_size: int,
         embedder: Embedder,
+        vector_size: int | None = None,
         embedding_type: EmbeddingType = EmbeddingType.TEXT,
         distance_method: str | None = None,
-        hnsw_params: dict | None = None,
+        is_hnsw: bool = True,
+        params: dict | None = None,
         default_options: VectorStoreOptions | None = None,
     ) -> None:
         """
@@ -66,12 +69,13 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         Args:
             client: The pgVector database connection pool.
             table_name: The name of the table.
-            vector_size: The size of the vectors.
             embedder: The embedder to use for converting entries to vectors.
+            vector_size: The size of the vectors. If None, will be determined automatically from the embedder.
             embedding_type: Which part of the entry to embed, either text or image. The other part will be ignored.
             distance_method: The distance method to use, default is "cosine" for dense vectors
                 and "sparsevec_l2" for sparse vectors.
-            hnsw_params: The parameters for the HNSW index. If None, the default parameters will be used.
+            is_hnsw: if hnsw or ivfflat indexing should be used
+            params: The parameters for the HNSW index. If None, the default parameters will be used.
             default_options: The default options for querying the vector store.
         """
         (
@@ -84,27 +88,34 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
 
         if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table_name):
             raise ValueError(f"Invalid table name: {table_name}")
-        if not isinstance(vector_size, int) or vector_size <= 0:
+        if vector_size is not None and (not isinstance(vector_size, int) or vector_size <= 0):
             raise ValueError("Vector size must be a positive integer.")
 
-        if hnsw_params is None:
-            hnsw_params = {"m": 4, "ef_construction": 10}
-        elif not isinstance(hnsw_params, dict):
-            raise ValueError("hnsw_params must be a dictionary.")
-        elif "m" not in hnsw_params or "ef_construction" not in hnsw_params:
-            raise ValueError("hnsw_params must contain 'm' and 'ef_construction' keys.")
-        elif not isinstance(hnsw_params["m"], int) or hnsw_params["m"] <= 0:
-            raise ValueError("m must be a positive integer.")
-        elif not isinstance(hnsw_params["ef_construction"], int) or hnsw_params["ef_construction"] <= 0:
-            raise ValueError("ef_construction must be a positive integer.")
+        if params is None and is_hnsw:
+            params = {"m": 4, "ef_construction": 10}
+        elif params is None and not is_hnsw:
+            params = {"lists": 100}
+        elif not isinstance(params, dict):
+            raise ValueError("params must be a dictionary.")
+        elif "m" not in params or "ef_construction" not in params and is_hnsw:
+            raise ValueError("params must contain 'm' and 'ef_construction' keys for hnsw indexing.")
+        elif not isinstance(params["m"], int) or params["m"] <= 0 and is_hnsw:
+            raise ValueError("m must be a positive integer for hnsw indexing.")
+        elif not isinstance(params["ef_construction"], int) or params["ef_construction"] <= 0 and is_hnsw:
+            raise ValueError("ef_construction must be a positive integer for hnsw indexing.")
+        elif "lists" not in params and not is_hnsw:
+            raise ValueError("params must contain 'lists' key for IVFFlat indexing.")
+        elif not isinstance(params["lists"], int) or params["lists"] <= 0 and not is_hnsw:
+            raise ValueError("lists must be a positive integer for IVFFlat indexing.")
 
         if distance_method is None:
             distance_method = "sparsevec_l2" if isinstance(embedder, SparseEmbedder) else "cosine"
         self._client = client
         self._table_name = table_name
         self._vector_size = vector_size
+        self._vector_size_info: VectorSize | None = None
         self._distance_method = distance_method
-        self._hnsw_params = hnsw_params
+        self._indexing_params = params
 
     def __reduce__(self) -> tuple:
         """
@@ -113,6 +124,32 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         # TODO: To be implemented. Required for Ray processing.
         raise NotImplementedError
 
+    async def _get_vector_size_info(self) -> VectorSize:
+        """
+        Get vector size information from the embedder if not already cached.
+
+        Returns:
+            VectorSize information including size and sparsity.
+        """
+        if self._vector_size_info is None:
+            self._vector_size_info = await self._embedder.get_vector_size()
+            # Update _vector_size for backward compatibility if it wasn't provided
+            if self._vector_size is None:
+                self._vector_size = self._vector_size_info.size
+        return self._vector_size_info
+
+    async def _get_vector_size(self) -> int:
+        """
+        Get the vector size, either from the constructor parameter or from the embedder.
+
+        Returns:
+            The vector size as an integer.
+        """
+        if self._vector_size is not None:
+            return self._vector_size
+        vector_size_info = await self._get_vector_size_info()
+        return vector_size_info.size
+
     def _vector_to_string(self, vector: list[float] | SparseVector) -> str:
         """
         Converts a vector to a string representation.
@@ -124,8 +161,13 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
             str: The string representation of the vector.
         """
         if isinstance(vector, SparseVector):
+            # For sparse vectors, we need the vector size to be available
+            # This will be resolved when this method is called from async context
+            vector_size = self._vector_size
+            if vector_size is None:
+                raise RuntimeError("Vector size must be determined before converting sparse vectors to string")
             points_str = ",".join(f"{i}:{v}" for i, v in zip(vector.indices, vector.values, strict=False))
-            return f"{{{points_str}}}/{self._vector_size}"
+            return f"{{{points_str}}}/{vector_size}"
         return json.dumps(vector)
 
     @staticmethod
@@ -173,13 +215,19 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         # _table_name has been validated in the class constructor, and it is a valid table name.
         query = f"SELECT *, vector {distance_operator} $1 as distance, {score_formula} as score FROM {self._table_name}"  # noqa S608
 
-        values: list[Any] = [
-            self._vector_to_string(vector),
-        ]
+        values: list[Any] = [self._vector_to_string(vector)]
+        where_clauses = []
 
         if query_options.score_threshold is not None:
-            query += " WHERE score >= $2"
-            values.extend([query_options.score_threshold])
+            where_clauses.append("score >= $" + str(len(values) + 1))
+            values.append(query_options.score_threshold)
+
+        if query_options.where:
+            where_clauses.append(f"metadata @> ${len(values) + 1}")
+            values.append(json.dumps(query_options.where))
+
+        if where_clauses:
+            query += " WHERE " + " AND ".join(where_clauses)
 
         query += " ORDER BY distance"
 
@@ -226,32 +274,55 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
 
     async def create_table(self) -> None:
         """
-        Create a pgVector table with an HNSW index for given similarity.
+        Create a pgVector table with an HNSW/IVFFlat index for given similarity.
         """
+        vector_size = await self._get_vector_size()
+
         with trace(
             table_name=self._table_name,
             distance_method=self._distance_method,
-            vector_size=self._vector_size,
-            hnsw_index_parameters=self._hnsw_params,
+            vector_size=vector_size,
+            hnsw_index_parameters=self._indexing_params,
         ):
             distance = DISTANCE_OPS[self._distance_method].function_name
             create_vector_extension = "CREATE EXTENSION IF NOT EXISTS vector;"
             # _table_name and has been validated in the class constructor, and it is a valid table name.
-            # _vector_size has been validated in the class constructor, and it is a valid vector size.
+            # vector_size has been validated in the class constructor or obtained from embedder,
+            # and it is a valid vector size.
 
             is_sparse = isinstance(self._embedder, SparseEmbedder)
-            vector_func = "VECTOR" if not is_sparse else "SPARSEVEC"
+
+            # Check vector size
+            # if greater than 2000 then choose type HALFVEC
+            # More info: https://github.com/pgvector/pgvector
+            vector_func = (
+                "HALFVEC"
+                if vector_size > MAX_VECTOR_SIZE and re.search("halfvec", distance)
+                else "VECTOR"
+                if not is_sparse
+                else "SPARSEVEC"
+            )
 
             create_table_query = f"""
             CREATE TABLE {self._table_name}
-            (id UUID, text TEXT, image_bytes BYTEA, vector {vector_func}({self._vector_size}), metadata JSONB);
+            (id UUID, text TEXT, image_bytes BYTEA, vector {vector_func}({vector_size}), metadata JSONB);
             """
-            # _hnsw_params has been validated in the class constructor, and it is valid dict[str,int].
+            # _idexing_params has been validated in the class constructor, and it is valid dict[str,int].
+            if "lists" in self._indexing_params:
+                index_type = "ivfflat"
+                index_params = f"(lists = {self._indexing_params['lists']});"
+            else:
+                index_type = "hnsw"
+                index_params = (
+                    f"(m = {self._indexing_params['m']}, ef_construction = {self._indexing_params['ef_construction']});"
+                )
+
             create_index_query = f"""
-                    CREATE INDEX {self._table_name + "_hnsw_idx"} ON {self._table_name}
-                    USING hnsw (vector {distance})
-                    WITH (m = {self._hnsw_params["m"]}, ef_construction = {self._hnsw_params["ef_construction"]});
-                    """
+            CREATE INDEX {self._table_name + "_" + index_type + "_idx"} ON {self._table_name}
+            USING {index_type} (vector {distance})
+            WITH {index_params}
+            """
+
             if await self._check_table_exists():
                 print(f"Table {self._table_name} already exist!")
                 return
@@ -277,6 +348,10 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         """
         if not entries:
             return
+
+        # Ensure vector size is determined before processing
+        vector_size = await self._get_vector_size()
+
         # _table_name has been validated in the class constructor, and it is a valid table name.
         insert_query = f"""
         INSERT INTO {self._table_name} (id, text, image_bytes, vector, metadata)
@@ -285,7 +360,7 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         with trace(
             table_name=self._table_name,
             entries=entries,
-            vector_size=self._vector_size,
+            vector_size=vector_size,
             embedder=repr(self._embedder),
             embedding_type=self._embedding_type,
         ):
@@ -351,25 +426,26 @@ class PgVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         Returns:
             The retrieved entries.
         """
-        query_options = (self.default_options | options) if options else self.default_options
+        merged_options = (self.default_options | options) if options else self.default_options
+
+        # Ensure vector size is determined before processing
+        vector_size = await self._get_vector_size()
+
         with trace(
             text=text,
+            options=merged_options.dict(),
             table_name=self._table_name,
-            query_options=query_options,
-            vector_size=self._vector_size,
+            vector_size=vector_size,
             distance_method=self._distance_method,
             embedder=repr(self._embedder),
             embedding_type=self._embedding_type,
         ) as outputs:
-            vector = (await self._embedder.embed_text([text]))[0]
-            vector = cast(list[float], vector)
-
-            query_options = (self.default_options | options) if options else self.default_options
-            retrieve_query, values = self._create_retrieve_query(vector, query_options)
+            query_vector = (await self._embedder.embed_text([text]))[0]
+            query, values = self._create_retrieve_query(query_vector, merged_options)
 
             try:
                 async with self._client.acquire() as conn:
-                    results = await conn.fetch(retrieve_query, *values)
+                    results = await conn.fetch(query, *values)
 
                 outputs.results = [
                     VectorStoreResult(

ragbits/core/vector_stores/qdrant.py

@@ -16,7 +16,7 @@ from qdrant_client.models import (
 )
 from typing_extensions import Self
 
-from ragbits.core.audit import trace
+from ragbits.core.audit.traces import trace
 from ragbits.core.embeddings import Embedder, SparseEmbedder, SparseVector
 from ragbits.core.utils.config_handling import ObjectConstructionConfig, import_by_path
 from ragbits.core.utils.dict_transformations import flatten_dict
@@ -24,7 +24,6 @@ from ragbits.core.vector_stores.base import (
     EmbeddingType,
     VectorStoreEntry,
     VectorStoreOptions,
-    VectorStoreOptionsT,
     VectorStoreResult,
     VectorStoreWithEmbedder,
     WhereQuery,
@@ -202,7 +201,7 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
                 models.PointStruct(
                     id=str(entry.id),
                     vector={self._vector_name: self._to_qdrant_vector(embeddings[entry.id])},  # type: ignore
-                    payload=entry.model_dump(exclude_none=True),
+                    payload=entry.model_dump(exclude_none=True, mode="json"),
                 )
                 for entry in entries
                 if entry.id in embeddings
@@ -214,7 +213,11 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
                 wait=True,
             )
 
-    async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None) -> list[VectorStoreResult]:
+    async def retrieve(
+        self,
+        text: str,
+        options: VectorStoreOptions | None = None,
+    ) -> list[VectorStoreResult]:
         """
         Retrieves entries from the Qdrant collection based on vector similarity.
 
@@ -236,7 +239,7 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
         )
         with trace(
             text=text,
-            options=merged_options,
+            options=merged_options.dict(),
             index_name=self._index_name,
             distance_method=self._distance_method,
             embedder=repr(self._embedder),
@@ -252,6 +255,7 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
                 score_threshold=score_threshold,
                 with_payload=True,
                 with_vectors=True,
+                query_filter=self._create_qdrant_filter(merged_options.where),
             )
 
             outputs.results = []
@@ -290,16 +294,19 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
             )
 
     @staticmethod
-    def _create_qdrant_filter(where: WhereQuery) -> Filter:
+    def _create_qdrant_filter(where: WhereQuery | None) -> Filter:
         """
         Creates the QdrantFilter from the given WhereQuery.
 
         Args:
-            where: The WhereQuery to filter.
+            where: The WhereQuery to filter. If None, returns an empty filter.
 
         Returns:
             The created filter.
         """
+        if where is None:
+            return Filter(must=[])
+
         where = flatten_dict(where)  # type: ignore
 
         return Filter(
@@ -336,6 +343,7 @@ class QdrantVectorStore(VectorStoreWithEmbedder[VectorStoreOptions]):
                 return []
 
             limit = limit or (await self._client.count(collection_name=self._index_name)).count
+            limit = max(1, limit)
 
             qdrant_filter = self._create_qdrant_filter(where) if where else None