graphiti/tests/utils/maintenance/test_temporal_operations_int.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.
"""

import os
from datetime import datetime, timedelta, timezone

import pytest
from dotenv import load_dotenv

from graphiti_core.edges import EntityEdge
from graphiti_core.llm_client import LLMConfig, OpenAIClient
from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
from graphiti_core.utils.maintenance.temporal_operations import (
    extract_edge_dates,
    get_edge_contradictions,
)

load_dotenv()


def setup_llm_client():
    return OpenAIClient(
        LLMConfig(
            api_key=os.getenv('TEST_OPENAI_API_KEY'),
            model=os.getenv('TEST_OPENAI_MODEL'),
            base_url='https://api.openai.com/v1',
        )
    )


def create_test_data():
    now = datetime.now(timezone.utc)

    # Create edges
    existing_edge = EntityEdge(
        uuid='e1',
        source_node_uuid='1',
        target_node_uuid='2',
        name='LIKES',
        fact='Alice likes Bob',
        created_at=now - timedelta(days=1),
        group_id='1',
    )
    new_edge = EntityEdge(
        uuid='e2',
        source_node_uuid='1',
        target_node_uuid='2',
        name='DISLIKES',
        fact='Alice dislikes Bob',
        created_at=now,
        group_id='1',
    )

    # Create current episode
    current_episode = EpisodicNode(
        name='Current Episode',
        content='Alice now dislikes Bob',
        created_at=now,
        valid_at=now,
        source=EpisodeType.message,
        source_description='Test episode for unit testing',
        group_id='1',
    )

    # Create previous episodes
    previous_episodes = [
        EpisodicNode(
            name='Previous Episode',
            content='Alice liked Bob',
            created_at=now - timedelta(days=1),
            valid_at=now - timedelta(days=1),
            source=EpisodeType.message,
            source_description='Test previous episode for unit testing',
            group_id='1',
        )
    ]

    return existing_edge, new_edge, current_episode, previous_episodes


@pytest.mark.asyncio
@pytest.mark.integration
async def test_get_edge_contradictions():
    existing_edge, new_edge, current_episode, previous_episodes = create_test_data()

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, [existing_edge])

    assert len(invalidated_edges) == 1
    assert invalidated_edges[0].uuid == existing_edge.uuid


@pytest.mark.asyncio
@pytest.mark.integration
async def test_get_edge_contradictions_no_contradictions():
    _, new_edge, current_episode, previous_episodes = create_test_data()

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, [])

    assert len(invalidated_edges) == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_get_edge_contradictions_multiple_existing():
    existing_edge1, new_edge, _, _ = create_test_data()
    existing_edge2, _, _, _ = create_test_data()
    existing_edge2.uuid = 'e3'
    existing_edge2.name = 'KNOWS'
    existing_edge2.fact = 'Alice knows Bob'

    invalidated_edges = await get_edge_contradictions(
        setup_llm_client(), new_edge, [existing_edge1, existing_edge2]
    )

    assert len(invalidated_edges) == 1
    assert invalidated_edges[0].uuid == existing_edge1.uuid


# Helper function to create more complex test data
def create_complex_test_data():
    now = datetime.now(timezone.utc)

    # Create nodes
    node1 = EntityNode(uuid='1', name='Alice', labels=['Person'], created_at=now, group_id='1')
    node2 = EntityNode(uuid='2', name='Bob', labels=['Person'], created_at=now, group_id='1')
    node3 = EntityNode(uuid='3', name='Charlie', labels=['Person'], created_at=now, group_id='1')
    node4 = EntityNode(
        uuid='4', name='Company XYZ', labels=['Organization'], created_at=now, group_id='1'
    )

    # Create edges
    existing_edge1 = EntityEdge(
        uuid='e1',
        source_node_uuid='1',
        target_node_uuid='2',
        name='LIKES',
        fact='Alice likes Bob',
        group_id='1',
        created_at=now - timedelta(days=5),
    )
    existing_edge2 = EntityEdge(
        uuid='e2',
        source_node_uuid='1',
        target_node_uuid='3',
        name='FRIENDS_WITH',
        fact='Alice is friends with Charlie',
        group_id='1',
        created_at=now - timedelta(days=3),
    )
    existing_edge3 = EntityEdge(
        uuid='e3',
        source_node_uuid='2',
        target_node_uuid='4',
        name='WORKS_FOR',
        fact='Bob works for Company XYZ',
        group_id='1',
        created_at=now - timedelta(days=2),
    )

    return [existing_edge1, existing_edge2, existing_edge3], [
        node1,
        node2,
        node3,
        node4,
    ]


@pytest.mark.asyncio
@pytest.mark.integration
async def test_invalidate_edges_complex():
    existing_edges, nodes = create_complex_test_data()

    # Create a new edge that contradicts an existing one
    new_edge = EntityEdge(
        uuid='e4',
        source_node_uuid='1',
        target_node_uuid='2',
        name='DISLIKES',
        fact='Alice dislikes Bob',
        group_id='1',
        created_at=datetime.now(timezone.utc),
    )

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, existing_edges)

    assert len(invalidated_edges) == 1
    assert invalidated_edges[0].uuid == 'e1'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_get_edge_contradictions_temporal_update():
    existing_edges, nodes = create_complex_test_data()

    # Create a new edge that updates an existing one with new information
    new_edge = EntityEdge(
        uuid='e5',
        source_node_uuid='2',
        target_node_uuid='4',
        name='LEFT_JOB',
        fact='Bob no longer works at at Company XYZ',
        group_id='1',
        created_at=datetime.now(timezone.utc),
    )

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, existing_edges)

    assert len(invalidated_edges) == 1
    assert invalidated_edges[0].uuid == 'e3'


@pytest.mark.asyncio
@pytest.mark.integration
async def test_get_edge_contradictions_no_effect():
    existing_edges, nodes = create_complex_test_data()

    # Create a new edge that doesn't invalidate any existing edges
    new_edge = EntityEdge(
        uuid='e8',
        source_node_uuid='3',
        target_node_uuid='4',
        name='APPLIED_TO',
        fact='Charlie applied to Company XYZ',
        group_id='1',
        created_at=datetime.now(timezone.utc),
    )

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, existing_edges)

    assert len(invalidated_edges) == 0


@pytest.mark.asyncio
@pytest.mark.integration
async def test_invalidate_edges_partial_update():
    existing_edges, nodes = create_complex_test_data()

    # Create a new edge that partially updates an existing one
    new_edge = EntityEdge(
        uuid='e9',
        source_node_uuid='2',
        target_node_uuid='4',
        name='CHANGED_POSITION',
        fact='Bob changed his position at Company XYZ',
        group_id='1',
        created_at=datetime.now(timezone.utc),
    )

    invalidated_edges = await get_edge_contradictions(setup_llm_client(), new_edge, existing_edges)

    assert len(invalidated_edges) == 0  # The existing edge is not invalidated, just updated


def create_data_for_temporal_extraction() -> tuple[EpisodicNode, list[EpisodicNode]]:
    now = datetime.now(timezone.utc)

    previous_episodes = [
        EpisodicNode(
            name='Previous Episode 1',
            content='Bob: I work at XYZ company',
            created_at=now - timedelta(days=2),
            valid_at=now - timedelta(days=2),
            source=EpisodeType.message,
            source_description='Test previous episode for unit testing',
            group_id='1',
        ),
        EpisodicNode(
            name='Previous Episode 2',
            content="Alice: That's really cool!",
            created_at=now - timedelta(days=1),
            valid_at=now - timedelta(days=1),
            source=EpisodeType.message,
            source_description='Test previous episode for unit testing',
            group_id='1',
        ),
    ]

    episode = EpisodicNode(
        name='Previous Episode',
        content='Bob: It was cool, but I no longer work at company XYZ',
        created_at=now,
        valid_at=now,
        source=EpisodeType.message,
        source_description='Test previous episode for unit testing',
        group_id='1',
    )

    return episode, previous_episodes


@pytest.mark.asyncio
@pytest.mark.integration
async def test_extract_edge_dates():
    episode, previous_episodes = create_data_for_temporal_extraction()

    # Create a new edge that partially updates an existing one
    new_edge = EntityEdge(
        uuid='e9',
        source_node_uuid='2',
        target_node_uuid='4',
        name='LEFT_JOB',
        fact='Bob no longer works at Company XYZ',
        group_id='1',
        created_at=datetime.now(timezone.utc),
    )

    valid_at, invalid_at = await extract_edge_dates(
        setup_llm_client(), new_edge, episode, previous_episodes
    )

    assert valid_at == episode.valid_at
    assert invalid_at is None


# Run the tests
if __name__ == '__main__':
    pytest.main([__file__])