vectoriz 0.0.5__py3-none-any.whl → 0.1.1b0__py3-none-any.whl

tests/test_files.py

@@ -1,40 +1,165 @@
+import docx
 import pytest
 import numpy as np
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 from vectoriz.files import FileArgument
+from vectoriz.files import FilesFeature
 
 
-def test_add_data_appends_to_lists():
-    file_arg = FileArgument([], [], [])
-    filename = "test.txt"
-    text = "Test content"
-    
-    with patch.object(FileArgument, '_create_embedding', return_value=np.array([0.1, 0.2, 0.3])):
-        file_arg.add_data(filename, text)
-        
-        assert file_arg.chunk_names == [filename]
-        assert file_arg.text_list == [text]
-        assert len(file_arg.embeddings) == 1
-        np.testing.assert_array_equal(file_arg.embeddings[0], np.array([0.1, 0.2, 0.3]))
-
-def test_add_data_multiple_entries():
-    file_arg = FileArgument(["existing.txt"], ["existing content"], [np.array([0.5, 0.5, 0.5])])
-    filename = "new.txt"
-    text = "New content"
-    
-    with patch.object(FileArgument, '_create_embedding', return_value=np.array([0.7, 0.8, 0.9])):
-        file_arg.add_data(filename, text)
-        assert file_arg.chunk_names == ["existing.txt", "new.txt"]
-        assert file_arg.text_list == ["existing content", "New content"]
-        assert len(file_arg.embeddings) == 2
-        np.testing.assert_array_equal(file_arg.embeddings[1], np.array([0.7, 0.8, 0.9]))
-
-def test_add_data_calls_create_embedding():
-    file_arg = FileArgument([], [], [])
-    filename = "test.txt"
-    text = "Test content"
-    
-    with patch.object(FileArgument, '_create_embedding') as mock_create_embedding:
-        mock_create_embedding.return_value = np.array([0.1, 0.2, 0.3])
-        file_arg.add_data(filename, text)
-        mock_create_embedding.assert_called_once_with(text)
+class TestFileArgument:
+    def test_add_data_appends_to_lists(self):
+        file_arg = FileArgument()
+        filename = "test.txt"
+        text = "Test content"
+
+        with patch.object(
+            FileArgument, "_create_embedding", return_value=np.array([0.1, 0.2, 0.3])
+        ):
+            file_arg.add_data(filename, text)
+
+            assert file_arg.chunk_names == [filename]
+            assert file_arg.text_list == [text]
+            assert len(file_arg.embeddings) == 1
+            np.testing.assert_array_equal(
+                file_arg.embeddings[0], np.array([0.1, 0.2, 0.3])
+            )
+
+    def test_add_data_multiple_entries(self):
+        file_arg = FileArgument(
+            ["existing.txt"], ["existing content"], [np.array([0.5, 0.5, 0.5])]
+        )
+        filename = "new.txt"
+        text = "New content"
+
+        with patch.object(
+            FileArgument, "_create_embedding", return_value=np.array([0.7, 0.8, 0.9])
+        ):
+            file_arg.add_data(filename, text)
+            assert file_arg.chunk_names == ["existing.txt", "new.txt"]
+            assert file_arg.text_list == ["existing content", "New content"]
+            assert len(file_arg.embeddings) == 2
+            np.testing.assert_array_equal(
+                file_arg.embeddings[1], np.array([0.7, 0.8, 0.9])
+            )
+
+    def test_add_data_calls_create_embedding(self):
+        file_arg = FileArgument()
+        filename = "test.txt"
+        text = "Test content"
+
+        with patch.object(FileArgument, "_create_embedding") as mock_create_embedding:
+            mock_create_embedding.return_value = np.array([0.1, 0.2, 0.3])
+            file_arg.add_data(filename, text)
+            mock_create_embedding.assert_called_once_with(text)
+
+    def test_create_embedding_returns_numpy_array(self):
+        file_arg = FileArgument()
+        text = "Test content"
+
+        with patch("vectoriz.files.TokenTransformer") as mock_transformer:
+            mock_instance = mock_transformer.return_value
+            mock_instance.text_to_embeddings.return_value = [np.array([0.1, 0.2, 0.3])]
+
+            result = file_arg._create_embedding(text)
+
+            assert isinstance(result, np.ndarray)
+            np.testing.assert_array_equal(result, np.array([0.1, 0.2, 0.3]))
+            mock_instance.text_to_embeddings.assert_called_once_with([text])
+
+    def test_create_embedding_handles_empty_text(self):
+        file_arg = FileArgument()
+        text = ""
+
+        with patch("vectoriz.files.TokenTransformer") as mock_transformer:
+            mock_instance = mock_transformer.return_value
+            mock_instance.text_to_embeddings.return_value = [np.array([0.0, 0.0, 0.0])]
+
+            result = file_arg._create_embedding(text)
+
+            assert isinstance(result, np.ndarray)
+            mock_instance.text_to_embeddings.assert_called_once_with([""])
+
+    def test_create_embedding_instantiates_token_transformer(self):
+        file_arg = FileArgument()
+        text = "Test content"
+
+        with patch("vectoriz.files.TokenTransformer") as mock_transformer:
+            mock_instance = mock_transformer.return_value
+            mock_instance.text_to_embeddings.return_value = [np.array([0.1, 0.2, 0.3])]
+
+            file_arg._create_embedding(text)
+
+            mock_transformer.assert_called_once()
+
+
+class TestFilesFeature:
+    def test_extract_txt_content_reads_file_correctly(self, tmp_path):
+        test_content = "This is test content"
+        test_file = tmp_path / "test.txt"
+        test_file.write_text(test_content)
+        files_feature = FilesFeature()
+        result = files_feature._extract_txt_content(str(tmp_path), "test.txt")
+        assert result == test_content
+
+    def test_extract_txt_content_with_unicode_chars(self, tmp_path):
+        test_content = "Unicode content: àáâãäåæç"
+        test_file = tmp_path / "unicode.txt"
+        test_file.write_text(test_content, encoding="utf-8")
+        files_feature = FilesFeature()
+        result = files_feature._extract_txt_content(str(tmp_path), "unicode.txt")
+        assert result == test_content
+
+    def test_extract_txt_content_raises_file_not_found(self):
+        files_feature = FilesFeature()
+        with pytest.raises(FileNotFoundError):
+            files_feature._extract_txt_content(
+                "/non_existent_dir", "non_existent_file.txt"
+            )
+
+    def test_extract_docx_content_reads_file_correctly(self, tmp_path, monkeypatch):
+        mock_doc = MagicMock()
+        mock_paragraph1 = MagicMock()
+        mock_paragraph1.text = "Paragraph 1"
+        mock_paragraph2 = MagicMock()
+        mock_paragraph2.text = "Paragraph 2"
+        mock_doc.paragraphs = [mock_paragraph1, mock_paragraph2]
+
+        monkeypatch.setattr(docx, "Document", lambda _: mock_doc)
+        files_feature = FilesFeature()
+        result = files_feature._extract_docx_content(str(tmp_path), "test.docx")
+
+        assert result == "Paragraph 1\nParagraph 2"
+
+    def test_extract_docx_content_skips_empty_paragraphs(self, tmp_path, monkeypatch):
+        mock_doc = MagicMock()
+        mock_paragraph1 = MagicMock()
+        mock_paragraph1.text = "Paragraph 1"
+        mock_paragraph2 = MagicMock()
+        mock_paragraph2.text = "   "
+        mock_paragraph3 = MagicMock()
+        mock_paragraph3.text = "Paragraph 3"
+        mock_doc.paragraphs = [mock_paragraph1, mock_paragraph2, mock_paragraph3]
+
+        monkeypatch.setattr(docx, "Document", lambda _: mock_doc)
+        files_feature = FilesFeature()
+        result = files_feature._extract_docx_content(str(tmp_path), "test.docx")
+
+        assert result == "Paragraph 1\nParagraph 3"
+
+    def test_extract_docx_content_exception_handling(self, tmp_path, monkeypatch):
+        def mock_document(_):
+            raise Exception("Failed to open document")
+
+        monkeypatch.setattr(docx, "Document", mock_document)
+
+        files_feature = FilesFeature()
+        with pytest.raises(Exception):
+            files_feature._extract_docx_content(str(tmp_path), "invalid.docx")
+
+    def test_extract_docx_content_with_no_paragraphs(self, tmp_path, monkeypatch):
+        mock_doc = MagicMock()
+        mock_doc.paragraphs = []
+        monkeypatch.setattr(docx, "Document", lambda _: mock_doc)
+        files_feature = FilesFeature()
+        result = files_feature._extract_docx_content(str(tmp_path), "empty.docx")
+        assert result == ""

