diff --git a/.github/workflows/mcp-server-lint.yml b/.github/workflows/mcp-server-lint.yml index a553682b..dde55c99 100644 --- a/.github/workflows/mcp-server-lint.yml +++ b/.github/workflows/mcp-server-lint.yml @@ -11,6 +11,9 @@ on: jobs: format-and-lint: runs-on: ubuntu-latest + permissions: + contents: read + id-token: write steps: - name: Checkout repository uses: actions/checkout@v4 diff --git a/.github/workflows/mcp-server-tests.yml b/.github/workflows/mcp-server-tests.yml index 51544ccd..8e80082a 100644 --- a/.github/workflows/mcp-server-tests.yml +++ b/.github/workflows/mcp-server-tests.yml @@ -11,6 +11,9 @@ on: jobs: test-mcp-server: runs-on: ubuntu-latest + permissions: + contents: read + id-token: write steps: - name: Checkout repository uses: actions/checkout@v4 diff --git a/mcp_server/src/config/embedder_config.py b/mcp_server/src/config/embedder_config.py deleted file mode 100644 index 0ddca7e2..00000000 --- a/mcp_server/src/config/embedder_config.py +++ /dev/null @@ -1,124 +0,0 @@ -"""Embedder configuration for Graphiti MCP Server.""" - -import logging -import os - -from graphiti_core.embedder.azure_openai import AzureOpenAIEmbedderClient -from graphiti_core.embedder.client import EmbedderClient -from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig -from openai import AsyncAzureOpenAI -from pydantic import BaseModel - -from utils import create_azure_credential_token_provider - -logger = logging.getLogger(__name__) - -DEFAULT_EMBEDDER_MODEL = 'text-embedding-3-small' - - -class GraphitiEmbedderConfig(BaseModel): - """Configuration for the embedder client. - - Centralizes all embedding-related configuration parameters. - """ - - model: str = DEFAULT_EMBEDDER_MODEL - api_key: str | None = None - azure_openai_endpoint: str | None = None - azure_openai_deployment_name: str | None = None - azure_openai_api_version: str | None = None - azure_openai_use_managed_identity: bool = False - - @classmethod - def from_env(cls) -> 'GraphitiEmbedderConfig': - """Create embedder configuration from environment variables.""" - # Get model from environment, or use default if not set or empty - model_env = os.environ.get('EMBEDDER_MODEL_NAME', '') - model = model_env if model_env.strip() else DEFAULT_EMBEDDER_MODEL - - azure_openai_endpoint = os.environ.get('AZURE_OPENAI_EMBEDDING_ENDPOINT', None) - azure_openai_api_version = os.environ.get('AZURE_OPENAI_EMBEDDING_API_VERSION', None) - azure_openai_deployment_name = os.environ.get( - 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME', None - ) - azure_openai_use_managed_identity = ( - os.environ.get('AZURE_OPENAI_USE_MANAGED_IDENTITY', 'false').lower() == 'true' - ) - - if azure_openai_endpoint is not None: - # Setup for Azure OpenAI API - # Log if empty deployment name was provided - azure_openai_deployment_name = os.environ.get( - 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME', None - ) - if azure_openai_deployment_name is None: - logger.error('AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME environment variable not set') - raise ValueError( - 'AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME environment variable not set' - ) - - if not azure_openai_use_managed_identity: - # api key - api_key = os.environ.get('AZURE_OPENAI_EMBEDDING_API_KEY', None) or os.environ.get( - 'OPENAI_API_KEY', None - ) - else: - # Managed identity - api_key = None - - return cls( - azure_openai_use_managed_identity=azure_openai_use_managed_identity, - azure_openai_endpoint=azure_openai_endpoint, - api_key=api_key, - azure_openai_api_version=azure_openai_api_version, - azure_openai_deployment_name=azure_openai_deployment_name, - model=model, - ) - else: - return cls( - model=model, - api_key=os.environ.get('OPENAI_API_KEY'), - ) - - def create_client(self) -> EmbedderClient | None: - """Create an embedder client based on this configuration. - - Returns: - EmbedderClient instance or None if configuration is invalid - """ - if self.azure_openai_endpoint is not None: - # Azure OpenAI API setup - if self.azure_openai_use_managed_identity: - # Use managed identity for authentication - token_provider = create_azure_credential_token_provider() - return AzureOpenAIEmbedderClient( - azure_client=AsyncAzureOpenAI( - azure_endpoint=self.azure_openai_endpoint, - azure_deployment=self.azure_openai_deployment_name, - api_version=self.azure_openai_api_version, - azure_ad_token_provider=token_provider, - ), - model=self.model, - ) - elif self.api_key: - # Use API key for authentication - return AzureOpenAIEmbedderClient( - azure_client=AsyncAzureOpenAI( - azure_endpoint=self.azure_openai_endpoint, - azure_deployment=self.azure_openai_deployment_name, - api_version=self.azure_openai_api_version, - api_key=self.api_key, - ), - model=self.model, - ) - else: - logger.error('OPENAI_API_KEY must be set when using Azure OpenAI API') - return None - else: - # OpenAI API setup - if not self.api_key: - return None - - embedder_config = OpenAIEmbedderConfig(api_key=self.api_key, embedding_model=self.model) - - return OpenAIEmbedder(config=embedder_config) diff --git a/mcp_server/src/config/llm_config.py b/mcp_server/src/config/llm_config.py deleted file mode 100644 index 2e42725d..00000000 --- a/mcp_server/src/config/llm_config.py +++ /dev/null @@ -1,182 +0,0 @@ -"""LLM configuration for Graphiti MCP Server.""" - -import argparse -import logging -import os -from typing import TYPE_CHECKING - -from graphiti_core.llm_client import LLMClient -from graphiti_core.llm_client.azure_openai_client import AzureOpenAILLMClient -from graphiti_core.llm_client.config import LLMConfig -from graphiti_core.llm_client.openai_client import OpenAIClient -from openai import AsyncAzureOpenAI -from pydantic import BaseModel - -from utils import create_azure_credential_token_provider - -if TYPE_CHECKING: - pass - -logger = logging.getLogger(__name__) - -DEFAULT_LLM_MODEL = 'gpt-4.1-mini' -SMALL_LLM_MODEL = 'gpt-4.1-nano' - - -class GraphitiLLMConfig(BaseModel): - """Configuration for the LLM client. - - Centralizes all LLM-specific configuration parameters including API keys and model selection. - """ - - api_key: str | None = None - model: str = DEFAULT_LLM_MODEL - small_model: str = SMALL_LLM_MODEL - temperature: float = 0.0 - azure_openai_endpoint: str | None = None - azure_openai_deployment_name: str | None = None - azure_openai_api_version: str | None = None - azure_openai_use_managed_identity: bool = False - - @classmethod - def from_env(cls) -> 'GraphitiLLMConfig': - """Create LLM configuration from environment variables.""" - # Get model from environment, or use default if not set or empty - model_env = os.environ.get('MODEL_NAME', '') - model = model_env if model_env.strip() else DEFAULT_LLM_MODEL - - # Get small_model from environment, or use default if not set or empty - small_model_env = os.environ.get('SMALL_MODEL_NAME', '') - small_model = small_model_env if small_model_env.strip() else SMALL_LLM_MODEL - - azure_openai_endpoint = os.environ.get('AZURE_OPENAI_ENDPOINT', None) - azure_openai_api_version = os.environ.get('AZURE_OPENAI_API_VERSION', None) - azure_openai_deployment_name = os.environ.get('AZURE_OPENAI_DEPLOYMENT_NAME', None) - azure_openai_use_managed_identity = ( - os.environ.get('AZURE_OPENAI_USE_MANAGED_IDENTITY', 'false').lower() == 'true' - ) - - if azure_openai_endpoint is None: - # Setup for OpenAI API - # Log if empty model was provided - if model_env == '': - logger.debug( - f'MODEL_NAME environment variable not set, using default: {DEFAULT_LLM_MODEL}' - ) - elif not model_env.strip(): - logger.warning( - f'Empty MODEL_NAME environment variable, using default: {DEFAULT_LLM_MODEL}' - ) - - return cls( - api_key=os.environ.get('OPENAI_API_KEY'), - model=model, - small_model=small_model, - temperature=float(os.environ.get('LLM_TEMPERATURE', '0.0')), - ) - else: - # Setup for Azure OpenAI API - # Log if empty deployment name was provided - if azure_openai_deployment_name is None: - logger.error('AZURE_OPENAI_DEPLOYMENT_NAME environment variable not set') - raise ValueError('AZURE_OPENAI_DEPLOYMENT_NAME environment variable not set') - - if not azure_openai_use_managed_identity: - # api key - api_key = os.environ.get('OPENAI_API_KEY', None) - else: - # Managed identity - api_key = None - - return cls( - azure_openai_use_managed_identity=azure_openai_use_managed_identity, - azure_openai_endpoint=azure_openai_endpoint, - api_key=api_key, - azure_openai_api_version=azure_openai_api_version, - azure_openai_deployment_name=azure_openai_deployment_name, - model=model, - small_model=small_model, - temperature=float(os.environ.get('LLM_TEMPERATURE', '0.0')), - ) - - @classmethod - def from_cli_and_env(cls, args: argparse.Namespace) -> 'GraphitiLLMConfig': - """Create LLM configuration from CLI arguments, falling back to environment variables.""" - # Start with environment-based config - config = cls.from_env() - - # CLI arguments override environment variables when provided - if hasattr(args, 'model') and args.model: - # Only use CLI model if it's not empty - if args.model.strip(): - config.model = args.model - else: - # Log that empty model was provided and default is used - logger.warning(f'Empty model name provided, using default: {DEFAULT_LLM_MODEL}') - - if hasattr(args, 'small_model') and args.small_model: - if args.small_model.strip(): - config.small_model = args.small_model - else: - logger.warning(f'Empty small_model name provided, using default: {SMALL_LLM_MODEL}') - - if hasattr(args, 'temperature') and args.temperature is not None: - config.temperature = args.temperature - - return config - - def create_client(self) -> LLMClient: - """Create an LLM client based on this configuration. - - Returns: - LLMClient instance - """ - if self.azure_openai_endpoint is not None: - # Azure OpenAI API setup - if self.azure_openai_use_managed_identity: - # Use managed identity for authentication - token_provider = create_azure_credential_token_provider() - return AzureOpenAILLMClient( - azure_client=AsyncAzureOpenAI( - azure_endpoint=self.azure_openai_endpoint, - azure_deployment=self.azure_openai_deployment_name, - api_version=self.azure_openai_api_version, - azure_ad_token_provider=token_provider, - ), - config=LLMConfig( - api_key=self.api_key, - model=self.model, - small_model=self.small_model, - temperature=self.temperature, - ), - ) - elif self.api_key: - # Use API key for authentication - return AzureOpenAILLMClient( - azure_client=AsyncAzureOpenAI( - azure_endpoint=self.azure_openai_endpoint, - azure_deployment=self.azure_openai_deployment_name, - api_version=self.azure_openai_api_version, - api_key=self.api_key, - ), - config=LLMConfig( - api_key=self.api_key, - model=self.model, - small_model=self.small_model, - temperature=self.temperature, - ), - ) - else: - raise ValueError('OPENAI_API_KEY must be set when using Azure OpenAI API') - - if not self.api_key: - raise ValueError('OPENAI_API_KEY must be set when using OpenAI API') - - llm_client_config = LLMConfig( - api_key=self.api_key, model=self.model, small_model=self.small_model - ) - - # Set temperature - llm_client_config.temperature = self.temperature - - return OpenAIClient(config=llm_client_config) diff --git a/mcp_server/src/config/manager.py b/mcp_server/src/config/manager.py deleted file mode 100644 index 85e683a4..00000000 --- a/mcp_server/src/config/manager.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Unified configuration manager for Graphiti MCP Server.""" - -import argparse - -from pydantic import BaseModel, Field - -from .embedder_config import GraphitiEmbedderConfig -from .llm_config import GraphitiLLMConfig -from .neo4j_config import Neo4jConfig - - -class GraphitiConfig(BaseModel): - """Configuration for Graphiti client. - - Centralizes all configuration parameters for the Graphiti client. - """ - - llm: GraphitiLLMConfig = Field(default_factory=GraphitiLLMConfig) - embedder: GraphitiEmbedderConfig = Field(default_factory=GraphitiEmbedderConfig) - neo4j: Neo4jConfig = Field(default_factory=Neo4jConfig) - group_id: str | None = None - use_custom_entities: bool = False - destroy_graph: bool = False - - @classmethod - def from_env(cls) -> 'GraphitiConfig': - """Create a configuration instance from environment variables.""" - return cls( - llm=GraphitiLLMConfig.from_env(), - embedder=GraphitiEmbedderConfig.from_env(), - neo4j=Neo4jConfig.from_env(), - ) - - @classmethod - def from_cli_and_env(cls, args: argparse.Namespace) -> 'GraphitiConfig': - """Create configuration from CLI arguments, falling back to environment variables.""" - # Start with environment configuration - config = cls.from_env() - - # Apply CLI overrides - if args.group_id: - config.group_id = args.group_id - else: - config.group_id = 'default' - - config.use_custom_entities = args.use_custom_entities - config.destroy_graph = args.destroy_graph - - # Update LLM config using CLI args - config.llm = GraphitiLLMConfig.from_cli_and_env(args) - - return config diff --git a/mcp_server/src/config/neo4j_config.py b/mcp_server/src/config/neo4j_config.py deleted file mode 100644 index 8d365f81..00000000 --- a/mcp_server/src/config/neo4j_config.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Neo4j database configuration for Graphiti MCP Server.""" - -import os - -from pydantic import BaseModel - - -class Neo4jConfig(BaseModel): - """Configuration for Neo4j database connection.""" - - uri: str = 'bolt://localhost:7687' - user: str = 'neo4j' - password: str = 'password' - - @classmethod - def from_env(cls) -> 'Neo4jConfig': - """Create Neo4j configuration from environment variables.""" - return cls( - uri=os.environ.get('NEO4J_URI', 'bolt://localhost:7687'), - user=os.environ.get('NEO4J_USER', 'neo4j'), - password=os.environ.get('NEO4J_PASSWORD', 'password'), - ) diff --git a/mcp_server/src/config/schema.py b/mcp_server/src/config/schema.py index 04b6d88f..a8d0c135 100644 --- a/mcp_server/src/config/schema.py +++ b/mcp_server/src/config/schema.py @@ -1,4 +1,4 @@ -"""Enhanced configuration with pydantic-settings and YAML support.""" +"""Configuration schemas with pydantic-settings and YAML support.""" import os from pathlib import Path @@ -206,6 +206,10 @@ class GraphitiConfig(BaseSettings): embedder: EmbedderConfig = Field(default_factory=EmbedderConfig) database: DatabaseConfig = Field(default_factory=DatabaseConfig) graphiti: GraphitiAppConfig = Field(default_factory=GraphitiAppConfig) + + # Additional server options + use_custom_entities: bool = Field(default=False, description='Enable custom entity types') + destroy_graph: bool = Field(default=False, description='Clear graph on startup') model_config = SettingsConfigDict( env_prefix='', diff --git a/mcp_server/src/config/server_config.py b/mcp_server/src/config/server_config.py deleted file mode 100644 index e80673e2..00000000 --- a/mcp_server/src/config/server_config.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Server configuration for Graphiti MCP Server.""" - -import argparse - -from pydantic import BaseModel - - -class MCPConfig(BaseModel): - """Configuration for MCP server.""" - - transport: str = 'sse' # Default to SSE transport - - @classmethod - def from_cli(cls, args: argparse.Namespace) -> 'MCPConfig': - """Create MCP configuration from CLI arguments.""" - return cls(transport=args.transport) diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py index ae92d86b..e6ef68ad 100644 --- a/mcp_server/src/graphiti_mcp_server.py +++ b/mcp_server/src/graphiti_mcp_server.py @@ -24,8 +24,7 @@ from graphiti_core.utils.maintenance.graph_data_operations import clear_data from mcp.server.fastmcp import FastMCP from pydantic import BaseModel -from config.schema import GraphitiConfig -from config.server_config import MCPConfig +from config.schema import GraphitiConfig, ServerConfig from models.entity_types import ENTITY_TYPES from models.response_types import ( EpisodeSearchResponse, @@ -628,7 +627,7 @@ async def get_status() -> StatusResponse: ) -async def initialize_server() -> MCPConfig: +async def initialize_server() -> ServerConfig: """Parse CLI arguments and initialize the Graphiti server configuration.""" global config, graphiti_service, queue_service, graphiti_client, semaphore @@ -761,7 +760,7 @@ async def initialize_server() -> MCPConfig: mcp.settings.port = config.server.port # Return MCP configuration for transport - return MCPConfig(transport=config.server.transport) + return config.server async def run_mcp_server(): diff --git a/mcp_server/src/services/factories.py b/mcp_server/src/services/factories.py index cd59ea49..7a169480 100644 --- a/mcp_server/src/services/factories.py +++ b/mcp_server/src/services/factories.py @@ -272,8 +272,14 @@ class DatabaseDriverFactory: ) 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') + + falkor_config = config.providers.falkordb + return { + 'driver': 'falkordb', + 'uri': falkor_config.uri, + 'password': falkor_config.password, + 'database': falkor_config.database, + } case _: raise ValueError(f'Unsupported Database provider: {provider}') diff --git a/mcp_server/tests/test_simple_validation.py b/mcp_server/tests/test_simple_validation.py index 22033b7f..f77c72cc 100644 --- a/mcp_server/tests/test_simple_validation.py +++ b/mcp_server/tests/test_simple_validation.py @@ -4,6 +4,7 @@ Simple validation test for the refactored Graphiti MCP Server. Tests basic functionality quickly without timeouts. """ +import os import subprocess import sys import time @@ -13,14 +14,30 @@ def test_server_startup(): """Test that the refactored server starts up successfully.""" print('🚀 Testing Graphiti MCP Server Startup...') + # Check if uv is available + uv_cmd = None + for potential_uv in ['uv', '/Users/danielchalef/.local/bin/uv', '/root/.local/bin/uv']: + try: + result = subprocess.run([potential_uv, '--version'], capture_output=True, timeout=5) + if result.returncode == 0: + uv_cmd = potential_uv + break + except (subprocess.TimeoutExpired, FileNotFoundError): + continue + + if not uv_cmd: + print(' ⚠️ uv not found in PATH, skipping server startup test') + return True + try: # Start the server and capture output process = subprocess.Popen( - ['uv', 'run', 'main.py', '--transport', 'stdio'], + [uv_cmd, 'run', 'main.py', '--transport', 'stdio'], env={ 'NEO4J_URI': 'bolt://localhost:7687', 'NEO4J_USER': 'neo4j', 'NEO4J_PASSWORD': 'demodemo', + 'PATH': os.environ.get('PATH', ''), }, stdout=subprocess.PIPE, stderr=subprocess.PIPE, @@ -29,6 +46,7 @@ def test_server_startup(): # Wait for startup logs startup_output = '' + success = False for _ in range(50): # Wait up to 5 seconds if process.poll() is not None: break @@ -49,9 +67,9 @@ def test_server_startup(): except Exception: continue - else: - print(' ⚠️ Timeout waiting for initialization') - success = False + + if not success: + print(' ⚠️ Timeout waiting for initialization or server startup failed') # Clean shutdown process.terminate() @@ -81,11 +99,8 @@ def test_syntax_validation(): files_to_test = [ 'src/graphiti_mcp_server.py', - 'src/config/manager.py', - 'src/config/llm_config.py', - 'src/config/embedder_config.py', - 'src/config/neo4j_config.py', - 'src/config/server_config.py', + 'src/config/schema.py', + 'src/services/factories.py', 'src/services/queue_service.py', 'src/models/entity_types.py', 'src/models/response_types.py',