graphiti-core 0.21.0rc6__py3-none-any.whl → 0.30.0rc1__py3-none-any.whl

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

Potentially problematic release.


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

@@ -43,8 +43,14 @@ from graphiti_core.models.nodes.node_db_queries import (
43
43
  get_entity_node_save_bulk_query,
44
44
  get_episode_node_save_bulk_query,
45
45
  )
46
- from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
46
+ from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
47
47
  from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
48
+ from graphiti_core.utils.maintenance.dedup_helpers import (
49
+ DedupResolutionState,
50
+ _build_candidate_indexes,
51
+ _normalize_string_exact,
52
+ _resolve_with_similarity,
53
+ )
48
54
  from graphiti_core.utils.maintenance.edge_operations import (
49
55
  extract_edges,
50
56
  resolve_extracted_edge,
@@ -63,6 +69,38 @@ logger = logging.getLogger(__name__)
63
69
  CHUNK_SIZE = 10
64
70
 
65
71
 
72
+ def _build_directed_uuid_map(pairs: list[tuple[str, str]]) -> dict[str, str]:
73
+ """Collapse alias -> canonical chains while preserving direction.
74
+
75
+ The incoming pairs represent directed mappings discovered during node dedupe. We use a simple
76
+ union-find with iterative path compression to ensure every source UUID resolves to its ultimate
77
+ canonical target, even if aliases appear lexicographically smaller than the canonical UUID.
78
+ """
79
+
80
+ parent: dict[str, str] = {}
81
+
82
+ def find(uuid: str) -> str:
83
+ """Directed union-find lookup using iterative path compression."""
84
+ parent.setdefault(uuid, uuid)
85
+ root = uuid
86
+ while parent[root] != root:
87
+ root = parent[root]
88
+
89
+ while parent[uuid] != root:
90
+ next_uuid = parent[uuid]
91
+ parent[uuid] = root
92
+ uuid = next_uuid
93
+
94
+ return root
95
+
96
+ for source_uuid, target_uuid in pairs:
97
+ parent.setdefault(source_uuid, source_uuid)
98
+ parent.setdefault(target_uuid, target_uuid)
99
+ parent[find(source_uuid)] = find(target_uuid)
100
+
101
+ return {uuid: find(uuid) for uuid in parent}
102
+
103
+
66
104
  class RawEpisode(BaseModel):
67
105
  name: str
68
106
  uuid: str | None = Field(default=None)
@@ -266,83 +304,111 @@ async def dedupe_nodes_bulk(
266
304
  episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]],
267
305
  entity_types: dict[str, type[BaseModel]] | None = None,
268
306
  ) -> tuple[dict[str, list[EntityNode]], dict[str, str]]:
269
- embedder = clients.embedder
270
- min_score = 0.8
271
-
272
- # generate embeddings
273
- await semaphore_gather(
274
- *[create_entity_node_embeddings(embedder, nodes) for nodes in extracted_nodes]
275
- )
276
-
277
- # Find similar results
278
- dedupe_tuples: list[tuple[list[EntityNode], list[EntityNode]]] = []
279
- for i, nodes_i in enumerate(extracted_nodes):
280
- existing_nodes: list[EntityNode] = []
281
- for j, nodes_j in enumerate(extracted_nodes):
282
- if i == j:
283
- continue
284
- existing_nodes += nodes_j
285
-
286
- candidates_i: list[EntityNode] = []
287
- for node in nodes_i:
288
- for existing_node in existing_nodes:
289
- # Approximate BM25 by checking for word overlaps (this is faster than creating many in-memory indices)
290
- # This approach will cast a wider net than BM25, which is ideal for this use case
291
- node_words = set(node.name.lower().split())
292
- existing_node_words = set(existing_node.name.lower().split())
293
- has_overlap = not node_words.isdisjoint(existing_node_words)
294
- if has_overlap:
295
- candidates_i.append(existing_node)
296
- continue
307
+ """Resolve entity duplicates across an in-memory batch using a two-pass strategy.
297
308
 
298
- # Check for semantic similarity even if there is no overlap
299
- similarity = np.dot(
300
- normalize_l2(node.name_embedding or []),
301
- normalize_l2(existing_node.name_embedding or []),
302
- )
303
- if similarity >= min_score:
304
- candidates_i.append(existing_node)
305
-
306
- dedupe_tuples.append((nodes_i, candidates_i))
309
+ 1. Run :func:`resolve_extracted_nodes` for every episode in parallel so each batch item is
310
+ reconciled against the live graph just like the non-batch flow.
311
+ 2. Re-run the deterministic similarity heuristics across the union of resolved nodes to catch
312
+ duplicates that only co-occur inside this batch, emitting a canonical UUID map that callers
313
+ can apply to edges and persistence.
314
+ """
307
315
 
