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

graphiti_core/prompts/eval.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 QueryExpansion(BaseModel):
+    query: str = Field(..., description='query optimized for database search')
+
+
+class QAResponse(BaseModel):
+    ANSWER: str = Field(..., description='how Alice would answer the question')
+
+
+class EvalResponse(BaseModel):
+    is_correct: bool = Field(..., description='boolean if the answer is correct or incorrect')
+    reasoning: str = Field(
+        ..., description='why you determined the response was correct or incorrect'
+    )
+
+
 class Prompt(Protocol):
     qa_prompt: PromptVersion
     eval_prompt: PromptVersion
@@ -41,10 +58,6 @@ def query_expansion(context: dict[str, Any]) -> list[Message]:
     <QUESTION>
     {json.dumps(context['query'])}
     </QUESTION>
-    respond with a JSON object in the following format:
-    {{
-        "query": "query optimized for database search"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),
@@ -67,10 +80,6 @@ def qa_prompt(context: dict[str, Any]) -> list[Message]:
     <QUESTION>
     {context['query']}
     </QUESTION>
-    respond with a JSON object in the following format:
-    {{
-        "ANSWER": "how Alice would answer the question"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),
@@ -96,12 +105,6 @@ def eval_prompt(context: dict[str, Any]) -> list[Message]:
     <RESPONSE>
     {context['response']}
     </RESPONSE>
-    
-    respond with a JSON object in the following format:
-    {{
-        "is_correct": "boolean if the answer is correct or incorrect"
-        "reasoning": "why you determined the response was correct or incorrect"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/prompts/extract_edge_dates.py

@@ -14,11 +14,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Any, Protocol, TypedDict
+from typing import Any, Optional, Protocol, TypedDict
+
+from pydantic import BaseModel, Field
 
 from .models import Message, PromptFunction, PromptVersion
 
 
+class EdgeDates(BaseModel):
+    valid_at: Optional[str] = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact became true or was established. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.',
+    )
+    invalid_at: Optional[str] = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact stopped being true or ended. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.',
+    )
+
+
 class Prompt(Protocol):
     v1: PromptVersion
 
@@ -60,7 +73,7 @@ def v1(context: dict[str, Any]) -> list[Message]:
             Analyze the conversation and determine if there are dates that are part of the edge fact. Only set dates if they explicitly relate to the formation or alteration of the relationship itself.
 
             Guidelines:
-            1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SSZ) for datetimes.
+            1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SS.SSSSSSZ) for datetimes.
             2. Use the reference timestamp as the current time when determining the valid_at and invalid_at dates.
             3. If the fact is written in the present tense, use the Reference Timestamp for the valid_at date
             4. If no temporal information is found that establishes or changes the relationship, leave the fields as null.
@@ -69,11 +82,6 @@ def v1(context: dict[str, Any]) -> list[Message]:
             7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
             8. If only year is mentioned, use January 1st of that year at 00:00:00.
             9. Always include the time zone offset (use Z for UTC if no specific time zone is mentioned).
-            Respond with a JSON object:
-            {{
-                "valid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
-                "invalid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
-            }}
             """,
         ),
     ]

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
@@ -45,11 +57,6 @@ def summarize_pair(context: dict[str, Any]) -> list[Message]:
 
         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"
-            }}
         """,
         ),
     ]
@@ -77,12 +84,6 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         <ENTITY>
         {context['node_name']}
         </ENTITY>
