graphiti-core 0.5.0rc2__tar.gz → 0.5.0rc4__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 (60) hide show
  1. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/PKG-INFO +1 -1
  2. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/edges.py +13 -10
  3. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/graphiti.py +6 -7
  4. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/helpers.py +6 -0
  5. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/openai_client.py +45 -4
  6. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/nodes.py +16 -12
  7. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/search/search_utils.py +6 -2
  8. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/bulk_utils.py +7 -4
  9. graphiti_core-0.5.0rc4/graphiti_core/utils/datetime_utils.py +42 -0
  10. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/community_operations.py +3 -3
  11. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/edge_operations.py +15 -11
  12. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/node_operations.py +2 -2
  13. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/temporal_operations.py +6 -3
  14. graphiti_core-0.5.0rc4/graphiti_core/utils/maintenance/utils.py +0 -0
  15. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/pyproject.toml +1 -1
  16. graphiti_core-0.5.0rc2/graphiti_core/utils/__init__.py +0 -15
  17. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/LICENSE +0 -0
  18. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/README.md +0 -0
  19. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/__init__.py +0 -0
  20. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/cross_encoder/__init__.py +0 -0
  21. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/cross_encoder/bge_reranker_client.py +0 -0
  22. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/cross_encoder/client.py +0 -0
  23. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/cross_encoder/openai_reranker_client.py +0 -0
  24. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/embedder/__init__.py +0 -0
  25. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/embedder/client.py +0 -0
  26. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/embedder/openai.py +0 -0
  27. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/embedder/voyage.py +0 -0
  28. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/errors.py +0 -0
  29. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/__init__.py +0 -0
  30. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/anthropic_client.py +0 -0
  31. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/client.py +0 -0
  32. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/config.py +0 -0
  33. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/errors.py +0 -0
  34. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/groq_client.py +0 -0
  35. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/llm_client/utils.py +0 -0
  36. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/models/__init__.py +0 -0
  37. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/models/edges/__init__.py +0 -0
  38. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/models/edges/edge_db_queries.py +0 -0
  39. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/models/nodes/__init__.py +0 -0
  40. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/models/nodes/node_db_queries.py +0 -0
  41. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/__init__.py +0 -0
  42. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/dedupe_edges.py +0 -0
  43. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/dedupe_nodes.py +0 -0
  44. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/eval.py +0 -0
  45. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/extract_edge_dates.py +0 -0
  46. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/extract_edges.py +0 -0
  47. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/extract_nodes.py +0 -0
  48. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/invalidate_edges.py +0 -0
  49. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/lib.py +0 -0
  50. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/models.py +0 -0
  51. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/prompt_helpers.py +0 -0
  52. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/prompts/summarize_nodes.py +0 -0
  53. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/py.typed +0 -0
  54. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/search/__init__.py +0 -0
  55. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/search/search.py +0 -0
  56. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/search/search_config.py +0 -0
  57. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/search/search_config_recipes.py +0 -0
  58. /graphiti_core-0.5.0rc2/graphiti_core/utils/maintenance/utils.py → /graphiti_core-0.5.0rc4/graphiti_core/utils/__init__.py +0 -0
  59. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/__init__.py +0 -0
  60. {graphiti_core-0.5.0rc2 → graphiti_core-0.5.0rc4}/graphiti_core/utils/maintenance/graph_data_operations.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: graphiti-core
3
- Version: 0.5.0rc2
3
+ Version: 0.5.0rc4
4
4
  Summary: A temporal graph building library
5
5
  License: Apache-2.0
6
6
  Author: Paul Paliychuk
@@ -27,7 +27,7 @@ from typing_extensions import LiteralString
27
27
 
28
28
  from graphiti_core.embedder import EmbedderClient
29
29
  from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
30
- from graphiti_core.helpers import DEFAULT_DATABASE, DEFAULT_PAGE_LIMIT, parse_db_date
30
+ from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
31
31
  from graphiti_core.models.edges.edge_db_queries import (
32
32
  COMMUNITY_EDGE_SAVE,
33
33
  ENTITY_EDGE_SAVE,
@@ -142,10 +142,11 @@ class EpisodicEdge(Edge):
142
142
  cls,
143
143
  driver: AsyncDriver,
144
144
  group_ids: list[str],
145
- limit: int = DEFAULT_PAGE_LIMIT,
145
+ limit: int | None = None,
146
146
  created_at: datetime | None = None,
147
147
  ):
