getzep--graphiti/graphiti_core/namespaces/edges.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 graphiti_core.driver.driver import GraphDriver
from graphiti_core.driver.operations.community_edge_ops import CommunityEdgeOperations
from graphiti_core.driver.operations.entity_edge_ops import EntityEdgeOperations
from graphiti_core.driver.operations.episodic_edge_ops import EpisodicEdgeOperations
from graphiti_core.driver.operations.has_episode_edge_ops import HasEpisodeEdgeOperations
from graphiti_core.driver.operations.next_episode_edge_ops import NextEpisodeEdgeOperations
from graphiti_core.driver.query_executor import Transaction
from graphiti_core.edges import (
    CommunityEdge,
    EntityEdge,
    EpisodicEdge,
    HasEpisodeEdge,
    NextEpisodeEdge,
)
from graphiti_core.embedder import EmbedderClient


class EntityEdgeNamespace:
    """Namespace for entity edge operations. Accessed as ``graphiti.edges.entity``."""

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

    async def save(
        self,
        edge: EntityEdge,
        tx: Transaction | None = None,
    ) -> EntityEdge:
        await edge.generate_embedding(self._embedder)
        await self._ops.save(self._driver, edge, tx=tx)
        return edge

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

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

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

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

    async def get_by_uuids(self, uuids: list[str]) -> list[EntityEdge]:
        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[EntityEdge]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)

    async def get_between_nodes(
        self,
        source_node_uuid: str,
        target_node_uuid: str,
    ) -> list[EntityEdge]:
        return await self._ops.get_between_nodes(self._driver, source_node_uuid, target_node_uuid)

    async def get_by_node_uuid(self, node_uuid: str) -> list[EntityEdge]:
        return await self._ops.get_by_node_uuid(self._driver, node_uuid)

    async def load_embeddings(self, edge: EntityEdge) -> None:
        await self._ops.load_embeddings(self._driver, edge)

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


class EpisodicEdgeNamespace:
    """Namespace for episodic edge operations. Accessed as ``graphiti.edges.episodic``."""

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

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

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

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

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

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

    async def get_by_uuids(self, uuids: list[str]) -> list[EpisodicEdge]:
        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[EpisodicEdge]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)


class CommunityEdgeNamespace:
    """Namespace for community edge operations. Accessed as ``graphiti.edges.community``."""

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

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

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

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

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

    async def get_by_uuids(self, uuids: list[str]) -> list[CommunityEdge]:
        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[CommunityEdge]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)


class HasEpisodeEdgeNamespace:
    """Namespace for has_episode edge operations. Accessed as ``graphiti.edges.has_episode``."""

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

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

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

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

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

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

    async def get_by_uuids(self, uuids: list[str]) -> list[HasEpisodeEdge]:
        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[HasEpisodeEdge]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)


class NextEpisodeEdgeNamespace:
    """Namespace for next_episode edge operations. Accessed as ``graphiti.edges.next_episode``."""

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

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

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

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

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

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

    async def get_by_uuids(self, uuids: list[str]) -> list[NextEpisodeEdge]:
        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[NextEpisodeEdge]:
        return await self._ops.get_by_group_ids(self._driver, group_ids, limit, uuid_cursor)


class EdgeNamespace:
    """Namespace for all edge operations. Accessed as ``graphiti.edges``.

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

    entity: EntityEdgeNamespace
    episodic: EpisodicEdgeNamespace
    community: CommunityEdgeNamespace
    has_episode: HasEpisodeEdgeNamespace
    next_episode: NextEpisodeEdgeNamespace

    _driver_name: str

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

        entity_edge_ops = driver.entity_edge_ops
        if entity_edge_ops is not None:
            self.entity = EntityEdgeNamespace(driver, entity_edge_ops, embedder)

        episodic_edge_ops = driver.episodic_edge_ops
        if episodic_edge_ops is not None:
            self.episodic = EpisodicEdgeNamespace(driver, episodic_edge_ops)

        community_edge_ops = driver.community_edge_ops
        if community_edge_ops is not None:
            self.community = CommunityEdgeNamespace(driver, community_edge_ops)

        has_episode_edge_ops = driver.has_episode_edge_ops
        if has_episode_edge_ops is not None:
            self.has_episode = HasEpisodeEdgeNamespace(driver, has_episode_edge_ops)

        next_episode_edge_ops = driver.next_episode_edge_ops
        if next_episode_edge_ops is not None:
            self.next_episode = NextEpisodeEdgeNamespace(driver, next_episode_edge_ops)

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