ForcePilot/backend/test/unit/backends/test_lightrag_serialization.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace

import pytest

from server.routers import knowledge_router
from yuxi.knowledge.base import FileStatus
from yuxi.knowledge.implementations import lightrag as lightrag_module
from yuxi.knowledge.implementations.lightrag import LightRagKB


pytestmark = pytest.mark.asyncio


def _build_kb_file(path: str) -> dict:
    return {
        "status": FileStatus.PARSED,
        "markdown_file": "mock://parsed.md",
        "path": path,
        "filename": path.rsplit("/", 1)[-1],
        "processing_params": {},
    }


class _FakeDocStatus:
    async def get_by_id(self, _file_id: str) -> dict:
        return {"status": "processed"}


@pytest.fixture
def light_rag_kb(tmp_path: pytest.TempPathFactory, monkeypatch: pytest.MonkeyPatch) -> LightRagKB:
    kb = LightRagKB(str(tmp_path))
    monkeypatch.setattr(kb, "_save_metadata", _async_noop)
    monkeypatch.setattr(kb, "_persist_file", _async_noop)
    monkeypatch.setattr(kb, "_read_markdown_from_minio", _fake_read_markdown)
    monkeypatch.setattr(kb, "delete_file_chunks_only", _async_noop)
    monkeypatch.setattr(
        lightrag_module,
        "chunk_markdown",
        lambda markdown, file_id, filename, params: [{"content": markdown or f"{file_id}:{filename}:{params}"}],
    )
    return kb


async def _async_noop(*_args, **_kwargs) -> None:
    return None


async def _fake_read_markdown(*_args, **_kwargs) -> str:
    return "mock markdown"


