graphiti-core 0.4.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

graphiti_core/prompts/extract_edges.py

@@ -17,9 +17,26 @@ limitations under the License.
 import json
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class Edge(BaseModel):
+    relation_type: str = Field(..., description='RELATION_TYPE_IN_CAPS')
+    source_entity_name: str = Field(..., description='name of the source entity')
+    target_entity_name: str = Field(..., description='name of the target entity')
+    fact: str = Field(..., description='extracted factual information')
+
+
+class ExtractedEdges(BaseModel):
+    edges: list[Edge]
+
+
+class MissingFacts(BaseModel):
+    missing_facts: list[str] = Field(..., description="facts that weren't extracted")
+
+
 class Prompt(Protocol):
     edge: PromptVersion
     reflexion: PromptVersion
@@ -54,25 +71,12 @@ def edge(context: dict[str, Any]) -> list[Message]:
 
         Given the above MESSAGES and ENTITIES, extract all facts pertaining to the listed ENTITIES from the CURRENT MESSAGE. 
         
-
         Guidelines:
         1. Extract facts only between the provided entities.
         2. Each fact should represent a clear relationship between two DISTINCT nodes.
         3. The relation_type should be a concise, all-caps description of the fact (e.g., LOVES, IS_FRIENDS_WITH, WORKS_FOR).
         4. Provide a more detailed fact containing all relevant information.
         5. Consider temporal aspects of relationships when relevant.
-
-        Respond with a JSON object in the following format:
-        {{
-            "edges": [
-                {{
-                    "relation_type": "RELATION_TYPE_IN_CAPS",
-                    "source_entity_name": "name of the source entity",
-                    "target_entity_name": "name of the target entity",
-                    "fact": "extracted factual information",
-                }}
-            ]
-        }}
         """,
         ),
     ]
@@ -98,12 +102,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
 </EXTRACTED FACTS>
 
 Given the above MESSAGES, list of EXTRACTED ENTITIES entities, and list of EXTRACTED FACTS; 
-determine if any facts haven't been extracted:
-
-Respond with a JSON object in the following format:
-{{
-    "missing_facts": [ "facts that weren't extracted", ...]
-}}
+determine if any facts haven't been extracted.
 """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/prompts/extract_nodes.py

@@ -17,9 +17,19 @@ limitations under the License.
 import json
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class ExtractedNodes(BaseModel):
+    extracted_node_names: list[str] = Field(..., description='Name of the extracted entity')
+
+
+class MissedEntities(BaseModel):
+    missed_entities: list[str] = Field(..., description="Names of entities that weren't extracted")
+
+
 class Prompt(Protocol):
     extract_message: PromptVersion
     extract_json: PromptVersion
@@ -56,11 +66,6 @@ Guidelines:
 4. DO NOT create nodes for temporal information like dates, times or years (these will be added to edges later).
 5. Be as explicit as possible in your node names, using full names.
 6. DO NOT extract entities mentioned only in PREVIOUS MESSAGES, those messages are only to provide context.
-
-Respond with a JSON object in the following format:
-{{
-    "extracted_node_names": ["Name of the extracted entity", ...],
-}}
 """
     return [
         Message(role='system', content=sys_prompt),
@@ -87,11 +92,6 @@ Given the above source description and JSON, extract relevant entity nodes from
 Guidelines:
 1. Always try to extract an entities that the JSON represents. This will often be something like a "name" or "user field
 2. Do NOT extract any properties that contain dates
-
-Respond with a JSON object in the following format:
-{{
-    "extracted_node_names": ["Name of the extracted entity", ...],
-}}
 """
     return [
         Message(role='system', content=sys_prompt),
@@ -116,11 +116,6 @@ Guidelines:
 2. Avoid creating nodes for relationships or actions.
 3. Avoid creating nodes for temporal information like dates, times or years (these will be added to edges later).
 4. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
-
-Respond with a JSON object in the following format:
-{{
-    "extracted_node_names": ["Name of the extracted entity", ...],
-}}
 """
     return [
         Message(role='system', content=sys_prompt),
@@ -144,12 +139,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
 </EXTRACTED ENTITIES>
 
 Given the above previous messages, current message, and list of extracted entities; determine if any entities haven't been
-extracted:
-
-Respond with a JSON object in the following format:
-{{
-    "missed_entities": [ "name of entity that wasn't extracted", ...]
-}}
+extracted.
 """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/prompts/invalidate_edges.py

@@ -16,9 +16,22 @@ limitations under the License.
 
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class InvalidatedEdge(BaseModel):
+    uuid: str = Field(..., description='The UUID of the edge to be invalidated')
+    fact: str = Field(..., description='Updated fact of the edge')
+
+
+class InvalidatedEdges(BaseModel):
+    invalidated_edges: list[InvalidatedEdge] = Field(
+        ..., description='List of edges that should be invalidated'
+    )
+
+
 class Prompt(Protocol):
     v1: PromptVersion
     v2: PromptVersion
@@ -56,18 +69,6 @@ def v1(context: dict[str, Any]) -> list[Message]:
                 {context['new_edges']}
 
                 Each edge is formatted as: "UUID | SOURCE_NODE - EDGE_NAME - TARGET_NODE (fact: EDGE_FACT), START_DATE (END_DATE, optional))"
-
-                For each existing edge that should be invalidated, respond with a JSON object in the following format:
-                {{
-                    "invalidated_edges": [
-                        {{
-                            "edge_uuid": "The UUID of the edge to be invalidated (the part before the | character)",
-                            "fact": "Updated fact of the edge"
-                        }}
-                    ]
-                }}
-
-                If no relationships need to be invalidated based on these strict criteria, return an empty list for "invalidated_edges".
             """,
         ),
     ]
@@ -89,19 +90,6 @@ def v2(context: dict[str, Any]) -> list[Message]:
 
                 New Edge:
                 {context['new_edge']}
-
-
-                For each existing edge that should be invalidated, respond with a JSON object in the following format:
-                {{
-                    "invalidated_edges": [
-                        {{
-                            "uuid": "The UUID of the edge to be invalidated",
-                            "fact": "Updated fact of the edge"
-                        }}
-                    ]
-                }}
-
-                If no relationships need to be invalidated based on these strict criteria, return an empty list for "invalidated_edges".
             """,
         ),
     ]

