tests/mem0_agentscope_adapter_test.py

# -*- coding: utf-8 -*-
# pylint: disable=wrong-import-order
"""Unit tests for the AgentScope ↔ mem0 adapters.

Verifies that ``AgentScopeLLM`` / ``AgentScopeEmbedding`` correctly
translate between mem0's sync OpenAI-style contract and AgentScope's
async ``Msg`` / ``ContentBlock`` / ``EmbeddingResponse`` shapes, and
that ``register_with_mem0`` plugs them into mem0's factories.
"""
import asyncio
import json
import unittest
from typing import Any

from agentscope.credential import DashScopeCredential
from agentscope.embedding import EmbeddingModelBase, EmbeddingResponse
from agentscope.message import (
    Msg,
    TextBlock,
    ThinkingBlock,
    ToolCallBlock,
)
from agentscope.middleware._longterm_memory._mem0._agentscope_adapter import (
    AgentScopeEmbedding,
    AgentScopeLLM,
    _convert_messages_to_agentscope,
    _parse_chat_response,
    build_mem0_config,
)
from agentscope.model import ChatResponse
from utils import MockModel


# ----------------------------------------------------------------------
# Pure helpers
# ----------------------------------------------------------------------


class TestConvertMessages(unittest.TestCase):
    """Tests for converting mem0 dict messages into AgentScope messages."""

    def test_three_roles_map_to_correct_role(self) -> None:
        """System, user, and assistant roles should be preserved."""
        msgs = _convert_messages_to_agentscope(
            [
                {"role": "system", "content": "you are helpful"},
                {"role": "user", "content": "hi"},
                {"role": "assistant", "content": "hello"},
            ],
        )
        self.assertEqual(
            [m.role for m in msgs],
            ["system", "user", "assistant"],
        )
        self.assertTrue(all(isinstance(m, Msg) for m in msgs))
        self.assertEqual(msgs[1].get_text_content(), "hi")

    def test_unknown_role_dropped(self) -> None:
        """Unsupported message roles should be skipped."""
        msgs = _convert_messages_to_agentscope(
            [
                {"role": "tool", "content": "noise"},
                {"role": "user", "content": "real"},
            ],
        )
        self.assertEqual(len(msgs), 1)
        self.assertEqual(msgs[0].role, "user")


class TestParseChatResponse(unittest.TestCase):
    """Tests for converting AgentScope chat responses into mem0 output."""

    def _resp(self, blocks: list) -> ChatResponse:
        """Build a final chat response from content blocks."""
        return ChatResponse(content=blocks, is_last=True)

    def test_text_only_returns_string(self) -> None:
        """Plain text responses should become plain strings."""
        resp = self._resp([TextBlock(type="text", text="answer")])
        self.assertEqual(_parse_chat_response(resp, has_tool=False), "answer")

    def test_thinking_prefixed_to_text(self) -> None:
        """Thinking blocks should be preserved before visible text."""
        resp = self._resp(
            [
                ThinkingBlock(type="thinking", thinking="hmm"),
                TextBlock(type="text", text="final"),
            ],
        )
        out = _parse_chat_response(resp, has_tool=False)
        self.assertIn("hmm", out)
        self.assertIn("final", out)
        # thinking comes first to mirror v1 order
        self.assertLess(out.index("hmm"), out.index("final"))

    def test_tool_call_with_json_string_input(self) -> None:
        """v2 stores ToolCallBlock.input as a JSON string — adapter
        must parse it back to a dict for mem0."""
        resp = self._resp(
            [
                TextBlock(type="text", text="calling tool"),
                ToolCallBlock(
                    type="tool_call",
                    id="call_1",
                    name="lookup",
                    input=json.dumps({"q": "alice"}),
                ),
            ],
        )
        out = _parse_chat_response(resp, has_tool=True)
        self.assertEqual(out["content"], "calling tool")
        self.assertEqual(
            out["tool_calls"],
            [{"name": "lookup", "arguments": {"q": "alice"}}],
        )

    def test_tool_call_with_malformed_input_keeps_raw(self) -> None:
        """Malformed JSON tool inputs should remain as raw strings."""
        resp = self._resp(
            [
                ToolCallBlock(
                    type="tool_call",
                    id="call_2",
                    name="lookup",
                    input="not json",
                ),
            ],
        )
        out = _parse_chat_response(resp, has_tool=True)
        self.assertEqual(out["tool_calls"][0]["arguments"], "not json")

    def test_empty_content(self) -> None:
        """Empty responses should convert to the empty mem0 shapes."""
        resp = self._resp([])
        self.assertEqual(_parse_chat_response(resp, has_tool=False), "")
        self.assertEqual(
            _parse_chat_response(resp, has_tool=True),
            {"content": "", "tool_calls": []},
        )


