From 71f1f66d112bca26c54109809293fc59855bf9ac Mon Sep 17 00:00:00 2001 From: Preston Rasmussen <109292228+prasmussen15@users.noreply.github.com> Date: Sun, 26 Oct 2025 22:07:36 -0400 Subject: [PATCH] Search client update (#1026) * update bulk interfae handling * bump version * format --- graphiti_core/driver/falkordb_driver.py | 10 +++++----- graphiti_core/driver/neo4j_driver.py | 4 ++-- graphiti_core/utils/bulk_utils.py | 6 ++---- mcp_server/graphiti_mcp_server.py | 16 ++++++++++++---- pyproject.toml | 2 +- uv.lock | 4 ++-- 6 files changed, 24 insertions(+), 18 deletions(-) diff --git a/graphiti_core/driver/falkordb_driver.py b/graphiti_core/driver/falkordb_driver.py index 793f0545..d0b4ffe8 100644 --- a/graphiti_core/driver/falkordb_driver.py +++ b/graphiti_core/driver/falkordb_driver.py @@ -14,8 +14,8 @@ See the License for the specific language governing permissions and limitations under the License. """ -import datetime import asyncio +import datetime import logging from typing import TYPE_CHECKING, Any @@ -231,17 +231,17 @@ class FalkorDriver(GraphDriver): """ cloned = FalkorDriver(falkor_db=self.client, database=database) - return cloned + return cloned async def health_check(self) -> None: """Check FalkorDB connectivity by running a simple query.""" try: - await self.execute_query("MATCH (n) RETURN 1 LIMIT 1") + await self.execute_query('MATCH (n) RETURN 1 LIMIT 1') return None except Exception as e: - print(f"FalkorDB health check failed: {e}") + print(f'FalkorDB health check failed: {e}') raise - + @staticmethod def convert_datetimes_to_strings(obj): if isinstance(obj, dict): diff --git a/graphiti_core/driver/neo4j_driver.py b/graphiti_core/driver/neo4j_driver.py index 4e114943..4a0baf79 100644 --- a/graphiti_core/driver/neo4j_driver.py +++ b/graphiti_core/driver/neo4j_driver.py @@ -72,12 +72,12 @@ class Neo4jDriver(GraphDriver): return self.client.execute_query( 'CALL db.indexes() YIELD name DROP INDEX name', ) - + async def health_check(self) -> None: """Check Neo4j connectivity by running the driver's verify_connectivity method.""" try: await self.client.verify_connectivity() return None except Exception as e: - print(f"Neo4j health check failed: {e}") + print(f'Neo4j health check failed: {e}') raise diff --git a/graphiti_core/utils/bulk_utils.py b/graphiti_core/utils/bulk_utils.py index 049aa53e..4a861b1b 100644 --- a/graphiti_core/utils/bulk_utils.py +++ b/graphiti_core/utils/bulk_utils.py @@ -214,12 +214,10 @@ async def add_nodes_and_edges_bulk_tx( edges.append(edge_data) if driver.graph_operations_interface: - await driver.graph_operations_interface.episodic_node_save_bulk( - None, driver, tx, episodic_nodes - ) + await driver.graph_operations_interface.episodic_node_save_bulk(None, driver, tx, episodes) await driver.graph_operations_interface.node_save_bulk(None, driver, tx, nodes) await driver.graph_operations_interface.episodic_edge_save_bulk( - None, driver, tx, episodic_edges + None, driver, tx, [edge.model_dump() for edge in episodic_edges] ) await driver.graph_operations_interface.edge_save_bulk(None, driver, tx, edges) diff --git a/mcp_server/graphiti_mcp_server.py b/mcp_server/graphiti_mcp_server.py index 5a650b24..3919fd78 100644 --- a/mcp_server/graphiti_mcp_server.py +++ b/mcp_server/graphiti_mcp_server.py @@ -467,6 +467,7 @@ class Neo4jConfig(BaseModel): password=os.environ.get('NEO4J_PASSWORD', 'password'), ) + class FalkorConfig(BaseModel): """Configuration for FalkorDB database connection.""" @@ -483,6 +484,7 @@ class FalkorConfig(BaseModel): password = os.environ.get('FALKORDB_PASSWORD', '') return cls(host=host, port=port, user=user, password=password) + class GraphitiConfig(BaseModel): """Configuration for Graphiti client. @@ -504,7 +506,9 @@ class GraphitiConfig(BaseModel): """Create a configuration instance from environment variables.""" db_type = os.environ.get('DATABASE_TYPE') if not db_type: - raise ValueError('DATABASE_TYPE environment variable must be set (e.g., "neo4j" or "falkordb")') + raise ValueError( + 'DATABASE_TYPE environment variable must be set (e.g., "neo4j" or "falkordb")' + ) if db_type == 'neo4j': return cls( llm=GraphitiLLMConfig.from_env(), @@ -622,7 +626,9 @@ async def initialize_graphiti(): raise ValueError('NEO4J_URI, NEO4J_USER, and NEO4J_PASSWORD must be set') # Validate FalkorDB configuration - if config.database_type == 'falkordb' and (not config.falkordb.host or not config.falkordb.port): + if config.database_type == 'falkordb' and ( + not config.falkordb.host or not config.falkordb.port + ): raise ValueError('FALKORDB_HOST and FALKORDB_PORT must be set for FalkorDB') embedder_client = config.embedder.create_client() @@ -637,6 +643,7 @@ async def initialize_graphiti(): ) elif config.database_type == 'falkordb': from graphiti_core.driver.falkordb_driver import FalkorDriver + host = config.falkordb.host if hasattr(config.falkordb, 'host') else 'localhost' port = int(config.falkordb.port) if hasattr(config.falkordb, 'port') else 6379 username = config.falkordb.user or None @@ -1205,10 +1212,11 @@ async def get_status() -> StatusResponse: client = cast(Graphiti, graphiti_client) # Test database connection - await client.driver.health_check() # type: ignore # type: ignore + await client.driver.health_check() # type: ignore # type: ignore return StatusResponse( - status='ok', message=f'Graphiti MCP server is running and connected to {config.database_type}' + status='ok', + message=f'Graphiti MCP server is running and connected to {config.database_type}', ) except Exception as e: error_msg = str(e) diff --git a/pyproject.toml b/pyproject.toml index 1932e770..59e88e21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "graphiti-core" description = "A temporal graph building library" -version = "0.22.0" +version = "0.22.1pre1" authors = [ { name = "Paul Paliychuk", email = "paul@getzep.com" }, { name = "Preston Rasmussen", email = "preston@getzep.com" }, diff --git a/uv.lock b/uv.lock index a489a57d..2bf14336 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10, <4" resolution-markers = [ "python_full_version >= '3.14'", @@ -783,7 +783,7 @@ wheels = [ [[package]] name = "graphiti-core" -version = "0.22.0rc5" +version = "0.22.1rc1" source = { editable = "." } dependencies = [ { name = "diskcache" },