graphiti-core 0.3.19__py3-none-any.whl → 0.3.21__py3-none-any.whl

graphiti_core/graphiti.py

@@ -51,6 +51,7 @@ from graphiti_core.utils import (
 )
 from graphiti_core.utils.bulk_utils import (
     RawEpisode,
+    add_nodes_and_edges_bulk,
     dedupe_edges_bulk,
     dedupe_nodes_bulk,
     extract_edge_dates_bulk,
@@ -451,10 +452,9 @@ class Graphiti:
             if not self.store_raw_episode_content:
                 episode.content = ''
 
-            await episode.save(self.driver)
-            await asyncio.gather(*[node.save(self.driver) for node in nodes])
-            await asyncio.gather(*[edge.save(self.driver) for edge in episodic_edges])
-            await asyncio.gather(*[edge.save(self.driver) for edge in entity_edges])
+            await add_nodes_and_edges_bulk(
+                self.driver, [episode], episodic_edges, nodes, entity_edges
+            )
 
             # Update any communities
             if update_communities:

graphiti_core/helpers.py

@@ -24,6 +24,7 @@ from neo4j import time as neo4j_time
 load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
+USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 
 
 def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:

graphiti_core/models/edges/edge_db_queries.py

@@ -5,6 +5,15 @@ EPISODIC_EDGE_SAVE = """
         SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
         RETURN r.uuid AS uuid"""
 
+EPISODIC_EDGE_SAVE_BULK = """
+    UNWIND $episodic_edges AS edge
+    MATCH (episode:Episodic {uuid: edge.source_node_uuid}) 
+    MATCH (node:Entity {uuid: edge.target_node_uuid}) 
+    MERGE (episode)-[r:MENTIONS {uuid: edge.uuid}]->(node)
+    SET r = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
+    RETURN r.uuid AS uuid
+"""
+
 ENTITY_EDGE_SAVE = """
         MATCH (source:Entity {uuid: $source_uuid}) 
         MATCH (target:Entity {uuid: $target_uuid}) 
@@ -14,6 +23,17 @@ ENTITY_EDGE_SAVE = """
         WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
         RETURN r.uuid AS uuid"""
 
+ENTITY_EDGE_SAVE_BULK = """
+    UNWIND $entity_edges AS edge
+    MATCH (source:Entity {uuid: edge.source_node_uuid}) 
+    MATCH (target:Entity {uuid: edge.target_node_uuid}) 
+    MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
+    SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes, 
+    created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at}
+    WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
+    RETURN r.uuid AS uuid
+"""
+
 COMMUNITY_EDGE_SAVE = """
         MATCH (community:Community {uuid: $community_uuid}) 
         MATCH (node:Entity | Community {uuid: $entity_uuid}) 

graphiti_core/models/nodes/node_db_queries.py

@@ -4,12 +4,29 @@ EPISODIC_NODE_SAVE = """
         entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
         RETURN n.uuid AS uuid"""
 
+EPISODIC_NODE_SAVE_BULK = """
+    UNWIND $episodes AS episode
+    MERGE (n:Episodic {uuid: episode.uuid})
+    SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, 
+        source: episode.source, content: episode.content, 
+    entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
+    RETURN n.uuid AS uuid
+"""
+
 ENTITY_NODE_SAVE = """
         MERGE (n:Entity {uuid: $uuid})
         SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
         WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid"""
 
