From 5d1ba3cf579a45a97ae6178e02c30e8de7631e2d Mon Sep 17 00:00:00 2001 From: prestonrasmussen Date: Wed, 5 Nov 2025 11:35:41 -0500 Subject: [PATCH] update kuze and neptune with build indices function --- graphiti_core/driver/kuzu_driver.py | 38 ++++++++++++++- graphiti_core/driver/neo4j_driver.py | 2 +- graphiti_core/driver/neptune_driver.py | 21 ++++++++- .../llm_client/openai_base_client.py | 8 +++- graphiti_core/llm_client/openai_client.py | 8 +++- mcp_server/src/graphiti_mcp_server.py | 46 +++++++++---------- 6 files changed, 93 insertions(+), 30 deletions(-) diff --git a/graphiti_core/driver/kuzu_driver.py b/graphiti_core/driver/kuzu_driver.py index 8a04a4ac..d3c99c39 100644 --- a/graphiti_core/driver/kuzu_driver.py +++ b/graphiti_core/driver/kuzu_driver.py @@ -15,11 +15,13 @@ limitations under the License. """ import logging -from typing import Any +from typing import Any, LiteralString import kuzu from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider +from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices +from graphiti_core.helpers import semaphore_gather logger = logging.getLogger(__name__) @@ -174,3 +176,37 @@ class KuzuDriverSession(GraphDriverSession): else: await self.driver.execute_query(query, **kwargs) return None + + async def build_indices_and_constraints(self, delete_existing: bool = False): + if delete_existing: + await self.delete_all_indexes() + + range_indices: list[LiteralString] = get_range_indices(self.provider) + + # Skip creating fulltext indices if they already exist. Need to do this manually + # until Kuzu supports `IF NOT EXISTS` for indices. + result, _, _ = await self.execute_query('CALL SHOW_INDEXES() RETURN *;') + if len(result) > 0: + fulltext_indices = [] + + # Only load the `fts` extension if it's not already loaded, otherwise throw an error. + result, _, _ = await self.execute_query('CALL SHOW_LOADED_EXTENSIONS() RETURN *;') + if len(result) == 0: + fulltext_indices.insert( + 0, + """ + INSTALL fts; + LOAD fts; + """, + ) + + index_queries: list[LiteralString] = range_indices + fulltext_indices + + await semaphore_gather( + *[ + self.execute_query( + query, + ) + for query in index_queries + ] + ) diff --git a/graphiti_core/driver/neo4j_driver.py b/graphiti_core/driver/neo4j_driver.py index 5d85c2b0..4fa73f57 100644 --- a/graphiti_core/driver/neo4j_driver.py +++ b/graphiti_core/driver/neo4j_driver.py @@ -106,7 +106,7 @@ class Neo4jDriver(GraphDriver): for query in index_queries ] ) - + async def health_check(self) -> None: """Check Neo4j connectivity by running the driver's verify_connectivity method.""" try: diff --git a/graphiti_core/driver/neptune_driver.py b/graphiti_core/driver/neptune_driver.py index d43d20a1..dd98cc32 100644 --- a/graphiti_core/driver/neptune_driver.py +++ b/graphiti_core/driver/neptune_driver.py @@ -18,13 +18,15 @@ import asyncio import datetime import logging from collections.abc import Coroutine -from typing import Any +from typing import Any, LiteralString import boto3 from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider +from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices +from graphiti_core.helpers import semaphore_gather logger = logging.getLogger(__name__) DEFAULT_SIZE = 10 @@ -267,6 +269,23 @@ class NeptuneDriver(GraphDriver): return 0 + async def build_indices_and_constraints(self, delete_existing: bool = False): + if delete_existing: + await self.delete_all_indexes() + + range_indices: list[LiteralString] = get_range_indices(self.provider) + + index_queries: list[LiteralString] = range_indices + + await semaphore_gather( + *[ + self.execute_query( + query, + ) + for query in index_queries + ] + ) + class NeptuneDriverSession(GraphDriverSession): provider = GraphProvider.NEPTUNE diff --git a/graphiti_core/llm_client/openai_base_client.py b/graphiti_core/llm_client/openai_base_client.py index 93e9c598..9580ff03 100644 --- a/graphiti_core/llm_client/openai_base_client.py +++ b/graphiti_core/llm_client/openai_base_client.py @@ -166,13 +166,17 @@ class BaseOpenAIClient(LLMClient): except openai.RateLimitError as e: raise RateLimitError from e except openai.AuthenticationError as e: - logger.error(f'OpenAI Authentication Error: {e}. Please verify your API key is correct.') + logger.error( + f'OpenAI Authentication Error: {e}. Please verify your API key is correct.' + ) raise except Exception as e: # Provide more context for connection errors error_msg = str(e) if 'Connection error' in error_msg or 'connection' in error_msg.lower(): - logger.error(f'Connection error communicating with OpenAI API. Please check your network connection and API key. Error: {e}') + logger.error( + f'Connection error communicating with OpenAI API. Please check your network connection and API key. Error: {e}' + ) else: logger.error(f'Error in generating LLM response: {e}') raise diff --git a/graphiti_core/llm_client/openai_client.py b/graphiti_core/llm_client/openai_client.py index a586f9f0..83410ff1 100644 --- a/graphiti_core/llm_client/openai_client.py +++ b/graphiti_core/llm_client/openai_client.py @@ -74,7 +74,9 @@ class OpenAIClient(BaseOpenAIClient): ): """Create a structured completion using OpenAI's beta parse API.""" # Reasoning models (gpt-5 family) don't support temperature - is_reasoning_model = model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3') + is_reasoning_model = ( + model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3') + ) response = await self.client.responses.parse( model=model, @@ -100,7 +102,9 @@ class OpenAIClient(BaseOpenAIClient): ): """Create a regular completion with JSON format.""" # Reasoning models (gpt-5 family) don't support temperature - is_reasoning_model = model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3') + is_reasoning_model = ( + model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3') + ) return await self.client.chat.completions.create( model=model, diff --git a/mcp_server/src/graphiti_mcp_server.py b/mcp_server/src/graphiti_mcp_server.py index 0c9a568a..a3bc1910 100644 --- a/mcp_server/src/graphiti_mcp_server.py +++ b/mcp_server/src/graphiti_mcp_server.py @@ -245,35 +245,35 @@ class GraphitiService: db_provider = self.config.database.provider if db_provider.lower() == 'falkordb': raise RuntimeError( - f"\n{'='*70}\n" - f"Database Connection Error: FalkorDB is not running\n" - f"{'='*70}\n\n" - f"FalkorDB at {db_config['host']}:{db_config['port']} is not accessible.\n\n" - f"To start FalkorDB:\n" - f" - Using Docker Compose: cd mcp_server && docker compose up\n" - f" - Or run FalkorDB manually: docker run -p 6379:6379 falkordb/falkordb\n\n" - f"{'='*70}\n" + f'\n{"=" * 70}\n' + f'Database Connection Error: FalkorDB is not running\n' + f'{"=" * 70}\n\n' + f'FalkorDB at {db_config["host"]}:{db_config["port"]} is not accessible.\n\n' + f'To start FalkorDB:\n' + f' - Using Docker Compose: cd mcp_server && docker compose up\n' + f' - Or run FalkorDB manually: docker run -p 6379:6379 falkordb/falkordb\n\n' + f'{"=" * 70}\n' ) from db_error elif db_provider.lower() == 'neo4j': raise RuntimeError( - f"\n{'='*70}\n" - f"Database Connection Error: Neo4j is not running\n" - f"{'='*70}\n\n" - f"Neo4j at {db_config.get('uri', 'unknown')} is not accessible.\n\n" - f"To start Neo4j:\n" - f" - Using Docker Compose: cd mcp_server && docker compose -f docker/docker-compose-neo4j.yml up\n" - f" - Or install Neo4j Desktop from: https://neo4j.com/download/\n" - f" - Or run Neo4j manually: docker run -p 7474:7474 -p 7687:7687 neo4j:latest\n\n" - f"{'='*70}\n" + f'\n{"=" * 70}\n' + f'Database Connection Error: Neo4j is not running\n' + f'{"=" * 70}\n\n' + f'Neo4j at {db_config.get("uri", "unknown")} is not accessible.\n\n' + f'To start Neo4j:\n' + f' - Using Docker Compose: cd mcp_server && docker compose -f docker/docker-compose-neo4j.yml up\n' + f' - Or install Neo4j Desktop from: https://neo4j.com/download/\n' + f' - Or run Neo4j manually: docker run -p 7474:7474 -p 7687:7687 neo4j:latest\n\n' + f'{"=" * 70}\n' ) from db_error else: raise RuntimeError( - f"\n{'='*70}\n" - f"Database Connection Error: {db_provider} is not running\n" - f"{'='*70}\n\n" - f"{db_provider} at {db_config.get('uri', 'unknown')} is not accessible.\n\n" - f"Please ensure {db_provider} is running and accessible.\n\n" - f"{'='*70}\n" + f'\n{"=" * 70}\n' + f'Database Connection Error: {db_provider} is not running\n' + f'{"=" * 70}\n\n' + f'{db_provider} at {db_config.get("uri", "unknown")} is not accessible.\n\n' + f'Please ensure {db_provider} is running and accessible.\n\n' + f'{"=" * 70}\n' ) from db_error # Re-raise other errors raise