graphiti-core 0.10.5__tar.gz → 0.11.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 (68) hide show
  1. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/PKG-INFO +14 -8
  2. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/README.md +13 -7
  3. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/edges.py +32 -57
  4. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/embedder/client.py +3 -0
  5. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/embedder/gemini.py +10 -0
  6. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/embedder/openai.py +6 -0
  7. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/embedder/voyage.py +7 -0
  8. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/graphiti.py +42 -138
  9. graphiti_core-0.11.0/graphiti_core/graphiti_types.py +31 -0
  10. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/helpers.py +6 -1
  11. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/models/edges/edge_db_queries.py +1 -1
  12. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/nodes.py +8 -2
  13. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/dedupe_edges.py +5 -7
  14. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/dedupe_nodes.py +8 -21
  15. graphiti_core-0.11.0/graphiti_core/prompts/extract_edges.py +148 -0
  16. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/extract_nodes.py +89 -18
  17. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/invalidate_edges.py +11 -11
  18. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search.py +13 -5
  19. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search_utils.py +204 -82
  20. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/bulk_utils.py +10 -7
  21. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/edge_operations.py +88 -40
  22. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/graph_data_operations.py +9 -3
  23. graphiti_core-0.11.0/graphiti_core/utils/maintenance/node_operations.py +461 -0
  24. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/temporal_operations.py +4 -11
  25. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/pyproject.toml +1 -1
  26. graphiti_core-0.10.5/graphiti_core/prompts/extract_edges.py +0 -113
  27. graphiti_core-0.10.5/graphiti_core/utils/maintenance/node_operations.py +0 -468
  28. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/LICENSE +0 -0
  29. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/__init__.py +0 -0
  30. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/cross_encoder/__init__.py +0 -0
  31. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/cross_encoder/bge_reranker_client.py +0 -0
  32. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/cross_encoder/client.py +0 -0
  33. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/cross_encoder/openai_reranker_client.py +0 -0
  34. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/embedder/__init__.py +0 -0
  35. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/errors.py +0 -0
  36. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/__init__.py +0 -0
  37. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/anthropic_client.py +0 -0
  38. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/client.py +0 -0
  39. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/config.py +0 -0
  40. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/errors.py +0 -0
  41. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/gemini_client.py +0 -0
  42. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/groq_client.py +0 -0
  43. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/openai_client.py +0 -0
  44. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/openai_generic_client.py +0 -0
  45. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/llm_client/utils.py +0 -0
  46. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/models/__init__.py +0 -0
  47. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/models/edges/__init__.py +0 -0
  48. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/models/nodes/__init__.py +0 -0
  49. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/models/nodes/node_db_queries.py +0 -0
  50. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/__init__.py +0 -0
  51. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/eval.py +0 -0
  52. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/extract_edge_dates.py +0 -0
  53. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/lib.py +0 -0
  54. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/models.py +0 -0
  55. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/prompt_helpers.py +0 -0
  56. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/prompts/summarize_nodes.py +0 -0
  57. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/py.typed +0 -0
  58. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/__init__.py +0 -0
  59. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search_config.py +0 -0
  60. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search_config_recipes.py +0 -0
  61. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search_filters.py +0 -0
  62. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/search/search_helpers.py +0 -0
  63. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/__init__.py +0 -0
  64. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/datetime_utils.py +0 -0
  65. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/__init__.py +0 -0
  66. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/community_operations.py +0 -0
  67. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/maintenance/utils.py +0 -0
  68. {graphiti_core-0.10.5 → graphiti_core-0.11.0}/graphiti_core/utils/ontology_utils/entity_types_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: graphiti-core
3
- Version: 0.10.5
3
+ Version: 0.11.0
4
4
  Summary: A temporal graph building library
5
5
  License: Apache-2.0
6
6
  Author: Paul Paliychuk
@@ -42,11 +42,11 @@ Graphiti
42
42
  <h2 align="center"> Build Real-Time Knowledge Graphs for AI Agents</h2>