tests/test_token_transformer.py

@@ -0,0 +1,193 @@
+import pytest
+import faiss
+import numpy as np
+from unittest.mock import Mock
+from vectoriz.token_transformer import TokenData, TokenTransformer
+
+
+class TestTokenData:
+
+    def test_from_vector_db(self):
+        mock_vector_data = Mock()
+
+        mock_index = faiss.IndexFlatL2(5)
+        mock_embeddings = np.random.random((3, 5)).astype("float32")
+        mock_texts = ["text1", "text2", "text3"]
+
+        mock_file_argument = Mock()
+        mock_file_argument.embeddings = mock_embeddings
+        mock_file_argument.text_list = mock_texts
+        mock_vector_data.faiss_index = mock_index
+        mock_vector_data.file_argument = mock_file_argument
+        token_data = TokenData.from_vector_db(mock_vector_data)
+
+        assert token_data.texts == mock_texts
+        assert token_data.index == mock_index
+        assert np.array_equal(token_data.embeddings, mock_embeddings)
+        assert isinstance(token_data, TokenData)
+
+    def test_from_file_argument(self):
+        mock_file_argument = Mock()
+        mock_file_argument.embeddings = np.random.random((3, 5)).astype("float32")
+        mock_file_argument.text_list = ["text1", "text2", "text3"]
+        mock_index = faiss.IndexFlatL2(5)
+        token_data = TokenData.from_file_argument(mock_file_argument, mock_index)
+
+        assert token_data.texts == mock_file_argument.text_list
+        assert token_data.index == mock_index
+        assert np.array_equal(token_data.embeddings, mock_file_argument.embeddings)
+        assert isinstance(token_data, TokenData)
+
+
+class TestTokenTransformer:
+
+    def test_text_to_embeddings(self):
+        transformer = TokenTransformer()
+        sentences = ["This is a test sentence.", "Another test sentence."]
+        embeddings = transformer.text_to_embeddings(sentences)
+
+        assert isinstance(embeddings, np.ndarray)
+        assert embeddings.shape[0] == len(sentences)
+        assert embeddings.shape[1] > 0
+
+    def test_text_to_embeddings_empty_list(self):
+        transformer = TokenTransformer()
+        sentences = []
+        embeddings = transformer.text_to_embeddings(sentences)
+
+        assert isinstance(embeddings, np.ndarray)
+
+    def test_get_np_vectors(self):
+        transformer = TokenTransformer()
+        embeddings = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
+        result = transformer.get_np_vectors(embeddings)
+
+        assert isinstance(result, np.ndarray)
+        assert result.dtype == np.float32
+        assert result.shape == (2, 3)
+        assert np.array_equal(
+            result, np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
+        )
+
+    def test_get_np_vectors_empty_list(self):
+        transformer = TokenTransformer()
+        embeddings = []
+        result = transformer.get_np_vectors(embeddings)
+
+        assert isinstance(result, np.ndarray)
+        assert result.dtype == np.float32
+        assert result.shape == (0,)
+
+    def test_get_np_vectors_single_element(self):
+        transformer = TokenTransformer()
+        embeddings = [[1.5, 2.5]]
+        result = transformer.get_np_vectors(embeddings)
+
+        assert isinstance(result, np.ndarray)
+        assert result.dtype == np.float32
+        assert result.shape == (1, 2)
+        assert np.array_equal(result, np.array([[1.5, 2.5]], dtype=np.float32))
+
+    def test_query_to_embeddings(self):
+        transformer = TokenTransformer()
+        query = "This is a test query"
+        result = transformer._query_to_embeddings(query)
+
+        assert isinstance(result, np.ndarray)
+        assert result.shape[0] == 1
+        assert result.shape[1] == 384
+
+    def test_query_to_embeddings_empty_string(self):
+        transformer = TokenTransformer()
+        query = ""
+        result = transformer._query_to_embeddings(query)
+
+        assert isinstance(result, np.ndarray)
+        assert result.shape[0] == 1
+        assert result.shape[1] == 384
+
+    def test_query_to_embeddings_returns_correct_shape(self):
+        transformer = TokenTransformer()
+        query1 = "First query"
+        query2 = "Second query with more words"
+        result1 = transformer._query_to_embeddings(query1)
+        result2 = transformer._query_to_embeddings(query2)
+
+        assert result1.shape == result2.shape
+        assert len(result1.shape) == 2
+        assert result1.shape[0] == 1
+
+    def test_search(self):
+        transformer = TokenTransformer()
+        texts = ["First document", "Second document", "Third document"]
+
+        embeddings = transformer.text_to_embeddings(texts)
+        index = transformer.embeddings_to_index(embeddings)
+
+        result = transformer.search("first", index, texts, context_amount=1)
+        assert isinstance(result, str)
+        assert "First document" in result
+
+        result = transformer.search("document", index, texts, context_amount=2)
+        assert isinstance(result, str)
+        assert len(result.strip().split("\n")) == 2
+
+    def test_search_with_empty_texts(self):
+        transformer = TokenTransformer()
+        texts = []
+
+        if texts:
+            embeddings = transformer.text_to_embeddings(texts)
+            index = transformer.embeddings_to_index(embeddings)
+        else:
+            index = faiss.IndexFlatL2(384)
+
+        result = transformer.search("query", index, texts)
+        assert result == ""
+
+    def test_search_with_different_context_amounts(self):
+        transformer = TokenTransformer()
+        texts = ["Doc 1", "Doc 2", "Doc 3", "Doc 4", "Doc 5"]
+
+        embeddings = transformer.text_to_embeddings(texts)
+        index = transformer.embeddings_to_index(embeddings)
+
+        result1 = transformer.search("Doc", index, texts, context_amount=1)
+        result3 = transformer.search("Doc", index, texts, context_amount=3)
+        result5 = transformer.search("Doc", index, texts, context_amount=5)
+
+        assert len(result1.strip().split("\n")) == 1
+        assert len(result3.strip().split("\n")) == 3
+        assert len(result5.strip().split("\n")) == 5
+
+    def test_create_index(self):
+        transformer = TokenTransformer()
+        texts = ["First document", "Second document", "Third document"]
+
+        token_data = transformer.create_index(texts)
+
+        assert isinstance(token_data, TokenData)
+        assert token_data.texts == texts
+        assert isinstance(token_data.index, faiss.IndexFlatL2)
+        assert isinstance(token_data.embeddings, np.ndarray)
+        assert token_data.embeddings.shape[0] == len(texts)
+        assert token_data.embeddings.shape[1] == 384
+
+    def test_create_index_empty_list(self):
+        transformer = TokenTransformer()
+        texts = []
+
+        with pytest.raises(ValueError, match="The input texts list is empty."):
+            transformer.create_index(texts)
+
+    def test_create_index_single_element(self):
+        transformer = TokenTransformer()
+        texts = ["Single document text"]
+
+        token_data = transformer.create_index(texts)
+
+        assert isinstance(token_data, TokenData)
+        assert token_data.texts == texts
+        assert isinstance(token_data.index, faiss.IndexFlatL2)
+        assert isinstance(token_data.embeddings, np.ndarray)
+        assert token_data.embeddings.shape == (1, 384)

