graphiti-core 0.21.0rc12__py3-none-any.whl → 0.22.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 (41) hide show
  1. graphiti_core/driver/driver.py +4 -211
  2. graphiti_core/driver/falkordb_driver.py +31 -3
  3. graphiti_core/driver/graph_operations/graph_operations.py +195 -0
  4. graphiti_core/driver/neo4j_driver.py +0 -49
  5. graphiti_core/driver/neptune_driver.py +43 -26
  6. graphiti_core/driver/search_interface/__init__.py +0 -0
  7. graphiti_core/driver/search_interface/search_interface.py +89 -0
  8. graphiti_core/edges.py +11 -34
  9. graphiti_core/graphiti.py +459 -326
  10. graphiti_core/graphiti_types.py +2 -0
  11. graphiti_core/llm_client/anthropic_client.py +64 -45
  12. graphiti_core/llm_client/client.py +67 -19
  13. graphiti_core/llm_client/gemini_client.py +73 -54
  14. graphiti_core/llm_client/openai_base_client.py +65 -43
  15. graphiti_core/llm_client/openai_generic_client.py +65 -43
  16. graphiti_core/models/edges/edge_db_queries.py +1 -0
  17. graphiti_core/models/nodes/node_db_queries.py +1 -0
  18. graphiti_core/nodes.py +26 -99
  19. graphiti_core/prompts/dedupe_edges.py +4 -4
  20. graphiti_core/prompts/dedupe_nodes.py +10 -10
  21. graphiti_core/prompts/extract_edges.py +4 -4
  22. graphiti_core/prompts/extract_nodes.py +26 -28
  23. graphiti_core/prompts/prompt_helpers.py +18 -2
  24. graphiti_core/prompts/snippets.py +29 -0
  25. graphiti_core/prompts/summarize_nodes.py +22 -24
  26. graphiti_core/search/search_filters.py +0 -38
  27. graphiti_core/search/search_helpers.py +4 -4
  28. graphiti_core/search/search_utils.py +84 -220
  29. graphiti_core/tracer.py +193 -0
  30. graphiti_core/utils/bulk_utils.py +16 -28
  31. graphiti_core/utils/maintenance/community_operations.py +4 -1
  32. graphiti_core/utils/maintenance/edge_operations.py +30 -15
  33. graphiti_core/utils/maintenance/graph_data_operations.py +6 -25
  34. graphiti_core/utils/maintenance/node_operations.py +99 -51
  35. graphiti_core/utils/maintenance/temporal_operations.py +4 -1
  36. graphiti_core/utils/text_utils.py +53 -0
  37. {graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0.dist-info}/METADATA +7 -3
  38. {graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0.dist-info}/RECORD +41 -35
  39. /graphiti_core/{utils/maintenance/utils.py → driver/graph_operations/__init__.py} +0 -0
  40. {graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0.dist-info}/WHEEL +0 -0
  41. {graphiti_core-0.21.0rc12.dist-info → graphiti_core-0.22.0.dist-info}/licenses/LICENSE +0 -0
@@ -24,9 +24,6 @@ from numpy._typing import NDArray
24
24
  from typing_extensions import LiteralString
25
25
 
26
26
  from graphiti_core.driver.driver import (
27
- ENTITY_EDGE_INDEX_NAME,
28
- ENTITY_INDEX_NAME,
29
- EPISODE_INDEX_NAME,
30
27
  GraphDriver,
31
28
  GraphProvider,
32
29
  )
@@ -57,8 +54,6 @@ from graphiti_core.nodes import (
57
54
  )
58
55
  from graphiti_core.search.search_filters import (
59
56
  SearchFilters,
60
- build_aoss_edge_filters,
61
- build_aoss_node_filters,
62
57
  edge_search_filter_query_constructor,
63
58
  node_search_filter_query_constructor,
64
59
  )
@@ -179,6 +174,11 @@ async def edge_fulltext_search(
179
174
  group_ids: list[str] | None = None,
180
175
  limit=RELEVANT_SCHEMA_LIMIT,
181
176
  ) -> list[EntityEdge]:
177
+ if driver.search_interface:
178
+ return await driver.search_interface.edge_fulltext_search(
179
+ driver, query, search_filter, group_ids, limit
180
+ )
181
+
182
182
  # fulltext search over facts
183
183
  fuzzy_query = fulltext_query(query, group_ids, driver)
184
184
 