43
43
  <div align="center">
44
44
 
45
-
46
45
  [![Lint](https://github.com/getzep/Graphiti/actions/workflows/lint.yml/badge.svg?style=flat)](https://github.com/getzep/Graphiti/actions/workflows/lint.yml)
47
46
  [![Unit Tests](https://github.com/getzep/Graphiti/actions/workflows/unit_tests.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/unit_tests.yml)
48
47
  [![MyPy Check](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml)
49
48
 
49
+ ![GitHub Repo stars](https://img.shields.io/github/stars/getzep/graphiti)
50
50
  [![Discord](https://dcbadge.vercel.app/api/server/W8Kw6bsgXQ?style=flat)](https://discord.com/invite/W8Kw6bsgXQ)
51
51
  [![arXiv](https://img.shields.io/badge/arXiv-2501.13956-b31b1b.svg?style=flat)](https://arxiv.org/abs/2501.13956)
52
52
  [![Release](https://img.shields.io/github/v/release/getzep/graphiti?style=flat&label=Release&color=limegreen)](https://github.com/getzep/graphiti/releases)
@@ -55,10 +55,16 @@ Graphiti
55
55
  <div align="center">
56
56
 
57
57
  <a href="https://trendshift.io/repositories/12986" target="_blank"><img src="https://trendshift.io/api/badge/repositories/12986" alt="getzep%2Fgraphiti | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
58
+
58
59
  </div>
60
+
59
61
  :star: _Help us reach more developers and grow the Graphiti community. Star this repo!_
62
+
60
63
  <br />
61
64
 
65
+ > [!TIP]
66
+ > Check out the new [MCP server for Graphiti](mcp_server/README.md)! Give Claude, Cursor, and other MCP clients powerful Knowledge Graph-based memory.
67
+
62
68
  Graphiti is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti continuously integrates user interactions, structured and unstructured enterprise data, and external information into a coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI applications.
63
69
 
64
70
  Use Graphiti to:
@@ -191,12 +197,6 @@ For a complete working example, see the [Quickstart Example](./examples/quicksta
191
197
 
192
198
  The example is fully documented with clear explanations of each functionality and includes a comprehensive README with setup instructions and next steps.
193
199
 
194
- ## Graph Service
195
-
196
- The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
197
-
198
- Please see the [server README](./server/README.md) for more information.
199
-
200
200
  ## MCP Server
201
201
 
202
202
  The `mcp_server` directory contains a Model Context Protocol (MCP) server implementation for Graphiti. This server allows AI assistants to interact with Graphiti's knowledge graph capabilities through the MCP protocol.
@@ -213,6 +213,12 @@ The MCP server can be deployed using Docker with Neo4j, making it easy to integr
213
213
 
214
214
  For detailed setup instructions and usage examples, see the [MCP server README](./mcp_server/README.md).
215
215
 
216
+ ## REST Service
217
+
218
+ The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
219
+
220
+ Please see the [server README](./server/README.md) for more information.
221
+
216
222
  ## Optional Environment Variables
217
223
 
218
224
  In addition to the Neo4j and OpenAi-compatible credentials, Graphiti also has a few optional environment variables.
@@ -10,11 +10,11 @@ Graphiti
10
10
  <h2 align="center"> Build Real-Time Knowledge Graphs for AI Agents</h2>
11
11
  <div align="center">
12
12
 
13
-
14
13
  [![Lint](https://github.com/getzep/Graphiti/actions/workflows/lint.yml/badge.svg?style=flat)](https://github.com/getzep/Graphiti/actions/workflows/lint.yml)
15
14
  [![Unit Tests](https://github.com/getzep/Graphiti/actions/workflows/unit_tests.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/unit_tests.yml)
16
15
  [![MyPy Check](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml/badge.svg)](https://github.com/getzep/Graphiti/actions/workflows/typecheck.yml)
17
16
 
17
+ ![GitHub Repo stars](https://img.shields.io/github/stars/getzep/graphiti)
18
18
  [![Discord](https://dcbadge.vercel.app/api/server/W8Kw6bsgXQ?style=flat)](https://discord.com/invite/W8Kw6bsgXQ)
19
19
  [![arXiv](https://img.shields.io/badge/arXiv-2501.13956-b31b1b.svg?style=flat)](https://arxiv.org/abs/2501.13956)
20
20
  [![Release](https://img.shields.io/github/v/release/getzep/graphiti?style=flat&label=Release&color=limegreen)](https://github.com/getzep/graphiti/releases)
@@ -23,10 +23,16 @@ Graphiti
23
23
  <div align="center">
24
24
 
25
25
  <a href="https://trendshift.io/repositories/12986" target="_blank"><img src="https://trendshift.io/api/badge/repositories/12986" alt="getzep%2Fgraphiti | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
26
+
26
27
  </div>
28
+
27
29
  :star: _Help us reach more developers and grow the Graphiti community. Star this repo!_
30
+
28
31
  <br />
29
32
 
33
+ > [!TIP]
34
+ > Check out the new [MCP server for Graphiti](mcp_server/README.md)! Give Claude, Cursor, and other MCP clients powerful Knowledge Graph-based memory.
35
+
30
36
  Graphiti is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti continuously integrates user interactions, structured and unstructured enterprise data, and external information into a coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI applications.
31
37
 
32
38
  Use Graphiti to:
@@ -159,12 +165,6 @@ For a complete working example, see the [Quickstart Example](./examples/quicksta
159
165
 
160
166
  The example is fully documented with clear explanations of each functionality and includes a comprehensive README with setup instructions and next steps.
161
167
 
162
- ## Graph Service
163
-
164
- The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
165
-
166
- Please see the [server README](./server/README.md) for more information.
167
-
168
168
  ## MCP Server
169
169
 
170
170
  The `mcp_server` directory contains a Model Context Protocol (MCP) server implementation for Graphiti. This server allows AI assistants to interact with Graphiti's knowledge graph capabilities through the MCP protocol.
@@ -181,6 +181,12 @@ The MCP server can be deployed using Docker with Neo4j, making it easy to integr
181
181
 
182
182
  For detailed setup instructions and usage examples, see the [MCP server README](./mcp_server/README.md).
183
183
 
184
+ ## REST Service
185
+
186
+ The `server` directory contains an API service for interacting with the Graphiti API. It is built using FastAPI.
187
+
188
+ Please see the [server README](./server/README.md) for more information.
189
+
184
190
  ## Optional Environment Variables
185
191
 
186
192
  In addition to the Neo4j and OpenAi-compatible credentials, Graphiti also has a few optional environment variables.
@@ -37,6 +37,21 @@ from graphiti_core.nodes import Node
37
37
 
38
38
  logger = logging.getLogger(__name__)
39
39
 
40
+ ENTITY_EDGE_RETURN: LiteralString = """
41
+ RETURN
42
+ e.uuid AS uuid,
43
+ startNode(e).uuid AS source_node_uuid,
44
+ endNode(e).uuid AS target_node_uuid,
45
+ e.created_at AS created_at,
46
+ e.name AS name,
47
+ e.group_id AS group_id,
48
+ e.fact AS fact,
49
+ e.fact_embedding AS fact_embedding,
50
+ e.episodes AS episodes,
51
+ e.expired_at AS expired_at,
52
+ e.valid_at AS valid_at,
53
+ e.invalid_at AS invalid_at"""
54
+
40
55
 
41
56
  class Edge(BaseModel, ABC):
42
57
  uuid: str = Field(default_factory=lambda: str(uuid4()))
@@ -234,20 +249,8 @@ class EntityEdge(Edge):
234
249
  records, _, _ = await driver.execute_query(
235
250
  """
236
251
  MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
237
- RETURN
238
- e.uuid AS uuid,
239
- n.uuid AS source_node_uuid,
240
- m.uuid AS target_node_uuid,
241
- e.created_at AS created_at,
242
- e.name AS name,
243
- e.group_id AS group_id,
244
- e.fact AS fact,
245
- e.fact_embedding AS fact_embedding,
246
- e.episodes AS episodes,
247
- e.expired_at AS expired_at,
248
- e.valid_at AS valid_at,
249
- e.invalid_at AS invalid_at
250
- """,
252
+ """
253
+ + ENTITY_EDGE_RETURN,
251
254
  uuid=uuid,
252
255
  database_=DEFAULT_DATABASE,
253
256
  routing_='r',
@@ -268,20 +271,8 @@ class EntityEdge(Edge):
268
271
  """
269
272
  MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
270
273
  WHERE e.uuid IN $uuids
271
- RETURN
272
- e.uuid AS uuid,
273
- n.uuid AS source_node_uuid,
274
- m.uuid AS target_node_uuid,
275
- e.created_at AS created_at,
276
- e.name AS name,
277
- e.group_id AS group_id,
278
- e.fact AS fact,
279
- e.fact_embedding AS fact_embedding,
280
- e.episodes AS episodes,
281
- e.expired_at AS expired_at,
282
- e.valid_at AS valid_at,
283
- e.invalid_at AS invalid_at
284
- """,
274
+ """
275
+ + ENTITY_EDGE_RETURN,
285
276
  uuids=uuids,
286
277
  database_=DEFAULT_DATABASE,
287
278
  routing_='r',
@@ -308,20 +299,8 @@ class EntityEdge(Edge):
308
299
  WHERE e.group_id IN $group_ids
309
300
  """
310
301
  + cursor_query
302
+ + ENTITY_EDGE_RETURN
311
303
  + """
312
- RETURN
313
- e.uuid AS uuid,
314
- n.uuid AS source_node_uuid,
315
- m.uuid AS target_node_uuid,
316
- e.created_at AS created_at,
317
- e.name AS name,
318
- e.group_id AS group_id,
319
- e.fact AS fact,
320
- e.fact_embedding AS fact_embedding,
321
- e.episodes AS episodes,
322
- e.expired_at AS expired_at,
323
- e.valid_at AS valid_at,
324
- e.invalid_at AS invalid_at
325
304
  ORDER BY e.uuid DESC
326
305
  """
327
306
  + limit_query,
@@ -340,22 +319,12 @@ class EntityEdge(Edge):
340
319
 
341
320
  @classmethod
342
321
  async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str):
343
- query: LiteralString = """
344
- MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
345
- RETURN DISTINCT
346
- e.uuid AS uuid,
347
- n.uuid AS source_node_uuid,
348
- m.uuid AS target_node_uuid,
349
- e.created_at AS created_at,
350
- e.name AS name,
351
- e.group_id AS group_id,
352
- e.fact AS fact,
353
- e.fact_embedding AS fact_embedding,
354
- e.episodes AS episodes,
355
- e.expired_at AS expired_at,
356
- e.valid_at AS valid_at,
357
- e.invalid_at AS invalid_at
358
- """
322
+ query: LiteralString = (
323
+ """
324
+ MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
325
+ """
326
+ + ENTITY_EDGE_RETURN
327
+ )
359
328
  records, _, _ = await driver.execute_query(
360
329
  query, node_uuid=node_uuid, database_=DEFAULT_DATABASE, routing_='r'
361
330
  )
@@ -499,3 +468,9 @@ def get_community_edge_from_record(record: Any):
499
468
  target_node_uuid=record['target_node_uuid'],
500
469
  created_at=record['created_at'].to_native(),
501
470
  )
471
+
472
+
473
+ async def create_entity_edge_embeddings(embedder: EmbedderClient, edges: list[EntityEdge]):
474
+ fact_embeddings = await embedder.create_batch([edge.fact for edge in edges])
475
+ for edge, fact_embedding in zip(edges, fact_embeddings, strict=True):
476
+ edge.fact_embedding = fact_embedding
@@ -32,3 +32,6 @@ class EmbedderClient(ABC):
32
32
  self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
33
33
  ) -> list[float]:
34
34
  pass
35
+
36
+ async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
37
+ raise NotImplementedError()
@@ -66,3 +66,13 @@ class GeminiEmbedder(EmbedderClient):
66
66
  )
67
67
 
68
68
  return result.embeddings[0].values
69
+
70
+ async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
71
+ # Generate embeddings
72
+ result = await self.client.aio.models.embed_content(
73
+ model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
74
+ contents=input_data_list,
75
+ config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
76
+ )
77
+
78
+ return [embedding.values for embedding in result.embeddings]
@@ -58,3 +58,9 @@ class OpenAIEmbedder(EmbedderClient):
58
58
  input=input_data, model=self.config.embedding_model
59
59
  )
60
60
  return result.data[0].embedding[: self.config.embedding_dim]
61
+
62
+ async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
63
+ result = await self.client.embeddings.create(
64
+ input=input_data_list, model=self.config.embedding_model
65
+ )
66
+ return [embedding.embedding[: self.config.embedding_dim] for embedding in result.data]
@@ -56,3 +56,10 @@ class VoyageAIEmbedder(EmbedderClient):
56
56
 
57
57
  result = await self.client.embed(input_list, model=self.config.embedding_model)
58
58
  return [float(x) for x in result.embeddings[0][: self.config.embedding_dim]]
59
+
60
+ async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
61
+ result = await self.client.embed(input_data_list, model=self.config.embedding_model)
62
+ return [
63
+ [float(x) for x in embedding[: self.config.embedding_dim]]
64
+ for embedding in result.embeddings
65
+ ]
@@ -27,6 +27,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
27
27
  from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
28
28
  from graphiti_core.edges import EntityEdge, EpisodicEdge
29
29
  from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
30
+ from graphiti_core.graphiti_types import GraphitiClients
30
31
  from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
31
32
  from graphiti_core.llm_client import LLMClient, OpenAIClient
32
33
  from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
@@ -42,7 +43,6 @@ from graphiti_core.search.search_utils import (
42
43
  RELEVANT_SCHEMA_LIMIT,
43
44
  get_mentioned_nodes,
44
45
  get_relevant_edges,
45
- get_relevant_nodes,
46
46
  )
47
47
  from graphiti_core.utils.bulk_utils import (
48
48
  RawEpisode,
@@ -72,7 +72,11 @@ from graphiti_core.utils.maintenance.graph_data_operations import (
72
72
  build_indices_and_constraints,
73
73
  retrieve_episodes,
74
74
  )
75
- from graphiti_core.utils.maintenance.node_operations import extract_nodes, resolve_extracted_nodes
75
+ from graphiti_core.utils.maintenance.node_operations import (
76
+ extract_attributes_from_nodes,
77
+ extract_nodes,
78
+ resolve_extracted_nodes,
79
+ )
76
80
  from graphiti_core.utils.maintenance.temporal_operations import get_edge_contradictions
77
81
  from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types
78
82
 
@@ -150,6 +154,13 @@ class Graphiti:
150
154
  else:
151
155
  self.cross_encoder = OpenAIRerankerClient()
152
156
 
157
+ self.clients = GraphitiClients(
158
+ driver=self.driver,
159
+ llm_client=self.llm_client,
160
+ embedder=self.embedder,
161
+ cross_encoder=self.cross_encoder,
162
+ )
163
+
153
164
  async def close(self):
154
165
  """
155
166
  Close the connection to the Neo4j database.
@@ -222,6 +233,7 @@ class Graphiti:
222
233
  reference_time: datetime,
223
234
  last_n: int = EPISODE_WINDOW_LEN,
224
235
  group_ids: list[str] | None = None,
236
+ source: EpisodeType | None = None,
225
237
  ) -> list[EpisodicNode]:
226
238
  """
227
239
  Retrieve the last n episodic nodes from the graph.
@@ -248,7 +260,7 @@ class Graphiti:
248
260
  The actual retrieval is performed by the `retrieve_episodes` function
249
261
  from the `graphiti_core.utils` module.
250
262
  """
251
- return await retrieve_episodes(self.driver, reference_time, last_n, group_ids)
263
+ return await retrieve_episodes(self.driver, reference_time, last_n, group_ids, source)
252
264
 
253
265
  async def add_episode(
254
266
  self,
@@ -314,15 +326,16 @@ class Graphiti:
314
326
  """
315
327
  try:
316
328
  start = time()
317
-
318
- entity_edges: list[EntityEdge] = []
319
329
  now = utc_now()
320
330
 
321
331
  validate_entity_types(entity_types)
322
332
 
323
333
  previous_episodes = (
324
334
  await self.retrieve_episodes(
325
- reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
335
+ reference_time,
336
+ last_n=RELEVANT_SCHEMA_LIMIT,
337
+ group_ids=[group_id],
338
+ source=source,
326
339
  )
327
340
  if previous_episode_uuids is None
328
341
  else await EpisodicNode.get_by_uuids(self.driver, previous_episode_uuids)
@@ -346,132 +359,36 @@ class Graphiti:
346
359
  # Extract entities as nodes
347
360
 
348
361
  extracted_nodes = await extract_nodes(
349
- self.llm_client, episode, previous_episodes, entity_types
362
+ self.clients, episode, previous_episodes, entity_types
350
363
  )
351
- logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
352
-
353
- # Calculate Embeddings
354
364
 
355
- await semaphore_gather(
356
- *[node.generate_name_embedding(self.embedder) for node in extracted_nodes]
357
- )
358
-
359
- # Find relevant nodes already in the graph
360
- existing_nodes_lists: list[list[EntityNode]] = list(
361
- await semaphore_gather(
362
- *[
363
- get_relevant_nodes(self.driver, SearchFilters(), [node])
364
- for node in extracted_nodes
365
- ]
366
- )
367
- )
368
-
369
- # Resolve extracted nodes with nodes already in the graph and extract facts
370
- logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
371
-
372
- (mentioned_nodes, uuid_map), extracted_edges = await semaphore_gather(
365
+ # Extract edges and resolve nodes
366
+ (nodes, uuid_map), extracted_edges = await semaphore_gather(
373
367
  resolve_extracted_nodes(
374
- self.llm_client,
368
+ self.clients,
375
369
  extracted_nodes,
376
- existing_nodes_lists,
377
370
  episode,
378
371
  previous_episodes,
379
372
  entity_types,
380
373
  ),
381
- extract_edges(
382
- self.llm_client, episode, extracted_nodes, previous_episodes, group_id
383
- ),
384
- )
385
- logger.debug(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
386
- nodes = mentioned_nodes
387
-
388
- extracted_edges_with_resolved_pointers = resolve_edge_pointers(
389
- extracted_edges, uuid_map
374
+ extract_edges(self.clients, episode, extracted_nodes, previous_episodes, group_id),
390
375
  )
391
376
 
392
- # calculate embeddings
393
- await semaphore_gather(
394
- *[
395
- edge.generate_embedding(self.embedder)
396
- for edge in extracted_edges_with_resolved_pointers
397
- ]
398
- )
377
+ edges = resolve_edge_pointers(extracted_edges, uuid_map)
399
378
 
400
- # Resolve extracted edges with related edges already in the graph
401
- related_edges_list: list[list[EntityEdge]] = list(
402
- await semaphore_gather(
403
- *[
404
- get_relevant_edges(
405
- self.driver,
406
- [edge],
407
- edge.source_node_uuid,
408
- edge.target_node_uuid,
409
- RELEVANT_SCHEMA_LIMIT,
410
- )
411
- for edge in extracted_edges_with_resolved_pointers
412
- ]
413
- )
414
- )
415
- logger.debug(
416
- f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_list for e in edges_lst]}'
417
- )
418
- logger.debug(
419
- f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges_with_resolved_pointers]}'
420
- )
421
-
422
- existing_source_edges_list: list[list[EntityEdge]] = list(
423
- await semaphore_gather(
424
- *[
425
- get_relevant_edges(
426
- self.driver,
427
- [edge],
428
- edge.source_node_uuid,
429
- None,
430
- RELEVANT_SCHEMA_LIMIT,
431
- )
432
- for edge in extracted_edges_with_resolved_pointers
433
- ]
434
- )
435
- )
436
-
437
- existing_target_edges_list: list[list[EntityEdge]] = list(
438
- await semaphore_gather(
439
- *[
440
- get_relevant_edges(
441
- self.driver,
442
- [edge],
443
- None,
444
- edge.target_node_uuid,
445
- RELEVANT_SCHEMA_LIMIT,
446
- )
447
- for edge in extracted_edges_with_resolved_pointers
448
- ]
449
- )
450
- )
451
-
452
- existing_edges_list: list[list[EntityEdge]] = [
453
- source_lst + target_lst
454
- for source_lst, target_lst in zip(
455
- existing_source_edges_list, existing_target_edges_list, strict=False
456
- )
457
- ]
458
-
459
- resolved_edges, invalidated_edges = await resolve_extracted_edges(
460
- self.llm_client,
461
- extracted_edges_with_resolved_pointers,
462
- related_edges_list,
463
- existing_edges_list,
464
- episode,
465
- previous_episodes,
379
+ (resolved_edges, invalidated_edges), hydrated_nodes = await semaphore_gather(
380
+ resolve_extracted_edges(
381
+ self.clients,
382
+ edges,
383
+ ),
384
+ extract_attributes_from_nodes(
385
+ self.clients, nodes, episode, previous_episodes, entity_types
386
+ ),
466
387
  )
467
388
 
468
- entity_edges.extend(resolved_edges + invalidated_edges)
389
+ entity_edges = resolved_edges + invalidated_edges
469
390
 
470
- logger.debug(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_edges]}')
471
-
472
- episodic_edges: list[EpisodicEdge] = build_episodic_edges(mentioned_nodes, episode, now)
473
-
474
- logger.debug(f'Built episodic edges: {episodic_edges}')
391
+ episodic_edges = build_episodic_edges(nodes, episode, now)
475
392
 
476
393
  episode.entity_edges = [edge.uuid for edge in entity_edges]
477
394
 
@@ -565,7 +482,7 @@ class Graphiti:
565
482
  extracted_nodes,
566
483
  extracted_edges,
567
484
  episodic_edges,
568
- ) = await extract_nodes_and_edges_bulk(self.llm_client, episode_pairs)
485
+ ) = await extract_nodes_and_edges_bulk(self.clients, episode_pairs)
569
486
 
570
487
  # Generate embeddings
571
488
  await semaphore_gather(
@@ -684,9 +601,7 @@ class Graphiti:
684
601
 
685
602
  edges = (
686
603
  await search(
687
- self.driver,
688
- self.embedder,
689
- self.cross_encoder,
604
+ self.clients,
690
605
  query,
691
606
  group_ids,
692
607
  search_config,
@@ -728,9 +643,7 @@ class Graphiti:
728
643
  """
729
644
 
730
645
  return await search(
731
- self.driver,
732
- self.embedder,
733
- self.cross_encoder,
646
+ self.clients,
734
647
  query,
735
648
  group_ids,
736
649
  config,
@@ -761,26 +674,17 @@ class Graphiti:
761
674
  await edge.generate_embedding(self.embedder)
762
675
 
763
676
  resolved_nodes, uuid_map = await resolve_extracted_nodes(
764
- self.llm_client,
677
+ self.clients,
765
678
  [source_node, target_node],
766
- [
767
- await get_relevant_nodes(self.driver, SearchFilters(), [source_node]),
768
- await get_relevant_nodes(self.driver, SearchFilters(), [target_node]),
769
- ],
770
679
  )
771
680
 
772
681
  updated_edge = resolve_edge_pointers([edge], uuid_map)[0]
773
682
 
774
- related_edges = await get_relevant_edges(
775
- self.driver,
776
- [updated_edge],
777
- source_node_uuid=resolved_nodes[0].uuid,
778
- target_node_uuid=resolved_nodes[1].uuid,
779
- )
683
+ related_edges = await get_relevant_edges(self.driver, [updated_edge], SearchFilters(), 0.8)
780
684
 
781
- resolved_edge = await dedupe_extracted_edge(self.llm_client, updated_edge, related_edges)
685
+ resolved_edge = await dedupe_extracted_edge(self.llm_client, updated_edge, related_edges[0])
782
686
 
783
- contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges)
687
+ contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges[0])
784
688
  invalidated_edges = resolve_edge_contradictions(resolved_edge, contradicting_edges)
785
689
 
786
690
  await add_nodes_and_edges_bulk(
@@ -0,0 +1,31 @@
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 neo4j import AsyncDriver
18
+ from pydantic import BaseModel, ConfigDict
19
+
20
+ from graphiti_core.cross_encoder import CrossEncoderClient
21
+ from graphiti_core.embedder import EmbedderClient
22
+ from graphiti_core.llm_client import LLMClient
23
+
24
+
25
+ class GraphitiClients(BaseModel):
26
+ driver: AsyncDriver
27
+ llm_client: LLMClient
28
+ embedder: EmbedderClient
29
+ cross_encoder: CrossEncoderClient
30
+
31
+ model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -22,15 +22,20 @@ from datetime import datetime
22
22
  import numpy as np
23
23
  from dotenv import load_dotenv
24
24
  from neo4j import time as neo4j_time
25
+ from typing_extensions import LiteralString
25
26
 
26
27
  load_dotenv()
27
28
 
28
29
  DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
29
30
  USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
30
31
  SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
31
- MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 2))
32
+ MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0))
32
33
  DEFAULT_PAGE_LIMIT = 20
33
34
 
35
+ RUNTIME_QUERY: LiteralString = (
36
+ 'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
37
+ )
38
+
34
39
 
35
40
  def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
36
41
  return neo_date.to_native() if neo_date else None
@@ -47,7 +47,7 @@ ENTITY_EDGE_SAVE_BULK = """
47
47
  SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
48
48
  created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at}
49
49
  WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
50
- RETURN r.uuid AS uuid
50
+ RETURN edge.uuid AS uuid
51
51
  """
52
52
 
53
53
  COMMUNITY_EDGE_SAVE = """
@@ -332,8 +332,8 @@ class EntityNode(Node):
332
332
  async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
333
333
  query = (
334
334
  """
335
- MATCH (n:Entity {uuid: $uuid})
336
- """
335
+ MATCH (n:Entity {uuid: $uuid})
336
+ """
337
337
  + ENTITY_NODE_RETURN
338
338
  )
339
339
  records, _, _ = await driver.execute_query(
@@ -560,3 +560,9 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
560
560
  created_at=record['created_at'].to_native(),
561
561
  summary=record['summary'],
562
562
  )
563
+
564
+
565
+ async def create_entity_node_embeddings(embedder: EmbedderClient, nodes: list[EntityNode]):
566
+ name_embeddings = await embedder.create_batch([node.name for node in nodes])
567
+ for node, name_embedding in zip(nodes, name_embeddings, strict=True):
568
+ node.name_embedding = name_embedding