# ----------------------------------------------------------------------
# AgentScopeLLM end-to-end (fake AgentScope model on caller event loop)
# ----------------------------------------------------------------------


class _CurrentEventLoopTestCase(unittest.TestCase):
    """Provides a current event loop for the adapter's sync bridge."""

    _event_loop: asyncio.AbstractEventLoop
    _previous_event_loop: asyncio.AbstractEventLoop | None

    def setUp(self) -> None:
        """Install a fresh event loop for each sync-bridge test."""
        super().setUp()
        try:
            self._previous_event_loop = asyncio.get_event_loop()
        except RuntimeError:
            self._previous_event_loop = None
        self._event_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self._event_loop)

    def tearDown(self) -> None:
        """Restore the previous event loop after each test."""
        asyncio.set_event_loop(self._previous_event_loop)
        self._event_loop.close()
        super().tearDown()


class _RecordingMockChatModel(MockModel):
    """Captures the ``messages`` arg so we can assert the v2 Msg
    objects mem0's dict messages were converted into."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        """Initialize the recording model."""
        super().__init__(*args, **kwargs)
        self.received_messages: list[list[Msg]] = []

    async def _call_api(self, *args: Any, **kwargs: Any) -> Any:
        """Record delivered messages before delegating to MockModel."""
        self.received_messages.append(list(kwargs.get("messages") or []))
        return await super()._call_api(*args, **kwargs)


class TestAgentScopeLLM(_CurrentEventLoopTestCase):
    """End-to-end tests for the mem0 LLM adapter."""

    def test_constructor_rejects_non_chatmodel(self) -> None:
        """The LLM adapter should reject non-AgentScope chat models."""
        with self.assertRaises(TypeError):
            AgentScopeLLM(config={"model": object()})

    def test_constructor_requires_model(self) -> None:
        """The LLM adapter should require a model config entry."""
        with self.assertRaises(ValueError):
            AgentScopeLLM(config={})

    def test_generate_response_routes_through_agentscope_model(self) -> None:
        """mem0 generation should call the wrapped AgentScope model."""
        model = _RecordingMockChatModel()
        model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(type="text", text="from agentscope")],
                    is_last=True,
                ),
            ],
        )
        llm = AgentScopeLLM(config={"model": model})

        result = llm.generate_response(
            [
                {"role": "system", "content": "sys"},
                {"role": "user", "content": "hello"},
            ],
        )
        self.assertEqual(result, "from agentscope")
        # The dict messages were converted to Msg objects with the
        # correct roles preserved.
        delivered = model.received_messages[0]
        self.assertEqual([m.role for m in delivered], ["system", "user"])

    def test_generate_response_with_tools(self) -> None:
        """Tool-call responses should be converted to mem0's tool shape."""
        model = _RecordingMockChatModel()
        model.set_responses(
            [
                ChatResponse(
                    content=[
                        ToolCallBlock(
                            type="tool_call",
                            id="call_3",
                            name="search",
                            input=json.dumps({"q": "x"}),
                        ),
                    ],
                    is_last=True,
                ),
            ],
        )
        llm = AgentScopeLLM(config={"model": model})

        result = llm.generate_response(
            [{"role": "user", "content": "find x"}],
            tools=[{"name": "search"}],
        )
        self.assertIsInstance(result, dict)
        self.assertEqual(
            result["tool_calls"],
            [{"name": "search", "arguments": {"q": "x"}}],
        )

    def test_streaming_response_drained_to_last_chunk(self) -> None:
        """When the AgentScope model streams, the adapter must
        consume all chunks and use the last (which carries the
        complete content per AgentScope's streaming contract)."""
        model = _RecordingMockChatModel()
        # MockModel.set_responses with a list-of-list triggers stream mode
        model.set_responses(
            [
                [
                    ChatResponse(
                        content=[TextBlock(type="text", text="part 1")],
                        is_last=False,
                    ),
                    ChatResponse(
                        content=[TextBlock(type="text", text="final")],
                        is_last=True,
                    ),
                ],
            ],
        )
        llm = AgentScopeLLM(config={"model": model})
        result = llm.generate_response(
            [{"role": "user", "content": "stream"}],
        )
        self.assertEqual(result, "final")

    def test_empty_messages_raises(self) -> None:
        """Dropping all unsupported messages should raise ValueError."""
        llm = AgentScopeLLM(config={"model": _RecordingMockChatModel()})
        with self.assertRaises(ValueError):
            llm.generate_response([{"role": "tool", "content": "ignored"}])


