graphiti-core 0.12.3__tar.gz → 0.13.0__tar.gz

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/PKG-INFO

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

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/errors.py

@@ -73,3 +73,11 @@ class EntityTypeValidationError(GraphitiError):
     def __init__(self, entity_type: str, entity_type_attribute: str):
         self.message = f'{entity_type_attribute} cannot be used as an attribute for {entity_type} as it is a protected attribute name.'
         super().__init__(self.message)
+
+
+class GroupIdValidationError(GraphitiError):
+    """Raised when a group_id contains invalid characters."""
+
+    def __init__(self, group_id: str):
+        self.message = f'group_id "{group_id}" must contain only alphanumeric characters, dashes, or underscores'
+        super().__init__(self.message)

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/graphiti.py

@@ -29,7 +29,7 @@ from graphiti_core.driver.neo4j_driver import Neo4jDriver
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.graphiti_types import GraphitiClients
-from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather, validate_group_id
 from graphiti_core.llm_client import LLMClient, OpenAIClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
@@ -103,6 +103,7 @@ class Graphiti:
         cross_encoder: CrossEncoderClient | None = None,
         store_raw_episode_content: bool = True,
         graph_driver: GraphDriver | None = None,
+        max_coroutines: int | None = None,
     ):
         """
         Initialize a Graphiti instance.
@@ -121,6 +122,20 @@ class Graphiti:
         llm_client : LLMClient | None, optional
             An instance of LLMClient for natural language processing tasks.
             If not provided, a default OpenAIClient will be initialized.
+        embedder : EmbedderClient | None, optional
+            An instance of EmbedderClient for embedding tasks.
+            If not provided, a default OpenAIEmbedder will be initialized.
+        cross_encoder : CrossEncoderClient | None, optional
+            An instance of CrossEncoderClient for reranking tasks.
+            If not provided, a default OpenAIRerankerClient will be initialized.
+        store_raw_episode_content : bool, optional
+            Whether to store the raw content of episodes. Defaults to True.
+        graph_driver : GraphDriver | None, optional
+            An instance of GraphDriver for database operations.
+            If not provided, a default Neo4jDriver will be initialized.
+        max_coroutines : int | None, optional
+            The maximum number of concurrent operations allowed. Overrides SEMAPHORE_LIMIT set in the environment.
+            If not set, the Graphiti default is used.
 
         Returns
         -------
@@ -145,6 +160,7 @@ class Graphiti:
 
         self.database = DEFAULT_DATABASE
         self.store_raw_episode_content = store_raw_episode_content
+        self.max_coroutines = max_coroutines
         if llm_client:
             self.llm_client = llm_client
         else:
@@ -335,6 +351,7 @@ class Graphiti:
             now = utc_now()
 
             validate_entity_types(entity_types)
+            validate_group_id(group_id)
 
             previous_episodes = (
                 await self.retrieve_episodes(
@@ -393,6 +410,7 @@ class Graphiti:
                     group_id,
                     edge_types,
                 ),
+                max_coroutines=self.max_coroutines,
             )
 
             edges = resolve_edge_pointers(extracted_edges, uuid_map)
@@ -409,6 +427,7 @@ class Graphiti:
                 extract_attributes_from_nodes(
                     self.clients, nodes, episode, previous_episodes, entity_types
                 ),
+                max_coroutines=self.max_coroutines,
             )
 
             duplicate_of_edges = build_duplicate_of_edges(episode, now, node_duplicates)
@@ -432,7 +451,8 @@ class Graphiti:
                     *[
                         update_community(self.driver, self.llm_client, self.embedder, node)
                         for node in nodes
-                    ]
+                    ],
+                    max_coroutines=self.max_coroutines,
                 )
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
@@ -484,6 +504,8 @@ class Graphiti:
             start = time()
             now = utc_now()
 
+            validate_group_id(group_id)
+
             episodes = [
                 EpisodicNode(
                     name=episode.name,
@@ -499,7 +521,10 @@ class Graphiti:
             ]
 
             # Save all the episodes
-            await semaphore_gather(*[episode.save(self.driver) for episode in episodes])
+            await semaphore_gather(
+                *[episode.save(self.driver) for episode in episodes],
+                max_coroutines=self.max_coroutines,
+            )
 
             # Get previous episode context for each episode
             episode_pairs = await retrieve_previous_episodes_bulk(self.driver, episodes)
@@ -515,16 +540,21 @@ class Graphiti:
             await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
                 *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
+                max_coroutines=self.max_coroutines,
             )
 
             # Dedupe extracted nodes, compress extracted edges
             (nodes, uuid_map), extracted_edges_timestamped = await semaphore_gather(
                 dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes),
                 extract_edge_dates_bulk(self.llm_client, extracted_edges, episode_pairs),
+                max_coroutines=self.max_coroutines,
             )
 
             # save nodes to KG
-            await semaphore_gather(*[node.save(self.driver) for node in nodes])
+            await semaphore_gather(
+                *[node.save(self.driver) for node in nodes],
+                max_coroutines=self.max_coroutines,
+            )
 
             # re-map edge pointers so that they don't point to discard dupe nodes
             extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
@@ -536,7 +566,8 @@ class Graphiti:
 
             # save episodic edges to KG
             await semaphore_gather(
-                *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers]
+                *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers],
+                max_coroutines=self.max_coroutines,
             )
 
             # Dedupe extracted edges
@@ -548,7 +579,10 @@ class Graphiti:
             # invalidate edges
 
             # save edges to KG
-            await semaphore_gather(*[edge.save(self.driver) for edge in edges])
+            await semaphore_gather(
+                *[edge.save(self.driver) for edge in edges],
+                max_coroutines=self.max_coroutines,
+            )
 
             end = time()
             logger.info(f'Completed add_episode_bulk in {(end - start) * 1000} ms')
@@ -572,11 +606,18 @@ class Graphiti:
         )
 
         await semaphore_gather(
-            *[node.generate_name_embedding(self.embedder) for node in community_nodes]
+            *[node.generate_name_embedding(self.embedder) for node in community_nodes],
+            max_coroutines=self.max_coroutines,
         )
 
-        await semaphore_gather(*[node.save(self.driver) for node in community_nodes])
-        await semaphore_gather(*[edge.save(self.driver) for edge in community_edges])
+        await semaphore_gather(
+            *[node.save(self.driver) for node in community_nodes],
+            max_coroutines=self.max_coroutines,
+        )
+        await semaphore_gather(
+            *[edge.save(self.driver) for edge in community_edges],
+            max_coroutines=self.max_coroutines,
+        )
 
         return community_nodes
 
@@ -683,7 +724,8 @@ class Graphiti:
         episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
         edges_list = await semaphore_gather(
-            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
+            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes],
+            max_coroutines=self.max_coroutines,
         )
 
         edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
@@ -759,6 +801,12 @@ class Graphiti:
                 if record['episode_count'] == 1:
                     nodes_to_delete.append(node)
 
-        await semaphore_gather(*[node.delete(self.driver) for node in nodes_to_delete])
-        await semaphore_gather(*[edge.delete(self.driver) for edge in edges_to_delete])
+        await semaphore_gather(
+            *[node.delete(self.driver) for node in nodes_to_delete],
+            max_coroutines=self.max_coroutines,
+        )
+        await semaphore_gather(
+            *[edge.delete(self.driver) for edge in edges_to_delete],
+            max_coroutines=self.max_coroutines,
+        )
         await episode.delete(self.driver)

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/helpers.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 import asyncio
 import os
+import re
 from collections.abc import Coroutine
 from datetime import datetime
 
@@ -25,6 +26,8 @@ from neo4j import time as neo4j_time
 from numpy._typing import NDArray
 from typing_extensions import LiteralString
 
+from graphiti_core.errors import GroupIdValidationError
+
 load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', 'neo4j')
@@ -94,12 +97,38 @@ def normalize_l2(embedding: list[float]) -> NDArray:
 # Use this instead of asyncio.gather() to bound coroutines
 async def semaphore_gather(
     *coroutines: Coroutine,
-    max_coroutines: int = SEMAPHORE_LIMIT,
+    max_coroutines: int | None = None,
 ):
-    semaphore = asyncio.Semaphore(max_coroutines)
+    semaphore = asyncio.Semaphore(max_coroutines or SEMAPHORE_LIMIT)
 
     async def _wrap_coroutine(coroutine):
         async with semaphore:
             return await coroutine
 
     return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines))
+
+
+def validate_group_id(group_id: str) -> bool:
+    """
+    Validate that a group_id contains only ASCII alphanumeric characters, dashes, and underscores.
+
+    Args:
+        group_id: The group_id to validate
+
+    Returns:
+        True if valid, False otherwise
+
+    Raises:
+        GroupIdValidationError: If group_id contains invalid characters
+    """
+
+    # Allow empty string (default case)
+    if not group_id:
+        return True
+
+    # Check if string contains only ASCII alphanumeric characters, dashes, or underscores
+    # Pattern matches: letters (a-z, A-Z), digits (0-9), hyphens (-), and underscores (_)
+    if not re.match(r'^[a-zA-Z0-9_-]+$', group_id):
+        raise GroupIdValidationError(group_id)
+
+    return True

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/models/edges/edge_db_queries.py

@@ -35,7 +35,7 @@ ENTITY_EDGE_SAVE = """
         MATCH (target:Entity {uuid: $target_uuid}) 
         MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target)
         SET r = $edge_data
