graphiti-core 0.3.20__py3-none-any.whl → 0.4.0__py3-none-any.whl

graphiti_core/prompts/extract_edges.py

@@ -21,120 +21,94 @@ from .models import Message, PromptFunction, PromptVersion
 
 
 class Prompt(Protocol):
-    v1: PromptVersion
-    v2: PromptVersion
+    edge: PromptVersion
+    reflexion: PromptVersion
 
 
 class Versions(TypedDict):
-    v1: PromptFunction
-    v2: PromptFunction
+    edge: PromptFunction
+    reflexion: PromptFunction
 
 
-def v1(context: dict[str, Any]) -> list[Message]:
+def edge(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
             role='system',
-            content='You are a helpful assistant that extracts graph edges from provided context.',
+            content='You are an expert fact extractor that extracts fact triples from text.',
         ),
         Message(
             role='user',
             content=f"""
-        Given the following context, extract new semantic edges (relationships) that need to be added to the knowledge graph:
-
-        Current Graph Structure:
-        {context['relevant_schema']}
-
-        New Nodes:
-        {json.dumps(context['new_nodes'], indent=2)}
-
-        New Episode:
-        Content: {context['episode_content']}
-        Timestamp: {context['episode_timestamp']}
-
-        Previous Episodes:
-        {json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
+        <PREVIOUS MESSAGES>
+        {json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+        </PREVIOUS MESSAGES>
+        <CURRENT MESSAGE>
+        {context["episode_content"]}
+        </CURRENT MESSAGE>
+        
+        <ENTITIES>
+        {context["nodes"]}
+        </ENTITIES>
+        
+        {context['custom_prompt']}
 
-        Extract new semantic edges based on the content of the current episode, considering the existing graph structure, new nodes, and context from previous episodes.
+        Given the above MESSAGES and ENTITIES, extract all facts pertaining to the listed ENTITIES from the CURRENT MESSAGE. 
+        
 
         Guidelines:
-        1. Create edges only between semantic nodes (not episodic nodes like messages).
-        2. Each edge should represent a clear relationship between two semantic nodes.
-        3. The relation_type should be a concise, all-caps description of the relationship (e.g., LOVES, IS_FRIENDS_WITH, WORKS_FOR).
-        4. Provide a more detailed fact describing the relationship.
-        5. If a relationship seems to update an existing one, create a new edge with the updated information.
-        6. Consider temporal aspects of relationships when relevant.
-        7. Do not create edges involving episodic nodes (like Message 1 or Message 2).
-        8. Use existing nodes from the current graph structure when appropriate.
+        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:
         {{
-            "new_edges": [
+            "edges": [
                 {{
                     "relation_type": "RELATION_TYPE_IN_CAPS",
-                    "source_node": "Name of the source semantic node",
-                    "target_node": "Name of the target semantic node",
-                    "fact": "Detailed description of the relationship",
-                    "valid_at": "YYYY-MM-DDTHH:MM:SSZ or null if not explicitly mentioned",
-                    "invalid_at": "YYYY-MM-DDTHH:MM:SSZ or null if ongoing or not explicitly mentioned"
+                    "source_entity_name": "name of the source entity",
+                    "target_entity_name": "name of the target entity",
+                    "fact": "extracted factual information",
                 }}
             ]
         }}
-
-        If no new edges need to be added, return an empty list for "new_edges".
         """,
         ),
     ]
 
 
