graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/prompts/extract_nodes.py

@@ -14,12 +14,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
 from typing import Any, Protocol, TypedDict
 
 from pydantic import BaseModel, Field
 
+from graphiti_core.utils.text_utils import MAX_SUMMARY_CHARS
+
 from .models import Message, PromptFunction, PromptVersion
+from .prompt_helpers import to_prompt_json
+from .snippets import summary_instructions
 
 
 class ExtractedEntity(BaseModel):
@@ -42,7 +45,8 @@ class EntityClassificationTriple(BaseModel):
     uuid: str = Field(description='UUID of the entity')
     name: str = Field(description='Name of the entity')
     entity_type: str | None = Field(
-        default=None, description='Type of the entity. Must be one of the provided types or None'
+        default=None,
+        description='Type of the entity. Must be one of the provided types or None',
     )
 
 
@@ -52,6 +56,13 @@ class EntityClassification(BaseModel):
     )
 
 
+class EntitySummary(BaseModel):
+    summary: str = Field(
+        ...,
+        description=f'Summary containing the important information about the entity. Under {MAX_SUMMARY_CHARS} characters.',
+    )
+
+
 class Prompt(Protocol):
     extract_message: PromptVersion
     extract_json: PromptVersion
@@ -59,6 +70,7 @@ class Prompt(Protocol):
     reflexion: PromptVersion
     classify_nodes: PromptVersion
     extract_attributes: PromptVersion
+    extract_summary: PromptVersion
 
 
 class Versions(TypedDict):
@@ -68,6 +80,7 @@ class Versions(TypedDict):
     reflexion: PromptFunction
     classify_nodes: PromptFunction
     extract_attributes: PromptFunction
+    extract_summary: PromptFunction
 
 
 def extract_message(context: dict[str, Any]) -> list[Message]:
@@ -75,21 +88,23 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
     Your primary task is to extract and classify the speaker and other significant entities mentioned in the conversation."""
 
     user_prompt = f"""
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
+
 <PREVIOUS MESSAGES>