graphiti_core/prompts/lib.py

@@ -74,6 +74,7 @@ from .invalidate_edges import (
     versions as invalidate_edges_versions,
 )
 from .models import Message, PromptFunction
+from .prompt_helpers import DO_NOT_ESCAPE_UNICODE
 from .summarize_nodes import Prompt as SummarizeNodesPrompt
 from .summarize_nodes import Versions as SummarizeNodesVersions
 from .summarize_nodes import versions as summarize_nodes_versions
@@ -106,7 +107,10 @@ class VersionWrapper:
         self.func = func
 
     def __call__(self, context: dict[str, Any]) -> list[Message]:
-        return self.func(context)
+        messages = self.func(context)
+        for message in messages:
+            message.content += DO_NOT_ESCAPE_UNICODE if message.role == 'system' else ''
+        return messages
 
 
 class PromptTypeWrapper:

graphiti_core/prompts/prompt_helpers.py

@@ -0,0 +1 @@
+DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n'

graphiti_core/prompts/summarize_nodes.py

@@ -17,9 +17,21 @@ limitations under the License.
 import json
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class Summary(BaseModel):
+    summary: str = Field(
+        ..., description='Summary containing the important information from both summaries'
+    )
+
+
+class SummaryDescription(BaseModel):
+    description: str = Field(..., description='One sentence description of the provided summary')
+
+
 class Prompt(Protocol):
     summarize_pair: PromptVersion
     summarize_context: PromptVersion
@@ -42,14 +54,11 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         Synthesize the information from the following two summaries into a single succinct summary.
+        
+        Summaries must be under 500 words.
 
         Summaries:
         {json.dumps(context['node_summaries'], indent=2)}
-
-        Respond with a JSON object in the following format:
-            {{
-                "summary": "Summary containing the important information from both summaries"
-            }}
         """,
         ),
     ]
@@ -74,15 +83,11 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         information from the provided MESSAGES. Your summary should also only contain information relevant to the
         provided ENTITY.
         
+        Summaries must be under 500 words.
+        
         <ENTITY>
         {context['node_name']}
         </ENTITY>
-        
-
-        Respond with a JSON object in the following format:
-            {{
-                "summary": "Entity summary"
-            }}
         """,
         ),
     ]
@@ -98,14 +103,10 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         Create a short one sentence description of the summary that explains what kind of information is summarized.
+        Summaries must be under 500 words.
 
         Summary:
         {json.dumps(context['summary'], indent=2)}
