From ac42992e2286a9432170e0987af6c0b27bc03dcf Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 10 Mar 2025 17:16:49 +0000
Subject: [PATCH] Add additional tests and fix RAGStorage.search() comparison

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/memory/storage/rag_storage.py      |  2 +-
 .../test_string_knowledge_source_fix.py       | 56 +++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/crewai/memory/storage/rag_storage.py b/src/crewai/memory/storage/rag_storage.py
index fd4c77838..a5a6e8cf6 100644
--- a/src/crewai/memory/storage/rag_storage.py
+++ b/src/crewai/memory/storage/rag_storage.py
@@ -130,7 +130,7 @@ class RAGStorage(BaseRAGStorage):
                     "context": response["documents"][0][i],
                     "score": response["distances"][0][i],
                 }
-                if result["score"] >= score_threshold:
+                if result["score"] < score_threshold:  # Lower distance values indicate higher similarity in ChromaDB
                     results.append(result)
 
             return results
diff --git a/tests/knowledge/test_string_knowledge_source_fix.py b/tests/knowledge/test_string_knowledge_source_fix.py
index 690b9f640..3099f2fb4 100644
--- a/tests/knowledge/test_string_knowledge_source_fix.py
+++ b/tests/knowledge/test_string_knowledge_source_fix.py
@@ -61,3 +61,59 @@ def test_string_knowledge_source_integration():
         results = mock_storage.search(["What city does John live in?"])
         assert len(results) > 0
         assert "San Francisco" in results[0]["context"]
+
+def test_knowledge_storage_search_empty_results():
+    """Test that KnowledgeStorage.search() correctly handles empty results."""
+    # Create a mock collection to simulate ChromaDB with empty results
+    mock_collection = MagicMock()
+    mock_collection.query.return_value = {
+        "ids": [[]],
+        "metadatas": [[]],
+        "documents": [[]],
+        "distances": [[]]
+    }
+    
+    # Create a KnowledgeStorage instance with the mock collection
+    storage = KnowledgeStorage()
+    storage.collection = mock_collection
+    
+    # Search with the fixed implementation
+    results = storage.search(["test query"], score_threshold=0.35)
+    
+    # Assert that no results are returned
+    assert len(results) == 0
+
+def test_knowledge_storage_search_threshold_boundary():
+    """Test that KnowledgeStorage.search() correctly handles boundary threshold values."""
+    # Create a mock collection to simulate ChromaDB with a result at the exact threshold
+    mock_collection = MagicMock()
+    mock_collection.query.return_value = {
+        "ids": [["1"]],
+        "metadatas": [[{}]],
+        "documents": [["Doc1"]],
+        "distances": [[0.35]]  # Exact threshold value
+    }
+    
+    # Create a KnowledgeStorage instance with the mock collection
+    storage = KnowledgeStorage()
+    storage.collection = mock_collection
+    
+    # Search with the fixed implementation
+    results = storage.search(["test query"], score_threshold=0.35)
+    
+    # Assert that exact threshold matches are excluded
+    assert len(results) == 0
+
+def test_knowledge_storage_search_error_handling():
+    """Test that KnowledgeStorage.search() correctly handles errors."""
+    # Create a mock collection that raises an exception
+    mock_collection = MagicMock()
+    mock_collection.query.side_effect = Exception("ChromaDB error")
+    
+    # Create a KnowledgeStorage instance with the mock collection
+    storage = KnowledgeStorage()
+    storage.collection = mock_collection
+    
+    # Assert that the exception is propagated
+    with pytest.raises(Exception):
+        storage.search(["test query"], score_threshold=0.35)