graphiti-core 0.11.6rc7__tar.gz → 0.12.0rc1__tar.gz

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.11.6rc7
+Version: 0.12.0rc1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/edges.py

@@ -49,7 +49,9 @@ ENTITY_EDGE_RETURN: LiteralString = """
             e.episodes AS episodes,
             e.expired_at AS expired_at,
             e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at"""
+            e.invalid_at AS invalid_at,
+            properties(e) AS attributes
+            """
 
 
 class Edge(BaseModel, ABC):
@@ -209,6 +211,9 @@ class EntityEdge(Edge):
     invalid_at: datetime | None = Field(
         default=None, description='datetime of when the fact stopped being true'
     )
+    attributes: dict[str, Any] = Field(
+        default={}, description='Additional attributes of the edge. Dependent on edge name'
+    )
 
     async def generate_embedding(self, embedder: EmbedderClient):
         start = time()
@@ -236,20 +241,26 @@ class EntityEdge(Edge):
         self.fact_embedding = records[0]['fact_embedding']
 
     async def save(self, driver: AsyncDriver):
+        edge_data: dict[str, Any] = {
+            'source_uuid': self.source_node_uuid,
+            'target_uuid': self.target_node_uuid,
+            'uuid': self.uuid,
+            'name': self.name,
+            'group_id': self.group_id,
+            'fact': self.fact,
+            'fact_embedding': self.fact_embedding,
+            'episodes': self.episodes,
+            'created_at': self.created_at,
+            'expired_at': self.expired_at,
+            'valid_at': self.valid_at,
+            'invalid_at': self.invalid_at,
+        }
+
+        edge_data.update(self.attributes or {})
+
         result = await driver.execute_query(
             ENTITY_EDGE_SAVE,
-            source_uuid=self.source_node_uuid,
-            target_uuid=self.target_node_uuid,
-            uuid=self.uuid,
-            name=self.name,
-            group_id=self.group_id,
-            fact=self.fact,
-            fact_embedding=self.fact_embedding,
-            episodes=self.episodes,
-            created_at=self.created_at,
-            expired_at=self.expired_at,
-            valid_at=self.valid_at,
-            invalid_at=self.invalid_at,
+            edge_data=edge_data,
             database_=DEFAULT_DATABASE,
         )
 
@@ -334,8 +345,8 @@ class EntityEdge(Edge):
     async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str):
         query: LiteralString = (
             """
