From 28e8e5000540aaa9321f26ad71cb9d737ae9f9bc Mon Sep 17 00:00:00 2001 From: LIghtJUNction Date: Fri, 20 Mar 2026 21:40:21 +0800 Subject: [PATCH] fix: improve dashboard and knowledge base file handling --- astrbot/cli/utils/dashboard.py | 7 ++++-- astrbot/core/knowledge_base/kb_helper.py | 31 +++++++++++++----------- astrbot/core/utils/io.py | 8 ------ 3 files changed, 22 insertions(+), 24 deletions(-) diff --git a/astrbot/cli/utils/dashboard.py b/astrbot/cli/utils/dashboard.py index 83dbc0178..d0cced998 100644 --- a/astrbot/cli/utils/dashboard.py +++ b/astrbot/cli/utils/dashboard.py @@ -15,7 +15,7 @@ class DashboardManager: from astrbot.core.utils.io import download_dashboard, get_dashboard_version if self._bundled_dist.is_dir(): - click.echo("Dashboard is bundled with the package – skipping download.") + click.echo("Dashboard is bundled with the package - skipping download.") return try: @@ -41,7 +41,10 @@ class DashboardManager: click.echo(f"Failed to install dashboard: {e}") case str(): - if VersionComparator.compare_version(VERSION, dashboard_version) <= 0: + if ( + VersionComparator.compare_version(VERSION, dashboard_version) + <= 0 + ): click.echo("Dashboard is already up to date") return try: diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index d304260af..4014b77cd 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -233,15 +233,12 @@ class KBHelper: await self._ensure_vec_db() doc_id = str(uuid.uuid4()) media_paths: list[Path] = [] + saved_file_path: Path | None = None file_size = 0 - # file_path = self.kb_files_dir / f"{doc_id}.{file_type}" - # async with aiofiles.open(file_path, "wb") as f: - # await f.write(file_content) - try: - chunks_text = [] - saved_media = [] + chunks_text: list[str] = [] + saved_media: list[KBMedia] = [] if pre_chunked_text is not None: # 如果提供了预分块文本,直接使用 @@ -256,6 +253,9 @@ class KBHelper: ) file_size = len(file_content) + saved_file_path = self.kb_files_dir / f"{doc_id}.{file_type}" + async with aiofiles.open(saved_file_path, "wb") as f: + await f.write(file_content) # 阶段1: 解析文档 if progress_callback: @@ -326,17 +326,15 @@ class KBHelper: doc_name=file_name, file_type=file_type, file_size=file_size, - # file_path=str(file_path), - file_path="", + file_path=str(saved_file_path) if saved_file_path else "", chunk_count=len(chunks_text), - media_count=0, + media_count=len(saved_media), ) async with self.kb_db.get_db() as session: async with session.begin(): session.add(doc) for media in saved_media: session.add(media) - await session.commit() await session.refresh(doc) @@ -347,8 +345,14 @@ class KBHelper: return doc except Exception as e: logger.error(f"上传文档失败: {e}") - # if file_path.exists(): - # file_path.unlink() + + if saved_file_path and saved_file_path.exists(): + try: + saved_file_path.unlink() + except Exception as file_error: + logger.warning( + f"清理原始文档文件失败 {saved_file_path}: {file_error}" + ) for media_path in media_paths: try: @@ -357,7 +361,7 @@ class KBHelper: except Exception as me: logger.warning(f"清理多媒体文件失败 {media_path}: {me}") - raise e + raise async def list_documents( self, @@ -412,7 +416,6 @@ class KBHelper: async with self.kb_db.get_db() as session: async with session.begin(): session.add(doc) - await session.commit() await session.refresh(doc) async def get_chunks_by_doc_id( diff --git a/astrbot/core/utils/io.py b/astrbot/core/utils/io.py index da1ea57af..b220136ff 100644 --- a/astrbot/core/utils/io.py +++ b/astrbot/core/utils/io.py @@ -58,7 +58,6 @@ def save_temp_img(img: Image.Image | bytes) -> str: timestamp = f"{int(time.time())}_{uuid.uuid4().hex[:8]}" p = os.path.join(temp_dir, f"io_temp_img_{timestamp}.jpg") - if isinstance(img, Image.Image): img.save(p) else: @@ -67,8 +66,6 @@ def save_temp_img(img: Image.Image | bytes) -> str: return p - - async def download_image_by_url( url: str, post: bool = False, @@ -207,7 +204,6 @@ async def download_file(url: str, path: str, show_progress: bool = False) -> Non logger.info("下载完成") - def file_to_base64(file_path: str) -> str: with open(file_path, "rb") as f: data_bytes = f.read() @@ -215,7 +211,6 @@ def file_to_base64(file_path: str) -> str: return "base64://" + base64_str - def get_local_ip_addresses() -> list[IPv4Address | IPv6Address]: net_interfaces = psutil.net_if_addrs() network_ips: list[IPv4Address | IPv6Address] = [] @@ -233,7 +228,6 @@ def get_local_ip_addresses() -> list[IPv4Address | IPv6Address]: return network_ips - async def get_public_ip_address() -> list[IPv4Address | IPv6Address]: urls = [ "https://api64.ipify.org", @@ -264,7 +258,6 @@ async def get_public_ip_address() -> list[IPv4Address | IPv6Address]: return list(found_ips.values()) - async def get_dashboard_version(): # First check user data directory (manually updated / downloaded dashboard). dist_dir = os.path.join(get_astrbot_data_path(), "dist") @@ -282,7 +275,6 @@ async def get_dashboard_version(): return None - async def download_dashboard( path: str | None = None, extract_path: str = "data",