fix: UT & code format

Author: zhongyu09

openchatbi/utils.py

@@ -532,75 +532,71 @@ def max_marginal_relevance_search(
 
         # Get initial candidates using BM25 similarity search
         candidates = self.similarity_search_with_score(query, k=fetch_k, **kwargs)
-        
+
         if not candidates:
             return []
-        
+
         if len(candidates) <= k:
             return [doc for doc, _ in candidates]
-        
+
         # Normalize BM25 scores to [0, 1] for proper MMR calculation
         scores = [score for _, score in candidates]
         min_score = min(scores) if scores else 0
         max_score = max(scores) if scores else 1
         score_range = max_score - min_score if max_score > min_score else 1
-        
-        normalized_candidates = [
-            (doc, (score - min_score) / score_range) 
-            for doc, score in candidates
-        ]
-        
+
+        normalized_candidates = [(doc, (score - min_score) / score_range) for doc, score in candidates]
+
         # MMR implementation following standard algorithm
         selected = []
         remaining = list(range(len(normalized_candidates)))
-        
+
         # Select documents iteratively using MMR formula
         while len(selected) < k and remaining:
-            best_mmr_score = float('-inf')
+            best_mmr_score = float("-inf")
             best_idx = -1
             best_remaining_idx = -1
-            
+
             for i, doc_idx in enumerate(remaining):
                 candidate_doc, relevance_score = normalized_candidates[doc_idx]
-                
+
                 # Calculate maximum similarity to already selected documents
                 max_similarity = 0.0
                 if selected:
                     max_similarity = max(
                         self._calculate_similarity(candidate_doc, normalized_candidates[sel_idx][0])
                         for sel_idx in selected
                     )
-                
+
                 # Standard MMR formula: λ * Sim(q, d) - (1-λ) * max(Sim(d, s)) for s in selected
                 mmr_score = lambda_mult * relevance_score - (1 - lambda_mult) * max_similarity
-                
+
                 if mmr_score > best_mmr_score:
                     best_mmr_score = mmr_score
                     best_idx = doc_idx
                     best_remaining_idx = i
-            
+
             if best_idx != -1:
                 selected.append(best_idx)
                 remaining.pop(best_remaining_idx)
-        
+
         return [normalized_candidates[idx][0] for idx in selected]
 
     def _calculate_similarity(self, doc1: Document, doc2: Document) -> float:
         """Calculate similarity between two documents using Jaccard similarity.
-        
+
         Args:
             doc1: First document.
             doc2: Second document.
-            
+
         Returns:
             Similarity score between 0 and 1 (higher means more similar).
         """
         tokens1 = set(self._tokenize(doc1.page_content))
         tokens2 = set(self._tokenize(doc2.page_content))
-        
+
         # Calculate Jaccard similarity
         intersection = len(tokens1 & tokens2)
         union = len(tokens1 | tokens2)
-        
-        return intersection / union if union > 0 else 0.0
 
+        return intersection / union if union > 0 else 0.0

tests/test_simple_store.py

@@ -208,89 +208,35 @@ def test_chinese_and_mixed_language(self):
                 "数据科学" in doc.page_content for doc in cn_results
             )
 
