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

reformat (#655)

Browse source
This commit is contained in:
Preston Rasmussen 2025-07-01 12:26:15 -04:00 • committed by GitHub
parent 537e19f44a
commit 71360d91fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 69 additions and 99 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
examples/quickstart/quickstart_falkordb.py
View file

@ -49,7 +49,7 @@ load_dotenv()
# Make sure FalkorDB (on-premises) is running — see https://docs.falkordb.com/
# By default, FalkorDB does not require a username or password,
# but you can set them via environment variables for added security.
#
#
# If you're using FalkorDB Cloud, set the environment variables accordingly.
# For on-premises use, you can leave them as None or set them to your preferred values.
#
@ -61,6 +61,7 @@ falkor_password = os.environ.get('FALKORDB_PASSWORD', None)
falkor_host = os.environ.get('FALKORDB_HOST', 'localhost')
falkor_port = os.environ.get('FALKORDB_PORT', '6379')
async def main():
#################################################
# INITIALIZATION
@ -71,7 +72,9 @@ async def main():
#################################################
# Initialize Graphiti with FalkorDB connection
falkor_driver = FalkorDriver(host=falkor_host, port=falkor_port, username=falkor_username, password=falkor_password)
falkor_driver = FalkorDriver(
host=falkor_host, port=falkor_port, username=falkor_username, password=falkor_password
)
graphiti = Graphiti(graph_driver=falkor_driver)
try:

2
graphiti_core/graphiti.py
View file

@ -166,7 +166,7 @@ class Graphiti:
self.driver = graph_driver
else:
if uri is None:
raise ValueError("uri must be provided when graph_driver is None")
raise ValueError('uri must be provided when graph_driver is None')
self.driver = Neo4jDriver(uri, user, password)
self.database = DEFAULT_DATABASE

8
graphiti_core/nodes.py
View file

@ -542,12 +542,12 @@ class CommunityNode(Node):
def get_episodic_node_from_record(record: Any) -> EpisodicNode:
created_at = parse_db_date(record['created_at'])
valid_at = parse_db_date(record['valid_at'])
if created_at is None:
raise ValueError(f"created_at cannot be None for episode {record.get('uuid', 'unknown')}")
raise ValueError(f'created_at cannot be None for episode {record.get("uuid", "unknown")}')
if valid_at is None:
raise ValueError(f"valid_at cannot be None for episode {record.get('uuid', 'unknown')}")
raise ValueError(f'valid_at cannot be None for episode {record.get("uuid", "unknown")}')
return EpisodicNode(
content=record['content'],
created_at=created_at,

3
graphiti_core/utils/maintenance/graph_data_operations.py
View file

@ -140,7 +140,8 @@ async def retrieve_episodes(
episodes = [
EpisodicNode(
content=record['content'],
created_at=parse_db_date(record['created_at']) or datetime.min.replace(tzinfo=timezone.utc),
created_at=parse_db_date(record['created_at'])
or datetime.min.replace(tzinfo=timezone.utc),
valid_at=parse_db_date(record['valid_at']) or datetime.min.replace(tzinfo=timezone.utc),
uuid=record['uuid'],
group_id=record['group_id'],

106
tests/driver/test_falkordb_driver.py
View file

@ -35,7 +35,7 @@ except ImportError:
class TestFalkorDriver:
"""Comprehensive test suite for FalkorDB driver."""
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def setup_method(self):
"""Set up test fixtures."""
self.mock_client = MagicMock()
@ -43,25 +43,19 @@ class TestFalkorDriver:
self.driver = FalkorDriver()
self.driver.client = self.mock_client
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_init_with_connection_params(self):
"""Test initialization with connection parameters."""
with patch('graphiti_core.driver.falkordb_driver.FalkorDB') as mock_falkor_db:
driver = FalkorDriver(
host='test-host',
port='1234',
username='test-user',
password='test-pass'
host='test-host', port='1234', username='test-user', password='test-pass'
)
assert driver.provider == 'falkordb'
mock_falkor_db.assert_called_once_with(
host='test-host',
port='1234',
username='test-user',
password='test-pass'
host='test-host', port='1234', username='test-user', password='test-pass'
)
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_init_with_falkor_db_instance(self):
"""Test initialization with a FalkorDB instance."""
with patch('graphiti_core.driver.falkordb_driver.FalkorDB') as mock_falkor_db_class:
@ -71,12 +65,12 @@ class TestFalkorDriver:
assert driver.client is mock_falkor_db
mock_falkor_db_class.assert_not_called()
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_provider(self):
"""Test driver provider identification."""
assert self.driver.provider == 'falkordb'
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_get_graph_with_name(self):
"""Test _get_graph with specific graph name."""
mock_graph = MagicMock()
@ -87,7 +81,7 @@ class TestFalkorDriver:
self.mock_client.select_graph.assert_called_once_with('test_graph')
assert result is mock_graph
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_get_graph_with_none_defaults_to_default_database(self):
"""Test _get_graph with None defaults to DEFAULT_DATABASE."""
mock_graph = MagicMock()
@ -99,7 +93,7 @@ class TestFalkorDriver:
assert result is mock_graph
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_execute_query_success(self):
"""Test successful query execution."""
mock_graph = MagicMock()
@ -110,16 +104,11 @@ class TestFalkorDriver:
self.mock_client.select_graph.return_value = mock_graph
result = await self.driver.execute_query(
'MATCH (n) RETURN n',
param1='value1',
database_='test_db'
'MATCH (n) RETURN n', param1='value1', database_='test_db'
)
self.mock_client.select_graph.assert_called_once_with('test_db')
mock_graph.query.assert_called_once_with(
'MATCH (n) RETURN n',
{'param1': 'value1'}
)
mock_graph.query.assert_called_once_with('MATCH (n) RETURN n', {'param1': 'value1'})
result_set, header, summary = result
assert result_set == [{'column1': 'row1col1', 'column2': 'row1col2'}]
@ -127,7 +116,7 @@ class TestFalkorDriver:
assert summary is None
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_execute_query_handles_index_already_exists_error(self):
"""Test handling of 'already indexed' error."""
mock_graph = MagicMock()
@ -141,7 +130,7 @@ class TestFalkorDriver:
assert result is None
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_execute_query_propagates_other_exceptions(self):
"""Test that other exceptions are properly propagated."""
mock_graph = MagicMock()
@ -155,7 +144,7 @@ class TestFalkorDriver:
mock_logger.error.assert_called_once()
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_execute_query_converts_datetime_parameters(self):
"""Test that datetime objects in kwargs are converted to ISO strings."""
mock_graph = MagicMock()
@ -168,14 +157,13 @@ class TestFalkorDriver:
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
await self.driver.execute_query(
'CREATE (n:Node) SET n.created_at = $created_at',
created_at=test_datetime
'CREATE (n:Node) SET n.created_at = $created_at', created_at=test_datetime
)
call_args = mock_graph.query.call_args[0]
assert call_args[1]['created_at'] == test_datetime.isoformat()
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_session_creation(self):
"""Test session creation with specific database."""
mock_graph = MagicMock()
@ -187,7 +175,7 @@ class TestFalkorDriver:
assert session.graph is mock_graph
self.mock_client.select_graph.assert_called_once_with('test_db')
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_session_creation_with_none_uses_default_database(self):
"""Test session creation with None uses default database."""
mock_graph = MagicMock()
@ -199,7 +187,7 @@ class TestFalkorDriver:
self.mock_client.select_graph.assert_called_once_with(DEFAULT_DATABASE)
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_close_calls_connection_close(self):
"""Test driver close method calls connection close."""
mock_connection = MagicMock()
@ -211,10 +199,10 @@ class TestFalkorDriver:
with patch('builtins.hasattr') as mock_hasattr:
# hasattr(self.client, 'aclose') returns False
# hasattr(self.client.connection, 'aclose') returns False
# hasattr(self.client.connection, 'aclose') returns False
# hasattr(self.client.connection, 'close') returns True
mock_hasattr.side_effect = lambda obj, attr: (
attr == 'close' and obj is mock_connection
attr == 'close' and obj is mock_connection
)
await self.driver.close()
@ -222,42 +210,41 @@ class TestFalkorDriver:
mock_connection.close.assert_called_once()
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_delete_all_indexes(self):
"""Test delete_all_indexes method."""
with patch.object(self.driver, 'execute_query', new_callable=AsyncMock) as mock_execute:
await self.driver.delete_all_indexes('test_db')
mock_execute.assert_called_once_with(
'CALL db.indexes() YIELD name DROP INDEX name',
database_='test_db'
'CALL db.indexes() YIELD name DROP INDEX name', database_='test_db'
)
class TestFalkorDriverSession:
"""Test FalkorDB driver session functionality."""
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def setup_method(self):
"""Set up test fixtures."""
self.mock_graph = MagicMock()
self.session = FalkorDriverSession(self.mock_graph)
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_session_async_context_manager(self):
"""Test session can be used as async context manager."""
async with self.session as s:
assert s is self.session
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_close_method(self):
"""Test session close method doesn't raise exceptions."""
await self.session.close() # Should not raise
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_execute_write_passes_session_and_args(self):
"""Test execute_write method passes session and arguments correctly."""
@ -267,37 +254,30 @@ class TestFalkorDriverSession:
assert kwargs == {'key': 'value'}
return 'result'
result = await self.session.execute_write(
test_func, 'arg1', 'arg2', key='value'
)
result = await self.session.execute_write(test_func, 'arg1', 'arg2', key='value')
assert result == 'result'
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_run_single_query_with_parameters(self):
"""Test running a single query with parameters."""
self.mock_graph.query = AsyncMock()
await self.session.run(
'MATCH (n) RETURN n',
param1='value1',
param2='value2'
)
await self.session.run('MATCH (n) RETURN n', param1='value1', param2='value2')
self.mock_graph.query.assert_called_once_with(
'MATCH (n) RETURN n',
{'param1': 'value1', 'param2': 'value2'}
'MATCH (n) RETURN n', {'param1': 'value1', 'param2': 'value2'}
)
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_run_multiple_queries_as_list(self):
"""Test running multiple queries passed as list."""
self.mock_graph.query = AsyncMock()
queries = [
('MATCH (n) RETURN n', {'param1': 'value1'}),
('CREATE (n:Node)', {'param2': 'value2'})
('CREATE (n:Node)', {'param2': 'value2'}),
]
await self.session.run(queries)
@ -308,15 +288,14 @@ class TestFalkorDriverSession:
assert calls[1][0] == ('CREATE (n:Node)', {'param2': 'value2'})
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_run_converts_datetime_objects_to_iso_strings(self):
"""Test that datetime objects are converted to ISO strings."""
self.mock_graph.query = AsyncMock()
test_datetime = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
await self.session.run(
'CREATE (n:Node) SET n.created_at = $created_at',
created_at=test_datetime
'CREATE (n:Node) SET n.created_at = $created_at', created_at=test_datetime
)
self.mock_graph.query.assert_called_once()
@ -327,7 +306,7 @@ class TestFalkorDriverSession:
class TestDatetimeConversion:
"""Test datetime conversion utility function."""
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_convert_datetime_dict(self):
"""Test datetime conversion in nested dictionary."""
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
@ -336,10 +315,7 @@ class TestDatetimeConversion:
input_dict = {
'string_val': 'test',
'datetime_val': test_datetime,
'nested_dict': {
'nested_datetime': test_datetime,
'nested_string': 'nested_test'
}
'nested_dict': {'nested_datetime': test_datetime, 'nested_string': 'nested_test'},
}
result = convert_datetimes_to_strings(input_dict)
@ -349,7 +325,7 @@ class TestDatetimeConversion:
assert result['nested_dict']['nested_datetime'] == test_datetime.isoformat()
assert result['nested_dict']['nested_string'] == 'nested_test'
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_convert_datetime_list_and_tuple(self):
"""Test datetime conversion in lists and tuples."""
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
@ -370,7 +346,7 @@ class TestDatetimeConversion:
assert result_tuple[0] == 'test'
assert result_tuple[1] == test_datetime.isoformat()
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_convert_single_datetime(self):
"""Test datetime conversion for single datetime object."""
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
@ -379,7 +355,7 @@ class TestDatetimeConversion:
result = convert_datetimes_to_strings(test_datetime)
assert result == test_datetime.isoformat()
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
def test_convert_other_types_unchanged(self):
"""Test that non-datetime types are returned unchanged."""
from graphiti_core.driver.falkordb_driver import convert_datetimes_to_strings
@ -396,7 +372,7 @@ class TestFalkorDriverIntegration:
@pytest.mark.integration
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_basic_integration_with_real_falkordb(self):
"""Basic integration test with real FalkorDB instance."""
pytest.importorskip('falkordb')
@ -418,4 +394,4 @@ class TestFalkorDriverIntegration:
await driver.close()
except Exception as e:
pytest.skip(f"FalkorDB not available for integration test: {e}")
pytest.skip(f'FalkorDB not available for integration test: {e}')

14
tests/test_graphiti_falkordb_int.py
View file

@ -71,15 +71,12 @@ def setup_logging():
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_graphiti_falkordb_init():
logger = setup_logging()
falkor_driver = FalkorDriver(
host=FALKORDB_HOST,
port=FALKORDB_PORT,
username=FALKORDB_USER,
password=FALKORDB_PASSWORD
host=FALKORDB_HOST, port=FALKORDB_PORT, username=FALKORDB_USER, password=FALKORDB_PASSWORD
)
graphiti = Graphiti(graph_driver=falkor_driver)
@ -94,13 +91,10 @@ async def test_graphiti_falkordb_init():
@pytest.mark.asyncio
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_graph_falkordb_integration():
falkor_driver = FalkorDriver(
host=FALKORDB_HOST,
port=FALKORDB_PORT,
username=FALKORDB_USER,
password=FALKORDB_PASSWORD
host=FALKORDB_HOST, port=FALKORDB_PORT, username=FALKORDB_USER, password=FALKORDB_PASSWORD
)
client = Graphiti(graph_driver=falkor_driver)

7
tests/test_graphiti_int.py
View file

@ -26,7 +26,9 @@ from graphiti_core.edges import EntityEdge, EpisodicEdge
from graphiti_core.graphiti import Graphiti
from graphiti_core.helpers import semaphore_gather
from graphiti_core.nodes import EntityNode, EpisodicNode
from graphiti_core.search.search_filters import ComparisonOperator, DateFilter, SearchFilters
from graphiti_core.search.search_helpers import search_results_to_context_string
from graphiti_core.utils.datetime_utils import utc_now
pytestmark = pytest.mark.integration
@ -64,8 +66,11 @@ def setup_logging():
async def test_graphiti_init():
logger = setup_logging()
graphiti = Graphiti(NEO4J_URI, NEO4j_USER, NEO4j_PASSWORD)
search_filter = SearchFilters(
created_at=[[DateFilter(date=utc_now(), comparison_operator=ComparisonOperator.less_than)]]
)
results = await graphiti.search_(query='Who is the user?')
results = await graphiti.search_(query='Who is Tania?', search_filter=search_filter)
pretty_results = search_results_to_context_string(results)

21
tests/test_node_falkordb_int.py
View file

@ -80,13 +80,10 @@ def sample_community_node():
@pytest.mark.asyncio
@pytest.mark.integration
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_entity_node_save_get_and_delete(sample_entity_node):
falkor_driver = FalkorDriver(
host=FALKORDB_HOST,
port=FALKORDB_PORT,
username=FALKORDB_USER,
password=FALKORDB_PASSWORD
host=FALKORDB_HOST, port=FALKORDB_PORT, username=FALKORDB_USER, password=FALKORDB_PASSWORD
)
await sample_entity_node.save(falkor_driver)
@ -102,13 +99,10 @@ async def test_entity_node_save_get_and_delete(sample_entity_node):
@pytest.mark.asyncio
@pytest.mark.integration
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_community_node_save_get_and_delete(sample_community_node):
falkor_driver = FalkorDriver(
host=FALKORDB_HOST,
port=FALKORDB_PORT,
username=FALKORDB_USER,
password=FALKORDB_PASSWORD
host=FALKORDB_HOST, port=FALKORDB_PORT, username=FALKORDB_USER, password=FALKORDB_PASSWORD
)
await sample_community_node.save(falkor_driver)
@ -125,13 +119,10 @@ async def test_community_node_save_get_and_delete(sample_community_node):
@pytest.mark.asyncio
@pytest.mark.integration
@unittest.skipIf(not HAS_FALKORDB, "FalkorDB is not installed")
@unittest.skipIf(not HAS_FALKORDB, 'FalkorDB is not installed')
async def test_episodic_node_save_get_and_delete(sample_episodic_node):
falkor_driver = FalkorDriver(
host=FALKORDB_HOST,
port=FALKORDB_PORT,
username=FALKORDB_USER,
password=FALKORDB_PASSWORD
host=FALKORDB_HOST, port=FALKORDB_PORT, username=FALKORDB_USER, password=FALKORDB_PASSWORD
)
await sample_episodic_node.save(falkor_driver)