getzep--graphiti/tests/test_add_triplet.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 unittest.mock import AsyncMock, Mock, patch

import pytest

from graphiti_core.cross_encoder.client import CrossEncoderClient
from graphiti_core.edges import EntityEdge
from graphiti_core.graphiti import Graphiti
from graphiti_core.llm_client import LLMClient
from graphiti_core.nodes import EntityNode
from tests.helpers_test import group_id

pytest_plugins = ('pytest_asyncio', 'tests.helpers_test')


@pytest.fixture
def mock_llm_client():
    """Create a mock LLM"""
    mock_llm = Mock(spec=LLMClient)
    mock_llm.config = Mock()
    mock_llm.model = 'test-model'
    mock_llm.small_model = 'test-small-model'
    mock_llm.temperature = 0.0
    mock_llm.max_tokens = 1000
    mock_llm.cache_enabled = False
    mock_llm.cache_dir = None

    # Mock the public method that's actually called
    mock_llm.generate_response = AsyncMock()
    mock_llm.generate_response.return_value = {
        'duplicate_facts': [],
        'invalidate_facts': [],
    }

    return mock_llm


@pytest.fixture
def mock_cross_encoder_client():
    """Create a mock cross encoder"""
    mock_ce = Mock(spec=CrossEncoderClient)
    mock_ce.config = Mock()
    mock_ce.rerank = AsyncMock()
    mock_ce.rerank.return_value = []

    return mock_ce


@pytest.mark.asyncio
async def test_add_triplet_merges_attributes(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that attributes are merged (not replaced) when adding a triplet."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing node with some attributes
    existing_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Existing summary',
        attributes={'age': 30, 'city': 'New York'},
    )
    await existing_source.generate_name_embedding(mock_embedder)
    await existing_source.save(graph_driver)

    # Create a user-provided node with additional attributes
    user_source = EntityNode(
        uuid=existing_source.uuid,  # Same UUID to trigger direct lookup
        name='Alice',
        group_id=group_id,
        labels=['Person', 'Employee'],
        created_at=now,
        summary='Updated summary',
        attributes={'age': 31, 'department': 'Engineering'},  # age updated, department added
    )

    # Create target node
    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={'age': 25},
    )

    # Create edge
    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='WORKS_WITH',
        fact='Alice works with Bob',
        group_id=group_id,
        created_at=now,
    )

    # Mock the search functions to return empty results
    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (edge, [], [])

        await graphiti.add_triplet(user_source, edge, user_target)

        # Verify attributes were merged (not replaced)
        # The resolved node should have both existing and new attributes
        retrieved_source = await EntityNode.get_by_uuid(graph_driver, existing_source.uuid)
        assert 'age' in retrieved_source.attributes
        assert retrieved_source.attributes['age'] == 31  # Updated value
        assert retrieved_source.attributes['city'] == 'New York'  # Preserved
        assert retrieved_source.attributes['department'] == 'Engineering'  # Added
        assert retrieved_source.summary == 'Updated summary'


