diff --git a/examples/quickstart/quickstart_falkordb.py b/examples/quickstart/quickstart_falkordb.py index d62c39d5..19b7cce8 100644 --- a/examples/quickstart/quickstart_falkordb.py +++ b/examples/quickstart/quickstart_falkordb.py @@ -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: diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py index 524b73b8..3ca6c316 100644 --- a/graphiti_core/graphiti.py +++ b/graphiti_core/graphiti.py @@ -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 diff --git a/graphiti_core/nodes.py b/graphiti_core/nodes.py index d162718d..3df7069c 100644 --- a/graphiti_core/nodes.py +++ b/graphiti_core/nodes.py @@ -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, diff --git a/graphiti_core/utils/maintenance/graph_data_operations.py b/graphiti_core/utils/maintenance/graph_data_operations.py index 89a9cc62..cccb9131 100644 --- a/graphiti_core/utils/maintenance/graph_data_operations.py +++ b/graphiti_core/utils/maintenance/graph_data_operations.py @@ -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'], diff --git a/tests/driver/test_falkordb_driver.py b/tests/driver/test_falkordb_driver.py index 38447927..70c0a0f2 100644 --- a/tests/driver/test_falkordb_driver.py +++ b/tests/driver/test_falkordb_driver.py @@ -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}') diff --git a/tests/test_graphiti_falkordb_int.py b/tests/test_graphiti_falkordb_int.py index 61564be0..c41ef47e 100644 --- a/tests/test_graphiti_falkordb_int.py +++ b/tests/test_graphiti_falkordb_int.py @@ -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) diff --git a/tests/test_graphiti_int.py b/tests/test_graphiti_int.py index a0ab2877..0e2522fc 100644 --- a/tests/test_graphiti_int.py +++ b/tests/test_graphiti_int.py @@ -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) diff --git a/tests/test_node_falkordb_int.py b/tests/test_node_falkordb_int.py index 62db9666..35eaa578 100644 --- a/tests/test_node_falkordb_int.py +++ b/tests/test_node_falkordb_int.py @@ -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)