getzep--graphiti/graphiti_core/namespaces/nodes.py

"""
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 datetime import datetime

from graphiti_core.driver.driver import GraphDriver
from graphiti_core.driver.operations.community_node_ops import CommunityNodeOperations
from graphiti_core.driver.operations.entity_node_ops import EntityNodeOperations
from graphiti_core.driver.operations.episode_node_ops import EpisodeNodeOperations
from graphiti_core.driver.operations.saga_node_ops import SagaNodeOperations
from graphiti_core.driver.query_executor import Transaction
from graphiti_core.embedder import EmbedderClient
from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode, SagaNode


class EntityNodeNamespace:
    """Namespace for entity node operations. Accessed as ``graphiti.nodes.entity``."""

    def __init__(
        self,
        driver: GraphDriver,
        ops: EntityNodeOperations,
        embedder: EmbedderClient,
    ):
        self._driver = driver
        self._ops = ops
        self._embedder = embedder

    async def save(
        self,
        node: EntityNode,
        tx: Transaction | None = None,
    ) -> EntityNode:
        await node.generate_name_embedding(self._embedder)
        await self._ops.save(self._driver, node, tx=tx)
        return node

    async def save_bulk(
        self,
        nodes: list[EntityNode],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.save_bulk(self._driver, nodes, tx=tx, batch_size=batch_size)

    async def delete(
        self,
        node: EntityNode,
        tx: Transaction | None = None,
    ) -> None:
        await self._ops.delete(self._driver, node, tx=tx)

    async def delete_by_group_id(
        self,
        group_id: str,
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_group_id(self._driver, group_id, tx=tx, batch_size=batch_size)

    async def delete_by_uuids(
        self,
        uuids: list[str],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_uuids(self._driver, uuids, tx=tx, batch_size=batch_size)

    async def get_by_uuid(self, uuid: str) -> EntityNode:
        return await self._ops.get_by_uuid(self._driver, uuid)

    async def get_by_uuids(self, uuids: list[str]) -> list[EntityNode]:
        return await self._ops.get_by_uuids(self._driver, uuids)

    async def get_by_group_ids(
        self,
        group_ids: list[str],
        limit: int | None = None,
        uuid_cursor: str | None = None,
    ) -> list[EntityNode]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)

    async def load_embeddings(self, node: EntityNode) -> None:
        await self._ops.load_embeddings(self._driver, node)

    async def load_embeddings_bulk(
        self,
        nodes: list[EntityNode],
        batch_size: int = 100,
    ) -> None:
        await self._ops.load_embeddings_bulk(self._driver, nodes, batch_size)


class EpisodeNodeNamespace:
    """Namespace for episode node operations. Accessed as ``graphiti.nodes.episode``."""

    def __init__(self, driver: GraphDriver, ops: EpisodeNodeOperations):
        self._driver = driver
        self._ops = ops

    async def save(
        self,
        node: EpisodicNode,
        tx: Transaction | None = None,
    ) -> EpisodicNode:
        await self._ops.save(self._driver, node, tx=tx)
        return node

    async def save_bulk(
        self,
        nodes: list[EpisodicNode],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.save_bulk(self._driver, nodes, tx=tx, batch_size=batch_size)

    async def delete(
        self,
        node: EpisodicNode,
        tx: Transaction | None = None,
    ) -> None:
        await self._ops.delete(self._driver, node, tx=tx)

    async def delete_by_group_id(
        self,
        group_id: str,
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_group_id(self._driver, group_id, tx=tx, batch_size=batch_size)

    async def delete_by_uuids(
        self,
        uuids: list[str],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_uuids(self._driver, uuids, tx=tx, batch_size=batch_size)

    async def get_by_uuid(self, uuid: str) -> EpisodicNode:
        return await self._ops.get_by_uuid(self._driver, uuid)

    async def get_by_uuids(self, uuids: list[str]) -> list[EpisodicNode]:
        return await self._ops.get_by_uuids(self._driver, uuids)

    async def get_by_group_ids(
        self,
        group_ids: list[str],
        limit: int | None = None,
        uuid_cursor: str | None = None,
    ) -> list[EpisodicNode]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)

    async def get_by_entity_node_uuid(
        self,
        entity_node_uuid: str,
    ) -> list[EpisodicNode]:
        return await self._ops.get_by_entity_node_uuid(self._driver, entity_node_uuid)

    async def retrieve_episodes(
        self,
        reference_time: datetime,
        last_n: int = 3,
        group_ids: list[str] | None = None,
        source: str | None = None,
        saga: str | None = None,
    ) -> list[EpisodicNode]:
        return await self._ops.retrieve_episodes(
            self._driver, reference_time, last_n, group_ids, source, saga
        )


class CommunityNodeNamespace:
    """Namespace for community node operations. Accessed as ``graphiti.nodes.community``."""

    def __init__(
        self,
        driver: GraphDriver,
        ops: CommunityNodeOperations,
        embedder: EmbedderClient,
    ):
        self._driver = driver
        self._ops = ops
        self._embedder = embedder

    async def save(
        self,
        node: CommunityNode,
        tx: Transaction | None = None,
    ) -> CommunityNode:
        await node.generate_name_embedding(self._embedder)
        await self._ops.save(self._driver, node, tx=tx)
        return node

    async def save_bulk(
        self,
        nodes: list[CommunityNode],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.save_bulk(self._driver, nodes, tx=tx, batch_size=batch_size)

    async def delete(
        self,
        node: CommunityNode,
        tx: Transaction | None = None,
    ) -> None:
        await self._ops.delete(self._driver, node, tx=tx)

    async def delete_by_group_id(
        self,
        group_id: str,
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_group_id(self._driver, group_id, tx=tx, batch_size=batch_size)

    async def delete_by_uuids(
        self,
        uuids: list[str],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_uuids(self._driver, uuids, tx=tx, batch_size=batch_size)

    async def get_by_uuid(self, uuid: str) -> CommunityNode:
        return await self._ops.get_by_uuid(self._driver, uuid)

    async def get_by_uuids(self, uuids: list[str]) -> list[CommunityNode]:
        return await self._ops.get_by_uuids(self._driver, uuids)

    async def get_by_group_ids(
        self,
        group_ids: list[str],
        limit: int | None = None,
        uuid_cursor: str | None = None,
    ) -> list[CommunityNode]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)

    async def load_name_embedding(self, node: CommunityNode) -> None:
        await self._ops.load_name_embedding(self._driver, node)


class SagaNodeNamespace:
    """Namespace for saga node operations. Accessed as ``graphiti.nodes.saga``."""

    def __init__(self, driver: GraphDriver, ops: SagaNodeOperations):
        self._driver = driver
        self._ops = ops

    async def save(
        self,
        node: SagaNode,
        tx: Transaction | None = None,
    ) -> SagaNode:
        await self._ops.save(self._driver, node, tx=tx)
        return node

    async def save_bulk(
        self,
        nodes: list[SagaNode],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.save_bulk(self._driver, nodes, tx=tx, batch_size=batch_size)

    async def delete(
        self,
        node: SagaNode,
        tx: Transaction | None = None,
    ) -> None:
        await self._ops.delete(self._driver, node, tx=tx)

    async def delete_by_group_id(
        self,
        group_id: str,
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_group_id(self._driver, group_id, tx=tx, batch_size=batch_size)

    async def delete_by_uuids(
        self,
        uuids: list[str],
        tx: Transaction | None = None,
        batch_size: int = 100,
    ) -> None:
        await self._ops.delete_by_uuids(self._driver, uuids, tx=tx, batch_size=batch_size)

    async def get_by_uuid(self, uuid: str) -> SagaNode:
        return await self._ops.get_by_uuid(self._driver, uuid)

    async def get_by_uuids(self, uuids: list[str]) -> list[SagaNode]:
        return await self._ops.get_by_uuids(self._driver, uuids)

    async def get_by_group_ids(
        self,
        group_ids: list[str],
        limit: int | None = None,
        uuid_cursor: str | None = None,
    ) -> list[SagaNode]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)


class NodeNamespace:
    """Namespace for all node operations. Accessed as ``graphiti.nodes``.

    Sub-namespaces are set only when the driver provides the corresponding
    operations implementation.  Accessing an unset attribute raises
    ``NotImplementedError`` with a clear message.
    """

    entity: EntityNodeNamespace
    episode: EpisodeNodeNamespace
    community: CommunityNodeNamespace
    saga: SagaNodeNamespace

    _driver_name: str

    def __init__(self, driver: GraphDriver, embedder: EmbedderClient):
        self._driver_name = type(driver).__name__

        entity_node_ops = driver.entity_node_ops
        if entity_node_ops is not None:
            self.entity = EntityNodeNamespace(driver, entity_node_ops, embedder)

        episode_node_ops = driver.episode_node_ops
        if episode_node_ops is not None:
            self.episode = EpisodeNodeNamespace(driver, episode_node_ops)

        community_node_ops = driver.community_node_ops
        if community_node_ops is not None:
            self.community = CommunityNodeNamespace(driver, community_node_ops, embedder)

        saga_node_ops = driver.saga_node_ops
        if saga_node_ops is not None:
            self.saga = SagaNodeNamespace(driver, saga_node_ops)

    def __getattr__(self, name: str) -> object:
        if name in ('entity', 'episode', 'community', 'saga'):
            raise NotImplementedError(f'{self._driver_name} does not implement {name}_node_ops')
        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")