graphiti-core 0.10.2__tar.gz → 0.10.4__tar.gz

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/PKG-INFO

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

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/graphiti.py

@@ -797,7 +797,7 @@ class Graphiti:
         # We should only delete edges created by the episode
         edges_to_delete: list[EntityEdge] = []
         for edge in edges:
-            if edge.episodes[0] == episode.uuid:
+            if edge.episodes and edge.episodes[0] == episode.uuid:
                 edges_to_delete.append(edge)
 
         # Find nodes mentioned by the episode

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/nodes.py

@@ -38,6 +38,20 @@ from graphiti_core.utils.datetime_utils import utc_now
 
 logger = logging.getLogger(__name__)
 
+ENTITY_NODE_RETURN: LiteralString = """
+        OPTIONAL MATCH (e:Episodic)-[r:MENTIONS]->(n)
+        WITH n, collect(e.uuid) AS episodes
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id,
+            n.created_at AS created_at, 
+            n.summary AS summary,
+            labels(n) AS labels,
+            properties(n) AS attributes,
+            episodes"""
+
 
 class EpisodeType(Enum):
     """
@@ -280,6 +294,9 @@ class EpisodicNode(Node):
 class EntityNode(Node):
     name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
     summary: str = Field(description='regional summary of surrounding edges', default_factory=str)
+    episodes: list[str] | None = Field(
+        default=None, description='List of episode uuids that mention this node.'
+    )
     attributes: dict[str, Any] = Field(
         default={}, description='Additional attributes of the node. Dependent on node labels'
     )
@@ -318,19 +335,14 @@ class EntityNode(Node):
 
     @classmethod
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
-        records, _, _ = await driver.execute_query(
+        query = (
             """
-        MATCH (n:Entity {uuid: $uuid})
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.name_embedding AS name_embedding,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes
-        """,
+                MATCH (n:Entity {uuid: $uuid})
+                """
+            + ENTITY_NODE_RETURN
+        )
+        records, _, _ = await driver.execute_query(
+            query,
             uuid=uuid,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -348,16 +360,8 @@ class EntityNode(Node):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity) WHERE n.uuid IN $uuids
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.name_embedding AS name_embedding,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes
-        """,
+        """
+            + ENTITY_NODE_RETURN,
             uuids=uuids,
             database_=DEFAULT_DATABASE,
             routing_='r',
@@ -383,16 +387,8 @@ class EntityNode(Node):
         MATCH (n:Entity) WHERE n.group_id IN $group_ids
         """
             + cursor_query
+            + ENTITY_NODE_RETURN
             + """
-        RETURN
-            n.uuid As uuid, 
-            n.name AS name,
-            n.name_embedding AS name_embedding,
-            n.group_id AS group_id,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes
         ORDER BY n.uuid DESC
         """
             + limit_query,
@@ -548,6 +544,7 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         created_at=record['created_at'].to_native(),
         summary=record['summary'],
         attributes=record['attributes'],
+        episodes=record['episodes'],
     )
 
     entity_node.attributes.pop('uuid', None)

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/prompts/extract_edge_dates.py

@@ -82,6 +82,7 @@ def v1(context: dict[str, Any]) -> list[Message]:
             7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
             8. If only year is mentioned, use January 1st of that year at 00:00:00.
             9. Always include the time zone offset (use Z for UTC if no specific time zone is mentioned).
+            10. A fact discussing that something is no longer true should have a valid_at according to when the negated fact became true.
             """,
         ),
     ]

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/search/search.py

@@ -210,7 +210,6 @@ async def edge_search(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
-            min_score=reranker_min_score,
         )
     elif config.reranker == EdgeReranker.cross_encoder:
         search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
@@ -310,7 +309,6 @@ async def node_search(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
-            min_score=reranker_min_score,
         )
     elif config.reranker == NodeReranker.cross_encoder:
         # use rrf as a preliminary reranker