@@ -217,11 +217,11 @@ async def edge_fulltext_search(
217
217
  # Match the edge ids and return the values
218
218
  query = (
219
219
  """
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
- """
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
+ """
225
225
  + filter_query
226
226
  + """
227
227
  AND id(e)=id
@@ -253,35 +253,6 @@ async def edge_fulltext_search(
253
253
  )
254
254
  else:
255
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 []
285
256
  else:
286
257
  query = (
287
258
  get_relationships_query('edge_name_and_fact', limit=limit, provider=driver.provider)
@@ -321,6 +292,18 @@ async def edge_similarity_search(
321
292
  limit: int = RELEVANT_SCHEMA_LIMIT,
322
293
  min_score: float = DEFAULT_MIN_SCORE,
323
294
  ) -> list[EntityEdge]:
295
+ if driver.search_interface:
296
+ return await driver.search_interface.edge_similarity_search(
297
+ driver,
298
+ search_vector,
299
+ source_node_uuid,
300
+ target_node_uuid,
301
+ search_filter,
302
+ group_ids,
303
+ limit,
304
+ min_score,
305
+ )
306
+
324
307
  match_query = """
325
308
  MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
326
309
  """
@@ -356,8 +339,8 @@ async def edge_similarity_search(
356
339
  if driver.provider == GraphProvider.NEPTUNE:
357
340
  query = (
358
341
  """
359
- MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
360
- """
342
+ MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
343
+ """
361
344
  + filter_query
362
345
  + """
363
346
  RETURN DISTINCT id(e) as id, e.fact_embedding as embedding
@@ -415,38 +398,6 @@ async def edge_similarity_search(
415
398
  )
416
399
  else:
417
400
  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
-
450
401
  else:
451
402
  query = (
452
403
  match_query
@@ -609,6 +560,11 @@ async def node_fulltext_search(
609
560
  group_ids: list[str] | None = None,
610
561
  limit=RELEVANT_SCHEMA_LIMIT,
611
562
  ) -> list[EntityNode]:
563
+ if driver.search_interface:
564
+ return await driver.search_interface.node_fulltext_search(
565
+ driver, query, search_filter, group_ids, limit
566
+ )
567
+
612
568
  # BM25 search to get top nodes
613
569
  fuzzy_query = fulltext_query(query, group_ids, driver)
614
570
  if fuzzy_query == '':
@@ -640,11 +596,11 @@ async def node_fulltext_search(
640
596
  # Match the edge ides and return the values
641
597
  query = (
642
598
  """
643
- UNWIND $ids as i
644
- MATCH (n:Entity)
645
- WHERE n.uuid=i.id
646
- RETURN
647
- """
599
+ UNWIND $ids as i
600
+ MATCH (n:Entity)
601
+ WHERE n.uuid=i.id
602
+ RETURN
603
+ """
648
604
  + get_entity_node_return_query(driver.provider)
649
605
  + """
650
606
  ORDER BY i.score DESC
@@ -661,43 +617,6 @@ async def node_fulltext_search(
661
617
  )
662
618
  else:
663
619
  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 []
701
620
  else:
702
621
  query = (
703
622
  get_nodes_query(
@@ -735,6 +654,11 @@ async def node_similarity_search(
735
654
  limit=RELEVANT_SCHEMA_LIMIT,
736
655
  min_score: float = DEFAULT_MIN_SCORE,
737
656
  ) -> list[EntityNode]:
657
+ if driver.search_interface:
658
+ return await driver.search_interface.node_similarity_search(
659
+ driver, search_vector, search_filter, group_ids, limit, min_score
660
+ )
661
+
738
662
  filter_queries, filter_params = node_search_filter_query_constructor(
739
663
  search_filter, driver.provider
740
664
  )
@@ -754,8 +678,8 @@ async def node_similarity_search(
754
678
  if driver.provider == GraphProvider.NEPTUNE:
755
679
  query = (
756
680
  """
757
- MATCH (n:Entity)
758
- """
681
+ MATCH (n:Entity)
682
+ """
759
683
  + filter_query
760
684
  + """
761
685
  RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -784,11 +708,11 @@ async def node_similarity_search(
784
708
  # Match the edge ides and return the values
785
709
  query = (
786
710
  """
787
- UNWIND $ids as i
788
- MATCH (n:Entity)
789
- WHERE id(n)=i.id
790
- RETURN
791
- """
711
+ UNWIND $ids as i
712
+ MATCH (n:Entity)
713
+ WHERE id(n)=i.id
714
+ RETURN
715
+ """
792
716
  + get_entity_node_return_query(driver.provider)
793
717
  + """
794
718
  ORDER BY i.score DESC
@@ -806,42 +730,11 @@ async def node_similarity_search(
806
730
  )
807
731
  else:
808
732
  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 []
840
733
  else:
841
734
  query = (
842
735
  """
843
- MATCH (n:Entity)
844
- """
736
+ MATCH (n:Entity)
737
+ """
845
738
  + filter_query
846
739
  + """
847
740
  WITH n, """
@@ -966,6 +859,11 @@ async def episode_fulltext_search(
966
859
  group_ids: list[str] | None = None,
967
860
  limit=RELEVANT_SCHEMA_LIMIT,
968
861
  ) -> list[EpisodicNode]:
862
+ if driver.search_interface:
863
+ return await driver.search_interface.episode_fulltext_search(
864
+ driver, query, _search_filter, group_ids, limit
865
+ )
866
+
969
867
  # BM25 search to get top episodes
970
868
  fuzzy_query = fulltext_query(query, group_ids, driver)
971
869
  if fuzzy_query == '':
@@ -1012,40 +910,6 @@ async def episode_fulltext_search(
1012
910
  )
1013
911
  else:
1014
912
  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 []
1049
913
  else:
1050
914
  query = (
1051
915
  get_nodes_query('episode_content', '$query', limit=limit, provider=driver.provider)
@@ -1173,8 +1037,8 @@ async def community_similarity_search(
1173
1037
  if driver.provider == GraphProvider.NEPTUNE:
1174
1038
  query = (
1175
1039
  """
1176
- MATCH (n:Community)
1177
- """
1040
+ MATCH (n:Community)
1041
+ """
1178
1042
  + group_filter_query
1179
1043
  + """
1180
1044
  RETURN DISTINCT id(n) as id, n.name_embedding as embedding
@@ -1233,8 +1097,8 @@ async def community_similarity_search(
1233
1097
 
1234
1098
  query = (
1235
1099
  """
1236
- MATCH (c:Community)
1237
- """
1100
+ MATCH (c:Community)
1101
+ """
1238
1102
  + group_filter_query
1239
1103
  + """
1240
1104
  WITH c,
@@ -1376,9 +1240,9 @@ async def get_relevant_nodes(
1376
1240
  # 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.
1377
1241
  query = (
1378
1242
  """
1379
- UNWIND $nodes AS node
1380
- MATCH (n:Entity {group_id: $group_id})
1381
- """
1243
+ UNWIND $nodes AS node
1244
+ MATCH (n:Entity {group_id: $group_id})
1245
+ """
1382
1246
  + filter_query
1383
1247
  + """
1384
1248
  WITH node, n, """
@@ -1423,9 +1287,9 @@ async def get_relevant_nodes(
1423
1287
  else:
1424
1288
  query = (
1425
1289
  """
1426
- UNWIND $nodes AS node
1427
- MATCH (n:Entity {group_id: $group_id})
1428
- """
1290
+ UNWIND $nodes AS node
1291
+ MATCH (n:Entity {group_id: $group_id})
1292
+ """
1429
1293
  + filter_query
1430
1294
  + """
1431
1295
  WITH node, n, """
@@ -1514,9 +1378,9 @@ async def get_relevant_edges(
1514
1378
  if driver.provider == GraphProvider.NEPTUNE:
1515
1379
  query = (
1516
1380
  """
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
- """
1381
+ UNWIND $edges AS edge
1382
+ MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1383
+ """
1520
1384
  + filter_query
1521
1385
  + """
1522
1386
  WITH e, edge
@@ -1586,9 +1450,9 @@ async def get_relevant_edges(
1586
1450
 
1587
1451
  query = (
1588
1452
  """
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
- """
1453
+ UNWIND $edges AS edge
1454
+ 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})
1455
+ """
1592
1456
  + filter_query
1593
1457
  + """
1594
1458
  WITH e, edge, n, m, """
@@ -1624,9 +1488,9 @@ async def get_relevant_edges(
1624
1488
  else:
1625
1489
  query = (
1626
1490
  """
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
- """
1491
+ UNWIND $edges AS edge
1492
+ MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid})
1493
+ """
1630
1494
  + filter_query
1631
1495
  + """
1632
1496
  WITH e, edge, """
@@ -1699,10 +1563,10 @@ async def get_edge_invalidation_candidates(
1699
1563
  if driver.provider == GraphProvider.NEPTUNE:
1700
1564
  query = (
1701
1565
  """
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
- """
1566
+ UNWIND $edges AS edge
1567
+ MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1568
+ WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1569
+ """
1706
1570
  + filter_query
1707
1571
  + """
1708
1572
  WITH e, edge
@@ -1772,10 +1636,10 @@ async def get_edge_invalidation_candidates(
1772
1636
 
1773
1637
  query = (
1774
1638
  """
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
- """
1639
+ UNWIND $edges AS edge
1640
+ MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity)
1641
+ WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid])
1642
+ """
1779
1643
  + filter_query
1780
1644
  + """
1781
1645
  WITH edge, e, n, m, """
@@ -1811,10 +1675,10 @@ async def get_edge_invalidation_candidates(
1811
1675
  else:
1812
1676
  query = (
1813
1677
  """
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
- """
1678
+ UNWIND $edges AS edge
1679
+ MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity)
1680
+ WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]
1681
+ """
1818
1682
  + filter_query
1819
1683
  + """
1820
1684
  WITH edge, e, """