308
- # Determine Node Resolutions
309
- bulk_node_resolutions: list[
310
- tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]
311
- ] = await semaphore_gather(
316
+ first_pass_results = await semaphore_gather(
312
317
  *[
313
318
  resolve_extracted_nodes(
314
319
  clients,
315
- dedupe_tuple[0],
320
+ nodes,
316
321
  episode_tuples[i][0],
317
322
  episode_tuples[i][1],
318
323
  entity_types,
319
- existing_nodes_override=dedupe_tuples[i][1],
320
324
  )
321
- for i, dedupe_tuple in enumerate(dedupe_tuples)
325
+ for i, nodes in enumerate(extracted_nodes)
322
326
  ]
323
327
  )
324
328
 
325
- # Collect all duplicate pairs sorted by uuid
329
+ episode_resolutions: list[tuple[str, list[EntityNode]]] = []
330
+ per_episode_uuid_maps: list[dict[str, str]] = []
326
331
  duplicate_pairs: list[tuple[str, str]] = []
327
- for _, _, duplicates in bulk_node_resolutions:
328
- for duplicate in duplicates:
329
- n, m = duplicate
330
- duplicate_pairs.append((n.uuid, m.uuid))
331
332
 
332
- # Now we compress the duplicate_map, so that 3 -> 2 and 2 -> becomes 3 -> 1 (sorted by uuid)
333
- compressed_map: dict[str, str] = compress_uuid_map(duplicate_pairs)
333
+ for (resolved_nodes, uuid_map, duplicates), (episode, _) in zip(
334
+ first_pass_results, episode_tuples, strict=True
335
+ ):
336
+ episode_resolutions.append((episode.uuid, resolved_nodes))
337
+ per_episode_uuid_maps.append(uuid_map)
338
+ duplicate_pairs.extend((source.uuid, target.uuid) for source, target in duplicates)
339
+
340
+ canonical_nodes: dict[str, EntityNode] = {}
341
+ for _, resolved_nodes in episode_resolutions:
342
+ for node in resolved_nodes:
343
+ # NOTE: this loop is O(n^2) in the number of nodes inside the batch because we rebuild
344
+ # the MinHash index for the accumulated canonical pool each time. The LRU-backed
345
+ # shingle cache keeps the constant factors low for typical batch sizes (≤ CHUNK_SIZE),
346
+ # but if batches grow significantly we should switch to an incremental index or chunked
347
+ # processing.
348
+ if not canonical_nodes:
349
+ canonical_nodes[node.uuid] = node
350
+ continue
334
351
 
335
- node_uuid_map: dict[str, EntityNode] = {
336
- node.uuid: node for nodes in extracted_nodes for node in nodes
337
- }
352
+ existing_candidates = list(canonical_nodes.values())
353
+ normalized = _normalize_string_exact(node.name)
354
+ exact_match = next(
355
+ (
356
+ candidate
357
+ for candidate in existing_candidates
358
+ if _normalize_string_exact(candidate.name) == normalized
359
+ ),
360
+ None,
361
+ )
362
+ if exact_match is not None:
363
+ if exact_match.uuid != node.uuid:
364
+ duplicate_pairs.append((node.uuid, exact_match.uuid))
365
+ continue
366
+
367
+ indexes = _build_candidate_indexes(existing_candidates)
368
+ state = DedupResolutionState(
369
+ resolved_nodes=[None],
370
+ uuid_map={},
371
+ unresolved_indices=[],
372
+ )
373
+ _resolve_with_similarity([node], indexes, state)
374
+
375
+ resolved = state.resolved_nodes[0]
376
+ if resolved is None:
377
+ canonical_nodes[node.uuid] = node
378
+ continue
379
+
380
+ canonical_uuid = resolved.uuid
381
+ canonical_nodes.setdefault(canonical_uuid, resolved)
382
+ if canonical_uuid != node.uuid:
383
+ duplicate_pairs.append((node.uuid, canonical_uuid))
384
+
385
+ union_pairs: list[tuple[str, str]] = []
386
+ for uuid_map in per_episode_uuid_maps:
387
+ union_pairs.extend(uuid_map.items())
388
+ union_pairs.extend(duplicate_pairs)
389
+
390
+ compressed_map: dict[str, str] = _build_directed_uuid_map(union_pairs)
338
391
 
339
392
  nodes_by_episode: dict[str, list[EntityNode]] = {}
340
- for i, nodes in enumerate(extracted_nodes):
341
- episode = episode_tuples[i][0]
393
+ for episode_uuid, resolved_nodes in episode_resolutions:
394
+ deduped_nodes: list[EntityNode] = []
395
+ seen: set[str] = set()
396
+ for node in resolved_nodes:
397
+ canonical_uuid = compressed_map.get(node.uuid, node.uuid)
398
+ if canonical_uuid in seen:
399
+ continue
400
+ seen.add(canonical_uuid)
401
+ canonical_node = canonical_nodes.get(canonical_uuid)
402
+ if canonical_node is None:
403
+ logger.error(
404
+ 'Canonical node %s missing during batch dedupe; falling back to %s',
405
+ canonical_uuid,
406
+ node.uuid,
407
+ )
408
+ canonical_node = node
409
+ deduped_nodes.append(canonical_node)
342
410
 
343
- nodes_by_episode[episode.uuid] = [
344
- node_uuid_map[compressed_map.get(node.uuid, node.uuid)] for node in nodes
345
- ]
411
+ nodes_by_episode[episode_uuid] = deduped_nodes
346
412
 
347
413
  return nodes_by_episode, compressed_map
348
414
 
@@ -0,0 +1,262 @@
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 __future__ import annotations
18
+
19
+ import math
20
+ import re
21
+ from collections import defaultdict
22
+ from collections.abc import Iterable
23
+ from dataclasses import dataclass, field
24
+ from functools import lru_cache
25
+ from hashlib import blake2b
26
+ from typing import TYPE_CHECKING
27
+
28
+ if TYPE_CHECKING:
29
+ from graphiti_core.nodes import EntityNode
30
+
31
+ _NAME_ENTROPY_THRESHOLD = 1.5
32
+ _MIN_NAME_LENGTH = 6
33
+ _MIN_TOKEN_COUNT = 2
34
+ _FUZZY_JACCARD_THRESHOLD = 0.9
35
+ _MINHASH_PERMUTATIONS = 32
36
+ _MINHASH_BAND_SIZE = 4
37
+
38
+
39
+ def _normalize_string_exact(name: str) -> str:
40
+ """Lowercase text and collapse whitespace so equal names map to the same key."""
41
+ normalized = re.sub(r'[\s]+', ' ', name.lower())
42
+ return normalized.strip()
43
+
44
+
45
+ def _normalize_name_for_fuzzy(name: str) -> str:
46
+ """Produce a fuzzier form that keeps alphanumerics and apostrophes for n-gram shingles."""
47
+ normalized = re.sub(r"[^a-z0-9' ]", ' ', _normalize_string_exact(name))
48
+ normalized = normalized.strip()
49
+ return re.sub(r'[\s]+', ' ', normalized)
50
+
51
+
52
+ def _name_entropy(normalized_name: str) -> float:
53
+ """Approximate text specificity using Shannon entropy over characters.
54
+
55
+ We strip spaces, count how often each character appears, and sum
56
+ probability * -log2(probability). Short or repetitive names yield low
57
+ entropy, which signals we should defer resolution to the LLM instead of
58
+ trusting fuzzy similarity.
59
+ """
60
+ if not normalized_name:
61
+ return 0.0
62
+
63
+ counts: dict[str, int] = {}
64
+ for char in normalized_name.replace(' ', ''):
65
+ counts[char] = counts.get(char, 0) + 1
66
+
67
+ total = sum(counts.values())
68
+ if total == 0:
69
+ return 0.0
70
+
71
+ entropy = 0.0
72
+ for count in counts.values():
73
+ probability = count / total
74
+ entropy -= probability * math.log2(probability)
75
+
76
+ return entropy
77
+
78
+
79
+ def _has_high_entropy(normalized_name: str) -> bool:
80
+ """Filter out very short or low-entropy names that are unreliable for fuzzy matching."""
81
+ token_count = len(normalized_name.split())
82
+ if len(normalized_name) < _MIN_NAME_LENGTH and token_count < _MIN_TOKEN_COUNT:
83
+ return False
84
+
85
+ return _name_entropy(normalized_name) >= _NAME_ENTROPY_THRESHOLD
86
+
87
+
88
+ def _shingles(normalized_name: str) -> set[str]:
89
+ """Create 3-gram shingles from the normalized name for MinHash calculations."""
90
+ cleaned = normalized_name.replace(' ', '')
91
+ if len(cleaned) < 2:
92
+ return {cleaned} if cleaned else set()
93
+
94
+ return {cleaned[i : i + 3] for i in range(len(cleaned) - 2)}
95
+
96
+
97
+ def _hash_shingle(shingle: str, seed: int) -> int:
98
+ """Generate a deterministic 64-bit hash for a shingle given the permutation seed."""
99
+ digest = blake2b(f'{seed}:{shingle}'.encode(), digest_size=8)
100
+ return int.from_bytes(digest.digest(), 'big')
101
+
102
+
103
+ def _minhash_signature(shingles: Iterable[str]) -> tuple[int, ...]:
104
+ """Compute the MinHash signature for the shingle set across predefined permutations."""
105
+ if not shingles:
106
+ return tuple()
107
+
108
+ seeds = range(_MINHASH_PERMUTATIONS)
109
+ signature: list[int] = []
110
+ for seed in seeds:
111
+ min_hash = min(_hash_shingle(shingle, seed) for shingle in shingles)
112
+ signature.append(min_hash)
113
+
114
+ return tuple(signature)
115
+
116
+
117
+ def _lsh_bands(signature: Iterable[int]) -> list[tuple[int, ...]]:
118
+ """Split the MinHash signature into fixed-size bands for locality-sensitive hashing."""
119
+ signature_list = list(signature)
120
+ if not signature_list:
121
+ return []
122
+
123
+ bands: list[tuple[int, ...]] = []
124
+ for start in range(0, len(signature_list), _MINHASH_BAND_SIZE):
125
+ band = tuple(signature_list[start : start + _MINHASH_BAND_SIZE])
126
+ if len(band) == _MINHASH_BAND_SIZE:
127
+ bands.append(band)
128
+ return bands
129
+
130
+
131
+ def _jaccard_similarity(a: set[str], b: set[str]) -> float:
132
+ """Return the Jaccard similarity between two shingle sets, handling empty edge cases."""
133
+ if not a and not b:
134
+ return 1.0
135
+ if not a or not b:
136
+ return 0.0
137
+
138
+ intersection = len(a.intersection(b))
139
+ union = len(a.union(b))
140
+ return intersection / union if union else 0.0
141
+
142
+
143
+ @lru_cache(maxsize=512)
144
+ def _cached_shingles(name: str) -> set[str]:
145
+ """Cache shingle sets per normalized name to avoid recomputation within a worker."""
146
+ return _shingles(name)
147
+
148
+
149
+ @dataclass
150
+ class DedupCandidateIndexes:
151
+ """Precomputed lookup structures that drive entity deduplication heuristics."""
152
+
153
+ existing_nodes: list[EntityNode]
154
+ nodes_by_uuid: dict[str, EntityNode]
155
+ normalized_existing: defaultdict[str, list[EntityNode]]
156
+ shingles_by_candidate: dict[str, set[str]]
157
+ lsh_buckets: defaultdict[tuple[int, tuple[int, ...]], list[str]]
158
+
159
+
160
+ @dataclass
161
+ class DedupResolutionState:
162
+ """Mutable resolution bookkeeping shared across deterministic and LLM passes."""
163
+
164
+ resolved_nodes: list[EntityNode | None]
165
+ uuid_map: dict[str, str]
166
+ unresolved_indices: list[int]
167
+ duplicate_pairs: list[tuple[EntityNode, EntityNode]] = field(default_factory=list)
168
+
169
+
170
+ def _build_candidate_indexes(existing_nodes: list[EntityNode]) -> DedupCandidateIndexes:
171
+ """Precompute exact and fuzzy lookup structures once per dedupe run."""
172
+ normalized_existing: defaultdict[str, list[EntityNode]] = defaultdict(list)
173
+ nodes_by_uuid: dict[str, EntityNode] = {}
174
+ shingles_by_candidate: dict[str, set[str]] = {}
175
+ lsh_buckets: defaultdict[tuple[int, tuple[int, ...]], list[str]] = defaultdict(list)
176
+
177
+ for candidate in existing_nodes:
178
+ normalized = _normalize_string_exact(candidate.name)
179
+ normalized_existing[normalized].append(candidate)
180
+ nodes_by_uuid[candidate.uuid] = candidate
181
+
182
+ shingles = _cached_shingles(_normalize_name_for_fuzzy(candidate.name))
183
+ shingles_by_candidate[candidate.uuid] = shingles
184
+
185
+ signature = _minhash_signature(shingles)
186
+ for band_index, band in enumerate(_lsh_bands(signature)):
187
+ lsh_buckets[(band_index, band)].append(candidate.uuid)
188
+
189
+ return DedupCandidateIndexes(
190
+ existing_nodes=existing_nodes,
191
+ nodes_by_uuid=nodes_by_uuid,
192
+ normalized_existing=normalized_existing,
193
+ shingles_by_candidate=shingles_by_candidate,
194
+ lsh_buckets=lsh_buckets,
195
+ )
196
+
197
+
198
+ def _resolve_with_similarity(
199
+ extracted_nodes: list[EntityNode],
200
+ indexes: DedupCandidateIndexes,
201
+ state: DedupResolutionState,
202
+ ) -> None:
203
+ """Attempt deterministic resolution using exact name hits and fuzzy MinHash comparisons."""
204
+ for idx, node in enumerate(extracted_nodes):
205
+ normalized_exact = _normalize_string_exact(node.name)
206
+ normalized_fuzzy = _normalize_name_for_fuzzy(node.name)
207
+
208
+ if not _has_high_entropy(normalized_fuzzy):
209
+ state.unresolved_indices.append(idx)
210
+ continue
211
+
212
+ existing_matches = indexes.normalized_existing.get(normalized_exact, [])
213
+ if len(existing_matches) == 1:
214
+ match = existing_matches[0]
215
+ state.resolved_nodes[idx] = match
216
+ state.uuid_map[node.uuid] = match.uuid
217
+ if match.uuid != node.uuid:
218
+ state.duplicate_pairs.append((node, match))
219
+ continue
220
+ if len(existing_matches) > 1:
221
+ state.unresolved_indices.append(idx)
222
+ continue
223
+
224
+ shingles = _cached_shingles(normalized_fuzzy)
225
+ signature = _minhash_signature(shingles)
226
+ candidate_ids: set[str] = set()
227
+ for band_index, band in enumerate(_lsh_bands(signature)):
228
+ candidate_ids.update(indexes.lsh_buckets.get((band_index, band), []))
229
+
230
+ best_candidate: EntityNode | None = None
231
+ best_score = 0.0
232
+ for candidate_id in candidate_ids:
233
+ candidate_shingles = indexes.shingles_by_candidate.get(candidate_id, set())
234
+ score = _jaccard_similarity(shingles, candidate_shingles)
235
+ if score > best_score:
236
+ best_score = score
237
+ best_candidate = indexes.nodes_by_uuid.get(candidate_id)
238
+
239
+ if best_candidate is not None and best_score >= _FUZZY_JACCARD_THRESHOLD:
240
+ state.resolved_nodes[idx] = best_candidate
241
+ state.uuid_map[node.uuid] = best_candidate.uuid
242
+ if best_candidate.uuid != node.uuid:
243
+ state.duplicate_pairs.append((node, best_candidate))
244
+ continue
245
+
246
+ state.unresolved_indices.append(idx)
247
+
248
+
249
+ __all__ = [
250
+ 'DedupCandidateIndexes',
251
+ 'DedupResolutionState',
252
+ '_normalize_string_exact',
253
+ '_normalize_name_for_fuzzy',
254
+ '_has_high_entropy',
255
+ '_minhash_signature',
256
+ '_lsh_bands',
257
+ '_jaccard_similarity',
258
+ '_cached_shingles',
259
+ '_FUZZY_JACCARD_THRESHOLD',
260
+ '_build_candidate_indexes',
261
+ '_resolve_with_similarity',
262
+ ]
@@ -41,6 +41,7 @@ from graphiti_core.search.search_config import SearchResults
41
41
  from graphiti_core.search.search_config_recipes import EDGE_HYBRID_SEARCH_RRF
42
42
  from graphiti_core.search.search_filters import SearchFilters
43
43
  from graphiti_core.utils.datetime_utils import ensure_utc, utc_now
44
+ from graphiti_core.utils.maintenance.dedup_helpers import _normalize_string_exact
44
45
 
45
46
  logger = logging.getLogger(__name__)
46
47
 
@@ -397,6 +398,19 @@ async def resolve_extracted_edge(
397
398
  if len(related_edges) == 0 and len(existing_edges) == 0:
398
399
  return extracted_edge, [], []
399
400
 
401
+ # Fast path: if the fact text and endpoints already exist verbatim, reuse the matching edge.
402
+ normalized_fact = _normalize_string_exact(extracted_edge.fact)
403
+ for edge in related_edges:
404
+ if (
405
+ edge.source_node_uuid == extracted_edge.source_node_uuid
406
+ and edge.target_node_uuid == extracted_edge.target_node_uuid
407
+ and _normalize_string_exact(edge.fact) == normalized_fact
408
+ ):
409
+ resolved = edge
410
+ if episode is not None and episode.uuid not in resolved.episodes:
411
+ resolved.episodes.append(episode.uuid)
412
+ return resolved, [], []
413
+
400
414
  start = time()
401
415
 
402
416
  # Prepare context for LLM
@@ -24,7 +24,12 @@ from graphiti_core.graphiti_types import GraphitiClients
24
24
  from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
25
25
  from graphiti_core.llm_client import LLMClient
26
26
  from graphiti_core.llm_client.config import ModelSize
27
- from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
27
+ from graphiti_core.nodes import (
28
+ EntityNode,
29
+ EpisodeType,
30
+ EpisodicNode,
31
+ create_entity_node_embeddings,
32
+ )
28
33
  from graphiti_core.prompts import prompt_library
29
34
  from graphiti_core.prompts.dedupe_nodes import NodeDuplicate, NodeResolutions
30
35
  from graphiti_core.prompts.extract_nodes import (
@@ -38,7 +43,15 @@ from graphiti_core.search.search_config import SearchResults
38
43
  from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF
39
44
  from graphiti_core.search.search_filters import SearchFilters
40
45
  from graphiti_core.utils.datetime_utils import utc_now
41
- from graphiti_core.utils.maintenance.edge_operations import filter_existing_duplicate_of_edges
46
+ from graphiti_core.utils.maintenance.dedup_helpers import (
47
+ DedupCandidateIndexes,
48
+ DedupResolutionState,
49
+ _build_candidate_indexes,
50
+ _resolve_with_similarity,
51
+ )
52
+ from graphiti_core.utils.maintenance.edge_operations import (
53
+ filter_existing_duplicate_of_edges,
54
+ )
42
55
 
43
56
  logger = logging.getLogger(__name__)
44
57
 
@@ -119,11 +132,13 @@ async def extract_nodes(
119
132
  )
120
133
  elif episode.source == EpisodeType.text:
121
134
  llm_response = await llm_client.generate_response(
122
- prompt_library.extract_nodes.extract_text(context), response_model=ExtractedEntities
135
+ prompt_library.extract_nodes.extract_text(context),
136
+ response_model=ExtractedEntities,
123
137
  )
124
138
  elif episode.source == EpisodeType.json:
125
139
  llm_response = await llm_client.generate_response(
126
- prompt_library.extract_nodes.extract_json(context), response_model=ExtractedEntities
140
+ prompt_library.extract_nodes.extract_json(context),
141
+ response_model=ExtractedEntities,
127
142
  )
128
143
 
129
144
  response_object = ExtractedEntities(**llm_response)
@@ -181,17 +196,12 @@ async def extract_nodes(
181
196
  return extracted_nodes
182
197
 
183
198
 
184
- async def resolve_extracted_nodes(
199
+ async def _collect_candidate_nodes(
185
200
  clients: GraphitiClients,
186
201
  extracted_nodes: list[EntityNode],
187
- episode: EpisodicNode | None = None,
188
- previous_episodes: list[EpisodicNode] | None = None,
189
- entity_types: dict[str, type[BaseModel]] | None = None,
190
- existing_nodes_override: list[EntityNode] | None = None,
191
- ) -> tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]:
192
- llm_client = clients.llm_client
193
- driver = clients.driver
194
-
202
+ existing_nodes_override: list[EntityNode] | None,
203
+ ) -> list[EntityNode]:
204
+ """Search per extracted name and return unique candidates with overrides honored in order."""
195
205
  search_results: list[SearchResults] = await semaphore_gather(
196
206
  *[
197
207
  search(
@@ -205,33 +215,40 @@ async def resolve_extracted_nodes(
205
215
  ]
206
216
  )
207
217
 
208
- candidate_nodes: list[EntityNode] = (
209
- [node for result in search_results for node in result.nodes]
210
- if existing_nodes_override is None
211
- else existing_nodes_override
212
- )
218
+ candidate_nodes: list[EntityNode] = [node for result in search_results for node in result.nodes]
213
219
 
214
- existing_nodes_dict: dict[str, EntityNode] = {node.uuid: node for node in candidate_nodes}
220
+ if existing_nodes_override is not None:
221
+ candidate_nodes.extend(existing_nodes_override)
215
222
 
216
- existing_nodes: list[EntityNode] = list(existing_nodes_dict.values())
223
+ seen_candidate_uuids: set[str] = set()
224
+ ordered_candidates: list[EntityNode] = []
225
+ for candidate in candidate_nodes:
226
+ if candidate.uuid in seen_candidate_uuids:
227
+ continue
228
+ seen_candidate_uuids.add(candidate.uuid)
229
+ ordered_candidates.append(candidate)
217
230
 
218
- existing_nodes_context = (
219
- [
220
- {
221
- **{
222
- 'idx': i,
223
- 'name': candidate.name,
224
- 'entity_types': candidate.labels,
225
- },
226
- **candidate.attributes,
227
- }
228
- for i, candidate in enumerate(existing_nodes)
229
- ],
230
- )
231
+ return ordered_candidates
232
+
233
+
234
+ async def _resolve_with_llm(
235
+ llm_client: LLMClient,
236
+ extracted_nodes: list[EntityNode],
237
+ indexes: DedupCandidateIndexes,
238
+ state: DedupResolutionState,
239
+ ensure_ascii: bool,
240
+ episode: EpisodicNode | None,
241
+ previous_episodes: list[EpisodicNode] | None,
242
+ entity_types: dict[str, type[BaseModel]] | None,
243
+ ) -> None:
244
+ """Escalate unresolved nodes to the dedupe prompt so the LLM can select or reject duplicates."""
245
+ if not state.unresolved_indices:
246
+ return
231
247
 
232
248
  entity_types_dict: dict[str, type[BaseModel]] = entity_types if entity_types is not None else {}
233
249
 
234
- # Prepare context for LLM
250
+ llm_extracted_nodes = [extracted_nodes[i] for i in state.unresolved_indices]
251
+
235
252
  extracted_nodes_context = [
236
253
  {
237
254
  'id': i,
@@ -242,17 +259,29 @@ async def resolve_extracted_nodes(
242
259
  ).__doc__
243
260
  or 'Default Entity Type',
244
261
  }
245
- for i, node in enumerate(extracted_nodes)
262
+ for i, node in enumerate(llm_extracted_nodes)
263
+ ]
264
+
265
+ existing_nodes_context = [
266
+ {
267
+ **{
268
+ 'idx': i,
269
+ 'name': candidate.name,
270
+ 'entity_types': candidate.labels,
271
+ },
272
+ **candidate.attributes,
273
+ }
274
+ for i, candidate in enumerate(indexes.existing_nodes)
246
275
  ]
247
276
 
248
277
  context = {
249
278
  'extracted_nodes': extracted_nodes_context,
250
279
  'existing_nodes': existing_nodes_context,
251
280
  'episode_content': episode.content if episode is not None else '',
252
- 'previous_episodes': [ep.content for ep in previous_episodes]
253
- if previous_episodes is not None
254
- else [],
255
- 'ensure_ascii': clients.ensure_ascii,
281
+ 'previous_episodes': (
282
+ [ep.content for ep in previous_episodes] if previous_episodes is not None else []
283
+ ),
284
+ 'ensure_ascii': ensure_ascii,
256
285
  }
257
286
 
258
287
  llm_response = await llm_client.generate_response(
@@ -262,33 +291,82 @@ async def resolve_extracted_nodes(
262
291
 
263
292
  node_resolutions: list[NodeDuplicate] = NodeResolutions(**llm_response).entity_resolutions
264
293
 
265
- resolved_nodes: list[EntityNode] = []
266
- uuid_map: dict[str, str] = {}
267
- node_duplicates: list[tuple[EntityNode, EntityNode]] = []
268
294
  for resolution in node_resolutions:
269
- resolution_id: int = resolution.id
295
+ relative_id: int = resolution.id
270
296
  duplicate_idx: int = resolution.duplicate_idx
271
297
 
272
- extracted_node = extracted_nodes[resolution_id]
298
+ original_index = state.unresolved_indices[relative_id]
299
+ extracted_node = extracted_nodes[original_index]
273
300
 
274
301
  resolved_node = (
275
- existing_nodes[duplicate_idx]
276
- if 0 <= duplicate_idx < len(existing_nodes)
302
+ indexes.existing_nodes[duplicate_idx]
303
+ if 0 <= duplicate_idx < len(indexes.existing_nodes)
277
304
  else extracted_node
278
305
  )
279
306
 
280
- # resolved_node.name = resolution.get('name')
307
+ state.resolved_nodes[original_index] = resolved_node
308
+ state.uuid_map[extracted_node.uuid] = resolved_node.uuid
309
+ if resolved_node.uuid != extracted_node.uuid:
310
+ state.duplicate_pairs.append((extracted_node, resolved_node))
311
+
312
+
313
+ async def resolve_extracted_nodes(
314
+ clients: GraphitiClients,
315
+ extracted_nodes: list[EntityNode],
316
+ episode: EpisodicNode | None = None,
317
+ previous_episodes: list[EpisodicNode] | None = None,
318
+ entity_types: dict[str, type[BaseModel]] | None = None,
319
+ existing_nodes_override: list[EntityNode] | None = None,
320
+ ) -> tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]:
321
+ """Search for existing nodes, resolve deterministic matches, then escalate holdouts to the LLM dedupe prompt."""
322
+ llm_client = clients.llm_client
323
+ driver = clients.driver
324
+ existing_nodes = await _collect_candidate_nodes(
325
+ clients,
326
+ extracted_nodes,
327
+ existing_nodes_override,
328
+ )
329
+
330
+ indexes: DedupCandidateIndexes = _build_candidate_indexes(existing_nodes)
281
331
 
282
- resolved_nodes.append(resolved_node)
283
- uuid_map[extracted_node.uuid] = resolved_node.uuid
332
+ state = DedupResolutionState(
333
+ resolved_nodes=[None] * len(extracted_nodes),
334
+ uuid_map={},
335
+ unresolved_indices=[],
336
+ )
337
+
338
+ _resolve_with_similarity(extracted_nodes, indexes, state)
339
+
340
+ await _resolve_with_llm(
341
+ llm_client,
342
+ extracted_nodes,
343
+ indexes,
344
+ state,
345
+ clients.ensure_ascii,
346
+ episode,
347
+ previous_episodes,
348
+ entity_types,
349
+ )
284
350
 
285
- logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
351
+ for idx, node in enumerate(extracted_nodes):
352
+ if state.resolved_nodes[idx] is None:
353
+ state.resolved_nodes[idx] = node
354
+ state.uuid_map[node.uuid] = node.uuid
355
+
356
+ logger.debug(
357
+ 'Resolved nodes: %s',
358
+ [(node.name, node.uuid) for node in state.resolved_nodes if node is not None],
359
+ )
286
360
 
287
361
  new_node_duplicates: list[
288
362
  tuple[EntityNode, EntityNode]
289
- ] = await filter_existing_duplicate_of_edges(driver, node_duplicates)
363
+ ] = await filter_existing_duplicate_of_edges(driver, state.duplicate_pairs)
290
364
 
291
- return resolved_nodes, uuid_map, new_node_duplicates
365
+ return (
366
+ [node for node in state.resolved_nodes if node is not None],
367
+ state.uuid_map,
368
+ new_node_duplicates,
369
+ )
292
370
 
293
371
 
294
372
  async def extract_attributes_from_nodes(
@@ -307,9 +385,11 @@ async def extract_attributes_from_nodes(
307
385
  node,
308
386
  episode,
309
387
  previous_episodes,
310
- entity_types.get(next((item for item in node.labels if item != 'Entity'), ''))
311
- if entity_types is not None
312
- else None,
388
+ (
389
+ entity_types.get(next((item for item in node.labels if item != 'Entity'), ''))
390
+ if entity_types is not None
391
+ else None
392
+ ),
313
393
  clients.ensure_ascii,
314
394
  )
315
395
  for node in nodes
@@ -339,18 +419,18 @@ async def extract_attributes_from_node(
339
419
  attributes_context: dict[str, Any] = {
340
420
  'node': node_context,
341
421
  'episode_content': episode.content if episode is not None else '',
342
- 'previous_episodes': [ep.content for ep in previous_episodes]
343
- if previous_episodes is not None
344
- else [],
422
+ 'previous_episodes': (
423
+ [ep.content for ep in previous_episodes] if previous_episodes is not None else []
424
+ ),
345
425
  'ensure_ascii': ensure_ascii,
346
426
  }
347
427
 
348
428
  summary_context: dict[str, Any] = {
349
429
  'node': node_context,
350
430
  'episode_content': episode.content if episode is not None else '',
351
- 'previous_episodes': [ep.content for ep in previous_episodes]
352
- if previous_episodes is not None
353
- else [],
431
+ 'previous_episodes': (
432
+ [ep.content for ep in previous_episodes] if previous_episodes is not None else []
433
+ ),
354
434
  'ensure_ascii': ensure_ascii,
355
435
  }
356
436
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: graphiti-core
3
- Version: 0.21.0rc6
3
+ Version: 0.30.0rc1
4
4
  Summary: A temporal graph building library
5
5
  Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
6
6
  Project-URL: Repository, https://github.com/getzep/graphiti
@@ -64,17 +64,18 @@ graphiti_core/search/search_utils.py,sha256=ak1aBeKNuxS7szydNHwva2ABWSRlQ0S_v8ZO
64
64
  graphiti_core/telemetry/__init__.py,sha256=5kALLDlU9bb2v19CdN7qVANsJWyfnL9E60J6FFgzm3o,226
65
65
  graphiti_core/telemetry/telemetry.py,sha256=47LrzOVBCcZxsYPsnSxWFiztHoxYKKxPwyRX0hnbDGc,3230
66
66
  graphiti_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
67
- graphiti_core/utils/bulk_utils.py,sha256=9XWXqjxiu2ydKMLKQRTbvzO6cO1o1HRjjpmaf5Ym51k,17633
67
+ graphiti_core/utils/bulk_utils.py,sha256=yeURrhy6whJQDYiFfwna7bPDy0HF04WNfRZRSpPv4E8,20255
68
68
  graphiti_core/utils/datetime_utils.py,sha256=J-zYSq7-H-2n9hYOXNIun12kM10vNX9mMATGR_egTmY,1806
69
69
  graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4wPx3Vt6aAPaXh78,301
70
70
  graphiti_core/utils/maintenance/community_operations.py,sha256=XMiokEemn96GlvjkOvbo9hIX04Fea3eVj408NHG5P4o,11042
71
- graphiti_core/utils/maintenance/edge_operations.py,sha256=rfFsqigWXNcUGKu1l1-RdSoFdEeioK78oo4VWOagqgs,20576
71
+ graphiti_core/utils/maintenance/dedup_helpers.py,sha256=B7k6KkB6Sii8PZCWNNTvsNiy4BNTNWpoLeGgrPLq6BE,9220
72
+ graphiti_core/utils/maintenance/edge_operations.py,sha256=fvWKJWzz4_d2Y8bOfZFjJpLnGmsFwnrutFW25LX-S08,21287
72
73
  graphiti_core/utils/maintenance/graph_data_operations.py,sha256=42icj3S_ELAJ-NK3jVS_rg_243dmnaZOyUitJj_uJ-M,6085
73
- graphiti_core/utils/maintenance/node_operations.py,sha256=xSeRK8cdit9GT9ZKNGpawg01Wcu1W3FsBH9moFH_mao,13423
74
+ graphiti_core/utils/maintenance/node_operations.py,sha256=msMev_W36uG5zXzlkFd00u4xq2SFGm6N5q3arGvaDJs,15788
74
75
  graphiti_core/utils/maintenance/temporal_operations.py,sha256=IIaVtShpVkOYe6haxz3a1x3v54-MzaEXG8VsxFUNeoY,3582
75
76
  graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
76
77
  graphiti_core/utils/ontology_utils/entity_types_utils.py,sha256=4eVgxLWY6Q8k9cRJ5pW59IYF--U4nXZsZIGOVb_yHfQ,1285
77
- graphiti_core-0.21.0rc6.dist-info/METADATA,sha256=55pEDr5ujjIjMUKQTqbJ_YRvCSAuUi41ejVjYXsQiSA,26933
78
- graphiti_core-0.21.0rc6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
79
- graphiti_core-0.21.0rc6.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
80
- graphiti_core-0.21.0rc6.dist-info/RECORD,,
78
+ graphiti_core-0.30.0rc1.dist-info/METADATA,sha256=mUraIWnv6SfxFTkCAOaYZQr9g4TpnH539mXewc36MrI,26933
79
+ graphiti_core-0.30.0rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
80
+ graphiti_core-0.30.0rc1.dist-info/licenses/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
81
+ graphiti_core-0.30.0rc1.dist-info/RECORD,,