-
-        Respond with a JSON object in the following format:
-            {{
-                "description": "One sentence description of the provided summary"
-            }}
         """,
         ),
     ]

graphiti_core/search/search.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
@@ -25,6 +24,7 @@ from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.edges import EntityEdge
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import SearchRerankerError
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.nodes import CommunityNode, EntityNode
 from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
@@ -78,7 +78,7 @@ async def search(
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
-    edges, nodes, communities = await asyncio.gather(
+    edges, nodes, communities = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
@@ -141,7 +141,7 @@ async def edge_search(
         return []
 
     search_results: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 edge_fulltext_search(driver, query, group_ids, 2 * limit),
                 edge_similarity_search(
@@ -226,7 +226,7 @@ async def node_search(
         return []
 
     search_results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 node_fulltext_search(driver, query, group_ids, 2 * limit),
                 node_similarity_search(
@@ -295,7 +295,7 @@ async def community_search(
         return []
 
     search_results: list[list[CommunityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 community_fulltext_search(driver, query, group_ids, 2 * limit),
                 community_similarity_search(

graphiti_core/search/search_utils.py

@@ -14,10 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
+from typing import Any
 
 import numpy as np
 from neo4j import AsyncDriver, Query
@@ -29,6 +29,7 @@ from graphiti_core.helpers import (
     USE_PARALLEL_RUNTIME,
     lucene_sanitize,
     normalize_l2,
+    semaphore_gather,
 )
 from graphiti_core.nodes import (
     CommunityNode,
@@ -40,7 +41,7 @@ from graphiti_core.nodes import (
 
 logger = logging.getLogger(__name__)
 
-RELEVANT_SCHEMA_LIMIT = 3
+RELEVANT_SCHEMA_LIMIT = 10
 DEFAULT_MIN_SCORE = 0.6
 DEFAULT_MMR_LAMBDA = 0.5
 MAX_SEARCH_DEPTH = 3
@@ -191,12 +192,27 @@ async def edge_similarity_search(
         '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])
-                AND ($target_uuid IS NULL OR m.uuid IN [$source_uuid, $target_uuid])
-                WITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE r.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+        query_params['source_node_uuid'] = source_node_uuid
+        query_params['target_node_uuid'] = target_node_uuid
+
+        if source_node_uuid is not None:
+            group_filter_query += '\nAND (n.uuid IN [$source_uuid, $target_uuid])'
+
+        if target_node_uuid is not None:
+            group_filter_query += '\nAND (m.uuid IN [$source_uuid, $target_uuid])'
+
+    query: LiteralString = (
+        """
+                        MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                        """
+        + group_filter_query
+        + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
                 WHERE score > $min_score
                 RETURN
                     r.uuid AS uuid,
@@ -214,9 +230,11 @@ async def edge_similarity_search(
                 ORDER BY score DESC
                 LIMIT $limit
         """
