graphiti-core 0.20.4__py3-none-any.whl → 0.21.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of graphiti-core might be problematic. Click here for more details.

Files changed (39) hide show
  1. graphiti_core/driver/driver.py +28 -0
  2. graphiti_core/driver/falkordb_driver.py +112 -0
  3. graphiti_core/driver/kuzu_driver.py +1 -0
  4. graphiti_core/driver/neo4j_driver.py +10 -2
  5. graphiti_core/driver/neptune_driver.py +4 -6
  6. graphiti_core/edges.py +67 -7
  7. graphiti_core/embedder/client.py +2 -1
  8. graphiti_core/graph_queries.py +35 -6
  9. graphiti_core/graphiti.py +27 -23
  10. graphiti_core/graphiti_types.py +0 -1
  11. graphiti_core/helpers.py +2 -2
  12. graphiti_core/llm_client/client.py +19 -4
  13. graphiti_core/llm_client/gemini_client.py +4 -2
  14. graphiti_core/llm_client/openai_base_client.py +3 -2
  15. graphiti_core/llm_client/openai_generic_client.py +3 -2
  16. graphiti_core/models/edges/edge_db_queries.py +36 -16
  17. graphiti_core/models/nodes/node_db_queries.py +30 -10
  18. graphiti_core/nodes.py +126 -25
  19. graphiti_core/prompts/dedupe_edges.py +40 -29
  20. graphiti_core/prompts/dedupe_nodes.py +51 -34
  21. graphiti_core/prompts/eval.py +3 -3
  22. graphiti_core/prompts/extract_edges.py +17 -9
  23. graphiti_core/prompts/extract_nodes.py +10 -9
  24. graphiti_core/prompts/prompt_helpers.py +3 -3
  25. graphiti_core/prompts/summarize_nodes.py +5 -5
  26. graphiti_core/search/search_filters.py +53 -0
  27. graphiti_core/search/search_helpers.py +5 -7
  28. graphiti_core/search/search_utils.py +227 -57
  29. graphiti_core/utils/bulk_utils.py +168 -69
  30. graphiti_core/utils/maintenance/community_operations.py +8 -20
  31. graphiti_core/utils/maintenance/dedup_helpers.py +262 -0
  32. graphiti_core/utils/maintenance/edge_operations.py +187 -50
  33. graphiti_core/utils/maintenance/graph_data_operations.py +9 -5
  34. graphiti_core/utils/maintenance/node_operations.py +244 -88
  35. graphiti_core/utils/maintenance/temporal_operations.py +0 -4
  36. {graphiti_core-0.20.4.dist-info → graphiti_core-0.21.0.dist-info}/METADATA +7 -1
  37. {graphiti_core-0.20.4.dist-info → graphiti_core-0.21.0.dist-info}/RECORD +39 -38
  38. {graphiti_core-0.20.4.dist-info → graphiti_core-0.21.0.dist-info}/WHEEL +0 -0
  39. {graphiti_core-0.20.4.dist-info → graphiti_core-0.21.0.dist-info}/licenses/LICENSE +0 -0
@@ -23,7 +23,13 @@ import numpy as np
23
23
  from numpy._typing import NDArray
24
24
  from typing_extensions import LiteralString
25
25
 
26
- from graphiti_core.driver.driver import GraphDriver, GraphProvider
26
+ from graphiti_core.driver.driver import (
27
+ ENTITY_EDGE_INDEX_NAME,
28
+ ENTITY_INDEX_NAME,
29
+ EPISODE_INDEX_NAME,
30
+ GraphDriver,
31
+ GraphProvider,
32
+ )
27
33
  from graphiti_core.edges import EntityEdge, get_entity_edge_from_record