+ENTITY_NODE_SAVE_BULK = """
+    UNWIND $nodes AS node
+    MERGE (n:Entity {uuid: node.uuid})
+    SET n = {uuid: node.uuid, name: node.name, group_id: node.group_id, summary: node.summary, created_at: node.created_at}
+    WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
+    RETURN n.uuid AS uuid
+"""
+
 COMMUNITY_NODE_SAVE = """
         MERGE (n:Community {uuid: $uuid})
         SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}

graphiti_core/search/search.py

@@ -66,6 +66,12 @@ async def search(
     bfs_origin_node_uuids: list[str] | None = None,
 ) -> SearchResults:
     start = time()
+    if query.strip() == '':
+        return SearchResults(
+            edges=[],
+            nodes=[],
+            communities=[],
+        )
     query_vector = await embedder.create(input_data=[query.replace('\n', ' ')])
 
     # if group_ids is empty, set it to None

graphiti_core/search/search_utils.py

@@ -21,9 +21,15 @@ from time import time
 
 import numpy as np
 from neo4j import AsyncDriver, Query
+from typing_extensions import LiteralString
 
 from graphiti_core.edges import EntityEdge, get_entity_edge_from_record
-from graphiti_core.helpers import DEFAULT_DATABASE, lucene_sanitize, normalize_l2
+from graphiti_core.helpers import (
+    DEFAULT_DATABASE,
+    USE_PARALLEL_RUNTIME,
+    lucene_sanitize,
+    normalize_l2,
+)
 from graphiti_core.nodes import (
     CommunityNode,
     EntityNode,
@@ -38,7 +44,7 @@ RELEVANT_SCHEMA_LIMIT = 3
 DEFAULT_MIN_SCORE = 0.6
 DEFAULT_MMR_LAMBDA = 0.5
 MAX_SEARCH_DEPTH = 3
-MAX_QUERY_LENGTH = 128
+MAX_QUERY_LENGTH = 32
 
 
 def fulltext_query(query: str, group_ids: list[str] | None = None):
@@ -187,8 +193,11 @@ async def edge_similarity_search(
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityEdge]:
     # vector similarity search over embedded facts
-    query = Query("""
-                CYPHER runtime = parallel parallelRuntimeSupport=all
+    runtime_query: LiteralString = (
+        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
+    )
+
+    query: LiteralString = """
                 MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
                 WHERE ($group_ids IS NULL OR r.group_id IN $group_ids)
                 AND ($source_uuid IS NULL OR n.uuid IN [$source_uuid, $target_uuid])
@@ -210,10 +219,10 @@ async def edge_similarity_search(
                     r.invalid_at AS invalid_at
                 ORDER BY score DESC
                 LIMIT $limit
-        """)
+        """
 
     records, _, _ = await driver.execute_query(
-        query,
+        runtime_query + query,
         search_vector=search_vector,
         source_uuid=source_node_uuid,
         target_uuid=target_node_uuid,
@@ -318,9 +327,13 @@ async def node_similarity_search(
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityNode]:
     # vector similarity search over entity names
+    runtime_query: LiteralString = (
+        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
+    )
+
     records, _, _ = await driver.execute_query(
-        """
-            CYPHER runtime = parallel parallelRuntimeSupport=all
+        runtime_query
+        + """
             MATCH (n:Entity)
             WHERE $group_ids IS NULL OR n.group_id IN $group_ids
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
@@ -425,23 +438,27 @@ async def community_similarity_search(
     min_score=DEFAULT_MIN_SCORE,
 ) -> list[CommunityNode]:
     # vector similarity search over entity names
+    runtime_query: LiteralString = (
+        'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
+    )
+
     records, _, _ = await driver.execute_query(
-        """
-            CYPHER runtime = parallel parallelRuntimeSupport=all
-            MATCH (comm:Community)
-            WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
-            WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
-            WHERE score > $min_score
-            RETURN
-                comm.uuid As uuid,
-                comm.group_id AS group_id,
-                comm.name AS name, 
-                comm.name_embedding AS name_embedding,
-                comm.created_at AS created_at, 
-                comm.summary AS summary
-            ORDER BY score DESC
-            LIMIT $limit
-            """,
+        runtime_query
+        + """
+           MATCH (comm:Community)
+           WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
+           WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
+           WHERE score > $min_score
+           RETURN
+               comm.uuid As uuid,
+               comm.group_id AS group_id,
+               comm.name AS name, 
+               comm.name_embedding AS name_embedding,
+               comm.created_at AS created_at, 
+               comm.summary AS summary
+           ORDER BY score DESC
+           LIMIT $limit
+        """,
         search_vector=search_vector,
         group_ids=group_ids,
         limit=limit,

graphiti_core/utils/bulk_utils.py

@@ -21,12 +21,20 @@ from collections import defaultdict
 from datetime import datetime
 from math import ceil
 
-from neo4j import AsyncDriver
+from neo4j import AsyncDriver, AsyncManagedTransaction
 from numpy import dot, sqrt
 from pydantic import BaseModel
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
 from graphiti_core.llm_client import LLMClient
+from graphiti_core.models.edges.edge_db_queries import (
+    ENTITY_EDGE_SAVE_BULK,
+    EPISODIC_EDGE_SAVE_BULK,
+)
+from graphiti_core.models.nodes.node_db_queries import (
+    ENTITY_NODE_SAVE_BULK,
+    EPISODIC_NODE_SAVE_BULK,
+)
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search_utils import get_relevant_edges, get_relevant_nodes
 from graphiti_core.utils import retrieve_episodes
@@ -75,6 +83,35 @@ async def retrieve_previous_episodes_bulk(
     return episode_tuples
 
 
+async def add_nodes_and_edges_bulk(
+    driver: AsyncDriver,
+    episodic_nodes: list[EpisodicNode],
+    episodic_edges: list[EpisodicEdge],
+    entity_nodes: list[EntityNode],
+    entity_edges: list[EntityEdge],
+):
+    async with driver.session() as session:
+        await session.execute_write(
+            add_nodes_and_edges_bulk_tx, episodic_nodes, episodic_edges, entity_nodes, entity_edges
+        )
+
+
+async def add_nodes_and_edges_bulk_tx(
+    tx: AsyncManagedTransaction,
+    episodic_nodes: list[EpisodicNode],
+    episodic_edges: list[EpisodicEdge],
+    entity_nodes: list[EntityNode],
+    entity_edges: list[EntityEdge],
+):
+    episodes = [dict(episode) for episode in episodic_nodes]
+    for episode in episodes:
+        episode['source'] = str(episode['source'].value)
+    await tx.run(EPISODIC_NODE_SAVE_BULK, episodes=episodes)
+    await tx.run(ENTITY_NODE_SAVE_BULK, nodes=[dict(entity) for entity in entity_nodes])
+    await tx.run(EPISODIC_EDGE_SAVE_BULK, episodic_edges=[dict(edge) for edge in episodic_edges])
+    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[dict(edge) for edge in entity_edges])
+
+
 async def extract_nodes_and_edges_bulk(
     llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:

{graphiti_core-0.3.19.dist-info → graphiti_core-0.3.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.19
+Version: 0.3.21
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: diskcache (>=5.6.3,<6.0.0)
 Requires-Dist: neo4j (>=5.23.0,<6.0.0)
 Requires-Dist: numpy (>=1.0.0)
-Requires-Dist: openai (>=1.52.2,<2.0.0)
+Requires-Dist: openai (>=1.53.0,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: tenacity (<9.0.0)
 Description-Content-Type: text/markdown
@@ -194,6 +194,17 @@ The `server` directory contains an API service for interacting with the Graphiti
 
 Please see the [server README](./server/README.md) for more information.
 
+## Optional Environment Variables
+
+In addition to the Neo4j and OpenAi-compatible credentials, Graphiti also has a few optional environment variables.
+If you are using one of our supported models, such as Anthropic or Voyage models, the necessary environment variables
+must be set.
+
+`USE_PARALLEL_RUNTIME` is an optional boolean variable that can be set to true if you wish
+to enable Neo4j's parallel runtime feature for several of our search queries.
+Note that this feature is not supported for Neo4j Community edition or for smaller AuraDB instances,
+as such this feature is off by default.
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).
@@ -207,11 +218,11 @@ Graphiti is under active development. We aim to maintain API stability while wor
 - [x] Implementing node and edge CRUD operations
 - [ ] Improving performance and scalability
 - [ ] Achieving good performance with different LLM and embedding models
-- [ ] Creating a dedicated embedder interface
+- [x] Creating a dedicated embedder interface
 - [ ] Supporting custom graph schemas:
     - Allow developers to provide their own defined node and edge classes when ingesting episodes
     - Enable more flexible knowledge representation tailored to specific use cases
-- [ ] Enhancing retrieval capabilities with more robust and configurable options
+- [x] Enhancing retrieval capabilities with more robust and configurable options
 - [ ] Expanding test coverage to ensure reliability and catch edge cases
 
 ## Contributing

{graphiti_core-0.3.19.dist-info → graphiti_core-0.3.21.dist-info}/RECORD

@@ -9,8 +9,8 @@ graphiti_core/embedder/client.py,sha256=gVr_xdN-d0UQfeE4Nyoa4pL3M8UxllNH4eFqmarx
 graphiti_core/embedder/openai.py,sha256=yYUYPymx_lBlxDTGrlc03yNhPFyGG-etM2sszRK2G2U,1618
 graphiti_core/embedder/voyage.py,sha256=_eGFI5_NjNG8z7qG3jTWCdE7sAs1Yb8fiSZSJlQLD9o,1879
 graphiti_core/errors.py,sha256=ddHrHGQxhwkVAtSph4AV84UoOlgwZufMczXPwB7uqPo,1795
-graphiti_core/graphiti.py,sha256=c9Rh777TrHYffPF6qvFAfm-m-PA4kD8a3ZW_ShsZGxE,27714
-graphiti_core/helpers.py,sha256=kqC2TD8Auwty4sG7KH4BuRMX413oTChGaAT_XUt9ZjU,2108
+graphiti_core/graphiti.py,sha256=MIwkprC9L9FVj7cGrAzrNbpgkpPtUPPofLlfCrtXLHU,27584
+graphiti_core/helpers.py,sha256=6N5nNgyXiHWsMXHxm9M_S6Og7_TuejE5GuyV6EtZgVw,2178
 graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
 graphiti_core/llm_client/anthropic_client.py,sha256=4l2PbCjIoeRr7UJ2DUh2grYLTtE2vNaWlo72IIRQDeI,2405
 graphiti_core/llm_client/client.py,sha256=WAnX0e4EuCFHXdFHeq_O1HZsW1STSByvDCFUHMAHEFU,3394
@@ -21,9 +21,9 @@ graphiti_core/llm_client/openai_client.py,sha256=xLkbpusRVFRK0zPr3kOqY31HK_XCXrp
 graphiti_core/llm_client/utils.py,sha256=zKpxXEbKa369m4W7RDEf-m56kH46V1Mx3RowcWZEWWs,1000
 graphiti_core/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/edges/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/models/edges/edge_db_queries.py,sha256=gUC5AZeXvobsIzG8Zqxj4Sa-j0PKV4CuEPHWupFTXQs,1145
+graphiti_core/models/edges/edge_db_queries.py,sha256=0T1OSSQeP63OjoMTrscqVhq0gNRRxh3-n1sNxo0kTJE,2102
 graphiti_core/models/nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/models/nodes/node_db_queries.py,sha256=HxZIjVLqVRnG7OJHeZv7t3KCEWFTWvdNOBJyLSN-1Ts,918
+graphiti_core/models/nodes/node_db_queries.py,sha256=zIEKhodGOEYQ2ZdH766t-QRSk2xZNmtmAaPv60pvepM,1688
 graphiti_core/nodes.py,sha256=xlObDAC6Osp6L0wwqREuzBlIFlOzuUs9lN3ZpIv4seM,13844
 graphiti_core/prompts/__init__.py,sha256=EA-x9xUki9l8wnu2l8ek_oNf75-do5tq5hVq7Zbv8Kw,101
 graphiti_core/prompts/dedupe_edges.py,sha256=DUNHdIudj50FAjkla4nc68tSFSD2yjmYHBw-Bb7ph20,6529
@@ -38,12 +38,12 @@ graphiti_core/prompts/models.py,sha256=cvx_Bv5RMFUD_5IUawYrbpOKLPHogai7_bm7YXrSz
 graphiti_core/prompts/summarize_nodes.py,sha256=FLuZpGTABgcxuIDkx_IKH115nHEw0rIaFhcGlWveAMc,2357
 graphiti_core/py.typed,sha256=vlmmzQOt7bmeQl9L3XJP4W6Ry0iiELepnOrinKz5KQg,79
 graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/search/search.py,sha256=Fz45sXToQ6KN_y2k94kg9UyUZeqsS14jfaGfQ5KGifQ,11176
+graphiti_core/search/search.py,sha256=u512ZK2qPL28SmXDIo_RELq_dTRndo-GUPlSJrQqo5A,11316
 graphiti_core/search/search_config.py,sha256=UZN8jFA4pBlw2O5N1cuhVRBdTwMLR9N3Oyo6sQ4MDVw,3117
 graphiti_core/search/search_config_recipes.py,sha256=20jS7veJExDnXA-ovJSUJfyDHKt7GW-nng-eoiT7ATA,5810
-graphiti_core/search/search_utils.py,sha256=l8BR4GOo-A2eIXx4ybC18n6t6CeerN_9KQbYzCB6ix0,22551
+graphiti_core/search/search_utils.py,sha256=gqBlhnb75srhymQ4_yzb7o019UxtE3e7eHHK9NL6nCc,22916
 graphiti_core/utils/__init__.py,sha256=cJAcMnBZdHBQmWrZdU1PQ1YmaL75bhVUkyVpIPuOyns,260
-graphiti_core/utils/bulk_utils.py,sha256=JtoYTZPCigPa3n2E43Oe7QhFZRTA_QKNGy1jVgklHag,12614
+graphiti_core/utils/bulk_utils.py,sha256=azn8xlUT9u5-CFDkLSAY4y5zt5uvOvKBzatZZKDvp1g,13994
 graphiti_core/utils/maintenance/__init__.py,sha256=TRY3wWWu5kn3Oahk_KKhltrWnh0NACw0FskjqF6OtlA,314
 graphiti_core/utils/maintenance/community_operations.py,sha256=Ltx8cQJcFDXcQs17lHDkh2r3qmJalzpRdA3Azl2U3m8,9804
 graphiti_core/utils/maintenance/edge_operations.py,sha256=htG9h32u0O1g7ZdR-jFSHuzDgiO4MChYJsNuvrRMTu8,11270
@@ -51,7 +51,7 @@ graphiti_core/utils/maintenance/graph_data_operations.py,sha256=w66_SLlvPapuG91Y
 graphiti_core/utils/maintenance/node_operations.py,sha256=h5nlRojbXOGJs-alpv6z6WnZ1UCixVGlAQYBQUqz8Bs,9030
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=MvaRLWrBlDeYw8CQrKish1xbYcY5ovpfdqA2hSX7v5k,3367
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core-0.3.19.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.3.19.dist-info/METADATA,sha256=uwulLozcE0fov_hyYb3rIWUtuvcy6BObIuOcfng0aaQ,9396
-graphiti_core-0.3.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-graphiti_core-0.3.19.dist-info/RECORD,,
+graphiti_core-0.3.21.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.3.21.dist-info/METADATA,sha256=dGRf7FLQ0HBD58uYXHMjHzT8uF2dcJG2mPgbyw2CobU,9990
+graphiti_core-0.3.21.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+graphiti_core-0.3.21.dist-info/RECORD,,