-    def test_similarity_search_by_vector(self, simple_store):
-        """Test similarity_search_by_vector method."""
-        # Test with dummy embedding vector
-        dummy_embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
-        results = simple_store.similarity_search_by_vector(dummy_embedding, k=2)
-        
-        assert len(results) == 2
-        assert all(hasattr(doc, "page_content") for doc in results)
-        
-        # Test k parameter bounds
-        results = simple_store.similarity_search_by_vector(dummy_embedding, k=10)
-        assert len(results) == 4  # Should return all documents
-        
-        # Test empty store
-        empty_store = SimpleStore([])
-        results = empty_store.similarity_search_by_vector(dummy_embedding, k=5)
-        assert results == []
-
-    def test_max_marginal_relevance_search_by_vector(self, simple_store):
-        """Test max_marginal_relevance_search_by_vector method."""
-        dummy_embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
-        
-        # Test basic functionality
-        results = simple_store.max_marginal_relevance_search_by_vector(
-            dummy_embedding, k=2, fetch_k=4, lambda_mult=0.5
-        )
-        assert len(results) == 2
-        assert all(hasattr(doc, "page_content") for doc in results)
-        
-        # Test with k >= fetch_k
-        results = simple_store.max_marginal_relevance_search_by_vector(
-            dummy_embedding, k=4, fetch_k=3
-        )
-        assert len(results) == 3  # Should return fetch_k documents
-        
-        # Test diversity (lambda_mult = 0 should prioritize diversity)
-        results_diverse = simple_store.max_marginal_relevance_search_by_vector(
-            dummy_embedding, k=2, fetch_k=4, lambda_mult=0.0
-        )
-        assert len(results_diverse) == 2
-        
-        # Test empty store
-        empty_store = SimpleStore([])
-        results = empty_store.max_marginal_relevance_search_by_vector(dummy_embedding, k=2)
-        assert results == []
-
     def test_max_marginal_relevance_search(self, simple_store):
         """Test max_marginal_relevance_search method."""
         query = "programming language"
-        
+
         # Test basic MMR search
-        results = simple_store.max_marginal_relevance_search(
-            query, k=2, fetch_k=4, lambda_mult=0.5
-        )
+        results = simple_store.max_marginal_relevance_search(query, k=2, fetch_k=4, lambda_mult=0.5)
         assert len(results) == 2
         assert all(hasattr(doc, "page_content") for doc in results)
-        
+
         # Test relevance-focused search (lambda_mult = 1.0)
-        results_relevant = simple_store.max_marginal_relevance_search(
-            query, k=3, fetch_k=4, lambda_mult=1.0
-        )
+        results_relevant = simple_store.max_marginal_relevance_search(query, k=3, fetch_k=4, lambda_mult=1.0)
         assert len(results_relevant) == 3
-        
+
         # Test diversity-focused search (lambda_mult = 0.0)
-        results_diverse = simple_store.max_marginal_relevance_search(
-            query, k=3, fetch_k=4, lambda_mult=0.0
-        )
+        results_diverse = simple_store.max_marginal_relevance_search(query, k=3, fetch_k=4, lambda_mult=0.0)
         assert len(results_diverse) == 3
-        
+
         # Verify different lambda values produce different results
         # (unless there are ties in scoring)
         assert len(results_relevant) == len(results_diverse)
-        
+
         # Test with k >= fetch_k
-        results = simple_store.max_marginal_relevance_search(
-            query, k=5, fetch_k=3, lambda_mult=0.5
-        )
+        results = simple_store.max_marginal_relevance_search(query, k=5, fetch_k=3, lambda_mult=0.5)
         assert len(results) == 3  # Should return fetch_k documents
-        
+
         # Test empty query
         results = simple_store.max_marginal_relevance_search("", k=2)
         assert len(results) <= 2
-        
+
         # Test empty store
         empty_store = SimpleStore([])
         results = empty_store.max_marginal_relevance_search(query, k=2)
@@ -302,17 +248,18 @@ def test_calculate_similarity(self, simple_store):
         doc1 = simple_store.documents[0]  # "Python is a programming language"
         doc2 = simple_store.documents[1]  # "Machine learning is a subset of AI"
         doc3 = simple_store.documents[0]  # Same as doc1
-        
+
         # Test similarity between different documents
         similarity_diff = simple_store._calculate_similarity(doc1, doc2)
         assert 0.0 <= similarity_diff <= 1.0
-        
+
         # Test similarity between identical documents
         similarity_same = simple_store._calculate_similarity(doc1, doc3)
         assert similarity_same == 1.0
-        
+
         # Test with empty documents
         from langchain_core.documents import Document
+
         empty_doc1 = Document(page_content="", metadata={})
         empty_doc2 = Document(page_content="", metadata={})
         similarity_empty = simple_store._calculate_similarity(empty_doc1, empty_doc2)