Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 20 additions & 9 deletions astrbot/core/knowledge_base/parsers/text_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,12 @@
支持解析 TXT 和 Markdown 文件。
"""

import codecs

from astrbot.core.knowledge_base.parsers.base import BaseParser, ParseResult

_BOM_UTF16_PAIR = (codecs.BOM_UTF16_LE, codecs.BOM_UTF16_BE)


class TextParser(BaseParser):
"""TXT/MD 文本解析器
Expand All @@ -15,7 +19,9 @@ class TextParser(BaseParser):
async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
"""解析文本文件

尝试使用多种编码解析文件内容。
尝试使用多种编码解析文件内容。带 BOM 的文件(如 Windows 记事本
「UTF-8 with BOM」/「Unicode」保存的 txt)按 BOM 直接解码,避免被
无 BOM 序列误判。

Args:
file_content: 文件内容
Expand All @@ -28,15 +34,20 @@ async def parse(self, file_content: bytes, file_name: str) -> ParseResult:
ValueError: 如果无法解码文件

"""
# 尝试多种编码
for encoding in ["utf-8", "gbk", "gb2312", "gb18030"]:
try:
text = file_content.decode(encoding)
break
except UnicodeDecodeError:
continue
if file_content.startswith(codecs.BOM_UTF8):
text = file_content.decode("utf-8-sig")
elif file_content.startswith(_BOM_UTF16_PAIR):
text = file_content.decode("utf-16")
else:
raise ValueError(f"无法解码文件: {file_name}")
# 尝试多种编码(无 BOM 文件,utf-8 优先,GBK 系兜底)
for encoding in ["utf-8", "gbk", "gb2312", "gb18030"]:
try:
text = file_content.decode(encoding)
break
except UnicodeDecodeError:
continue
else:
raise ValueError(f"无法解码文件: {file_name}")

# 文本文件无多媒体资源
return ParseResult(text=text, media=[])
71 changes: 71 additions & 0 deletions tests/test_text_parser.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
from __future__ import annotations

import codecs

import pytest

from astrbot.core.knowledge_base.parsers.text_parser import TextParser


@pytest.mark.asyncio
async def test_parse_utf8() -> None:
parser = TextParser()

result = await parser.parse("你好, world".encode(), "a.txt")

assert result.text == "你好, world"
assert result.media == []


@pytest.mark.asyncio
async def test_parse_utf8_with_bom() -> None:
parser = TextParser()

result = await parser.parse(
codecs.BOM_UTF8 + "带 BOM 的内容".encode(),
"a.txt",
)

assert result.text == "带 BOM 的内容"


@pytest.mark.asyncio
async def test_parse_utf16_le_with_bom() -> None:
# Windows Notepad "Unicode" (UTF-16 LE with BOM) saved txt files.
parser = TextParser()

result = await parser.parse(
codecs.BOM_UTF16_LE + "Windows 记事本 Unicode".encode("utf-16-le"),
"a.txt",
)

assert result.text == "Windows 记事本 Unicode"


@pytest.mark.asyncio
async def test_parse_utf16_be_with_bom() -> None:
parser = TextParser()

result = await parser.parse(
codecs.BOM_UTF16_BE + "big endian".encode("utf-16-be"),
"a.txt",
)

assert result.text == "big endian"


@pytest.mark.asyncio
async def test_parse_gbk() -> None:
parser = TextParser()

result = await parser.parse("中文内容".encode("gbk"), "a.txt")

assert result.text == "中文内容"


@pytest.mark.asyncio
async def test_parse_undecodable_raises_value_error() -> None:
parser = TextParser()

with pytest.raises(ValueError, match="无法解码文件"):
await parser.parse(b"\x81\x7f\x81\x7f\xff", "a.txt")
Loading