-        
-
-        Respond with a JSON object in the following format:
-            {{
-                "summary": "Entity summary"
-            }}
         """,
         ),
     ]
@@ -101,11 +102,6 @@ def summary_description(context: dict[str, Any]) -> list[Message]:
 
         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_utils.py

@@ -40,7 +40,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

graphiti_core/utils/maintenance/community_operations.py

@@ -16,6 +16,7 @@ from graphiti_core.nodes import (
     get_community_node_from_record,
 )
 from graphiti_core.prompts import prompt_library
+from graphiti_core.prompts.summarize_nodes import Summary, SummaryDescription
 from graphiti_core.utils.maintenance.edge_operations import build_community_edges
 
 MAX_COMMUNITY_BUILD_CONCURRENCY = 10
@@ -131,7 +132,7 @@ async def summarize_pair(llm_client: LLMClient, summary_pair: tuple[str, str]) -
     context = {'node_summaries': [{'summary': summary} for summary in summary_pair]}
 
     llm_response = await llm_client.generate_response(
-        prompt_library.summarize_nodes.summarize_pair(context)
+        prompt_library.summarize_nodes.summarize_pair(context), response_model=Summary
     )
 
     pair_summary = llm_response.get('summary', '')
@@ -143,7 +144,8 @@ async def generate_summary_description(llm_client: LLMClient, summary: str) -> s
     context = {'summary': summary}
 
     llm_response = await llm_client.generate_response(
-        prompt_library.summarize_nodes.summary_description(context)
+        prompt_library.summarize_nodes.summary_description(context),
+        response_model=SummaryDescription,
     )
 
     description = llm_response.get('description', '')

graphiti_core/utils/maintenance/edge_operations.py

@@ -18,13 +18,14 @@ import asyncio
 import logging
 from datetime import datetime, timezone
 from time import time
-from typing import List
 
 from graphiti_core.edges import CommunityEdge, EntityEdge, EpisodicEdge
 from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
 from graphiti_core.prompts import prompt_library
+from graphiti_core.prompts.dedupe_edges import EdgeDuplicate, UniqueFacts
+from graphiti_core.prompts.extract_edges import ExtractedEdges, MissingFacts
 from graphiti_core.utils.maintenance.temporal_operations import (
     extract_edge_dates,
     get_edge_contradictions,
@@ -34,11 +35,11 @@ logger = logging.getLogger(__name__)
 
 
 def build_episodic_edges(
-    entity_nodes: List[EntityNode],
+    entity_nodes: list[EntityNode],
     episode: EpisodicNode,
     created_at: datetime,
-) -> List[EpisodicEdge]:
-    edges: List[EpisodicEdge] = [
+) -> list[EpisodicEdge]:
+    edges: list[EpisodicEdge] = [
         EpisodicEdge(
             source_node_uuid=episode.uuid,
             target_node_uuid=node.uuid,
@@ -52,11 +53,11 @@ def build_episodic_edges(
 
 
 def build_community_edges(
-    entity_nodes: List[EntityNode],
+    entity_nodes: list[EntityNode],
     community_node: CommunityNode,
     created_at: datetime,
-) -> List[CommunityEdge]:
-    edges: List[CommunityEdge] = [
+) -> list[CommunityEdge]:
+    edges: list[CommunityEdge] = [
         CommunityEdge(
             source_node_uuid=community_node.uuid,
             target_node_uuid=node.uuid,
@@ -92,7 +93,7 @@ async def extract_edges(
     reflexion_iterations = 0
     while facts_missed and reflexion_iterations < MAX_REFLEXION_ITERATIONS:
         llm_response = await llm_client.generate_response(
-            prompt_library.extract_edges.edge(context)
+            prompt_library.extract_edges.edge(context), response_model=ExtractedEdges
         )
         edges_data = llm_response.get('edges', [])
 
@@ -101,7 +102,7 @@ async def extract_edges(
         reflexion_iterations += 1
         if reflexion_iterations < MAX_REFLEXION_ITERATIONS:
             reflexion_response = await llm_client.generate_response(
-                prompt_library.extract_edges.reflexion(context)
+                prompt_library.extract_edges.reflexion(context), response_model=MissingFacts
             )
 
             missing_facts = reflexion_response.get('missing_facts', [])
@@ -318,7 +319,9 @@ async def dedupe_extracted_edge(
         'extracted_edges': extracted_edge_context,
     }
 
-    llm_response = await llm_client.generate_response(prompt_library.dedupe_edges.edge(context))
+    llm_response = await llm_client.generate_response(
+        prompt_library.dedupe_edges.edge(context), response_model=EdgeDuplicate
+    )
 
     is_duplicate: bool = llm_response.get('is_duplicate', False)
     uuid: str | None = llm_response.get('uuid', None)
@@ -353,7 +356,7 @@ async def dedupe_edge_list(
     context = {'edges': [{'uuid': edge.uuid, 'fact': edge.fact} for edge in edges]}
 
     llm_response = await llm_client.generate_response(
-        prompt_library.dedupe_edges.edge_list(context)
+        prompt_library.dedupe_edges.edge_list(context), response_model=UniqueFacts
     )
     unique_edges_data = llm_response.get('unique_facts', [])
 

graphiti_core/utils/maintenance/node_operations.py

@@ -23,6 +23,9 @@ from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.prompts import prompt_library
+from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
+from graphiti_core.prompts.extract_nodes import ExtractedNodes, MissedEntities
+from graphiti_core.prompts.summarize_nodes import Summary
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +45,7 @@ async def extract_message_nodes(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.extract_message(context)
+        prompt_library.extract_nodes.extract_message(context), response_model=ExtractedNodes
     )
     extracted_node_names = llm_response.get('extracted_node_names', [])
     return extracted_node_names
@@ -63,7 +66,7 @@ async def extract_text_nodes(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.extract_text(context)
+        prompt_library.extract_nodes.extract_text(context), ExtractedNodes
     )
     extracted_node_names = llm_response.get('extracted_node_names', [])
     return extracted_node_names
@@ -81,7 +84,7 @@ async def extract_json_nodes(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.extract_json(context)
+        prompt_library.extract_nodes.extract_json(context), ExtractedNodes
     )
     extracted_node_names = llm_response.get('extracted_node_names', [])
     return extracted_node_names
@@ -101,7 +104,7 @@ async def extract_nodes_reflexion(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.reflexion(context)
+        prompt_library.extract_nodes.reflexion(context), MissedEntities
     )
     missed_entities = llm_response.get('missed_entities', [])
 
@@ -273,9 +276,12 @@ async def resolve_extracted_node(
     }
 
     llm_response, node_summary_response = await asyncio.gather(
-        llm_client.generate_response(prompt_library.dedupe_nodes.node(context)),
         llm_client.generate_response(
-            prompt_library.summarize_nodes.summarize_context(summary_context)
+            prompt_library.dedupe_nodes.node(context), response_model=NodeDuplicate
+        ),
+        llm_client.generate_response(
+            prompt_library.summarize_nodes.summarize_context(summary_context),
+            response_model=Summary,
         ),
     )
 
@@ -294,7 +300,8 @@ async def resolve_extracted_node(
             summary_response = await llm_client.generate_response(
                 prompt_library.summarize_nodes.summarize_pair(
                     {'node_summaries': [extracted_node.summary, existing_node.summary]}
-                )
+                ),
+                response_model=Summary,
             )
             node = existing_node
             node.name = name

graphiti_core/utils/maintenance/temporal_operations.py

@@ -17,12 +17,13 @@ limitations under the License.
 import logging
 from datetime import datetime
 from time import time
-from typing import List
 
 from graphiti_core.edges import EntityEdge
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import EpisodicNode
 from graphiti_core.prompts import prompt_library
+from graphiti_core.prompts.extract_edge_dates import EdgeDates
+from graphiti_core.prompts.invalidate_edges import InvalidatedEdges
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +32,7 @@ async def extract_edge_dates(
     llm_client: LLMClient,
     edge: EntityEdge,
     current_episode: EpisodicNode,
-    previous_episodes: List[EpisodicNode],
+    previous_episodes: list[EpisodicNode],
 ) -> tuple[datetime | None, datetime | None]:
     context = {
         'edge_fact': edge.fact,
@@ -39,7 +40,9 @@ async def extract_edge_dates(
         'previous_episodes': [ep.content for ep in previous_episodes],
         'reference_timestamp': current_episode.valid_at.isoformat(),
     }
-    llm_response = await llm_client.generate_response(prompt_library.extract_edge_dates.v1(context))
+    llm_response = await llm_client.generate_response(
+        prompt_library.extract_edge_dates.v1(context), response_model=EdgeDates
+    )
 
     valid_at = llm_response.get('valid_at')
     invalid_at = llm_response.get('invalid_at')
@@ -76,7 +79,9 @@ async def get_edge_contradictions(
 
     context = {'new_edge': new_edge_context, 'existing_edges': existing_edge_context}
 
-    llm_response = await llm_client.generate_response(prompt_library.invalidate_edges.v2(context))
+    llm_response = await llm_client.generate_response(
+        prompt_library.invalidate_edges.v2(context), response_model=InvalidatedEdges
+    )
 
     contradicted_edge_data = llm_response.get('invalidated_edges', [])
 

{graphiti_core-0.4.2.dist-info → graphiti_core-0.5.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.4.2
+Version: 0.5.0rc1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk