graphiti-core 0.5.2__tar.gz → 0.6.0__tar.gz

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 (61) hide show
  1. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/PKG-INFO +20 -16
  2. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/README.md +19 -15
  3. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/graphiti.py +38 -0
  4. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/config.py +1 -1
  5. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/nodes.py +7 -3
  6. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/search/search.py +19 -4
  7. graphiti_core-0.6.0/graphiti_core/search/search_filters.py +152 -0
  8. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/search/search_utils.py +49 -23
  9. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/pyproject.toml +1 -1
  10. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/LICENSE +0 -0
  11. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/__init__.py +0 -0
  12. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/cross_encoder/__init__.py +0 -0
  13. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/cross_encoder/bge_reranker_client.py +0 -0
  14. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/cross_encoder/client.py +0 -0
  15. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/cross_encoder/openai_reranker_client.py +0 -0
  16. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/edges.py +0 -0
  17. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/embedder/__init__.py +0 -0
  18. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/embedder/client.py +0 -0
  19. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/embedder/openai.py +0 -0
  20. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/embedder/voyage.py +0 -0
  21. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/errors.py +0 -0
  22. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/helpers.py +0 -0
  23. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/__init__.py +0 -0
  24. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/anthropic_client.py +0 -0
  25. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/client.py +0 -0
  26. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/errors.py +0 -0
  27. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/groq_client.py +0 -0
  28. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/openai_client.py +0 -0
  29. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/openai_generic_client.py +0 -0
  30. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/llm_client/utils.py +0 -0
  31. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/models/__init__.py +0 -0
  32. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/models/edges/__init__.py +0 -0
  33. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/models/edges/edge_db_queries.py +0 -0
  34. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/models/nodes/__init__.py +0 -0
  35. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/models/nodes/node_db_queries.py +0 -0
  36. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/__init__.py +0 -0
  37. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/dedupe_edges.py +0 -0
  38. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/dedupe_nodes.py +0 -0
  39. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/eval.py +0 -0
  40. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/extract_edge_dates.py +0 -0
  41. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/extract_edges.py +0 -0
  42. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/extract_nodes.py +0 -0
  43. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/invalidate_edges.py +0 -0
  44. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/lib.py +0 -0
  45. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/models.py +0 -0
  46. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/prompt_helpers.py +0 -0
  47. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/prompts/summarize_nodes.py +0 -0
  48. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/py.typed +0 -0
  49. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/search/__init__.py +0 -0
  50. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/search/search_config.py +0 -0
  51. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/search/search_config_recipes.py +0 -0
  52. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/__init__.py +0 -0
  53. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/bulk_utils.py +0 -0
  54. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/datetime_utils.py +0 -0
  55. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/__init__.py +0 -0
  56. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/community_operations.py +0 -0
  57. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/edge_operations.py +0 -0
  58. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/graph_data_operations.py +0 -0
  59. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/node_operations.py +0 -0
  60. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/temporal_operations.py +0 -0
  61. {graphiti_core-0.5.2 → graphiti_core-0.6.0}/graphiti_core/utils/maintenance/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: graphiti-core
3
- Version: 0.5.2
3
+ Version: 0.6.0
4
4
  Summary: A temporal graph building library
5
5
  License: Apache-2.0
6
6
  Author: Paul Paliychuk
@@ -22,7 +22,7 @@ Description-Content-Type: text/markdown
22
22
 
23
23
  <div align="center">
24
24
 
25
- <img width="350" alt="Graphiti-ts-small" src="https://github.com/user-attachments/assets/bbd02947-e435-4a05-b25a-bbbac36d52c8">
25
+ # Graphiti
26
26
 
27
27
  ## Temporal Knowledge Graphs for Agentic Applications
28
28
 
@@ -45,7 +45,7 @@ a fusion of time, full-text, semantic, and graph algorithm approaches.
45
45
  <br />
46
46
 
47
47
  <p align="center">
48
- <img src="/images/graphiti-graph-intro.gif" alt="Graphiti temporal walkthrough" width="700px">
48
+ <img src="images/graphiti-graph-intro.gif" alt="Graphiti temporal walkthrough" width="700px">
49
49
  </p>