tests/test_vector_db.py

@@ -0,0 +1,474 @@
+import os
+import pytest
+import numpy as np
+import faiss
+from unittest.mock import patch, MagicMock
+from vectoriz.vector_db import VectorDB, VectorDBClient
+from vectoriz.files import FileArgument
+
+
+class TestVectorDB:
+
+    def test_load_saved_data_successful(self, tmp_path):
+        vector_db = VectorDB()
+        faiss_path = str(tmp_path / "test.index")
+        np_path = str(tmp_path / "test.npz")
+
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+        mock_file_argument = MagicMock(spec=FileArgument)
+
+        with patch.object(
+            vector_db, "load_faiss_index", return_value=mock_index
+        ) as mock_load_index:
+            with patch.object(
+                vector_db, "load_numpy_embeddings", return_value=mock_file_argument
+            ) as mock_load_embeddings:
+
+                result = vector_db.load_saved_data(faiss_path, np_path)
+
+                assert isinstance(result, VectorDBClient)
+                assert result.faiss_index == mock_index
+                assert result.file_argument == mock_file_argument
+                mock_load_index.assert_called_once_with(faiss_path)
+                mock_load_embeddings.assert_called_once_with(np_path)
+
+    def test_load_saved_data_missing_faiss_index(self, tmp_path):
+        vector_db = VectorDB()
+        faiss_path = str(tmp_path / "nonexistent.index")
+        np_path = str(tmp_path / "test.npz")
+
+        mock_file_argument = MagicMock(spec=FileArgument)
+
+        with patch.object(
+            vector_db, "load_faiss_index", return_value=None
+        ) as mock_load_index:
+            with patch.object(
+                vector_db, "load_numpy_embeddings", return_value=mock_file_argument
+            ) as mock_load_embeddings:
+
+                result = vector_db.load_saved_data(faiss_path, np_path)
+
+                assert result is None
+                mock_load_index.assert_called_once_with(faiss_path)
+                mock_load_embeddings.assert_called_once_with(np_path)
+
+    def test_load_saved_data_missing_numpy_embeddings(self, tmp_path):
+        vector_db = VectorDB()
+        faiss_path = str(tmp_path / "test.index")
+        np_path = str(tmp_path / "nonexistent.npz")
+
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+
+        with patch.object(
+            vector_db, "load_faiss_index", return_value=mock_index
+        ) as mock_load_index:
+            with patch.object(
+                vector_db, "load_numpy_embeddings", return_value=None
+            ) as mock_load_embeddings:
+
+                result = vector_db.load_saved_data(faiss_path, np_path)
+
+                assert result is None
+                mock_load_index.assert_called_once_with(faiss_path)
+                mock_load_embeddings.assert_called_once_with(np_path)
+
+    def test_load_saved_data_both_missing(self, tmp_path):
+        vector_db = VectorDB()
+        faiss_path = str(tmp_path / "nonexistent.index")
+        np_path = str(tmp_path / "nonexistent.npz")
+
+        with patch.object(
+            vector_db, "load_faiss_index", return_value=None
+        ) as mock_load_index:
+            with patch.object(
+                vector_db, "load_numpy_embeddings", return_value=None
+            ) as mock_load_embeddings:
+
+                result = vector_db.load_saved_data(faiss_path, np_path)
+
+                assert result is None
+                mock_load_index.assert_called_once_with(faiss_path)
+                mock_load_embeddings.assert_called_once_with(np_path)
+
+    def test_save_faiss_index(self, tmp_path):
+        vector_db = VectorDB()
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+        faiss_path = str(tmp_path / "test")
+
+        with patch("faiss.write_index") as mock_write_index:
+            vector_db.save_faiss_index(mock_index, faiss_path)
+
+            mock_write_index.assert_called_once_with(mock_index, faiss_path + ".index")
+
+    def test_save_faiss_index_with_extension(self, tmp_path):
+        vector_db = VectorDB()
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+        faiss_path = str(tmp_path / "test.index")
+
+        with patch("faiss.write_index") as mock_write_index:
+            vector_db.save_faiss_index(mock_index, faiss_path)
+
+            mock_write_index.assert_called_once_with(mock_index, faiss_path)
+
+    def test_save_faiss_index_integration(self, tmp_path):
+        vector_db = VectorDB()
+        dimension = 128
+        index = faiss.IndexFlatL2(dimension)
+
+        sample_vectors = np.random.random((10, dimension)).astype("float32")
+        index.add(sample_vectors)
+        faiss_path = str(tmp_path / "test.index")
+
+        vector_db.save_faiss_index(index, faiss_path)
+
+        assert os.path.exists(faiss_path)
+
+        loaded_index = faiss.read_index(faiss_path)
+        assert loaded_index.ntotal == index.ntotal
+
+    def test_load_faiss_index_successful(self, tmp_path):
+        vector_db = VectorDB()
+        dimension = 128
+        index = faiss.IndexFlatL2(dimension)
+
+        sample_vectors = np.random.random((5, dimension)).astype("float32")
+        index.add(sample_vectors)
+        faiss_path = str(tmp_path / "test.index")
+
+        faiss.write_index(index, faiss_path)
+
+        result = vector_db.load_faiss_index(faiss_path)
+
+        assert result is not None
+        assert result.ntotal == index.ntotal
+        assert result.d == index.d
+
+    def test_load_faiss_index_missing_file(self, tmp_path):
+        vector_db = VectorDB()
+        nonexistent_path = str(tmp_path / "nonexistent.index")
+
+        result = vector_db.load_faiss_index(nonexistent_path)
+
+        assert result is None
+
+    def test_load_faiss_index_with_mock(self):
+        vector_db = VectorDB()
+        test_path = "/mock/path/test.index"
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+
+        with patch("os.path.exists", return_value=True) as mock_exists:
+            with patch("faiss.read_index", return_value=mock_index) as mock_read:
+                result = vector_db.load_faiss_index(test_path)
+
+                mock_exists.assert_called_once_with(test_path)
+                mock_read.assert_called_once_with(test_path)
+                assert result == mock_index
+
+    def test_save_numpy_embeddings_with_ndarray(self, tmp_path):
+        vector_db = VectorDB()
+        np_path = str(tmp_path / "test")
+
+        embeddings_np = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        chunk_names = np.array(["chunk1", "chunk2"])
+        texts = np.array(["text1", "text2"])
+
+        file_arg = FileArgument(
+            chunk_names=chunk_names,
+            text_list=texts,
+            embeddings=[],
+            ndarray_data=embeddings_np,
+        )
+
+        with patch("numpy.savez") as mock_savez:
+            vector_db.save_numpy_embeddings(file_arg, np_path)
+
+            mock_savez.assert_called_once_with(
+                np_path + ".npz",
+                embeddings=embeddings_np,
+                chunk_names=chunk_names,
+                texts=texts,
+            )
+
+    def test_save_numpy_embeddings_with_embeddings_list(self, tmp_path):
+        vector_db = VectorDB()
+        np_path = str(tmp_path / "test")
+
+        embeddings_list = [[[0.1, 0.2]], [[0.3, 0.4]]]
+        chunk_names = np.array(["chunk1", "chunk2"])
+        texts = np.array(["text1", "text2"])
+
+        file_arg = FileArgument(
+            chunk_names=chunk_names,
+            text_list=texts,
+            embeddings=embeddings_list,
+            ndarray_data=None,
+        )
+
+        transformed_np = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+
+        with patch(
+            "vectoriz.token_transformer.TokenTransformer.get_np_vectors",
+            return_value=transformed_np,
+        ) as mock_transform:
+            with patch("numpy.savez") as mock_savez:
+                vector_db.save_numpy_embeddings(file_arg, np_path)
+
+                mock_transform.assert_called_once_with(embeddings_list)
+                mock_savez.assert_called_once_with(
+                    np_path + ".npz",
+                    embeddings=transformed_np,
+                    chunk_names=chunk_names,
+                    texts=texts,
+                )
+
+    def test_save_numpy_embeddings_with_extension(self, tmp_path):
+        vector_db = VectorDB()
+        np_path = str(tmp_path / "test.npz")
+
+        embeddings_np = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        chunk_names = np.array(["chunk1", "chunk2"])
+        texts = np.array(["text1", "text2"])
+
+        file_arg = FileArgument(
+            chunk_names=chunk_names,
+            text_list=texts,
+            embeddings=[],
+            ndarray_data=embeddings_np,
+        )
+
+        with patch("numpy.savez") as mock_savez:
+            vector_db.save_numpy_embeddings(file_arg, np_path)
+
+            mock_savez.assert_called_once_with(
+                np_path, embeddings=embeddings_np, chunk_names=chunk_names, texts=texts
+            )
+
+    def test_save_numpy_embeddings_integration(self, tmp_path):
+        vector_db = VectorDB()
+        np_path = str(tmp_path / "test.npz")
+
+        embeddings_np = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        chunk_names = np.array(["chunk1", "chunk2"])
+        texts = np.array(["text1", "text2"])
+
+        file_arg = FileArgument(
+            chunk_names=chunk_names,
+            text_list=texts,
+            embeddings=[],
+            ndarray_data=embeddings_np,
+        )
+
+        vector_db.save_numpy_embeddings(file_arg, np_path)
+
+        assert os.path.exists(np_path)
+        loaded_data = np.load(np_path)
+        assert "embeddings" in loaded_data
+        assert "chunk_names" in loaded_data
+        assert "texts" in loaded_data
+        np.testing.assert_array_equal(loaded_data["embeddings"], embeddings_np)
+        np.testing.assert_array_equal(loaded_data["chunk_names"], chunk_names)
+        np.testing.assert_array_equal(loaded_data["texts"], texts)
+
+    def test_load_numpy_embeddings_successful(self, tmp_path):
+        vector_db = VectorDB()
+        np_path = str(tmp_path / "test.npz")
+
+        embeddings_np = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        chunk_names = np.array(["chunk1", "chunk2"])
+        texts = np.array(["text1", "text2"])
+
+        np.savez(
+            np_path, embeddings=embeddings_np, chunk_names=chunk_names, texts=texts
+        )
+
+        result = vector_db.load_numpy_embeddings(np_path)
+
+        assert result is not None
+        np.testing.assert_array_equal(result.ndarray_data, embeddings_np)
+        np.testing.assert_array_equal(result.chunk_names, chunk_names)
+        np.testing.assert_array_equal(result.text_list, texts)
+        assert result.embeddings == []
+
+    def test_load_numpy_embeddings_missing_file(self, tmp_path):
+        vector_db = VectorDB()
+        nonexistent_path = str(tmp_path / "nonexistent.npz")
+
+        result = vector_db.load_numpy_embeddings(nonexistent_path)
+
+        assert result is None
+
+    def test_load_numpy_embeddings_with_mock(self):
+        vector_db = VectorDB()
+        test_path = "/mock/path/test.npz"
+
+        mock_data = {
+            "embeddings": np.array([[0.1, 0.2], [0.3, 0.4]]),
+            "chunk_names": np.array(["name1", "name2"]),
+            "texts": np.array(["text1", "text2"]),
+        }
+
+        with patch("os.path.exists", return_value=True) as mock_exists:
+            with patch("numpy.load", return_value=mock_data) as mock_load:
+                result = vector_db.load_numpy_embeddings(test_path)
+
+                mock_exists.assert_called_once_with(test_path)
+                mock_load.assert_called_once_with(test_path)
+                assert isinstance(result, FileArgument)
+                np.testing.assert_array_equal(
+                    result.ndarray_data, mock_data["embeddings"]
+                )
+                np.testing.assert_array_equal(
+                    result.chunk_names, mock_data["chunk_names"]
+                )
+                np.testing.assert_array_equal(result.text_list, mock_data["texts"])
+
+
+class TestVectorDBClient:
+    def test_save_data_successful(self, tmp_path):
+        # Setup
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+        mock_file_argument = MagicMock(spec=FileArgument)
+        client = VectorDBClient(mock_index, mock_file_argument)
+        
+        faiss_path = str(tmp_path / "test")
+        np_path = str(tmp_path / "test")
+        
+        # Test with mocks to verify function calls
+        with patch.object(VectorDB, 'save_faiss_index') as mock_save_faiss:
+            with patch.object(VectorDB, 'save_numpy_embeddings') as mock_save_numpy:
+                client.save_data(faiss_path, np_path)
+                
+                mock_save_faiss.assert_called_once_with(mock_index, faiss_path)
+                mock_save_numpy.assert_called_once_with(mock_file_argument, np_path)
+    
+    def test_save_data_not_initialized(self):
+        # Test with None values
+        client1 = VectorDBClient(None, MagicMock(spec=FileArgument))
+        client2 = VectorDBClient(MagicMock(spec=faiss.IndexFlatL2), None)
+        client3 = VectorDBClient(None, None)
+        
+        with pytest.raises(ValueError, match="FAISS index or file argument is not initialized."):
+            client1.save_data("test.index", "test.npz")
+            
+        with pytest.raises(ValueError, match="FAISS index or file argument is not initialized."):
+            client2.save_data("test.index", "test.npz")
+            
+        with pytest.raises(ValueError, match="FAISS index or file argument is not initialized."):
+            client3.save_data("test.index", "test.npz")
+            
+    def test_save_data_integration(self, tmp_path):
+        # Setup real objects for integration test
+        dimension = 128
+        index = faiss.IndexFlatL2(dimension)
+        
+        # Add some vectors to the index
+        sample_vectors = np.random.random((5, dimension)).astype("float32")
+        index.add(sample_vectors)
+        
+        # Create file argument
+        chunk_names = np.array(["chunk1", "chunk2", "chunk3", "chunk4", "chunk5"])
+        texts = np.array(["text1", "text2", "text3", "text4", "text5"])
+        file_arg = FileArgument(
+            chunk_names=chunk_names,
+            text_list=texts,
+            embeddings=[],
+            ndarray_data=sample_vectors
+        )
+        
+        # Create client
+        client = VectorDBClient(index, file_arg)
+        
+        # Define paths
+        faiss_path = str(tmp_path / "test.index")
+        np_path = str(tmp_path / "test.npz")
+        
+        # Save the data
+        client.save_data(faiss_path, np_path)
+        
+        # Verify files were created
+        assert os.path.exists(faiss_path)
+        assert os.path.exists(np_path)
+        
+        # Load and verify the data
+        loaded_index = faiss.read_index(faiss_path)
+        loaded_data = np.load(np_path)
+        
+        assert loaded_index.ntotal == index.ntotal
+        assert loaded_index.d == index.d
+        assert "embeddings" in loaded_data
+        assert "chunk_names" in loaded_data
+        assert "texts" in loaded_data
+        np.testing.assert_array_equal(loaded_data["embeddings"], sample_vectors)
+        np.testing.assert_array_equal(loaded_data["chunk_names"], chunk_names)
+        np.testing.assert_array_equal(loaded_data["texts"], texts)
+    
+    def test_load_data_successful(self, tmp_path):
+        # Setup - create data files first
+        dimension = 128
+        index = faiss.IndexFlatL2(dimension)
+        sample_vectors = np.random.random((5, dimension)).astype("float32")
+        index.add(sample_vectors)
+        
+        chunk_names = np.array(["chunk1", "chunk2", "chunk3", "chunk4", "chunk5"])
+        texts = np.array(["text1", "text2", "text3", "text4", "text5"])
+        
+        # Save test data
+        faiss_path = str(tmp_path / "test.index")
+        np_path = str(tmp_path / "test.npz")
+        faiss.write_index(index, faiss_path)
+        np.savez(np_path, embeddings=sample_vectors, chunk_names=chunk_names, texts=texts)
+        
+        # Create client and load data
+        client = VectorDBClient()
+        client.load_data(faiss_path, np_path)
+        
+        # Verify data was loaded correctly
+        assert client.faiss_index is not None
+        assert client.file_argument is not None
+        assert client.faiss_index.ntotal == 5
+        assert client.faiss_index.d == dimension
+        np.testing.assert_array_equal(client.file_argument.ndarray_data, sample_vectors)
+        np.testing.assert_array_equal(client.file_argument.chunk_names, chunk_names)
+        np.testing.assert_array_equal(client.file_argument.text_list, texts)
+
+    def test_load_data_with_mocks(self):
+        client = VectorDBClient()
+        mock_index = MagicMock(spec=faiss.IndexFlatL2)
+        mock_file_arg = MagicMock(spec=FileArgument)
+        
+        with patch.object(VectorDB, 'load_faiss_index', return_value=mock_index) as mock_load_faiss:
+            with patch.object(VectorDB, 'load_numpy_embeddings', return_value=mock_file_arg) as mock_load_numpy:
+                client.load_data("test.index", "test.npz")
+                
+                mock_load_faiss.assert_called_once_with("test.index")
+                mock_load_numpy.assert_called_once_with("test.npz")
+                assert client.faiss_index == mock_index
+                assert client.file_argument == mock_file_arg
+
+    def test_load_data_missing_files(self, tmp_path):
+        client = VectorDBClient()
+        nonexistent_faiss = str(tmp_path / "nonexistent.index")
+        nonexistent_np = str(tmp_path / "nonexistent.npz")
+        
+        # Test loading non-existent files
+        client.load_data(nonexistent_faiss, nonexistent_np)
+        
+        # When files don't exist, the values should be None
+        assert client.faiss_index is None
+        assert client.file_argument is None
+
+    def test_load_data_partial_missing(self, tmp_path):
+        # Setup - create only one file
+        dimension = 128
+        index = faiss.IndexFlatL2(dimension)
+        faiss_path = str(tmp_path / "test.index")
+        nonexistent_np = str(tmp_path / "nonexistent.npz")
+        faiss.write_index(index, faiss_path)
+        
+        # Test loading with one missing file
+        client = VectorDBClient()
+        client.load_data(faiss_path, nonexistent_np)
+        
+        # Verify partial loading
+        assert client.faiss_index is not None
+        assert client.file_argument is None

