graphiti-core 0.11.6rc7__py3-none-any.whl → 0.12.0rc1__py3-none-any.whl

graphiti_core/search/search.py

@@ -50,6 +50,9 @@ from graphiti_core.search.search_utils import (
     edge_similarity_search,
     episode_fulltext_search,
     episode_mentions_reranker,
+    get_embeddings_for_communities,
+    get_embeddings_for_edges,
+    get_embeddings_for_nodes,
     maximal_marginal_relevance,
     node_bfs_search,
     node_distance_reranker,
@@ -209,26 +212,17 @@ async def edge_search(
 
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == EdgeReranker.mmr:
-        await semaphore_gather(
-            *[edge.load_fact_embedding(driver) for result in search_results for edge in result]
+        search_result_uuids_and_vectors = await get_embeddings_for_edges(
+            driver, list(edge_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (edge.uuid, edge.fact_embedding if edge.fact_embedding is not None else [0.0] * 1024)
-            for result in search_results
-            for edge in result
-        ]
         reranked_uuids = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
+            reranker_min_score,
         )
     elif config.reranker == EdgeReranker.cross_encoder:
-        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
-
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
-        rrf_edges = [edge_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
-
-        fact_to_uuid_map = {edge.fact: edge.uuid for edge in rrf_edges}
+        fact_to_uuid_map = {edge.fact: edge.uuid for edge in list(edge_uuid_map.values())[:limit]}
         reranked_facts = await cross_encoder.rank(query, list(fact_to_uuid_map.keys()))
         reranked_uuids = [
             fact_to_uuid_map[fact] for fact, score in reranked_facts if score >= reranker_min_score
@@ -311,30 +305,23 @@ async def node_search(
     if config.reranker == NodeReranker.rrf:
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == NodeReranker.mmr:
-        await semaphore_gather(
-            *[node.load_name_embedding(driver) for result in search_results for node in result]
+        search_result_uuids_and_vectors = await get_embeddings_for_nodes(
+            driver, list(node_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (node.uuid, node.name_embedding if node.name_embedding is not None else [0.0] * 1024)
-            for result in search_results
-            for node in result
-        ]
+
         reranked_uuids = maximal_marginal_relevance(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
+            reranker_min_score,
         )
     elif config.reranker == NodeReranker.cross_encoder:
-        # use rrf as a preliminary reranker
-        rrf_result_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
-        rrf_results = [node_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
+        name_to_uuid_map = {node.name: node.uuid for node in list(node_uuid_map.values())}
 
-        summary_to_uuid_map = {node.summary: node.uuid for node in rrf_results}
-
-        reranked_summaries = await cross_encoder.rank(query, list(summary_to_uuid_map.keys()))
+        reranked_node_names = await cross_encoder.rank(query, list(name_to_uuid_map.keys()))
         reranked_uuids = [
-            summary_to_uuid_map[fact]
-            for fact, score in reranked_summaries
+            name_to_uuid_map[name]
+            for name, score in reranked_node_names
             if score >= reranker_min_score
         ]
     elif config.reranker == NodeReranker.episode_mentions:
@@ -437,25 +424,12 @@ async def community_search(
     if config.reranker == CommunityReranker.rrf:
         reranked_uuids = rrf(search_result_uuids, min_score=reranker_min_score)
     elif config.reranker == CommunityReranker.mmr:
-        await semaphore_gather(
-            *[
-                community.load_name_embedding(driver)
-                for result in search_results
-                for community in result
-            ]
+        search_result_uuids_and_vectors = await get_embeddings_for_communities(
+            driver, list(community_uuid_map.values())
         )
-        search_result_uuids_and_vectors = [
-            (
-                community.uuid,
-                community.name_embedding if community.name_embedding is not None else [0.0] * 1024,
-            )
-            for result in search_results
-            for community in result
-        ]
+
         reranked_uuids = maximal_marginal_relevance(
-            query_vector,
-            search_result_uuids_and_vectors,
-            config.mmr_lambda,
+            query_vector, search_result_uuids_and_vectors, config.mmr_lambda, reranker_min_score
         )
     elif config.reranker == CommunityReranker.cross_encoder:
         name_to_uuid_map = {node.name: node.uuid for result in search_results for node in result}

graphiti_core/search/search_utils.py

@@ -21,6 +21,7 @@ from typing import Any
 
 import numpy as np
 from neo4j import AsyncDriver, Query
+from numpy._typing import NDArray
 from typing_extensions import LiteralString
 
 from graphiti_core.edges import EntityEdge, get_entity_edge_from_record
@@ -173,7 +174,8 @@ async def edge_fulltext_search(
                      r.episodes AS episodes,
                      r.expired_at AS expired_at,
                      r.valid_at AS valid_at,
-                     r.invalid_at AS invalid_at
+                     r.invalid_at AS invalid_at,
+                     properties(r) AS attributes
                  ORDER BY score DESC LIMIT $limit
                  """
     )
@@ -242,7 +244,8 @@ async def edge_similarity_search(
                     r.episodes AS episodes,
                     r.expired_at AS expired_at,
                     r.valid_at AS valid_at,
-                    r.invalid_at AS invalid_at
+                    r.invalid_at AS invalid_at,
+                    properties(r) AS attributes
                 ORDER BY score DESC
                 LIMIT $limit
         """
@@ -300,7 +303,8 @@ async def edge_bfs_search(
                     r.episodes AS episodes,
                     r.expired_at AS expired_at,
                     r.valid_at AS valid_at,
-                    r.invalid_at AS invalid_at
+                    r.invalid_at AS invalid_at,
+                    properties(r) AS attributes
                 LIMIT $limit
         """
     )
@@ -336,10 +340,10 @@ async def node_fulltext_search(
 
     query = (
         """
-                                                CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
-                                                YIELD node AS n, score
-                                                WHERE n:Entity
-                                                """
+                                                                                        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
+                                                                                        YIELD node AS n, score
+                                                                                        WHERE n:Entity
+                                                                                        """
         + filter_query
         + ENTITY_NODE_RETURN
         + """
@@ -770,7 +774,8 @@ async def get_relevant_edges(
                     episodes: e.episodes,
                     expired_at: e.expired_at,
                     valid_at: e.valid_at,
-                    invalid_at: e.invalid_at
+                    invalid_at: e.invalid_at,
+                    attributes: properties(e)
                 })[..$limit] AS matches
         """
     )
@@ -836,7 +841,8 @@ async def get_edge_invalidation_candidates(
                     episodes: e.episodes,
                     expired_at: e.expired_at,
                     valid_at: e.valid_at,
-                    invalid_at: e.invalid_at
+                    invalid_at: e.invalid_at,
+                    attributes: properties(e)
                 })[..$limit] AS matches
         """
     )
@@ -899,6 +905,7 @@ async def node_distance_reranker(
         node_uuids=filtered_uuids,
         center_uuid=center_node_uuid,
         database_=DEFAULT_DATABASE,
+        routing_='r',
     )
 
     for result in path_results:
@@ -939,6 +946,7 @@ async def episode_mentions_reranker(
         query,
         node_uuids=sorted_uuids,
         database_=DEFAULT_DATABASE,
+        routing_='r',
     )
 
     for result in results:
@@ -952,15 +960,116 @@ async def episode_mentions_reranker(
 
 def maximal_marginal_relevance(
     query_vector: list[float],
-    candidates: list[tuple[str, list[float]]],
+    candidates: dict[str, list[float]],
     mmr_lambda: float = DEFAULT_MMR_LAMBDA,
-):
-    candidates_with_mmr: list[tuple[str, float]] = []
-    for candidate in candidates:
-        max_sim = max([np.dot(normalize_l2(candidate[1]), normalize_l2(c[1])) for c in candidates])
-        mmr = mmr_lambda * np.dot(candidate[1], query_vector) - (1 - mmr_lambda) * max_sim
-        candidates_with_mmr.append((candidate[0], mmr))
+    min_score: float = -2.0,
+) -> list[str]:
+    start = time()
+    query_array = np.array(query_vector)
+    candidate_arrays: dict[str, NDArray] = {}
+    for uuid, embedding in candidates.items():
+        candidate_arrays[uuid] = normalize_l2(embedding)
+
+    uuids: list[str] = list(candidate_arrays.keys())
+
+    similarity_matrix = np.zeros((len(uuids), len(uuids)))
+
+    for i, uuid_1 in enumerate(uuids):
+        for j, uuid_2 in enumerate(uuids[:i]):
+            u = candidate_arrays[uuid_1]
+            v = candidate_arrays[uuid_2]
+            similarity = np.dot(u, v)
+
+            similarity_matrix[i, j] = similarity
+            similarity_matrix[j, i] = similarity
+
+    mmr_scores: dict[str, float] = {}
+    for i, uuid in enumerate(uuids):
+        max_sim = np.max(similarity_matrix[i, :])
+        mmr = mmr_lambda * np.dot(query_array, candidate_arrays[uuid]) + (mmr_lambda - 1) * max_sim
+        mmr_scores[uuid] = mmr
+
+    uuids.sort(reverse=True, key=lambda c: mmr_scores[c])
+
+    end = time()
+    logger.debug(f'Completed MMR reranking in {(end - start) * 1000} ms')
+
+    return [uuid for uuid in uuids if mmr_scores[uuid] >= min_score]
+
+
+async def get_embeddings_for_nodes(
+    driver: AsyncDriver, nodes: list[EntityNode]
+) -> dict[str, list[float]]:
+    query: LiteralString = """MATCH (n:Entity)
+                              WHERE n.uuid IN $node_uuids
+                              RETURN DISTINCT
+                                n.uuid AS uuid,
+                                n.name_embedding AS name_embedding
+                    """
+
+    results, _, _ = await driver.execute_query(
+        query, node_uuids=[node.uuid for node in nodes], database_=DEFAULT_DATABASE, routing_='r'
+    )
+
+    embeddings_dict: dict[str, list[float]] = {}
+    for result in results:
+        uuid: str = result.get('uuid')
+        embedding: list[float] = result.get('name_embedding')
+        if uuid is not None and embedding is not None:
+            embeddings_dict[uuid] = embedding
 
-    candidates_with_mmr.sort(reverse=True, key=lambda c: c[1])
+    return embeddings_dict
+
+
+async def get_embeddings_for_communities(
+    driver: AsyncDriver, communities: list[CommunityNode]
+) -> dict[str, list[float]]:
+    query: LiteralString = """MATCH (c:Community)
+                              WHERE c.uuid IN $community_uuids
+                              RETURN DISTINCT
+                                c.uuid AS uuid,
+                                c.name_embedding AS name_embedding
+                    """
+
+    results, _, _ = await driver.execute_query(
+        query,
+        community_uuids=[community.uuid for community in communities],
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+
+    embeddings_dict: dict[str, list[float]] = {}
+    for result in results:
+        uuid: str = result.get('uuid')
+        embedding: list[float] = result.get('name_embedding')
+        if uuid is not None and embedding is not None:
+            embeddings_dict[uuid] = embedding
+
+    return embeddings_dict
+
+
+async def get_embeddings_for_edges(
+    driver: AsyncDriver, edges: list[EntityEdge]
+) -> dict[str, list[float]]:
+    query: LiteralString = """MATCH (n:Entity)-[e:RELATES_TO]-(m:Entity)
+                              WHERE e.uuid IN $edge_uuids
+                              RETURN DISTINCT
+                                e.uuid AS uuid,
+                                e.fact_embedding AS fact_embedding
+                    """
+
+    results, _, _ = await driver.execute_query(
+        query,
+        edge_uuids=[edge.uuid for edge in edges],
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+
+    embeddings_dict: dict[str, list[float]] = {}
+    for result in results:
+        uuid: str = result.get('uuid')
+        embedding: list[float] = result.get('fact_embedding')
+        if uuid is not None and embedding is not None:
+            embeddings_dict[uuid] = embedding
 
-    return list(set([candidate[0] for candidate in candidates_with_mmr]))
+    return embeddings_dict

graphiti_core/utils/bulk_utils.py

@@ -137,16 +137,34 @@ async def add_nodes_and_edges_bulk_tx(
         entity_data['labels'] = list(set(node.labels + ['Entity']))
         nodes.append(entity_data)
 
+    edges: list[dict[str, Any]] = []
     for edge in entity_edges:
         if edge.fact_embedding is None:
             await edge.generate_embedding(embedder)
+        edge_data: dict[str, Any] = {
+            'uuid': edge.uuid,
+            'source_node_uuid': edge.source_node_uuid,
+            'target_node_uuid': edge.target_node_uuid,
+            'name': edge.name,
+            'fact': edge.fact,
+            'fact_embedding': edge.fact_embedding,
+            'group_id': edge.group_id,
+            'episodes': edge.episodes,
+            'created_at': edge.created_at,
+            'expired_at': edge.expired_at,
+            'valid_at': edge.valid_at,
+            'invalid_at': edge.invalid_at,
+        }
+
+        edge_data.update(edge.attributes or {})
+        edges.append(edge_data)
 
     await tx.run(EPISODIC_NODE_SAVE_BULK, episodes=episodes)
     await tx.run(ENTITY_NODE_SAVE_BULK, nodes=nodes)
     await tx.run(
         EPISODIC_EDGE_SAVE_BULK, episodic_edges=[edge.model_dump() for edge in episodic_edges]
     )
-    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[edge.model_dump() for edge in entity_edges])
+    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=edges)
 
 
 async def extract_nodes_and_edges_bulk(

graphiti_core/utils/maintenance/edge_operations.py

@@ -18,6 +18,8 @@ import logging
 from datetime import datetime
 from time import time
 
+from pydantic import BaseModel
+
 from graphiti_core.edges import (
     CommunityEdge,
     EntityEdge,
@@ -35,9 +37,6 @@ from graphiti_core.prompts.extract_edges import ExtractedEdges, MissingFacts
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.search.search_utils import get_edge_invalidation_candidates, get_relevant_edges
 from graphiti_core.utils.datetime_utils import ensure_utc, utc_now
-from graphiti_core.utils.maintenance.temporal_operations import (
-    get_edge_contradictions,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -86,6 +85,7 @@ async def extract_edges(
     nodes: list[EntityNode],
     previous_episodes: list[EpisodicNode],
     group_id: str = '',
+    edge_types: dict[str, BaseModel] | None = None,
 ) -> list[EntityEdge]:
     start = time()
 
@@ -94,12 +94,25 @@ async def extract_edges(
 
     node_uuids_by_name_map = {node.name: node.uuid for node in nodes}
 
+    edge_types_context = (
+        [
+            {
+                'fact_type_name': type_name,
+                'fact_type_description': type_model.__doc__,
+            }
+            for type_name, type_model in edge_types.items()
+        ]
+        if edge_types is not None
+        else []
+    )
+
     # Prepare context for LLM
     context = {
         'episode_content': episode.content,
         'nodes': [node.name for node in nodes],
         'previous_episodes': [ep.content for ep in previous_episodes],
         'reference_time': episode.valid_at,
+        'edge_types': edge_types_context,
         'custom_prompt': '',
     }
 
@@ -236,6 +249,9 @@ async def resolve_extracted_edges(
     clients: GraphitiClients,
     extracted_edges: list[EntityEdge],
     episode: EpisodicNode,
+    entities: list[EntityNode],
+    edge_types: dict[str, BaseModel],
+    edge_type_map: dict[tuple[str, str], list[str]],
 ) -> tuple[list[EntityEdge], list[EntityEdge]]:
     driver = clients.driver
     llm_client = clients.llm_client
@@ -245,7 +261,7 @@ async def resolve_extracted_edges(
 
     search_results: tuple[list[list[EntityEdge]], list[list[EntityEdge]]] = await semaphore_gather(
         get_relevant_edges(driver, extracted_edges, SearchFilters()),
-        get_edge_invalidation_candidates(driver, extracted_edges, SearchFilters()),
+        get_edge_invalidation_candidates(driver, extracted_edges, SearchFilters(), 0.2),
     )
 
     related_edges_lists, edge_invalidation_candidates = search_results
@@ -254,15 +270,50 @@ async def resolve_extracted_edges(
         f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_lists for e in edges_lst]}'
     )
 
+    # Build entity hash table
+    uuid_entity_map: dict[str, EntityNode] = {entity.uuid: entity for entity in entities}
+
+    # Determine which edge types are relevant for each edge
+    edge_types_lst: list[dict[str, BaseModel]] = []
+    for extracted_edge in extracted_edges:
+        source_node_labels = uuid_entity_map[extracted_edge.source_node_uuid].labels
+        target_node_labels = uuid_entity_map[extracted_edge.target_node_uuid].labels
+        label_tuples = [
+            (source_label, target_label)
+            for source_label in source_node_labels
+            for target_label in target_node_labels
+        ]
+
+        extracted_edge_types = {}
+        for label_tuple in label_tuples:
+            type_names = edge_type_map.get(label_tuple, [])
+            for type_name in type_names:
+                type_model = edge_types.get(type_name)
+                if type_model is None:
+                    continue
+
+                extracted_edge_types[type_name] = type_model
+
+        edge_types_lst.append(extracted_edge_types)
+
     # resolve edges with related edges in the graph and find invalidation candidates
     results: list[tuple[EntityEdge, list[EntityEdge]]] = list(
         await semaphore_gather(
             *[
                 resolve_extracted_edge(
-                    llm_client, extracted_edge, related_edges, existing_edges, episode
+                    llm_client,
+                    extracted_edge,
+                    related_edges,
+                    existing_edges,
+                    episode,
+                    extracted_edge_types,
                 )
-                for extracted_edge, related_edges, existing_edges in zip(
-                    extracted_edges, related_edges_lists, edge_invalidation_candidates, strict=True
+                for extracted_edge, related_edges, existing_edges, extracted_edge_types in zip(
+                    extracted_edges,
+                    related_edges_lists,
+                    edge_invalidation_candidates,
+                    edge_types_lst,
+                    strict=True,
                 )
             ]
         )
@@ -326,10 +377,86 @@ async def resolve_extracted_edge(
     related_edges: list[EntityEdge],
     existing_edges: list[EntityEdge],
     episode: EpisodicNode,
+    edge_types: dict[str, BaseModel] | None = None,
 ) -> tuple[EntityEdge, list[EntityEdge]]:
-    resolved_edge, invalidation_candidates = await semaphore_gather(
-        dedupe_extracted_edge(llm_client, extracted_edge, related_edges, episode),
-        get_edge_contradictions(llm_client, extracted_edge, existing_edges),
+    if len(related_edges) == 0 and len(existing_edges) == 0:
+        return extracted_edge, []
+
+    start = time()
+
+    # Prepare context for LLM
+    related_edges_context = [
+        {'id': edge.uuid, 'fact': edge.fact} for i, edge in enumerate(related_edges)
+    ]
+
+    invalidation_edge_candidates_context = [
+        {'id': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges)
+    ]
+
+    edge_types_context = (
+        [
+            {
+                'fact_type_id': i,
+                'fact_type_name': type_name,
+                'fact_type_description': type_model.__doc__,
+            }
+            for i, (type_name, type_model) in enumerate(edge_types.items())
+        ]
+        if edge_types is not None
+        else []
+    )
+
+    context = {
+        'existing_edges': related_edges_context,
+        'new_edge': extracted_edge.fact,
+        'edge_invalidation_candidates': invalidation_edge_candidates_context,
+        'edge_types': edge_types_context,
+    }
+
+    llm_response = await llm_client.generate_response(
+        prompt_library.dedupe_edges.resolve_edge(context),
+        response_model=EdgeDuplicate,
+        model_size=ModelSize.small,
+    )
+
+    duplicate_fact_id: int = llm_response.get('duplicate_fact_id', -1)
+
+    resolved_edge = (
+        related_edges[duplicate_fact_id]
+        if 0 <= duplicate_fact_id < len(related_edges)
+        else extracted_edge
+    )
+
+    if duplicate_fact_id >= 0 and episode is not None:
+        resolved_edge.episodes.append(episode.uuid)
+
+    contradicted_facts: list[int] = llm_response.get('contradicted_facts', [])
+
+    invalidation_candidates: list[EntityEdge] = [existing_edges[i] for i in contradicted_facts]
+
+    fact_type: str = str(llm_response.get('fact_type'))
+    if fact_type.upper() != 'DEFAULT' and edge_types is not None:
+        resolved_edge.name = fact_type
+
+        edge_attributes_context = {
+            'message': episode.content,
+            'reference_time': episode.valid_at,
+            'fact': resolved_edge.fact,
+        }
+
+        edge_model = edge_types.get(fact_type)
+
+        edge_attributes_response = await llm_client.generate_response(
+            prompt_library.extract_edges.extract_attributes(edge_attributes_context),
+            response_model=edge_model,  # type: ignore
+            model_size=ModelSize.small,
+        )
+
+        resolved_edge.attributes = edge_attributes_response
+
+    end = time()
+    logger.debug(
+        f'Resolved Edge: {extracted_edge.name} is {resolved_edge.name}, in {(end - start) * 1000} ms'
     )
 
     now = utc_now()

graphiti_core/utils/maintenance/node_operations.py

@@ -29,7 +29,7 @@ from graphiti_core.llm_client import LLMClient
 from graphiti_core.llm_client.config import ModelSize
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
 from graphiti_core.prompts import prompt_library
-from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
+from graphiti_core.prompts.dedupe_nodes import NodeDuplicate, NodeResolutions
 from graphiti_core.prompts.extract_nodes import (
     ExtractedEntities,
     ExtractedEntity,
@@ -243,28 +243,65 @@ async def resolve_extracted_nodes(
 
     existing_nodes_lists: list[list[EntityNode]] = [result.nodes for result in search_results]
 
-    resolved_nodes: list[EntityNode] = await semaphore_gather(
-        *[
-            resolve_extracted_node(
-                llm_client,
-                extracted_node,
-                existing_nodes,
-                episode,
-                previous_episodes,
-                entity_types.get(
-                    next((item for item in extracted_node.labels if item != 'Entity'), '')
-                )
-                if entity_types is not None
-                else None,
-            )
-            for extracted_node, existing_nodes in zip(
-                extracted_nodes, existing_nodes_lists, strict=True
-            )
-        ]
+    entity_types_dict: dict[str, BaseModel] = entity_types if entity_types is not None else {}
+
+    # Prepare context for LLM
+    extracted_nodes_context = [
+        {
+            'id': i,
+            'name': node.name,
+            'entity_type': node.labels,
+            'entity_type_description': entity_types_dict.get(
+                next((item for item in node.labels if item != 'Entity'), '')
+            ).__doc__
+            or 'Default Entity Type',
+            'duplication_candidates': [
+                {
+                    **{
+                        'idx': j,
+                        'name': candidate.name,
+                        'entity_types': candidate.labels,
+                    },
+                    **candidate.attributes,
+                }
+                for j, candidate in enumerate(existing_nodes_lists[i])
+            ],
+        }
+        for i, node in enumerate(extracted_nodes)
+    ]
+
+    context = {
+        'extracted_nodes': extracted_nodes_context,
+        'episode_content': episode.content if episode is not None else '',
+        'previous_episodes': [ep.content for ep in previous_episodes]
+        if previous_episodes is not None
+        else [],
+    }
+
+    llm_response = await llm_client.generate_response(
+        prompt_library.dedupe_nodes.nodes(context),
+        response_model=NodeResolutions,
     )
 
+    node_resolutions: list = llm_response.get('entity_resolutions', [])
+
+    resolved_nodes: list[EntityNode] = []
     uuid_map: dict[str, str] = {}
-    for extracted_node, resolved_node in zip(extracted_nodes, resolved_nodes, strict=True):
+    for resolution in node_resolutions:
+        resolution_id = resolution.get('id', -1)
+        duplicate_idx = resolution.get('duplicate_idx', -1)
+
+        extracted_node = extracted_nodes[resolution_id]
+
+        resolved_node = (
+            existing_nodes_lists[resolution_id][duplicate_idx]
+            if 0 <= duplicate_idx < len(existing_nodes_lists[resolution_id])
+            else extracted_node
+        )
+
+        resolved_node.name = resolution.get('name')
+
+        resolved_nodes.append(resolved_node)
         uuid_map[extracted_node.uuid] = resolved_node.uuid
 
     logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
@@ -410,6 +447,7 @@ async def extract_attributes_from_node(
     llm_response = await llm_client.generate_response(
         prompt_library.extract_nodes.extract_attributes(summary_context),
         response_model=entity_attributes_model,
+        model_size=ModelSize.small,
     )
 
     node.summary = llm_response.get('summary', node.summary)

{graphiti_core-0.11.6rc7.dist-info → graphiti_core-0.12.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.11.6rc7
+Version: 0.12.0rc1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk