graphiti/core/utils/maintenance/edge_operations.py

import json
import logging
from datetime import datetime
from time import time
from typing import List

from core.edges import EntityEdge, EpisodicEdge
from core.llm_client import LLMClient
from core.nodes import EntityNode, EpisodicNode
from core.prompts import prompt_library

logger = logging.getLogger(__name__)


def build_episodic_edges(
    entity_nodes: List[EntityNode],
    episode: EpisodicNode,
    created_at: datetime,
) -> List[EpisodicEdge]:
    edges: List[EpisodicEdge] = []

    for node in entity_nodes:
        edge = EpisodicEdge(
            source_node_uuid=episode.uuid,
            target_node_uuid=node.uuid,
            created_at=created_at,
        )
        edges.append(edge)

    return edges


async def extract_new_edges(
    llm_client: LLMClient,
    episode: EpisodicNode,
    new_nodes: list[EntityNode],
    relevant_schema: dict[str, any],
    previous_episodes: list[EpisodicNode],
) -> tuple[list[EntityEdge], list[EntityNode]]:
    # Prepare context for LLM
    context = {
        "episode_content": episode.content,
        "episode_timestamp": (
            episode.valid_at.isoformat() if episode.valid_at else None
        ),
        "relevant_schema": json.dumps(relevant_schema, indent=2),
        "new_nodes": [
            {"name": node.name, "summary": node.summary} for node in new_nodes
        ],
        "previous_episodes": [
            {
                "content": ep.content,
                "timestamp": ep.valid_at.isoformat() if ep.valid_at else None,
            }
            for ep in previous_episodes
        ],
    }

    llm_response = await llm_client.generate_response(
        prompt_library.extract_edges.v1(context)
    )
    new_edges_data = llm_response.get("new_edges", [])
    logger.info(f"Extracted new edges: {new_edges_data}")

    # Convert the extracted data into EntityEdge objects
    new_edges = []
    for edge_data in new_edges_data:
        source_node = next(
            (node for node in new_nodes if node.name == edge_data["source_node"]),
            None,
        )
        target_node = next(
            (node for node in new_nodes if node.name == edge_data["target_node"]),
            None,
        )

        # If source or target is not in new_nodes, check if it's an existing node
        if source_node is None and edge_data["source_node"] in relevant_schema["nodes"]:
            existing_node_data = relevant_schema["nodes"][edge_data["source_node"]]
            source_node = EntityNode(
                uuid=existing_node_data["uuid"],
                name=edge_data["source_node"],
                labels=[existing_node_data["label"]],
                summary="",
                created_at=datetime.now(),
            )
        if target_node is None and edge_data["target_node"] in relevant_schema["nodes"]:
            existing_node_data = relevant_schema["nodes"][edge_data["target_node"]]
            target_node = EntityNode(
                uuid=existing_node_data["uuid"],
                name=edge_data["target_node"],
                labels=[existing_node_data["label"]],
                summary="",
                created_at=datetime.now(),
            )

        if (
            source_node
            and target_node
            and not (
                source_node.name.startswith("Message")
                or target_node.name.startswith("Message")
            )
        ):
            valid_at = (
                datetime.fromisoformat(edge_data["valid_at"])
                if edge_data["valid_at"]
                else episode.valid_at or datetime.now()
            )
            invalid_at = (
                datetime.fromisoformat(edge_data["invalid_at"])
                if edge_data["invalid_at"]
                else None
            )

            new_edge = EntityEdge(
                source_node=source_node,
                target_node=target_node,
                name=edge_data["relation_type"],
                fact=edge_data["fact"],
                episodes=[episode.uuid],
                created_at=datetime.now(),
                valid_at=valid_at,
                invalid_at=invalid_at,
            )
            new_edges.append(new_edge)
            logger.info(
                f"Created new edge: {new_edge.name} from {source_node.name} (UUID: {source_node.uuid}) to {target_node.name} (UUID: {target_node.uuid})"
            )

    affected_nodes = set()

    for edge in new_edges:
        affected_nodes.add(edge.source_node)
        affected_nodes.add(edge.target_node)
    return new_edges, list(affected_nodes)


async def extract_edges(
    llm_client: LLMClient,
    episode: EpisodicNode,
    nodes: list[EntityNode],
    previous_episodes: list[EpisodicNode],
) -> list[EntityEdge]:
    start = time()

    # Prepare context for LLM
    context = {
        "episode_content": episode.content,
        "episode_timestamp": (
            episode.valid_at.isoformat() if episode.valid_at else None
        ),
        "nodes": [
            {"uuid": node.uuid, "name": node.name, "summary": node.summary}
            for node in nodes
        ],
        "previous_episodes": [
            {
                "content": ep.content,
                "timestamp": ep.valid_at.isoformat() if ep.valid_at else None,
            }
            for ep in previous_episodes
        ],
    }

    llm_response = await llm_client.generate_response(
        prompt_library.extract_edges.v2(context)
    )
    edges_data = llm_response.get("edges", [])

    end = time()
    logger.info(f"Extracted new edges: {edges_data} in {(end - start) * 1000} ms")

    # Convert the extracted data into EntityEdge objects
    edges = []
    for edge_data in edges_data:
        if edge_data["target_node_uuid"] and edge_data["source_node_uuid"]:
            edge = EntityEdge(
                source_node_uuid=edge_data["source_node_uuid"],
                target_node_uuid=edge_data["target_node_uuid"],
                name=edge_data["relation_type"],
                fact=edge_data["fact"],
                episodes=[episode.uuid],
                created_at=datetime.now(),
                valid_at=None,
                invalid_at=None,
            )
            edges.append(edge)
            logger.info(
                f"Created new edge: {edge.name} from (UUID: {edge.source_node_uuid}) to (UUID: {edge.target_node_uuid})"
            )

    return edges


async def dedupe_extracted_edges(
    llm_client: LLMClient,
    extracted_edges: list[EntityEdge],
    existing_edges: list[EntityEdge],
) -> list[EntityEdge]:
    # Create edge map
    edge_map = {}
    for edge in existing_edges:
        edge_map[edge.fact] = edge
    for edge in extracted_edges:
        if edge.fact in edge_map:
            continue
        edge_map[edge.fact] = edge

    # Prepare context for LLM
    context = {
        "extracted_edges": [
            {"name": edge.name, "fact": edge.fact} for edge in extracted_edges
        ],
        "existing_edges": [
            {"name": edge.name, "fact": edge.fact} for edge in extracted_edges
        ],
    }

    llm_response = await llm_client.generate_response(
        prompt_library.dedupe_edges.v1(context)
    )
    new_edges_data = llm_response.get("new_edges", [])
    logger.info(f"Extracted new edges: {new_edges_data}")

    # Get full edge data
    edges = []
    for edge_data in new_edges_data:
        edge = edge_map[edge_data["fact"]]
        edges.append(edge)

    return edges


async def dedupe_edge_list(
    llm_client: LLMClient,
    edges: list[EntityEdge],
) -> list[EntityEdge]:
    start = time()

    # Create edge map
    edge_map = {}
    for edge in edges:
        edge_map[edge.fact] = edge

    # Prepare context for LLM
    context = {"edges": [{"name": edge.name, "fact": edge.fact} for edge in edges]}

    llm_response = await llm_client.generate_response(
        prompt_library.dedupe_edges.edge_list(context)
    )
    unique_edges_data = llm_response.get("unique_edges", [])

    end = time()
    logger.info(
        f"Extracted edge duplicates: {unique_edges_data} in {(end - start)*1000} ms "
    )

    # Get full edge data
    unique_edges = []
    for edge_data in unique_edges_data:
        fact = edge_data["fact"]
        unique_edges.append(edge_map[fact])

    return unique_edges