-def v2(context: dict[str, Any]) -> list[Message]:
-    return [
-        Message(
-            role='system',
-            content='You are a helpful assistant that extracts graph edges from provided context.',
-        ),
-        Message(
-            role='user',
-            content=f"""
-        Given the following context, extract edges (relationships) that need to be added to the knowledge graph:
-        Nodes:
-        {json.dumps(context['nodes'], indent=2)}
-
-        
+def reflexion(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an AI assistant that determines which facts have not been extracted from the given context"""
 
-        Episodes:
-        {json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
-        {context['episode_content']} <-- New Episode
-        
+    user_prompt = f"""
+<PREVIOUS MESSAGES>
+{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+</PREVIOUS MESSAGES>
+<CURRENT MESSAGE>
+{context["episode_content"]}
+</CURRENT MESSAGE>
 
-        Extract entity edges based on the content of the current episode, the given nodes, and context from previous episodes.
+<EXTRACTED ENTITIES>
+{context["nodes"]}
+</EXTRACTED ENTITIES>
 
-        Guidelines:
-        1. Create edges only between the provided nodes.
-        2. Each edge should represent a clear relationship between two DISTINCT nodes.
-        3. The relation_type should be a concise, all-caps description of the relationship (e.g., LOVES, IS_FRIENDS_WITH, WORKS_FOR).
-        4. Provide a more detailed fact describing the relationship.
-        5. The fact should include any specific relevant information, including numeric information
-        6. Consider temporal aspects of relationships when relevant.
-        7. Avoid using the same node as the source and target of a relationship
+<EXTRACTED FACTS>
+{context["extracted_facts"]}
+</EXTRACTED FACTS>
 
-        Respond with a JSON object in the following format:
-        {{
-            "edges": [
-                {{
-                    "relation_type": "RELATION_TYPE_IN_CAPS",
-                    "source_node_uuid": "uuid of the source entity node",
-                    "target_node_uuid": "uuid of the target entity node",
-                    "fact": "brief description of the relationship",
-                    "valid_at": "YYYY-MM-DDTHH:MM:SSZ or null if not explicitly mentioned",
-                    "invalid_at": "YYYY-MM-DDTHH:MM:SSZ or null if ongoing or not explicitly mentioned"
-                }}
-            ]
-        }}
+Given the above MESSAGES, list of EXTRACTED ENTITIES entities, and list of EXTRACTED FACTS; 
+determine if any facts haven't been extracted:
 
-        If no edges need to be added, return an empty list for "edges".
-        """,
-        ),
+Respond with a JSON object in the following format:
+{{
+    "missing_facts": [ "facts that weren't extracted", ...]
+}}
+"""
+    return [
+        Message(role='system', content=sys_prompt),
+        Message(role='user', content=user_prompt),
     ]
 
 
-versions: Versions = {'v1': v1, 'v2': v2}
+versions: Versions = {'edge': edge, 'reflexion': reflexion}

graphiti_core/prompts/extract_nodes.py

@@ -21,89 +21,45 @@ from .models import Message, PromptFunction, PromptVersion
 
 
 class Prompt(Protocol):
-    v1: PromptVersion
-    v2: PromptVersion
+    extract_message: PromptVersion
     extract_json: PromptVersion
     extract_text: PromptVersion
+    reflexion: PromptVersion
 
 
 class Versions(TypedDict):
-    v1: PromptFunction
-    v2: PromptFunction
+    extract_message: PromptFunction
     extract_json: PromptFunction
     extract_text: PromptFunction
+    reflexion: PromptFunction
 
 
-def v1(context: dict[str, Any]) -> list[Message]:
-    return [
-        Message(
-            role='system',
-            content='You are a helpful assistant that extracts graph nodes from provided context.',
-        ),
-        Message(
-            role='user',
-            content=f"""
-        Given the following context, extract new entity nodes that need to be added to the knowledge graph:
-
-        Previous Episodes:
-        {json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
-
-        New Episode:
-        Content: {context["episode_content"]}
-
-        Extract new entity nodes based on the content of the current episode, while considering the context from previous episodes.
-
-        Guidelines:
-        1. Focus on entities, concepts, or actors that are central to the current episode.
-        2. Avoid creating nodes for relationships or actions (these will be handled as edges later).
-        3. Provide a brief but informative summary for each node.
-        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:
-        {{
-            "new_nodes": [
-                {{
-                    "name": "Unique identifier for the node",
-                    "labels": ["Entity", "OptionalAdditionalLabel"],
-                    "summary": "Brief summary of the node's role or significance"
-                }}
-            ]
-        }}
-
-        If no new nodes need to be added, return an empty list for "new_nodes".
-        """,
-        ),
-    ]
-
-
-def v2(context: dict[str, Any]) -> list[Message]:
-    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational text. Your primary task is to identify and extract the speaker and other significant entities mentioned in the conversation."""
+def extract_message(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational messages. Your primary task is to identify and extract the speaker and other significant entities mentioned in the conversation."""
 
     user_prompt = f"""
-Given the following conversation, extract entity nodes from the CURRENT MESSAGE that are explicitly or implicitly mentioned:
-
-Conversation:
-{json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
+<PREVIOUS MESSAGES>
+{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+</PREVIOUS MESSAGES>
 <CURRENT MESSAGE>
 {context["episode_content"]}
+</CURRENT MESSAGE>
+
+{context['custom_prompt']}
+
+Given the above conversation, extract entity nodes from the CURRENT MESSAGE that are explicitly or implicitly mentioned:
 
 Guidelines:
 1. ALWAYS extract the speaker/actor as the first node. The speaker is the part before the colon in each line of dialogue.
-2. Extract other significant entities, concepts, or actors mentioned in the conversation.
-3. Provide concise but informative summaries for each extracted node.
-4. Avoid creating nodes for relationships or actions.
-5. Avoid creating nodes for temporal information like dates, times or years (these will be added to edges later).
-6. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
+2. Extract other significant entities, concepts, or actors mentioned in the CURRENT MESSAGE.
+3. DO NOT create nodes for relationships or actions.
+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_nodes": [
-        {{
-            "name": "Unique identifier for the node (use the speaker's name for speaker nodes)",
-            "labels": ["Entity", "Speaker" for speaker nodes, "OptionalAdditionalLabel"],
-            "summary": "Brief summary of the node's role or significance"
-        }}
-    ]
+    "extracted_node_names": ["Name of the extracted entity", ...],
 }}
 """
     return [
@@ -113,17 +69,20 @@ Respond with a JSON object in the following format:
 
 
 def extract_json(context: dict[str, Any]) -> list[Message]:
-    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational text. 
+    sys_prompt = """You are an AI assistant that extracts entity nodes from JSON. 
     Your primary task is to identify and extract relevant entities from JSON files"""
 
     user_prompt = f"""
-Given the following source description, extract relevant entity nodes from the provided JSON:
-
-Source Description:
+<SOURCE DESCRIPTION>:
 {context["source_description"]}
-
-JSON:
+</SOURCE DESCRIPTION>
+<JSON>
 {context["episode_content"]}
+</JSON>
+
+{context['custom_prompt']}
+
+Given the above source description and JSON, extract relevant entity nodes from the provided JSON:
 
 Guidelines:
 1. Always try to extract an entities that the JSON represents. This will often be something like a "name" or "user field
@@ -131,13 +90,7 @@ Guidelines:
 
 Respond with a JSON object in the following format:
 {{
-    "extracted_nodes": [
-        {{
-            "name": "Unique identifier for the node (use the speaker's name for speaker nodes)",
-            "labels": ["Entity", "Speaker" for speaker nodes, "OptionalAdditionalLabel"],
-            "summary": "Brief summary of the node's role or significance"
-        }}
-    ]
+    "extracted_node_names": ["Name of the extracted entity", ...],
 }}
 """
     return [
@@ -147,32 +100,55 @@ Respond with a JSON object in the following format:
 
 
 def extract_text(context: dict[str, Any]) -> list[Message]:
-    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational text. Your primary task is to identify and extract the speaker and other significant entities mentioned in the conversation."""
+    sys_prompt = """You are an AI assistant that extracts entity nodes from text. Your primary task is to identify and extract the speaker and other significant entities mentioned in the provided text."""
 
     user_prompt = f"""
-Given the following conversation, extract entity nodes from the CURRENT MESSAGE that are explicitly or implicitly mentioned:
-
-Conversation:
-{json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
-<CURRENT MESSAGE>
+<TEXT>
 {context["episode_content"]}
+</TEXT>
+
+{context['custom_prompt']}
+
+Given the following text, extract entity nodes from the TEXT that are explicitly or implicitly mentioned:
 
 Guidelines:
-2. Extract significant entities, concepts, or actors mentioned in the conversation.
-3. Provide concise but informative summaries for each extracted node.
-4. Avoid creating nodes for relationships or actions.
-5. Avoid creating nodes for temporal information like dates, times or years (these will be added to edges later).
-6. Be as explicit as possible in your node names, using full names and avoiding abbreviations.
+1. Extract significant entities, concepts, or actors mentioned in the conversation.
+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_nodes": [
-        {{
-            "name": "Unique identifier for the node (use the speaker's name for speaker nodes)",
-            "labels": ["Entity", "OptionalAdditionalLabel"],
-            "summary": "Brief summary of the node's role or significance"
-        }}
+    "extracted_node_names": ["Name of the extracted entity", ...],
+}}
+"""
+    return [
+        Message(role='system', content=sys_prompt),
+        Message(role='user', content=user_prompt),
     ]
+
+
+def reflexion(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an AI assistant that determines which entities have not been extracted from the given context"""
+
+    user_prompt = f"""
+<PREVIOUS MESSAGES>
+{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+</PREVIOUS MESSAGES>
+<CURRENT MESSAGE>
+{context["episode_content"]}
+</CURRENT MESSAGE>
+
+<EXTRACTED ENTITIES>
+{context["extracted_entities"]}
+</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", ...]
 }}
 """
     return [
@@ -182,8 +158,8 @@ Respond with a JSON object in the following format:
 
 
 versions: Versions = {
-    'v1': v1,
-    'v2': v2,
+    'extract_message': extract_message,
     'extract_json': extract_json,
     'extract_text': extract_text,
+    'reflexion': reflexion,
 }

graphiti_core/prompts/summarize_nodes.py

@@ -22,11 +22,13 @@ from .models import Message, PromptFunction, PromptVersion
 
 class Prompt(Protocol):
     summarize_pair: PromptVersion
+    summarize_context: PromptVersion
     summary_description: PromptVersion
 
 
 class Versions(TypedDict):
     summarize_pair: PromptFunction
+    summarize_context: PromptFunction
     summary_description: PromptFunction
 
 
@@ -53,6 +55,39 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def summarize_context(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that combines summaries with new conversation context.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+            
+        <MESSAGES>
+        {json.dumps(context['previous_episodes'], indent=2)}
+        {json.dumps(context['episode_content'], indent=2)}
+        </MESSAGES>
+        
+        Given the above MESSAGES and the following ENTITY name, create a summary for the ENTITY. Your summary must only use
+        information from the provided MESSAGES. Your summary should also only contain information relevant to the
+        provided ENTITY.
+        
+        <ENTITY>
+        {context['node_name']}
+        </ENTITY>
+        
+
+        Respond with a JSON object in the following format:
+            {{
+                "summary": "Entity summary"
+            }}
+        """,
+        ),
+    ]
+
+
 def summary_description(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
@@ -76,4 +111,8 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'summarize_pair': summarize_pair, 'summary_description': summary_description}
+versions: Versions = {
+    'summarize_pair': summarize_pair,
+    'summarize_context': summarize_context,
+    'summary_description': summary_description,
+}

graphiti_core/search/search.py

@@ -32,8 +32,10 @@ from graphiti_core.search.search_config import (
     CommunitySearchConfig,
     EdgeReranker,
     EdgeSearchConfig,
+    EdgeSearchMethod,
     NodeReranker,
     NodeSearchConfig,
+    NodeSearchMethod,
     SearchConfig,
     SearchResults,
 )
@@ -66,6 +68,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
@@ -144,6 +152,12 @@ async def edge_search(
         )
     )
 
