graphiti-core 0.3.8__tar.gz → 0.3.11__tar.gz

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.8
+Version: 0.3.11
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -17,6 +17,7 @@ Requires-Dist: numpy (>=1.0.0)
 Requires-Dist: openai (>=1.50.2,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: tenacity (<9.0.0)
+Requires-Dist: voyageai (>=0.2.3,<0.3.0)
 Description-Content-Type: text/markdown
 
 <div align="center">

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/edges.py

@@ -51,7 +51,7 @@ class Edge(BaseModel, ABC):
             uuid=self.uuid,
         )
 
-        logger.info(f'Deleted Edge: {self.uuid}')
+        logger.debug(f'Deleted Edge: {self.uuid}')
 
         return result
 
@@ -83,7 +83,7 @@ class EpisodicEdge(Edge):
             created_at=self.created_at,
         )
 
-        logger.info(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to neo4j: {self.uuid}')
 
         return result
 
@@ -178,7 +178,7 @@ class EntityEdge(Edge):
         self.fact_embedding = await embedder.create(input=[text])
 
         end = time()
-        logger.info(f'embedded {text} in {end - start} ms')
+        logger.debug(f'embedded {text} in {end - start} ms')
 
         return self.fact_embedding
 
@@ -188,9 +188,9 @@ class EntityEdge(Edge):
         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, fact_embedding: $fact_embedding, 
-        episodes: $episodes, created_at: $created_at, expired_at: $expired_at, 
-        valid_at: $valid_at, invalid_at: $invalid_at}
+        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}
+        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
         RETURN r.uuid AS uuid""",
             source_uuid=self.source_node_uuid,
             target_uuid=self.target_node_uuid,
@@ -206,7 +206,7 @@ class EntityEdge(Edge):
             invalid_at=self.invalid_at,
         )
 
-        logger.info(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to neo4j: {self.uuid}')
 
         return result
 
@@ -313,7 +313,7 @@ class CommunityEdge(Edge):
             created_at=self.created_at,
         )
 
-        logger.info(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to neo4j: {self.uuid}')
 
         return result
 

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/errors.py

@@ -35,6 +35,14 @@ class GroupsEdgesNotFoundError(GraphitiError):
         super().__init__(self.message)
 
 
+class GroupsNodesNotFoundError(GraphitiError):
+    """Raised when no nodes are found for a list of group ids."""
+
+    def __init__(self, group_ids: list[str]):
+        self.message = f'no nodes found for group ids {group_ids}'
+        super().__init__(self.message)
+
+
 class NodeNotFoundError(GraphitiError):
     """Raised when a node is not found."""
 

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/graphiti.py

@@ -21,11 +21,12 @@ from time import time
 
 from dotenv import load_dotenv
 from neo4j import AsyncGraphDatabase
+from pydantic import BaseModel
 
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.llm_client import LLMClient, OpenAIClient
-from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
+from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
 from graphiti_core.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults
 from graphiti_core.search.search_config_recipes import (
@@ -77,6 +78,12 @@ logger = logging.getLogger(__name__)
 load_dotenv()
 
 
+class AddEpisodeResults(BaseModel):
+    episode: EpisodicNode
+    nodes: list[EntityNode]
+    edges: list[EntityEdge]
+
+
 class Graphiti:
     def __init__(
         self,
@@ -245,7 +252,7 @@ class Graphiti:
         group_id: str = '',
         uuid: str | None = None,
         update_communities: bool = False,
-    ):
+    ) -> AddEpisodeResults:
         """
         Process an episode and update the graph.
 
@@ -312,13 +319,11 @@ class Graphiti:
                 valid_at=reference_time,
             )
             episode.uuid = uuid if uuid is not None else episode.uuid
-            if not self.store_raw_episode_content:
-                episode.content = ''
 
             # Extract entities as nodes
 
             extracted_nodes = await extract_nodes(self.llm_client, episode, previous_episodes)