@pytest.mark.asyncio
async def test_add_triplet_updates_summary(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that summary is updated when provided by user."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing node with a summary
    existing_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Old summary',
        attributes={},
    )
    await existing_target.generate_name_embedding(mock_embedder)
    await existing_target.save(graph_driver)

    # Create user-provided nodes
    user_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={},
    )

    user_target = EntityNode(
        uuid=existing_target.uuid,
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='New summary for Bob',
        attributes={},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (edge, [], [])

        await graphiti.add_triplet(user_source, edge, user_target)

        # Verify summary was updated
        retrieved_target = await EntityNode.get_by_uuid(graph_driver, existing_target.uuid)
        assert retrieved_target.summary == 'New summary for Bob'


@pytest.mark.asyncio
async def test_add_triplet_updates_labels(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that labels are updated when provided by user."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing node with labels
    existing_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='',
        attributes={},
    )
    await existing_source.generate_name_embedding(mock_embedder)
    await existing_source.save(graph_driver)

    # Create user-provided node with different labels
    user_source = EntityNode(
        uuid=existing_source.uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person', 'Employee', 'Manager'],
        created_at=now,
        summary='',
        attributes={},
    )

    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='',
        attributes={},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='MANAGES',
        fact='Alice manages Bob',
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (edge, [], [])

        await graphiti.add_triplet(user_source, edge, user_target)

        # Verify labels were updated
        retrieved_source = await EntityNode.get_by_uuid(graph_driver, existing_source.uuid)
        # Labels should be set to user-provided labels (not merged)
        assert set(retrieved_source.labels) == {'Person', 'Employee', 'Manager'}


@pytest.mark.asyncio
async def test_add_triplet_with_new_nodes_no_uuid(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test add_triplet with nodes that don't have UUIDs (will be resolved)."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create user-provided nodes without UUIDs
    user_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={'age': 30},
    )

    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={'age': 25},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    with patch('graphiti_core.graphiti.search') as mock_search:
        mock_search.return_value = Mock(edges=[])
        with patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge:
            mock_resolve_edge.return_value = (edge, [], [])

            result = await graphiti.add_triplet(user_source, edge, user_target)

            # Verify nodes were created with user-provided attributes
            assert len(result.nodes) >= 2
            # Find the nodes in the result
            source_in_result = next((n for n in result.nodes if n.name == 'Alice'), None)
            target_in_result = next((n for n in result.nodes if n.name == 'Bob'), None)

            if source_in_result:
                assert source_in_result.attributes.get('age') == 30
                assert source_in_result.summary == 'Alice summary'
            if target_in_result:
                assert target_in_result.attributes.get('age') == 25
                assert target_in_result.summary == 'Bob summary'


@pytest.mark.asyncio
async def test_add_triplet_preserves_existing_attributes(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that existing attributes are preserved when merging new ones."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing node with multiple attributes
    existing_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Existing summary',
        attributes={
            'age': 30,
            'city': 'New York',
            'country': 'USA',
            'email': 'alice@example.com',
        },
    )
    await existing_source.generate_name_embedding(mock_embedder)
    await existing_source.save(graph_driver)

    # Create user-provided node with only some attributes
    user_source = EntityNode(
        uuid=existing_source.uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Updated summary',
        attributes={'age': 31, 'city': 'Boston'},  # Only updating age and city
    )

    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='',
        attributes={},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (edge, [], [])

        await graphiti.add_triplet(user_source, edge, user_target)

        # Verify all attributes are preserved/updated correctly
        retrieved_source = await EntityNode.get_by_uuid(graph_driver, existing_source.uuid)
        assert retrieved_source.attributes['age'] == 31  # Updated
        assert retrieved_source.attributes['city'] == 'Boston'  # Updated
        assert retrieved_source.attributes['country'] == 'USA'  # Preserved
        assert retrieved_source.attributes['email'] == 'alice@example.com'  # Preserved
        assert retrieved_source.summary == 'Updated summary'


@pytest.mark.asyncio
async def test_add_triplet_empty_attributes_preserved(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that nodes with empty attributes don't overwrite existing attributes."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing node with attributes
    existing_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Existing summary',
        attributes={'age': 30, 'city': 'New York'},
    )
    await existing_source.generate_name_embedding(mock_embedder)
    await existing_source.save(graph_driver)

    # Create user-provided node with empty attributes
    user_source = EntityNode(
        uuid=existing_source.uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='',  # Empty summary should not overwrite
        attributes={},  # Empty attributes should not overwrite
    )

    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='',
        attributes={},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (edge, [], [])

        await graphiti.add_triplet(user_source, edge, user_target)

        # Verify existing attributes are preserved when user provides empty dict
        retrieved_source = await EntityNode.get_by_uuid(graph_driver, existing_source.uuid)
        # Empty attributes dict should not clear existing attributes
        assert 'age' in retrieved_source.attributes
        assert 'city' in retrieved_source.attributes
        # Empty summary should not overwrite existing summary
        assert retrieved_source.summary == 'Existing summary'


@pytest.mark.asyncio
async def test_add_triplet_invalid_source_uuid(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that ValueError is raised when source_node has a UUID that doesn't exist."""
    from uuid import uuid4

    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create a node with a UUID that doesn't exist in the database
    invalid_uuid = str(uuid4())
    user_source = EntityNode(
        uuid=invalid_uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={'age': 30},
    )

    user_target = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={'age': 25},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    # Should raise ValueError for invalid source UUID
    with pytest.raises(ValueError, match=f'Node with UUID {invalid_uuid} not found'):
        await graphiti.add_triplet(user_source, edge, user_target)


@pytest.mark.asyncio
async def test_add_triplet_invalid_target_uuid(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that ValueError is raised when target_node has a UUID that doesn't exist."""
    from uuid import uuid4

    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create an existing source node
    existing_source = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={'age': 30},
    )
    await existing_source.generate_name_embedding(mock_embedder)
    await existing_source.save(graph_driver)

    # Create a target node with a UUID that doesn't exist in the database
    invalid_uuid = str(uuid4())
    user_source = EntityNode(
        uuid=existing_source.uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={'age': 30},
    )

    user_target = EntityNode(
        uuid=invalid_uuid,
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={'age': 25},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    # Should raise ValueError for invalid target UUID
    with pytest.raises(ValueError, match=f'Node with UUID {invalid_uuid} not found'):
        await graphiti.add_triplet(user_source, edge, user_target)


@pytest.mark.asyncio
async def test_add_triplet_invalid_both_uuids(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that ValueError is raised for source_node first when both UUIDs are invalid."""
    from uuid import uuid4

    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create nodes with UUIDs that don't exist in the database
    invalid_source_uuid = str(uuid4())
    invalid_target_uuid = str(uuid4())

    user_source = EntityNode(
        uuid=invalid_source_uuid,
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={'age': 30},
    )

    user_target = EntityNode(
        uuid=invalid_target_uuid,
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={'age': 25},
    )

    edge = EntityEdge(
        source_node_uuid=user_source.uuid,
        target_node_uuid=user_target.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )

    # Should raise ValueError for source UUID first (source is checked before target)
    with pytest.raises(ValueError, match=f'Node with UUID {invalid_source_uuid} not found'):
        await graphiti.add_triplet(user_source, edge, user_target)


@pytest.mark.asyncio
async def test_add_triplet_edge_uuid_with_different_nodes_creates_new_edge(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that providing an edge UUID with different src/dst nodes creates a new edge."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create existing nodes: Alice and Bob
    alice = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={},
    )
    await alice.generate_name_embedding(mock_embedder)
    await alice.save(graph_driver)

    bob = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={},
    )
    await bob.generate_name_embedding(mock_embedder)
    await bob.save(graph_driver)

    # Create a third node: Charlie
    charlie = EntityNode(
        name='Charlie',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Charlie summary',
        attributes={},
    )
    await charlie.generate_name_embedding(mock_embedder)
    await charlie.save(graph_driver)

    # Create an existing edge between Alice and Bob
    existing_edge = EntityEdge(
        source_node_uuid=alice.uuid,
        target_node_uuid=bob.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )
    await existing_edge.generate_embedding(mock_embedder)
    await existing_edge.save(graph_driver)

    # Now try to add a triplet using the existing edge UUID but with different nodes (Alice -> Charlie)
    new_edge_with_same_uuid = EntityEdge(
        uuid=existing_edge.uuid,  # Reuse the existing edge's UUID
        source_node_uuid=alice.uuid,
        target_node_uuid=charlie.uuid,  # Different target!
        name='KNOWS',
        fact='Alice knows Charlie',
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        # Return the edge as-is (simulating no deduplication)
        mock_resolve_edge.return_value = (new_edge_with_same_uuid, [], [])

        result = await graphiti.add_triplet(alice, new_edge_with_same_uuid, charlie)

        # The original edge (Alice -> Bob) should still exist
        original_edge = await EntityEdge.get_by_uuid(graph_driver, existing_edge.uuid)
        assert original_edge.source_node_uuid == alice.uuid
        assert original_edge.target_node_uuid == bob.uuid
        assert original_edge.fact == 'Alice knows Bob'

        # The new edge should have a different UUID
        new_edge = result.edges[0]
        assert new_edge.uuid != existing_edge.uuid
        assert new_edge.source_node_uuid == alice.uuid
        assert new_edge.target_node_uuid == charlie.uuid


@pytest.mark.asyncio
async def test_add_triplet_edge_uuid_with_same_nodes_updates_edge(
    graph_driver, mock_llm_client, mock_embedder, mock_cross_encoder_client
):
    """Test that providing an edge UUID with same src/dst nodes allows updating the edge."""
    graphiti = Graphiti(
        graph_driver=graph_driver,
        llm_client=mock_llm_client,
        embedder=mock_embedder,
        cross_encoder=mock_cross_encoder_client,
    )

    await graphiti.build_indices_and_constraints()

    now = datetime.now()

    # Create existing nodes: Alice and Bob
    alice = EntityNode(
        name='Alice',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Alice summary',
        attributes={},
    )
    await alice.generate_name_embedding(mock_embedder)
    await alice.save(graph_driver)

    bob = EntityNode(
        name='Bob',
        group_id=group_id,
        labels=['Person'],
        created_at=now,
        summary='Bob summary',
        attributes={},
    )
    await bob.generate_name_embedding(mock_embedder)
    await bob.save(graph_driver)

    # Create an existing edge between Alice and Bob
    existing_edge = EntityEdge(
        source_node_uuid=alice.uuid,
        target_node_uuid=bob.uuid,
        name='KNOWS',
        fact='Alice knows Bob',
        group_id=group_id,
        created_at=now,
    )
    await existing_edge.generate_embedding(mock_embedder)
    await existing_edge.save(graph_driver)

    # Now update the edge with the same source/target but different fact
    updated_edge = EntityEdge(
        uuid=existing_edge.uuid,  # Reuse the existing edge's UUID
        source_node_uuid=alice.uuid,
        target_node_uuid=bob.uuid,  # Same target
        name='WORKS_WITH',
        fact='Alice works with Bob',  # Updated fact
        group_id=group_id,
        created_at=now,
    )

    with (
        patch('graphiti_core.graphiti.search') as mock_search,
        patch('graphiti_core.graphiti.resolve_extracted_edge') as mock_resolve_edge,
    ):
        mock_search.return_value = Mock(edges=[])
        mock_resolve_edge.return_value = (updated_edge, [], [])

        result = await graphiti.add_triplet(alice, updated_edge, bob)

        # The edge should keep the same UUID (update allowed)
        result_edge = result.edges[0]
        assert result_edge.uuid == existing_edge.uuid