graphiti/tests/config/test_factory.py

"""
Copyright 2024, Zep Software, Inc.

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

    http://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.
"""

import pytest

from graphiti_core.config import (
    DatabaseConfig,
    EmbedderConfig,
    EmbedderProvider,
    LLMProvider,
    LLMProviderConfig,
    RerankerConfig,
)
from graphiti_core.config.factory import (
    create_database_driver,
    create_embedder,
    create_llm_client,
    create_reranker,
)
from graphiti_core.config.providers import DatabaseProvider, RerankerProvider
from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
from graphiti_core.embedder.openai import OpenAIEmbedder
from graphiti_core.llm_client.openai_client import OpenAIClient


class TestCreateLLMClient:
    def test_create_openai_client(self, monkeypatch):
        """Test creating OpenAI LLM client."""
        monkeypatch.setenv('OPENAI_API_KEY', 'test-key')

        config = LLMProviderConfig(provider=LLMProvider.OPENAI)
        client = create_llm_client(config)

        assert isinstance(client, OpenAIClient)
        assert client.model == 'gpt-4.1-mini'
        assert client.small_model == 'gpt-4.1-nano'

    def test_create_azure_openai_client(self, monkeypatch):
        """Test creating Azure OpenAI LLM client."""
        from graphiti_core.llm_client.azure_openai_client import AzureOpenAILLMClient

        monkeypatch.setenv('AZURE_OPENAI_API_KEY', 'test-key')

        config = LLMProviderConfig(
            provider=LLMProvider.AZURE_OPENAI,
            base_url='https://my-resource.openai.azure.com',
            azure_deployment='gpt-4-deployment',
        )
        client = create_llm_client(config)

        assert isinstance(client, AzureOpenAILLMClient)

    def test_create_anthropic_client_missing_dep(self):
        """Test creating Anthropic client without dependency raises error."""
        config = LLMProviderConfig(provider=LLMProvider.ANTHROPIC)

        # This test will pass if anthropic is installed (in dev env)
        # or raise ImportError if not installed
        try:
            client = create_llm_client(config)
            from graphiti_core.llm_client.anthropic_client import AnthropicClient

            assert isinstance(client, AnthropicClient)
        except ImportError as e:
            assert 'anthropic' in str(e).lower()

    def test_create_litellm_client_missing_dep(self):
        """Test creating LiteLLM client without dependency raises error."""
        config = LLMProviderConfig(
            provider=LLMProvider.LITELLM,
            litellm_model='gpt-4',
        )

        # This will raise ImportError if litellm is not installed
        try:
            client = create_llm_client(config)
            from graphiti_core.llm_client.litellm_client import LiteLLMClient

            assert isinstance(client, LiteLLMClient)
        except ImportError as e:
            assert 'litellm' in str(e).lower()

    def test_unsupported_provider_raises_error(self):
        """Test unsupported provider raises ValueError."""
        # We need to bypass pydantic validation to test this
        config = LLMProviderConfig(provider=LLMProvider.OPENAI)
        config.provider = 'unsupported'  # type: ignore

        with pytest.raises(ValueError, match='Unsupported LLM provider'):
            create_llm_client(config)


class TestCreateEmbedder:
    def test_create_openai_embedder(self, monkeypatch):
        """Test creating OpenAI embedder."""
        monkeypatch.setenv('OPENAI_API_KEY', 'test-key')

        config = EmbedderConfig(provider=EmbedderProvider.OPENAI)
        embedder = create_embedder(config)

        assert isinstance(embedder, OpenAIEmbedder)
        assert embedder.config.embedding_model == 'text-embedding-3-small'

    def test_create_voyage_embedder_missing_dep(self):
        """Test creating Voyage embedder without dependency."""
        config = EmbedderConfig(provider=EmbedderProvider.VOYAGE)

        try:
            embedder = create_embedder(config)
            from graphiti_core.embedder.voyage import VoyageEmbedder

            assert isinstance(embedder, VoyageEmbedder)
        except ImportError as e:
            assert 'voyageai' in str(e).lower()

    def test_create_azure_openai_embedder(self, monkeypatch):
        """Test creating Azure OpenAI embedder."""
        from graphiti_core.embedder.azure_openai import AzureOpenAIEmbedderClient

        monkeypatch.setenv('AZURE_OPENAI_API_KEY', 'test-key')

        config = EmbedderConfig(
            provider=EmbedderProvider.AZURE_OPENAI,
            base_url='https://my-resource.openai.azure.com',
            azure_deployment='embedding-deployment',
        )
        embedder = create_embedder(config)

        assert isinstance(embedder, AzureOpenAIEmbedderClient)


class TestCreateReranker:
    def test_create_openai_reranker(self):
        """Test creating OpenAI reranker."""
        config = RerankerConfig(provider=RerankerProvider.OPENAI)
        reranker = create_reranker(config)

        assert isinstance(reranker, OpenAIRerankerClient)


class TestCreateDatabaseDriver:
    def test_create_neo4j_driver(self):
        """Test creating Neo4j driver."""
        from graphiti_core.driver.neo4j_driver import Neo4jDriver

        config = DatabaseConfig(
            provider=DatabaseProvider.NEO4J,
            uri='bolt://localhost:7687',
            user='neo4j',
            password='password',
        )
        driver = create_database_driver(config)

        assert isinstance(driver, Neo4jDriver)

    def test_neo4j_driver_missing_uri(self):
        """Test Neo4j driver without URI raises error."""
        config = DatabaseConfig(provider=DatabaseProvider.NEO4J)

        with pytest.raises(ValueError, match='uri is required'):
            create_database_driver(config)

    def test_create_falkordb_driver_missing_dep(self):
        """Test creating FalkorDB driver without dependency."""
        config = DatabaseConfig(
            provider=DatabaseProvider.FALKORDB,
            uri='redis://localhost:6379',
        )

        try:
            driver = create_database_driver(config)
            from graphiti_core.driver.falkor_driver import FalkorDriver

            assert isinstance(driver, FalkorDriver)
        except ImportError as e:
            assert 'falkordb' in str(e).lower()