ducklm/tests/smoke/test_memory_store.py

from fastapi.testclient import TestClient

from duck_core.api import create_app
from duck_core.context_builder import ContextBuilder
from duck_core.memory.store import MemoryStore
from duck_core.tasks.state import TaskState


def test_memory_api_stores_workspace_scoped_notes(tmp_path, monkeypatch):
    monkeypatch.setenv("DUCK_DB_PATH", str(tmp_path / "duck.sqlite3"))
    client = TestClient(create_app())

    first = client.post(
        "/v1/memory",
        json={
            "text": "User prefers concise Russian answers.",
            "workspace": "/tmp/project-a",
            "conversation_id": "chat_a",
            "memory_type": "preference",
            "importance": 0.8,
        },
    ).json()
    client.post(
        "/v1/memory",
        json={
            "text": "Different workspace note.",
            "workspace": "/tmp/project-b",
            "conversation_id": "chat_b",
        },
    )

    listed = client.get("/v1/memory", params={"workspace": "/tmp/project-a"}).json()
    search = client.get(
        "/v1/memory/search",
        params={"q": "concise Russian", "workspace": "/tmp/project-a"},
    ).json()

    assert first["memory_id"].startswith("mem_")
    assert [item["text"] for item in listed["results"]] == [
        "User prefers concise Russian answers."
    ]
    assert search["results"][0]["memory_id"] == first["memory_id"]
    assert search["results"][0]["memory_type"] == "preference"
    assert "Different workspace note." not in str(search)


async def test_memory_store_searches_text_and_metadata(tmp_path):
    store = MemoryStore(str(tmp_path / "duck.sqlite3"))
    await store.init()
    await store.add(
        text="RX580 should use Vulkan builds.",
        workspace="/tmp/duck",
        conversation_id="chat_runtime",
        memory_type="fact",
        importance=0.9,
        metadata={"topic": "gpu"},
    )

    results = await store.search("vulkan", workspace="/tmp/duck")

    assert len(results) == 1
    assert results[0].workspace == "/tmp/duck"
    assert results[0].metadata["topic"] == "gpu"


async def test_memory_store_returns_relevant_global_workspace_and_chat_memory(tmp_path):
    store = MemoryStore(str(tmp_path / "duck.sqlite3"))
    await store.init()
    await store.add("Global preference", scope="global", workspace="")
    await store.add("Workspace fact", scope="workspace", workspace="/tmp/duck")
    await store.add(
        "Chat fact",
        scope="conversation",
        workspace="/tmp/duck",
        conversation_id="chat_1",
    )
    await store.add("Other workspace fact", scope="workspace", workspace="/tmp/other")

    results = await store.relevant(
        workspace="/tmp/duck", conversation_id="chat_1", query="fact"
    )

    assert [record.text for record in results] == [
        "Global preference",
        "Chat fact",
        "Workspace fact",
    ]


def test_context_builder_injects_memory_context_before_user_message():
    task = TaskState(
        task_id="task_1",
        status="running",
        user_message="Что помнить?",
        workspace="/tmp/duck",
        debug=True,
        created_at="now",
        updated_at="now",
    )

    messages = ContextBuilder().build_basic_messages(
        task,
        memory_records=[
            {"scope": "global", "text": "Use Russian."},
            {"scope": "workspace", "text": "DuckLM uses Vulkan."},
        ],
    )

    assert messages[0]["role"] == "system"
    assert "Relevant memory" in messages[0]["content"]
    assert "global: Use Russian." in messages[0]["content"]
    assert messages[-1]["content"] == "Что помнить?"


def test_chat_api_injects_relevant_memory_into_model_context(tmp_path, monkeypatch):
    monkeypatch.setenv("DUCK_DB_PATH", str(tmp_path / "duck.sqlite3"))
    seen_messages = []

    async def fake_chat(self, role, messages, temperature=None, max_output_tokens=None, response_format=None):
        if role == "action":
            return ModelResponse(
                role=role,
                model="local-main",
                content='{"kind":"action_directive","intent":"answer","risk_level":"none","actions":[]}',
                reasoning_content=None,
                raw={},
                latency_ms=1.0,
            )
        seen_messages.append(messages)
        return ModelResponse(
            role=role,
            model="local-main",
            content="remembered",
            reasoning_content=None,
            raw={},
            latency_ms=1.0,
        )

    from duck_core.model_client import ModelResponse

    monkeypatch.setattr("duck_core.model_client.ModelClient.chat", fake_chat)
    client = TestClient(create_app())
    conversation = client.post(
        "/v1/conversations",
        json={"title": "Memory chat", "workspace": str(tmp_path)},
    ).json()
    client.post(
        "/v1/memory",
        json={
            "text": "User prefers direct Russian answers.",
            "workspace": str(tmp_path),
            "conversation_id": conversation["conversation_id"],
        },
    )

    client.post(
        "/v1/chat",
        json={
            "conversation_id": conversation["conversation_id"],
            "message": "Как отвечать?",
        },
    )

    assert seen_messages[0][0]["role"] == "system"
    assert "User prefers direct Russian answers." in seen_messages[0][0]["content"]