graphiti-core 0.10.5__py3-none-any.whl → 0.11.0__py3-none-any.whl

graphiti_core/prompts/extract_edges.py

@@ -23,10 +23,18 @@ 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')
+    relation_type: str = Field(..., description='FACT_PREDICATE_IN_SCREAMING_SNAKE_CASE')
+    source_entity_name: str = Field(..., description='The name of the source entity of the fact.')
+    target_entity_name: str = Field(..., description='The name of the target entity of the fact.')
+    fact: str = Field(..., description='')
+    valid_at: str | None = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact became true or was established. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SS.SSSSSSZ)',
+    )
+    invalid_at: str | None = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact stopped being true or ended. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SS.SSSSSSZ)',
+    )
 
 
 class ExtractedEdges(BaseModel):
@@ -51,32 +59,59 @@ def edge(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
             role='system',
-            content='You are an expert fact extractor that extracts fact triples from text.',
+            content='You are an expert fact extractor that extracts fact triples from text. '
+            '1. Extracted fact triples should also be extracted with relevant date information.'
+            '2. Treat the CURRENT TIME as the time the CURRENT MESSAGE was sent. All temporal information should be extracted relative to this time.',
         ),
         Message(
             role='user',
             content=f"""
-        <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']}
-
-        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.
+<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']}  # Each has: id, label (e.g., Person, Org), name, aliases
+</ENTITIES>
+
+<REFERENCE_TIME>
+{context['reference_time']}  # ISO 8601 (UTC); used to resolve relative time mentions
+</REFERENCE_TIME>
+
+# TASK
+Extract all factual relationships between the given ENTITIES based on the CURRENT MESSAGE.
+Only extract facts that:
+- involve two DISTINCT ENTITIES from the ENTITIES list,
+- are clearly stated or unambiguously implied in the CURRENT MESSAGE,
+- and can be represented as edges in a knowledge graph.
+
+You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
+
+
+{context['custom_prompt']}
+
+# EXTRACTION RULES
+
+1. Only emit facts where both the subject and object match IDs in ENTITIES.
+2. Each fact must involve two **distinct** entities.
+3. Use a SCREAMING_SNAKE_CASE string as the `relation_type` (e.g., FOUNDED, WORKS_AT).
+4. Do not emit duplicate or semantically redundant facts.
+5. The `fact_text` should quote or closely paraphrase the original source sentence(s).
+6. Use `REFERENCE_TIME` to resolve vague or relative temporal expressions (e.g., "last week").
+7. Do **not** hallucinate or infer temporal bounds from unrelated events.
+
+# DATETIME RULES
+
+- Use ISO 8601 with “Z” suffix (UTC) (e.g., 2025-04-30T00:00:00Z).
+- If the fact is ongoing (present tense), set `valid_at` to REFERENCE_TIME.
+- If a change/termination is expressed, set `invalid_at` to the relevant timestamp.
+- Leave both fields `null` if no explicit or resolvable time is stated.
+- If only a date is mentioned (no time), assume 00:00:00.
+- If only a year is mentioned, use January 1st at 00:00:00.
         """,
         ),
     ]

graphiti_core/prompts/extract_nodes.py

@@ -22,8 +22,16 @@ 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 ExtractedEntity(BaseModel):
+    name: str = Field(..., description='Name of the extracted entity')
+    entity_type_id: int = Field(
+        description='ID of the classified entity type. '
+        'Must be one of the provided entity_type_id integers.',
+    )
+
+
+class ExtractedEntities(BaseModel):
+    extracted_entities: list[ExtractedEntity] = Field(..., description='List of extracted entities')
 
 
 class MissedEntities(BaseModel):
@@ -50,6 +58,7 @@ class Prompt(Protocol):
     extract_text: PromptVersion
     reflexion: PromptVersion
     classify_nodes: PromptVersion
+    extract_attributes: PromptVersion
 
 
 class Versions(TypedDict):
@@ -58,31 +67,49 @@ class Versions(TypedDict):
     extract_text: PromptFunction
     reflexion: PromptFunction
     classify_nodes: PromptFunction
+    extract_attributes: PromptFunction
 
 
 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."""
+    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational messages. 
+    Your primary task is to extract and classify the speaker and other significant entities mentioned in the conversation."""
 
     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>
 
-{context['custom_prompt']}
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
 
-Given the above conversation, extract entity nodes from the CURRENT MESSAGE that are explicitly or implicitly mentioned:
+Instructions:
 
-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 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.
-7. Extract preferences as their own nodes
+You are given a conversation context and a CURRENT MESSAGE. Your task is to extract **entity nodes** mentioned **explicitly or implicitly** in the CURRENT MESSAGE.
+
+1. **Speaker Extraction**: Always extract the speaker (the part before the colon `:` in each dialogue line) as the first entity node.
+   - If the speaker is mentioned again in the message, treat both mentions as a **single entity**.
+
+2. **Entity Identification**:
+   - Extract all significant entities, concepts, or actors that are **explicitly or implicitly** mentioned in the CURRENT MESSAGE.
+   - **Exclude** entities mentioned only in the PREVIOUS MESSAGES (they are for context only).
+
+3. **Entity Classification**:
+   - Use the descriptions in ENTITY TYPES to classify each extracted entity.
+   - Assign the appropriate `entity_type_id` for each one.
+
+4. **Exclusions**:
+   - Do NOT extract entities representing relationships or actions.
+   - Do NOT extract dates, times, or other temporal information—these will be handled separately.
+
+5. **Formatting**:
+   - Be **explicit and unambiguous** in naming entities (e.g., use full names when available).
+
+{context['custom_prompt']}
 """
     return [
         Message(role='system', content=sys_prompt),
@@ -92,7 +119,7 @@ Guidelines:
 
 def extract_json(context: dict[str, Any]) -> list[Message]:
     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"""