50
50
 
51
51
  <br />
@@ -65,6 +65,20 @@ With Graphiti, you can build LLM applications such as:
65
65
  Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term
66
66
  recall and state-based reasoning for both assistants and agents.
67
67
 
68
+ ## Graphiti and Zep Memory
69
+
70
+ Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
71
+
72
+ Using Graphiti, we've demonstrated Zep is the [State of the Art in Agent Memory](https://blog.getzep.com/state-of-the-art-agent-memory/).
73
+
74
+ Read our paper: [Zep: A Temporal Knowledge Graph Architecture for Agent Memory](https://arxiv.org/abs/2501.13956).
75
+
76
+ We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
77
+
78
+ <p align="center">
79
+ <a href="https://arxiv.org/abs/2501.13956"><img src="images/arxiv-screenshot.png" alt="Zep: A Temporal Knowledge Graph Architecture for Agent Memory" width="700px"></a>
80
+ </p>
81
+
68
82
  ## Why Graphiti?
69
83
 
70
84
  We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a
@@ -89,12 +103,6 @@ scale:
89
103
  <img src="/images/graphiti-intro-slides-stock-2.gif" alt="Graphiti structured + unstructured demo" width="700px">
90
104
  </p>
91
105
 
92
- ## Graphiti and Zep Memory
93
-
94
- Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
95
-
96
- We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
97
-
98
106
  ## Installation
99
107
 
100
108
  Requirements:
@@ -125,7 +133,7 @@ poetry add graphiti-core
125
133
 
126
134
  > [!IMPORTANT]
127
135
  > Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment.
128
- > Support for Anthropic and Groq LLM inferences is available, too.
136
+ > Support for Anthropic and Groq LLM inferences is available, too. Other LLM providers may be supported via OpenAI compatible APIs.
129
137
 
130
138
  ```python
131
139
  from graphiti_core import Graphiti
@@ -216,13 +224,9 @@ as such this feature is off by default.
216
224
 
217
225
  Graphiti is under active development. We aim to maintain API stability while working on:
218
226
 
219
- - [x] Implementing node and edge CRUD operations
220
- - [ ] Improving performance and scalability
221
- - [ ] Achieving good performance with different LLM and embedding models
222
- - [x] Creating a dedicated embedder interface
223
227
  - [ ] Supporting custom graph schemas:
224
- - Allow developers to provide their own defined node and edge classes when ingesting episodes
225
- - Enable more flexible knowledge representation tailored to specific use cases
228
+ - Allow developers to provide their own defined node and edge classes when ingesting episodes
229
+ - Enable more flexible knowledge representation tailored to specific use cases
226
230
  - [x] Enhancing retrieval capabilities with more robust and configurable options
227
231
  - [ ] Expanding test coverage to ensure reliability and catch edge cases
228
232
 
@@ -1,6 +1,6 @@
1
1
  <div align="center">
2
2
 
3
- <img width="350" alt="Graphiti-ts-small" src="https://github.com/user-attachments/assets/bbd02947-e435-4a05-b25a-bbbac36d52c8">
3
+ # Graphiti
4
4
 
5
5
  ## Temporal Knowledge Graphs for Agentic Applications
6
6
 
@@ -23,7 +23,7 @@ a fusion of time, full-text, semantic, and graph algorithm approaches.
23
23
  <br />
24
24
 
25
25
  <p align="center">
26
- <img src="/images/graphiti-graph-intro.gif" alt="Graphiti temporal walkthrough" width="700px">
26
+ <img src="images/graphiti-graph-intro.gif" alt="Graphiti temporal walkthrough" width="700px">
27
27
  </p>
28
28
 
29
29
  <br />
@@ -43,6 +43,20 @@ With Graphiti, you can build LLM applications such as:
43
43
  Graphiti supports a wide range of applications in sales, customer service, health, finance, and more, enabling long-term
44
44
  recall and state-based reasoning for both assistants and agents.
45
45
 
46
+ ## Graphiti and Zep Memory
47
+
48
+ Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
49
+
50
+ Using Graphiti, we've demonstrated Zep is the [State of the Art in Agent Memory](https://blog.getzep.com/state-of-the-art-agent-memory/).
51
+
52
+ Read our paper: [Zep: A Temporal Knowledge Graph Architecture for Agent Memory](https://arxiv.org/abs/2501.13956).
53
+
54
+ We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
55
+
56
+ <p align="center">
57
+ <a href="https://arxiv.org/abs/2501.13956"><img src="images/arxiv-screenshot.png" alt="Zep: A Temporal Knowledge Graph Architecture for Agent Memory" width="700px"></a>
58
+ </p>
59
+
46
60
  ## Why Graphiti?
47
61
 
48
62
  We were intrigued by Microsoft’s GraphRAG, which expanded on RAG text chunking by using a graph to better model a
@@ -67,12 +81,6 @@ scale:
67
81
  <img src="/images/graphiti-intro-slides-stock-2.gif" alt="Graphiti structured + unstructured demo" width="700px">
68
82
  </p>
69
83
 
70
- ## Graphiti and Zep Memory
71
-
72
- Graphiti powers the core of [Zep's memory layer](https://www.getzep.com) for LLM-powered Assistants and Agents.
73
-
74
- We're excited to open-source Graphiti, believing its potential reaches far beyond memory applications.
75
-
76
84
  ## Installation
77
85
 
78
86
  Requirements:
@@ -103,7 +111,7 @@ poetry add graphiti-core
103
111
 
104
112
  > [!IMPORTANT]
105
113
  > Graphiti uses OpenAI for LLM inference and embedding. Ensure that an `OPENAI_API_KEY` is set in your environment.
106
- > Support for Anthropic and Groq LLM inferences is available, too.
114
+ > Support for Anthropic and Groq LLM inferences is available, too. Other LLM providers may be supported via OpenAI compatible APIs.
107
115
 
108
116
  ```python
109
117
  from graphiti_core import Graphiti
@@ -194,13 +202,9 @@ as such this feature is off by default.
194
202
 
195
203
  Graphiti is under active development. We aim to maintain API stability while working on:
196
204
 
197
- - [x] Implementing node and edge CRUD operations
198
- - [ ] Improving performance and scalability
199
- - [ ] Achieving good performance with different LLM and embedding models
200
- - [x] Creating a dedicated embedder interface
201
205
  - [ ] Supporting custom graph schemas:
202
- - Allow developers to provide their own defined node and edge classes when ingesting episodes
203
- - Enable more flexible knowledge representation tailored to specific use cases
206
+ - Allow developers to provide their own defined node and edge classes when ingesting episodes
207
+ - Enable more flexible knowledge representation tailored to specific use cases
204
208
  - [x] Enhancing retrieval capabilities with more robust and configurable options
205
209
  - [ ] Expanding test coverage to ensure reliability and catch edge cases
206
210
 
@@ -21,6 +21,7 @@ from time import time
21
21
  from dotenv import load_dotenv
22
22
  from neo4j import AsyncGraphDatabase
23
23
  from pydantic import BaseModel
24
+ from typing_extensions import LiteralString
24
25
 
25
26
  from graphiti_core.cross_encoder.client import CrossEncoderClient
26
27
  from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
@@ -35,6 +36,7 @@ from graphiti_core.search.search_config_recipes import (
35
36
  EDGE_HYBRID_SEARCH_NODE_DISTANCE,
36
37
  EDGE_HYBRID_SEARCH_RRF,
37
38
  )
39
+ from graphiti_core.search.search_filters import SearchFilters
38
40
  from graphiti_core.search.search_utils import (
39
41
  RELEVANT_SCHEMA_LIMIT,
40
42
  get_communities_by_nodes,
@@ -481,6 +483,7 @@ class Graphiti:
481
483
  except Exception as e:
482
484
  raise e
483
485
 
486
+ #### WIP: USE AT YOUR OWN RISK ####
484
487
  async def add_episode_bulk(self, bulk_episodes: list[RawEpisode], group_id: str = ''):
485
488
  """
486
489
  Process multiple episodes in bulk and update the graph.
@@ -624,6 +627,7 @@ class Graphiti:
624
627
  center_node_uuid: str | None = None,
625
628
  group_ids: list[str] | None = None,
626
629
  num_results=DEFAULT_SEARCH_LIMIT,
630
+ search_filter: SearchFilters | None = None,
627
631
  ) -> list[EntityEdge]:
628
632
  """
629
633
  Perform a hybrid search on the knowledge graph.
@@ -669,6 +673,7 @@ class Graphiti:
669
673
  query,
670
674
  group_ids,
671
675
  search_config,
676
+ search_filter if search_filter is not None else SearchFilters(),
672
677
  center_node_uuid,
673
678
  )
674
679
  ).edges
@@ -682,6 +687,7 @@ class Graphiti:
682
687
  group_ids: list[str] | None = None,
683
688
  center_node_uuid: str | None = None,
684
689
  bfs_origin_node_uuids: list[str] | None = None,
690
+ search_filter: SearchFilters | None = None,
685
691
  ) -> SearchResults:
686
692
  return await search(
687
693
  self.driver,
@@ -690,6 +696,7 @@ class Graphiti:
690
696
  query,
691
697
  group_ids,
692
698
  config,
699
+ search_filter if search_filter is not None else SearchFilters(),
693
700
  center_node_uuid,
694
701
  bfs_origin_node_uuids,
695
702
  )
@@ -741,3 +748,34 @@ class Graphiti:
741
748
  await add_nodes_and_edges_bulk(
742
749
  self.driver, [], [], resolved_nodes, [resolved_edge] + invalidated_edges
743
750
  )
751
+
752
+ async def remove_episode(self, episode_uuid: str):
753
+ # Find the episode to be deleted
754
+ episode = await EpisodicNode.get_by_uuid(self.driver, episode_uuid)
755
+
756
+ # Find edges mentioned by the episode
757
+ edges = await EntityEdge.get_by_uuids(self.driver, episode.entity_edges)
758
+
759
+ # We should only delete edges created by the episode
760
+ edges_to_delete: list[EntityEdge] = []
761
+ for edge in edges:
762
+ if edge.episodes[0] == episode.uuid:
763
+ edges_to_delete.append(edge)
764
+
765
+ # Find nodes mentioned by the episode
766
+ nodes = await get_mentioned_nodes(self.driver, [episode])
767
+ # We should delete all nodes that are only mentioned in the deleted episode
768
+ nodes_to_delete: list[EntityNode] = []
769
+ for node in nodes:
770
+ query: LiteralString = 'MATCH (e:Episodic)-[:MENTIONS]->(n:Entity {uuid: $uuid}) RETURN count(*) AS episode_count'
771
+ records, _, _ = await self.driver.execute_query(
772
+ query, uuid=node.uuid, database_=DEFAULT_DATABASE, routing_='r'
773
+ )
774
+
775
+ for record in records:
776
+ if record['episode_count'] == 1:
777
+ nodes_to_delete.append(node)
778
+
779
+ await semaphore_gather(*[node.delete(self.driver) for node in nodes_to_delete])
780
+ await semaphore_gather(*[edge.delete(self.driver) for edge in edges_to_delete])
781
+ await episode.delete(self.driver)
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
14
14
  limitations under the License.