vectoriz/files.py

@@ -2,14 +2,15 @@ import os
 import docx
 import numpy as np
 from typing import Optional
-from token_transformer import TokenTransformer
+
+from vectoriz.token_transformer import TokenTransformer
 
 class FileArgument:
     def __init__(
         self,
-        chunk_names: list[str],
-        text_list: list[str],
-        embeddings: list[float],
+        chunk_names: list[str] = [],
+        text_list: list[str] = [],
+        embeddings: list[float] = [],
         ndarray_data: Optional[np.ndarray] = None,
     ) -> None:
         """

vectoriz/token_transformer.py

@@ -77,6 +77,22 @@ class TokenTransformer:
         texts: list[str],
         context_amount: int = 1,
     ) -> str:
+        """
+        Searches for the most similar texts to the given query using the provided FAISS index.
+        This method converts the query into an embedding, searches for the k nearest neighbors
+        in the index, and returns the corresponding texts as context.
+        Args:
+            query (str): The search query text
+            index (faiss.IndexFlatL2): A FAISS index containing embeddings for the texts
+            texts (list[str]): A list of texts corresponding to the embeddings in the index
+            context_amount (int, optional): The number of texts to retrieve. Defaults to 1.
+        Returns:
+            str: The concatenated text of the most similar documents, separated by newlines.
+                 Returns an empty string if texts or query is empty or None.
+        """
+        if texts is None or len(texts) == 0 or query is None or len(query) == 0:
+            return ""
+        
         query_embedding = self._query_to_embeddings(query)
         _, I = index.search(query_embedding, k=context_amount)
         context = ""
@@ -99,6 +115,9 @@ class TokenTransformer:
         Returns:
             faiss.IndexFlatL2: A FAISS index containing the embeddings of the input texts.
         """