-                                        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
-                                        """
+                                            MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+                                            """
             + ENTITY_EDGE_RETURN
         )
         records, _, _ = await driver.execute_query(
@@ -457,7 +468,7 @@ def get_episodic_edge_from_record(record: Any) -> EpisodicEdge:
 
 
 def get_entity_edge_from_record(record: Any) -> EntityEdge:
-    return EntityEdge(
+    edge = EntityEdge(
         uuid=record['uuid'],
         source_node_uuid=record['source_node_uuid'],
         target_node_uuid=record['target_node_uuid'],
@@ -469,8 +480,23 @@ def get_entity_edge_from_record(record: Any) -> EntityEdge:
         expired_at=parse_db_date(record['expired_at']),
         valid_at=parse_db_date(record['valid_at']),
         invalid_at=parse_db_date(record['invalid_at']),
+        attributes=record['attributes'],
     )
 
+    edge.attributes.pop('uuid', None)
+    edge.attributes.pop('source_node_uuid', None)
+    edge.attributes.pop('target_node_uuid', None)
+    edge.attributes.pop('fact', None)
+    edge.attributes.pop('name', None)
+    edge.attributes.pop('group_id', None)
+    edge.attributes.pop('episodes', None)
+    edge.attributes.pop('created_at', None)
+    edge.attributes.pop('expired_at', None)
+    edge.attributes.pop('valid_at', None)
+    edge.attributes.pop('invalid_at', None)
+
+    return edge
+
 
 def get_community_edge_from_record(record: Any):
     return CommunityEdge(

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/embedder/gemini.py

@@ -61,18 +61,29 @@ class GeminiEmbedder(EmbedderClient):
         # Generate embeddings
         result = await self.client.aio.models.embed_content(
             model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
-            contents=[input_data],
+            contents=[input_data],  # type: ignore[arg-type]  # mypy fails on broad union type
             config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
         )
 
+        if not result.embeddings or len(result.embeddings) == 0 or not result.embeddings[0].values:
+            raise ValueError('No embeddings returned from Gemini API in create()')
+
         return result.embeddings[0].values
 
     async def create_batch(self, input_data_list: list[str]) -> list[list[float]]:
         # Generate embeddings
         result = await self.client.aio.models.embed_content(
             model=self.config.embedding_model or DEFAULT_EMBEDDING_MODEL,
-            contents=input_data_list,
+            contents=input_data_list,  # type: ignore[arg-type]  # mypy fails on broad union type
             config=types.EmbedContentConfig(output_dimensionality=self.config.embedding_dim),
         )
 
-        return [embedding.values for embedding in result.embeddings]
+        if not result.embeddings or len(result.embeddings) == 0:
+            raise Exception('No embeddings returned')
+
+        embeddings = []
+        for embedding in result.embeddings:
+            if not embedding.values:
+                raise ValueError('Empty embedding values returned')
+            embeddings.append(embedding.values)
+        return embeddings

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/graphiti.py

@@ -41,6 +41,7 @@ from graphiti_core.search.search_config_recipes import (
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
+    get_edge_invalidation_candidates,
     get_mentioned_nodes,
     get_relevant_edges,
 )
@@ -62,9 +63,8 @@ from graphiti_core.utils.maintenance.community_operations import (
 )
 from graphiti_core.utils.maintenance.edge_operations import (
     build_episodic_edges,
-    dedupe_extracted_edge,
     extract_edges,
-    resolve_edge_contradictions,
+    resolve_extracted_edge,
     resolve_extracted_edges,
 )
 from graphiti_core.utils.maintenance.graph_data_operations import (
@@ -77,7 +77,6 @@ from graphiti_core.utils.maintenance.node_operations import (
     extract_nodes,
     resolve_extracted_nodes,
 )
-from graphiti_core.utils.maintenance.temporal_operations import get_edge_contradictions
 from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types
 
 logger = logging.getLogger(__name__)
@@ -274,6 +273,8 @@ class Graphiti:
         update_communities: bool = False,
         entity_types: dict[str, BaseModel] | None = None,
         previous_episode_uuids: list[str] | None = None,
+        edge_types: dict[str, BaseModel] | None = None,
+        edge_type_map: dict[tuple[str, str], list[str]] | None = None,
     ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
@@ -356,6 +357,13 @@ class Graphiti:
                 )
             )
 
+            # Create default edge type map
+            edge_type_map_default = (
+                {('Entity', 'Entity'): list(edge_types.keys())}
+                if edge_types is not None
+                else {('Entity', 'Entity'): []}
+            )
+
             # Extract entities as nodes
 
             extracted_nodes = await extract_nodes(
@@ -371,7 +379,9 @@ class Graphiti:
                     previous_episodes,
                     entity_types,
                 ),
-                extract_edges(self.clients, episode, extracted_nodes, previous_episodes, group_id),
+                extract_edges(
+                    self.clients, episode, extracted_nodes, previous_episodes, group_id, edge_types
+                ),
             )
 
             edges = resolve_edge_pointers(extracted_edges, uuid_map)
@@ -381,6 +391,9 @@ class Graphiti:
                     self.clients,
                     edges,
                     episode,
+                    nodes,
+                    edge_types or {},
+                    edge_type_map or edge_type_map_default,
                 ),
                 extract_attributes_from_nodes(
                     self.clients, nodes, episode, previous_episodes, entity_types
@@ -681,17 +694,27 @@ class Graphiti:
 
         updated_edge = resolve_edge_pointers([edge], uuid_map)[0]
 
-        related_edges = await get_relevant_edges(self.driver, [updated_edge], SearchFilters(), 0.8)
+        related_edges = (await get_relevant_edges(self.driver, [updated_edge], SearchFilters()))[0]
+        existing_edges = (
+            await get_edge_invalidation_candidates(self.driver, [updated_edge], SearchFilters())
+        )[0]
 
-        resolved_edge = await dedupe_extracted_edge(
+        resolved_edge, invalidated_edges = await resolve_extracted_edge(
             self.llm_client,
             updated_edge,
-            related_edges[0],
+            related_edges,
+            existing_edges,
+            EpisodicNode(
+                name='',
+                source=EpisodeType.text,
+                source_description='',
+                content='',
+                valid_at=edge.valid_at or utc_now(),
+                entity_edges=[],
+                group_id=edge.group_id,
+            ),
         )
 
-        contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges[0])
-        invalidated_edges = resolve_edge_contradictions(resolved_edge, contradicting_edges)
-
         await add_nodes_and_edges_bulk(
             self.driver, [], [], resolved_nodes, [resolved_edge] + invalidated_edges, self.embedder
         )

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/helpers.py

@@ -18,11 +18,11 @@ import asyncio
 import os
 from collections.abc import Coroutine
 from datetime import datetime
-from typing import Any
 
 import numpy as np
 from dotenv import load_dotenv
 from neo4j import time as neo4j_time
+from numpy._typing import NDArray
 from typing_extensions import LiteralString
 
 load_dotenv()
@@ -79,16 +79,10 @@ def lucene_sanitize(query: str) -> str:
     return sanitized
 
 
-def normalize_l2(embedding: list[float]):
+def normalize_l2(embedding: list[float]) -> NDArray:
     embedding_array = np.array(embedding)
-    if embedding_array.ndim == 1:
-        norm = np.linalg.norm(embedding_array)
-        if norm == 0:
-            return [0.0] * len(embedding)
-        return (embedding_array / norm).tolist()
-    else:
-        norm = np.linalg.norm(embedding_array, 2, axis=1, keepdims=True)
-        return (np.where(norm == 0, embedding_array, embedding_array / norm)).tolist()
+    norm = np.linalg.norm(embedding_array, 2, axis=0, keepdims=True)
+    return np.where(norm == 0, embedding_array, embedding_array / norm)
 
 
 # Use this instead of asyncio.gather() to bound coroutines
@@ -98,21 +92,8 @@ async def semaphore_gather(
 ):
     semaphore = asyncio.Semaphore(max_coroutines)
 
-    async def _wrap(coro: Coroutine) -> Any:
+    async def _wrap_coroutine(coroutine):
         async with semaphore:
-            return await coro
-
-    results = []
-    batch = []
-    for coroutine in coroutines:
-        batch.append(_wrap(coroutine))
-        # once we hit max_coroutines, gather and clear the batch
-        if len(batch) >= max_coroutines:
-            results.extend(await asyncio.gather(*batch))
-            batch.clear()
-
-    # gather any remaining coroutines in the final batch
-    if batch:
-        results.extend(await asyncio.gather(*batch))
-
-    return results
+            return await coroutine
+
+    return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines))

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/llm_client/gemini_client.py