+    )
 
     records, _, _ = await driver.execute_query(
         runtime_query + query,
+        query_params,
         search_vector=search_vector,
         source_uuid=source_node_uuid,
         target_uuid=target_node_uuid,
@@ -325,11 +343,20 @@ async def node_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE n.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
             MATCH (n:Entity)
-            WHERE $group_ids IS NULL OR n.group_id IN $group_ids
+            """
+        + group_filter_query
+        + """
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
             WHERE score > $min_score
             RETURN
@@ -342,6 +369,7 @@ async def node_similarity_search(
             ORDER BY score DESC
             LIMIT $limit
             """,
+        query_params,
         search_vector=search_vector,
         group_ids=group_ids,
         limit=limit,
@@ -436,11 +464,20 @@ async def community_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE comm.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
            MATCH (comm:Community)
-           WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
+           """
+        + group_filter_query
+        + """
            WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
            WHERE score > $min_score
            RETURN
@@ -512,7 +549,7 @@ async def hybrid_node_search(
 
     start = time()
     results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[node_fulltext_search(driver, q, group_ids, 2 * limit) for q in queries],
             *[node_similarity_search(driver, e, group_ids, 2 * limit) for e in embeddings],
         )
@@ -582,7 +619,7 @@ async def get_relevant_edges(
     relevant_edges: list[EntityEdge] = []
     relevant_edge_uuids = set()
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             edge_similarity_search(
                 driver,
@@ -631,7 +668,7 @@ async def node_distance_reranker(
 ) -> list[str]:
     # filter out node_uuid center node node uuid
     filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
-    scores: dict[str, float] = {}
+    scores: dict[str, float] = {center_node_uuid: 0.0}
 
     # Find the shortest path to center node
     query = Query("""
@@ -649,9 +686,13 @@ async def node_distance_reranker(
 
     for result in path_results:
         uuid = result['uuid']
-        score = result['score'] if 'score' in result else float('inf')
+        score = result['score']
         scores[uuid] = score
 
+    for uuid in filtered_uuids:
+        if uuid not in scores:
+            scores[uuid] = float('inf')
+
     # rerank on shortest distance
     filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
 

graphiti_core/utils/__init__.py

@@ -1,15 +0,0 @@
-from .maintenance import (
-    build_episodic_edges,
-    clear_data,
-    extract_edges,
-    extract_nodes,
-    retrieve_episodes,
-)
-
-__all__ = [
-    'extract_edges',
-    'build_episodic_edges',
-    'extract_nodes',
-    'clear_data',
-    'retrieve_episodes',
-]

graphiti_core/utils/bulk_utils.py

@@ -14,11 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 import typing
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime
 from math import ceil
 
 from neo4j import AsyncDriver, AsyncManagedTransaction
@@ -26,6 +25,7 @@ from numpy import dot, sqrt
 from pydantic import BaseModel
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.models.edges.edge_db_queries import (
     ENTITY_EDGE_SAVE_BULK,
@@ -37,14 +37,17 @@ from graphiti_core.models.nodes.node_db_queries import (
 )
 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
+from graphiti_core.utils.datetime_utils import utc_now
 from graphiti_core.utils.maintenance.edge_operations import (
     build_episodic_edges,
     dedupe_edge_list,
     dedupe_extracted_edges,
     extract_edges,
 )
-from graphiti_core.utils.maintenance.graph_data_operations import EPISODE_WINDOW_LEN
+from graphiti_core.utils.maintenance.graph_data_operations import (
+    EPISODE_WINDOW_LEN,
+    retrieve_episodes,
+)
 from graphiti_core.utils.maintenance.node_operations import (
     dedupe_extracted_nodes,
     dedupe_node_list,
@@ -68,7 +71,7 @@ class RawEpisode(BaseModel):
 async def retrieve_previous_episodes_bulk(
     driver: AsyncDriver, episodes: list[EpisodicNode]
 ) -> list[tuple[EpisodicNode, list[EpisodicNode]]]:
-    previous_episodes_list = await asyncio.gather(
+    previous_episodes_list = await semaphore_gather(
         *[
             retrieve_episodes(
                 driver, episode.valid_at, last_n=EPISODE_WINDOW_LEN, group_ids=[episode.group_id]
@@ -115,7 +118,7 @@ async def add_nodes_and_edges_bulk_tx(
 async def extract_nodes_and_edges_bulk(
     llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
-    extracted_nodes_bulk = await asyncio.gather(
+    extracted_nodes_bulk = await semaphore_gather(
         *[
             extract_nodes(llm_client, episode, previous_episodes)
             for episode, previous_episodes in episode_tuples
@@ -127,7 +130,7 @@ async def extract_nodes_and_edges_bulk(
         [episode[1] for episode in episode_tuples],
     )
 
-    extracted_edges_bulk = await asyncio.gather(
+    extracted_edges_bulk = await semaphore_gather(
         *[
             extract_edges(
                 llm_client,
@@ -168,13 +171,13 @@ async def dedupe_nodes_bulk(
     node_chunks = [nodes[i : i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_nodes(driver, node_chunk) for node_chunk in node_chunks]
         )
     )
 
     results: list[tuple[list[EntityNode], dict[str, str]]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_nodes(llm_client, node_chunk, existing_nodes_chunks[i])
                 for i, node_chunk in enumerate(node_chunks)
@@ -202,13 +205,13 @@ async def dedupe_edges_bulk(
     ]
 
     relevant_edges_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_edges(driver, edge_chunk, None, None) for edge_chunk in edge_chunks]
         )
     )
 
     resolved_edge_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_edges(llm_client, edge_chunk, relevant_edges_chunks[i])
                 for i, edge_chunk in enumerate(edge_chunks)
@@ -289,7 +292,9 @@ async def compress_nodes(
             # add both nodes to the shortest chunk
             node_chunks[-1].extend([n, m])
 
-    results = await asyncio.gather(*[dedupe_node_list(llm_client, chunk) for chunk in node_chunks])
+    results = await semaphore_gather(
+        *[dedupe_node_list(llm_client, chunk) for chunk in node_chunks]
+    )
 
     extended_map = dict(uuid_map)
     compressed_nodes: list[EntityNode] = []
@@ -312,7 +317,9 @@ async def compress_edges(llm_client: LLMClient, edges: list[EntityEdge]) -> list
     # We build a map of the edges based on their source and target nodes.
     edge_chunks = chunk_edges_by_nodes(edges)
 
-    results = await asyncio.gather(*[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks])
+    results = await semaphore_gather(
+        *[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks]
+    )
 
     compressed_edges: list[EntityEdge] = []
     for edge_chunk in results:
@@ -365,7 +372,7 @@ async def extract_edge_dates_bulk(
         episode.uuid: (episode, previous_episodes) for episode, previous_episodes in episode_pairs
     }
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             extract_edge_dates(
                 llm_client,
@@ -385,7 +392,7 @@ async def extract_edge_dates_bulk(
         edge.valid_at = valid_at
         edge.invalid_at = invalid_at
         if edge.invalid_at:
-            edge.expired_at = datetime.now(timezone.utc)
+            edge.expired_at = utc_now()
 
     return edges