+        if len(texts) == 0:
+            raise ValueError("The input texts list is empty.")
+        
         embeddings = self.text_to_embeddings(texts)
         index = self.embeddings_to_index(embeddings)
         return TokenData(texts, index, embeddings)

vectoriz/vector_db.py

@@ -3,8 +3,8 @@ import faiss
 import numpy as np
 from typing import Optional
 
-from files import FileArgument
-from token_transformer import TokenTransformer
+from vectoriz.files import FileArgument
+from vectoriz.token_transformer import TokenTransformer
 
 
 class VectorDBClient:
@@ -182,7 +182,7 @@ class VectorDB:
         if not os.path.exists(np_db_path):
             return None
 
-        data = np.load(self.np_db_path)
+        data = np.load(np_db_path)
         embeddings_np = data["embeddings"]
         chunk_names = data["chunk_names"]
         texts = data["texts"]

{vectoriz-0.0.5.dist-info → vectoriz-0.1.1b0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vectoriz
-Version: 0.0.5
+Version: 0.1.1b0
 Summary: Python library for creating vectorized data from text or files.
 Home-page: https://github.com/PedroHenriqueDevBR/vectoriz
 Author: PedroHenriqueDevBR

vectoriz-0.1.1b0.dist-info/RECORD

@@ -0,0 +1,12 @@
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tests/test_files.py,sha256=EFXN9GChf9widEb3OvUcTXtOeU9X3naohMWyTZVPTJs,6559
+tests/test_token_transformer.py,sha256=LoLA9t_7owaghB5jS2hJrM1LYk3VSxa3Xo-qrWM2QZY,7152
+tests/test_vector_db.py,sha256=4vFxM6nhFFtI4ERuEY61dnQGsc7B90JBcn2_mvT8bWA,18369
+vectoriz/__init__.py,sha256=fnnle0EjVejiZQ8t243kvFiqcTTFh9dzmZbNwayjh4U,156
+vectoriz/files.py,sha256=jTwNBs1A_nqo0WWzLFNDGaBnyAPvLep283q3GuOH8bk,9056
+vectoriz/token_transformer.py,sha256=B7fPt-A-RzJjIoYns7wL_yyxQIj0UBRsnJIGCY_Ae2Q,6828
+vectoriz/vector_db.py,sha256=EqjKOTK1P4zP7wCmMo_Y2GsPzVP02UOzvurX-nTVuqI,6830
+vectoriz-0.1.1b0.dist-info/METADATA,sha256=wFmCkyq0v8xLZD9CUBkjqgiJyQUqs7eaEHynii_BUFE,3696
+vectoriz-0.1.1b0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+vectoriz-0.1.1b0.dist-info/top_level.txt,sha256=Tcfk3kazBwJ_yySjjhlIhLoTWLQGSb5xV006X18O6Nk,15
+vectoriz-0.1.1b0.dist-info/RECORD,,

vectoriz-0.0.5.dist-info/RECORD

@@ -1,10 +0,0 @@
-tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/test_files.py,sha256=jNIcwdDrPGKv78zGJReb6s0kPZLr3mTvkwwl3cP6MlM,1622
-vectoriz/__init__.py,sha256=fnnle0EjVejiZQ8t243kvFiqcTTFh9dzmZbNwayjh4U,156
-vectoriz/files.py,sha256=qpuD9CUZ4UFJQM1rokP4u23dnLz-3Du3fASDqmBpssk,9031
-vectoriz/token_transformer.py,sha256=1KIGL6EAiuqFhFzh3grkJcdEjjwTC3kH6RNqRkYMkQU,5811
-vectoriz/vector_db.py,sha256=cfcwN_QDbnbuBFqcJ_HyJy8jcRWpm8_pfsaug9JAiqo,6817
-vectoriz-0.0.5.dist-info/METADATA,sha256=FxQu0qlmjgRMPKq4OcmrdNKEYd2SnBa0LHwWExjxPQU,3694
-vectoriz-0.0.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-vectoriz-0.0.5.dist-info/top_level.txt,sha256=Tcfk3kazBwJ_yySjjhlIhLoTWLQGSb5xV006X18O6Nk,15
-vectoriz-0.0.5.dist-info/RECORD,,