@@ -437,7 +435,6 @@ async def community_search(
             query_vector,
             search_result_uuids_and_vectors,
             config.mmr_lambda,
-            min_score=reranker_min_score,
         )
     elif config.reranker == CommunityReranker.cross_encoder:
         summary_to_uuid_map = {

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/search/search_utils.py

@@ -32,6 +32,7 @@ from graphiti_core.helpers import (
     semaphore_gather,
 )
 from graphiti_core.nodes import (
+    ENTITY_NODE_RETURN,
     CommunityNode,
     EntityNode,
     EpisodicNode,
@@ -53,6 +54,20 @@ DEFAULT_MMR_LAMBDA = 0.5
 MAX_SEARCH_DEPTH = 3
 MAX_QUERY_LENGTH = 32
 
+SEARCH_ENTITY_NODE_RETURN: LiteralString = """
+        OPTIONAL MATCH (e:Episodic)-[r:MENTIONS]->(n)
+        WITH n, score, collect(e.uuid) AS episodes
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id,
+            n.created_at AS created_at, 
+            n.summary AS summary,
+            labels(n) AS labels,
+            properties(n) AS attributes,
+            episodes"""
+
 
 def fulltext_query(query: str, group_ids: list[str] | None = None):
     group_ids_filter_list = (
@@ -230,8 +245,8 @@ async def edge_similarity_search(
 
     query: LiteralString = (
         """
-                                                                                                                    MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
-                                                                                                                    """
+                                                                                                                                            MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                                                                                                                                            """
         + group_filter_query
         + filter_query
         + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
@@ -341,27 +356,21 @@ async def node_fulltext_search(
 
     filter_query, filter_params = node_search_filter_query_constructor(search_filter)
 
-    records, _, _ = await driver.execute_query(
-        """
-        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
-        YIELD node AS node, score
-        MATCH (n:Entity)
-        WHERE n.uuid = node.uuid
+    query = (
         """
+                CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
+                YIELD node AS n, score
+                WHERE n:Entity
+                """
         + filter_query
+        + SEARCH_ENTITY_NODE_RETURN
         + """
-        RETURN
-            n.uuid AS uuid,
-            n.group_id AS group_id, 
-            n.name AS name, 
-            n.name_embedding AS name_embedding,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes
         ORDER BY score DESC
-        LIMIT $limit
-        """,
+        """
+    )
+
+    records, _, _ = await driver.execute_query(
+        query,
         filter_params,
         query=fuzzy_query,
         group_ids=group_ids,
@@ -406,19 +415,12 @@ async def node_similarity_search(
         + filter_query
         + """
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
-            WHERE score > $min_score
-            RETURN
-                n.uuid As uuid,
-                n.group_id AS group_id,
-                n.name AS name, 
-                n.name_embedding AS name_embedding,
-                n.created_at AS created_at, 
-                n.summary AS summary,
-                labels(n) AS labels,
-                properties(n) AS attributes
-            ORDER BY score DESC
-            LIMIT $limit
-            """,
+            WHERE score > $min_score"""
+        + SEARCH_ENTITY_NODE_RETURN
+        + """
+        ORDER BY score DESC
+        LIMIT $limit
+        """,
         query_params,
         search_vector=search_vector,
         group_ids=group_ids,
@@ -452,16 +454,8 @@ async def node_bfs_search(
             WHERE n.group_id = origin.group_id
             """
         + filter_query
+        + ENTITY_NODE_RETURN
         + """
-        RETURN DISTINCT
-            n.uuid As uuid,
-            n.group_id AS group_id,
-            n.name AS name, 
-            n.name_embedding AS name_embedding,
-            n.created_at AS created_at, 
-            n.summary AS summary,
-            labels(n) AS labels,
-            properties(n) AS attributes
         LIMIT $limit
         """,
         filter_params,
@@ -852,7 +846,6 @@ def maximal_marginal_relevance(
     query_vector: list[float],
     candidates: list[tuple[str, list[float]]],
     mmr_lambda: float = DEFAULT_MMR_LAMBDA,
-    min_score: float = 0,
 ):
     candidates_with_mmr: list[tuple[str, float]] = []
     for candidate in candidates:
@@ -862,6 +855,4 @@ def maximal_marginal_relevance(
 
     candidates_with_mmr.sort(reverse=True, key=lambda c: c[1])
 
-    return list(
-        set([candidate[0] for candidate in candidates_with_mmr if candidate[1] >= min_score])
-    )
+    return list(set([candidate[0] for candidate in candidates_with_mmr]))

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/utils/bulk_utils.py

@@ -26,7 +26,7 @@ from pydantic import BaseModel
 from typing_extensions import Any
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
-from graphiti_core.helpers import semaphore_gather
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.models.edges.edge_db_queries import (
     ENTITY_EDGE_SAVE_BULK,
@@ -95,7 +95,7 @@ async def add_nodes_and_edges_bulk(
     entity_nodes: list[EntityNode],
     entity_edges: list[EntityEdge],
 ):
-    async with driver.session() as session:
+    async with driver.session(database=DEFAULT_DATABASE) as session:
         await session.execute_write(
             add_nodes_and_edges_bulk_tx, episodic_nodes, episodic_edges, entity_nodes, entity_edges
         )

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/utils/maintenance/community_operations.py

@@ -101,9 +101,11 @@ def label_propagation(projection: dict[str, list[Neighbor]]) -> list[list[str]]:
             ]
 
             community_lst.sort(reverse=True)
-            community_candidate = community_lst[0][1] if len(community_lst) > 0 else -1
-
-            new_community = max(community_candidate, curr_community)
+            candidate_rank, community_candidate = community_lst[0] if community_lst else (0, -1)
+            if community_candidate != -1 and candidate_rank > 1:
+                new_community = community_candidate
+            else:
+                new_community = max(community_candidate, curr_community)
 
             new_community_map[uuid] = new_community
 

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/utils/maintenance/graph_data_operations.py

@@ -42,7 +42,7 @@ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bo
                 driver.execute_query(
                     """DROP INDEX $name""",
                     name=name,
-                    _database=DEFAULT_DATABASE,
+                    database_=DEFAULT_DATABASE,
                 )
                 for name in index_names
             ]
@@ -87,7 +87,7 @@ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bo
         *[
             driver.execute_query(
                 query,
-                _database=DEFAULT_DATABASE,
+                database_=DEFAULT_DATABASE,
             )
             for query in index_queries
         ]
@@ -95,7 +95,7 @@ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bo
 
 
 async def clear_data(driver: AsyncDriver, group_ids: list[str] | None = None):
-    async with driver.session() as session:
+    async with driver.session(database=DEFAULT_DATABASE) as session:
 
         async def delete_all(tx):
             await tx.run('MATCH (n) DETACH DELETE n')
@@ -150,7 +150,7 @@ async def retrieve_episodes(
         reference_time=reference_time,
         num_episodes=last_n,
         group_ids=group_ids,
-        _database=DEFAULT_DATABASE,
+        database_=DEFAULT_DATABASE,
     )
     episodes = [
         EpisodicNode(

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/graphiti_core/utils/maintenance/node_operations.py

@@ -306,7 +306,7 @@ async def resolve_extracted_node(
 
     # Prepare context for LLM
     existing_nodes_context = [
-        {'uuid': node.uuid, 'name': node.name, 'attributes': node.attributes}
+        {**{'uuid': node.uuid, 'name': node.name, 'summary': node.summary}, **node.attributes}
         for node in existing_nodes
     ]
 
@@ -408,6 +408,7 @@ async def resolve_extracted_node(
                 if new_attributes.get(attribute_name) is None:
                     new_attributes[attribute_name] = attribute_value
             node.attributes = new_attributes
+            node.labels = list(set(existing_node.labels + extracted_node.labels))
 
             uuid_map[extracted_node.uuid] = existing_node.uuid
 
@@ -432,7 +433,8 @@ async def dedupe_node_list(
 
     # Prepare context for LLM
     nodes_context = [
-        {'uuid': node.uuid, 'name': node.name, 'summary': node.summary} for node in nodes
+        {'uuid': node.uuid, 'name': node.name, 'summary': node.summary}.update(node.attributes)
+        for node in nodes
     ]
 
     context = {

{graphiti_core-0.10.2 → graphiti_core-0.10.4}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.10.2"
+version = "0.10.4"
 authors = [
     { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" },
     { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" },