@@ -139,13 +139,16 @@ class GeminiClient(LLMClient):
             # Generate content using the simple string approach
             response = await self.client.aio.models.generate_content(
                 model=self.model or DEFAULT_MODEL,
-                contents=gemini_messages,
+                contents=gemini_messages,  # type: ignore[arg-type]  # mypy fails on broad union type
                 config=generation_config,
             )
 
             # If this was a structured output request, parse the response into the Pydantic model
             if response_model is not None:
                 try:
+                    if not response.text:
+                        raise ValueError('No response text')
+
                     validated_model = response_model.model_validate(json.loads(response.text))
 
                     # Return as a dictionary for API consistency

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/models/edges/edge_db_queries.py

@@ -34,8 +34,7 @@ ENTITY_EDGE_SAVE = """
         MATCH (source:Entity {uuid: $source_uuid}) 
         MATCH (target:Entity {uuid: $target_uuid}) 
         MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target)
-        SET r = {uuid: $uuid, name: $name, group_id: $group_id, fact: $fact, episodes: $episodes, 
-        created_at: $created_at, expired_at: $expired_at, valid_at: $valid_at, invalid_at: $invalid_at}
+        SET r = $edge_data
         WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
         RETURN r.uuid AS uuid"""
 
@@ -44,8 +43,7 @@ ENTITY_EDGE_SAVE_BULK = """
     MATCH (source:Entity {uuid: edge.source_node_uuid}) 
     MATCH (target:Entity {uuid: edge.target_node_uuid}) 
     MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
-    SET r = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes, 
-    created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at}
+    SET r = edge
     WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
     RETURN edge.uuid AS uuid
 """

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/prompts/dedupe_edges.py