-{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS MESSAGES>
 
 <CURRENT MESSAGE>
 {context['episode_content']}
 </CURRENT MESSAGE>
 
-<ENTITY TYPES>
-{context['entity_types']}
-</ENTITY TYPES>
-
 Instructions:
 
 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.
+Pronoun references such as he/she/they or this/that/those should be disambiguated to the names of the 
+reference entities. Only extract distinct entities from the CURRENT MESSAGE. Don't extract pronouns like you, me, he/she/they, we/us as entities.
 
 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**.
@@ -122,15 +137,16 @@ def extract_json(context: dict[str, Any]) -> list[Message]:
     Your primary task is to extract and classify relevant entities from JSON files"""
 
     user_prompt = f"""
+<ENTITY TYPES>
+{context['entity_types']}
+</ENTITY TYPES>
+
 <SOURCE DESCRIPTION>:
 {context['source_description']}
 </SOURCE DESCRIPTION>
 <JSON>
 {context['episode_content']}
 </JSON>
-<ENTITY TYPES>
-{context['entity_types']}
-</ENTITY TYPES>
 
 {context['custom_prompt']}
 
@@ -139,8 +155,9 @@ For each entity extracted, also determine its entity type based on the provided
 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
-2. Do NOT extract any properties that contain dates
+1. Extract all entities that the JSON represents. This will often be something like a "name" or "user" field
+2. Extract all entities mentioned in all other properties throughout the JSON structure
+3. Do NOT extract any properties that contain dates
 """
     return [
         Message(role='system', content=sys_prompt),
@@ -153,13 +170,14 @@ def extract_text(context: dict[str, Any]) -> list[Message]:
     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>
 
+<TEXT>
+{context['episode_content']}
+</TEXT>
+
 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.
@@ -183,7 +201,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
 
     user_prompt = f"""
 <PREVIOUS MESSAGES>
-{json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS MESSAGES>
 <CURRENT MESSAGE>
 {context['episode_content']}
@@ -207,22 +225,22 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
 
     user_prompt = f"""
     <PREVIOUS MESSAGES>
-    {json.dumps([ep for ep in context['previous_episodes']], indent=2)}
+    {to_prompt_json([ep for ep in context['previous_episodes']])}
     </PREVIOUS MESSAGES>
     <CURRENT MESSAGE>
     {context['episode_content']}
     </CURRENT MESSAGE>
-    
+
     <EXTRACTED ENTITIES>
     {context['extracted_entities']}
     </EXTRACTED ENTITIES>
-    
+
     <ENTITY TYPES>
     {context['entity_types']}
     </ENTITY TYPES>
-    
+
     Given the above conversation, extracted entities, and provided entity types and their descriptions, classify the extracted entities.
-    
+
     Guidelines:
     1. Each entity must have exactly one type
     2. Only use the provided ENTITY TYPES as types, do not use additional types to classify entities.
@@ -243,21 +261,45 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
         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
+        Given the 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 250 words.
-        
+
+        <MESSAGES>
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
+        </MESSAGES>
+
+        <ENTITY>
+        {context['node']}
+        </ENTITY>
+        """,
+        ),
+    ]
+
+
+def extract_summary(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that extracts entity summaries from the provided text.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        Given the MESSAGES and the ENTITY, update the summary that combines relevant information about the entity
+        from the messages and relevant information from the existing summary.
+
+        {summary_instructions}
+
+        <MESSAGES>
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
+        </MESSAGES>
+
         <ENTITY>
         {context['node']}
         </ENTITY>
@@ -271,6 +313,7 @@ versions: Versions = {
     'extract_json': extract_json,
     'extract_text': extract_text,
     'reflexion': reflexion,
+    'extract_summary': extract_summary,
     'classify_nodes': classify_nodes,
     'extract_attributes': extract_attributes,
 }

graphiti_core/prompts/prompt_helpers.py

@@ -1 +1,40 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import json
+from typing import Any
+
 DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n'
+
+
+def to_prompt_json(data: Any, ensure_ascii: bool = False, indent: int | None = None) -> str:
+    """
+    Serialize data to JSON for use in prompts.
+
+    Args:
+        data: The data to serialize
+        ensure_ascii: If True, escape non-ASCII characters. If False (default), preserve them.
+        indent: Number of spaces for indentation. Defaults to None (minified).
+
+    Returns:
+        JSON string representation of the data
+
+    Notes:
+        By default (ensure_ascii=False), non-ASCII characters (e.g., Korean, Japanese, Chinese)
+        are preserved in their original form in the prompt, making them readable
+        in LLM logs and improving model understanding.
+    """
+    return json.dumps(data, ensure_ascii=ensure_ascii, indent=indent)

graphiti_core/prompts/snippets.py

@@ -0,0 +1,29 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+summary_instructions = """Guidelines:
+        1. Output only factual content. Never explain what you're doing, why, or mention limitations/constraints. 
+        2. Only use the provided messages, entity, and entity context to set attribute values.
+        3. Keep the summary concise and to the point. STATE FACTS DIRECTLY IN UNDER 250 CHARACTERS.
+
+        Example summaries:
+        BAD: "This is the only activity in the context. The user listened to this song. No other details were provided to include in this summary."
+        GOOD: "User played 'Blue Monday' by New Order (electronic genre) on 2024-12-03 at 14:22 UTC."
+        BAD: "Based on the messages provided, the user attended a meeting. This summary focuses on that event as it was the main topic discussed."
+        GOOD: "User attended Q3 planning meeting with sales team on March 15."
+        BAD: "The context shows John ordered pizza. Due to length constraints, other details are omitted from this summary."
+        GOOD: "John ordered pepperoni pizza from Mario's at 7:30 PM, delivered to office."
+        """

graphiti_core/prompts/summarize_nodes.py

@@ -14,18 +14,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
 from typing import Any, Protocol, TypedDict
 
 from pydantic import BaseModel, Field
 
 from .models import Message, PromptFunction, PromptVersion
+from .prompt_helpers import to_prompt_json
+from .snippets import summary_instructions
 
 
 class Summary(BaseModel):
     summary: str = Field(
         ...,
-        description='Summary containing the important information about the entity. Under 250 words',
+        description='Summary containing the important information about the entity. Under 250 characters',
     )
 
 
@@ -55,11 +56,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 250 words.
+
+        IMPORTANT: Keep the summary concise and to the point. SUMMARIES MUST BE LESS THAN 250 CHARACTERS.
 
         Summaries:
-        {json.dumps(context['node_summaries'], indent=2)}
+        {to_prompt_json(context['node_summaries'])}
         """,
         ),
     ]
@@ -69,38 +70,35 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
             role='system',
-            content='You are a helpful assistant that extracts entity properties from the provided text.',
+            content='You are a helpful assistant that generates a summary and attributes from 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 name, create a summary for the ENTITY. Your summary must only use
+        Given the MESSAGES and the 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. Summaries must be under 250 words.
-        
+        provided ENTITY.
+
         In addition, extract any values for the provided entity properties based on their descriptions.
         If the value of the entity property cannot be found in the current context, set the value of the property to the Python value None.
-        
-        Guidelines:
-        1. Do not hallucinate entity property values if they cannot be found in the current context.
-        2. Only use the provided messages, entity, and entity context to set attribute values.
-        
+
+        {summary_instructions}
+
+        <MESSAGES>
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
+        </MESSAGES>
+
         <ENTITY>
         {context['node_name']}
         </ENTITY>
-        
+
         <ENTITY CONTEXT>
         {context['node_summary']}
         </ENTITY CONTEXT>
-        
+
         <ATTRIBUTES>
-        {json.dumps(context['attributes'], indent=2)}
+        {to_prompt_json(context['attributes'])}
         </ATTRIBUTES>
         """,
         ),
@@ -117,10 +115,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 250 words.
+        Summaries must be under 250 characters.
 
         Summary:
-        {json.dumps(context['summary'], indent=2)}
+        {to_prompt_json(context['summary'])}
         """,
         ),
     ]