diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py index a314e6a2..3405e0c0 100644 --- a/mcp_server/src/graphiti_mcp_server.py +++ b/mcp_server/src/graphiti_mcp_server.py @@ -144,13 +144,7 @@ For optimal performance, ensure the database is properly configured and accessib API keys are provided for any language model operations. """ -# MCP server instance -mcp = FastMCP( - 'Graphiti Agent Memory', - instructions=GRAPHITI_MCP_INSTRUCTIONS, -) - -# Global services +# Global services (set during lifespan initialization) graphiti_service: Optional['GraphitiService'] = None queue_service: QueueService | None = None @@ -159,6 +153,60 @@ graphiti_client: Graphiti | None = None semaphore: asyncio.Semaphore +@asynccontextmanager +async def graphiti_lifespan(app): + """Lifespan context manager for FastMCP Cloud deployment. + + This function initializes the Graphiti service when the server starts. + FastMCP Cloud calls this automatically - it does NOT run the if __name__ == '__main__' block. + """ + global config, graphiti_service, queue_service, graphiti_client, semaphore + + logger.info('Initializing Graphiti service via lifespan...') + + try: + # Load configuration from environment variables (FastMCP Cloud sets these) + config = GraphitiConfig() + + # Log configuration details + logger.info('Using configuration:') + logger.info(f' - LLM: {config.llm.provider} / {config.llm.model}') + logger.info(f' - Embedder: {config.embedder.provider} / {config.embedder.model}') + logger.info(f' - Database: {config.database.provider}') + logger.info(f' - Group ID: {config.graphiti.group_id}') + + # Initialize services (GraphitiService is defined below, but that's OK + # since this function is only called at runtime, not definition time) + graphiti_service = GraphitiService(config, SEMAPHORE_LIMIT) + queue_service = QueueService() + await graphiti_service.initialize() + + # Set global client for backward compatibility + graphiti_client = await graphiti_service.get_client() + semaphore = graphiti_service.semaphore + + # Initialize queue service with the client + await queue_service.initialize(graphiti_client) + + logger.info('Graphiti service initialized successfully via lifespan') + + yield # Server runs here + + except Exception as e: + logger.error(f'Failed to initialize Graphiti service: {e}') + raise + finally: + logger.info('Shutting down Graphiti service...') + + +# MCP server instance - pass lifespan to constructor for FastMCP Cloud +mcp = FastMCP( + 'Graphiti Agent Memory', + instructions=GRAPHITI_MCP_INSTRUCTIONS, + lifespan=graphiti_lifespan, +) + + class GraphitiService: """Graphiti service using the unified configuration system.""" @@ -329,55 +377,6 @@ class GraphitiService: return self.client -@asynccontextmanager -async def graphiti_lifespan(app): - """Lifespan context manager for FastMCP Cloud deployment. - - This function initializes the Graphiti service when the server starts. - FastMCP Cloud calls this automatically - it does NOT run the if __name__ == '__main__' block. - """ - global config, graphiti_service, queue_service, graphiti_client, semaphore - - logger.info('Initializing Graphiti service via lifespan...') - - try: - # Load configuration from environment variables (FastMCP Cloud sets these) - config = GraphitiConfig() - - # Log configuration details - logger.info('Using configuration:') - logger.info(f' - LLM: {config.llm.provider} / {config.llm.model}') - logger.info(f' - Embedder: {config.embedder.provider} / {config.embedder.model}') - logger.info(f' - Database: {config.database.provider}') - logger.info(f' - Group ID: {config.graphiti.group_id}') - - # Initialize services - graphiti_service = GraphitiService(config, SEMAPHORE_LIMIT) - queue_service = QueueService() - await graphiti_service.initialize() - - # Set global client for backward compatibility - graphiti_client = await graphiti_service.get_client() - semaphore = graphiti_service.semaphore - - # Initialize queue service with the client - await queue_service.initialize(graphiti_client) - - logger.info('Graphiti service initialized successfully via lifespan') - - yield # Server runs here - - except Exception as e: - logger.error(f'Failed to initialize Graphiti service: {e}') - raise - finally: - logger.info('Shutting down Graphiti service...') - - -# Update the MCP server to use the lifespan -mcp.lifespan = graphiti_lifespan - - @mcp.tool() async def add_memory( name: str,