cognee-mcp

#!/usr/bin/env python3 """ Test client for Cognee MCP Server functionality. This script tests all the tools and functions available in the Cognee MCP server, including cognify, codify, search, prune, status checks, and utility functions. Usage: # Set your OpenAI API key first export OPENAI_API_KEY="your-api-key-here" # Run the test client python src/test_client.py # Or use LLM_API_KEY instead of OPENAI_API_KEY export LLM_API_KEY="your-api-key-here" python src/test_client.py """ import asyncio import os import tempfile import time from contextlib import asynccontextmanager from cognee.shared.logging_utils import setup_logging from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from cognee.modules.pipelines.models.PipelineRun import PipelineRunStatus from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError from src.server import ( node_to_string, retrieved_edges_to_string, load_class, ) # Set timeout for cognify/codify to complete in TIMEOUT = 5 * 60 # 5 min in seconds class CogneeTestClient: """Test client for Cognee MCP Server functionality.""" def __init__(self): self.test_results = {} self.temp_files = [] async def setup(self): """Setup test environment.""" print("🔧 Setting up test environment...") # Check for required API keys api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("LLM_API_KEY") if not api_key: print("⚠️ Warning: No OPENAI_API_KEY or LLM_API_KEY found in environment.") print(" Some tests may fail without proper LLM API configuration.") print(" Set OPENAI_API_KEY environment variable for full functionality.") else: print("✅ API key configured.") # Create temporary test files self.test_data_dir = tempfile.mkdtemp(prefix="cognee_test_") # Create a test text file self.test_text_file = os.path.join(self.test_data_dir, "test.txt") with open(self.test_text_file, "w") as f: f.write( "This is a test document for Cognee testing. It contains information about AI and knowledge graphs." ) # Create a test code repository structure self.test_repo_dir = os.path.join(self.test_data_dir, "test_repo") os.makedirs(self.test_repo_dir) # Create test Python files test_py_file = os.path.join(self.test_repo_dir, "main.py") with open(test_py_file, "w") as f: f.write(""" def hello_world(): '''A simple hello world function.''' return "Hello, World!" class TestClass: '''A test class for demonstration.''' def __init__(self, name): self.name = name def greet(self): return f"Hello, {self.name}!" """) # Create a test configuration file config_file = os.path.join(self.test_repo_dir, "config.py") with open(config_file, "w") as f: f.write(""" # Configuration settings DATABASE_URL = "sqlite:///test.db" DEBUG = True """) # Create test developer rules files cursorrules_file = os.path.join(self.test_data_dir, ".cursorrules") with open(cursorrules_file, "w") as f: f.write("# Test cursor rules\nUse Python best practices.") self.temp_files.extend([self.test_text_file, test_py_file, config_file, cursorrules_file]) print(f"✅ Test environment created at: {self.test_data_dir}") async def cleanup(self): """Clean up test environment.""" print("🧹 Cleaning up test environment...") import shutil if os.path.exists(self.test_data_dir): shutil.rmtree(self.test_data_dir) print("✅ Cleanup completed") @asynccontextmanager async def mcp_server_session(self): """Context manager to start and manage MCP server session.""" # Get the path to the server script server_script = os.path.join(os.path.dirname(__file__), "server.py") # Pass current environment variables to the server process # This ensures OpenAI API key and other config is available server_env = os.environ.copy() # Start the server process server_params = StdioServerParameters( command="python", args=[server_script, "--transport", "stdio"], env=server_env, ) async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: # Initialize the session await session.initialize() yield session async def test_mcp_server_startup_and_tools(self): """Test that the MCP server starts properly and returns tool results.""" print("\n🧪 Testing MCP server startup and tool execution...") try: async with self.mcp_server_session() as session: # Test 1: List available tools print(" 🔍 Testing tool discovery...") tools_result = await session.list_tools() expected_tools = { "cognify", "codify", "search", "prune", "cognify_status", "codify_status", "cognee_add_developer_rules", "list_data", "delete", } available_tools = {tool.name for tool in tools_result.tools} if not expected_tools.issubset(available_tools): missing_tools = expected_tools - available_tools raise AssertionError(f"Missing expected tools: {missing_tools}") print( f" ✅ Found {len(available_tools)} tools: {', '.join(sorted(available_tools))}" ) except Exception as e: self.test_results["mcp_server_integration"] = { "status": "FAIL", "error": str(e), "message": "MCP server integration test failed", } print(f"❌ MCP server integration test failed: {e}") async def test_prune(self): """Test the prune functionality using MCP client.""" print("\n🧪 Testing prune functionality...") try: async with self.mcp_server_session() as session: result = await session.call_tool("prune", arguments={}) self.test_results["prune"] = { "status": "PASS", "result": result, "message": "Prune executed successfully", } print("✅ Prune test passed") except Exception as e: self.test_results["prune"] = { "status": "FAIL", "error": str(e), "message": "Prune test failed", } print(f"❌ Prune test failed: {e}") raise e async def test_cognify(self, test_text, test_name): """Test the cognify functionality using MCP client.""" print("\n🧪 Testing cognify functionality...") try: # Test with simple text using MCP client async with self.mcp_server_session() as session: cognify_result = await session.call_tool("cognify", arguments={"data": test_text}) start = time.time() # mark the start while True: try: # Wait a moment await asyncio.sleep(5) # Check if cognify processing is finished status_result = await session.call_tool("cognify_status", arguments={}) if hasattr(status_result, "content") and status_result.content: status_text = ( status_result.content[0].text if status_result.content else str(status_result) ) else: status_text = str(status_result) if str(PipelineRunStatus.DATASET_PROCESSING_COMPLETED) in status_text: break elif time.time() - start > TIMEOUT: raise TimeoutError("Cognify did not complete in 5min") except DatabaseNotCreatedError: if time.time() - start > TIMEOUT: raise TimeoutError("Database was not created in 5min") self.test_results[test_name] = { "status": "PASS", "result": cognify_result, "message": f"{test_name} executed successfully", } print(f"✅ {test_name} test passed") except Exception as e: self.test_results[test_name] = { "status": "FAIL", "error": str(e), "message": f"{test_name} test failed", } print(f"❌ {test_name} test failed: {e}") async def test_codify(self): """Test the codify functionality using MCP client.""" print("\n🧪 Testing codify functionality...") try: async with self.mcp_server_session() as session: codify_result = await session.call_tool( "codify", arguments={"repo_path": self.test_repo_dir} ) start = time.time() # mark the start while True: try: # Wait a moment await asyncio.sleep(5) # Check if codify processing is finished status_result = await session.call_tool("codify_status", arguments={}) if hasattr(status_result, "content") and status_result.content: status_text = ( status_result.content[0].text if status_result.content else str(status_result) ) else: status_text = str(status_result) if str(PipelineRunStatus.DATASET_PROCESSING_COMPLETED) in status_text: break elif time.time() - start > TIMEOUT: raise TimeoutError("Codify did not complete in 5min") except DatabaseNotCreatedError: if time.time() - start > TIMEOUT: raise TimeoutError("Database was not created in 5min") self.test_results["codify"] = { "status": "PASS", "result": codify_result, "message": "Codify executed successfully", } print("✅ Codify test passed") except Exception as e: self.test_results["codify"] = { "status": "FAIL", "error": str(e), "message": "Codify test failed", } print(f"❌ Codify test failed: {e}") async def test_cognee_add_developer_rules(self): """Test the cognee_add_developer_rules functionality using MCP client.""" print("\n🧪 Testing cognee_add_developer_rules functionality...") try: async with self.mcp_server_session() as session: result = await session.call_tool( "cognee_add_developer_rules", arguments={"base_path": self.test_data_dir} ) start = time.time() # mark the start while True: try: # Wait a moment await asyncio.sleep(5) # Check if developer rule cognify processing is finished status_result = await session.call_tool("cognify_status", arguments={}) if hasattr(status_result, "content") and status_result.content: status_text = ( status_result.content[0].text if status_result.content else str(status_result) ) else: status_text = str(status_result) if str(PipelineRunStatus.DATASET_PROCESSING_COMPLETED) in status_text: break elif time.time() - start > TIMEOUT: raise TimeoutError( "Cognify of developer rules did not complete in 5min" ) except DatabaseNotCreatedError: if time.time() - start > TIMEOUT: raise TimeoutError("Database was not created in 5min") self.test_results["cognee_add_developer_rules"] = { "status": "PASS", "result": result, "message": "Developer rules addition executed successfully", } print("✅ Developer rules test passed") except Exception as e: self.test_results["cognee_add_developer_rules"] = { "status": "FAIL", "error": str(e), "message": "Developer rules test failed", } print(f"❌ Developer rules test failed: {e}") async def test_search_functionality(self): """Test the search functionality with different search types using MCP client.""" print("\n🧪 Testing search functionality...") search_query = "What is artificial intelligence?" # Test if all search types will execute from cognee import SearchType # Go through all Cognee search types for search_type in SearchType: # Don't test these search types if search_type in [SearchType.NATURAL_LANGUAGE, SearchType.CYPHER]: break try: async with self.mcp_server_session() as session: result = await session.call_tool( "search", arguments={"search_query": search_query, "search_type": search_type.value}, ) self.test_results[f"search_{search_type}"] = { "status": "PASS", "result": result, "message": f"Search with {search_type} successful", } print(f"✅ Search {search_type} test passed") except Exception as e: self.test_results[f"search_{search_type}"] = { "status": "FAIL", "error": str(e), "message": f"Search with {search_type} failed", } print(f"❌ Search {search_type} test failed: {e}") async def test_list_data(self): """Test the list_data functionality.""" print("\n🧪 Testing list_data functionality...") try: async with self.mcp_server_session() as session: # Test listing all datasets result = await session.call_tool("list_data", arguments={}) if result.content and len(result.content) > 0: content = result.content[0].text # Check if the output contains expected elements if "Available Datasets:" in content or "No datasets found" in content: self.test_results["list_data_all"] = { "status": "PASS", "result": content[:200] + "..." if len(content) > 200 else content, "message": "list_data (all datasets) successful", } print("✅ list_data (all datasets) test passed") # If there are datasets, try to list data for the first one if "Dataset ID:" in content: # Extract the first dataset ID from the output lines = content.split("\n") dataset_id = None for line in lines: if "Dataset ID:" in line: dataset_id = line.split("Dataset ID:")[1].strip() break if dataset_id: # Test listing data for specific dataset specific_result = await session.call_tool( "list_data", arguments={"dataset_id": dataset_id} ) if specific_result.content and len(specific_result.content) > 0: specific_content = specific_result.content[0].text if "Dataset:" in specific_content: self.test_results["list_data_specific"] = { "status": "PASS", "result": specific_content[:200] + "..." if len(specific_content) > 200 else specific_content, "message": "list_data (specific dataset) successful", } print("✅ list_data (specific dataset) test passed") else: raise Exception( "Specific dataset listing returned unexpected format" ) else: raise Exception("Specific dataset listing returned no content") else: raise Exception("list_data returned unexpected format") else: raise Exception("list_data returned no content") except Exception as e: self.test_results["list_data"] = { "status": "FAIL", "error": str(e), "message": "list_data test failed", } print(f"❌ list_data test failed: {e}") async def test_delete(self): """Test the delete functionality.""" print("\n🧪 Testing delete functionality...") try: async with self.mcp_server_session() as session: # First, let's get available data to delete list_result = await session.call_tool("list_data", arguments={}) if not (list_result.content and len(list_result.content) > 0): raise Exception("No data available for delete test - list_data returned empty") content = list_result.content[0].text # Look for data IDs and dataset IDs in the content lines = content.split("\n") dataset_id = None data_id = None for line in lines: if "Dataset ID:" in line: dataset_id = line.split("Dataset ID:")[1].strip() elif "Data ID:" in line: data_id = line.split("Data ID:")[1].strip() break # Get the first data item if dataset_id and data_id: # Test soft delete (default) delete_result = await session.call_tool( "delete", arguments={"data_id": data_id, "dataset_id": dataset_id, "mode": "soft"}, ) if delete_result.content and len(delete_result.content) > 0: delete_content = delete_result.content[0].text if "Delete operation completed successfully" in delete_content: self.test_results["delete_soft"] = { "status": "PASS", "result": delete_content[:200] + "..." if len(delete_content) > 200 else delete_content, "message": "delete (soft mode) successful", } print("✅ delete (soft mode) test passed") else: # Check if it's an expected error (like document not found) if "not found" in delete_content.lower(): self.test_results["delete_soft"] = { "status": "PASS", "result": delete_content, "message": "delete test passed with expected 'not found' error", } print("✅ delete test passed (expected 'not found' error)") else: raise Exception( f"Delete returned unexpected content: {delete_content}" ) else: raise Exception("Delete returned no content") else: # Test with invalid UUIDs to check error handling invalid_result = await session.call_tool( "delete", arguments={ "data_id": "invalid-uuid", "dataset_id": "another-invalid-uuid", "mode": "soft", }, ) if invalid_result.content and len(invalid_result.content) > 0: invalid_content = invalid_result.content[0].text if "Invalid UUID format" in invalid_content: self.test_results["delete_error_handling"] = { "status": "PASS", "result": invalid_content, "message": "delete error handling works correctly", } print("✅ delete error handling test passed") else: raise Exception(f"Expected UUID error not found: {invalid_content}") else: raise Exception("Delete error test returned no content") except Exception as e: self.test_results["delete"] = { "status": "FAIL", "error": str(e), "message": "delete test failed", } print(f"❌ delete test failed: {e}") def test_utility_functions(self): """Test utility functions.""" print("\n🧪 Testing utility functions...") # Test node_to_string try: test_node = {"id": "test_id", "name": "test_name", "type": "test_type"} result = node_to_string(test_node) expected = 'Node(id: "test_id", name: "test_name")' if result == expected: self.test_results["node_to_string"] = { "status": "PASS", "result": result, "message": "node_to_string function works correctly", } print("✅ node_to_string test passed") else: self.test_results["node_to_string"] = { "status": "FAIL", "result": result, "expected": expected, "message": "node_to_string function output mismatch", } print(f"❌ node_to_string test failed: expected {expected}, got {result}") except Exception as e: self.test_results["node_to_string"] = { "status": "FAIL", "error": str(e), "message": "node_to_string test failed", } print(f"❌ node_to_string test failed: {e}") # Test retrieved_edges_to_string try: test_triplet = [ ( {"id": "node1", "name": "Node1"}, {"relationship_name": "CONNECTS_TO"}, {"id": "node2", "name": "Node2"}, ) ] result = retrieved_edges_to_string(test_triplet) expected = ( 'Node(id: "node1", name: "Node1") CONNECTS_TO Node(id: "node2", name: "Node2")' ) if result == expected: self.test_results["retrieved_edges_to_string"] = { "status": "PASS", "result": result, "message": "retrieved_edges_to_string function works correctly", } print("✅ retrieved_edges_to_string test passed") else: self.test_results["retrieved_edges_to_string"] = { "status": "FAIL", "result": result, "expected": expected, "message": "retrieved_edges_to_string function output mismatch", } print( f"❌ retrieved_edges_to_string test failed: expected {expected}, got {result}" ) except Exception as e: self.test_results["retrieved_edges_to_string"] = { "status": "FAIL", "error": str(e), "message": "retrieved_edges_to_string test failed", } print(f"❌ retrieved_edges_to_string test failed: {e}") def test_load_class_function(self): """Test load_class function.""" print("\n🧪 Testing load_class function...") try: # Create a temporary Python file with a test class test_module_file = os.path.join(self.test_data_dir, "test_model.py") with open(test_module_file, "w") as f: f.write(""" class TestModel: def __init__(self): self.name = "TestModel" def get_name(self): return self.name """) # Test loading the class loaded_class = load_class(test_module_file, "TestModel") instance = loaded_class() if hasattr(instance, "get_name") and instance.get_name() == "TestModel": self.test_results["load_class"] = { "status": "PASS", "message": "load_class function works correctly", } print("✅ load_class test passed") else: self.test_results["load_class"] = { "status": "FAIL", "message": "load_class function did not load class correctly", } print("❌ load_class test failed: class not loaded correctly") except Exception as e: self.test_results["load_class"] = { "status": "FAIL", "error": str(e), "message": "load_class test failed", } print(f"❌ load_class test failed: {e}") async def run_all_tests(self): """Run all tests.""" print("🚀 Starting Cognee MCP Server Test Suite") print("=" * 50) await self.setup() # Test MCP server integration first await self.test_mcp_server_startup_and_tools() # Run tests in logical order await self.test_prune() # Start with clean slate # Test cognify twice to make sure updating a dataset with new docs is working as expected await self.test_cognify( test_text="Artificial Intelligence is transforming the world through machine learning and deep learning technologies.", test_name="Cognify1", ) await self.test_cognify( test_text="Natural language processing (NLP) is an interdisciplinary subfield of computer science and information retrieval.", test_name="Cognify2", ) await self.test_codify() await self.test_cognee_add_developer_rules() # Test list_data and delete functionality await self.test_list_data() await self.test_delete() await self.test_search_functionality() # Test utility functions (synchronous) self.test_utility_functions() self.test_load_class_function() await self.cleanup() # Print summary self.print_test_summary() def print_test_summary(self): """Print test results summary.""" print("\n" + "=" * 50) print("📊 TEST RESULTS SUMMARY") print("=" * 50) passed = 0 failed = 0 for test_name, result in self.test_results.items(): if result["status"] == "PASS": status_emoji = "✅" passed += 1 else: status_emoji = "❌" failed += 1 print(f"{status_emoji} {test_name}: {result['status']}") if result["status"] == "FAIL" and "error" in result: print(f" Error: {result['error']}") print("\n" + "-" * 50) total_tests = passed + failed print(f"Total Tests: {total_tests}") print(f"Passed: {passed}") print(f"Failed: {failed}") print(f"Success Rate: {(passed / total_tests * 100):.1f}%") if failed > 0: print(f"\n ⚠️ {failed} test(s) failed - review results above for details") async def main(): """Main function to run the test suite.""" client = CogneeTestClient() await client.run_all_tests() if __name__ == "__main__": from logging import ERROR logger = setup_logging(log_level=ERROR) asyncio.run(main())