Run tests in vendor-specific environments

.github/workflows/pr.yaml

@@ -43,3 +43,29 @@ jobs:
         run: |
           poetry run coverage run -m pytest tests/unit
           poetry run coverage report --fail-under=90
+  test_contextual:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [ '3.9', '3.12' ]
+        vendor: ["anthropic", "cohere", "mistralai", "openai", "sentence_transformers", "google"]
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+      - name: Set Python version for Poetry
+        run: poetry env use python${{ matrix.python-version }}
+      - name: Install dependencies
+        run: poetry install --extras ${{ matrix.vendor }}
+      - name: Run unit tests and check coverage
+        run: |
+          poetry run pytest -k ${{ matrix.vendor }} tests/unit

pyproject.toml

@@ -71,7 +71,7 @@ openai = ["openai"]
 mistralai = ["mistralai"]
 qdrant = ["qdrant-client"]
 kg_creation_tools = ["pygraphviz"]
-sentence-transformers = ["sentence-transformers"]
+sentence_transformers = ["sentence-transformers"]
 experimental = ["pypdf", "fsspec", "langchain-text-splitters", "pygraphviz", "llama-index"]
 examples = ["langchain-openai", "langchain-huggingface"]
 
@@ -84,6 +84,15 @@ testpaths = ["tests"]
 filterwarnings = [
     "",
 ]
+markers = [
+    "anthropic: marks tests for anthropic env only",
+    "cohere: marks tests for cohere env only",
+    "mistralai: marks tests for mistral env only",
+    "sentence_transformers: marks tests for sentence_transformers env only",
+    "openai: marks tests for openai env only",
+    "google: marks tests for vertexai env only",
+    "experimental: marks tests for experimental features",
+]
 
 [tool.coverage.paths]
 source = ["src"]

tests/unit/embeddings/test_cohere_embedder.py

@@ -17,6 +17,8 @@
 import pytest
 from neo4j_graphrag.embeddings.cohere import CohereEmbeddings
 
+pytestmark = pytest.mark.cohere
+
 
 @patch("neo4j_graphrag.embeddings.cohere.cohere", None)
 def test_cohere_embedder_missing_cohere_dependency() -> None:

tests/unit/embeddings/test_mistralai_embedder.py

@@ -17,6 +17,8 @@
 import pytest
 from neo4j_graphrag.embeddings import MistralAIEmbeddings
 
+pytestmark = pytest.mark.mistralai
+
 
 @patch("neo4j_graphrag.embeddings.mistral.Mistral", None)
 def test_mistralai_embedder_missing_dependency() -> None:

tests/unit/embeddings/test_openai_embedder.py

@@ -14,8 +14,14 @@
 #  limitations under the License.
 from unittest.mock import MagicMock, Mock, patch
 
-import openai
 import pytest