148
148
  cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
149
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
149
150
 
150
151
  records, _, _ = await driver.execute_query(
151
152
  """
@@ -161,8 +162,8 @@ class EpisodicEdge(Edge):
161
162
  m.uuid AS target_node_uuid,
162
163
  e.created_at AS created_at
163
164
  ORDER BY e.uuid DESC
164
- LIMIT $limit
165
- """,
165
+ """
166
+ + limit_query,
166
167
  group_ids=group_ids,
167
168
  created_at=created_at,
168
169
  limit=limit,
@@ -294,10 +295,11 @@ class EntityEdge(Edge):
294
295
  cls,
295
296
  driver: AsyncDriver,
296
297
  group_ids: list[str],
297
- limit: int = DEFAULT_PAGE_LIMIT,
298
+ limit: int | None = None,
298
299
  created_at: datetime | None = None,
299
300
  ):
300
301
  cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
302
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
301
303
 
302
304
  records, _, _ = await driver.execute_query(
303
305
  """
@@ -320,8 +322,8 @@ class EntityEdge(Edge):
320
322
  e.valid_at AS valid_at,
321
323
  e.invalid_at AS invalid_at
322
324
  ORDER BY e.uuid DESC
323
- LIMIT $limit
324
- """,
325
+ """
326
+ + limit_query,
325
327
  group_ids=group_ids,
326
328
  created_at=created_at,
327
329
  limit=limit,
@@ -400,10 +402,11 @@ class CommunityEdge(Edge):
400
402
  cls,
401
403
  driver: AsyncDriver,
402
404
  group_ids: list[str],
403
- limit: int = DEFAULT_PAGE_LIMIT,
405
+ limit: int | None = None,
404
406
  created_at: datetime | None = None,
405
407
  ):
406
408
  cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
409
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
407
410
 
