deer-flow/backend/tests/test_memory_updater.py

from unittest.mock import MagicMock, patch

from deerflow.agents.memory.prompt import format_conversation_for_update
from deerflow.agents.memory.updater import MemoryUpdater, _extract_text
from deerflow.config.memory_config import MemoryConfig


def _make_memory(facts: list[dict[str, object]] | None = None) -> dict[str, object]:
    return {
        "version": "1.0",
        "lastUpdated": "",
        "user": {
            "workContext": {"summary": "", "updatedAt": ""},
            "personalContext": {"summary": "", "updatedAt": ""},
            "topOfMind": {"summary": "", "updatedAt": ""},
        },
        "history": {
            "recentMonths": {"summary": "", "updatedAt": ""},
            "earlierContext": {"summary": "", "updatedAt": ""},
            "longTermBackground": {"summary": "", "updatedAt": ""},
        },
        "facts": facts or [],
    }


def _memory_config(**overrides: object) -> MemoryConfig:
    config = MemoryConfig()
    for key, value in overrides.items():
        setattr(config, key, value)
    return config


def test_apply_updates_skips_existing_duplicate_and_preserves_removals() -> None:
    updater = MemoryUpdater()
    current_memory = _make_memory(
        facts=[
            {
                "id": "fact_existing",
                "content": "User likes Python",
                "category": "preference",
                "confidence": 0.9,
                "createdAt": "2026-03-18T00:00:00Z",
                "source": "thread-a",
            },
            {
                "id": "fact_remove",
                "content": "Old context to remove",
                "category": "context",
                "confidence": 0.8,
                "createdAt": "2026-03-18T00:00:00Z",
                "source": "thread-a",
            },
        ]
    )
    update_data = {
        "factsToRemove": ["fact_remove"],
        "newFacts": [
            {"content": "User likes Python", "category": "preference", "confidence": 0.95},
        ],
    }

    with patch(
        "deerflow.agents.memory.updater.get_memory_config",
        return_value=_memory_config(max_facts=100, fact_confidence_threshold=0.7),
    ):
        result = updater._apply_updates(current_memory, update_data, thread_id="thread-b")

    assert [fact["content"] for fact in result["facts"]] == ["User likes Python"]
    assert all(fact["id"] != "fact_remove" for fact in result["facts"])


def test_apply_updates_skips_same_batch_duplicates_and_keeps_source_metadata() -> None:
    updater = MemoryUpdater()
    current_memory = _make_memory()
    update_data = {
        "newFacts": [
            {"content": "User prefers dark mode", "category": "preference", "confidence": 0.91},
            {"content": "User prefers dark mode", "category": "preference", "confidence": 0.92},
            {"content": "User works on DeerFlow", "category": "context", "confidence": 0.87},
        ],
    }

    with patch(
        "deerflow.agents.memory.updater.get_memory_config",
        return_value=_memory_config(max_facts=100, fact_confidence_threshold=0.7),
    ):
        result = updater._apply_updates(current_memory, update_data, thread_id="thread-42")

    assert [fact["content"] for fact in result["facts"]] == [
        "User prefers dark mode",
        "User works on DeerFlow",
    ]
    assert all(fact["id"].startswith("fact_") for fact in result["facts"])
    assert all(fact["source"] == "thread-42" for fact in result["facts"])


def test_apply_updates_preserves_threshold_and_max_facts_trimming() -> None:
    updater = MemoryUpdater()
    current_memory = _make_memory(
        facts=[
            {
                "id": "fact_python",
                "content": "User likes Python",
                "category": "preference",
                "confidence": 0.95,
                "createdAt": "2026-03-18T00:00:00Z",
                "source": "thread-a",
            },
            {
                "id": "fact_dark_mode",
                "content": "User prefers dark mode",
                "category": "preference",
                "confidence": 0.8,
                "createdAt": "2026-03-18T00:00:00Z",
                "source": "thread-a",
            },
        ]
    )
    update_data = {
        "newFacts": [
            {"content": "User prefers dark mode", "category": "preference", "confidence": 0.9},
            {"content": "User uses uv", "category": "context", "confidence": 0.85},
            {"content": "User likes noisy logs", "category": "behavior", "confidence": 0.6},
        ],
    }

    with patch(
        "deerflow.agents.memory.updater.get_memory_config",
        return_value=_memory_config(max_facts=2, fact_confidence_threshold=0.7),
    ):
        result = updater._apply_updates(current_memory, update_data, thread_id="thread-9")

    assert [fact["content"] for fact in result["facts"]] == [
        "User likes Python",
        "User uses uv",
    ]
    assert all(fact["content"] != "User likes noisy logs" for fact in result["facts"])
    assert result["facts"][1]["source"] == "thread-9"


# ---------------------------------------------------------------------------
# _extract_text — LLM response content normalization
# ---------------------------------------------------------------------------


class TestExtractText:
    """_extract_text should normalize all content shapes to plain text."""

    def test_string_passthrough(self):
        assert _extract_text("hello world") == "hello world"

    def test_list_single_text_block(self):
        assert _extract_text([{"type": "text", "text": "hello"}]) == "hello"

    def test_list_multiple_text_blocks_joined(self):
        content = [
            {"type": "text", "text": "part one"},
            {"type": "text", "text": "part two"},
        ]
        assert _extract_text(content) == "part one\npart two"

    def test_list_plain_strings(self):
        assert _extract_text(["raw string"]) == "raw string"

    def test_list_string_chunks_join_without_separator(self):
        content = ['{"user"', ': "alice"}']
        assert _extract_text(content) == '{"user": "alice"}'

    def test_list_mixed_strings_and_blocks(self):
        content = [
            "raw text",
            {"type": "text", "text": "block text"},
        ]
        assert _extract_text(content) == "raw text\nblock text"

    def test_list_adjacent_string_chunks_then_block(self):
        content = [
            "prefix",
            "-continued",
            {"type": "text", "text": "block text"},
        ]
        assert _extract_text(content) == "prefix-continued\nblock text"

    def test_list_skips_non_text_blocks(self):
        content = [
            {"type": "image_url", "image_url": {"url": "http://img.png"}},
            {"type": "text", "text": "actual text"},
        ]
        assert _extract_text(content) == "actual text"

    def test_empty_list(self):
        assert _extract_text([]) == ""

    def test_list_no_text_blocks(self):
        assert _extract_text([{"type": "image_url", "image_url": {}}]) == ""

    def test_non_str_non_list(self):
        assert _extract_text(42) == "42"


# ---------------------------------------------------------------------------
# format_conversation_for_update — handles mixed list content
# ---------------------------------------------------------------------------


class TestFormatConversationForUpdate:
    def test_plain_string_messages(self):
        human_msg = MagicMock()
        human_msg.type = "human"
        human_msg.content = "What is Python?"

        ai_msg = MagicMock()
        ai_msg.type = "ai"
        ai_msg.content = "Python is a programming language."

        result = format_conversation_for_update([human_msg, ai_msg])
        assert "User: What is Python?" in result
        assert "Assistant: Python is a programming language." in result

    def test_list_content_with_plain_strings(self):
        """Plain strings in list content should not be lost."""
        msg = MagicMock()
        msg.type = "human"
        msg.content = ["raw user text", {"type": "text", "text": "structured text"}]

        result = format_conversation_for_update([msg])
        assert "raw user text" in result
        assert "structured text" in result


# ---------------------------------------------------------------------------
# update_memory — structured LLM response handling
# ---------------------------------------------------------------------------


class TestUpdateMemoryStructuredResponse:
    """update_memory should handle LLM responses returned as list content blocks."""

    def _make_mock_model(self, content):
        model = MagicMock()
        response = MagicMock()
        response.content = content
        model.invoke.return_value = response
        return model

    def test_string_response_parses(self):
        updater = MemoryUpdater()
        valid_json = '{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}'

        with (
            patch.object(updater, "_get_model", return_value=self._make_mock_model(valid_json)),
            patch("deerflow.agents.memory.updater.get_memory_config", return_value=_memory_config(enabled=True)),
            patch("deerflow.agents.memory.updater.get_memory_data", return_value=_make_memory()),
            patch("deerflow.agents.memory.updater._save_memory_to_file", return_value=True),
        ):
            msg = MagicMock()
            msg.type = "human"
            msg.content = "Hello"
            ai_msg = MagicMock()
            ai_msg.type = "ai"
            ai_msg.content = "Hi there"
            ai_msg.tool_calls = []
            result = updater.update_memory([msg, ai_msg])

        assert result is True

    def test_list_content_response_parses(self):
        """LLM response as list-of-blocks should be extracted, not repr'd."""
        updater = MemoryUpdater()
        valid_json = '{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}'
        list_content = [{"type": "text", "text": valid_json}]

        with (
            patch.object(updater, "_get_model", return_value=self._make_mock_model(list_content)),
            patch("deerflow.agents.memory.updater.get_memory_config", return_value=_memory_config(enabled=True)),
            patch("deerflow.agents.memory.updater.get_memory_data", return_value=_make_memory()),
            patch("deerflow.agents.memory.updater._save_memory_to_file", return_value=True),
        ):
            msg = MagicMock()
            msg.type = "human"
            msg.content = "Hello"
            ai_msg = MagicMock()
            ai_msg.type = "ai"
            ai_msg.content = "Hi"
            ai_msg.tool_calls = []
            result = updater.update_memory([msg, ai_msg])

        assert result is True