+    Your primary task is to extract and classify relevant entities from JSON files"""
 
     user_prompt = f"""
 <SOURCE DESCRIPTION>:
@@ -101,10 +128,15 @@ def extract_json(context: dict[str, Any]) -> list[Message]:
 <JSON>
 {context['episode_content']}
 </JSON>
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
 
 {context['custom_prompt']}
 
-Given the above source description and JSON, extract relevant entity nodes from the provided JSON:
+Given the above source description and JSON, extract relevant entities from the provided JSON.
+For each entity extracted, also determine its entity type based on the provided ENTITY TYPES and their descriptions.
+Indicate the classified entity type by providing its entity_type_id.
 
 Guidelines:
 1. Always try to extract an entities that the JSON represents. This will often be something like a "name" or "user field
@@ -117,16 +149,22 @@ Guidelines:
 
 
 def extract_text(context: dict[str, Any]) -> list[Message]:
-    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."""
+    sys_prompt = """You are an AI assistant that extracts entity nodes from text. 
+    Your primary task is to extract and classify the speaker and other significant entities mentioned in the provided text."""
 
     user_prompt = f"""
 <TEXT>
 {context['episode_content']}
 </TEXT>
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
 
-{context['custom_prompt']}
+Given the above text, extract entities from the TEXT that are explicitly or implicitly mentioned.
+For each entity extracted, also determine its entity type based on the provided ENTITY TYPES and their descriptions.
+Indicate the classified entity type by providing its entity_type_id.
 
-Given the above text, extract entity nodes from the TEXT that are explicitly or implicitly mentioned:
+{context['custom_prompt']}
 
 Guidelines:
 1. Extract significant entities, concepts, or actors mentioned in the conversation.
@@ -196,10 +234,43 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def extract_attributes(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that extracts entity properties from the provided text.',
+        ),
+        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, update any of its attributes based on the information provided
+        in MESSAGES. Use the provided attribute descriptions to better understand how each attribute should be determined.
+
+        Guidelines:
+        1. Do not hallucinate entity property values if they cannot be found in the current context.
+        2. Only use the provided MESSAGES and ENTITY to set attribute values.
+        3. The summary attribute represents a summary of the ENTITY, and should be updated with new information about the Entity from the MESSAGES. 
+            Summaries must be no longer than 200 words.
+        
+        <ENTITY>
+        {context['node']}
+        </ENTITY>
+        """,
+        ),
+    ]
+
+
 versions: Versions = {
     'extract_message': extract_message,
     'extract_json': extract_json,
     'extract_text': extract_text,
     'reflexion': reflexion,
     'classify_nodes': classify_nodes,
+    'extract_attributes': extract_attributes,
 }

graphiti_core/prompts/invalidate_edges.py

@@ -21,14 +21,10 @@ 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'
+    contradicted_facts: list[int] = Field(
+        ...,
+        description='List of ids of facts that be should invalidated. If no facts should be invalidated, the list should be empty.',
     )
 
 
@@ -78,18 +74,22 @@ def v2(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
             role='system',
-            content='You are an AI assistant that helps determine which relationships in a knowledge graph should be invalidated based solely on explicit contradictions in newer information.',
+            content='You are an AI assistant that determines which facts contradict each other.',
         ),
         Message(
             role='user',
             content=f"""
-               Based on the provided Existing Edges and a New Edge, determine which existing edges, if any, should be marked as invalidated due to invalidations with the New Edge.
+               Based on the provided EXISTING FACTS and a NEW FACT, determine which existing facts the new fact contradicts.
+               Return a list containing all ids of the facts that are contradicted by the NEW FACT.
+               If there are no contradicted facts, return an empty list.
 
-                Existing Edges:
+                <EXISTING FACTS>
                 {context['existing_edges']}
+                </EXISTING FACTS>
 
-                New Edge:
+                <NEW FACT>
                 {context['new_edge']}
+                </NEW FACT>
             """,
         ),
     ]

graphiti_core/search/search.py

@@ -22,8 +22,8 @@ from neo4j import AsyncDriver
 
 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.graphiti_types import GraphitiClients
 from graphiti_core.helpers import semaphore_gather
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
 from graphiti_core.search.search_config import (
@@ -62,17 +62,21 @@ logger = logging.getLogger(__name__)
 
 
 async def search(
-    driver: AsyncDriver,
-    embedder: EmbedderClient,
-    cross_encoder: CrossEncoderClient,
+    clients: GraphitiClients,
     query: str,
     group_ids: list[str] | None,
     config: SearchConfig,
     search_filter: SearchFilters,
     center_node_uuid: str | None = None,
     bfs_origin_node_uuids: list[str] | None = None,
+    query_vector: list[float] | None = None,
 ) -> SearchResults:
     start = time()
+
+    driver = clients.driver
+    embedder = clients.embedder
+    cross_encoder = clients.cross_encoder
+
     if query.strip() == '':
         return SearchResults(
             edges=[],
@@ -80,7 +84,11 @@ async def search(
             episodes=[],
             communities=[],
         )
-    query_vector = await embedder.create(input_data=[query.replace('\n', ' ')])
+    query_vector = (
+        query_vector
+        if query_vector is not None
+        else await embedder.create(input_data=[query.replace('\n', ' ')])
+    )
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None