+
+pytestmark = pytest.mark.openai
+
+try:
+    import openai
+except ImportError:
+    pass
 from neo4j_graphrag.embeddings.openai import (
     AzureOpenAIEmbeddings,
     OpenAIEmbeddings,

tests/unit/embeddings/test_sentence_transformers.py

@@ -1,8 +1,14 @@
 from unittest.mock import MagicMock, Mock, patch
 
-import numpy as np
 import pytest
-import torch
+
+pytestmark = pytest.mark.sentence_transformers
+
+try:
+    import numpy as np
+    import torch
+except ImportError:
+    np = None
 from neo4j_graphrag.embeddings.base import Embedder
 from neo4j_graphrag.embeddings.sentence_transformers import (
     SentenceTransformerEmbeddings,
@@ -18,6 +24,7 @@ def get_mock_sentence_transformers() -> MagicMock:
     return mock
 
 
+@pytest.mark.skipif(np is None, reason="numpy is not installed")
 @patch("builtins.__import__")
 def test_initialization(mock_import: Mock) -> None:
     MockSentenceTransformer = get_mock_sentence_transformers()
@@ -27,6 +34,7 @@ def test_initialization(mock_import: Mock) -> None:
     assert isinstance(instance, Embedder)
 
 
+@pytest.mark.skipif(np is None, reason="numpy is not installed")
 @patch("builtins.__import__")
 def test_initialization_with_custom_model(mock_import: Mock) -> None:
     MockSentenceTransformer = get_mock_sentence_transformers()
@@ -36,6 +44,7 @@ def test_initialization_with_custom_model(mock_import: Mock) -> None:
     MockSentenceTransformer.SentenceTransformer.assert_called_with(custom_model)
 
 
+@pytest.mark.skipif(np is None, reason="numpy is not installed")
 @patch("builtins.__import__")
 def test_embed_query(mock_import: Mock) -> None:
     MockSentenceTransformer = get_mock_sentence_transformers()

tests/unit/embeddings/test_vertexai_embedder.py

@@ -17,6 +17,8 @@
 import pytest
 from neo4j_graphrag.embeddings.vertexai import VertexAIEmbeddings
 
+pytestmark = pytest.mark.google
+
 
 @patch("neo4j_graphrag.embeddings.vertexai.vertexai", None)
 def test_vertexai_embedder_missing_dependency() -> None:

tests/unit/experimental/pipeline/test_kg_builder.py

@@ -18,11 +18,18 @@
 import neo4j
 import pytest
 from neo4j_graphrag.embeddings import Embedder
-from neo4j_graphrag.experimental.components.entity_relation_extractor import OnError
-from neo4j_graphrag.experimental.components.schema import SchemaEntity, SchemaRelation
-from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
-from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
-from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+
+try:
+    from neo4j_graphrag.experimental.components.entity_relation_extractor import OnError
+    from neo4j_graphrag.experimental.components.schema import (
+        SchemaEntity,
+        SchemaRelation,
+    )
+    from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
+    from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline
+    from neo4j_graphrag.experimental.pipeline.pipeline import PipelineResult
+except ImportError:
+    pass
 from neo4j_graphrag.llm.base import LLMInterface
 
 

tests/unit/llm/test_anthropic_llm.py

@@ -17,10 +17,15 @@
 from typing import Generator
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
-import anthropic
+try:
+    import anthropic
+except ImportError:
+    anthropic = None
 import pytest
 from neo4j_graphrag.llm.anthropic_llm import AnthropicLLM
 
+pytestmark = pytest.mark.anthropic
+
 
 @pytest.fixture
 def mock_anthropic() -> Generator[MagicMock, None, None]:

tests/unit/llm/test_cohere_llm.py

@@ -16,12 +16,17 @@
 from typing import Generator
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
-import cohere.core
+try:
+    import cohere.core
+except ImportError:
+    pass
 import pytest
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm import LLMResponse
 from neo4j_graphrag.llm.cohere_llm import CohereLLM
 
+pytestmark = pytest.mark.cohere
+
 
 @pytest.fixture
 def mock_cohere() -> Generator[MagicMock, None, None]:

tests/unit/llm/test_mistralaillm.py → tests/unit/llm/test_mistralai_llm.py

@@ -20,6 +20,8 @@
 from neo4j_graphrag.exceptions import LLMGenerationError
 from neo4j_graphrag.llm import LLMResponse, MistralAILLM
 
+pytestmark = pytest.mark.mistralai
+
 
 @patch("neo4j_graphrag.llm.mistralai_llm.Mistral", None)
 def test_mistral_ai_llm_missing_dependency() -> None:

tests/unit/llm/test_openai_llm.py

@@ -14,7 +14,10 @@
 #  limitations under the License.
 from unittest.mock import MagicMock, Mock, patch
 
-import openai
+try:
+    import openai
+except ImportError:
+    openai = None
 import pytest
 from neo4j_graphrag.llm import LLMResponse
 from neo4j_graphrag.llm.openai_llm import AzureOpenAILLM, OpenAILLM

tests/unit/llm/test_vertexai_llm.py

@@ -18,6 +18,8 @@
 import pytest
 from neo4j_graphrag.llm.vertexai_llm import VertexAILLM
 
+pytestmark = pytest.mark.google
+
 
 @patch("neo4j_graphrag.llm.vertexai_llm.GenerativeModel", None)
 def test_vertexai_llm_missing_dependency() -> None:

tests/unit/retrievers/external/test_pinecone.py

@@ -17,6 +17,9 @@
 
 import neo4j
 import pytest
+
+pytest.importorskip("pinecone", "Pinecone is not installed")
+
 from neo4j_graphrag.exceptions import RetrieverInitializationError
 from neo4j_graphrag.retrievers import PineconeNeo4jRetriever
 from neo4j_graphrag.retrievers.external.utils import get_match_query

tests/unit/retrievers/external/test_qdrant.py

@@ -17,12 +17,19 @@
 
 import neo4j
 import pytest
+
+pytest.importorskip("qdrant_client", "Qdrant client is not installed")
+
 from neo4j_graphrag.exceptions import RetrieverInitializationError
 from neo4j_graphrag.retrievers import QdrantNeo4jRetriever
 from neo4j_graphrag.retrievers.external.utils import get_match_query
 from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem
-from qdrant_client import QdrantClient
-from qdrant_client.http.models import QueryResponse, ScoredPoint
+
+try:
+    from qdrant_client import QdrantClient
+    from qdrant_client.http.models import QueryResponse, ScoredPoint
+except ImportError:
+    pass
 
 
 @pytest.fixture(scope="function")

tests/unit/retrievers/external/test_weaviate.py

@@ -19,6 +19,9 @@
 
 import neo4j
 import pytest
+
+pytest.importorskip("weaviate", "Weaviate is not installed")
+
 from neo4j_graphrag.exceptions import RetrieverInitializationError
 from neo4j_graphrag.retrievers import WeaviateNeo4jRetriever
 from neo4j_graphrag.retrievers.external.utils import get_match_query