diff --git a/sweepai/core/test_vector_db.py b/sweepai/core/test_vector_db.py new file mode 100644 index 0000000000..19f9d88aeb --- /dev/null +++ b/sweepai/core/test_vector_db.py @@ -0,0 +1,81 @@ +import pytest +import numpy as np +from unittest.mock import patch, MagicMock +from sweepai.core.vector_db import ( + embed_text_array, + openai_call_embedding, + multi_get_query_texts_similarity, + cosine_similarity, + normalize_l2, +) + +@pytest.fixture +def mock_openai_embedding(): + with patch('sweepai.core.vector_db.openai_call_embedding') as mock: + mock.return_value = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + yield mock + +@pytest.fixture +def mock_redis_client(): + with patch('sweepai.core.vector_db.vector_cache') as mock: + yield mock + +def test_embed_text_array(mock_openai_embedding): + texts = ["Hello, world!", "This is a test."] + result = embed_text_array(texts) + assert isinstance(result, list) + assert len(result) == 1 + assert isinstance(result[0], np.ndarray) + assert result[0].shape == (2, 3) + mock_openai_embedding.assert_called_once_with(texts) + +def test_openai_call_embedding(): + batch = ["Test text 1", "Test text 2"] + result = openai_call_embedding(batch) + assert isinstance(result, np.ndarray) + assert result.shape[0] == len(batch) + assert result.shape[1] == 512 # Assuming we're using the default 512 dimensions + +def test_normalize_l2(): + x = np.array([3, 4]) + result = normalize_l2(x) + assert np.allclose(result, np.array([0.6, 0.8])) + + x_2d = np.array([[3, 4], [6, 8]]) + result_2d = normalize_l2(x_2d) + assert np.allclose(result_2d, np.array([[0.6, 0.8], [0.6, 0.8]])) + +def test_embed_text_array_with_cache(mock_redis_client, mock_openai_embedding): + texts = ["Cached text", "New text"] + cached_embedding = np.array([0.7, 0.8, 0.9]) + mock_redis_client.get.side_effect = [cached_embedding.tobytes(), None] + + result = embed_text_array(texts) + + assert len(result) == 1 + assert np.allclose(result[0][0], cached_embedding) + assert np.allclose(result[0][1], mock_openai_embedding.return_value[0]) + mock_openai_embedding.assert_called_once_with(["New text"]) + mock_redis_client.set.assert_called_once() + +def test_cosine_similarity(): + a = np.array([[1, 0]]) + B = np.array([[1, 0], [0, 1], [-1, 0]]) + result = cosine_similarity(a, B) + expected = np.array([[1, 0, -1]]) + assert np.allclose(result, expected) + +def test_multi_get_query_texts_similarity(mock_openai_embedding): + queries = ["Query 1", "Query 2"] + documents = ["Doc 1", "Doc 2", "Doc 3"] + mock_openai_embedding.side_effect = [ + np.array([[0.1, 0.2], [0.3, 0.4]]), # Query embeddings + np.array([[0.5, 0.6], [0.7, 0.8], [0.9, 1.0]]) # Document embeddings + ] + + result = multi_get_query_texts_similarity(queries, documents) + + assert isinstance(result, list) + assert len(result) == len(queries) + assert len(result[0]) == len(documents) + mock_openai_embedding.assert_called() \ No newline at end of file