408
411
  records, _, _ = await driver.execute_query(
409
412
  """
@@ -419,8 +422,8 @@ class CommunityEdge(Edge):
419
422
  m.uuid AS target_node_uuid,
420
423
  e.created_at AS created_at
421
424
  ORDER BY e.uuid DESC
422
- LIMIT $limit
423
- """,
425
+ """
426
+ + limit_query,
424
427
  group_ids=group_ids,
425
428
  created_at=created_at,
426
429
  limit=limit,
@@ -16,7 +16,7 @@ limitations under the License.
16
16
 
17
17
  import asyncio
18
18
  import logging
19
- from datetime import datetime, timezone
19
+ from datetime import datetime
20
20
  from time import time
21
21
 
22
22
  from dotenv import load_dotenv
@@ -43,10 +43,6 @@ from graphiti_core.search.search_utils import (
43
43
  get_relevant_edges,
44
44
  get_relevant_nodes,
45
45
  )
46
- from graphiti_core.utils import (
47
- build_episodic_edges,
48
- retrieve_episodes,
49
- )
50
46
  from graphiti_core.utils.bulk_utils import (
51
47
  RawEpisode,
52
48
  add_nodes_and_edges_bulk,
@@ -57,12 +53,14 @@ from graphiti_core.utils.bulk_utils import (
57
53
  resolve_edge_pointers,
58
54
  retrieve_previous_episodes_bulk,
59
55
  )
56
+ from graphiti_core.utils.datetime_utils import utc_now
60
57
  from graphiti_core.utils.maintenance.community_operations import (
61
58
  build_communities,
62
59
  remove_communities,
63
60
  update_community,
64
61
  )
65
62
  from graphiti_core.utils.maintenance.edge_operations import (
63
+ build_episodic_edges,
66
64
  dedupe_extracted_edge,
67
65
  extract_edges,
68
66
  resolve_edge_contradictions,
@@ -71,6 +69,7 @@ from graphiti_core.utils.maintenance.edge_operations import (
71
69
  from graphiti_core.utils.maintenance.graph_data_operations import (
72
70
  EPISODE_WINDOW_LEN,
73
71
  build_indices_and_constraints,
72
+ retrieve_episodes,
74
73
  )
75
74
  from graphiti_core.utils.maintenance.node_operations import (
76
75
  extract_nodes,
@@ -313,7 +312,7 @@ class Graphiti:
313
312
  start = time()
314
313
 
315
314
  entity_edges: list[EntityEdge] = []
316
- now = datetime.now(timezone.utc)
315
+ now = utc_now()
317
316
 
318
317
  previous_episodes = await self.retrieve_episodes(
319
318
  reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
@@ -522,7 +521,7 @@ class Graphiti:
522
521
  """
523
522
  try:
524
523
  start = time()
525
- now = datetime.now(timezone.utc)
524
+ now = utc_now()
526
525
 
527
526
  episodes = [
528
527
  EpisodicNode(
@@ -57,6 +57,12 @@ def lucene_sanitize(query: str) -> str:
57
57
  ':': r'\:',
58
58
  '\\': r'\\',
59
59
  '/': r'\/',
60
+ 'O': r'\O',
61
+ 'R': r'\R',
62
+ 'N': r'\N',
63
+ 'T': r'\T',
64
+ 'A': r'\A',
65
+ 'D': r'\D',
60
66
  }
61
67
  )
62
68
 
@@ -16,6 +16,7 @@ limitations under the License.
16
16
 
17
17
  import logging
18
18
  import typing
19
+ from typing import ClassVar
19
20
 
20
21
  import openai
21
22
  from openai import AsyncOpenAI
@@ -53,6 +54,9 @@ class OpenAIClient(LLMClient):
53
54
  Generates a response from the language model based on the provided messages.
54
55
  """
55
56
 
57
+ # Class-level constants
58
+ MAX_RETRIES: ClassVar[int] = 2
59
+
56
60
  def __init__(
57
61
  self, config: LLMConfig | None = None, cache: bool = False, client: typing.Any = None
58
62
  ):
@@ -104,7 +108,7 @@ class OpenAIClient(LLMClient):
104
108
  elif response_object.refusal:
105
109
  raise RefusalError(response_object.refusal)
106
110
  else:
107
- raise Exception('No response from LLM')
111
+ raise Exception(f'Invalid response from LLM: {response_object.model_dump()}')
108
112
  except openai.LengthFinishReasonError as e:
109
113
  raise Exception(f'Output length exceeded max tokens {self.max_tokens}: {e}') from e
110
114
  except openai.RateLimitError as e:
@@ -116,6 +120,43 @@ class OpenAIClient(LLMClient):
116
120
  async def generate_response(
117
121
  self, messages: list[Message], response_model: type[BaseModel] | None = None
118
122
  ) -> dict[str, typing.Any]:
119
- response = await self._generate_response(messages, response_model)
120
-
121
- return response
123
+ retry_count = 0
124
+ last_error = None
125
+
126
+ while retry_count <= self.MAX_RETRIES:
127
+ try:
128
+ response = await self._generate_response(messages, response_model)
129
+ return response
130
+ except (RateLimitError, RefusalError):
131
+ # These errors should not trigger retries
132
+ raise
133
+ except (openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError):
134
+ # Let OpenAI's client handle these retries
135
+ raise
136
+ except Exception as e:
137
+ last_error = e
138
+
139
+ # Don't retry if we've hit the max retries
140
+ if retry_count >= self.MAX_RETRIES:
141
+ logger.error(f'Max retries ({self.MAX_RETRIES}) exceeded. Last error: {e}')
142
+ raise
143
+
144
+ retry_count += 1
145
+
146
+ # Construct a detailed error message for the LLM
147
+ error_context = (
148
+ f'The previous response attempt was invalid. '
149
+ f'Error type: {e.__class__.__name__}. '
150
+ f'Error details: {str(e)}. '
151
+ f'Please try again with a valid response, ensuring the output matches '
152
+ f'the expected format and constraints.'
153
+ )
154
+
155
+ error_message = Message(role='user', content=error_context)
156
+ messages.append(error_message)
157
+ logger.warning(
158
+ f'Retrying after application error (attempt {retry_count}/{self.MAX_RETRIES}): {e}'
159
+ )
160
+
161
+ # If we somehow get here, raise the last error
162
+ raise last_error or Exception('Max retries exceeded with no specific error')
@@ -16,7 +16,7 @@ limitations under the License.
16
16
 
17
17
  import logging
18
18
  from abc import ABC, abstractmethod
19
- from datetime import datetime, timezone
19
+ from datetime import datetime
20
20
  from enum import Enum
21
21
  from time import time
22
22
  from typing import Any
@@ -28,12 +28,13 @@ from typing_extensions import LiteralString
28
28
 
29
29
  from graphiti_core.embedder import EmbedderClient
30
30
  from graphiti_core.errors import NodeNotFoundError
31
- from graphiti_core.helpers import DEFAULT_DATABASE, DEFAULT_PAGE_LIMIT
31
+ from graphiti_core.helpers import DEFAULT_DATABASE
32
32
  from graphiti_core.models.nodes.node_db_queries import (
33
33
  COMMUNITY_NODE_SAVE,
34
34
  ENTITY_NODE_SAVE,
35
35
  EPISODIC_NODE_SAVE,
36
36
  )
37
+ from graphiti_core.utils.datetime_utils import utc_now
37
38
 
38
39
  logger = logging.getLogger(__name__)
39
40
 
@@ -79,7 +80,7 @@ class Node(BaseModel, ABC):
79
80
  name: str = Field(description='name of the node')
80
81
  group_id: str = Field(description='partition of the graph')
81
82
  labels: list[str] = Field(default_factory=list)
82
- created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
83
+ created_at: datetime = Field(default_factory=lambda: utc_now())
83
84
 
84
85
  @abstractmethod
85
86
  async def save(self, driver: AsyncDriver): ...
@@ -212,10 +213,11 @@ class EpisodicNode(Node):
212
213
  cls,
213
214
  driver: AsyncDriver,
214
215
  group_ids: list[str],
215
- limit: int = DEFAULT_PAGE_LIMIT,
216
+ limit: int | None = None,
216
217
  created_at: datetime | None = None,
217
218
  ):
218
219
  cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
220
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
219
221
 
220
222
  records, _, _ = await driver.execute_query(
221
223
  """
@@ -233,8 +235,8 @@ class EpisodicNode(Node):
233
235
  e.source_description AS source_description,
234
236
  e.source AS source
235
237
  ORDER BY e.uuid DESC
236
- LIMIT $limit
237
- """,
238
+ """
239
+ + limit_query,
238
240
  group_ids=group_ids,
239
241
  created_at=created_at,
240
242
  limit=limit,
@@ -328,10 +330,11 @@ class EntityNode(Node):
328
330
  cls,
329
331
  driver: AsyncDriver,
330
332
  group_ids: list[str],
331
- limit: int = DEFAULT_PAGE_LIMIT,
333
+ limit: int | None = None,
332
334
  created_at: datetime | None = None,
333
335
  ):
