graphiti-core 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

graphiti_core/nodes.py

@@ -25,6 +25,7 @@ from uuid import uuid4
 from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
 
+from graphiti_core.errors import NodeNotFoundError
 from graphiti_core.llm_client.config import EMBEDDING_DIM
 
 logger = logging.getLogger(__name__)
@@ -76,8 +77,18 @@ class Node(BaseModel, ABC):
     @abstractmethod
     async def save(self, driver: AsyncDriver): ...
 
-    @abstractmethod
-    async def delete(self, driver: AsyncDriver): ...
+    async def delete(self, driver: AsyncDriver):
+        result = await driver.execute_query(
+            """
+        MATCH (n {uuid: $uuid})
+        DETACH DELETE n
+        """,
+            uuid=self.uuid,
+        )
+
+        logger.info(f'Deleted Node: {self.uuid}')
+
+        return result
 
     def __hash__(self):
         return hash(self.uuid)
@@ -90,6 +101,9 @@ class Node(BaseModel, ABC):
     @classmethod
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str): ...
 
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]): ...
+
 
 class EpisodicNode(Node):
     source: EpisodeType = Field(description='source type')
@@ -125,24 +139,37 @@ class EpisodicNode(Node):
 
         return result
 
-    async def delete(self, driver: AsyncDriver):
-        result = await driver.execute_query(
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+        records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Episodic {uuid: $uuid})
-        DETACH DELETE n
+        MATCH (e:Episodic {uuid: $uuid})
+            RETURN e.content AS content,
+            e.created_at AS created_at,
+            e.valid_at AS valid_at,
+            e.uuid AS uuid,
+            e.name AS name,
+            e.group_id AS group_id,
+            e.source_description AS source_description,
+            e.source AS source
         """,
-            uuid=self.uuid,
+            uuid=uuid,
         )
 
-        logger.info(f'Deleted Node: {self.uuid}')
+        episodes = [get_episodic_node_from_record(record) for record in records]
 
-        return result
+        logger.info(f'Found Node: {uuid}')
+
+        if len(episodes) == 0:
+            raise NodeNotFoundError(uuid)
+
+        return episodes[0]
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (e:Episodic {uuid: $uuid})
+        MATCH (e:Episodic) WHERE e.uuid IN $uuids
             RETURN e.content AS content,
             e.created_at AS created_at,
             e.valid_at AS valid_at,
@@ -152,14 +179,14 @@ class EpisodicNode(Node):
             e.source_description AS source_description,
             e.source AS source
         """,
-            uuid=uuid,
+            uuids=uuids,
         )
 
         episodes = [get_episodic_node_from_record(record) for record in records]
 
-        logger.info(f'Found Node: {uuid}')
+        logger.info(f'Found Nodes: {uuids}')
 
-        return episodes[0]
+        return episodes
 
 
 class EntityNode(Node):
@@ -194,24 +221,88 @@ class EntityNode(Node):
 
         return result
 
-    async def delete(self, driver: AsyncDriver):
-        result = await driver.execute_query(
+    @classmethod
+    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+        records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity {uuid: $uuid})
-        DETACH DELETE n
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id
+            n.created_at AS created_at, 
+            n.summary AS summary
+        """,
+            uuid=uuid,
+        )
+
+        nodes = [get_entity_node_from_record(record) for record in records]
+
+        logger.info(f'Found Node: {uuid}')
+
+        return nodes[0]
+
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Entity) WHERE n.uuid IN $uuids
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id,
+            n.created_at AS created_at, 
+            n.summary AS summary
         """,