28
34
  from graphiti_core.graph_queries import (
29
35
  get_nodes_query,
@@ -51,6 +57,8 @@ from graphiti_core.nodes import (
51
57
  )
52
58
  from graphiti_core.search.search_filters import (
53
59
  SearchFilters,
60
+ build_aoss_edge_filters,
61
+ build_aoss_node_filters,
54
62
  edge_search_filter_query_constructor,
55
63
  node_search_filter_query_constructor,
56
64
  )
@@ -84,6 +92,8 @@ def fulltext_query(query: str, group_ids: list[str] | None, driver: GraphDriver)
84
92
  if len(query.split(' ')) > MAX_QUERY_LENGTH:
85
93
  return ''
86
94
  return query
95
+ elif driver.provider == GraphProvider.FALKORDB:
96
+ return driver.build_fulltext_query(query, group_ids, MAX_QUERY_LENGTH)
87
97
  group_ids_filter_list = (
88
98
  [driver.fulltext_syntax + f'group_id:"{g}"' for g in group_ids]
89
99
  if group_ids is not None
@@ -200,7 +210,6 @@ async def edge_fulltext_search(
200
210
  if driver.provider == GraphProvider.NEPTUNE:
201
211
  res = driver.run_aoss_query('edge_name_and_fact', query) # pyright: ignore reportAttributeAccessIssue
202
212
  if res['hits']['total']['value'] > 0:
203
- # Calculate Cosine similarity then return the edge ids
204
213
  input_ids = []
205
214
  for r in res['hits']['hits']:
206
215
  input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
@@ -208,11 +217,11 @@ async def edge_fulltext_search(
208
217
  # Match the edge ids and return the values
209
218
  query = (
210
219
  """
211
- UNWIND $ids as id
212
- MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
213
- WHERE e.group_id IN $group_ids
214
- AND id(e)=id
215
- """
220
+ UNWIND $ids as id
221
+ MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
222
+ WHERE e.group_id IN $group_ids
223
+ AND id(e)=id
224
+ """
216
225
  + filter_query
217
226
  + """
218
227
  AND id(e)=id
@@ -244,6 +253,35 @@ async def edge_fulltext_search(
244
253
  )
245
254
  else:
246
255
  return []
256
+ elif driver.aoss_client:
257
+ route = group_ids[0] if group_ids else None
258
+ filters = build_aoss_edge_filters(group_ids or [], search_filter)
259
+ res = await driver.aoss_client.search(
260
+ index=ENTITY_EDGE_INDEX_NAME,
261
+ params={'routing': route},
262
+ body={
263
+ 'size': limit,
264
+ '_source': ['uuid'],
265
+ 'query': {
266
+ 'bool': {
267
+ 'filter': filters,
268
+ 'must': [{'match': {'fact': {'query': query, 'operator': 'or'}}}],
269
+ }
270
+ },
271
+ },
272
+ )
273
+
274
+ if res['hits']['total']['value'] > 0:
275
+ input_uuids = {}
276
+ for r in res['hits']['hits']:
277
+ input_uuids[r['_source']['uuid']] = r['_score']
278
+
279
+ # Get edges
280
+ entity_edges = await EntityEdge.get_by_uuids(driver, list(input_uuids.keys()))
281
+ entity_edges.sort(key=lambda e: input_uuids.get(e.uuid, 0), reverse=True)
282
+ return entity_edges
283
+ else:
284
+ return []
247
285
  else:
248
286
  query = (
249
287
  get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
@@ -318,8 +356,8 @@ async def edge_similarity_search(
318
356
  if driver.provider == GraphProvider.NEPTUNE:
319
357
  query = (
320
358
  """
321
- MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
322
- """
359
+ MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
360
+ """
323
361
  + filter_query
324
362
  + """
325
363
  RETURN DISTINCT id(e) as id, e.fact_embedding as embedding
@@ -377,6 +415,38 @@ async def edge_similarity_search(
377
415
  )
378
416
  else:
379
417
  return []
418
+ elif driver.aoss_client:
419
+ route = group_ids[0] if group_ids else None
420
+ filters = build_aoss_edge_filters(group_ids or [], search_filter)
421
+ res = await driver.aoss_client.search(
422
+ index=ENTITY_EDGE_INDEX_NAME,
423
+ params={'routing': route},
424
+ body={
425
+ 'size': limit,
426
+ '_source': ['uuid'],
427
+ 'query': {
428
+ 'knn': {
429
+ 'fact_embedding': {
430
+ 'vector': list(map(float, search_vector)),
431
+ 'k': limit,
432
+ 'filter': {'bool': {'filter': filters}},
433
+ }
434
+ }
435
+ },
436
+ },
437
+ )
438
+
439
+ if res['hits']['total']['value'] > 0:
440
+ input_uuids = {}
441
+ for r in res['hits']['hits']:
442
+ input_uuids[r['_source']['uuid']] = r['_score']
443
+
444
+ # Get edges
445
+ entity_edges = await EntityEdge.get_by_uuids(driver, list(input_uuids.keys()))
446
+ entity_edges.sort(key=lambda e: input_uuids.get(e.uuid, 0), reverse=True)
447
+ return entity_edges
448
+ return []
449
+
380
450
  else:
381
451
  query = (
382
452
  match_query
@@ -563,7 +633,6 @@ async def node_fulltext_search(
563
633
  if driver.provider == GraphProvider.NEPTUNE:
564
634
  res = driver.run_aoss_query('node_name_and_summary', query, limit=limit) # pyright: ignore reportAttributeAccessIssue
565
635
  if res['hits']['total']['value'] > 0:
566
- # Calculate Cosine similarity then return the edge ids
567
636
  input_ids = []
568
637
  for r in res['hits']['hits']:
569
638
  input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
@@ -571,11 +640,11 @@ async def node_fulltext_search(
571
640
  # Match the edge ides and return the values
572
641
  query = (
573
642
  """
574
- UNWIND $ids as i
575
- MATCH (n:Entity)
576
- WHERE n.uuid=i.id
577
- RETURN
578
- """
643
+ UNWIND $ids as i
644
+ MATCH (n:Entity)
645
+ WHERE n.uuid=i.id
646
+ RETURN
647
+ """
579
648
  + get_entity_node_return_query(driver.provider)
580
649
  + """
581
650
  ORDER BY i.score DESC
@@ -592,6 +661,43 @@ async def node_fulltext_search(
592
661
  )
593
662
  else:
594
663
  return []
664
+ elif driver.aoss_client:
665
+ route = group_ids[0] if group_ids else None
666
+ filters = build_aoss_node_filters(group_ids or [], search_filter)
667
+ res = await driver.aoss_client.search(
668
+ index=ENTITY_INDEX_NAME,
669
+ params={'routing': route},
670
+ body={
671
+ '_source': ['uuid'],
672
+ 'size': limit,
673
+ 'query': {
674
+ 'bool': {
675
+ 'filter': filters,
676
+ 'must': [
677
+ {
678
+ 'multi_match': {
679
+ 'query': query,
680
+ 'fields': ['name', 'summary'],
681
+ 'operator': 'or',
682
+ }
683
+ }
684
+ ],
685
+ }
686
+ },
687
+ },
688
+ )
689
+
690
+ if res['hits']['total']['value'] > 0:
691
+ input_uuids = {}
692
+ for r in res['hits']['hits']:
693
+ input_uuids[r['_source']['uuid']] = r['_score']
694
+
695
+ # Get nodes
696
+ entities = await EntityNode.get_by_uuids(driver, list(input_uuids.keys()))
697
+ entities.sort(key=lambda e: input_uuids.get(e.uuid, 0), reverse=True)
698
+ return entities
699
+ else:
700
+ return []
595
701
  else:
596
702
  query = (
597
703
  get_nodes_query(
@@ -648,8 +754,8 @@ async def node_similarity_search(
648
754
  if driver.provider == GraphProvider.NEPTUNE:
649
755
  query = (
650
756
  """
651
- MATCH (n:Entity)
652
- """
757
+ MATCH (n:Entity)
758
+ """
653
759
  + filter_query
654
760
  + """
655
761
  RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -678,11 +784,11 @@ async def node_similarity_search(
678
784
  # Match the edge ides and return the values
679
785
  query = (
680
786
  """
681
- UNWIND $ids as i
682
- MATCH (n:Entity)
683
- WHERE id(n)=i.id
684
- RETURN
685
- """
787
+ UNWIND $ids as i
788
+ MATCH (n:Entity)
789
+ WHERE id(n)=i.id
790
+ RETURN
791
+ """
686
792
  + get_entity_node_return_query(driver.provider)
687
793
  + """
688
794
  ORDER BY i.score DESC
@@ -700,11 +806,42 @@ async def node_similarity_search(
700
806
  )
701
807
  else:
702
808
  return []
809
+ elif driver.aoss_client:
810
+ route = group_ids[0] if group_ids else None
811
+ filters = build_aoss_node_filters(group_ids or [], search_filter)
812
+ res = await driver.aoss_client.search(
813
+ index=ENTITY_INDEX_NAME,
814
+ params={'routing': route},
815
+ body={
816
+ 'size': limit,
817
+ '_source': ['uuid'],
818
+ 'query': {
819
+ 'knn': {
820
+ 'name_embedding': {
821
+ 'vector': list(map(float, search_vector)),
822
+ 'k': limit,
823
+ 'filter': {'bool': {'filter': filters}},
824
+ }
825
+ }
826
+ },
827
+ },
828
+ )
829
+
830
+ if res['hits']['total']['value'] > 0:
831
+ input_uuids = {}
832
+ for r in res['hits']['hits']:
833
+ input_uuids[r['_source']['uuid']] = r['_score']
834
+
835
+ # Get edges
836
+ entity_nodes = await EntityNode.get_by_uuids(driver, list(input_uuids.keys()))
837
+ entity_nodes.sort(key=lambda e: input_uuids.get(e.uuid, 0), reverse=True)
838
+ return entity_nodes
839
+ return []
703
840
  else:
704
841
  query = (
705
842
  """
706
- MATCH (n:Entity)
707
- """
843
+ MATCH (n:Entity)
844
+ """
708
845
  + filter_query
709
846
  + """
710
847
  WITH n, """
@@ -843,7 +980,6 @@ async def episode_fulltext_search(
843
980
  if driver.provider == GraphProvider.NEPTUNE:
844
981
  res = driver.run_aoss_query('episode_content', query, limit=limit) # pyright: ignore reportAttributeAccessIssue
845
982
  if res['hits']['total']['value'] > 0:
846
- # Calculate Cosine similarity then return the edge ids
847
983
  input_ids = []
848
984
  for r in res['hits']['hits']:
849
985
  input_ids.append({'id': r['_source']['uuid'], 'score': r['_score']})
@@ -852,7 +988,7 @@ async def episode_fulltext_search(
852
988
  query = """
853
989
  UNWIND $ids as i
854
990
  MATCH (e:Episodic)
855
- WHERE e.uuid=i.id
991
+ WHERE e.uuid=i.uuid
856
992
  RETURN
857
993
  e.content AS content,
858
994
  e.created_at AS created_at,
@@ -876,6 +1012,40 @@ async def episode_fulltext_search(
876
1012
  )
877
1013
  else:
878
1014
  return []
1015
+ elif driver.aoss_client:
1016
+ route = group_ids[0] if group_ids else None
1017
+ res = await driver.aoss_client.search(
1018
+ index=EPISODE_INDEX_NAME,
1019
+ params={'routing': route},
1020
+ body={
1021
+ 'size': limit,
1022
+ '_source': ['uuid'],
1023
+ 'bool': {
1024
+ 'filter': {'terms': group_ids},
1025
+ 'must': [
1026
+ {
1027
+ 'multi_match': {
1028
+ 'query': query,
1029
+ 'field': ['name', 'content'],
1030
+ 'operator': 'or',
1031
+ }
1032
+ }
1033
+ ],
1034
+ },
1035
+ },
1036
+ )
1037
+
1038
+ if res['hits']['total']['value'] > 0:
1039
+ input_uuids = {}
1040
+ for r in res['hits']['hits']:
1041
+ input_uuids[r['_source']['uuid']] = r['_score']
1042
+
1043
+ # Get nodes
1044
+ episodes = await EpisodicNode.get_by_uuids(driver, list(input_uuids.keys()))
1045
+ episodes.sort(key=lambda e: input_uuids.get(e.uuid, 0), reverse=True)
1046
+ return episodes
1047
+ else:
1048
+ return []
879
1049
  else:
880
1050
  query = (
881
1051
  get_nodes_query('episode_content', '$query', limit=limit, provider=driver.provider)
@@ -1003,8 +1173,8 @@ async def community_similarity_search(
1003
1173
  if driver.provider == GraphProvider.NEPTUNE:
1004
1174
  query = (
1005
1175
  """
1006
- MATCH (n:Community)
1007
- """
1176
+ MATCH (n:Community)
1177
+ """
1008
1178
  + group_filter_query
1009
1179
  + """
1010
1180
  RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -1063,8 +1233,8 @@ async def community_similarity_search(
1063
1233
 
1064
1234
  query = (
1065
1235
  """
1066
- MATCH (c:Community)
1067
- """
1236
+ MATCH (c:Community)
1237
+ """
1068
1238
  + group_filter_query
1069
1239
  + """
1070
1240
  WITH c,
@@ -1206,9 +1376,9 @@ async def get_relevant_nodes(
1206
1376
  # FIXME: Kuzu currently does not support using variables such as `node.fulltext_query` as an input to FTS, which means `get_relevant_nodes()` won't work with Kuzu as the graph driver.
1207
1377
  query = (
1208
1378
  """
1209
- UNWIND $nodes AS node
1210
- MATCH (n:Entity {group_id: $group_id})
1211
- """
1379
+ UNWIND $nodes AS node
1380
+ MATCH (n:Entity {group_id: $group_id})
1381
+ """
1212
1382
  + filter_query
1213
1383
  + """
1214
1384
  WITH node, n, """
@@ -1253,9 +1423,9 @@ async def get_relevant_nodes(
1253
1423
  else:
1254
1424
  query = (
1255
1425
  """
1256
- UNWIND $nodes AS node
1257
- MATCH (n:Entity {group_id: $group_id})
1258
- """
1426
+ UNWIND $nodes AS node
1427
+ MATCH (n:Entity {group_id: $group_id})
1428
+ """
1259
1429
  + filter_query
1260
1430
  + """
1261
1431
  WITH node, n, """
@@ -1344,9 +1514,9 @@ async def get_relevant_edges(
1344
1514
  if driver.provider == GraphProvider.NEPTUNE:
1345
1515
  query = (
1346
1516
  """
1347
- UNWIND $edges AS edge
1348
- MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1349
- """
1517
+ UNWIND $edges AS edge
1518
+ MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1519
+ """
1350
1520
  + filter_query
1351
1521
  + """
1352
1522
  WITH e, edge
@@ -1416,9 +1586,9 @@ async def get_relevant_edges(
1416
1586
 
1417
1587
  query = (
1418
1588
  """
1419
- UNWIND $edges AS edge
1420
- MATCH (n:Entity {uuid: edge.source_node_uuid})-[:RELATES_TO]-(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]-(m:Entity {uuid: edge.target_node_uuid})
1421
- """
1589
+ UNWIND $edges AS edge
1590
+ MATCH (n:Entity {uuid: edge.source_node_uuid})-[:RELATES_TO]-(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]-(m:Entity {uuid: edge.target_node_uuid})
1591
+ """
1422
1592
  + filter_query
1423
1593
  + """
1424
1594
  WITH e, edge, n, m, """
@@ -1454,9 +1624,9 @@ async def get_relevant_edges(
1454
1624
  else:
1455
1625
  query = (
1456
1626
  """
1457
- UNWIND $edges AS edge
1458
- MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1459
- """
1627
+ UNWIND $edges AS edge
1628
+ MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1629
+ """
1460
1630
  + filter_query
1461
1631
  + """
1462
1632
  WITH e, edge, """
@@ -1529,10 +1699,10 @@ async def get_edge_invalidation_candidates(
1529
1699
  if driver.provider == GraphProvider.NEPTUNE:
1530
1700
  query = (
1531
1701
  """
1532
- UNWIND $edges AS edge
1533
- MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1534
- WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1535
- """
1702
+ UNWIND $edges AS edge
1703
+ MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1704
+ WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1705
+ """
1536
1706
  + filter_query
1537
1707
  + """
1538
1708
  WITH e, edge
@@ -1602,10 +1772,10 @@ async def get_edge_invalidation_candidates(
1602
1772
 
1603
1773
  query = (
1604
1774
  """
1605
- UNWIND $edges AS edge
1606
- MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity)
1607
- WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid])
1608
- """
1775
+ UNWIND $edges AS edge
1776
+ MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity)
1777
+ WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid])
1778
+ """
1609
1779
  + filter_query
1610
1780
  + """
1611
1781
  WITH edge, e, n, m, """
@@ -1641,10 +1811,10 @@ async def get_edge_invalidation_candidates(
1641
1811
  else:
1642
1812
  query = (
1643
1813
  """
1644
- UNWIND $edges AS edge
1645
- MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1646
- WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1647
- """
1814
+ UNWIND $edges AS edge
1815
+ MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1816
+ WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1817
+ """
1648
1818
  + filter_query
1649
1819
  + """
1650
1820
  WITH edge, e, """