From cff148860afb96b1baf0fe59d9ef9864fabe80e1 Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Fri, 10 Apr 2026 16:31:53 +0800
Subject: [PATCH] feat: enhance file reading capabilities to support PDF and
DOCX parsing, including workspace storage for long documents
---
astrbot/core/computer/file_read_utils.py | 226 ++++++++++++++++++++++-
astrbot/core/tools/computer_tools/fs.py | 26 ++-
tests/test_computer_fs_tools.py | 70 ++++++-
3 files changed, 311 insertions(+), 11 deletions(-)
diff --git a/astrbot/core/computer/file_read_utils.py b/astrbot/core/computer/file_read_utils.py
index 3c553434a..31b9e5ea0 100644
--- a/astrbot/core/computer/file_read_utils.py
+++ b/astrbot/core/computer/file_read_utils.py
@@ -1,7 +1,10 @@
from __future__ import annotations
import base64
+import hashlib
+import io
import json
+import zipfile
from asyncio import to_thread
from dataclasses import dataclass
from pathlib import Path
@@ -45,11 +48,13 @@ _UTF_BOMS = (
b"\xff\xfe\x00\x00",
b"\x00\x00\xfe\xff",
)
-_BINARY_MAGIC_PREFIXES = (
- b"%PDF-",
+_ZIP_MAGIC_PREFIXES = (
b"PK\x03\x04",
b"PK\x05\x06",
b"PK\x07\x08",
+)
+_BINARY_MAGIC_PREFIXES = (
+ b"%PDF-",
b"\x1f\x8b",
b"7z\xbc\xaf\x27\x1c",
b"Rar!\x1a\x07",
@@ -66,6 +71,13 @@ class FileProbe:
size_bytes: int
+@dataclass(frozen=True)
+class ParsedDocument:
+ kind: Literal["docx", "pdf"]
+ file_bytes: bytes
+ text: str
+
+
def _build_probe_script(path: str) -> str:
return f"""
import base64
@@ -269,6 +281,10 @@ async def _read_local_image_base64(path: str) -> dict[str, str | int]:
return await to_thread(_run)
+async def _read_local_file_bytes(path: str) -> bytes:
+ return await to_thread(Path(path).read_bytes)
+
+
async def _compress_image_bytes_to_base64(data: bytes) -> dict[str, str | int]:
def _run() -> dict[str, str | int]:
temp_dir = Path(get_astrbot_temp_path())
@@ -320,6 +336,63 @@ def _looks_like_known_binary(sample: bytes) -> bool:
return any(sample.startswith(prefix) for prefix in _BINARY_MAGIC_PREFIXES)
+def _looks_like_pdf(path: str, sample: bytes) -> bool:
+ return Path(path).suffix.lower() == ".pdf" or sample.startswith(b"%PDF-")
+
+
+def _looks_like_zip_container(sample: bytes) -> bool:
+ return any(sample.startswith(prefix) for prefix in _ZIP_MAGIC_PREFIXES)
+
+
+def _is_docx_bytes(file_bytes: bytes) -> bool:
+ try:
+ with zipfile.ZipFile(io.BytesIO(file_bytes)) as archive:
+ names = set(archive.namelist())
+ except (OSError, zipfile.BadZipFile):
+ return False
+
+ if "[Content_Types].xml" not in names:
+ return False
+
+ return any(name.startswith("word/") for name in names)
+
+
+async def _parse_local_docx_text(file_bytes: bytes, file_name: str) -> str:
+ from astrbot.core.knowledge_base.parsers.markitdown_parser import (
+ MarkitdownParser,
+ )
+
+ result = await MarkitdownParser().parse(file_bytes, file_name)
+ return result.text
+
+
+async def _parse_local_pdf_text(file_bytes: bytes, file_name: str) -> str:
+ from astrbot.core.knowledge_base.parsers.pdf_parser import PDFParser
+
+ result = await PDFParser().parse(file_bytes, file_name)
+ return result.text
+
+
+async def _parse_local_supported_document(
+ path: str,
+ sample: bytes,
+) -> ParsedDocument | None:
+ file_name = Path(path).name
+ if _looks_like_pdf(path, sample):
+ file_bytes = await _read_local_file_bytes(path)
+ text = await _parse_local_pdf_text(file_bytes, file_name)
+ return ParsedDocument(kind="pdf", file_bytes=file_bytes, text=text)
+
+ if Path(path).suffix.lower() == ".docx" or _looks_like_zip_container(sample):
+ file_bytes = await _read_local_file_bytes(path)
+ if not _is_docx_bytes(file_bytes):
+ return None
+ text = await _parse_local_docx_text(file_bytes, file_name)
+ return ParsedDocument(kind="docx", file_bytes=file_bytes, text=text)
+
+ return None
+
+
def _probe_file(sample: bytes, *, size_bytes: int) -> FileProbe:
if image_mime := _detect_image_mime(sample):
return FileProbe(
@@ -375,6 +448,10 @@ def _validate_text_output(content: str) -> str | None:
return None
+def _text_exceeds_read_thresholds(content: str) -> bool:
+ return _validate_text_output(content) is not None
+
+
def _validate_full_text_read_request(probe: FileProbe) -> str | None:
if probe.size_bytes > _MAX_TEXT_FILE_FULL_READ_BYTES:
return (
@@ -385,6 +462,135 @@ def _validate_full_text_read_request(probe: FileProbe) -> str | None:
return None
+def _slice_text_by_lines(
+ content: str,
+ *,
+ offset: int | None,
+ limit: int | None,
+) -> str:
+ if offset is None and limit is None:
+ return content
+
+ lines = content.splitlines(keepends=True)
+ start = 0 if offset is None else offset
+ end = None if limit is None else start + limit
+ return "".join(lines[start:end])
+
+
+async def _store_converted_text_for_workspace(
+ *,
+ workspace_dir: str,
+ original_path: str,
+ original_bytes: bytes,
+ content: str,
+) -> str:
+ def _run() -> str:
+ original_name = Path(original_path).name
+ digest_suffix = hashlib.md5(original_bytes).hexdigest()[-6:]
+ target_dir = (
+ Path(workspace_dir) / "converted_files" / f"{original_name}_{digest_suffix}"
+ )
+ target_dir.mkdir(parents=True, exist_ok=True)
+ target_path = target_dir / "text.txt"
+ target_path.write_text(content, encoding="utf-8")
+ return str(target_path)
+
+ return await to_thread(_run)
+
+
+def _build_converted_text_notice(
+ converted_text_path: str,
+ *,
+ selection_returned: bool,
+ selection_too_large: bool = False,
+) -> str:
+ if selection_too_large:
+ return (
+ "Converted text was saved to "
+ f"`{converted_text_path}`. The requested output is still too large to "
+ "return directly. Read that text file with a narrower `offset`/`limit` window."
+ )
+
+ if selection_returned:
+ return (
+ "Full converted text is also available at "
+ f"`{converted_text_path}`. Use that file with a narrow `offset`/`limit` "
+ "window for additional reads."
+ )
+
+ return (
+ "Converted text was saved to "
+ f"`{converted_text_path}` because the parsed document is too large to "
+ "return directly. Read that text file with a narrow `offset`/`limit` window."
+ )
+
+
+async def _read_local_supported_document_result(
+ *,
+ path: str,
+ parsed_document: ParsedDocument,
+ workspace_dir: str | None,
+ offset: int | None,
+ limit: int | None,
+) -> ToolExecResult:
+ content = parsed_document.text
+ if not content:
+ return "No content found at the requested line offset."
+
+ if not _text_exceeds_read_thresholds(content):
+ selected_content = _slice_text_by_lines(content, offset=offset, limit=limit)
+ if not selected_content:
+ return "No content found at the requested line offset."
+ if validation_error := _validate_text_output(selected_content):
+ return validation_error
+ return selected_content
+
+ if not workspace_dir:
+ return (
+ "Error reading file: parsed document exceeds the read output limit and "
+ "no workspace is available for storing converted text."
+ )
+
+ converted_text_path = await _store_converted_text_for_workspace(
+ workspace_dir=workspace_dir,
+ original_path=path,
+ original_bytes=parsed_document.file_bytes,
+ content=content,
+ )
+
+ if offset is None and limit is None:
+ return _build_converted_text_notice(
+ converted_text_path,
+ selection_returned=False,
+ )
+
+ selected_content = _slice_text_by_lines(content, offset=offset, limit=limit)
+ if not selected_content:
+ return (
+ "No content found at the requested line offset. "
+ + _build_converted_text_notice(
+ converted_text_path,
+ selection_returned=False,
+ )
+ )
+
+ notice = _build_converted_text_notice(
+ converted_text_path,
+ selection_returned=True,
+ )
+ combined_output = f"{selected_content}\n\n[{notice}]"
+ if _validate_text_output(combined_output):
+ if _validate_text_output(selected_content):
+ return _build_converted_text_notice(
+ converted_text_path,
+ selection_returned=False,
+ selection_too_large=True,
+ )
+ return selected_content
+
+ return combined_output
+
+
async def read_file_tool_result(
booter: ComputerBooter,
*,
@@ -392,6 +598,7 @@ async def read_file_tool_result(
path: str,
offset: int | None,
limit: int | None,
+ workspace_dir: str | None = None,
) -> ToolExecResult:
if local_mode:
probe_payload = await _probe_local_file(path)
@@ -406,6 +613,21 @@ async def read_file_tool_result(
size_bytes = int(probe_payload.get("size_bytes", 0) or 0)
probe = _probe_file(sample, size_bytes=size_bytes)
+ if local_mode:
+ try:
+ parsed_document = await _parse_local_supported_document(path, sample)
+ except Exception as exc:
+ return f"Error reading file: failed to parse document: {exc}"
+
+ if parsed_document is not None:
+ return await _read_local_supported_document_result(
+ path=path,
+ parsed_document=parsed_document,
+ workspace_dir=workspace_dir,
+ offset=offset,
+ limit=limit,
+ )
+
if probe.kind == "binary":
return "Error reading file: binary files are not supported by this tool."
diff --git a/astrbot/core/tools/computer_tools/fs.py b/astrbot/core/tools/computer_tools/fs.py
index 94f987fcd..d18701f7b 100644
--- a/astrbot/core/tools/computer_tools/fs.py
+++ b/astrbot/core/tools/computer_tools/fs.py
@@ -50,11 +50,11 @@ from astrbot.core.utils.astrbot_path import (
)
from ..registry import builtin_tool
+from . import util as computer_util
from .util import (
check_admin_permission,
is_local_runtime,
normalize_umo_for_workspace,
- workspace_root,
)
_COMPUTER_RUNTIME_TOOL_CONFIG = {
@@ -75,11 +75,22 @@ def _restricted_env_path_labels(umo: str) -> list[str]:
]
+def get_astrbot_workspaces_path() -> str:
+ """Compatibility wrapper for tests and older module-level monkeypatches."""
+ return computer_util.get_astrbot_workspaces_path()
+
+
+def _workspace_root(umo: str) -> Path:
+ """Workspace root that follows both util-level and fs-level getter monkeypatches."""
+ normalized_umo = normalize_umo_for_workspace(umo)
+ return (Path(get_astrbot_workspaces_path()) / normalized_umo).resolve(strict=False)
+
+
def _read_allowed_roots(umo: str) -> tuple[Path, ...]:
"""Non-admin users can only read files within these directories (and their subdirectories)"""
return (
Path(get_astrbot_skills_path()).resolve(strict=False),
- workspace_root(umo),
+ _workspace_root(umo),
Path("/tmp/.astrbot").resolve(strict=False),
)
@@ -103,7 +114,7 @@ def _resolve_tool_path(path: str, *, local_env: bool, umo: str) -> str:
if candidate.is_absolute():
return str(candidate.resolve(strict=False))
if local_env:
- return str((workspace_root(umo) / candidate).resolve(strict=False))
+ return str((_workspace_root(umo) / candidate).resolve(strict=False))
return normalized_path
@@ -112,7 +123,7 @@ def _resolve_user_path(path: str, *, local_env: bool, umo: str) -> Path:
if candidate.is_absolute():
return candidate.resolve(strict=False)
if local_env:
- return (workspace_root(umo) / candidate).resolve(strict=False)
+ return (_workspace_root(umo) / candidate).resolve(strict=False)
return (Path.cwd() / candidate).resolve(strict=False)
@@ -225,6 +236,11 @@ class FileReadTool(FunctionTool):
path=normalized_path,
offset=offset,
limit=limit,
+ workspace_dir=(
+ str(_workspace_root(context.context.event.unified_msg_origin))
+ if local_env
+ else None
+ ),
)
except PermissionError as exc:
return f"Error: {exc}"
@@ -505,7 +521,7 @@ class GrepTool(FunctionTool):
if restricted:
return [str(root) for root in _read_allowed_roots(umo)]
if local_env:
- return [str(workspace_root(umo))]
+ return [str(_workspace_root(umo))]
return ["."]
if restricted:
diff --git a/tests/test_computer_fs_tools.py b/tests/test_computer_fs_tools.py
index 6e718dfbd..aaaed11a9 100644
--- a/tests/test_computer_fs_tools.py
+++ b/tests/test_computer_fs_tools.py
@@ -1,6 +1,8 @@
from __future__ import annotations
import base64
+import io
+import zipfile
from types import SimpleNamespace
from typing import Any
@@ -12,6 +14,7 @@ from astrbot.core.agent.run_context import ContextWrapper
from astrbot.core.computer import file_read_utils
from astrbot.core.computer.booters.local import LocalBooter
from astrbot.core.tools.computer_tools import fs as fs_tools
+from astrbot.core.tools.computer_tools import util as computer_util
def _make_context(
@@ -52,7 +55,7 @@ def _setup_local_fs_tools(
temp_root.mkdir()
monkeypatch.setattr(
- fs_tools,
+ computer_util,
"get_astrbot_workspaces_path",
lambda: str(workspaces_root),
)
@@ -79,7 +82,7 @@ def _setup_local_fs_tools(
monkeypatch.setattr(fs_tools, "get_booter", _fake_get_booter)
- normalized_umo = fs_tools._normalize_umo_for_workspace(umo)
+ normalized_umo = computer_util.normalize_umo_for_workspace(umo)
workspace = workspaces_root / normalized_umo
workspace.mkdir(parents=True, exist_ok=True)
return workspace
@@ -174,7 +177,7 @@ async def test_file_read_tool_treats_svg_as_text(
@pytest.mark.asyncio
-async def test_file_read_tool_rejects_pdf_as_binary(
+async def test_file_read_tool_reads_pdf_via_parser(
monkeypatch: pytest.MonkeyPatch,
tmp_path,
):
@@ -182,12 +185,71 @@ async def test_file_read_tool_rejects_pdf_as_binary(
pdf_path = workspace / "doc.pdf"
pdf_path.write_bytes(b"%PDF-1.7\n%\xe2\xe3\xcf\xd3\n1 0 obj\n<<>>\nendobj\n")
+ async def _fake_parse_pdf(_file_bytes: bytes, _file_name: str) -> str:
+ return "page-1\npage-2\n"
+
+ monkeypatch.setattr(file_read_utils, "_parse_local_pdf_text", _fake_parse_pdf)
+
result = await fs_tools.FileReadTool().call(
_make_context(),
path="doc.pdf",
)
- assert result == "Error reading file: binary files are not supported by this tool."
+ assert result == "page-1\npage-2\n"
+
+
+@pytest.mark.asyncio
+async def test_file_read_tool_reads_docx_via_parser_and_magic(
+ monkeypatch: pytest.MonkeyPatch,
+ tmp_path,
+):
+ workspace = _setup_local_fs_tools(monkeypatch, tmp_path)
+ docx_path = workspace / "report.bin"
+ buffer = io.BytesIO()
+ with zipfile.ZipFile(buffer, mode="w") as archive:
+ archive.writestr("[Content_Types].xml", "")
+ archive.writestr("word/document.xml", "")
+ docx_path.write_bytes(buffer.getvalue())
+
+ async def _fake_parse_docx(_file_bytes: bytes, _file_name: str) -> str:
+ return "doc-line-1\ndoc-line-2\n"
+
+ monkeypatch.setattr(file_read_utils, "_parse_local_docx_text", _fake_parse_docx)
+
+ result = await fs_tools.FileReadTool().call(
+ _make_context(),
+ path="report.bin",
+ )
+
+ assert result == "doc-line-1\ndoc-line-2\n"
+
+
+@pytest.mark.asyncio
+async def test_file_read_tool_stores_long_converted_document_in_workspace(
+ monkeypatch: pytest.MonkeyPatch,
+ tmp_path,
+):
+ workspace = _setup_local_fs_tools(monkeypatch, tmp_path)
+ pdf_path = workspace / "manual.pdf"
+ pdf_path.write_bytes(b"%PDF-1.7\nfake\n")
+ long_text = _make_large_text()
+
+ async def _fake_parse_pdf(_file_bytes: bytes, _file_name: str) -> str:
+ return long_text
+
+ monkeypatch.setattr(file_read_utils, "_parse_local_pdf_text", _fake_parse_pdf)
+
+ result = await fs_tools.FileReadTool().call(
+ _make_context(),
+ path="manual.pdf",
+ )
+
+ converted_root = workspace / "converted_files"
+ converted_files = list(converted_root.glob("manual.pdf_*/text.txt"))
+ assert len(converted_files) == 1
+ assert converted_files[0].read_text(encoding="utf-8") == long_text
+ assert str(converted_files[0]) in result
+ assert "narrow `offset`/`limit` window" in result
@pytest.mark.asyncio