@@ -27,6 +27,11 @@ class EdgeDuplicate(BaseModel):
         ...,
         description='id of the duplicate fact. If no duplicate facts are found, default to -1.',
     )
+    contradicted_facts: list[int] = Field(
+        ...,
+        description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
+    )
+    fact_type: str = Field(..., description='One of the provided fact types or DEFAULT')
 
 
 class UniqueFact(BaseModel):
@@ -41,11 +46,13 @@ class UniqueFacts(BaseModel):
 class Prompt(Protocol):
     edge: PromptVersion
     edge_list: PromptVersion
+    resolve_edge: PromptVersion
 
 
 class Versions(TypedDict):
     edge: PromptFunction
     edge_list: PromptFunction
+    resolve_edge: PromptFunction
 
 
 def edge(context: dict[str, Any]) -> list[Message]:
@@ -106,4 +113,48 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'edge': edge, 'edge_list': edge_list}
+def resolve_edge(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that de-duplicates facts from fact lists and determines which existing '
+            'facts are contradicted by the new fact.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        <NEW FACT>
+        {context['new_edge']}
+        </NEW FACT>
+        
+        <EXISTING FACTS>
+        {context['existing_edges']}
+        </EXISTING FACTS>
+        <FACT INVALIDATION CANDIDATES>
+        {context['edge_invalidation_candidates']}
+        </FACT INVALIDATION CANDIDATES>
+        
+        <FACT TYPES>
+        {context['edge_types']}
+        </FACT TYPES>
+        
+
+        Task:
+        If the NEW FACT represents the same factual information as any fact in EXISTING FACTS, return the idx of the duplicate fact.
+        If the NEW FACT is not a duplicate of any of the EXISTING FACTS, return -1.
+        
+        Given the predefined FACT TYPES, determine if the NEW FACT should be classified as one of these types.
+        Return the fact type as fact_type or DEFAULT if NEW FACT is not one of the FACT TYPES.
+        
+        Based on the provided FACT INVALIDATION CANDIDATES and NEW FACT, determine which existing facts the new fact contradicts.
+        Return a list containing all idx's of the facts that are contradicted by the NEW FACT.
+        If there are no contradicted facts, return an empty list.
+
+        Guidelines:
+        1. The facts do not need to be completely identical to be duplicates, they just need to express the same information.
+        """,
+        ),
+    ]
+
+
+versions: Versions = {'edge': edge, 'edge_list': edge_list, 'resolve_edge': resolve_edge}

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/prompts/dedupe_nodes.py

@@ -23,21 +23,31 @@ from .models import Message, PromptFunction, PromptVersion
 
 
 class NodeDuplicate(BaseModel):
-    duplicate_node_id: int = Field(
+    id: int = Field(..., description='integer id of the entity')
+    duplicate_idx: int = Field(
         ...,
-        description='id of the duplicate node. If no duplicate nodes are found, default to -1.',
+        description='idx of the duplicate node. If no duplicate nodes are found, default to -1.',
     )
-    name: str = Field(..., description='Name of the entity.')
+    name: str = Field(
+        ...,
+        description='Name of the entity. Should be the most complete and descriptive name possible.',
+    )
+
+
+class NodeResolutions(BaseModel):
+    entity_resolutions: list[NodeDuplicate] = Field(..., description='List of resolved nodes')
 
 
 class Prompt(Protocol):
     node: PromptVersion
     node_list: PromptVersion
+    nodes: PromptVersion
 
 
 class Versions(TypedDict):
     node: PromptFunction
     node_list: PromptFunction
+    nodes: PromptFunction
 
 
 def node(context: dict[str, Any]) -> list[Message]:
@@ -89,6 +99,67 @@ def node(context: dict[str, Any]) -> list[Message]:
     ]
 
 
+def nodes(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that determines whether or not ENTITIES extracted from a conversation are duplicates'
+            'of existing entities.',
+        ),
+        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>
+        
+        
+        Each of the following ENTITIES were extracted from the CURRENT MESSAGE.
+        Each entity in ENTITIES is represented as a JSON object with the following structure:
+        {{
+            id: integer id of the entity,
+            name: "name of the entity",
+            entity_type: "ontological classification of the entity",
+            entity_type_description: "Description of what the entity type represents",
+            duplication_candidates: [
+                {{
+                    idx: integer index of the candidate entity,
+                    name: "name of the candidate entity",
+                    entity_type: "ontological classification of the candidate entity",
+                    ...<additional attributes>
+                }}
+            ]
+        }}
+        
+        <ENTITIES>
+        {json.dumps(context['extracted_nodes'], indent=2)}
+        </ENTITIES>
+
+        For each of the above ENTITIES, determine if the entity is a duplicate of any of its duplication candidates.
+
+        Entities should only be considered duplicates if they refer to the *same real-world object or concept*.
+
+        Do NOT mark entities as duplicates if:
+        - They are related but distinct.
+        - They have similar names or purposes but refer to separate instances or concepts.
+
+        Task:
+        Your response will be a list called entity_resolutions which contains one entry for each entity.
+        
+        For each entity, return the id of the entity as id, the name of the entity as name, and the duplicate_idx
+        as an integer.
+        
+        - If an entity is a duplicate of one of its duplication_candidates, return the idx of the candidate it is a 
+        duplicate of.
+        - If an entity is not a duplicate of one of its duplication candidates, return the -1 as the duplication_idx
+        """,
+        ),
+    ]
+
+
 def node_list(context: dict[str, Any]) -> list[Message]:
     return [
         Message(
@@ -126,4 +197,4 @@ def node_list(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'node': node, 'node_list': node_list}
+versions: Versions = {'node': node, 'node_list': node_list, 'nodes': nodes}

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/prompts/extract_edges.py

@@ -48,11 +48,13 @@ class MissingFacts(BaseModel):
 class Prompt(Protocol):
     edge: PromptVersion
     reflexion: PromptVersion
+    extract_attributes: PromptVersion
 
 
 class Versions(TypedDict):
     edge: PromptFunction
     reflexion: PromptFunction
+    extract_attributes: PromptFunction
 
 
 def edge(context: dict[str, Any]) -> list[Message]:
@@ -82,12 +84,18 @@ def edge(context: dict[str, Any]) -> list[Message]:
 {context['reference_time']}  # ISO 8601 (UTC); used to resolve relative time mentions
 </REFERENCE_TIME>
 
+<FACT TYPES>
+{context['edge_types']}
+</FACT TYPES>
+
 # 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.
+    and can be represented as edges in a knowledge graph.
+- The FACT TYPES provide a list of the most important types of facts, make sure to extract any facts that
+    could be classified into one of the provided fact types
 
 You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
 
@@ -145,4 +153,40 @@ determine if any facts haven't been extracted.
     ]
 
 
-versions: Versions = {'edge': edge, 'reflexion': reflexion}
+def extract_attributes(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that extracts fact properties from the provided text.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+
+        <MESSAGE>
+        {json.dumps(context['episode_content'], indent=2)}
+        </MESSAGE>
+        <REFERENCE TIME>
+        {context['reference_time']}
+        </REFERENCE TIME>
+
+        Given the above MESSAGE, its REFERENCE TIME, and the following FACT, update any of its attributes based on the information provided
+        in MESSAGE. 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 FACT to set attribute values.
+
+        <FACT>
+        {context['fact']}
+        </FACT>
+        """,
+        ),
+    ]
+
+
+versions: Versions = {
+    'edge': edge,
+    'reflexion': reflexion,
+    'extract_attributes': extract_attributes,
+}

{graphiti_core-0.11.6rc7 → graphiti_core-0.12.0rc1}/graphiti_core/prompts/invalidate_edges.py

@@ -24,7 +24,7 @@ from .models import Message, PromptFunction, PromptVersion
 class InvalidatedEdges(BaseModel):
     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.',
+        description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
     )