-            logger.info(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
+            logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
             # Calculate Embeddings
 
@@ -333,7 +338,7 @@ class Graphiti:
                 )
             )
 
-            logger.info(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
+            logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
             (mentioned_nodes, uuid_map), extracted_edges = await asyncio.gather(
                 resolve_extracted_nodes(self.llm_client, extracted_nodes, existing_nodes_lists),
@@ -341,7 +346,7 @@ class Graphiti:
                     self.llm_client, episode, extracted_nodes, previous_episodes, group_id
                 ),
             )
-            logger.info(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
+            logger.debug(f'Adjusted mentioned nodes: {[(n.name, n.uuid) for n in mentioned_nodes]}')
             nodes = mentioned_nodes
 
             extracted_edges_with_resolved_pointers = resolve_edge_pointers(
@@ -371,10 +376,10 @@ class Graphiti:
                     ]
                 )
             )
-            logger.info(
+            logger.debug(
                 f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_list for e in edges_lst]}'
             )
-            logger.info(
+            logger.debug(
                 f'Extracted edges: {[(e.name, e.uuid) for e in extracted_edges_with_resolved_pointers]}'
             )
 
@@ -426,15 +431,18 @@ class Graphiti:
 
             entity_edges.extend(resolved_edges + invalidated_edges)
 
-            logger.info(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_edges]}')
+            logger.debug(f'Resolved edges: {[(e.name, e.uuid) for e in resolved_edges]}')
 
             episodic_edges: list[EpisodicEdge] = build_episodic_edges(mentioned_nodes, episode, now)
 
-            logger.info(f'Built episodic edges: {episodic_edges}')
+            logger.debug(f'Built episodic edges: {episodic_edges}')
 
             episode.entity_edges = [edge.uuid for edge in entity_edges]
 
             # Future optimization would be using batch operations to save nodes and edges
+            if not self.store_raw_episode_content:
+                episode.content = ''
+
             await episode.save(self.driver)
             await asyncio.gather(*[node.save(self.driver) for node in nodes])
             await asyncio.gather(*[edge.save(self.driver) for edge in episodic_edges])
@@ -451,6 +459,8 @@ class Graphiti:
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
 
+            return AddEpisodeResults(episode=episode, nodes=nodes, edges=entity_edges)
+
         except Exception as e:
             raise e
 
@@ -554,7 +564,7 @@ class Graphiti:
             edges = await dedupe_edges_bulk(
                 self.driver, self.llm_client, extracted_edges_with_resolved_pointers
             )
-            logger.info(f'extracted edge length: {len(edges)}')
+            logger.debug(f'extracted edge length: {len(edges)}')
 
             # invalidate edges
 
@@ -567,11 +577,20 @@ class Graphiti:
         except Exception as e:
             raise e
 
-    async def build_communities(self):
+    async def build_communities(self, group_ids: list[str] | None = None) -> list[CommunityNode]:
+        """
+        Use a community clustering algorithm to find communities of nodes. Create community nodes summarising
+        the content of these communities.
+        ----------
+        query : list[str] | None
+            Optional. Create communities only for the listed group_ids. If blank the entire graph will be used.
+        """
         # Clear existing communities
         await remove_communities(self.driver)
 
