참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_bedrock_embedder.py

@@ -0,0 +1,148 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from __future__ import annotations
+
+import io
+import json
+from typing import Any, Generator
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from neo4j_graphrag.embeddings.bedrock import BedrockEmbeddings
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+@pytest.fixture
+def mock_boto3() -> Generator[MagicMock, None, None]:
+    with patch("neo4j_graphrag.embeddings.bedrock.boto3") as mock_boto:
+        mock_client = MagicMock()
+        mock_boto.client.return_value = mock_client
+        yield mock_boto
+
+
+def _make_invoke_response(embedding: list[float]) -> dict[str, Any]:
+    body_bytes = json.dumps({"embedding": embedding}).encode()
+    return {"body": io.BytesIO(body_bytes)}
+
+
+@patch("neo4j_graphrag.embeddings.bedrock.boto3", None)
+def test_bedrock_embedder_missing_dependency() -> None:
+    with pytest.raises(ImportError) as exc:
+        BedrockEmbeddings()
+    assert "Could not import boto3 python client" in str(exc.value)
+
+
+def test_bedrock_embedder_default_model_from_env(mock_boto3: MagicMock) -> None:
+    with patch.dict(
+        "os.environ",
+        {"BEDROCK_EMBED_MODEL_ID": "custom-model", "BEDROCK_EMBED_DIMENSIONS": "256"},
+    ):
+        import importlib
+        import sys
+
+        # Ensure reload picks up the mock instead of real boto3
+        original_boto3 = sys.modules.get("boto3")
+        sys.modules["boto3"] = mock_boto3
+
+        try:
+            import neo4j_graphrag.embeddings.bedrock as bedrock_mod
+
+            importlib.reload(bedrock_mod)
+
+            assert bedrock_mod.DEFAULT_MODEL_ID == "custom-model"
+            assert bedrock_mod.DEFAULT_DIMENSIONS == 256
+
+            embedder = bedrock_mod.BedrockEmbeddings()
+            assert embedder.model_id == "custom-model"
+            assert embedder.dimensions == 256
+        finally:
+            # Restore real boto3 and reload to reset defaults
+            if original_boto3 is not None:
+                sys.modules["boto3"] = original_boto3
+            importlib.reload(bedrock_mod)
+
+
+def test_bedrock_embed_query_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.invoke_model.return_value = _make_invoke_response([0.1, 0.2, 0.3])
+
+    embedder = BedrockEmbeddings()
+    res = embedder.embed_query("hello")
+
+    assert res == [0.1, 0.2, 0.3]
+    mock_client.invoke_model.assert_called_once()
+    call_kwargs = mock_client.invoke_model.call_args[1]
+    body = json.loads(call_kwargs["body"])
+    assert body["inputText"] == "hello"
+    assert body["dimensions"] == 1024
+    assert body["normalize"] is True
+
+
+@pytest.mark.asyncio
+async def test_bedrock_async_embed_query_happy_path(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.invoke_model.return_value = _make_invoke_response([0.4, 0.5, 0.6])
+
+    embedder = BedrockEmbeddings()
+    res = await embedder.async_embed_query("hello")
+
+    assert res == [0.4, 0.5, 0.6]
+    mock_client.invoke_model.assert_called_once()
+
+
+def test_bedrock_embed_query_error(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.invoke_model.side_effect = Exception("API error")
+
+    embedder = BedrockEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with Bedrock"
+    ):
+        embedder.embed_query("hello")
+
+    assert mock_client.invoke_model.call_count == 1
+
+
+def test_bedrock_embed_query_custom_params(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    mock_client.invoke_model.return_value = _make_invoke_response([1.0, 2.0])
+
+    embedder = BedrockEmbeddings(
+        model_id="amazon.titan-embed-text-v1",
+        dimensions=512,
+        normalize=False,
+        region_name="eu-west-1",
+    )
+    res = embedder.embed_query("test")
+
+    assert res == [1.0, 2.0]
+    call_kwargs = mock_client.invoke_model.call_args[1]
+    assert call_kwargs["modelId"] == "amazon.titan-embed-text-v1"
+    body = json.loads(call_kwargs["body"])
+    assert body["dimensions"] == 512
+    assert body["normalize"] is False
+
+
+def test_bedrock_embed_query_empty_response(mock_boto3: MagicMock) -> None:
+    mock_client = mock_boto3.client.return_value
+    body_bytes = json.dumps({"embedding": None}).encode()
+    mock_client.invoke_model.return_value = {"body": io.BytesIO(body_bytes)}
+
+    embedder = BedrockEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with Bedrock"
+    ):
+        embedder.embed_query("hello")

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_cohere_embedder.py

@@ -0,0 +1,92 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from tenacity import RetryError
+from neo4j_graphrag.embeddings.cohere import CohereEmbeddings
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+@patch("neo4j_graphrag.embeddings.cohere.cohere", None)
+def test_cohere_embedder_missing_cohere_dependency() -> None:
+    with pytest.raises(ImportError):
+        CohereEmbeddings()
+
+
+@patch("neo4j_graphrag.embeddings.cohere.cohere")
+def test_cohere_embedder_happy_path(mock_cohere: Mock) -> None:
+    mock_cohere.Client.return_value.embed.return_value = MagicMock(
+        embeddings=[[1.0, 2.0]]
+    )
+    embedder = CohereEmbeddings()
+    res = embedder.embed_query("my text")
+    assert res == [1.0, 2.0]
+
+
+@patch("neo4j_graphrag.embeddings.cohere.cohere")
+def test_cohere_embedder_non_retryable_error_handling(mock_cohere: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    mock_embeddings = mock_cohere.Client.return_value.embed
+    mock_embeddings.side_effect = Exception("API Error")
+    embedder = CohereEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with Cohere"
+    ):
+        embedder.embed_query("my text")
+
+    # Verify the API was called only once (no retries for non-rate-limit errors)
+    assert mock_embeddings.call_count == 1
+
+
+@patch("neo4j_graphrag.embeddings.cohere.cohere")
+def test_cohere_embedder_rate_limit_error_retries(mock_cohere: Mock) -> None:
+    """Test that rate limit errors are retried the expected number of times."""
+    # Rate limit error that should trigger retries (matches "too many requests" pattern)
+    # Create separate exception instances for each retry attempt
+    mock_embeddings = mock_cohere.Client.return_value.embed
+    mock_embeddings.side_effect = [
+        Exception("too many requests - please try again later"),
+        Exception("too many requests - please try again later"),
+        Exception("too many requests - please try again later"),
+    ]
+    embedder = CohereEmbeddings()
+
+    # After exhausting retries, tenacity raises RetryError
+    with pytest.raises(RetryError):
+        embedder.embed_query("my text")
+
+    # Verify the API was called 3 times (default max_attempts for RetryRateLimitHandler)
+    assert mock_cohere.Client.return_value.embed.call_count == 3
+
+
+@patch("neo4j_graphrag.embeddings.cohere.cohere")
+def test_cohere_embedder_rate_limit_error_eventual_success(mock_cohere: Mock) -> None:
+    """Test that rate limit errors eventually succeed after retries."""
+    # First two calls fail with rate limit, third succeeds
+    mock_embeddings = mock_cohere.Client.return_value.embed
+    mock_embeddings.side_effect = [
+        Exception("too many requests - please try again later"),
+        Exception("too many requests - please try again later"),
+        MagicMock(embeddings=[[1.0, 2.0]]),
+    ]
+    embedder = CohereEmbeddings()
+
+    result = embedder.embed_query("my text")
+
+    # Verify successful result
+    assert result == [1.0, 2.0]
+    # Verify the API was called 3 times before succeeding
+    assert mock_embeddings.call_count == 3

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_mistralai_embedder.py

@@ -0,0 +1,145 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from tenacity import RetryError
+from neo4j_graphrag.embeddings import MistralAIEmbeddings
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral", None)
+def test_mistralai_embedder_missing_dependency() -> None:
+    with pytest.raises(ImportError):
+        MistralAIEmbeddings()
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+def test_mistralai_embedder_happy_path(mock_mistralai: Mock) -> None:
+    mock_mistral_instance = mock_mistralai.return_value
+    embeddings_batch_response_mock = MagicMock()
+    embeddings_batch_response_mock.data = [MagicMock(embedding=[1.0, 2.0])]
+    mock_mistral_instance.embeddings.create.return_value = (
+        embeddings_batch_response_mock
+    )
+    embedder = MistralAIEmbeddings()
+
+    res = embedder.embed_query("my text")
+
+    assert isinstance(res, list)
+    assert res == [1.0, 2.0]
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+def test_mistralai_embedder_api_key_via_kwargs(mock_mistral: Mock) -> None:
+    mock_mistral_instance = mock_mistral.return_value
+    embeddings_batch_response_mock = MagicMock()
+    embeddings_batch_response_mock.data = [MagicMock(embedding=[1.0, 2.0])]
+    mock_mistral_instance.embeddings.create.return_value = (
+        embeddings_batch_response_mock
+    )
+    api_key = "test_api_key"
+
+    MistralAIEmbeddings(api_key=api_key)
+
+    mock_mistral.assert_called_with(api_key=api_key)
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+@patch("os.getenv")
+def test_mistralai_embedder_api_key_from_env(
+    mock_getenv: Mock, mock_mistral: Mock
+) -> None:
+    mock_getenv.return_value = "env_api_key"
+    mock_mistral_instance = mock_mistral.return_value
+    embeddings_batch_response_mock = MagicMock()
+    embeddings_batch_response_mock.data = [MagicMock(embedding=[1.0, 2.0])]
+    mock_mistral_instance.embeddings.create.return_value = (
+        embeddings_batch_response_mock
+    )
+
+    MistralAIEmbeddings()
+
+    mock_getenv.assert_called_with("MISTRAL_API_KEY", "")
+    mock_mistral.assert_called_with(api_key="env_api_key")
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+def test_mistralai_embedder_non_retryable_error_handling(mock_mistral: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    mock_mistral_instance = mock_mistral.return_value
+    mock_embeddings = mock_mistral_instance.embeddings.create
+    mock_embeddings.side_effect = Exception("API Error")
+
+    embedder = MistralAIEmbeddings()
+    # MistralAI now wraps exceptions, so we expect EmbeddingsGenerationError
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with MistralAI"
+    ):
+        embedder.embed_query("my text")
+
+    # Verify the API was called only once (no retries for non-rate-limit errors)
+    assert mock_embeddings.call_count == 1
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+def test_mistralai_embedder_rate_limit_error_retries(mock_mistral: Mock) -> None:
+    """Test that rate limit errors are retried the expected number of times."""
+    mock_mistral_instance = mock_mistral.return_value
+
+    # Rate limit error that should trigger retries (matches "too many requests" pattern)
+    # Create separate exception instances for each retry attempt
+    mock_embeddings = mock_mistral_instance.embeddings.create
+    mock_embeddings.side_effect = [
+        Exception("too many requests - rate limit exceeded"),
+        Exception("too many requests - rate limit exceeded"),
+        Exception("too many requests - rate limit exceeded"),
+    ]
+
+    embedder = MistralAIEmbeddings()
+
+    # After exhausting retries, tenacity raises RetryError
+    with pytest.raises(RetryError):
+        embedder.embed_query("my text")
+
+    # Verify the API was called 3 times (default max_attempts for RetryRateLimitHandler)
+    assert mock_embeddings.call_count == 3
+
+
+@patch("neo4j_graphrag.embeddings.mistral.Mistral")
+def test_mistralai_embedder_rate_limit_error_eventual_success(
+    mock_mistral: Mock,
+) -> None:
+    """Test that rate limit errors eventually succeed after retries."""
+    mock_mistral_instance = mock_mistral.return_value
+
+    # First two calls fail with rate limit, third succeeds
+    embeddings_batch_response_mock = MagicMock()
+    embeddings_batch_response_mock.data = [MagicMock(embedding=[1.0, 2.0])]
+
+    mock_embeddings = mock_mistral_instance.embeddings.create
+    mock_embeddings.side_effect = [
+        Exception("too many requests - rate limit exceeded"),
+        Exception("too many requests - rate limit exceeded"),
+        embeddings_batch_response_mock,
+    ]
+
+    embedder = MistralAIEmbeddings()
+    result = embedder.embed_query("my text")
+
+    # Verify successful result
+    assert result == [1.0, 2.0]
+    # Verify the API was called 3 times before succeeding
+    assert mock_embeddings.call_count == 3

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_ollama_embedder.py

@@ -0,0 +1,46 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from neo4j_graphrag.embeddings.ollama import OllamaEmbeddings
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_ollama_embedder_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        OllamaEmbeddings(model="test")
+
+
+@patch("builtins.__import__")
+def test_ollama_embedder_happy_path(mock_import: Mock) -> None:
+    mock_import.return_value.Client.return_value.embed.return_value = MagicMock(
+        embeddings=[[1.0, 2.0]],
+    )
+    embedder = OllamaEmbeddings(model="test")
+    res = embedder.embed_query("my text")
+    assert isinstance(res, list)
+    assert res == [1.0, 2.0]
+
+
+@patch("builtins.__import__")
+def test_ollama_embedder_empty_list(mock_import: Mock) -> None:
+    mock_import.return_value.Client.return_value.embed.return_value = MagicMock(
+        embeddings=[],
+    )
+    embedder = OllamaEmbeddings(model="test")
+    with pytest.raises(EmbeddingsGenerationError):
+        embedder.embed_query("my text")

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_openai_embedder.py

@@ -0,0 +1,163 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock, Mock, patch
+
+import openai
+import pytest
+from tenacity import RetryError
+from neo4j_graphrag.embeddings.openai import (
+    AzureOpenAIEmbeddings,
+    OpenAIEmbeddings,
+)
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+def get_mock_openai() -> MagicMock:
+    mock = MagicMock()
+    mock.OpenAIError = openai.OpenAIError
+    return mock
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_openai_embedder_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        OpenAIEmbeddings()
+
+
+@patch("builtins.__import__")
+def test_openai_embedder_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    mock_openai.OpenAI.return_value.embeddings.create.return_value = MagicMock(
+        data=[MagicMock(embedding=[1.0, 2.0])],
+    )
+    embedder = OpenAIEmbeddings(api_key="my key")
+    res = embedder.embed_query("my text")
+    assert isinstance(res, list)
+    assert res == [1.0, 2.0]
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_azure_openai_embedder_missing_dependency(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        AzureOpenAIEmbeddings()
+
+
+@patch("builtins.__import__")
+def test_azure_openai_embedder_happy_path(mock_import: Mock) -> None:
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    mock_openai.AzureOpenAI.return_value.embeddings.create.return_value = MagicMock(
+        data=[MagicMock(embedding=[1.0, 2.0])],
+    )
+    embedder = AzureOpenAIEmbeddings(
+        model_name="gpt",
+        azure_endpoint="https://test.openai.azure.com/",
+        api_key="my key",
+        api_version="version",
+    )
+    res = embedder.embed_query("my text")
+    assert isinstance(res, list)
+    assert res == [1.0, 2.0]
+
+
+def test_azure_openai_embedder_does_not_call_openai_client() -> None:
+    from unittest.mock import patch
+
+    mock_openai = get_mock_openai()
+
+    with patch.dict("sys.modules", {"openai": mock_openai}):
+        AzureOpenAIEmbeddings(
+            model="text-embedding-ada-002",
+            azure_endpoint="https://test.openai.azure.com/",
+            api_key="my_key",
+            api_version="2023-05-15",
+        )
+
+        mock_openai.OpenAI.assert_not_called()
+        mock_openai.AzureOpenAI.assert_called_once_with(
+            azure_endpoint="https://test.openai.azure.com/",
+            api_key="my_key",
+            api_version="2023-05-15",
+        )
+
+
+@patch("builtins.__import__")
+def test_openai_embedder_non_retryable_error_handling(mock_import: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Generic API error that doesn't match rate limit patterns - should not be retried
+    mock_embeddings = mock_openai.OpenAI.return_value.embeddings.create
+    mock_embeddings.side_effect = Exception("API Error")
+    embedder = OpenAIEmbeddings(api_key="my key")
+
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with OpenAI"
+    ):
+        embedder.embed_query("my text")
+
+    # Verify the API was called only once (no retries for non-rate-limit errors)
+    assert mock_embeddings.call_count == 1
+
+
+@patch("builtins.__import__")
+def test_openai_embedder_rate_limit_error_retries(mock_import: Mock) -> None:
+    """Test that rate limit errors are retried the expected number of times."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # Rate limit error that should trigger retries (matches "429" pattern)
+    # Create separate exception instances for each retry attempt
+    mock_embeddings = mock_openai.OpenAI.return_value.embeddings.create
+    mock_embeddings.side_effect = [
+        Exception("Error code: 429 - Too many requests"),
+        Exception("Error code: 429 - Too many requests"),
+        Exception("Error code: 429 - Too many requests"),
+    ]
+    embedder = OpenAIEmbeddings(api_key="my key")
+
+    # After exhausting retries, tenacity raises RetryError
+    with pytest.raises(RetryError):
+        embedder.embed_query("my text")
+
+    # Verify the API was called 3 times (default max_attempts for RetryRateLimitHandler)
+    assert mock_embeddings.call_count == 3
+
+
+@patch("builtins.__import__")
+def test_openai_embedder_rate_limit_error_eventual_success(mock_import: Mock) -> None:
+    """Test that rate limit errors eventually succeed after retries."""
+    mock_openai = get_mock_openai()
+    mock_import.return_value = mock_openai
+
+    # First two calls fail with rate limit, third succeeds
+    mock_embeddings = mock_openai.OpenAI.return_value.embeddings.create
+    mock_embeddings.side_effect = [
+        Exception("Error code: 429 - Too many requests"),
+        Exception("Error code: 429 - Too many requests"),
+        MagicMock(data=[MagicMock(embedding=[1.0, 2.0])]),
+    ]
+    embedder = OpenAIEmbeddings(api_key="my key")
+
+    result = embedder.embed_query("my text")
+
+    # Verify successful result
+    assert result == [1.0, 2.0]
+    # Verify the API was called 3 times before succeeding
+    assert mock_embeddings.call_count == 3

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_sentence_transformers.py

@@ -0,0 +1,77 @@
+from unittest.mock import MagicMock, Mock, patch
+
+import numpy as np
+import pytest
+import torch
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.embeddings.sentence_transformers import (
+    SentenceTransformerEmbeddings,
+)
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+def get_mock_sentence_transformers() -> MagicMock:
+    mock = MagicMock()
+    # I know, I know... ¯\_(ツ)_/¯
+    # This is to cover the if type checks in the embed_query method
+    mock.Tensor = torch.Tensor
+    mock.ndarray = np.ndarray
+    return mock
+
+
+@patch("builtins.__import__")
+def test_initialization(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    instance = SentenceTransformerEmbeddings()
+    MockSentenceTransformer.SentenceTransformer.assert_called_with("all-MiniLM-L6-v2")
+    assert isinstance(instance, Embedder)
+
+
+@patch("builtins.__import__")
+def test_initialization_with_custom_model(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    custom_model = "distilbert-base-nli-stsb-mean-tokens"
+    SentenceTransformerEmbeddings(model=custom_model)
+    MockSentenceTransformer.SentenceTransformer.assert_called_with(custom_model)
+
+
+@patch("builtins.__import__")
+def test_embed_query(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    mock_model = MockSentenceTransformer.SentenceTransformer.return_value
+    mock_model.encode.return_value = np.array([[0.1, 0.2, 0.3]])
+
+    instance = SentenceTransformerEmbeddings()
+    result = instance.embed_query("test query")
+
+    mock_model.encode.assert_called_with(["test query"])
+    assert isinstance(result, list)
+    assert result == [0.1, 0.2, 0.3]
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_import_error(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        SentenceTransformerEmbeddings()
+
+
+@patch("builtins.__import__")
+def test_embed_query_non_retryable_error_handling(mock_import: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    mock_model = MockSentenceTransformer.SentenceTransformer.return_value
+    mock_model.encode.side_effect = Exception("Model error")
+
+    instance = SentenceTransformerEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError,
+        match="Failed to generate embedding with SentenceTransformer",
+    ):
+        instance.embed_query("test query")
+
+    # Verify the model was called only once (no retries for non-rate-limit errors)
+    assert mock_model.encode.call_count == 1

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_vertexai_embedder.py

@@ -0,0 +1,94 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from tenacity import RetryError
+from neo4j_graphrag.embeddings.vertexai import VertexAIEmbeddings
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+@patch("neo4j_graphrag.embeddings.vertexai.TextEmbeddingModel", None)
+def test_vertexai_embedder_missing_dependency() -> None:
+    with pytest.raises(ImportError):
+        VertexAIEmbeddings()
+
+
+@patch("neo4j_graphrag.embeddings.vertexai.TextEmbeddingModel")
+def test_vertexai_embedder_happy_path(mock_vertexai: Mock) -> None:
+    mock_vertexai.from_pretrained.return_value.get_embeddings.return_value = [
+        MagicMock(values=[1.0, 2.0])
+    ]
+    embedder = VertexAIEmbeddings()
+    res = embedder.embed_query("my text")
+    assert isinstance(res, list)
+    assert res == [1.0, 2.0]
+
+
+@patch("neo4j_graphrag.embeddings.vertexai.TextEmbeddingModel")
+def test_vertexai_embedder_non_retryable_error_handling(mock_vertexai: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    mock_embeddings = mock_vertexai.from_pretrained.return_value.get_embeddings
+    mock_embeddings.side_effect = Exception("API Error")
+    embedder = VertexAIEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError, match="Failed to generate embedding with VertexAI"
+    ):
+        embedder.embed_query("my text")
+
+    # Verify the API was called only once (no retries for non-rate-limit errors)
+    assert mock_embeddings.call_count == 1
+
+
+@patch("neo4j_graphrag.embeddings.vertexai.TextEmbeddingModel")
+def test_vertexai_embedder_rate_limit_error_retries(mock_vertexai: Mock) -> None:
+    """Test that rate limit errors are retried the expected number of times."""
+    # Rate limit error that should trigger retries (matches "resource exhausted" pattern)
+    mock_embeddings = mock_vertexai.from_pretrained.return_value.get_embeddings
+    mock_embeddings.side_effect = [
+        Exception("resource exhausted - quota exceeded"),
+        Exception("resource exhausted - quota exceeded"),
+        Exception("resource exhausted - quota exceeded"),
+    ]
+    embedder = VertexAIEmbeddings()
+
+    # After exhausting retries, tenacity raises RetryError
+    with pytest.raises(RetryError):
+        embedder.embed_query("my text")
+
+    # Verify the API was called 3 times (default max_attempts for RetryRateLimitHandler)
+    assert mock_embeddings.call_count == 3
+
+
+@patch("neo4j_graphrag.embeddings.vertexai.TextEmbeddingModel")
+def test_vertexai_embedder_rate_limit_error_eventual_success(
+    mock_vertexai: Mock,
+) -> None:
+    """Test that rate limit errors eventually succeed after retries."""
+    # First two calls fail with rate limit, third succeeds
+    mock_embeddings = mock_vertexai.from_pretrained.return_value.get_embeddings
+    mock_embeddings.side_effect = [
+        Exception("resource exhausted - quota exceeded"),
+        Exception("resource exhausted - quota exceeded"),
+        [MagicMock(values=[1.0, 2.0])],
+    ]
+    embedder = VertexAIEmbeddings()
+
+    result = embedder.embed_query("my text")
+
+    # Verify successful result
+    assert result == [1.0, 2.0]
+    # Verify the API was called 3 times before succeeding
+    assert mock_embeddings.call_count == 3