+    if EdgeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
+        source_node_uuids = [edge.source_node_uuid for result in search_results for edge in result]
+        search_results.append(
+            await edge_bfs_search(driver, source_node_uuids, config.bfs_max_depth, 2 * limit)
+        )
+
     edge_uuid_map = {edge.uuid: edge for result in search_results for edge in result}
 
     reranked_uuids: list[str] = []
@@ -223,6 +237,12 @@ async def node_search(
         )
     )
 
+    if NodeSearchMethod.bfs in config.search_methods and bfs_origin_node_uuids is None:
+        origin_node_uuids = [node.uuid for result in search_results for node in result]
+        search_results.append(
+            await node_bfs_search(driver, origin_node_uuids, config.bfs_max_depth, 2 * limit)
+        )
+
     search_result_uuids = [[node.uuid for node in result] for result in search_results]
     node_uuid_map = {node.uuid: node for result in search_results for node in result}
 

graphiti_core/search/search_config_recipes.py

@@ -118,6 +118,19 @@ EDGE_HYBRID_SEARCH_EPISODE_MENTIONS = SearchConfig(
     )
 )
 
+# performs a hybrid search over edges with cross encoder reranking
+EDGE_HYBRID_SEARCH_CROSS_ENCODER = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[
+            EdgeSearchMethod.bm25,
+            EdgeSearchMethod.cosine_similarity,
+            EdgeSearchMethod.bfs,
+        ],
+        reranker=EdgeReranker.cross_encoder,
+    ),
+    limit=10,
+)
+
 # performs a hybrid search over nodes with rrf reranking
 NODE_HYBRID_SEARCH_RRF = SearchConfig(
     node_config=NodeSearchConfig(
@@ -150,6 +163,19 @@ NODE_HYBRID_SEARCH_EPISODE_MENTIONS = SearchConfig(
     )
 )
 
+# performs a hybrid search over nodes with episode mentions reranking
+NODE_HYBRID_SEARCH_CROSS_ENCODER = SearchConfig(
+    node_config=NodeSearchConfig(
+        search_methods=[
+            NodeSearchMethod.bm25,
+            NodeSearchMethod.cosine_similarity,
+            NodeSearchMethod.bfs,
+        ],
+        reranker=NodeReranker.cross_encoder,
+    ),
+    limit=10,
+)
+
 # performs a hybrid search over communities with rrf reranking
 COMMUNITY_HYBRID_SEARCH_RRF = SearchConfig(
     community_config=CommunitySearchConfig(
@@ -165,3 +191,12 @@ COMMUNITY_HYBRID_SEARCH_MMR = SearchConfig(
         reranker=CommunityReranker.mmr,
     )
 )
+
+# performs a hybrid search over communities with mmr reranking
+COMMUNITY_HYBRID_SEARCH_CROSS_ENCODER = SearchConfig(
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.cross_encoder,
+    ),
+    limit=3,
+)

graphiti_core/search/search_utils.py

@@ -146,7 +146,7 @@ async def edge_fulltext_search(
         return []
 
     cypher_query = Query("""
-              CALL db.index.fulltext.queryRelationships("edge_name_and_fact", $query) 
+              CALL db.index.fulltext.queryRelationships("edge_name_and_fact", $query, {limit: $limit}) 
               YIELD relationship AS rel, score
               MATCH (n:Entity)-[r {uuid: rel.uuid}]->(m:Entity)
               WHERE ($source_uuid IS NULL OR n.uuid IN [$source_uuid, $target_uuid])
@@ -296,7 +296,7 @@ async def node_fulltext_search(
 
     records, _, _ = await driver.execute_query(
         """
-        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query) 
+        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
         YIELD node AS n, score
         RETURN
             n.uuid AS uuid,
@@ -407,7 +407,7 @@ async def community_fulltext_search(
 
     records, _, _ = await driver.execute_query(
         """
-        CALL db.index.fulltext.queryNodes("community_name", $query) 
+        CALL db.index.fulltext.queryNodes("community_name", $query, {limit: $limit}) 
         YIELD node AS comm, score
         RETURN
             comm.uuid AS uuid,
@@ -539,8 +539,8 @@ async def hybrid_node_search(
 
 
 async def get_relevant_nodes(
-    nodes: list[EntityNode],
     driver: AsyncDriver,
+    nodes: list[EntityNode],
 ) -> list[EntityNode]:
     """
     Retrieve relevant nodes based on the provided list of EntityNodes.
@@ -573,6 +573,7 @@ async def get_relevant_nodes(
         driver,
         [node.group_id for node in nodes],
     )
+
     return relevant_nodes
 
 

graphiti_core/utils/bulk_utils.py

@@ -18,7 +18,7 @@ import asyncio
 import logging
 import typing
 from collections import defaultdict
-from datetime import datetime
+from datetime import datetime, timezone
 from math import ceil
 
 from neo4j import AsyncDriver, AsyncManagedTransaction
@@ -169,7 +169,7 @@ async def dedupe_nodes_bulk(
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
         await asyncio.gather(
-            *[get_relevant_nodes(node_chunk, driver) for node_chunk in node_chunks]
+            *[get_relevant_nodes(driver, node_chunk) for node_chunk in node_chunks]
         )
     )
 
@@ -385,7 +385,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()
+            edge.expired_at = datetime.now(timezone.utc)
 
     return edges