15
15
  """
16
16
 
17
- DEFAULT_MAX_TOKENS = 1024
17
+ DEFAULT_MAX_TOKENS = 2048
18
18
  DEFAULT_TEMPERATURE = 0
19
19
 
20
20
 
@@ -170,7 +170,8 @@ class EpisodicNode(Node):
170
170
  e.name AS name,
171
171
  e.group_id AS group_id,
172
172
  e.source_description AS source_description,
173
- e.source AS source
173
+ e.source AS source,
174
+ e.entity_edges AS entity_edges
174
175
  """,
175
176
  uuid=uuid,
176
177
  database_=DEFAULT_DATABASE,
@@ -197,7 +198,8 @@ class EpisodicNode(Node):
197
198
  e.name AS name,
198
199
  e.group_id AS group_id,
199
200
  e.source_description AS source_description,
200
- e.source AS source
201
+ e.source AS source,
202
+ e.entity_edges AS entity_edges
201
203
  """,
202
204
  uuids=uuids,
203
205
  database_=DEFAULT_DATABASE,
@@ -233,7 +235,8 @@ class EpisodicNode(Node):
233
235
  e.name AS name,
234
236
  e.group_id AS group_id,
235
237
  e.source_description AS source_description,
236
- e.source AS source
238
+ e.source AS source,
239
+ e.entity_edges AS entity_edges
237
240
  ORDER BY e.uuid DESC
238
241
  """
