diff --git a/astrbot/core/knowledge_base/parsers/util.py b/astrbot/core/knowledge_base/parsers/util.py index a98ed60ed2..0a017267ff 100644 --- a/astrbot/core/knowledge_base/parsers/util.py +++ b/astrbot/core/knowledge_base/parsers/util.py @@ -1,10 +1,32 @@ +from astrbot.core.exceptions import KnowledgeBaseUploadError + from .base import BaseParser +# Formats that are plain text by nature and need no external dependency. +_TEXT_PARSER_EXTS = {".md", ".txt", ".markdown"} +# Formats handled by markitdown-no-magika (an optional heavy dependency). +_MARKITDOWN_EXTS = {".rst", ".adoc", ".xlsx", ".docx", ".xls"} + +_MARKITDOWN_MISSING_HINT = ( + "文档解析失败:处理该格式需要 markitdown-no-magika 依赖," + "请先安装:pip install 'markitdown-no-magika[docx,xls,xlsx]'" +) + async def select_parser(ext: str) -> BaseParser: - if ext in {".md", ".txt", ".markdown", ".rst", ".adoc", ".xlsx", ".docx", ".xls"}: - from .markitdown_parser import MarkitdownParser + if ext in _TEXT_PARSER_EXTS: + from .text_parser import TextParser + return TextParser() + if ext in _MARKITDOWN_EXTS: + try: + from .markitdown_parser import MarkitdownParser + except ImportError as exc: + raise KnowledgeBaseUploadError( + stage="parsing", + user_message=_MARKITDOWN_MISSING_HINT, + details={"file_ext": ext}, + ) from exc return MarkitdownParser() if ext == ".epub": from .epub_parser import EpubParser diff --git a/tests/test_kb_select_parser.py b/tests/test_kb_select_parser.py new file mode 100644 index 0000000000..8856e98909 --- /dev/null +++ b/tests/test_kb_select_parser.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import sys + +import pytest + +from astrbot.core.exceptions import KnowledgeBaseUploadError +from astrbot.core.knowledge_base.parsers.markitdown_parser import MarkitdownParser +from astrbot.core.knowledge_base.parsers.text_parser import TextParser +from astrbot.core.knowledge_base.parsers.util import select_parser + + +@pytest.mark.parametrize("ext", [".txt", ".md", ".markdown"]) +@pytest.mark.asyncio +@pytest.mark.asyncio +async def test_text_formats_use_text_parser(ext: str) -> None: + parser = await select_parser(ext) + + assert isinstance(parser, TextParser) + + +@pytest.mark.parametrize("ext", [".rst", ".adoc", ".xlsx", ".docx", ".xls"]) +@pytest.mark.asyncio +@pytest.mark.asyncio +async def test_office_formats_use_markitdown_parser(ext: str) -> None: + parser = await select_parser(ext) + + assert isinstance(parser, MarkitdownParser) + + +@pytest.mark.asyncio +async def test_epub_and_pdf_routing() -> None: + from astrbot.core.knowledge_base.parsers.epub_parser import EpubParser + from astrbot.core.knowledge_base.parsers.pdf_parser import PDFParser + + assert isinstance(await select_parser(".epub"), EpubParser) + assert isinstance(await select_parser(".pdf"), PDFParser) + + +@pytest.mark.asyncio +async def test_unsupported_ext_raises() -> None: + with pytest.raises(ValueError, match="暂时不支持的文件格式"): + await select_parser(".exe") + + +@pytest.mark.asyncio +async def test_missing_markitdown_dependency_raises_clear_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Drop cached modules so the lazy import runs again, then make + # `import markitdown_no_magika` fail as if the package were absent. + for name in list(sys.modules): + if name == "markitdown_no_magika" or name.endswith("markitdown_parser"): + monkeypatch.delitem(sys.modules, name) + monkeypatch.setitem(sys.modules, "markitdown_no_magika", None) + + with pytest.raises(KnowledgeBaseUploadError, match="markitdown-no-magika"): + await select_parser(".docx")