mirror of
https://github.com/countbot-ai/CountBot.git
synced 2026-09-14 20:46:47 +08:00
573 lines
22 KiB
Python
573 lines
22 KiB
Python
"""文件系统工具 - 支持行号显示、行范围读取、追加写入、按行编辑"""
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from loguru import logger
|
|
|
|
from backend.modules.tools.base import Tool
|
|
|
|
|
|
class WorkspaceValidator:
|
|
"""工作空间验证器 - 确保文件操作在工作空间内"""
|
|
|
|
def __init__(self, workspace: Path, restrict_to_workspace: bool = True):
|
|
self._workspace = workspace.resolve()
|
|
self._restrict_to_workspace = restrict_to_workspace
|
|
logger.debug(f"WorkspaceValidator initialized with workspace: {self._workspace}")
|
|
|
|
@property
|
|
def workspace(self) -> Path:
|
|
try:
|
|
from backend.modules.config.loader import config_loader
|
|
workspace_path = config_loader.config.workspace.path
|
|
if workspace_path:
|
|
return Path(workspace_path).resolve()
|
|
except Exception:
|
|
pass
|
|
return self._workspace
|
|
|
|
@property
|
|
def restrict_to_workspace(self) -> bool:
|
|
try:
|
|
from backend.modules.config.loader import config_loader
|
|
return config_loader.config.security.restrict_to_workspace
|
|
except Exception:
|
|
return self._restrict_to_workspace
|
|
|
|
def validate_path(self, path: str) -> Path:
|
|
current_workspace = self.workspace
|
|
if Path(path).is_absolute():
|
|
resolved = Path(path).resolve()
|
|
else:
|
|
resolved = (current_workspace / path).resolve()
|
|
if not self.restrict_to_workspace:
|
|
return resolved
|
|
try:
|
|
resolved.relative_to(current_workspace)
|
|
except ValueError:
|
|
error_msg = f"Path '{path}' is outside workspace '{current_workspace}'"
|
|
logger.error(error_msg)
|
|
raise ValueError(error_msg)
|
|
return resolved
|
|
|
|
|
|
class ReadFileTool(Tool):
|
|
"""读取文件工具 - 支持行号显示和按行范围读取"""
|
|
|
|
def __init__(self, workspace: Path, skills_loader=None, restrict_to_workspace: bool = True):
|
|
self.validator = WorkspaceValidator(workspace, restrict_to_workspace)
|
|
self.skills_loader = skills_loader
|
|
logger.debug("ReadFileTool initialized")
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "read_file"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Read one file or many files. Single-file mode supports optional 1-based line ranges."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"description": "Single file path.",
|
|
},
|
|
"paths": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": "Multiple file paths.",
|
|
},
|
|
"start_line": {
|
|
"type": "integer",
|
|
"description": "1-based start line.",
|
|
},
|
|
"end_line": {
|
|
"type": "integer",
|
|
"description": "1-based end line.",
|
|
},
|
|
"show_line_numbers": {
|
|
"type": "boolean",
|
|
"description": "Show line numbers.",
|
|
},
|
|
},
|
|
"oneOf": [
|
|
{"required": ["path"]},
|
|
{"required": ["paths"]}
|
|
],
|
|
}
|
|
|
|
async def execute(self, **kwargs: Any) -> str:
|
|
path_str = kwargs.get("path")
|
|
paths_list = kwargs.get("paths")
|
|
start_line = kwargs.get("start_line")
|
|
end_line = kwargs.get("end_line")
|
|
show_line_numbers = kwargs.get("show_line_numbers", True)
|
|
|
|
# 参数验证:必须提供 path 或 paths 之一
|
|
if not path_str and not paths_list:
|
|
return "Error: Either 'path' or 'paths' parameter is required"
|
|
|
|
if path_str and paths_list:
|
|
return "Error: Provide either 'path' or 'paths', not both"
|
|
|
|
# 批量模式:读取多个文件
|
|
if paths_list:
|
|
if not isinstance(paths_list, list):
|
|
return "Error: 'paths' must be an array of strings"
|
|
|
|
if not paths_list:
|
|
return "Error: 'paths' array is empty"
|
|
|
|
# 批量模式不支持行范围参数
|
|
if start_line is not None or end_line is not None:
|
|
return "Error: Line range parameters (start_line/end_line) are not supported in batch mode"
|
|
|
|
return await self._read_multiple_files(paths_list, show_line_numbers)
|
|
|
|
# 单文件模式:保持原有逻辑
|
|
return await self._read_single_file(path_str, start_line, end_line, show_line_numbers)
|
|
|
|
async def _read_single_file(
|
|
self,
|
|
path_str: str,
|
|
start_line: Optional[int],
|
|
end_line: Optional[int],
|
|
show_line_numbers: bool
|
|
) -> str:
|
|
"""读取单个文件(原有逻辑)"""
|
|
if not path_str:
|
|
return "Error: Path parameter is required"
|
|
|
|
try:
|
|
# 禁用技能检查
|
|
if self.skills_loader:
|
|
file_path_check = Path(path_str)
|
|
if not file_path_check.is_absolute():
|
|
file_path_check = (self.validator.workspace / path_str).resolve()
|
|
else:
|
|
file_path_check = file_path_check.resolve()
|
|
if file_path_check.name == "SKILL.md":
|
|
skill_name = file_path_check.parent.name
|
|
skill = self.skills_loader.get_skill(skill_name)
|
|
if skill and not skill.enabled:
|
|
logger.warning(f"Blocked read of disabled skill: {skill_name}")
|
|
return f"Error: Skill '{skill_name}' is disabled. Enable it first."
|
|
|
|
file_path = self.validator.validate_path(path_str)
|
|
if not file_path.exists():
|
|
return f"Error: File not found: {path_str}"
|
|
if not file_path.is_file():
|
|
return f"Error: Not a file: {path_str}"
|
|
|
|
content = self._read_text_content(file_path, path_str)
|
|
lines = content.splitlines()
|
|
total = len(lines)
|
|
|
|
# 解析行范围
|
|
s = max(1, int(start_line)) if start_line is not None else 1
|
|
e = min(total, int(end_line)) if end_line is not None else total
|
|
|
|
if s > total:
|
|
return f"Error: start_line ({s}) exceeds total lines ({total})"
|
|
if s > e:
|
|
return f"Error: start_line ({s}) > end_line ({e})"
|
|
|
|
selected = lines[s - 1:e]
|
|
|
|
# 格式化
|
|
if show_line_numbers:
|
|
w = len(str(e))
|
|
output_lines = [f"{s + i:>{w}}| {line}" for i, line in enumerate(selected)]
|
|
else:
|
|
output_lines = selected
|
|
|
|
header = f"[File: {path_str} | Lines: {total}"
|
|
if s != 1 or e != total:
|
|
header += f" | Showing: {s}-{e}"
|
|
header += "]"
|
|
|
|
logger.info(f"Read file: {path_str} (lines {s}-{e} of {total})")
|
|
return header + "\n" + "\n".join(output_lines)
|
|
|
|
except ValueError as ve:
|
|
logger.error(f"Failed to read file '{path_str}': {ve}")
|
|
return f"Error: {ve}"
|
|
except Exception as ex:
|
|
logger.error(f"Unexpected error reading file '{path_str}': {ex}")
|
|
return f"Error reading file: {str(ex)}"
|
|
|
|
async def _read_multiple_files(self, paths_list: List[str], show_line_numbers: bool) -> str:
|
|
"""批量读取多个文件"""
|
|
results = []
|
|
success_count = 0
|
|
error_count = 0
|
|
|
|
for path_str in paths_list:
|
|
try:
|
|
# 禁用技能检查
|
|
if self.skills_loader:
|
|
file_path_check = Path(path_str)
|
|
if not file_path_check.is_absolute():
|
|
file_path_check = (self.validator.workspace / path_str).resolve()
|
|
else:
|
|
file_path_check = file_path_check.resolve()
|
|
if file_path_check.name == "SKILL.md":
|
|
skill_name = file_path_check.parent.name
|
|
skill = self.skills_loader.get_skill(skill_name)
|
|
if skill and not skill.enabled:
|
|
logger.warning(f"Blocked read of disabled skill: {skill_name}")
|
|
results.append(f"[File: {path_str}]\nError: Skill '{skill_name}' is disabled. Enable it first.")
|
|
error_count += 1
|
|
continue
|
|
|
|
file_path = self.validator.validate_path(path_str)
|
|
|
|
if not file_path.exists():
|
|
results.append(f"[File: {path_str}]\nError: File not found")
|
|
error_count += 1
|
|
continue
|
|
|
|
if not file_path.is_file():
|
|
results.append(f"[File: {path_str}]\nError: Not a file")
|
|
error_count += 1
|
|
continue
|
|
|
|
content = self._read_text_content(file_path, path_str)
|
|
lines = content.splitlines()
|
|
total = len(lines)
|
|
|
|
# 格式化输出
|
|
if show_line_numbers:
|
|
w = len(str(total))
|
|
output_lines = [f"{i + 1:>{w}}| {line}" for i, line in enumerate(lines)]
|
|
else:
|
|
output_lines = lines
|
|
|
|
header = f"[File: {path_str} | Lines: {total}]"
|
|
results.append(header + "\n" + "\n".join(output_lines))
|
|
success_count += 1
|
|
logger.info(f"Read file (batch): {path_str} ({total} lines)")
|
|
|
|
except ValueError as ve:
|
|
logger.error(f"Failed to read file '{path_str}' in batch: {ve}")
|
|
results.append(f"[File: {path_str}]\nError: {ve}")
|
|
error_count += 1
|
|
except Exception as ex:
|
|
logger.error(f"Unexpected error reading file '{path_str}' in batch: {ex}")
|
|
results.append(f"[File: {path_str}]\nError: {str(ex)}")
|
|
error_count += 1
|
|
|
|
# 添加批量读取摘要
|
|
summary = f"\n{'='*60}\n[Batch Read Summary: {success_count} succeeded, {error_count} failed]\n{'='*60}"
|
|
|
|
logger.info(f"Batch read completed: {success_count} succeeded, {error_count} failed out of {len(paths_list)} files")
|
|
|
|
return "\n\n".join(results) + summary
|
|
|
|
@staticmethod
|
|
def _read_text_content(file_path: Path, path_str: str) -> str:
|
|
"""读取 UTF-8 文本文件,遇到二进制内容时返回更明确的错误。"""
|
|
data = file_path.read_bytes()
|
|
if b"\x00" in data:
|
|
raise ValueError(
|
|
f"File appears to be binary and cannot be read as UTF-8 text: {path_str}"
|
|
)
|
|
|
|
try:
|
|
return data.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
raise ValueError(
|
|
f"File appears to be binary or uses a non-UTF-8 encoding and cannot be read as text: {path_str}"
|
|
) from None
|
|
|
|
|
|
class WriteFileTool(Tool):
|
|
"""写入文件工具 - 支持覆盖和追加模式,大文件请分段追加写入"""
|
|
|
|
def __init__(self, workspace: Path, restrict_to_workspace: bool = True):
|
|
self.validator = WorkspaceValidator(workspace, restrict_to_workspace)
|
|
logger.debug("WriteFileTool initialized")
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "write_file"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return (
|
|
"Write text to a file. Modes: `overwrite` or `append`. "
|
|
"For large HTML/code/text, write in small chunks and use `append` "
|
|
"after the first chunk."
|
|
)
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"description": "File path.",
|
|
},
|
|
"content": {
|
|
"type": "string",
|
|
"description": (
|
|
"Text to write. For large HTML/code/text, keep each chunk small "
|
|
"(recommended <= 800 chars) and append subsequent chunks."
|
|
),
|
|
},
|
|
"mode": {
|
|
"type": "string",
|
|
"enum": ["overwrite", "append"],
|
|
"description": (
|
|
"Write mode. Use `overwrite` for the first chunk and `append` "
|
|
"for later chunks."
|
|
),
|
|
},
|
|
},
|
|
"required": ["path", "content"],
|
|
}
|
|
|
|
async def execute(self, **kwargs: Any) -> str:
|
|
path_str = kwargs.get("path", "")
|
|
content = kwargs.get("content", "")
|
|
mode = kwargs.get("mode", "overwrite")
|
|
|
|
if not path_str:
|
|
return "Error: Path parameter is required"
|
|
|
|
try:
|
|
file_path = self.validator.validate_path(path_str)
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
if mode == "append":
|
|
if file_path.exists():
|
|
existing = file_path.read_text(encoding="utf-8")
|
|
file_path.write_text(existing + content, encoding="utf-8")
|
|
total = len(existing) + len(content)
|
|
logger.info(f"Appended to file: {path_str} (+{len(content)} chars, total: {total})")
|
|
return f"Appended {len(content)} chars to {path_str} (total: {total})"
|
|
else:
|
|
file_path.write_text(content, encoding="utf-8")
|
|
logger.info(f"Created file (append, new): {path_str} ({len(content)} chars)")
|
|
return f"Created {path_str} with {len(content)} chars (file was new)"
|
|
else:
|
|
file_path.write_text(content, encoding="utf-8")
|
|
logger.info(f"Wrote file: {path_str} ({len(content)} chars)")
|
|
return f"Wrote {len(content)} chars to {path_str}"
|
|
|
|
except ValueError as ve:
|
|
logger.error(f"Failed to write file '{path_str}': {ve}")
|
|
return f"Error: {ve}"
|
|
except Exception as ex:
|
|
logger.error(f"Unexpected error writing file '{path_str}': {ex}")
|
|
return f"Error writing file: {str(ex)}"
|
|
|
|
|
|
class EditFileTool(Tool):
|
|
"""编辑文件工具 - 支持文本替换和按行号编辑"""
|
|
|
|
def __init__(self, workspace: Path, restrict_to_workspace: bool = True):
|
|
self.validator = WorkspaceValidator(workspace, restrict_to_workspace)
|
|
logger.debug("EditFileTool initialized")
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "edit_file"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Edit a file by text replace or 1-based line range."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"description": "File path.",
|
|
},
|
|
"old_text": {
|
|
"type": "string",
|
|
"description": "Text to replace.",
|
|
},
|
|
"new_text": {
|
|
"type": "string",
|
|
"description": "New text.",
|
|
},
|
|
"start_line": {
|
|
"type": "integer",
|
|
"description": "1-based start line.",
|
|
},
|
|
"end_line": {
|
|
"type": "integer",
|
|
"description": "1-based end line.",
|
|
},
|
|
"insert": {
|
|
"type": "boolean",
|
|
"description": "Insert before start_line.",
|
|
},
|
|
},
|
|
"required": ["path"],
|
|
}
|
|
|
|
async def execute(self, **kwargs: Any) -> str:
|
|
path_str = kwargs.get("path", "")
|
|
old_text = kwargs.get("old_text")
|
|
new_text = kwargs.get("new_text", "")
|
|
start_line = kwargs.get("start_line")
|
|
end_line = kwargs.get("end_line")
|
|
insert_mode = kwargs.get("insert", False)
|
|
|
|
if not path_str:
|
|
return "Error: Path parameter is required"
|
|
|
|
try:
|
|
file_path = self.validator.validate_path(path_str)
|
|
if not file_path.exists():
|
|
return f"Error: File not found: {path_str}"
|
|
|
|
content = file_path.read_text(encoding="utf-8")
|
|
|
|
if start_line is not None:
|
|
return self._edit_by_lines(
|
|
file_path, path_str, content,
|
|
int(start_line),
|
|
int(end_line) if end_line is not None else None,
|
|
new_text, insert_mode
|
|
)
|
|
elif old_text is not None:
|
|
return self._edit_by_text(file_path, path_str, content, old_text, new_text)
|
|
else:
|
|
return "Error: Provide 'old_text' (text mode) or 'start_line' (line mode)"
|
|
|
|
except ValueError as ve:
|
|
logger.error(f"Failed to edit file '{path_str}': {ve}")
|
|
return f"Error: {ve}"
|
|
except Exception as ex:
|
|
logger.error(f"Unexpected error editing file '{path_str}': {ex}")
|
|
return f"Error editing file: {str(ex)}"
|
|
|
|
def _edit_by_text(self, file_path: Path, path_str: str, content: str,
|
|
old_text: str, new_text: str) -> str:
|
|
if not old_text:
|
|
return "Error: old_text is required for text replace mode"
|
|
if old_text not in content:
|
|
total = len(content.splitlines())
|
|
return (
|
|
f"Error: old_text not found in file ({total} lines, {len(content)} chars). "
|
|
f"Use read_file to check exact content."
|
|
)
|
|
count = content.count(old_text)
|
|
if count > 1:
|
|
return f"Warning: old_text found {count} times. Add more context to make it unique."
|
|
|
|
new_content = content.replace(old_text, new_text, 1)
|
|
file_path.write_text(new_content, encoding="utf-8")
|
|
logger.info(f"Edited file (text replace): {path_str}")
|
|
return f"Edited {path_str} (replaced 1 occurrence)"
|
|
|
|
def _edit_by_lines(self, file_path: Path, path_str: str, content: str,
|
|
start_line: int, end_line: Optional[int],
|
|
new_text: str, insert_mode: bool) -> str:
|
|
lines = content.splitlines(keepends=True)
|
|
total = len(lines)
|
|
|
|
if start_line < 1 or start_line > total + 1:
|
|
return f"Error: start_line ({start_line}) out of range (1-{total})"
|
|
|
|
if insert_mode:
|
|
new_lines = new_text.splitlines(keepends=True)
|
|
if new_lines and not new_lines[-1].endswith("\n"):
|
|
new_lines[-1] += "\n"
|
|
idx = start_line - 1
|
|
result = lines[:idx] + new_lines + lines[idx:]
|
|
file_path.write_text("".join(result), encoding="utf-8")
|
|
logger.info(f"Inserted {len(new_lines)} lines before line {start_line}: {path_str}")
|
|
return f"Inserted {len(new_lines)} lines before line {start_line} in {path_str}"
|
|
|
|
if end_line is None:
|
|
end_line = start_line
|
|
if end_line < start_line:
|
|
return f"Error: end_line ({end_line}) < start_line ({start_line})"
|
|
if end_line > total:
|
|
return f"Error: end_line ({end_line}) exceeds total lines ({total})"
|
|
|
|
if new_text:
|
|
new_lines = new_text.splitlines(keepends=True)
|
|
if new_lines and not new_lines[-1].endswith("\n"):
|
|
new_lines[-1] += "\n"
|
|
else:
|
|
new_lines = []
|
|
|
|
result = lines[:start_line - 1] + new_lines + lines[end_line:]
|
|
file_path.write_text("".join(result), encoding="utf-8")
|
|
|
|
if not new_text:
|
|
action = f"Deleted lines {start_line}-{end_line}"
|
|
else:
|
|
action = f"Replaced lines {start_line}-{end_line} with {len(new_lines)} lines"
|
|
logger.info(f"Line edit: {path_str} - {action}")
|
|
return f"Edited {path_str}: {action}"
|
|
|
|
|
|
class ListDirTool(Tool):
|
|
"""列出目录工具"""
|
|
|
|
def __init__(self, workspace: Path, restrict_to_workspace: bool = True):
|
|
self.validator = WorkspaceValidator(workspace, restrict_to_workspace)
|
|
logger.debug("ListDirTool initialized")
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "list_dir"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "List files and subdirectories in a directory."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {
|
|
"type": "string",
|
|
"description": "Directory path.",
|
|
"default": ".",
|
|
},
|
|
},
|
|
}
|
|
|
|
async def execute(self, **kwargs: Any) -> str:
|
|
path_str = kwargs.get("path", ".")
|
|
try:
|
|
dir_path = self.validator.validate_path(path_str)
|
|
if not dir_path.exists():
|
|
raise FileNotFoundError(f"Directory not found: {path_str}")
|
|
if not dir_path.is_dir():
|
|
raise ValueError(f"Not a directory: {path_str}")
|
|
|
|
items = []
|
|
for item in sorted(dir_path.iterdir()):
|
|
item_type = "dir" if item.is_dir() else "file"
|
|
size = item.stat().st_size if item.is_file() else 0
|
|
items.append(f"{item_type:4} {item.name:40} {size:>10} bytes")
|
|
|
|
result = f"Contents of {path_str}:\n" + "\n".join(items)
|
|
logger.info(f"Listed directory: {path_str} ({len(items)} items)")
|
|
return result
|
|
|
|
except (ValueError, FileNotFoundError) as e:
|
|
logger.error(f"Failed to list directory '{path_str}': {e}")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error listing directory '{path_str}': {e}")
|
|
raise
|