239
242
  + limit_query,
@@ -490,6 +493,7 @@ def get_episodic_node_from_record(record: Any) -> EpisodicNode:
490
493
  source=EpisodeType.from_str(record['source']),
491
494
  name=record['name'],
492
495
  source_description=record['source_description'],
496
+ entity_edges=record['entity_edges'],
493
497
  )
494
498
 
495
499
 
@@ -39,6 +39,7 @@ from graphiti_core.search.search_config import (
39
39
  SearchConfig,
40
40
  SearchResults,
41
41
  )
42
+ from graphiti_core.search.search_filters import SearchFilters
42
43
  from graphiti_core.search.search_utils import (
43
44
  community_fulltext_search,
44
45
  community_similarity_search,
@@ -64,6 +65,7 @@ async def search(
64
65
  query: str,
65
66
  group_ids: list[str] | None,
66
67
  config: SearchConfig,
68
+ search_filter: SearchFilters,
67
69
  center_node_uuid: str | None = None,
68
70
  bfs_origin_node_uuids: list[str] | None = None,
69
71
  ) -> SearchResults:
@@ -86,6 +88,7 @@ async def search(
86
88
  query_vector,
87
89
  group_ids,
88
90
  config.edge_config,
91
+ search_filter,
89
92
  center_node_uuid,
90
93
  bfs_origin_node_uuids,
91
94
  config.limit,
@@ -133,6 +136,7 @@ async def edge_search(
133
136
  query_vector: list[float],
134
137
  group_ids: list[str] | None,
135
138
  config: EdgeSearchConfig | None,
139
+ search_filter: SearchFilters,
136
140
  center_node_uuid: str | None = None,
137
141
  bfs_origin_node_uuids: list[str] | None = None,
138
142
  limit=DEFAULT_SEARCH_LIMIT,
@@ -143,11 +147,20 @@ async def edge_search(
143
147
  search_results: list[list[EntityEdge]] = list(
144
148
  await semaphore_gather(
145
149
  *[
146
- edge_fulltext_search(driver, query, group_ids, 2 * limit),
150
+ edge_fulltext_search(driver, query, search_filter, group_ids, 2 * limit),
147
151
  edge_similarity_search(
148
- driver, query_vector, None, None, group_ids, 2 * limit, config.sim_min_score
152
+ driver,
153
+ query_vector,
154
+ None,
155
+ None,
156
+ search_filter,
157
+ group_ids,
158
+ 2 * limit,
159
+ config.sim_min_score,
160
+ ),
161
+ edge_bfs_search(
162
+ driver, bfs_origin_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
149
163
  ),
150
- edge_bfs_search(driver, bfs_origin_node_uuids, config.bfs_max_depth, 2 * limit),
151
164
  ]
152
165
  )
153
166
  )
@@ -155,7 +168,9 @@ async def edge_search(
155
168
  if EdgeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
156
169
  source_node_uuids = [edge.source_node_uuid for result in search_results for edge in result]
157
170
  search_results.append(
158
- await edge_bfs_search(driver, source_node_uuids, config.bfs_max_depth, 2 * limit)
171
+ await edge_bfs_search(
172
+ driver, source_node_uuids, config.bfs_max_depth, search_filter, 2 * limit
173
+ )
159
174
  )
160
175
 
161
176
  edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
@@ -0,0 +1,152 @@
1
+ """
2
+ Copyright 2024, Zep Software, Inc.
3
+
4
+ Licensed under the Apache License, Version 2.0 (the "License");
5
+ you may not use this file except in compliance with the License.
6
+ You may obtain a copy of the License at
7
+
8
+ http://www.apache.org/licenses/LICENSE-2.0
9
+
10
+ Unless required by applicable law or agreed to in writing, software
11
+ distributed under the License is distributed on an "AS IS" BASIS,
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ See the License for the specific language governing permissions and
14
+ limitations under the License.
15
+ """
16
+
17
+ from datetime import datetime
18
+ from enum import Enum
19
+ from typing import Any
20
+
21
+ from pydantic import BaseModel, Field
22
+ from typing_extensions import LiteralString
23
+
24
+
25
+ class ComparisonOperator(Enum):
26
+ equals = '='
27
+ not_equals = '<>'
28
+ greater_than = '>'
29
+ less_than = '<'
30
+ greater_than_equal = '>='
31
+ less_than_equal = '<='
32
+
33
+
34
+ class DateFilter(BaseModel):
35
+ date: datetime = Field(description='A datetime to filter on')
36
+ comparison_operator: ComparisonOperator = Field(
37
+ description='Comparison operator for date filter'
38
+ )
39
+
40
+
41
+ class SearchFilters(BaseModel):
42
+ valid_at: list[list[DateFilter]] | None = Field(default=None)
43
+ invalid_at: list[list[DateFilter]] | None = Field(default=None)
44
+ created_at: list[list[DateFilter]] | None = Field(default=None)
45
+ expired_at: list[list[DateFilter]] | None = Field(default=None)
46
+
47
+
48
+ def search_filter_query_constructor(filters: SearchFilters) -> tuple[LiteralString, dict[str, Any]]:
49
+ filter_query: LiteralString = ''
50
+ filter_params: dict[str, Any] = {}
51
+
52
+ if filters.valid_at is not None:
53
+ valid_at_filter = 'AND ('
54
+ for i, or_list in enumerate(filters.valid_at):
55
+ for j, date_filter in enumerate(or_list):
56
+ filter_params['valid_at_' + str(j)] = date_filter.date
57
+
58
+ and_filters = [
59
+ '(r.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
60
+ for j, date_filter in enumerate(or_list)
61
+ ]
62
+ and_filter_query = ''
63
+ for j, and_filter in enumerate(and_filters):
64
+ and_filter_query += and_filter
65
+ if j != len(and_filter_query) - 1:
66
+ and_filter_query += ' AND '
67
+
68
+ valid_at_filter += and_filter_query
69
+
70
+ if i == len(or_list) - 1:
71
+ valid_at_filter += ')'
72
+ else:
73
+ valid_at_filter += ' OR '
74
+
75
+ filter_query += valid_at_filter
76
+
77
+ if filters.invalid_at is not None:
78
+ invalid_at_filter = 'AND ('
79
+ for i, or_list in enumerate(filters.invalid_at):
80
+ for j, date_filter in enumerate(or_list):
81
+ filter_params['invalid_at_' + str(j)] = date_filter.date
82
+
83
+ and_filters = [
84
+ '(r.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
85
+ for j, date_filter in enumerate(or_list)
86
+ ]
87
+ and_filter_query = ''
88
+ for j, and_filter in enumerate(and_filters):
89
+ and_filter_query += and_filter
90
+ if j != len(and_filter_query) - 1:
91
+ and_filter_query += ' AND '
92
+
93
+ invalid_at_filter += and_filter_query
94
+
95
+ if i == len(or_list) - 1:
96
+ invalid_at_filter += ')'
97
+ else:
98
+ invalid_at_filter += ' OR '
99
+
100
+ filter_query += invalid_at_filter
101
+
102
+ if filters.created_at is not None:
103
+ created_at_filter = 'AND ('
104
+ for i, or_list in enumerate(filters.created_at):
105
+ for j, date_filter in enumerate(or_list):
106
+ filter_params['created_at_' + str(j)] = date_filter.date
107
+
108
+ and_filters = [
109
+ '(r.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
110
+ for j, date_filter in enumerate(or_list)
111
+ ]
112
+ and_filter_query = ''
113
+ for j, and_filter in enumerate(and_filters):
114
+ and_filter_query += and_filter
115
+ if j != len(and_filter_query) - 1:
116
+ and_filter_query += ' AND '
117
+
118
+ created_at_filter += and_filter_query
119
+
120
+ if i == len(or_list) - 1:
121
+ created_at_filter += ')'
122
+ else:
123
+ created_at_filter += ' OR '
124
+
125
+ filter_query += created_at_filter
126
+
127
+ if filters.expired_at is not None:
128
+ expired_at_filter = 'AND ('
129
+ for i, or_list in enumerate(filters.expired_at):
130
+ for j, date_filter in enumerate(or_list):
131
+ filter_params['expired_at_' + str(j)] = date_filter.date
132
+
133
+ and_filters = [
134
+ '(r.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
135
+ for j, date_filter in enumerate(or_list)
136
+ ]
137
+ and_filter_query = ''
138
+ for j, and_filter in enumerate(and_filters):
139
+ and_filter_query += and_filter
140
+ if j != len(and_filter_query) - 1:
141
+ and_filter_query += ' AND '
142
+
143
+ expired_at_filter += and_filter_query
144
+
145
+ if i == len(or_list) - 1:
146
+ expired_at_filter += ')'
147
+ else:
148
+ expired_at_filter += ' OR '
149
+
150
+ filter_query += expired_at_filter
151
+
152
+ return filter_query, filter_params
@@ -38,6 +38,7 @@ from graphiti_core.nodes import (
38
38
  get_community_node_from_record,
39
39
  get_entity_node_from_record,
40
40
  )
41
+ from graphiti_core.search.search_filters import SearchFilters, search_filter_query_constructor
41
42
 
42
43
  logger = logging.getLogger(__name__)
43
44
 
@@ -136,6 +137,7 @@ async def get_communities_by_nodes(
136
137
  async def edge_fulltext_search(
137
138
  driver: AsyncDriver,
138
139
  query: str,
140
+ search_filter: SearchFilters,
139
141
  group_ids: list[str] | None = None,
140
142
  limit=RELEVANT_SCHEMA_LIMIT,
141
143
  ) -> list[EntityEdge]:
@@ -144,28 +146,36 @@ async def edge_fulltext_search(
144
146
  if fuzzy_query == '':
145
147
  return []
146
148
 
147
- cypher_query = Query("""
149
+ filter_query, filter_params = search_filter_query_constructor(search_filter)
150
+
151
+ cypher_query = Query(
152
+ """
148
153
  CALL db.index.fulltext.queryRelationships("edge_name_and_fact", $query, {limit: $limit})
149
- YIELD relationship AS r, score
150
- WITH r, score, startNode(r) AS n, endNode(r) AS m
151
- RETURN
152
- r.uuid AS uuid,
153
- r.group_id AS group_id,
154
- n.uuid AS source_node_uuid,
155
- m.uuid AS target_node_uuid,
156
- r.created_at AS created_at,
157
- r.name AS name,
158
- r.fact AS fact,
159
- r.fact_embedding AS fact_embedding,
160
- r.episodes AS episodes,
161
- r.expired_at AS expired_at,
162
- r.valid_at AS valid_at,
163
- r.invalid_at AS invalid_at
164
- ORDER BY score DESC LIMIT $limit
165
- """)
154
+ YIELD relationship AS rel, score
155
+ MATCH (:ENTITY)-[r:RELATES_TO]->(:ENTITY)
156
+ WHERE r.group_id IN $group_ids"""
157
+ + filter_query
158
+ + """\nWITH r, score, startNode(r) AS n, endNode(r) AS m
159
+ RETURN
160
+ r.uuid AS uuid,
161
+ r.group_id AS group_id,
162
+ n.uuid AS source_node_uuid,
163
+ m.uuid AS target_node_uuid,
164
+ r.created_at AS created_at,
165
+ r.name AS name,
166
+ r.fact AS fact,
167
+ r.fact_embedding AS fact_embedding,
168
+ r.episodes AS episodes,
169
+ r.expired_at AS expired_at,
170
+ r.valid_at AS valid_at,
171
+ r.invalid_at AS invalid_at
172
+ ORDER BY score DESC LIMIT $limit
173
+ """
174
+ )
166
175
 
167
176
  records, _, _ = await driver.execute_query(
168
177
  cypher_query,
178
+ filter_params,
169
179
  query=fuzzy_query,
170
180
  group_ids=group_ids,
171
181
  limit=limit,
@@ -183,6 +193,7 @@ async def edge_similarity_search(
183
193
  search_vector: list[float],
184
194
  source_node_uuid: str | None,
185
195
  target_node_uuid: str | None,
196
+ search_filter: SearchFilters,
186
197
  group_ids: list[str] | None = None,
187
198
  limit: int = RELEVANT_SCHEMA_LIMIT,
188
199
  min_score: float = DEFAULT_MIN_SCORE,
@@ -194,6 +205,9 @@ async def edge_similarity_search(
194
205
 
195
206
  query_params: dict[str, Any] = {}
196
207
 
208
+ filter_query, filter_params = search_filter_query_constructor(search_filter)
209
+ query_params.update(filter_params)
210
+
197
211
  group_filter_query: LiteralString = ''
198
212
  if group_ids is not None:
199
213
  group_filter_query += 'WHERE r.group_id IN $group_ids'
@@ -209,9 +223,10 @@ async def edge_similarity_search(
209
223
 
210
224
  query: LiteralString = (
211
225
  """
212
- MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
213
- """
226
+ MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
227
+ """
214
228
  + group_filter_query
229
+ + filter_query
215
230
  + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
216
231
  WHERE score > $min_score
217
232
  RETURN
@@ -254,17 +269,25 @@ async def edge_bfs_search(
254
269
  driver: AsyncDriver,
255
270
  bfs_origin_node_uuids: list[str] | None,
256
271
  bfs_max_depth: int,
272
+ search_filter: SearchFilters,
257
273
  limit: int,
258
274
  ) -> list[EntityEdge]:
259
275
  # vector similarity search over embedded facts
260
276
  if bfs_origin_node_uuids is None:
261
277
  return []
262
278
 
263
- query = Query("""
279
+ filter_query, filter_params = search_filter_query_constructor(search_filter)
280
+
281
+ query = Query(
282
+ """
264
283
  UNWIND $bfs_origin_node_uuids AS origin_uuid
265
284
  MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
266
285
  UNWIND relationships(path) AS rel
267
- MATCH ()-[r:RELATES_TO {uuid: rel.uuid}]-()
286
+ MATCH ()-[r:RELATES_TO]-()
287
+ WHERE r.uuid = rel.uuid
288
+ """
289
+ + filter_query
290
+ + """
268
291
  RETURN DISTINCT
269
292
  r.uuid AS uuid,
270
293
  r.group_id AS group_id,
@@ -279,10 +302,12 @@ async def edge_bfs_search(
279
302
  r.valid_at AS valid_at,
280
303
  r.invalid_at AS invalid_at
281
304
  LIMIT $limit
282
- """)
305
+ """
306
+ )
283
307
 
284
308
  records, _, _ = await driver.execute_query(
285
309
  query,
310
+ filter_params,
286
311
  bfs_origin_node_uuids=bfs_origin_node_uuids,
287
312
  depth=bfs_max_depth,
288
313
  limit=limit,
@@ -626,6 +651,7 @@ async def get_relevant_edges(
626
651
  edge.fact_embedding,
627
652
  source_node_uuid,
628
653
  target_node_uuid,
654
+ SearchFilters(),
629
655
  [edge.group_id],
630
656
  limit,
631
657
  )
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "graphiti-core"
3
- version = "0.5.2"
3
+ version = "0.6.0"
4
4
  description = "A temporal graph building library"
5
5
  authors = [
6
6
  "Paul Paliychuk <paul@getzep.com>",
File without changes