new code added

milvus_manager.py

@@ -99,17 +99,17 @@ class MilvusManager:
         self.client.create_index(
             collection_name=self.collection_name, index_params=index_params, sync=True
         )
-
-    def search(self, data, topk):
+    def search(self, data, topk, threshold=0.7):
         """
-        Search for the top-k most similar vectors in the collection.
+        Search for the top-k most similar vectors in the collection, filtered by a relevance threshold.
 
         Args:
             data (array-like): Query vector.
             topk (int): Number of top results to return.
+            threshold (float): Minimum score threshold for relevance (default is 0.5).
 
         Returns:
-            list: Sorted list of top-k results.
+            list: Sorted list of top-k results that meet the threshold.
         """
         search_params = {"metric_type": "IP", "params": {}}  # Search parameters for Inner Product
         results = self.client.search(
@@ -155,9 +155,71 @@ class MilvusManager:
                 score, doc_id = future.result()
                 scores.append((score, doc_id))
 
+        # Filter scores by threshold
+        filtered_scores = [item for item in scores if item[0] >= threshold]
+
         # Sort scores in descending order and return the top-k results
-        scores.sort(key=lambda x: x[0], reverse=True)
-        return scores[:topk] if len(scores) >= topk else scores
+        filtered_scores.sort(key=lambda x: x[0], reverse=True)
+        return filtered_scores[:topk] if len(filtered_scores) >= topk else filtered_scores
+
+    # def search(self, data, topk):
+    #     """
+    #     Search for the top-k most similar vectors in the collection.
+
+    #     Args:
+    #         data (array-like): Query vector.
+    #         topk (int): Number of top results to return.
+
+    #     Returns:
+    #         list: Sorted list of top-k results.
+    #     """
+    #     search_params = {"metric_type": "IP", "params": {}}  # Search parameters for Inner Product
+    #     results = self.client.search(
+    #         self.collection_name,
+    #         data,
+    #         limit=50,  # Initial retrieval limit
+    #         output_fields=["vector", "seq_id", "doc_id"],  # Fields to include in the output
+    #         search_params=search_params,
+    #     )
+
+    #     # Collect unique document IDs from the search results
+    #     doc_ids = set()
+    #     for r_id in range(len(results)):
+    #         for r in range(len(results[r_id])):
+    #             doc_ids.add(results[r_id][r]["entity"]["doc_id"])
+
+    #     scores = []
+
+    #     # Function to rerank a single document based on its relevance to the query
+    #     def rerank_single_doc(doc_id, data, client, collection_name):
+    #         doc_colbert_vecs = client.query(
+    #             collection_name=collection_name,
+    #             filter=f"doc_id in [{doc_id}, {doc_id + 1}]",  # Query documents by ID
+    #             output_fields=["seq_id", "vector", "doc"],  # Fields to retrieve
+    #             limit=1000,  # Retrieve a maximum of 1000 vectors per document
+    #         )
+    #         # Compute the maximum similarity score for the document
+    #         doc_vecs = np.vstack(
+    #             [doc_colbert_vecs[i]["vector"] for i in range(len(doc_colbert_vecs))]
+    #         )
+    #         score = np.dot(data, doc_vecs.T).max(1).sum()
+    #         return (score, doc_id)
+
+    #     # Use multithreading to rerank documents in parallel
+    #     with concurrent.futures.ThreadPoolExecutor(max_workers=300) as executor:
+    #         futures = {
+    #             executor.submit(
+    #                 rerank_single_doc, doc_id, data, self.client, self.collection_name
+    #             ): doc_id
+    #             for doc_id in doc_ids
+    #         }
+    #         for future in concurrent.futures.as_completed(futures):
+    #             score, doc_id = future.result()
+    #             scores.append((score, doc_id))
+
+    #     # Sort scores in descending order and return the top-k results
+    #     scores.sort(key=lambda x: x[0], reverse=True)
+    #     return scores[:topk] if len(scores) >= topk else scores
 
     def insert(self, data):
         """