# ----------------------------------------------------------------------
# AgentScopeEmbedding
# ----------------------------------------------------------------------


class _FakeEmbeddingModel(EmbeddingModelBase):
    """Minimal embedding model that records requested texts."""

    def __init__(self) -> None:
        """Initialize the fake embedding model."""
        super().__init__(
            credential=DashScopeCredential(api_key="fake"),
            model="fake-embed",
            parameters=self.Parameters(),
            context_size=8192,
            batch_size=10,
            max_retries=0,
            retry_delay=0.0,
            dimensions=3,
        )
        self.received: list[list[str]] = []

    async def _call_api(
        self,
        inputs: list[str],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Return a fixed vector for every input text."""
        self.received.append(list(inputs))
        return EmbeddingResponse(
            embeddings=[[0.1, 0.2, 0.3] for _ in inputs],
            source="api",
        )


class TestAgentScopeEmbedding(_CurrentEventLoopTestCase):
    """Tests for the mem0 embedding adapter."""

    def test_constructor_validation(self) -> None:
        """The embedding adapter should validate model config."""
        with self.assertRaises(ValueError):
            AgentScopeEmbedding(config={})
        with self.assertRaises(TypeError):
            AgentScopeEmbedding(config={"model": object()})

    def test_embed_single_string(self) -> None:
        """Single-string inputs should be wrapped before model calls."""
        model = _FakeEmbeddingModel()
        emb = AgentScopeEmbedding(config={"model": model})
        result = emb.embed("hello")
        self.assertEqual(result, [0.1, 0.2, 0.3])
        # The string was wrapped into a list before reaching the model.
        self.assertEqual(model.received, [["hello"]])

    def test_embed_list_of_strings_returns_first(self) -> None:
        """mem0's contract is that ``embed`` returns ONE vector. We
        return the first."""
        model = _FakeEmbeddingModel()
        emb = AgentScopeEmbedding(config={"model": model})
        result = emb.embed(["a", "b"])
        self.assertEqual(result, [0.1, 0.2, 0.3])
        self.assertEqual(model.received, [["a", "b"]])


# ----------------------------------------------------------------------
# Factory registration
# ----------------------------------------------------------------------


class TestBuildMem0Config(unittest.TestCase):
    """``build_mem0_config`` must bypass mem0's hardcoded provider
    whitelist and emit a config that names the AgentScope adapter."""

    def test_models_only_produces_fresh_config(self) -> None:
        """Models-only construction should create AgentScope providers."""
        chat = MockModel()
        emb = _FakeEmbeddingModel()
        cfg = build_mem0_config(chat_model=chat, embedding_model=emb)

        self.assertEqual(cfg.llm.provider, "agentscope")
        self.assertEqual(cfg.embedder.provider, "agentscope")
        self.assertIs(cfg.llm.config["model"], chat)
        self.assertIs(cfg.embedder.config["model"], emb)

    def test_models_only_requires_both(self) -> None:
        """Models-only construction should require chat and embedding."""
        with self.assertRaises(ValueError):
            build_mem0_config(chat_model=MockModel())
        with self.assertRaises(ValueError):
            build_mem0_config(embedding_model=_FakeEmbeddingModel())
        with self.assertRaises(ValueError):
            build_mem0_config()

    def test_base_config_with_both_models_preserves_other_fields(
        self,
    ) -> None:
        """When a ``mem0_config`` base is given, vector_store /
        history_db / etc. should survive — only .llm and .embedder
        are rewired to the AgentScope adapters."""
        from mem0.configs.base import MemoryConfig

        base = MemoryConfig(history_db_path="/tmp/custom_history.db")
        original_vs = base.vector_store

        chat = MockModel()
        emb = _FakeEmbeddingModel()
        cfg = build_mem0_config(
            chat_model=chat,
            embedding_model=emb,
            mem0_config=base,
        )

        self.assertEqual(cfg.llm.provider, "agentscope")
        self.assertIs(cfg.llm.config["model"], chat)
        self.assertEqual(cfg.embedder.provider, "agentscope")
        self.assertIs(cfg.embedder.config["model"], emb)
        # Non-llm/embedder fields preserved.
        self.assertEqual(cfg.history_db_path, "/tmp/custom_history.db")
        self.assertIs(cfg.vector_store, original_vs)

    def test_base_config_with_only_chat_model_partial_override(
        self,
    ) -> None:
        """Partial override: chat_model alone replaces .llm but
        leaves .embedder untouched (whatever the base config had)."""
        from mem0.configs.base import MemoryConfig

        base = MemoryConfig()  # base has the default openai embedder
        original_embedder = base.embedder

        cfg = build_mem0_config(
            chat_model=MockModel(),
            mem0_config=base,
        )
        self.assertEqual(cfg.llm.provider, "agentscope")
        # Embedder unchanged from base — still mem0's openai default.
        self.assertIs(cfg.embedder, original_embedder)

    def test_base_config_alone_is_pass_through(self) -> None:
        """``mem0_config=base`` with no models is just a pass-through —
        registration still happens (cheap) but no fields change."""
        from mem0.configs.base import MemoryConfig

        base = MemoryConfig()
        cfg = build_mem0_config(mem0_config=base)
        self.assertIs(cfg, base)

    def test_registers_adapter_in_factory(self) -> None:
        """The helper side-effect: mem0's factory dicts now know how
        to resolve provider='agentscope'."""
        from mem0.utils.factory import EmbedderFactory, LlmFactory

        build_mem0_config(
            chat_model=MockModel(),
            embedding_model=_FakeEmbeddingModel(),
        )
        self.assertIn("agentscope", LlmFactory.provider_to_class)
        self.assertIn("agentscope", EmbedderFactory.provider_to_class)

    def test_naive_from_config_path_still_rejected(self) -> None:
        """Sanity check: confirm WHY the helper exists — calling
        ``MemoryConfig`` with ``provider='agentscope'`` directly DOES
        raise, which is the failure ``build_mem0_config`` works around."""
        from mem0.configs.base import MemoryConfig
        from pydantic import ValidationError

        with self.assertRaises(ValidationError):
            MemoryConfig(
                llm={
                    "provider": "agentscope",
                    "config": {"model": MockModel()},
                },
                embedder={
                    "provider": "agentscope",
                    "config": {"model": _FakeEmbeddingModel()},
                },
            )