gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

update kuze and neptune with build indices function

Browse source
This commit is contained in:
prestonrasmussen 2025-11-05 11:35:41 -05:00
parent c144ff5995
commit 5d1ba3cf57
6 changed files with 93 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
graphiti_core/driver/kuzu_driver.py
View file

@ -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
]
)

2
graphiti_core/driver/neo4j_driver.py
View file

@ -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:

21
graphiti_core/driver/neptune_driver.py
View file

@ -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

8
graphiti_core/llm_client/openai_base_client.py
View file

@ -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

8
graphiti_core/llm_client/openai_client.py
View file

@ -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,

46
mcp_server/src/graphiti_mcp_server.py
View file

@ -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