beaver-db 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

beaver/collections.py

@@ -175,6 +175,9 @@ class CollectionWrapper:
         if not self._kdtree:
             return []
 
+        if top_k > len(self._doc_ids):
+            top_k = len(self._doc_ids)
+
         distances, indices = self._kdtree.query(
             np.array(vector, dtype=np.float32), k=top_k
         )
@@ -325,3 +328,58 @@ class CollectionWrapper:
             )
             for row in rows
         ]
+
+
+def rerank(
+    results: list[list[Document]],
+    weights: list[float] | None = None,
+    k: int = 60
+) -> list[Document]:
+    """
+    Reranks documents from multiple search result lists using Reverse Rank Fusion (RRF).
+    This function is specifically designed to work with beaver.collections.Document objects.
+
+    Args:
+        results (list[list[Document]]): A list of search result lists, where each
+            inner list contains Document objects.
+        weights (list[float], optional): A list of weights corresponding to each
+            result list. If None, all lists are weighted equally. Defaults to None.
+        k (int, optional): A constant used in the RRF formula. Defaults to 60.
+
+    Returns:
+        list[Document]: A single, reranked list of unique Document objects, sorted
+        by their fused rank score in descending order.
+    """
+    if not results:
+        return []
+
+    # Assign a default weight of 1.0 if none are provided
+    if weights is None:
+        weights = [1.0] * len(results)
+
+    if len(results) != len(weights):
+        raise ValueError("The number of result lists must match the number of weights.")
+
+    # Use dictionaries to store scores and unique documents by their ID
+    rrf_scores: dict[str, float] = {}
+    doc_store: dict[str, Document] = {}
+
+    # Iterate through each list of Document objects and its weight
+    for result_list, weight in zip(results, weights):
+        for rank, doc in enumerate(result_list):
+            # Use the .id attribute from the Document object
+            doc_id = doc.id
+            if doc_id not in doc_store:
+                doc_store[doc_id] = doc
+
+            # Calculate the reciprocal rank score, scaled by the weight
+            score = weight * (1 / (k + rank))
+
+            # Add the score to the document's running total
+            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + score
+
+    # Sort the document IDs by their final aggregated scores
+    sorted_doc_ids = sorted(rrf_scores.keys(), key=rrf_scores.get, reverse=True)
+
+    # Return the final list of Document objects in the new, reranked order
+    return [doc_store[doc_id] for doc_id in sorted_doc_ids]

{beaver_db-0.5.0.dist-info → beaver_db-0.5.2.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: beaver-db
-Version: 0.5.0
-Summary: Asynchronous, embedded, modern DB based on SQLite.
+Version: 0.5.2
+Summary: Fast, embedded, and multi-modal DB based on SQLite for AI-powered applications.
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 Requires-Dist: numpy>=2.3.3
@@ -11,7 +11,7 @@ Requires-Dist: scipy>=1.16.2
 
 A fast, single-file, multi-modal database for Python, built with the standard `sqlite3` library.
 
-`beaver` is the **B**ackend for **E**mbedded **A**synchronous **V**ector & **E**vent Retrieval. It's an industrious, all-in-one database designed to manage complex, modern data types without requiring a database server.
+`beaver` is the **B**ackend for **E**mbedded, **A**ll-in-one **V**ector, **E**ntity, and **R**elationship storage. It's a simple, local, and embedded database designed to manage complex, modern data types without requiring a database server, built on top of SQLite.
 
 ## Design Philosophy
 

{beaver_db-0.5.0.dist-info → beaver_db-0.5.2.dist-info}/RECORD

@@ -1,9 +1,9 @@
 beaver/__init__.py,sha256=-z5Gj6YKMOswpJOOn5Gej8z5i6k3c0Xs00DIYLA-bMI,75
-beaver/collections.py,sha256=fP1xkmo-XXlk3H_lPRiqFhtizINQn8192wOtXFlkTK4,12811
+beaver/collections.py,sha256=4rdGMTD7ex4SQUH52WIHOZhOdeWe7Nqvm9TPg7flv_g,15059
 beaver/core.py,sha256=sk0Z_k7EcORe6bN8CfPukGX7eAfmCGSX_B37KpJmQJ4,7279
 beaver/lists.py,sha256=JG1JOkaYCUldADUzPJhaNi93w-k3S8mUzcCw574uht4,5915
 beaver/subscribers.py,sha256=tCty2iDbeE9IXcPicbxj2CB5gqfLufMB9-nLQwqNBUU,1944
-beaver_db-0.5.0.dist-info/METADATA,sha256=VpZkov21mBFSTvgSsw6evkQZ9awqrMI42CQJ_RD3nDo,5829
-beaver_db-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-beaver_db-0.5.0.dist-info/top_level.txt,sha256=FxA4XnX5Qm5VudEXCduFriqi4dQmDWpQ64d7g69VQKI,7
-beaver_db-0.5.0.dist-info/RECORD,,
+beaver_db-0.5.2.dist-info/METADATA,sha256=ij6szwN6Ee0MutGlFgszF2vqddR5y8OlA9O4R5Y4Nbo,5904
+beaver_db-0.5.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+beaver_db-0.5.2.dist-info/top_level.txt,sha256=FxA4XnX5Qm5VudEXCduFriqi4dQmDWpQ64d7g69VQKI,7
+beaver_db-0.5.2.dist-info/RECORD,,