-        community_nodes, community_edges = await build_communities(self.driver, self.llm_client)
+        community_nodes, community_edges = await build_communities(
+            self.driver, self.llm_client, group_ids
+        )
 
         await asyncio.gather(
             *[node.generate_name_embedding(self.embedder) for node in community_nodes]
@@ -580,6 +599,8 @@ class Graphiti:
         await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
         await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
 
+        return community_nodes
+
     async def search(
         self,
         query: str,
@@ -700,18 +721,17 @@ class Graphiti:
         ).nodes
         return nodes
 
+    async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
+        episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
-async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
-    episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
-
-    edges_list = await asyncio.gather(
-        *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
-    )
+        edges_list = await asyncio.gather(
+            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
+        )
 
-    edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
+        edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
 
-    nodes = await get_mentioned_nodes(self.driver, episodes)
+        nodes = await get_mentioned_nodes(self.driver, episodes)
 
-    communities = await get_communities_by_nodes(self.driver, nodes)
+        communities = await get_communities_by_nodes(self.driver, nodes)
 
-    return SearchResults(edges=edges, nodes=nodes, communities=communities)
+        return SearchResults(edges=edges, nodes=nodes, communities=communities)

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/helpers.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 from datetime import datetime
 
+import numpy as np
 from neo4j import time as neo4j_time
 
 
@@ -25,7 +26,7 @@ def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
 
 def lucene_sanitize(query: str) -> str:
     # Escape special characters from a query before passing into Lucene
-    # + - && || ! ( ) { } [ ] ^ " ~ * ? : \
+    # + - && || ! ( ) { } [ ] ^ " ~ * ? : \ /
     escape_map = str.maketrans(
         {
             '+': r'\+',
@@ -46,8 +47,21 @@ def lucene_sanitize(query: str) -> str:
             '?': r'\?',
             ':': r'\:',
             '\\': r'\\',
+            '/': r'\/',
         }
     )
 
     sanitized = query.translate(escape_map)
     return sanitized
+
+
+def normalize_l2(embedding: list[float]) -> list[float]:
+    embedding_array = np.array(embedding)
+    if embedding_array.ndim == 1:
+        norm = np.linalg.norm(embedding_array)
+        if norm == 0:
+            return embedding_array.tolist()
+        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()

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/nodes.py

@@ -86,7 +86,7 @@ class Node(BaseModel, ABC):
             uuid=self.uuid,
         )
 
-        logger.info(f'Deleted Node: {self.uuid}')
+        logger.debug(f'Deleted Node: {self.uuid}')
 
         return result
 
@@ -135,7 +135,7 @@ class EpisodicNode(Node):
             source=self.source.value,
         )
 
-        logger.info(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to neo4j: {self.uuid}')
 
         return result
 
@@ -217,7 +217,7 @@ class EntityNode(Node):
         text = self.name.replace('\n', ' ')
         self.name_embedding = await embedder.create(input=[text])
         end = time()
-        logger.info(f'embedded {text} in {end - start} ms')
+        logger.debug(f'embedded {text} in {end - start} ms')
 
         return self.name_embedding
 
@@ -225,7 +225,8 @@ class EntityNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Entity {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, group_id: $group_id, summary: $summary, created_at: $created_at}
+        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,
@@ -235,7 +236,7 @@ class EntityNode(Node):
             created_at=self.created_at,
         )
 
-        logger.info(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to neo4j: {self.uuid}')
 
         return result
 
@@ -257,6 +258,9 @@ class EntityNode(Node):
 
         nodes = [get_entity_node_from_record(record) for record in records]
 
+        if len(nodes) == 0:
+            raise NodeNotFoundError(uuid)
+
         return nodes[0]
 
     @classmethod
@@ -308,7 +312,8 @@ class CommunityNode(Node):
         result = await driver.execute_query(
             """
         MERGE (n:Community {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, name_embedding: $name_embedding, group_id: $group_id, summary: $summary, created_at: $created_at}
+        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
         RETURN n.uuid AS uuid""",
             uuid=self.uuid,
             name=self.name,
@@ -318,7 +323,7 @@ class CommunityNode(Node):
             created_at=self.created_at,
         )
 
-        logger.info(f'Saved Node to neo4j: {self.uuid}')
+        logger.debug(f'Saved Node to neo4j: {self.uuid}')
 
         return result
 
@@ -327,7 +332,7 @@ class CommunityNode(Node):
         text = self.name.replace('\n', ' ')
         self.name_embedding = await embedder.create(input=[text])
         end = time()
-        logger.info(f'embedded {text} in {end - start} ms')
+        logger.debug(f'embedded {text} in {end - start} ms')
 
         return self.name_embedding
 
@@ -349,6 +354,9 @@ class CommunityNode(Node):
 
         nodes = [get_community_node_from_record(record) for record in records]
 
+        if len(nodes) == 0:
+            raise NodeNotFoundError(uuid)
+
         return nodes[0]
 
     @classmethod

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/prompts/eval.py

@@ -23,11 +23,33 @@ from .models import Message, PromptFunction, PromptVersion
 class Prompt(Protocol):
     qa_prompt: PromptVersion
     eval_prompt: PromptVersion
+    query_expansion: PromptVersion
 
 
 class Versions(TypedDict):
     qa_prompt: PromptFunction
     eval_prompt: PromptFunction
+    query_expansion: PromptFunction
+
+
+def query_expansion(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an expert at rephrasing questions into queries used in a database retrieval system"""
+
+    user_prompt = f"""
+    Bob is asking Alice a question, are you able to rephrase the question into a simpler one about Alice in the third person
+    that maintains the relevant context?
+    <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),
+        Message(role='user', content=user_prompt),
+    ]
 
 
 def qa_prompt(context: dict[str, Any]) -> list[Message]:
