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