-        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
+        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $edge_data.fact_embedding)
         RETURN r.uuid AS uuid"""
 
 ENTITY_EDGE_SAVE_BULK = """

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/prompts/dedupe_nodes.py

@@ -30,7 +30,7 @@ class NodeDuplicate(BaseModel):
     )
     name: str = Field(
         ...,
-        description='Name of the entity. Should be the most complete and descriptive name possible.',
+        description='Name of the entity. Should be the most complete and descriptive name of the entity. Do not include any JSON formatting in the Entity name such as {}.',
     )
     additional_duplicates: list[int] = Field(
         ...,
@@ -84,19 +84,19 @@ def node(context: dict[str, Any]) -> list[Message]:
         is a duplicate entity of one of the EXISTING ENTITIES.
         
         Entities should only be considered duplicates if they refer to the *same real-world object or concept*.
+        Semantic Equivalence: if a descriptive label in existing_entities clearly refers to a named entity in context, treat them as duplicates.
 
         Do NOT mark entities as duplicates if:
         - They are related but distinct.
         - They have similar names or purposes but refer to separate instances or concepts.
 
-        Task:
-        If the NEW ENTITY represents a duplicate entity of any entity in EXISTING ENTITIES, set duplicate_entity_id to the
-        id of the EXISTING ENTITY that is the duplicate. 
-        
-        If the NEW ENTITY is not a duplicate of any of the EXISTING ENTITIES,
-        duplicate_entity_id should be set to -1.
+         TASK:
+         1. Compare `new_entity` against each item in `existing_entities`.
+         2. If it refers to the same real‐world object or concept, collect its index.
+         3. Let `duplicate_idx` = the *first* collected index, or –1 if none.
+         4. Let `additional_duplicates` = the list of *any other* collected indices (empty list if none).
         
-        Also return the name that best describes the NEW ENTITY (whether it is the name of the NEW ENTITY, a node it
+        Also return the full name of the NEW ENTITY (whether it is the name of the NEW ENTITY, a node it
         is a duplicate of, or a combination of the two).
         """,
         ),

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/prompts/extract_edges.py

@@ -97,8 +97,7 @@ Only extract facts that:
 - The FACT TYPES provide a list of the most important types of facts, make sure to extract facts of these types
 - The FACT TYPES are not an exhaustive list, extract all facts from the message even if they do not fit into one
     of the FACT TYPES
-- The FACT TYPES each contain their fact_type_signature which represents the entity types which that fact_type is defined for.
-    A Type of Entity in the signature represents any extracted entity (it is a generic universal type for all entities).
+- The FACT TYPES each contain their fact_type_signature which represents the source and target entity types.
 
 You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
 

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/graphiti_core/utils/maintenance/node_operations.py

@@ -310,7 +310,7 @@ async def resolve_extracted_nodes(
             else extracted_node
         )
 
-        resolved_node.name = resolution.get('name')
+        # resolved_node.name = resolution.get('name')
 
         resolved_nodes.append(resolved_node)
         uuid_map[extracted_node.uuid] = resolved_node.uuid

{graphiti_core-0.12.3 → graphiti_core-0.13.0}/pyproject.toml

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