+            uuids=uuids,
+        )
+
+        nodes = [get_entity_node_from_record(record) for record in records]
+
+        logger.info(f'Found Nodes: {uuids}')
+
+        return nodes
+
+
+class CommunityNode(Node):
+    name_embedding: list[float] | None = Field(default=None, description='embedding of the name')
+    summary: str = Field(description='region summary of member nodes', default_factory=str)
+
+    async def save(self, driver: AsyncDriver):
+        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}
+        RETURN n.uuid AS uuid""",
             uuid=self.uuid,
+            name=self.name,
+            group_id=self.group_id,
+            summary=self.summary,
+            name_embedding=self.name_embedding,
+            created_at=self.created_at,
         )
 
-        logger.info(f'Deleted Node: {self.uuid}')
+        logger.info(f'Saved Node to neo4j: {self.uuid}')
 
         return result
 
+    async def generate_name_embedding(self, embedder, model='text-embedding-3-small'):
+        start = time()
+        text = self.name.replace('\n', ' ')
+        embedding = (await embedder.create(input=[text], model=model)).data[0].embedding
+        self.name_embedding = embedding[:EMBEDDING_DIM]
+        end = time()
+        logger.info(f'embedded {text} in {end - start} ms')
+
+        return embedding
+
     @classmethod
     async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-        MATCH (n:Entity {uuid: $uuid})
+        MATCH (n:Community {uuid: $uuid})
         RETURN
             n.uuid As uuid, 
             n.name AS name,
@@ -223,12 +314,34 @@ class EntityNode(Node):
             uuid=uuid,
         )
 
-        nodes = [get_entity_node_from_record(record) for record in records]
+        nodes = [get_community_node_from_record(record) for record in records]
 
         logger.info(f'Found Node: {uuid}')
 
         return nodes[0]
 
+    @classmethod
+    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+        records, _, _ = await driver.execute_query(
+            """
+        MATCH (n:Community) WHERE n.uuid IN $uuids
+        RETURN
+            n.uuid As uuid, 
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.group_id AS group_id
+            n.created_at AS created_at, 
+            n.summary AS summary
+        """,
+            uuids=uuids,
+        )
+
+        nodes = [get_community_node_from_record(record) for record in records]
+
+        logger.info(f'Found Nodes: {uuids}')
+
+        return nodes
+
 
 # Node helpers
 def get_episodic_node_from_record(record: Any) -> EpisodicNode:
@@ -254,3 +367,14 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         created_at=record['created_at'].to_native(),
         summary=record['summary'],
     )
+
+
+def get_community_node_from_record(record: Any) -> CommunityNode:
+    return CommunityNode(
+        uuid=record['uuid'],
+        name=record['name'],
+        group_id=record['group_id'],
+        name_embedding=record['name_embedding'],
+        created_at=record['created_at'].to_native(),
+        summary=record['summary'],
+    )

graphiti_core/prompts/extract_edge_dates.py

@@ -1,3 +1,19 @@
+"""
+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.
+"""
+
 from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion

graphiti_core/prompts/extract_nodes.py

@@ -24,12 +24,14 @@ class Prompt(Protocol):
     v1: PromptVersion
     v2: PromptVersion
     extract_json: PromptVersion
+    extract_text: PromptVersion
 
 
 class Versions(TypedDict):
     v1: PromptFunction
     v2: PromptFunction
     extract_json: PromptFunction
+    extract_text: PromptFunction
 
 
 def v1(context: dict[str, Any]) -> list[Message]:
@@ -144,4 +146,44 @@ Respond with a JSON object in the following format:
     ]
 
 
-versions: Versions = {'v1': v1, 'v2': v2, 'extract_json': extract_json}
+def extract_text(context: dict[str, Any]) -> list[Message]:
+    sys_prompt = """You are an AI assistant that extracts entity nodes from conversational text. Your primary task is to identify and extract the speaker and other significant entities mentioned in the conversation."""
+
+    user_prompt = f"""
+Given the following conversation, extract entity nodes from the CURRENT MESSAGE that are explicitly or implicitly mentioned:
+
+Conversation:
+{json.dumps([ep['content'] for ep in context['previous_episodes']], indent=2)}
+<CURRENT MESSAGE>
+{context["episode_content"]}
+
+Guidelines:
+2. Extract significant entities, concepts, or actors mentioned in the conversation.
+3. Provide concise but informative summaries for each extracted node.
+4. Avoid creating nodes for relationships or actions.
+5. Avoid creating nodes for temporal information like dates, times or years (these will be added to edges later).
+6. 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_nodes": [
+        {{
+            "name": "Unique identifier for the node (use the speaker's name for speaker nodes)",
+            "labels": ["Entity", "OptionalAdditionalLabel"],
+            "summary": "Brief summary of the node's role or significance"
+        }}
+    ]
+}}
+"""
+    return [
+        Message(role='system', content=sys_prompt),
+        Message(role='user', content=user_prompt),
+    ]
+
+
+versions: Versions = {
+    'v1': v1,
+    'v2': v2,
+    'extract_json': extract_json,
+    'extract_text': extract_text,
+}

graphiti_core/prompts/lib.py

@@ -71,6 +71,9 @@ from .invalidate_edges import (
     versions as invalidate_edges_versions,
 )
 from .models import Message, PromptFunction
+from .summarize_nodes import Prompt as SummarizeNodesPrompt
+from .summarize_nodes import Versions as SummarizeNodesVersions
+from .summarize_nodes import versions as summarize_nodes_versions
 
 
 class PromptLibrary(Protocol):
@@ -80,6 +83,7 @@ class PromptLibrary(Protocol):
     dedupe_edges: DedupeEdgesPrompt
     invalidate_edges: InvalidateEdgesPrompt
     extract_edge_dates: ExtractEdgeDatesPrompt
+    summarize_nodes: SummarizeNodesPrompt
 
 
 class PromptLibraryImpl(TypedDict):
@@ -89,6 +93,7 @@ class PromptLibraryImpl(TypedDict):
     dedupe_edges: DedupeEdgesVersions
     invalidate_edges: InvalidateEdgesVersions
     extract_edge_dates: ExtractEdgeDatesVersions
+    summarize_nodes: SummarizeNodesVersions
 
 
 class VersionWrapper:
@@ -118,5 +123,6 @@ PROMPT_LIBRARY_IMPL: PromptLibraryImpl = {
     'dedupe_edges': dedupe_edges_versions,
     'invalidate_edges': invalidate_edges_versions,
     'extract_edge_dates': extract_edge_dates_versions,
+    'summarize_nodes': summarize_nodes_versions,
 }
 prompt_library: PromptLibrary = PromptLibraryWrapper(PROMPT_LIBRARY_IMPL)  # type: ignore[assignment]

graphiti_core/prompts/summarize_nodes.py

@@ -0,0 +1,79 @@
+"""
+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, Protocol, TypedDict
+
+from .models import Message, PromptFunction, PromptVersion
+
+
+class Prompt(Protocol):
+    summarize_pair: PromptVersion
+    summary_description: PromptVersion
+
+
+class Versions(TypedDict):
+    summarize_pair: PromptFunction
+    summary_description: PromptFunction
+
+
+def summarize_pair(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that combines summaries.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        Synthesize the information from the following two summaries into a single succinct summary.
+
+        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"
+            }}
+        """,
+        ),
+    ]
+
+
+def summary_description(context: dict[str, Any]) -> list[Message]:
+    return [
+        Message(
+            role='system',
+            content='You are a helpful assistant that describes provided contents in a single sentence.',
+        ),
+        Message(
+            role='user',
+            content=f"""
+        Create a short one sentence description of the summary that explains what kind of information is summarized.
+
+        Summary:
+        {json.dumps(context['summary'], indent=2)}
+
+        Respond with a JSON object in the following format:
+            {{
+                "description": "One sentence description of the provided summary"
+            }}
+        """,
+        ),
+    ]
+
+
+versions: Versions = {'summarize_pair': summarize_pair, 'summary_description': summary_description}

graphiti_core/py.typed

@@ -0,0 +1 @@
+# This file is intentionally left empty to indicate that the package is typed.