334
336
  cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
337
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
335
338
 
336
339
  records, _, _ = await driver.execute_query(
337
340
  """
@@ -347,8 +350,8 @@ class EntityNode(Node):
347
350
  n.created_at AS created_at,
348
351
  n.summary AS summary
349
352
  ORDER BY n.uuid DESC
350
- LIMIT $limit
351
- """,
353
+ """
354
+ + limit_query,
352
355
  group_ids=group_ids,
353
356
  created_at=created_at,
354
357
  limit=limit,
@@ -442,10 +445,11 @@ class CommunityNode(Node):
442
445
  cls,
443
446
  driver: AsyncDriver,
444
447
  group_ids: list[str],
445
- limit: int = DEFAULT_PAGE_LIMIT,
448
+ limit: int | None = None,
446
449
  created_at: datetime | None = None,
447
450
  ):
448
451
  cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
452
+ limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
449
453
 
450
454
  records, _, _ = await driver.execute_query(
451
455
  """
@@ -461,8 +465,8 @@ class CommunityNode(Node):
461
465
  n.created_at AS created_at,
462
466
  n.summary AS summary
463
467
  ORDER BY n.uuid DESC
464
- LIMIT $limit
465
- """,
468
+ """
469
+ + limit_query,
466
470
  group_ids=group_ids,
467
471
  created_at=created_at,
468
472
  limit=limit,
@@ -631,7 +631,7 @@ async def node_distance_reranker(
631
631
  ) -> list[str]:
632
632
  # filter out node_uuid center node node uuid
633
633
  filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
634
- scores: dict[str, float] = {}
634
+ scores: dict[str, float] = {center_node_uuid: 0.0}
635
635
 
636
636
  # Find the shortest path to center node
637
637
  query = Query("""
@@ -649,9 +649,13 @@ async def node_distance_reranker(
649
649
 
650
650
  for result in path_results:
651
651
  uuid = result['uuid']
652
- score = result['score'] if 'score' in result else float('inf')
652
+ score = result['score']
653
653
  scores[uuid] = score
654
654
 
655
+ for uuid in filtered_uuids:
656
+ if uuid not in scores:
657
+ scores[uuid] = float('inf')
658
+
655
659
  # rerank on shortest distance
656
660
  filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
657
661
 
@@ -18,7 +18,7 @@ import asyncio
18
18
  import logging
19
19
  import typing
20
20
  from collections import defaultdict
21
- from datetime import datetime, timezone
21
+ from datetime import datetime
22
22
  from math import ceil
23
23
 
24
24
  from neo4j import AsyncDriver, AsyncManagedTransaction
@@ -37,14 +37,17 @@ from graphiti_core.models.nodes.node_db_queries import (
37
37
  )
38
38
  from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
39
39
  from graphiti_core.search.search_utils import get_relevant_edges, get_relevant_nodes
40
- from graphiti_core.utils import retrieve_episodes
40
+ from graphiti_core.utils.datetime_utils import utc_now
41
41
  from graphiti_core.utils.maintenance.edge_operations import (
42
42
  build_episodic_edges,
43
43
  dedupe_edge_list,
44
44
  dedupe_extracted_edges,
45
45
  extract_edges,
46
46
  )
47
- from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
47
+ from graphiti_core.utils.maintenance.graph_data_operations import (
48
+ EPISODE_WINDOW_LEN,
49
+ retrieve_episodes,
50
+ )
48
51
  from graphiti_core.utils.maintenance.node_operations import (
49
52
  dedupe_extracted_nodes,
50
53
  dedupe_node_list,
@@ -385,7 +388,7 @@ async def extract_edge_dates_bulk(
385
388
  edge.valid_at = valid_at
386
389
  edge.invalid_at = invalid_at
387
390
  if edge.invalid_at:
388
- edge.expired_at = datetime.now(timezone.utc)
391
+ edge.expired_at = utc_now()
389
392
 
390
393
  return edges
391
394
 
@@ -0,0 +1,42 @@
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, timezone
18
+
19
+
20
+ def utc_now() -> datetime:
21
+ """Returns the current UTC datetime with timezone information."""
22
+ return datetime.now(timezone.utc)
23
+
24
+
25
+ def ensure_utc(dt: datetime | None) -> datetime | None:
26
+ """
27
+ Ensures a datetime is timezone-aware and in UTC.
28
+ If the datetime is naive (no timezone), assumes it's in UTC.
29
+ If the datetime has a different timezone, converts it to UTC.
30
+ Returns None if input is None.
31
+ """
32
+ if dt is None:
33
+ return None
34
+
35
+ if dt.tzinfo is None:
36
+ # If datetime is naive, assume it's UTC
37
+ return dt.replace(tzinfo=timezone.utc)
38
+ elif dt.tzinfo != timezone.utc:
39
+ # If datetime has a different timezone, convert to UTC
40
+ return dt.astimezone(timezone.utc)
41
+
42
+ return dt
@@ -1,7 +1,6 @@
1
1
  import asyncio
2
2
  import logging
3
3
  from collections import defaultdict
4
- from datetime import datetime, timezone
5
4
 
6
5
  from neo4j import AsyncDriver
7
6
  from pydantic import BaseModel
@@ -17,6 +16,7 @@ from graphiti_core.nodes import (
17
16
  )
18
17
  from graphiti_core.prompts import prompt_library
19
18
  from graphiti_core.prompts.summarize_nodes import Summary, SummaryDescription
19
+ from graphiti_core.utils.datetime_utils import utc_now
20
20
  from graphiti_core.utils.maintenance.edge_operations import build_community_edges
21
21
 
22
22
  MAX_COMMUNITY_BUILD_CONCURRENCY = 10
@@ -180,7 +180,7 @@ async def build_community(
180
180
 
181
181
  summary = summaries[0]
182
182
  name = await generate_summary_description(llm_client, summary)
183
- now = datetime.now(timezone.utc)
183
+ now = utc_now()
184
184
  community_node = CommunityNode(
185
185
  name=name,
186
186
  group_id=community_cluster[0].group_id,
@@ -307,7 +307,7 @@ async def update_community(
307
307
  community.name = new_name
308
308
 
309
309
  if is_new:
310
- community_edge = (build_community_edges([entity], community, datetime.now(timezone.utc)))[0]
310
+ community_edge = (build_community_edges([entity], community, utc_now()))[0]
311
311
  await community_edge.save(driver)
312
312
 
313
313
  await community.generate_name_embedding(embedder)
@@ -16,7 +16,7 @@ limitations under the License.
16
16
 
17
17
  import asyncio
18
18
  import logging
19
- from datetime import datetime, timezone
19
+ from datetime import datetime
20
20
  from time import time
21
21
 
22
22
  from graphiti_core.edges import CommunityEdge, EntityEdge, EpisodicEdge
@@ -26,6 +26,7 @@ from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
26
26
  from graphiti_core.prompts import prompt_library
27
27
  from graphiti_core.prompts.dedupe_edges import EdgeDuplicate, UniqueFacts
28
28
  from graphiti_core.prompts.extract_edges import ExtractedEdges, MissingFacts
29
+ from graphiti_core.utils.datetime_utils import utc_now
29
30
  from graphiti_core.utils.maintenance.temporal_operations import (
30
31
  extract_edge_dates,
31
32
  get_edge_contradictions,
@@ -132,7 +133,7 @@ async def extract_edges(
132
133
  group_id=group_id,
133
134
  fact=edge_data.get('fact', ''),
134
135
  episodes=[episode.uuid],
135
- created_at=datetime.now(timezone.utc),
136
+ created_at=utc_now(),
136
137
  valid_at=None,
137
138
  invalid_at=None,
138
139
  )
@@ -251,9 +252,7 @@ def resolve_edge_contradictions(
251
252
  and edge.valid_at < resolved_edge.valid_at
252
253
  ):
253
254
  edge.invalid_at = resolved_edge.valid_at
254
- edge.expired_at = (
255
- edge.expired_at if edge.expired_at is not None else datetime.now(timezone.utc)
256
- )
255
+ edge.expired_at = edge.expired_at if edge.expired_at is not None else utc_now()
257
256
  invalidated_edges.append(edge)
258
257
 
259
258
  return invalidated_edges
@@ -273,11 +272,12 @@ async def resolve_extracted_edge(
273
272
  get_edge_contradictions(llm_client, extracted_edge, existing_edges),
274
273
  )
275
274
 
276
- now = datetime.now(timezone.utc)
275
+ now = utc_now()
276
+
277
+ resolved_edge.valid_at = valid_at if valid_at else resolved_edge.valid_at
278
+ resolved_edge.invalid_at = invalid_at if invalid_at else resolved_edge.invalid_at
277
279
 
278
- resolved_edge.valid_at = valid_at if valid_at is not None else resolved_edge.valid_at
279
- resolved_edge.invalid_at = invalid_at if invalid_at is not None else resolved_edge.invalid_at
280
- if invalid_at is not None and resolved_edge.expired_at is None:
280
+ if invalid_at and not resolved_edge.expired_at:
281
281
  resolved_edge.expired_at = now
282
282
 
283
283
  # Determine if the new_edge needs to be expired
@@ -285,8 +285,12 @@ async def resolve_extracted_edge(
285
285
  invalidation_candidates.sort(key=lambda c: (c.valid_at is None, c.valid_at))
286
286
  for candidate in invalidation_candidates:
287
287
  if (
288
- candidate.valid_at is not None and resolved_edge.valid_at is not None
289
- ) and candidate.valid_at > resolved_edge.valid_at:
288
+ candidate.valid_at
289
+ and resolved_edge.valid_at
290
+ and candidate.valid_at.tzinfo
291
+ and resolved_edge.valid_at.tzinfo
292
+ and candidate.valid_at > resolved_edge.valid_at
293
+ ):
290
294
  # Expire new edge since we have information about more recent events
291
295
  resolved_edge.invalid_at = candidate.valid_at
292
296
  resolved_edge.expired_at = now
@@ -16,7 +16,6 @@ limitations under the License.
16
16
 
17
17
  import asyncio
18
18
  import logging
19
- from datetime import datetime, timezone
20
19
  from time import time
21
20
 
22
21
  from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS
@@ -26,6 +25,7 @@ from graphiti_core.prompts import prompt_library
26
25
  from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
27
26
  from graphiti_core.prompts.extract_nodes import ExtractedNodes, MissedEntities
28
27
  from graphiti_core.prompts.summarize_nodes import Summary
28
+ from graphiti_core.utils.datetime_utils import utc_now
29
29
 
30
30
  logger = logging.getLogger(__name__)
31
31
 
@@ -155,7 +155,7 @@ async def extract_nodes(
155
155
  group_id=episode.group_id,
156
156
  labels=['Entity'],
157
157
  summary='',
158
- created_at=datetime.now(timezone.utc),
158
+ created_at=utc_now(),
159
159
  )
160
160
  new_nodes.append(new_node)
161
161
  logger.debug(f'Created new node: {new_node.name} (UUID: {new_node.uuid})')
@@ -9,7 +9,7 @@ You may obtain a copy of the License at
9
9
 
10
10
  Unless required by applicable law or agreed to in writing, software
11
11
  distributed under the License is distributed on an "AS IS" BASIS,
12
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
13
  See the License for the specific language governing permissions and
14
14
  limitations under the License.
15
15
  """
@@ -24,6 +24,7 @@ from graphiti_core.nodes import EpisodicNode
24
24
  from graphiti_core.prompts import prompt_library
25
25
  from graphiti_core.prompts.extract_edge_dates import EdgeDates
26
26
  from graphiti_core.prompts.invalidate_edges import InvalidatedEdges
27
+ from graphiti_core.utils.datetime_utils import ensure_utc
27
28
 
28
29
  logger = logging.getLogger(__name__)
29
30
 
@@ -52,13 +53,15 @@ async def extract_edge_dates(
52
53
 
53
54
  if valid_at:
54
55
  try:
55
- valid_at_datetime = datetime.fromisoformat(valid_at.replace('Z', '+00:00'))
56
+ valid_at_datetime = ensure_utc(datetime.fromisoformat(valid_at.replace('Z', '+00:00')))
56
57
  except ValueError as e:
57
58
  logger.error(f'Error parsing valid_at date: {e}. Input: {valid_at}')
58
59
 
59
60
  if invalid_at:
60
61
  try:
61
- invalid_at_datetime = datetime.fromisoformat(invalid_at.replace('Z', '+00:00'))
62
+ invalid_at_datetime = ensure_utc(
63
+ datetime.fromisoformat(invalid_at.replace('Z', '+00:00'))
64
+ )
62
65
  except ValueError as e:
63
66
  logger.error(f'Error parsing invalid_at date: {e}. Input: {invalid_at}')
64
67
 
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "graphiti-core"
3
- version = "0.5.0pre2"
3
+ version = "0.5.0pre4"
4
4
  description = "A temporal graph building library"
5
5
  authors = [
6
6
  "Paul Paliychuk <paul@getzep.com>",
@@ -1,15 +0,0 @@
1
- from .maintenance import (
2
- build_episodic_edges,
3
- clear_data,
4
- extract_edges,
5
- extract_nodes,
6
- retrieve_episodes,
7
- )
8
-
9
- __all__ = [
10
- 'extract_edges',
11
- 'build_episodic_edges',
12
- 'extract_nodes',
13
- 'clear_data',
14
- 'retrieve_episodes',
15
- ]