@@ -38,7 +60,7 @@ def qa_prompt(context: dict[str, Any]) -> list[Message]:
     You are given the following entity summaries and facts to help you determine the answer to your question.
     <ENTITY_SUMMARIES>
     {json.dumps(context['entity_summaries'])}
-    </ENTITY_SUMMARIES
+    </ENTITY_SUMMARIES>
     <FACTS>
     {json.dumps(context['facts'])}
     </FACTS>
@@ -87,4 +109,8 @@ def eval_prompt(context: dict[str, Any]) -> list[Message]:
     ]
 
 
-versions: Versions = {'qa_prompt': qa_prompt, 'eval_prompt': eval_prompt}
+versions: Versions = {
+    'qa_prompt': qa_prompt,
+    'eval_prompt': eval_prompt,
+    'query_expansion': query_expansion,
+}

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/prompts/extract_edge_dates.py

@@ -37,7 +37,6 @@ def v1(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
             Edge:
-            Edge Name: {context['edge_name']}
             Fact: {context['edge_fact']}
 
             Current Episode: {context['current_episode']}
@@ -56,17 +55,17 @@ def v1(context: dict[str, Any]) -> list[Message]:
             Guidelines:
             1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SSZ) for datetimes.
             2. Use the reference timestamp as the current time when determining the valid_at and invalid_at dates.
-            3. If no temporal information is found that establishes or changes the relationship, leave the fields as null.
-            4. Do not infer dates from related events. Only use dates that are directly stated to establish or change the relationship.
-			5. For relative time mentions directly related to the relationship, calculate the actual datetime based on the reference timestamp.
-            6. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
-            7. If only a year is mentioned, use January 1st of that year at 00:00:00.
+            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.
+            5. Do not infer dates from related events. Only use dates that are directly stated to establish or change the relationship.
+			6. For relative time mentions directly related to the relationship, calculate the actual datetime based on the reference timestamp.
+            7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
+            8. If only a 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:SSZ or null",
-                "invalid_at": "YYYY-MM-DDTHH:MM:SSZ or null",
-                "explanation": "Brief explanation of why these dates were chosen or why they were set to null"
+                "valid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
+                "invalid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
             }}
             """,
         ),

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/prompts/extract_edges.py

@@ -113,8 +113,9 @@ def v2(context: dict[str, Any]) -> list[Message]:
         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. Consider temporal aspects of relationships when relevant.
-        6. Avoid using the same node as the source and target of a 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
 
         Respond with a JSON object in the following format:
         {{

{graphiti_core-0.3.8 → graphiti_core-0.3.11}/graphiti_core/prompts/invalidate_edges.py

@@ -82,7 +82,7 @@ def v2(context: dict[str, Any]) -> list[Message]:
         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 contradictions with the New Edge.
+               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.
 
                 Existing Edges:
                 {context['existing_edges']}