async def test_index_file_marks_file_unparsed_when_markdown_file_is_missing(
    light_rag_kb: LightRagKB,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    persisted: list[str] = []
    db_id = "kb_missing_markdown"
    file_id = "file-missing"
    light_rag_kb.databases_meta[db_id] = {"metadata": {}}
    light_rag_kb.files_meta[file_id] = {
        "status": FileStatus.PARSED,
        "path": "/tmp/file.md",
        "filename": "file.md",
        "processing_params": {},
    }

    async def fake_persist_file(saved_file_id: str) -> None:
        persisted.append(saved_file_id)

    monkeypatch.setattr(light_rag_kb, "_persist_file", fake_persist_file)
    monkeypatch.setattr(light_rag_kb, "_get_lightrag_instance", _make_async_return(SimpleNamespace()))

    with pytest.raises(ValueError, match="no markdown_file"):
        await light_rag_kb.index_file(db_id, file_id, operator_id="user-1")

    assert light_rag_kb.files_meta[file_id]["status"] == FileStatus.UPLOADED
    assert "markdown_file" not in light_rag_kb.files_meta[file_id]
    assert light_rag_kb.files_meta[file_id]["updated_by"] == "user-1"
    assert persisted == [file_id]


async def test_index_file_serializes_writes_within_same_database(
    light_rag_kb: LightRagKB,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    db_id = "kb_same"
    light_rag_kb.databases_meta[db_id] = {"metadata": {}}
    light_rag_kb.files_meta["file-1"] = _build_kb_file("/tmp/file-1.md")
    light_rag_kb.files_meta["file-2"] = _build_kb_file("/tmp/file-2.md")

    started_first = asyncio.Event()
    release_first = asyncio.Event()
    started_second = asyncio.Event()
    call_order: list[str] = []

    async def fake_ainsert(*, ids: str, **_kwargs) -> None:
        if ids == "file-1":
            call_order.append("start:file-1")
            started_first.set()
            await release_first.wait()
            call_order.append("end:file-1")
            return
        call_order.append("start:file-2")
        started_second.set()
        call_order.append("end:file-2")

    rag = SimpleNamespace(ainsert=fake_ainsert, doc_status=_FakeDocStatus())
    monkeypatch.setattr(light_rag_kb, "_get_lightrag_instance", _make_async_return(rag))

    task1 = asyncio.create_task(light_rag_kb.index_file(db_id, "file-1"))
    await asyncio.wait_for(started_first.wait(), timeout=1)

    task2 = asyncio.create_task(light_rag_kb.index_file(db_id, "file-2"))
    await asyncio.sleep(0.05)
    assert not started_second.is_set()

    release_first.set()
    await asyncio.gather(task1, task2)

    assert call_order == ["start:file-1", "end:file-1", "start:file-2", "end:file-2"]


async def test_index_file_allows_parallel_writes_for_different_databases(
    light_rag_kb: LightRagKB,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    light_rag_kb.databases_meta["kb_a"] = {"metadata": {}}
    light_rag_kb.databases_meta["kb_b"] = {"metadata": {}}
    light_rag_kb.files_meta["file-a"] = _build_kb_file("/tmp/file-a.md")
    light_rag_kb.files_meta["file-b"] = _build_kb_file("/tmp/file-b.md")

    started_a = asyncio.Event()
    started_b = asyncio.Event()
    release_both = asyncio.Event()

    async def fake_ainsert(*, ids: str, **_kwargs) -> None:
        if ids == "file-a":
            started_a.set()
        elif ids == "file-b":
            started_b.set()
        await release_both.wait()

    rag = SimpleNamespace(ainsert=fake_ainsert, doc_status=_FakeDocStatus())
    monkeypatch.setattr(light_rag_kb, "_get_lightrag_instance", _make_async_return(rag))

    task_a = asyncio.create_task(light_rag_kb.index_file("kb_a", "file-a"))
    task_b = asyncio.create_task(light_rag_kb.index_file("kb_b", "file-b"))

    await asyncio.wait_for(started_a.wait(), timeout=1)
    await asyncio.wait_for(started_b.wait(), timeout=1)

    release_both.set()
    await asyncio.gather(task_a, task_b)


async def test_add_documents_auto_index_uses_latest_parsed_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
    calls: list[tuple[str, str]] = []

    class FakeKnowledgeBase:
        async def get_database_info(self, _db_id: str) -> dict:
            return {"name": "pytest-db"}

        async def add_file_record(
            self,
            _db_id: str,
            _item: str,
            params: dict | None = None,
            operator_id: str | None = None,
        ) -> dict:
            return {"file_id": "file-1", "status": FileStatus.UPLOADED, "params": params, "operator_id": operator_id}

        async def parse_file(self, _db_id: str, file_id: str, operator_id: str | None = None) -> dict:
            calls.append(("parse", file_id))
            return {"file_id": file_id, "status": FileStatus.PARSED, "operator_id": operator_id}

        async def update_file_params(
            self, _db_id: str, file_id: str, params: dict, operator_id: str | None = None
        ) -> None:
            calls.append(("update_params", file_id))

        async def index_file(
            self,
            _db_id: str,
            file_id: str,
            operator_id: str | None = None,
            params: dict | None = None,
        ) -> dict:
            calls.append(("index", file_id))
            return {"file_id": file_id, "status": FileStatus.INDEXED, "operator_id": operator_id, "params": params}

    class FakeTaskContext:
        async def set_message(self, _message: str) -> None:
            return None

        async def set_progress(self, _progress: float, _message: str | None = None) -> None:
            return None

        async def set_result(self, _result) -> None:
            return None

        async def raise_if_cancelled(self) -> None:
            return None

    class FakeTasker:
        async def enqueue(self, *, coroutine, **_kwargs):
            await coroutine(FakeTaskContext())
            return SimpleNamespace(id="task-1")

    async def fake_ensure_database_not_dify(_db_id: str, _operation: str) -> None:
        return None

    monkeypatch.setattr("yuxi.knowledge.utils.kb_utils.validate_file_path", lambda _item, _db_id: None)
    monkeypatch.setattr(knowledge_router, "knowledge_base", FakeKnowledgeBase())
    monkeypatch.setattr(knowledge_router, "tasker", FakeTasker())
    monkeypatch.setattr(knowledge_router, "_ensure_database_not_dify", fake_ensure_database_not_dify)

    current_user = SimpleNamespace(user_id="user-1", id="user-1")
    result = await knowledge_router.add_documents(
        "kb_auto_index",
        items=["/tmp/example.md"],
        params={"content_type": "file", "auto_index": True},
        current_user=current_user,
    )

    assert result["status"] == "queued"
    assert calls == [("parse", "file-1"), ("update_params", "file-1"), ("index", "file-1")]


def _make_async_return(value):
    async def _inner(*_args, **_kwargs):
        return value

    return _inner