"""Factory classes for creating LLM, Embedder, and Database clients.""" from config_schema import ( DatabaseConfig, EmbedderConfig, LLMConfig, ) # Try to import FalkorDriver if available try: from graphiti_core.driver import FalkorDriver # noqa: F401 HAS_FALKOR = True except ImportError: HAS_FALKOR = False from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder from graphiti_core.llm_client import LLMClient, OpenAIClient # Try to import additional providers if available try: from graphiti_core.embedder import AzureOpenAIEmbedderClient HAS_AZURE_EMBEDDER = True except ImportError: HAS_AZURE_EMBEDDER = False try: from graphiti_core.embedder import GeminiEmbedder HAS_GEMINI_EMBEDDER = True except ImportError: HAS_GEMINI_EMBEDDER = False try: from graphiti_core.embedder import VoyageEmbedder HAS_VOYAGE_EMBEDDER = True except ImportError: HAS_VOYAGE_EMBEDDER = False try: from graphiti_core.llm_client import AzureOpenAILLMClient HAS_AZURE_LLM = True except ImportError: HAS_AZURE_LLM = False try: from graphiti_core.llm_client import AnthropicClient HAS_ANTHROPIC = True except ImportError: HAS_ANTHROPIC = False try: from graphiti_core.llm_client import GeminiClient HAS_GEMINI = True except ImportError: HAS_GEMINI = False try: from graphiti_core.llm_client import GroqClient HAS_GROQ = True except ImportError: HAS_GROQ = False from utils import create_azure_credential_token_provider class LLMClientFactory: """Factory for creating LLM clients based on configuration.""" @staticmethod def create(config: LLMConfig) -> LLMClient: """Create an LLM client based on the configured provider.""" provider = config.provider.lower() match provider: case 'openai': if not config.providers.openai: raise ValueError('OpenAI provider configuration not found') from graphiti_core.llm_client.config import LLMConfig as CoreLLMConfig llm_config = CoreLLMConfig( api_key=config.providers.openai.api_key, model=config.model, temperature=config.temperature, max_tokens=config.max_tokens, ) return OpenAIClient(config=llm_config) case 'azure_openai': if not HAS_AZURE_LLM: raise ValueError( 'Azure OpenAI LLM client not available in current graphiti-core version' ) if not config.providers.azure_openai: raise ValueError('Azure OpenAI provider configuration not found') azure_config = config.providers.azure_openai # Handle Azure AD authentication if enabled api_key: str | None = None azure_ad_token_provider = None if azure_config.use_azure_ad: azure_ad_token_provider = create_azure_credential_token_provider() else: api_key = azure_config.api_key return AzureOpenAILLMClient( api_key=api_key, api_url=azure_config.api_url, api_version=azure_config.api_version, azure_deployment=azure_config.deployment_name, azure_ad_token_provider=azure_ad_token_provider, model=config.model, temperature=config.temperature, max_tokens=config.max_tokens, ) case 'anthropic': if not HAS_ANTHROPIC: raise ValueError( 'Anthropic client not available in current graphiti-core version' ) if not config.providers.anthropic: raise ValueError('Anthropic provider configuration not found') return AnthropicClient( api_key=config.providers.anthropic.api_key, model=config.model, temperature=config.temperature, max_tokens=config.max_tokens, ) case 'gemini': if not HAS_GEMINI: raise ValueError('Gemini client not available in current graphiti-core version') if not config.providers.gemini: raise ValueError('Gemini provider configuration not found') return GeminiClient( api_key=config.providers.gemini.api_key, model=config.model, temperature=config.temperature, max_tokens=config.max_tokens, ) case 'groq': if not HAS_GROQ: raise ValueError('Groq client not available in current graphiti-core version') if not config.providers.groq: raise ValueError('Groq provider configuration not found') return GroqClient( api_key=config.providers.groq.api_key, api_url=config.providers.groq.api_url, model=config.model, temperature=config.temperature, max_tokens=config.max_tokens, ) case _: raise ValueError(f'Unsupported LLM provider: {provider}') class EmbedderFactory: """Factory for creating Embedder clients based on configuration.""" @staticmethod def create(config: EmbedderConfig) -> EmbedderClient: """Create an Embedder client based on the configured provider.""" provider = config.provider.lower() match provider: case 'openai': if not config.providers.openai: raise ValueError('OpenAI provider configuration not found') from graphiti_core.embedder.openai import OpenAIEmbedderConfig embedder_config = OpenAIEmbedderConfig( api_key=config.providers.openai.api_key, model=config.model, dimensions=config.dimensions, ) return OpenAIEmbedder(config=embedder_config) case 'azure_openai': if not HAS_AZURE_EMBEDDER: raise ValueError( 'Azure OpenAI embedder not available in current graphiti-core version' ) if not config.providers.azure_openai: raise ValueError('Azure OpenAI provider configuration not found') azure_config = config.providers.azure_openai # Handle Azure AD authentication if enabled api_key: str | None = None azure_ad_token_provider = None if azure_config.use_azure_ad: azure_ad_token_provider = create_azure_credential_token_provider() else: api_key = azure_config.api_key return AzureOpenAIEmbedderClient( api_key=api_key, api_url=azure_config.api_url, api_version=azure_config.api_version, azure_deployment=azure_config.deployment_name, azure_ad_token_provider=azure_ad_token_provider, model=config.model, dimensions=config.dimensions, ) case 'gemini': if not HAS_GEMINI_EMBEDDER: raise ValueError( 'Gemini embedder not available in current graphiti-core version' ) if not config.providers.gemini: raise ValueError('Gemini provider configuration not found') return GeminiEmbedder( api_key=config.providers.gemini.api_key, model=config.model, dimensions=config.dimensions, ) case 'voyage': if not HAS_VOYAGE_EMBEDDER: raise ValueError( 'Voyage embedder not available in current graphiti-core version' ) if not config.providers.voyage: raise ValueError('Voyage provider configuration not found') return VoyageEmbedder( api_key=config.providers.voyage.api_key, model=config.providers.voyage.model, ) case _: raise ValueError(f'Unsupported Embedder provider: {provider}') class DatabaseDriverFactory: """Factory for creating Database drivers based on configuration. Note: This returns configuration dictionaries that can be passed to Graphiti(), not driver instances directly, as the drivers require complex initialization. """ @staticmethod def create_config(config: DatabaseConfig) -> dict: """Create database configuration dictionary based on the configured provider.""" provider = config.provider.lower() match provider: case 'neo4j': if not config.providers.neo4j: raise ValueError('Neo4j provider configuration not found') neo4j_config = config.providers.neo4j return { 'uri': neo4j_config.uri, 'user': neo4j_config.username, 'password': neo4j_config.password, # Note: database and use_parallel_runtime would need to be passed # to the driver after initialization if supported } case 'falkordb': if not HAS_FALKOR: raise ValueError( 'FalkorDB driver not available in current graphiti-core version' ) if not config.providers.falkordb: raise ValueError('FalkorDB provider configuration not found') # FalkorDB support would need to be added to Graphiti core raise NotImplementedError('FalkorDB support requires graphiti-core updates') case _: raise ValueError(f'Unsupported Database provider: {provider}')