diff --git a/camel/retrievers/vector_retriever.py b/camel/retrievers/vector_retriever.py index 5d95b81e0..c2887cf9c 100644 --- a/camel/retrievers/vector_retriever.py +++ b/camel/retrievers/vector_retriever.py @@ -195,6 +195,13 @@ def query( db_query = VectorDBQuery(query_vector=query_vector, top_k=top_k) query_results = self.storage.query(query=db_query) + # If no results found, raise an error + if not query_results: + raise ValueError( + "Query result is empty, please check if " + "the vector storage is empty." + ) + if query_results[0].record.payload is None: raise ValueError( "Payload of vector storage is None, please check the " diff --git a/test/retrievers/test_vector_retriever.py b/test/retrievers/test_vector_retriever.py index 1c61c2ec9..8367a4773 100644 --- a/test/retrievers/test_vector_retriever.py +++ b/test/retrievers/test_vector_retriever.py @@ -98,3 +98,39 @@ def test_query(vector_retriever): results = vector_retriever.query(query, top_k=top_k) assert len(results) == 1 assert results[0]['similarity score'] == '0.8' + + +# Test query with no results found +def test_query_no_results(vector_retriever): + query = "test query" + top_k = 1 + # Setup mock behavior for vector storage query + vector_retriever.storage.load = Mock() + vector_retriever.storage.query.return_value = [] + + with pytest.raises( + ValueError, + match="Query result is empty, please check if the " + "vector storage is empty.", + ): + vector_retriever.query(query, top_k=top_k) + + +# Test query with payload None +def test_query_payload_none(vector_retriever): + query = "test query" + top_k = 1 + # Setup mock behavior for vector storage query + vector_retriever.storage.load = Mock() + vector_retriever.storage.query.return_value = [ + Mock(similarity=0.8, record=Mock(payload=None)) + ] + + with pytest.raises( + ValueError, + match=( + "Payload of vector storage is None, " + "please check the collection." + ), + ): + vector_retriever.query(query, top_k=top_k)