diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 7977e439..ee9ab945 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -134,8 +134,11 @@ DEFAULT_CONFIG = { "persona": [], # deprecated "timezone": "Asia/Shanghai", "callback_api_base": "", - "default_kb_collection": "", # 默认知识库名称 + "default_kb_collection": "", # 默认知识库名称, 已经过时 "plugin_set": ["*"], # "*" 表示使用所有可用的插件, 空列表表示不使用任何插件 + "kb_names": [], # 默认知识库名称列表 + "kb_fusion_top_k": 20, # 知识库检索融合阶段返回结果数量 + "kb_final_top_k": 5, # 知识库检索最终返回结果数量 } @@ -1414,6 +1417,7 @@ CONFIG_METADATA_2 = { "description": "嵌入维度", "type": "int", "hint": "嵌入向量的维度。根据模型不同，可能需要调整，请参考具体模型的文档。此配置项请务必填写正确，否则将导致向量数据库无法正常工作。", + "_special": "get_embedding_dim", }, "embedding_model": { "description": "嵌入模型", @@ -2061,6 +2065,9 @@ CONFIG_METADATA_2 = { "default_kb_collection": { "type": "string", }, + "kb_names": {"type": "list", "items": {"type": "string"}}, + "kb_fusion_top_k": {"type": "int", "default": 20}, + "kb_final_top_k": {"type": "int", "default": 5}, }, }, } @@ -2139,10 +2146,22 @@ CONFIG_METADATA_3 = { "description": "知识库", "type": "object", "items": { - "default_kb_collection": { - "description": "默认使用的知识库", - "type": "string", + "kb_names": { + "description": "知识库列表", + "type": "list", + "items": {"type": "string"}, "_special": "select_knowledgebase", + "hint": "支持多选", + }, + "kb_fusion_top_k": { + "description": "融合检索结果数", + "type": "int", + "hint": "多个知识库检索结果融合后的返回结果数量", + }, + "kb_final_top_k": { + "description": "最终返回结果数", + "type": "int", + "hint": "从知识库中检索到的结果数量，越大可能获得越多相关信息，但也可能引入噪音。建议根据实际需求调整", }, }, }, diff --git a/astrbot/core/conversation_mgr.py b/astrbot/core/conversation_mgr.py index a6a2710f..8f8e2e0e 100644 --- a/astrbot/core/conversation_mgr.py +++ b/astrbot/core/conversation_mgr.py @@ -7,7 +7,7 @@ AstrBot 会话-对话管理器, 维护两个本地存储, 其中一个是 json import json from astrbot.core import sp -from typing import Dict, List +from typing import Dict, List, Callable, Awaitable from astrbot.core.db import BaseDatabase from astrbot.core.db.po import Conversation, ConversationV2 @@ -20,6 +20,38 @@ class ConversationManager: self.db = db_helper self.save_interval = 60 # 每 60 秒保存一次 + # 会话删除回调函数列表（用于级联清理，如知识库配置） + self._on_session_deleted_callbacks: List[Callable[[str], Awaitable[None]]] = [] + + def register_on_session_deleted( + self, callback: Callable[[str], Awaitable[None]] + ) -> None: + """注册会话删除回调函数 + + 其他模块可以注册回调来响应会话删除事件，实现级联清理。 + 例如：知识库模块可以注册回调来清理会话的知识库配置。 + + Args: + callback: 回调函数，接收会话ID (unified_msg_origin) 作为参数 + """ + self._on_session_deleted_callbacks.append(callback) + + async def _trigger_session_deleted(self, unified_msg_origin: str) -> None: + """触发会话删除回调 + + Args: + unified_msg_origin: 会话ID + """ + for callback in self._on_session_deleted_callbacks: + try: + await callback(unified_msg_origin) + except Exception as e: + from astrbot.core import logger + + logger.error( + f"会话删除回调执行失败 (session: {unified_msg_origin}): {e}" + ) + def _convert_conv_from_v2_to_v1(self, conv_v2: ConversationV2) -> Conversation: """将 ConversationV2 对象转换为 Conversation 对象""" created_at = int(conv_v2.created_at.timestamp()) @@ -106,6 +138,9 @@ class ConversationManager: self.session_conversations.pop(unified_msg_origin, None) await sp.session_remove(unified_msg_origin, "sel_conv_id") + # 触发会话删除回调（级联清理） + await self._trigger_session_deleted(unified_msg_origin) + async def get_curr_conversation_id(self, unified_msg_origin: str) -> str | None: """获取会话当前的对话 ID diff --git a/astrbot/core/core_lifecycle.py b/astrbot/core/core_lifecycle.py index 49963734..3d4b28c0 100644 --- a/astrbot/core/core_lifecycle.py +++ b/astrbot/core/core_lifecycle.py @@ -17,7 +17,6 @@ import os from .event_bus import EventBus from . import astrbot_config, html_renderer from asyncio import Queue -from typing import List from astrbot.core.pipeline.scheduler import PipelineScheduler, PipelineContext from astrbot.core.star import PluginManager from astrbot.core.platform.manager import PlatformManager @@ -36,6 +35,7 @@ from astrbot.core.umop_config_router import UmopConfigRouter from astrbot.core.astrbot_config_mgr import AstrBotConfigManager from astrbot.core.star.star_handler import star_handlers_registry, EventType from astrbot.core.star.star_handler import star_map +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager class AstrBotCoreLifecycle: @@ -122,6 +122,9 @@ class AstrBotCoreLifecycle: # 初始化平台消息历史管理器 self.platform_message_history_manager = PlatformMessageHistoryManager(self.db) + # 初始化知识库管理器 + self.kb_manager = KnowledgeBaseManager(self.provider_manager) + # 初始化提供给插件的上下文 self.star_context = Context( self.event_queue, @@ -133,6 +136,7 @@ class AstrBotCoreLifecycle: self.platform_message_history_manager, self.persona_mgr, self.astrbot_config_mgr, + self.kb_manager, ) # 初始化插件管理器 @@ -144,8 +148,9 @@ class AstrBotCoreLifecycle: # 根据配置实例化各个 Provider await self.provider_manager.initialize() - # 初始化消息事件流水线调度器 + await self.kb_manager.initialize() + # 初始化消息事件流水线调度器 self.pipeline_scheduler_mapping = await self.load_pipeline_scheduler() # 初始化更新器 @@ -160,7 +165,7 @@ class AstrBotCoreLifecycle: self.start_time = int(time.time()) # 初始化当前任务列表 - self.curr_tasks: List[asyncio.Task] = [] + self.curr_tasks: list[asyncio.Task] = [] # 根据配置实例化各个平台适配器 await self.platform_manager.initialize() @@ -245,6 +250,7 @@ class AstrBotCoreLifecycle: await self.provider_manager.terminate() await self.platform_manager.terminate() + await self.kb_manager.terminate() self.dashboard_shutdown_event.set() # 再次遍历curr_tasks等待每个任务真正结束 @@ -260,12 +266,13 @@ class AstrBotCoreLifecycle: """重启 AstrBot 核心生命周期管理类, 终止各个管理器并重新加载平台实例""" await self.provider_manager.terminate() await self.platform_manager.terminate() + await self.kb_manager.terminate() self.dashboard_shutdown_event.set() threading.Thread( target=self.astrbot_updator._reboot, name="restart", daemon=True ).start() - def load_platform(self) -> List[asyncio.Task]: + def load_platform(self) -> list[asyncio.Task]: """加载平台实例并返回所有平台实例的异步任务列表""" tasks = [] platform_insts = self.platform_manager.get_insts() diff --git a/astrbot/core/db/vec_db/base.py b/astrbot/core/db/vec_db/base.py index d6e5e56b..27fc9f3f 100644 --- a/astrbot/core/db/vec_db/base.py +++ b/astrbot/core/db/vec_db/base.py @@ -16,14 +16,42 @@ class BaseVecDB: pass @abc.abstractmethod - async def insert(self, content: str, metadata: dict = None, id: str = None) -> int: + async def insert( + self, content: str, metadata: dict | None = None, id: str | None = None + ) -> int: """ 插入一条文本和其对应向量，自动生成 ID 并保持一致性。 """ ... @abc.abstractmethod - async def retrieve(self, query: str, top_k: int = 5) -> list[Result]: + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> int: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + """ + ... + + @abc.abstractmethod + async def retrieve( + self, + query: str, + top_k: int = 5, + fetch_k: int = 20, + rerank: bool = False, + metadata_filters: dict | None = None, + ) -> list[Result]: """ 搜索最相似的文档。 Args: @@ -44,3 +72,6 @@ class BaseVecDB: bool: 删除是否成功 """ ... + + @abc.abstractmethod + async def close(self): ... diff --git a/astrbot/core/db/vec_db/faiss_impl/document_storage.py b/astrbot/core/db/vec_db/faiss_impl/document_storage.py index ee44da66..265c0cc4 100644 --- a/astrbot/core/db/vec_db/faiss_impl/document_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/document_storage.py @@ -1,59 +1,224 @@ -import aiosqlite import os +import json +from datetime import datetime +from contextlib import asynccontextmanager + +from sqlalchemy import Text, Column +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine +from sqlalchemy.orm import sessionmaker +from sqlmodel import Field, SQLModel, select, col, func, text, MetaData +from astrbot.core import logger + + +class BaseDocModel(SQLModel, table=False): + metadata = MetaData() + + +class Document(BaseDocModel, table=True): + """SQLModel for documents table.""" + + __tablename__ = "documents" # type: ignore + + id: int | None = Field( + default=None, primary_key=True, sa_column_kwargs={"autoincrement": True} + ) + doc_id: str = Field(nullable=False) + text: str = Field(nullable=False) + metadata_: str | None = Field(default=None, sa_column=Column("metadata", Text)) + created_at: datetime | None = Field(default=None) + updated_at: datetime | None = Field(default=None) class DocumentStorage: def __init__(self, db_path: str): self.db_path = db_path - self.connection = None + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.engine: AsyncEngine | None = None + self.async_session_maker: sessionmaker | None = None self.sqlite_init_path = os.path.join( os.path.dirname(__file__), "sqlite_init.sql" ) async def initialize(self): """Initialize the SQLite database and create the documents table if it doesn't exist.""" - if not os.path.exists(self.db_path): - await self.connect() - async with self.connection.cursor() as cursor: - with open(self.sqlite_init_path, "r", encoding="utf-8") as f: - sql_script = f.read() - await cursor.executescript(sql_script) - await self.connection.commit() - else: - await self.connect() + await self.connect() + async with self.engine.begin() as conn: # type: ignore + # Create tables using SQLModel + await conn.run_sync(BaseDocModel.metadata.create_all) + + try: + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN kb_doc_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.kb_doc_id')) STORED" + ) + ) + await conn.execute( + text( + "ALTER TABLE documents ADD COLUMN user_id TEXT " + "GENERATED ALWAYS AS (json_extract(metadata, '$.user_id')) STORED" + ) + ) + + # Create indexes + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_kb_doc_id ON documents(kb_doc_id)" + ) + ) + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_documents_user_id ON documents(user_id)" + ) + ) + except BaseException: + pass + + await conn.commit() async def connect(self): """Connect to the SQLite database.""" - self.connection = await aiosqlite.connect(self.db_path) + if self.engine is None: + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + future=True, + ) + self.async_session_maker = sessionmaker( + self.engine, # type: ignore + class_=AsyncSession, + expire_on_commit=False, + ) # type: ignore - async def get_documents(self, metadata_filters: dict, ids: list = None): + @asynccontextmanager + async def get_session(self): + """Context manager for database sessions.""" + async with self.async_session_maker() as session: # type: ignore + yield session + + async def get_documents( + self, + metadata_filters: dict, + ids: list | None = None, + offset: int | None = 0, + limit: int | None = 100, + ) -> list[dict]: """Retrieve documents by metadata filters and ids. Args: metadata_filters (dict): The metadata filters to apply. + ids (list | None): Optional list of document IDs to filter. + offset (int | None): Offset for pagination. + limit (int | None): Limit for pagination. Returns: - list: The list of document IDs(primary key, not doc_id) that match the filters. + list: The list of documents that match the filters. """ - # metadata filter -> SQL WHERE clause - where_clauses = [] - values = [] - for key, val in metadata_filters.items(): - where_clauses.append(f"json_extract(metadata, '$.{key}') = ?") - values.append(val) - if ids is not None and len(ids) > 0: - ids = [str(i) for i in ids if i != -1] - where_clauses.append("id IN ({})".format(",".join("?" * len(ids)))) - values.extend(ids) - where_sql = " AND ".join(where_clauses) or "1=1" + if self.engine is None: + logger.warning( + "Database connection is not initialized, returning empty result" + ) + return [] - result = [] - async with self.connection.cursor() as cursor: - sql = "SELECT * FROM documents WHERE " + where_sql - await cursor.execute(sql, values) - for row in await cursor.fetchall(): - result.append(await self.tuple_to_dict(row)) - return result + async with self.get_session() as session: + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + if ids is not None and len(ids) > 0: + valid_ids = [int(i) for i in ids if i != -1] + if valid_ids: + query = query.where(col(Document.id).in_(valid_ids)) + + if limit is not None: + query = query.limit(limit) + if offset is not None: + query = query.offset(offset) + + result = await session.execute(query) + documents = result.scalars().all() + + return [self._document_to_dict(doc) for doc in documents] + + async def insert_document(self, doc_id: str, text: str, metadata: dict) -> int: + """Insert a single document and return its integer ID. + + Args: + doc_id (str): The document ID (UUID string). + text (str): The document text. + metadata (dict): The document metadata. + + Returns: + int: The integer ID of the inserted document. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + session.add(document) + await session.flush() # Flush to get the ID + return document.id # type: ignore + + async def insert_documents_batch( + self, doc_ids: list[str], texts: list[str], metadatas: list[dict] + ) -> list[int]: + """Batch insert documents and return their integer IDs. + + Args: + doc_ids (list[str]): List of document IDs (UUID strings). + texts (list[str]): List of document texts. + metadatas (list[dict]): List of document metadata. + + Returns: + list[int]: List of integer IDs of the inserted documents. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + import json + + documents = [] + for doc_id, text, metadata in zip(doc_ids, texts, metadatas): + document = Document( + doc_id=doc_id, + text=text, + metadata_=json.dumps(metadata), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + documents.append(document) + session.add(document) + + await session.flush() # Flush to get all IDs + return [doc.id for doc in documents] # type: ignore + + async def delete_document_by_doc_id(self, doc_id: str): + """Delete a document by its doc_id. + + Args: + doc_id (str): The doc_id of the document to delete. + """ + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + await session.delete(document) async def get_document_by_doc_id(self, doc_id: str): """Retrieve a document by its doc_id. @@ -62,28 +227,91 @@ class DocumentStorage: doc_id (str): The doc_id of the document to retrieve. Returns: - dict: The document data. + dict: The document data or None if not found. """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT * FROM documents WHERE doc_id = ?", (doc_id,)) - row = await cursor.fetchone() - if row: - return await self.tuple_to_dict(row) - else: - return None + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + return self._document_to_dict(document) + return None async def update_document_by_doc_id(self, doc_id: str, new_text: str): - """Retrieve a document by its doc_id. + """Update a document by its doc_id. Args: doc_id (str): The doc_id. new_text (str): The new text to update the document with. """ - async with self.connection.cursor() as cursor: - await cursor.execute( - "UPDATE documents SET text = ? WHERE doc_id = ?", (new_text, doc_id) + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + async with session.begin(): + query = select(Document).where(col(Document.doc_id) == doc_id) + result = await session.execute(query) + document = result.scalar_one_or_none() + + if document: + document.text = new_text + document.updated_at = datetime.now() + session.add(document) + + async def delete_documents(self, metadata_filters: dict): + """Delete documents by their metadata filters. + + Args: + metadata_filters (dict): The metadata filters to apply. + """ + if self.engine is None: + logger.warning( + "Database connection is not initialized, skipping delete operation" ) - await self.connection.commit() + return + + async with self.get_session() as session: + async with session.begin(): + query = select(Document) + + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + documents = result.scalars().all() + + for doc in documents: + await session.delete(doc) + + async def count_documents(self, metadata_filters: dict | None = None) -> int: + """Count documents in the database. + + Args: + metadata_filters (dict | None): Metadata filters to apply. + + Returns: + int: The count of documents. + """ + if self.engine is None: + logger.warning("Database connection is not initialized, returning 0") + return 0 + + async with self.get_session() as session: + query = select(func.count(col(Document.id))) + + if metadata_filters: + for key, val in metadata_filters.items(): + query = query.where( + text(f"json_extract(metadata, '$.{key}') = :filter_{key}") + ).params(**{f"filter_{key}": val}) + + result = await session.execute(query) + count = result.scalar_one_or_none() + return count if count is not None else 0 async def get_user_ids(self) -> list[str]: """Retrieve all user IDs from the documents table. @@ -91,11 +319,38 @@ class DocumentStorage: Returns: list: A list of user IDs. """ - async with self.connection.cursor() as cursor: - await cursor.execute("SELECT DISTINCT user_id FROM documents") - rows = await cursor.fetchall() + assert self.engine is not None, "Database connection is not initialized." + + async with self.get_session() as session: + query = text( + "SELECT DISTINCT user_id FROM documents WHERE user_id IS NOT NULL" + ) + result = await session.execute(query) + rows = result.fetchall() return [row[0] for row in rows] + def _document_to_dict(self, document: Document) -> dict: + """Convert a Document model to a dictionary. + + Args: + document (Document): The document to convert. + + Returns: + dict: The converted dictionary. + """ + return { + "id": document.id, + "doc_id": document.doc_id, + "text": document.text, + "metadata": document.metadata_, + "created_at": document.created_at.isoformat() + if isinstance(document.created_at, datetime) + else document.created_at, + "updated_at": document.updated_at.isoformat() + if isinstance(document.updated_at, datetime) + else document.updated_at, + } + async def tuple_to_dict(self, row): """Convert a tuple to a dictionary. @@ -104,6 +359,8 @@ class DocumentStorage: Returns: dict: The converted dictionary. + + Note: This method is kept for backward compatibility but is no longer used internally. """ return { "id": row[0], @@ -116,6 +373,7 @@ class DocumentStorage: async def close(self): """Close the connection to the SQLite database.""" - if self.connection: - await self.connection.close() - self.connection = None + if self.engine: + await self.engine.dispose() + self.engine = None + self.async_session_maker = None diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index 18ee3189..2c0cc8df 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -9,7 +9,7 @@ import numpy as np class EmbeddingStorage: - def __init__(self, dimension: int, path: str = None): + def __init__(self, dimension: int, path: str | None = None): self.dimension = dimension self.path = path self.index = None @@ -18,7 +18,6 @@ class EmbeddingStorage: else: base_index = faiss.IndexFlatL2(dimension) self.index = faiss.IndexIDMap(base_index) - self.storage = {} async def insert(self, vector: np.ndarray, id: int): """插入向量 @@ -29,12 +28,29 @@ class EmbeddingStorage: Raises: ValueError: 如果向量的维度与存储的维度不匹配 """ + assert self.index is not None, "FAISS index is not initialized." if vector.shape[0] != self.dimension: raise ValueError( f"向量维度不匹配, 期望: {self.dimension}, 实际: {vector.shape[0]}" ) self.index.add_with_ids(vector.reshape(1, -1), np.array([id])) - self.storage[id] = vector + await self.save_index() + + async def insert_batch(self, vectors: np.ndarray, ids: list[int]): + """批量插入向量 + + Args: + vectors (np.ndarray): 要插入的向量数组 + ids (list[int]): 向量的ID列表 + Raises: + ValueError: 如果向量的维度与存储的维度不匹配 + """ + assert self.index is not None, "FAISS index is not initialized." + if vectors.shape[1] != self.dimension: + raise ValueError( + f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}" + ) + self.index.add_with_ids(vectors, np.array(ids)) await self.save_index() async def search(self, vector: np.ndarray, k: int) -> tuple: @@ -46,10 +62,22 @@ class EmbeddingStorage: Returns: tuple: (距离, 索引) """ + assert self.index is not None, "FAISS index is not initialized." faiss.normalize_L2(vector) distances, indices = self.index.search(vector, k) return distances, indices + async def delete(self, ids: list[int]): + """删除向量 + + Args: + ids (list[int]): 要删除的向量ID列表 + """ + assert self.index is not None, "FAISS index is not initialized." + id_array = np.array(ids, dtype=np.int64) + self.index.remove_ids(id_array) + await self.save_index() + async def save_index(self): """保存索引 diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index 7c2ae1c0..8a21538e 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -1,11 +1,12 @@ import uuid -import json +import time import numpy as np from .document_storage import DocumentStorage from .embedding_storage import EmbeddingStorage from ..base import Result, BaseVecDB from astrbot.core.provider.provider import EmbeddingProvider from astrbot.core.provider.provider import RerankProvider +from astrbot import logger class FaissVecDB(BaseVecDB): @@ -44,18 +45,56 @@ class FaissVecDB(BaseVecDB): vector = await self.embedding_provider.get_embedding(content) vector = np.array(vector, dtype=np.float32) - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute( - "INSERT INTO documents (doc_id, text, metadata) VALUES (?, ?, ?)", - (str_id, content, json.dumps(metadata)), - ) - await self.document_storage.connection.commit() - result = await self.document_storage.get_document_by_doc_id(str_id) - int_id = result["id"] - # 插入向量到 FAISS - await self.embedding_storage.insert(vector, int_id) - return int_id + # 使用 DocumentStorage 的方法插入文档 + int_id = await self.document_storage.insert_document(str_id, content, metadata) + + # 插入向量到 FAISS + await self.embedding_storage.insert(vector, int_id) + return int_id + + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> list[int]: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + + Args: + progress_callback: 进度回调函数，接收参数 (current, total) + """ + metadatas = metadatas or [{} for _ in contents] + ids = ids or [str(uuid.uuid4()) for _ in contents] + + start = time.time() + logger.debug(f"Generating embeddings for {len(contents)} contents...") + vectors = await self.embedding_provider.get_embeddings_batch( + contents, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + end = time.time() + logger.debug( + f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds." + ) + + # 使用 DocumentStorage 的批量插入方法 + int_ids = await self.document_storage.insert_documents_batch( + ids, contents, metadatas + ) + + # 批量插入向量到 FAISS + vectors_array = np.array(vectors).astype("float32") + await self.embedding_storage.insert_batch(vectors_array, int_ids) + return int_ids async def retrieve( self, @@ -119,23 +158,42 @@ class FaissVecDB(BaseVecDB): return top_k_results - async def delete(self, doc_id: int): + async def delete(self, doc_id: str): """ - 删除一条文档 + 删除一条文档块（chunk） """ - await self.document_storage.connection.execute( - "DELETE FROM documents WHERE doc_id = ?", (doc_id,) - ) - await self.document_storage.connection.commit() + # 获得对应的 int id + result = await self.document_storage.get_document_by_doc_id(doc_id) + int_id = result["id"] if result else None + if int_id is None: + return + + # 使用 DocumentStorage 的删除方法 + await self.document_storage.delete_document_by_doc_id(doc_id) + await self.embedding_storage.delete([int_id]) async def close(self): await self.document_storage.close() - async def count_documents(self) -> int: + async def count_documents(self, metadata_filter: dict | None = None) -> int: """ 计算文档数量 + + Args: + metadata_filter (dict | None): 元数据过滤器 """ - async with self.document_storage.connection.cursor() as cursor: - await cursor.execute("SELECT COUNT(*) FROM documents") - count = await cursor.fetchone() - return count[0] if count else 0 + count = await self.document_storage.count_documents( + metadata_filters=metadata_filter or {} + ) + return count + + async def delete_documents(self, metadata_filters: dict): + """ + 根据元数据过滤器删除文档 + """ + docs = await self.document_storage.get_documents( + metadata_filters=metadata_filters, offset=None, limit=None + ) + doc_ids: list[int] = [doc["id"] for doc in docs] + await self.embedding_storage.delete(doc_ids) + await self.document_storage.delete_documents(metadata_filters=metadata_filters) diff --git a/astrbot/core/knowledge_base/chunking/__init__.py b/astrbot/core/knowledge_base/chunking/__init__.py new file mode 100644 index 00000000..3124afe8 --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/__init__.py @@ -0,0 +1,11 @@ +""" +文档分块模块 +""" + +from .base import BaseChunker +from .fixed_size import FixedSizeChunker + +__all__ = [ + "BaseChunker", + "FixedSizeChunker", +] diff --git a/astrbot/core/knowledge_base/chunking/base.py b/astrbot/core/knowledge_base/chunking/base.py new file mode 100644 index 00000000..5aaf84ba --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/base.py @@ -0,0 +1,24 @@ +"""文档分块器基类 + +定义了文档分块处理的抽象接口。 +""" + +from abc import ABC, abstractmethod + + +class BaseChunker(ABC): + """分块器基类 + + 所有分块器都应该继承此类并实现 chunk 方法。 + """ + + @abstractmethod + async def chunk(self, text: str, **kwargs) -> list[str]: + """将文本分块 + + Args: + text: 输入文本 + + Returns: + list[str]: 分块后的文本列表 + """ diff --git a/astrbot/core/knowledge_base/chunking/fixed_size.py b/astrbot/core/knowledge_base/chunking/fixed_size.py new file mode 100644 index 00000000..c9b35d7d --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/fixed_size.py @@ -0,0 +1,57 @@ +"""固定大小分块器 + +按照固定的字符数将文本分块,支持重叠区域。 +""" + +from .base import BaseChunker + + +class FixedSizeChunker(BaseChunker): + """固定大小分块器 + + 按照固定的字符数分块,并支持块之间的重叠。 + """ + + def __init__(self, chunk_size: int = 512, chunk_overlap: int = 50): + """初始化分块器 + + Args: + chunk_size: 块的大小(字符数) + chunk_overlap: 块之间的重叠字符数 + """ + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + + async def chunk(self, text: str, **kwargs) -> list[str]: + """固定大小分块 + + Args: + text: 输入文本 + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + + Returns: + list[str]: 分块后的文本列表 + """ + chunk_size = kwargs.get("chunk_size", self.chunk_size) + chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + + chunks = [] + start = 0 + text_len = len(text) + + while start < text_len: + end = start + chunk_size + chunk = text[start:end] + + if chunk: + chunks.append(chunk) + + # 移动窗口,保留重叠部分 + start = end - chunk_overlap + + # 防止无限循环: 如果重叠过大,直接移到end + if start >= end or chunk_overlap >= chunk_size: + start = end + + return chunks diff --git a/astrbot/core/knowledge_base/chunking/recursive.py b/astrbot/core/knowledge_base/chunking/recursive.py new file mode 100644 index 00000000..21b76cba --- /dev/null +++ b/astrbot/core/knowledge_base/chunking/recursive.py @@ -0,0 +1,155 @@ +from collections.abc import Callable +from .base import BaseChunker + + +class RecursiveCharacterChunker(BaseChunker): + def __init__( + self, + chunk_size: int = 500, + chunk_overlap: int = 100, + length_function: Callable[[str], int] = len, + is_separator_regex: bool = False, + separators: list[str] | None = None, + ): + """ + 初始化递归字符文本分割器 + + Args: + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + length_function: 计算文本长度的函数 + is_separator_regex: 分隔符是否为正则表达式 + separators: 用于分割文本的分隔符列表，按优先级排序 + """ + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + self.length_function = length_function + self.is_separator_regex = is_separator_regex + + # 默认分隔符列表，按优先级从高到低 + self.separators = separators or [ + "\n\n", # 段落 + "\n", # 换行 + "。", # 中文句子 + "，", # 中文逗号 + ". ", # 句子 + ", ", # 逗号分隔 + " ", # 单词 + "", # 字符 + ] + + async def chunk(self, text: str, **kwargs) -> list[str]: + """ + 递归地将文本分割成块 + + Args: + text: 要分割的文本 + chunk_size: 每个文本块的最大大小 + chunk_overlap: 每个文本块之间的重叠部分大小 + + Returns: + 分割后的文本块列表 + """ + if not text: + return [] + + overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + chunk_size = kwargs.get("chunk_size", self.chunk_size) + + text_length = self.length_function(text) + if text_length <= chunk_size: + return [text] + + for separator in self.separators: + if separator == "": + return self._split_by_character(text, chunk_size, overlap) + + if separator in text: + splits = text.split(separator) + # 重新添加分隔符（除了最后一个片段） + splits = [s + separator for s in splits[:-1]] + [splits[-1]] + splits = [s for s in splits if s] + if len(splits) == 1: + continue + + # 递归合并分割后的文本块 + final_chunks = [] + current_chunk = [] + current_chunk_length = 0 + + for split in splits: + split_length = self.length_function(split) + + # 如果单个分割部分已经超过了chunk_size，需要递归分割 + if split_length > chunk_size: + # 先处理当前积累的块 + if current_chunk: + combined_text = "".join(current_chunk) + final_chunks.extend( + await self.chunk( + combined_text, + chunk_size=chunk_size, + chunk_overlap=overlap, + ) + ) + current_chunk = [] + current_chunk_length = 0 + + # 递归分割过大的部分 + final_chunks.extend( + await self.chunk( + split, chunk_size=chunk_size, chunk_overlap=overlap + ) + ) + # 如果添加这部分会使当前块超过chunk_size + elif current_chunk_length + split_length > chunk_size: + # 合并当前块并添加到结果中 + combined_text = "".join(current_chunk) + final_chunks.append(combined_text) + + # 处理重叠部分 + overlap_start = max(0, len(combined_text) - overlap) + if overlap_start > 0: + overlap_text = combined_text[overlap_start:] + current_chunk = [overlap_text, split] + current_chunk_length = ( + self.length_function(overlap_text) + split_length + ) + else: + current_chunk = [split] + current_chunk_length = split_length + else: + # 添加到当前块 + current_chunk.append(split) + current_chunk_length += split_length + + # 处理剩余的块 + if current_chunk: + final_chunks.append("".join(current_chunk)) + + return final_chunks + + return [text] + + def _split_by_character( + self, text: str, chunk_size: int | None = None, overlap: int | None = None + ) -> list[str]: + """ + 按字符级别分割文本 + + Args: + text: 要分割的文本 + + Returns: + 分割后的文本块列表 + """ + chunk_size = chunk_size or self.chunk_size + overlap = overlap or self.chunk_overlap + result = [] + for i in range(0, len(text), chunk_size - overlap): + end = min(i + chunk_size, len(text)) + result.append(text[i:end]) + if end == len(text): + break + + return result diff --git a/astrbot/core/knowledge_base/kb_db_sqlite.py b/astrbot/core/knowledge_base/kb_db_sqlite.py new file mode 100644 index 00000000..827d621d --- /dev/null +++ b/astrbot/core/knowledge_base/kb_db_sqlite.py @@ -0,0 +1,299 @@ +from contextlib import asynccontextmanager +from pathlib import Path + +from sqlmodel import col, desc +from sqlalchemy import text, func, select, update, delete +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine + +from astrbot.core import logger +from astrbot.core.knowledge_base.models import ( + BaseKBModel, + KBDocument, + KBMedia, + KnowledgeBase, +) +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB + + +class KBSQLiteDatabase: + def __init__(self, db_path: str = "data/knowledge_base/kb.db") -> None: + """初始化知识库数据库 + + Args: + db_path: 数据库文件路径, 默认为 data/knowledge_base/kb.db + """ + self.db_path = db_path + self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" + self.inited = False + + # 确保目录存在 + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + + # 创建异步引擎 + self.engine = create_async_engine( + self.DATABASE_URL, + echo=False, + pool_pre_ping=True, + pool_recycle=3600, + ) + + # 创建会话工厂 + self.async_session = async_sessionmaker( + self.engine, + class_=AsyncSession, + expire_on_commit=False, + ) + + @asynccontextmanager + async def get_db(self): + """获取数据库会话 + + 用法: + async with kb_db.get_db() as session: + # 执行数据库操作 + result = await session.execute(stmt) + """ + async with self.async_session() as session: + yield session + + async def initialize(self) -> None: + """初始化数据库,创建表并配置 SQLite 参数""" + async with self.engine.begin() as conn: + # 创建所有知识库相关表 + await conn.run_sync(BaseKBModel.metadata.create_all) + + # 配置 SQLite 性能优化参数 + await conn.execute(text("PRAGMA journal_mode=WAL")) + await conn.execute(text("PRAGMA synchronous=NORMAL")) + await conn.execute(text("PRAGMA cache_size=20000")) + await conn.execute(text("PRAGMA temp_store=MEMORY")) + await conn.execute(text("PRAGMA mmap_size=134217728")) + await conn.execute(text("PRAGMA optimize")) + await conn.commit() + + self.inited = True + + async def migrate_to_v1(self) -> None: + """执行知识库数据库 v1 迁移 + + 创建所有必要的索引以优化查询性能 + """ + async with self.get_db() as session: + session: AsyncSession + async with session.begin(): + # 创建知识库表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_kb_id " + "ON knowledge_bases(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_name " + "ON knowledge_bases(kb_name)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_kb_created_at " + "ON knowledge_bases(created_at)" + ) + ) + + # 创建文档表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_doc_id " + "ON kb_documents(doc_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_kb_id " + "ON kb_documents(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_name " + "ON kb_documents(doc_name)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_type " + "ON kb_documents(file_type)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_doc_created_at " + "ON kb_documents(created_at)" + ) + ) + + # 创建多媒体表索引 + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_media_id " + "ON kb_media(media_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_doc_id " + "ON kb_media(doc_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_kb_id ON kb_media(kb_id)" + ) + ) + await session.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_media_type " + "ON kb_media(media_type)" + ) + ) + + await session.commit() + + async def close(self) -> None: + """关闭数据库连接""" + await self.engine.dispose() + logger.info(f"知识库数据库已关闭: {self.db_path}") + + async def get_kb_by_id(self, kb_id: str) -> KnowledgeBase | None: + """根据 ID 获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_id) == kb_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def get_kb_by_name(self, kb_name: str) -> KnowledgeBase | None: + """根据名称获取知识库""" + async with self.get_db() as session: + stmt = select(KnowledgeBase).where(col(KnowledgeBase.kb_name) == kb_name) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_kbs(self, offset: int = 0, limit: int = 100) -> list[KnowledgeBase]: + """列出所有知识库""" + async with self.get_db() as session: + stmt = ( + select(KnowledgeBase) + .offset(offset) + .limit(limit) + .order_by(desc(KnowledgeBase.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_kbs(self) -> int: + """统计知识库数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KnowledgeBase.id))) + result = await session.execute(stmt) + return result.scalar() or 0 + + # ===== 文档查询 ===== + + async def get_document_by_id(self, doc_id: str) -> KBDocument | None: + """根据 ID 获取文档""" + async with self.get_db() as session: + stmt = select(KBDocument).where(col(KBDocument.doc_id) == doc_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def list_documents_by_kb( + self, kb_id: str, offset: int = 0, limit: int = 100 + ) -> list[KBDocument]: + """列出知识库的所有文档""" + async with self.get_db() as session: + stmt = ( + select(KBDocument) + .where(col(KBDocument.kb_id) == kb_id) + .offset(offset) + .limit(limit) + .order_by(desc(KBDocument.created_at)) + ) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def count_documents_by_kb(self, kb_id: str) -> int: + """统计知识库的文档数量""" + async with self.get_db() as session: + stmt = select(func.count(col(KBDocument.id))).where( + col(KBDocument.kb_id) == kb_id + ) + result = await session.execute(stmt) + return result.scalar() or 0 + + async def get_document_with_metadata(self, doc_id: str) -> dict | None: + async with self.get_db() as session: + stmt = ( + select(KBDocument, KnowledgeBase) + .join(KnowledgeBase, col(KBDocument.kb_id) == col(KnowledgeBase.kb_id)) + .where(col(KBDocument.doc_id) == doc_id) + ) + result = await session.execute(stmt) + row = result.first() + + if not row: + return None + + return { + "document": row[0], + "knowledge_base": row[1], + } + + async def delete_document_by_id(self, doc_id: str, vec_db: FaissVecDB): + """删除单个文档及其相关数据""" + # 在知识库表中删除 + async with self.get_db() as session: + async with session.begin(): + # 删除文档记录 + delete_stmt = delete(KBDocument).where(col(KBDocument.doc_id) == doc_id) + await session.execute(delete_stmt) + await session.commit() + + # 在 vec db 中删除相关向量 + await vec_db.delete_documents(metadata_filters={"kb_doc_id": doc_id}) + + # ===== 多媒体查询 ===== + + async def list_media_by_doc(self, doc_id: str) -> list[KBMedia]: + """列出文档的所有多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.doc_id) == doc_id) + result = await session.execute(stmt) + return list(result.scalars().all()) + + async def get_media_by_id(self, media_id: str) -> KBMedia | None: + """根据 ID 获取多媒体资源""" + async with self.get_db() as session: + stmt = select(KBMedia).where(col(KBMedia.media_id) == media_id) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + async def update_kb_stats(self, kb_id: str, vec_db: FaissVecDB) -> None: + """更新知识库统计信息""" + chunk_cnt = await vec_db.count_documents() + + async with self.get_db() as session: + async with session.begin(): + update_stmt = ( + update(KnowledgeBase) + .where(col(KnowledgeBase.kb_id) == kb_id) + .values( + doc_count=select(func.count(col(KBDocument.id))) + .where(col(KBDocument.kb_id) == kb_id) + .scalar_subquery(), + chunk_count=chunk_cnt, + ) + ) + + await session.execute(update_stmt) + await session.commit() diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py new file mode 100644 index 00000000..5c07c93b --- /dev/null +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -0,0 +1,352 @@ +import uuid +import aiofiles +import json +from pathlib import Path +from .models import KnowledgeBase, KBDocument, KBMedia +from .kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.db.vec_db.base import BaseVecDB +from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB +from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider +from astrbot.core.provider.manager import ProviderManager +from .parsers.base import BaseParser +from .chunking.base import BaseChunker +from astrbot.core import logger + + +class KBHelper: + vec_db: BaseVecDB + kb: KnowledgeBase + + def __init__( + self, + kb_db: KBSQLiteDatabase, + kb: KnowledgeBase, + provider_manager: ProviderManager, + kb_root_dir: str, + chunker: BaseChunker, + parsers: dict[str, BaseParser], + ): + self.kb_db = kb_db + self.kb = kb + self.prov_mgr = provider_manager + self.kb_root_dir = kb_root_dir + self.parsers = parsers + self.chunker = chunker + + self.kb_dir = Path(self.kb_root_dir) / self.kb.kb_id + self.kb_medias_dir = Path(self.kb_dir) / "medias" / self.kb.kb_id + self.kb_files_dir = Path(self.kb_dir) / "files" / self.kb.kb_id + + self.kb_medias_dir.mkdir(parents=True, exist_ok=True) + self.kb_files_dir.mkdir(parents=True, exist_ok=True) + + async def initialize(self): + await self._ensure_vec_db() + + async def get_ep(self) -> EmbeddingProvider: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + ep: EmbeddingProvider = await self.prov_mgr.get_provider_by_id( + self.kb.embedding_provider_id + ) # type: ignore + if not ep: + raise ValueError( + f"无法找到 ID 为 {self.kb.embedding_provider_id} 的 Embedding Provider" + ) + return ep + + async def get_rp(self) -> RerankProvider | None: + if not self.kb.rerank_provider_id: + return None + rp: RerankProvider = await self.prov_mgr.get_provider_by_id( + self.kb.rerank_provider_id + ) # type: ignore + if not rp: + raise ValueError( + f"无法找到 ID 为 {self.kb.rerank_provider_id} 的 Rerank Provider" + ) + return rp + + async def _ensure_vec_db(self) -> FaissVecDB: + if not self.kb.embedding_provider_id: + raise ValueError(f"知识库 {self.kb.kb_name} 未配置 Embedding Provider") + + ep = await self.get_ep() + rp = await self.get_rp() + + vec_db = FaissVecDB( + doc_store_path=str(self.kb_dir / "doc.db"), + index_store_path=str(self.kb_dir / "index.faiss"), + embedding_provider=ep, + rerank_provider=rp, + ) + await vec_db.initialize() + self.vec_db = vec_db + return vec_db + + async def delete_vec_db(self): + """删除知识库的向量数据库和所有相关文件""" + import shutil + + await self.terminate() + if self.kb_dir.exists(): + shutil.rmtree(self.kb_dir) + + async def terminate(self): + if self.vec_db: + await self.vec_db.close() + + async def upload_document( + self, + file_name: str, + file_content: bytes, + file_type: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> KBDocument: + """上传并处理文档（带原子性保证和失败清理） + + 流程: + 1. 保存原始文件 + 2. 解析文档内容 + 3. 提取多媒体资源 + 4. 分块处理 + 5. 生成向量并存储 + 6. 保存元数据（事务） + 7. 更新统计 + + Args: + progress_callback: 进度回调函数，接收参数 (stage, current, total) + - stage: 当前阶段 ('parsing', 'chunking', 'embedding') + - current: 当前进度 + - total: 总数 + """ + await self._ensure_vec_db() + doc_id = str(uuid.uuid4()) + media_paths: list[Path] = [] + + # 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: + # 阶段1: 解析文档 + if progress_callback: + await progress_callback("parsing", 0, 100) + + parser = self.parsers.get(file_type) + if not parser: + raise ValueError(f"不支持的文件类型: {file_type}") + parse_result = await parser.parse(file_content, file_name) + text_content = parse_result.text + media_items = parse_result.media + + if progress_callback: + await progress_callback("parsing", 100, 100) + + # 保存媒体文件 + saved_media = [] + for media_item in media_items: + media = await self._save_media( + doc_id=doc_id, + media_type=media_item.media_type, + file_name=media_item.file_name, + content=media_item.content, + mime_type=media_item.mime_type, + ) + saved_media.append(media) + media_paths.append(Path(media.file_path)) + + # 阶段2: 分块 + if progress_callback: + await progress_callback("chunking", 0, 100) + + chunks_text = await self.chunker.chunk( + text_content, chunk_size=chunk_size, chunk_overlap=chunk_overlap + ) + contents = [] + metadatas = [] + for idx, chunk_text in enumerate(chunks_text): + contents.append(chunk_text) + metadatas.append( + { + "kb_id": self.kb.kb_id, + "kb_doc_id": doc_id, + "chunk_index": idx, + } + ) + + if progress_callback: + await progress_callback("chunking", 100, 100) + + # 阶段3: 生成向量（带进度回调） + async def embedding_progress_callback(current, total): + if progress_callback: + await progress_callback("embedding", current, total) + + await self.vec_db.insert_batch( + contents=contents, + metadatas=metadatas, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=embedding_progress_callback, + ) + + # 保存文档的元数据 + doc = KBDocument( + doc_id=doc_id, + kb_id=self.kb.kb_id, + doc_name=file_name, + file_type=file_type, + file_size=len(file_content), + # file_path=str(file_path), + file_path="", + chunk_count=len(chunks_text), + media_count=0, + ) + 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) + + vec_db: FaissVecDB = self.vec_db # type: ignore + await self.kb_db.update_kb_stats(kb_id=self.kb.kb_id, vec_db=vec_db) + await self.refresh_kb() + await self.refresh_document(doc_id) + return doc + except Exception as e: + logger.error(f"上传文档失败: {e}") + # if file_path.exists(): + # file_path.unlink() + + for media_path in media_paths: + try: + if media_path.exists(): + media_path.unlink() + except Exception as me: + logger.warning(f"清理多媒体文件失败 {media_path}: {me}") + + raise e + + async def list_documents( + self, offset: int = 0, limit: int = 100 + ) -> list[KBDocument]: + """列出知识库的所有文档""" + docs = await self.kb_db.list_documents_by_kb(self.kb.kb_id, offset, limit) + return docs + + async def get_document(self, doc_id: str) -> KBDocument | None: + """获取单个文档""" + doc = await self.kb_db.get_document_by_id(doc_id) + return doc + + async def delete_document(self, doc_id: str): + """删除单个文档及其相关数据""" + await self.kb_db.delete_document_by_id( + doc_id=doc_id, + vec_db=self.vec_db, # type: ignore + ) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + + async def delete_chunk(self, chunk_id: str, doc_id: str): + """删除单个文本块及其相关数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + await vec_db.delete(chunk_id) + await self.kb_db.update_kb_stats( + kb_id=self.kb.kb_id, + vec_db=self.vec_db, # type: ignore + ) + await self.refresh_kb() + await self.refresh_document(doc_id) + + async def refresh_kb(self): + if self.kb: + kb = await self.kb_db.get_kb_by_id(self.kb.kb_id) + if kb: + self.kb = kb + + async def refresh_document(self, doc_id: str) -> None: + """更新文档的元数据""" + doc = await self.get_document(doc_id) + if not doc: + raise ValueError(f"无法找到 ID 为 {doc_id} 的文档") + chunk_count = await self.get_chunk_count_by_doc_id(doc_id) + doc.chunk_count = chunk_count + 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( + self, doc_id: str, offset: int = 0, limit: int = 100 + ) -> list[dict]: + """获取文档的所有块及其元数据""" + vec_db: FaissVecDB = self.vec_db # type: ignore + chunks = await vec_db.document_storage.get_documents( + metadata_filters={"kb_doc_id": doc_id}, offset=offset, limit=limit + ) + result = [] + for chunk in chunks: + chunk_md = json.loads(chunk["metadata"]) + result.append( + { + "chunk_id": chunk["doc_id"], + "doc_id": chunk_md["kb_doc_id"], + "kb_id": chunk_md["kb_id"], + "chunk_index": chunk_md["chunk_index"], + "content": chunk["text"], + "char_count": len(chunk["text"]), + } + ) + return result + + async def get_chunk_count_by_doc_id(self, doc_id: str) -> int: + """获取文档的块数量""" + vec_db: FaissVecDB = self.vec_db # type: ignore + count = await vec_db.count_documents(metadata_filter={"kb_doc_id": doc_id}) + return count + + async def _save_media( + self, + doc_id: str, + media_type: str, + file_name: str, + content: bytes, + mime_type: str, + ) -> KBMedia: + """保存多媒体资源""" + media_id = str(uuid.uuid4()) + ext = Path(file_name).suffix + + # 保存文件 + file_path = self.kb_medias_dir / doc_id / f"{media_id}{ext}" + file_path.parent.mkdir(parents=True, exist_ok=True) + async with aiofiles.open(file_path, "wb") as f: + await f.write(content) + + media = KBMedia( + media_id=media_id, + doc_id=doc_id, + kb_id=self.kb.kb_id, + media_type=media_type, + file_name=file_name, + file_path=str(file_path), + file_size=len(content), + mime_type=mime_type, + ) + + return media diff --git a/astrbot/core/knowledge_base/kb_mgr.py b/astrbot/core/knowledge_base/kb_mgr.py new file mode 100644 index 00000000..4f64383c --- /dev/null +++ b/astrbot/core/knowledge_base/kb_mgr.py @@ -0,0 +1,298 @@ +import traceback +from pathlib import Path +from astrbot.core import logger +from astrbot.core.provider.manager import ProviderManager + +from .retrieval.manager import RetrievalManager, RetrievalResult +from .retrieval.sparse_retriever import SparseRetriever +from .retrieval.rank_fusion import RankFusion +from .kb_db_sqlite import KBSQLiteDatabase + +from .parsers.text_parser import TextParser +from .parsers.pdf_parser import PDFParser + +# from .chunking.fixed_size import FixedSizeChunker +from .chunking.recursive import RecursiveCharacterChunker +from .kb_helper import KBHelper + +from .models import KnowledgeBase + + +FILES_PATH = "data/knowledge_base" +DB_PATH = Path(FILES_PATH) / "kb.db" +"""Knowledge Base storage root directory""" +PARSERS = { + "txt": TextParser(), + "md": TextParser(), + "markdown": TextParser(), + "pdf": PDFParser(), +} +CHUNKER = RecursiveCharacterChunker() + + +class KnowledgeBaseManager: + kb_db: KBSQLiteDatabase + retrieval_manager: RetrievalManager + + def __init__( + self, + provider_manager: ProviderManager, + ): + Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True) + self.provider_manager = provider_manager + self._session_deleted_callback_registered = False + + self.kb_insts: dict[str, KBHelper] = {} + + async def initialize(self): + """初始化知识库模块""" + try: + logger.info("正在初始化知识库模块...") + + # 初始化数据库 + await self._init_kb_database() + + # 初始化检索管理器 + sparse_retriever = SparseRetriever(self.kb_db) + rank_fusion = RankFusion(self.kb_db) + self.retrieval_manager = RetrievalManager( + sparse_retriever=sparse_retriever, + rank_fusion=rank_fusion, + kb_db=self.kb_db, + ) + await self.load_kbs() + + except ImportError as e: + logger.error(f"知识库模块导入失败: {e}") + logger.warning("请确保已安装所需依赖: pypdf, aiofiles, Pillow, rank-bm25") + except Exception as e: + logger.error(f"知识库模块初始化失败: {e}") + logger.error(traceback.format_exc()) + + async def _init_kb_database(self): + self.kb_db = KBSQLiteDatabase(DB_PATH.as_posix()) + await self.kb_db.initialize() + await self.kb_db.migrate_to_v1() + logger.info(f"KnowledgeBase database initialized: {DB_PATH}") + + async def load_kbs(self): + """加载所有知识库实例""" + kb_records = await self.kb_db.list_kbs() + for record in kb_records: + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=record, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + parsers=PARSERS, + ) + await kb_helper.initialize() + self.kb_insts[record.kb_id] = kb_helper + + async def create_kb( + self, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper: + """创建新的知识库实例""" + kb = KnowledgeBase( + kb_name=kb_name, + description=description, + emoji=emoji or "📚", + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size if chunk_size is not None else 512, + chunk_overlap=chunk_overlap if chunk_overlap is not None else 50, + top_k_dense=top_k_dense if top_k_dense is not None else 50, + top_k_sparse=top_k_sparse if top_k_sparse is not None else 50, + top_m_final=top_m_final if top_m_final is not None else 5, + ) + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + kb_helper = KBHelper( + kb_db=self.kb_db, + kb=kb, + provider_manager=self.provider_manager, + kb_root_dir=FILES_PATH, + chunker=CHUNKER, + parsers=PARSERS, + ) + await kb_helper.initialize() + self.kb_insts[kb.kb_id] = kb_helper + return kb_helper + + async def get_kb(self, kb_id: str) -> KBHelper | None: + """获取知识库实例""" + if kb_id in self.kb_insts: + return self.kb_insts[kb_id] + + async def get_kb_by_name(self, kb_name: str) -> KBHelper | None: + """通过名称获取知识库实例""" + for kb_helper in self.kb_insts.values(): + if kb_helper.kb.kb_name == kb_name: + return kb_helper + return None + + async def delete_kb(self, kb_id: str) -> bool: + """删除知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return False + + await kb_helper.delete_vec_db() + async with self.kb_db.get_db() as session: + await session.delete(kb_helper.kb) + await session.commit() + + self.kb_insts.pop(kb_id, None) + return True + + async def list_kbs(self) -> list[KnowledgeBase]: + """列出所有知识库实例""" + kbs = [kb_helper.kb for kb_helper in self.kb_insts.values()] + return kbs + + async def update_kb( + self, + kb_id: str, + kb_name: str, + description: str | None = None, + emoji: str | None = None, + embedding_provider_id: str | None = None, + rerank_provider_id: str | None = None, + chunk_size: int | None = None, + chunk_overlap: int | None = None, + top_k_dense: int | None = None, + top_k_sparse: int | None = None, + top_m_final: int | None = None, + ) -> KBHelper | None: + """更新知识库实例""" + kb_helper = await self.get_kb(kb_id) + if not kb_helper: + return None + + kb = kb_helper.kb + if kb_name is not None: + kb.kb_name = kb_name + if description is not None: + kb.description = description + if emoji is not None: + kb.emoji = emoji + if embedding_provider_id is not None: + kb.embedding_provider_id = embedding_provider_id + kb.rerank_provider_id = rerank_provider_id # 允许设置为 None + if chunk_size is not None: + kb.chunk_size = chunk_size + if chunk_overlap is not None: + kb.chunk_overlap = chunk_overlap + if top_k_dense is not None: + kb.top_k_dense = top_k_dense + if top_k_sparse is not None: + kb.top_k_sparse = top_k_sparse + if top_m_final is not None: + kb.top_m_final = top_m_final + async with self.kb_db.get_db() as session: + session.add(kb) + await session.commit() + await session.refresh(kb) + + return kb_helper + + async def retrieve( + self, + query: str, + kb_names: list[str], + top_k_fusion: int = 20, + top_m_final: int = 5, + ) -> dict | None: + """从指定知识库中检索相关内容""" + kb_ids = [] + kb_id_helper_map = {} + for kb_name in kb_names: + if kb_helper := await self.get_kb_by_name(kb_name): + kb_ids.append(kb_helper.kb.kb_id) + kb_id_helper_map[kb_helper.kb.kb_id] = kb_helper + + if not kb_ids: + return {} + + results = await self.retrieval_manager.retrieve( + query=query, + kb_ids=kb_ids, + kb_id_helper_map=kb_id_helper_map, + top_k_fusion=top_k_fusion, + top_m_final=top_m_final, + ) + if not results: + return None + + context_text = self._format_context(results) + + results_dict = [ + { + "chunk_id": r.chunk_id, + "doc_id": r.doc_id, + "kb_id": r.kb_id, + "kb_name": r.kb_name, + "doc_name": r.doc_name, + "chunk_index": r.metadata.get("chunk_index", 0), + "content": r.content, + "score": r.score, + "char_count": r.metadata.get("char_count", 0), + } + for r in results + ] + + return { + "context_text": context_text, + "results": results_dict, + } + + def _format_context(self, results: list[RetrievalResult]) -> str: + """格式化知识上下文 + + Args: + results: 检索结果列表 + + Returns: + str: 格式化的上下文文本 + """ + lines = ["以下是相关的知识库内容,请参考这些信息回答用户的问题:\n"] + + for i, result in enumerate(results, 1): + lines.append(f"【知识 {i}】") + lines.append(f"来源: {result.kb_name} / {result.doc_name}") + lines.append(f"内容: {result.content}") + lines.append(f"相关度: {result.score:.2f}") + lines.append("") + + return "\n".join(lines) + + async def terminate(self): + """终止所有知识库实例,关闭数据库连接""" + for kb_id, kb_helper in self.kb_insts.items(): + try: + await kb_helper.terminate() + except Exception as e: + logger.error(f"关闭知识库 {kb_id} 失败: {e}") + + self.kb_insts.clear() + + # 关闭元数据数据库 + if hasattr(self, "kb_db") and self.kb_db: + try: + await self.kb_db.close() + except Exception as e: + logger.error(f"关闭知识库元数据数据库失败: {e}") diff --git a/astrbot/core/knowledge_base/models.py b/astrbot/core/knowledge_base/models.py new file mode 100644 index 00000000..010d6113 --- /dev/null +++ b/astrbot/core/knowledge_base/models.py @@ -0,0 +1,114 @@ +import uuid +from datetime import datetime, timezone + +from sqlmodel import Field, SQLModel, Text, UniqueConstraint, MetaData + + +class BaseKBModel(SQLModel, table=False): + metadata = MetaData() + + +class KnowledgeBase(BaseKBModel, table=True): + """知识库表 + + 存储知识库的基本信息和统计数据。 + """ + + __tablename__ = "knowledge_bases" # type: ignore + + id: int | None = Field( + primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + ) + kb_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + kb_name: str = Field(max_length=100, nullable=False) + description: str | None = Field(default=None, sa_type=Text) + emoji: str | None = Field(default="📚", max_length=10) + embedding_provider_id: str | None = Field(default=None, max_length=100) + rerank_provider_id: str | None = Field(default=None, max_length=100) + # 分块配置参数 + chunk_size: int | None = Field(default=512, nullable=True) + chunk_overlap: int | None = Field(default=50, nullable=True) + # 检索配置参数 + top_k_dense: int | None = Field(default=50, nullable=True) + top_k_sparse: int | None = Field(default=50, nullable=True) + top_m_final: int | None = Field(default=5, nullable=True) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, + ) + doc_count: int = Field(default=0, nullable=False) + chunk_count: int = Field(default=0, nullable=False) + + __table_args__ = ( + UniqueConstraint( + "kb_name", + name="uix_kb_name", + ), + ) + + +class KBDocument(BaseKBModel, table=True): + """文档表 + + 存储上传到知识库的文档元数据。 + """ + + __tablename__ = "kb_documents" # type: ignore + + id: int | None = Field( + primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + ) + doc_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + kb_id: str = Field(max_length=36, nullable=False, index=True) + doc_name: str = Field(max_length=255, nullable=False) + file_type: str = Field(max_length=20, nullable=False) + file_size: int = Field(nullable=False) + file_path: str = Field(max_length=512, nullable=False) + chunk_count: int = Field(default=0, nullable=False) + media_count: int = Field(default=0, nullable=False) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc), + sa_column_kwargs={"onupdate": datetime.now(timezone.utc)}, + ) + + +class KBMedia(BaseKBModel, table=True): + """多媒体资源表 + + 存储从文档中提取的图片、视频等多媒体资源。 + """ + + __tablename__ = "kb_media" # type: ignore + + id: int | None = Field( + primary_key=True, sa_column_kwargs={"autoincrement": True}, default=None + ) + media_id: str = Field( + max_length=36, + nullable=False, + unique=True, + default_factory=lambda: str(uuid.uuid4()), + index=True, + ) + doc_id: str = Field(max_length=36, nullable=False, index=True) + kb_id: str = Field(max_length=36, nullable=False, index=True) + media_type: str = Field(max_length=20, nullable=False) + file_name: str = Field(max_length=255, nullable=False) + file_path: str = Field(max_length=512, nullable=False) + file_size: int = Field(nullable=False) + mime_type: str = Field(max_length=100, nullable=False) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) diff --git a/astrbot/core/knowledge_base/parsers/__init__.py b/astrbot/core/knowledge_base/parsers/__init__.py new file mode 100644 index 00000000..6851edeb --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/__init__.py @@ -0,0 +1,15 @@ +""" +文档解析器模块 +""" + +from .base import BaseParser, MediaItem, ParseResult +from .text_parser import TextParser +from .pdf_parser import PDFParser + +__all__ = [ + "BaseParser", + "MediaItem", + "ParseResult", + "TextParser", + "PDFParser", +] diff --git a/astrbot/core/knowledge_base/parsers/base.py b/astrbot/core/knowledge_base/parsers/base.py new file mode 100644 index 00000000..1c571db2 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/base.py @@ -0,0 +1,50 @@ +"""文档解析器基类和数据结构 + +定义了文档解析器的抽象接口和相关数据类。 +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class MediaItem: + """多媒体项 + + 表示从文档中提取的多媒体资源。 + """ + + media_type: str # image, video + file_name: str + content: bytes + mime_type: str + + +@dataclass +class ParseResult: + """解析结果 + + 包含解析后的文本内容和提取的多媒体资源。 + """ + + text: str + media: list[MediaItem] + + +class BaseParser(ABC): + """文档解析器基类 + + 所有文档解析器都应该继承此类并实现 parse 方法。 + """ + + @abstractmethod + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析文档 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 解析结果 + """ diff --git a/astrbot/core/knowledge_base/parsers/pdf_parser.py b/astrbot/core/knowledge_base/parsers/pdf_parser.py new file mode 100644 index 00000000..fca62687 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/pdf_parser.py @@ -0,0 +1,100 @@ +"""PDF 文件解析器 + +支持解析 PDF 文件中的文本和图片资源。 +""" + +import io + +from pypdf import PdfReader + +from astrbot.core.knowledge_base.parsers.base import ( + BaseParser, + MediaItem, + ParseResult, +) + + +class PDFParser(BaseParser): + """PDF 文档解析器 + + 提取 PDF 中的文本内容和嵌入的图片资源。 + """ + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析 PDF 文件 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 包含文本和图片的解析结果 + """ + pdf_file = io.BytesIO(file_content) + reader = PdfReader(pdf_file) + + text_parts = [] + media_items = [] + + # 提取文本 + for page in reader.pages: + text = page.extract_text() + if text: + text_parts.append(text) + + # 提取图片 + image_counter = 0 + for page_num, page in enumerate(reader.pages): + try: + # 安全检查 Resources + if "/Resources" not in page: + continue + + resources = page["/Resources"] + if not resources or "/XObject" not in resources: # type: ignore + continue + + xobjects = resources["/XObject"].get_object() # type: ignore + if not xobjects: + continue + + for obj_name in xobjects: + try: + obj = xobjects[obj_name] + + if obj.get("/Subtype") != "/Image": + continue + + # 提取图片数据 + image_data = obj.get_data() + + # 确定格式 + filter_type = obj.get("/Filter", "") + if filter_type == "/DCTDecode": + ext = "jpg" + mime_type = "image/jpeg" + elif filter_type == "/FlateDecode": + ext = "png" + mime_type = "image/png" + else: + ext = "png" + mime_type = "image/png" + + image_counter += 1 + media_items.append( + MediaItem( + media_type="image", + file_name=f"page_{page_num}_img_{image_counter}.{ext}", + content=image_data, + mime_type=mime_type, + ) + ) + except Exception: + # 单个图片提取失败不影响整体 + continue + except Exception: + # 页面处理失败不影响其他页面 + continue + + full_text = "\n\n".join(text_parts) + return ParseResult(text=full_text, media=media_items) diff --git a/astrbot/core/knowledge_base/parsers/text_parser.py b/astrbot/core/knowledge_base/parsers/text_parser.py new file mode 100644 index 00000000..49a95a95 --- /dev/null +++ b/astrbot/core/knowledge_base/parsers/text_parser.py @@ -0,0 +1,41 @@ +"""文本文件解析器 + +支持解析 TXT 和 Markdown 文件。 +""" + +from astrbot.core.knowledge_base.parsers.base import BaseParser, ParseResult + + +class TextParser(BaseParser): + """TXT/MD 文本解析器 + + 支持多种字符编码的自动检测。 + """ + + async def parse(self, file_content: bytes, file_name: str) -> ParseResult: + """解析文本文件 + + 尝试使用多种编码解析文件内容。 + + Args: + file_content: 文件内容 + file_name: 文件名 + + Returns: + ParseResult: 解析结果,不包含多媒体资源 + + Raises: + ValueError: 如果无法解码文件 + """ + # 尝试多种编码 + 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=[]) diff --git a/astrbot/core/knowledge_base/retrieval/__init__.py b/astrbot/core/knowledge_base/retrieval/__init__.py new file mode 100644 index 00000000..16a5e664 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/__init__.py @@ -0,0 +1,16 @@ +""" +检索模块 +""" + +from .manager import RetrievalManager, RetrievalResult +from .sparse_retriever import SparseRetriever, SparseResult +from .rank_fusion import RankFusion, FusedResult + +__all__ = [ + "RetrievalManager", + "RetrievalResult", + "SparseRetriever", + "SparseResult", + "RankFusion", + "FusedResult", +] diff --git a/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt new file mode 100644 index 00000000..84b26283 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/hit_stopwords.txt @@ -0,0 +1,767 @@ +——— +》）， +）÷（１－ +”， +）、 +＝（ +: +→ +℃ +& +* +一一 +~~~~ +’ +. +『 +.一 +./ +-- +』 +＝″ +【 +［＊］ +｝＞ +［⑤］］ +［①Ｄ］ +ｃ］ +ｎｇ昉 +＊ +// +［ +］ +［②ｅ］ +［②ｇ］ +＝｛ +} +，也 +‘ +Ａ +［①⑥］ +［②Ｂ］ +［①ａ］ +［④ａ］ +［①③］ +［③ｈ］ +③］ +１． +－－ +［②ｂ］ +’‘ +××× +［①⑧］ +０：２ +＝［ +［⑤ｂ］ +［②ｃ］ +［④ｂ］ +［②③］ +［③ａ］ +［④ｃ］ +［①⑤］ +［①⑦］ +［①ｇ］ +∈［ +［①⑨］ +［①④］ +［①ｃ］ +［②ｆ］ +［②⑧］ +［②①］ +［①Ｃ］ +［③ｃ］ +［③ｇ］ +［②⑤］ +［②②］ +一. +［①ｈ］ +.数 +［］ +［①Ｂ］ +数/ +［①ｉ］ +［③ｅ］ +［①①］ +［④ｄ］ +［④ｅ］ +［③ｂ］ +［⑤ａ］ +［①Ａ］ +［②⑧］ +［②⑦］ +［①ｄ］ +［②ｊ］ +〕〔 +］［ +:// +′∈ +［②④ +［⑤ｅ］ +１２％ +ｂ］ +... +................... +…………………………………………………③ +ＺＸＦＩＴＬ +［③Ｆ］ +」 +［①ｏ］ +］∧′＝［ +∪φ∈ +′｜ +｛－ +②ｃ +｝ +［③①］ +Ｒ．Ｌ． +［①Ｅ］ +Ψ +－［＊］－ +↑ +.日 +［②ｄ］ +［② +［②⑦］ +［②②］ +［③ｅ］ +［①ｉ］ +［①Ｂ］ +［①ｈ］ +［①ｄ］ +［①ｇ］ +［①②］ +［②ａ］ +ｆ］ +［⑩］ +ａ］ +［①ｅ］ +［②ｈ］ +［②⑥］ +［③ｄ］ +［②⑩］ +ｅ］ +〉 +】 +元／吨 +［②⑩］ +２．３％ +５：０ +［①］ +:: +［②］ +［③］ +［④］ +［⑤］ +［⑥］ +［⑦］ +［⑧］ +［⑨］ +…… +—— +? +、 +。 +“ +” +《 +》 +！ +， +： +； +？ +． +, +． +' +? +· +——— +── +? +— +< +> +（ +） +〔 +〕 +[ +] +( +) +- ++ +～ +× +／ +/ +① +② +③ +④ +⑤ +⑥ +⑦ +⑧ +⑨ +⑩ +Ⅲ +В +" +; +# +@ +γ +μ +φ +φ． +× +Δ +■ +▲ +sub +exp +sup +sub +Lex +＃ +％ +＆ +＇ +＋ +＋ξ +＋＋ +－ +－β +＜ +＜± +＜Δ +＜λ +＜φ +＜＜ += +＝ +＝☆ +＝－ +＞ +＞λ +＿ +～± +～＋ +［⑤ｆ］ +［⑤ｄ］ +［②ｉ］ +≈ +［②Ｇ］ +［①ｆ］ +ＬＩ +㈧ +［－ +...... +〉 +［③⑩］ +第二 +一番 +一直 +一个 +一些 +许多 +种 +有的是 +也就是说 +末##末 +啊 +阿 +哎 +哎呀 +哎哟 +唉 +俺 +俺们 +按 +按照 +吧 +吧哒 +把 +罢了 +被 +本 +本着 +比 +比方 +比如 +鄙人 +彼 +彼此 +边 +别 +别的 +别说 +并 +并且 +不比 +不成 +不单 +不但 +不独 +不管 +不光 +不过 +不仅 +不拘 +不论 +不怕 +不然 +不如 +不特 +不惟 +不问 +不只 +朝 +朝着 +趁 +趁着 +乘 +冲 +除 +除此之外 +除非 +除了 +此 +此间 +此外 +从 +从而 +打 +待 +但 +但是 +当 +当着 +到 +得 +的 +的话 +等 +等等 +地 +第 +叮咚 +对 +对于 +多 +多少 +而 +而况 +而且 +而是 +而外 +而言 +而已 +尔后 +反过来 +反过来说 +反之 +非但 +非徒 +否则 +嘎 +嘎登 +该 +赶 +个 +各 +各个 +各位 +各种 +各自 +给 +根据 +跟 +故 +故此 +固然 +关于 +管 +归 +果然 +果真 +过 +哈 +哈哈 +呵 +和 +何 +何处 +何况 +何时 +嘿 +哼 +哼唷 +呼哧 +乎 +哗 +还是 +还有 +换句话说 +换言之 +或 +或是 +或者 +极了 +及 +及其 +及至 +即 +即便 +即或 +即令 +即若 +即使 +几 +几时 +己 +既 +既然 +既是 +继而 +加之 +假如 +假若 +假使 +鉴于 +将 +较 +较之 +叫 +接着 +结果 +借 +紧接着 +进而 +尽 +尽管 +经 +经过 +就 +就是 +就是说 +据 +具体地说 +具体说来 +开始 +开外 +靠 +咳 +可 +可见 +可是 +可以 +况且 +啦 +来 +来着 +离 +例如 +哩 +连 +连同 +两者 +了 +临 +另 +另外 +另一方面 +论 +嘛 +吗 +慢说 +漫说 +冒 +么 +每 +每当 +们 +莫若 +某 +某个 +某些 +拿 +哪 +哪边 +哪儿 +哪个 +哪里 +哪年 +哪怕 +哪天 +哪些 +哪样 +那 +那边 +那儿 +那个 +那会儿 +那里 +那么 +那么些 +那么样 +那时 +那些 +那样 +乃 +乃至 +呢 +能 +你 +你们 +您 +宁 +宁可 +宁肯 +宁愿 +哦 +呕 +啪达 +旁人 +呸 +凭 +凭借 +其 +其次 +其二 +其他 +其它 +其一 +其余 +其中 +起 +起见 +起见 +岂但 +恰恰相反 +前后 +前者 +且 +然而 +然后 +然则 +让 +人家 +任 +任何 +任凭 +如 +如此 +如果 +如何 +如其 +如若 +如上所述 +若 +若非 +若是 +啥 +上下 +尚且 +设若 +设使 +甚而 +甚么 +甚至 +省得 +时候 +什么 +什么样 +使得 +是 +是的 +首先 +谁 +谁知 +顺 +顺着 +似的 +虽 +虽然 +虽说 +虽则 +随 +随着 +所 +所以 +他 +他们 +他人 +它 +它们 +她 +她们 +倘 +倘或 +倘然 +倘若 +倘使 +腾 +替 +通过 +同 +同时 +哇 +万一 +往 +望 +为 +为何 +为了 +为什么 +为着 +喂 +嗡嗡 +我 +我们 +呜 +呜呼 +乌乎 +无论 +无宁 +毋宁 +嘻 +吓 +相对而言 +像 +向 +向着 +嘘 +呀 +焉 +沿 +沿着 +要 +要不 +要不然 +要不是 +要么 +要是 +也 +也罢 +也好 +一 +一般 +一旦 +一方面 +一来 +一切 +一样 +一则 +依 +依照 +矣 +以 +以便 +以及 +以免 +以至 +以至于 +以致 +抑或 +因 +因此 +因而 +因为 +哟 +用 +由 +由此可见 +由于 +有 +有的 +有关 +有些 +又 +于 +于是 +于是乎 +与 +与此同时 +与否 +与其 +越是 +云云 +哉 +再说 +再者 +在 +在下 +咱 +咱们 +则 +怎 +怎么 +怎么办 +怎么样 +怎样 +咋 +照 +照着 +者 +这 +这边 +这儿 +这个 +这会儿 +这就是说 +这里 +这么 +这么点儿 +这么些 +这么样 +这时 +这些 +这样 +正如 +吱 +之 +之类 +之所以 +之一 +只是 +只限 +只要 +只有 +至 +至于 +诸位 +着 +着呢 +自 +自从 +自个儿 +自各儿 +自己 +自家 +自身 +综上所述 +总的来看 +总的来说 +总的说来 +总而言之 +总之 +纵 +纵令 +纵然 +纵使 +遵照 +作为 +兮 +呃 +呗 +咚 +咦 +喏 +啐 +喔唷 +嗬 +嗯 +嗳 \ No newline at end of file diff --git a/astrbot/core/knowledge_base/retrieval/manager.py b/astrbot/core/knowledge_base/retrieval/manager.py new file mode 100644 index 00000000..278e4da2 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/manager.py @@ -0,0 +1,273 @@ +"""检索管理器 + +协调稠密检索、稀疏检索和 Rerank,提供统一的检索接口 +""" + +import time + +from dataclasses import dataclass +from typing import List + +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.retrieval.rank_fusion import RankFusion +from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseRetriever +from astrbot.core.provider.provider import RerankProvider +from astrbot.core.db.vec_db.base import Result +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB +from ..kb_helper import KBHelper +from astrbot import logger + + +@dataclass +class RetrievalResult: + """检索结果""" + + chunk_id: str + doc_id: str + doc_name: str + kb_id: str + kb_name: str + content: str + score: float + metadata: dict + + +class RetrievalManager: + """检索管理器 + + 职责: + - 协调稠密检索、稀疏检索和 Rerank + - 结果融合和排序 + """ + + def __init__( + self, + sparse_retriever: SparseRetriever, + rank_fusion: RankFusion, + kb_db: KBSQLiteDatabase, + ): + """初始化检索管理器 + + Args: + vec_db_factory: 向量数据库工厂 + sparse_retriever: 稀疏检索器 + rank_fusion: 结果融合器 + kb_db: 知识库数据库实例 + """ + self.sparse_retriever = sparse_retriever + self.rank_fusion = rank_fusion + self.kb_db = kb_db + + async def retrieve( + self, + query: str, + kb_ids: List[str], + kb_id_helper_map: dict[str, KBHelper], + top_k_fusion: int = 20, + top_m_final: int = 5, + ) -> List[RetrievalResult]: + """混合检索 + + 流程: + 1. 稠密检索 (向量相似度) + 2. 稀疏检索 (BM25) + 3. 结果融合 (RRF) + 4. Rerank 重排序 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + top_m_final: 最终返回数量 + enable_rerank: 是否启用 Rerank + + Returns: + List[RetrievalResult]: 检索结果列表 + """ + if not kb_ids: + return [] + + kb_options: dict = {} + new_kb_ids = [] + for kb_id in kb_ids: + kb_helper = kb_id_helper_map.get(kb_id) + if kb_helper: + kb = kb_helper.kb + kb_options[kb_id] = { + "top_k_dense": kb.top_k_dense or 50, + "top_k_sparse": kb.top_k_sparse or 50, + "top_m_final": kb.top_m_final or 5, + "vec_db": kb_helper.vec_db, + "rerank_provider_id": kb.rerank_provider_id, + } + new_kb_ids.append(kb_id) + else: + logger.warning(f"知识库 ID {kb_id} 实例未找到, 已跳过该知识库的检索") + + kb_ids = new_kb_ids + + # 1. 稠密检索 + time_start = time.time() + dense_results = await self._dense_retrieve( + query=query, + kb_ids=kb_ids, + kb_options=kb_options, + ) + time_end = time.time() + logger.debug( + f"Dense retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(dense_results)} results." + ) + + # 2. 稀疏检索 + time_start = time.time() + sparse_results = await self.sparse_retriever.retrieve( + query=query, + kb_ids=kb_ids, + kb_options=kb_options, + ) + time_end = time.time() + logger.debug( + f"Sparse retrieval across {len(kb_ids)} bases took {time_end - time_start:.2f}s and returned {len(sparse_results)} results." + ) + + # 3. 结果融合 + time_start = time.time() + fused_results = await self.rank_fusion.fuse( + dense_results=dense_results, + sparse_results=sparse_results, + top_k=top_k_fusion, + ) + time_end = time.time() + logger.debug( + f"Rank fusion took {time_end - time_start:.2f}s and returned {len(fused_results)} results." + ) + + # 4. 转换为 RetrievalResult (获取元数据) + retrieval_results = [] + for fr in fused_results: + metadata_dict = await self.kb_db.get_document_with_metadata(fr.doc_id) + if metadata_dict: + retrieval_results.append( + RetrievalResult( + chunk_id=fr.chunk_id, + doc_id=fr.doc_id, + doc_name=metadata_dict["document"].doc_name, + kb_id=fr.kb_id, + kb_name=metadata_dict["knowledge_base"].kb_name, + content=fr.content, + score=fr.score, + metadata={ + "chunk_index": fr.chunk_index, + "char_count": len(fr.content), + }, + ) + ) + + # 5. Rerank + first_rerank = None + for kb_id in kb_ids: + vec_db: FaissVecDB = kb_options[kb_id]["vec_db"] + rerank_pi = kb_options[kb_id]["rerank_provider_id"] + if ( + vec_db + and vec_db.rerank_provider + and rerank_pi + and rerank_pi == vec_db.rerank_provider.meta().id + ): + first_rerank = vec_db.rerank_provider + break + if first_rerank and retrieval_results: + retrieval_results = await self._rerank( + query=query, + results=retrieval_results, + top_k=top_m_final, + rerank_provider=first_rerank, + ) + + return retrieval_results[:top_m_final] + + async def _dense_retrieve( + self, + query: str, + kb_ids: List[str], + kb_options: dict, + ): + """稠密检索 (向量相似度) + + 为每个知识库使用独立的向量数据库进行检索,然后合并结果。 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + top_k: 返回结果数量 + + Returns: + List[Result]: 检索结果列表 + """ + all_results: list[Result] = [] + for kb_id in kb_ids: + if kb_id not in kb_options: + continue + try: + vec_db: FaissVecDB = kb_options[kb_id]["vec_db"] + dense_k = int(kb_options[kb_id]["top_k_dense"]) + vec_results = await vec_db.retrieve( + query=query, + k=dense_k, + fetch_k=dense_k * 2, + rerank=False, # 稠密检索阶段不进行 rerank + metadata_filters={"kb_id": kb_id}, + ) + + all_results.extend(vec_results) + except Exception as e: + from astrbot.core import logger + + logger.warning(f"知识库 {kb_id} 稠密检索失败: {e}") + continue + + # 按相似度排序并返回 top_k + all_results.sort(key=lambda x: x.similarity, reverse=True) + # return all_results[: len(all_results) // len(kb_ids)] + return all_results + + async def _rerank( + self, + query: str, + results: List[RetrievalResult], + top_k: int, + rerank_provider: RerankProvider, + ) -> List[RetrievalResult]: + """Rerank 重排序 + + Args: + query: 查询文本 + results: 检索结果列表 + top_k: 返回结果数量 + + Returns: + List[RetrievalResult]: 重排序后的结果列表 + """ + if not results: + return [] + + # 准备文档列表 + docs = [r.content for r in results] + + # 调用 Rerank Provider + rerank_results = await rerank_provider.rerank( + query=query, + documents=docs, + ) + + # 更新分数并重新排序 + reranked_list = [] + for rerank_result in rerank_results: + idx = rerank_result.index + if idx < len(results): + result = results[idx] + result.score = rerank_result.relevance_score + reranked_list.append(result) + + reranked_list.sort(key=lambda x: x.score, reverse=True) + + return reranked_list[:top_k] diff --git a/astrbot/core/knowledge_base/retrieval/rank_fusion.py b/astrbot/core/knowledge_base/retrieval/rank_fusion.py new file mode 100644 index 00000000..3ceba4ff --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/rank_fusion.py @@ -0,0 +1,138 @@ +"""检索结果融合器 + +使用 Reciprocal Rank Fusion (RRF) 算法融合稠密检索和稀疏检索的结果 +""" + +import json +from dataclasses import dataclass + +from astrbot.core.db.vec_db.base import Result +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.knowledge_base.retrieval.sparse_retriever import SparseResult + + +@dataclass +class FusedResult: + """融合后的检索结果""" + + chunk_id: str + chunk_index: int + doc_id: str + kb_id: str + content: str + score: float + + +class RankFusion: + """检索结果融合器 + + 职责: + - 融合稠密检索和稀疏检索的结果 + - 使用 Reciprocal Rank Fusion (RRF) 算法 + """ + + def __init__(self, kb_db: KBSQLiteDatabase, k: int = 60): + """初始化结果融合器 + + Args: + kb_db: 知识库数据库实例 + k: RRF 参数,用于平滑排名 + """ + self.kb_db = kb_db + self.k = k + + async def fuse( + self, + dense_results: list[Result], + sparse_results: list[SparseResult], + top_k: int = 20, + ) -> list[FusedResult]: + """融合稠密和稀疏检索结果 + + RRF 公式: + score(doc) = sum(1 / (k + rank_i)) + + Args: + dense_results: 稠密检索结果 + sparse_results: 稀疏检索结果 + top_k: 返回结果数量 + + Returns: + List[FusedResult]: 融合后的结果列表 + """ + # 1. 构建排名映射 + dense_ranks = { + r.data["doc_id"]: (idx + 1) for idx, r in enumerate(dense_results) + } # 这里的 doc_id 实际上是 chunk_id + sparse_ranks = {r.chunk_id: (idx + 1) for idx, r in enumerate(sparse_results)} + + # 2. 收集所有唯一的 ID + # 需要统一为 chunk_id + all_chunk_ids = set() + vec_doc_id_to_dense: dict[str, Result] = {} # vec_doc_id -> Result + chunk_id_to_sparse: dict[str, SparseResult] = {} # chunk_id -> SparseResult + + # 处理稀疏检索结果 + for r in sparse_results: + all_chunk_ids.add(r.chunk_id) + chunk_id_to_sparse[r.chunk_id] = r + + # 处理稠密检索结果 (需要转换 vec_doc_id 到 chunk_id) + for r in dense_results: + vec_doc_id = r.data["doc_id"] + all_chunk_ids.add(vec_doc_id) + vec_doc_id_to_dense[vec_doc_id] = r + + # 3. 计算 RRF 分数 + rrf_scores: dict[str, float] = {} + + for identifier in all_chunk_ids: + score = 0.0 + + # 来自稠密检索的贡献 + if identifier in dense_ranks: + score += 1.0 / (self.k + dense_ranks[identifier]) + + # 来自稀疏检索的贡献 + if identifier in sparse_ranks: + score += 1.0 / (self.k + sparse_ranks[identifier]) + + rrf_scores[identifier] = score + + # 4. 排序 + sorted_ids = sorted( + rrf_scores.keys(), key=lambda cid: rrf_scores[cid], reverse=True + )[:top_k] + + # 5. 构建融合结果 + fused_results = [] + for identifier in sorted_ids: + # 优先从稀疏检索获取完整信息 + if identifier in chunk_id_to_sparse: + sr = chunk_id_to_sparse[identifier] + fused_results.append( + FusedResult( + chunk_id=sr.chunk_id, + chunk_index=sr.chunk_index, + doc_id=sr.doc_id, + kb_id=sr.kb_id, + content=sr.content, + score=rrf_scores[identifier], + ) + ) + elif identifier in vec_doc_id_to_dense: + # 从向量检索获取信息,需要从数据库获取块的详细信息 + vec_result = vec_doc_id_to_dense[identifier] + chunk_md = json.loads(vec_result.data["metadata"]) + fused_results.append( + FusedResult( + chunk_id=identifier, + chunk_index=chunk_md["chunk_index"], + doc_id=chunk_md["kb_doc_id"], + kb_id=chunk_md["kb_id"], + content=vec_result.data["text"], + score=rrf_scores[identifier], + ) + ) + + return fused_results diff --git a/astrbot/core/knowledge_base/retrieval/sparse_retriever.py b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py new file mode 100644 index 00000000..315930b3 --- /dev/null +++ b/astrbot/core/knowledge_base/retrieval/sparse_retriever.py @@ -0,0 +1,130 @@ +"""稀疏检索器 + +使用 BM25 算法进行基于关键词的文档检索 +""" + +import jieba +import os +import json +from dataclasses import dataclass +from rank_bm25 import BM25Okapi +from astrbot.core.knowledge_base.kb_db_sqlite import KBSQLiteDatabase +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB + + +@dataclass +class SparseResult: + """稀疏检索结果""" + + chunk_index: int + chunk_id: str + doc_id: str + kb_id: str + content: str + score: float + + +class SparseRetriever: + """BM25 稀疏检索器 + + 职责: + - 基于关键词的文档检索 + - 使用 BM25 算法计算相关度 + """ + + def __init__(self, kb_db: KBSQLiteDatabase): + """初始化稀疏检索器 + + Args: + kb_db: 知识库数据库实例 + """ + self.kb_db = kb_db + self._index_cache = {} # 缓存 BM25 索引 + + with open( + os.path.join(os.path.dirname(__file__), "hit_stopwords.txt"), + encoding="utf-8", + ) as f: + self.hit_stopwords = { + word.strip() for word in set(f.read().splitlines()) if word.strip() + } + + async def retrieve( + self, + query: str, + kb_ids: list[str], + kb_options: dict, + ) -> list[SparseResult]: + """执行稀疏检索 + + Args: + query: 查询文本 + kb_ids: 知识库 ID 列表 + kb_options: 每个知识库的检索选项 + + Returns: + List[SparseResult]: 检索结果列表 + """ + # 1. 获取所有相关块 + top_k_sparse = 0 + chunks = [] + for kb_id in kb_ids: + vec_db: FaissVecDB = kb_options.get(kb_id, {}).get("vec_db") + if not vec_db: + continue + result = await vec_db.document_storage.get_documents( + metadata_filters={}, limit=None, offset=None + ) + chunk_mds = [json.loads(doc["metadata"]) for doc in result] + result = [ + { + "chunk_id": doc["doc_id"], + "chunk_index": chunk_md["chunk_index"], + "doc_id": chunk_md["kb_doc_id"], + "kb_id": kb_id, + "text": doc["text"], + } + for doc, chunk_md in zip(result, chunk_mds) + ] + chunks.extend(result) + top_k_sparse += kb_options.get(kb_id, {}).get("top_k_sparse", 50) + + if not chunks: + return [] + + # 2. 准备文档和索引 + corpus = [chunk["text"] for chunk in chunks] + tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus] + tokenized_corpus = [ + [word for word in doc if word not in self.hit_stopwords] + for doc in tokenized_corpus + ] + + # 3. 构建 BM25 索引 + bm25 = BM25Okapi(tokenized_corpus) + + # 4. 执行检索 + tokenized_query = list(jieba.cut(query)) + tokenized_query = [ + word for word in tokenized_query if word not in self.hit_stopwords + ] + scores = bm25.get_scores(tokenized_query) + + # 5. 排序并返回 Top-K + results = [] + for idx, score in enumerate(scores): + chunk = chunks[idx] + results.append( + SparseResult( + chunk_id=chunk["chunk_id"], + chunk_index=chunk["chunk_index"], + doc_id=chunk["doc_id"], + kb_id=chunk["kb_id"], + content=chunk["text"], + score=float(score), + ) + ) + + results.sort(key=lambda x: x.score, reverse=True) + # return results[: len(results) // len(kb_ids)] + return results[:top_k_sparse] diff --git a/astrbot/core/pipeline/process_stage/method/llm_request.py b/astrbot/core/pipeline/process_stage/method/llm_request.py index 514a87a9..703b3681 100644 --- a/astrbot/core/pipeline/process_stage/method/llm_request.py +++ b/astrbot/core/pipeline/process_stage/method/llm_request.py @@ -7,7 +7,7 @@ import copy import json import traceback from datetime import timedelta -from typing import AsyncGenerator, Union +from collections.abc import AsyncGenerator from astrbot.core.conversation_mgr import Conversation from astrbot.core import logger from astrbot.core.message.components import Image @@ -33,6 +33,7 @@ from astrbot.core.star.star_handler import EventType from astrbot.core.utils.metrics import Metric from ...context import PipelineContext, call_event_hook, call_handler from ..stage import Stage +from ..utils import inject_kb_context from astrbot.core.provider.register import llm_tools from astrbot.core.star.star_handler import star_map from astrbot.core.astr_agent_context import AstrAgentContext @@ -367,7 +368,7 @@ class LLMRequestSubStage(Stage): async def process( self, event: AstrMessageEvent, _nested: bool = False - ) -> Union[None, AsyncGenerator[None, None]]: + ) -> None | AsyncGenerator[None, None]: req: ProviderRequest | None = None if not self.ctx.astrbot_config["provider_settings"]["enable"]: @@ -419,6 +420,14 @@ class LLMRequestSubStage(Stage): if not req.prompt and not req.image_urls: return + # 应用知识库 + try: + await inject_kb_context( + umo=event.unified_msg_origin, p_ctx=self.ctx, req=req + ) + except Exception as e: + logger.error(f"调用知识库时遇到问题: {e}") + # 执行请求 LLM 前事件钩子。 if await call_event_hook(event, EventType.OnLLMRequestEvent, req): return diff --git a/astrbot/core/pipeline/process_stage/utils.py b/astrbot/core/pipeline/process_stage/utils.py new file mode 100644 index 00000000..e799ad4d --- /dev/null +++ b/astrbot/core/pipeline/process_stage/utils.py @@ -0,0 +1,80 @@ +from ..context import PipelineContext +from astrbot.core.provider.entities import ProviderRequest +from astrbot.api import logger, sp + + +async def inject_kb_context( + umo: str, + p_ctx: PipelineContext, + req: ProviderRequest, +) -> None: + """inject knowledge base context into the provider request + + Args: + umo: Unique message object (session ID) + p_ctx: Pipeline context + req: Provider request + """ + + kb_mgr = p_ctx.plugin_manager.context.kb_manager + + # 1. 优先读取会话级配置 + session_config = await sp.session_get(umo, "kb_config", default={}) + + if session_config and "kb_ids" in session_config: + # 会话级配置 + kb_ids = session_config.get("kb_ids", []) + + # 如果配置为空列表，明确表示不使用知识库 + if not kb_ids: + logger.info(f"[知识库] 会话 {umo} 已被配置为不使用知识库") + return + + top_k = session_config.get("top_k", 5) + + # 将 kb_ids 转换为 kb_names + kb_names = [] + invalid_kb_ids = [] + for kb_id in kb_ids: + kb_helper = await kb_mgr.get_kb(kb_id) + if kb_helper: + kb_names.append(kb_helper.kb.kb_name) + else: + logger.warning(f"[知识库] 知识库不存在或未加载: {kb_id}") + invalid_kb_ids.append(kb_id) + + if invalid_kb_ids: + logger.warning( + f"[知识库] 会话 {umo} 配置的以下知识库无效: {invalid_kb_ids}" + ) + + if not kb_names: + return + + logger.debug(f"[知识库] 使用会话级配置，知识库数量: {len(kb_names)}") + else: + kb_names = p_ctx.astrbot_config.get("kb_names", []) + top_k = p_ctx.astrbot_config.get("kb_final_top_k", 5) + logger.debug(f"[知识库] 使用全局配置，知识库数量: {len(kb_names)}") + + top_k_fusion = p_ctx.astrbot_config.get("kb_fusion_top_k", 20) + + if not kb_names: + return + + logger.debug(f"[知识库] 开始检索知识库，数量: {len(kb_names)}, top_k={top_k}") + kb_context = await kb_mgr.retrieve( + query=req.prompt, + kb_names=kb_names, + top_k_fusion=top_k_fusion, + top_m_final=top_k, + ) + + if not kb_context: + return + + formatted = kb_context.get("context_text", "") + if formatted: + results = kb_context.get("results", []) + logger.debug(f"[知识库] 为会话 {umo} 注入了 {len(results)} 条相关知识块") + req.system_prompt = f"{formatted}\n\n{req.system_prompt or ''}" diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index 6666b33e..edb9b767 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -1,6 +1,5 @@ import asyncio import traceback -from typing import List from astrbot.core import logger, sp from astrbot.core.astrbot_config_mgr import AstrBotConfigManager @@ -28,7 +27,7 @@ class ProviderManager: self.persona_mgr = persona_mgr self.acm = acm config = acm.confs["default"] - self.providers_config: List = config["provider"] + self.providers_config: list = config["provider"] self.provider_settings: dict = config["provider_settings"] self.provider_stt_settings: dict = config.get("provider_stt_settings", {}) self.provider_tts_settings: dict = config.get("provider_tts_settings", {}) @@ -36,15 +35,15 @@ class ProviderManager: # 人格相关属性，v4.0.0 版本后被废弃，推荐使用 PersonaManager self.default_persona_name = persona_mgr.default_persona - self.provider_insts: List[Provider] = [] + self.provider_insts: list[Provider] = [] """加载的 Provider 的实例""" - self.stt_provider_insts: List[STTProvider] = [] + self.stt_provider_insts: list[STTProvider] = [] """加载的 Speech To Text Provider 的实例""" - self.tts_provider_insts: List[TTSProvider] = [] + self.tts_provider_insts: list[TTSProvider] = [] """加载的 Text To Speech Provider 的实例""" - self.embedding_provider_insts: List[EmbeddingProvider] = [] + self.embedding_provider_insts: list[EmbeddingProvider] = [] """加载的 Embedding Provider 的实例""" - self.rerank_provider_insts: List[RerankProvider] = [] + self.rerank_provider_insts: list[RerankProvider] = [] """加载的 Rerank Provider 的实例""" self.inst_map: dict[ str, @@ -175,7 +174,11 @@ class ProviderManager: async def initialize(self): # 逐个初始化提供商 for provider_config in self.providers_config: - await self.load_provider(provider_config) + try: + await self.load_provider(provider_config) + except Exception as e: + logger.error(traceback.format_exc()) + logger.error(e) # 设置默认提供商 selected_provider_id = sp.get( @@ -404,10 +407,12 @@ class ProviderManager: self.inst_map[provider_config["id"]] = inst except Exception as e: - logger.error(traceback.format_exc()) logger.error( f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}" ) + raise Exception( + f"实例化 {provider_config['type']}({provider_config['id']}) 提供商适配器失败：{e}" + ) async def reload(self, provider_config: dict): await self.terminate_provider(provider_config["id"]) diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 901036b1..9953e9f1 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -1,4 +1,5 @@ import abc +import asyncio from typing import List from typing import AsyncGenerator from astrbot.core.agent.tool import ToolSet @@ -203,6 +204,72 @@ class EmbeddingProvider(AbstractProvider): """获取向量的维度""" ... + async def get_embeddings_batch( + self, + texts: list[str], + batch_size: int = 16, + tasks_limit: int = 3, + max_retries: int = 3, + progress_callback=None, + ) -> list[list[float]]: + """批量获取文本的向量，分批处理以节省内存 + + Args: + texts: 文本列表 + batch_size: 每批处理的文本数量 + tasks_limit: 并发任务数量限制 + max_retries: 失败时的最大重试次数 + progress_callback: 进度回调函数，接收参数 (current, total) + + Returns: + 向量列表 + """ + semaphore = asyncio.Semaphore(tasks_limit) + all_embeddings: list[list[float]] = [] + failed_batches: list[tuple[int, list[str]]] = [] + completed_count = 0 + total_count = len(texts) + + async def process_batch(batch_idx: int, batch_texts: list[str]): + nonlocal completed_count + async with semaphore: + for attempt in range(max_retries): + try: + batch_embeddings = await self.get_embeddings(batch_texts) + all_embeddings.extend(batch_embeddings) + completed_count += len(batch_texts) + if progress_callback: + await progress_callback(completed_count, total_count) + return + except Exception as e: + if attempt == max_retries - 1: + # 最后一次重试失败，记录失败的批次 + failed_batches.append((batch_idx, batch_texts)) + raise Exception( + f"批次 {batch_idx} 处理失败，已重试 {max_retries} 次: {str(e)}" + ) + # 等待一段时间后重试，使用指数退避 + await asyncio.sleep(2**attempt) + + tasks = [] + for i in range(0, len(texts), batch_size): + batch_texts = texts[i : i + batch_size] + batch_idx = i // batch_size + tasks.append(process_batch(batch_idx, batch_texts)) + + # 收集所有任务的结果，包括失败的任务 + results = await asyncio.gather(*tasks, return_exceptions=True) + + # 检查是否有失败的任务 + errors = [r for r in results if isinstance(r, Exception)] + if errors: + error_msg = ( + f"有 {len(errors)} 个批次处理失败: {'; '.join(str(e) for e in errors)}" + ) + raise Exception(error_msg) + + return all_embeddings + class RerankProvider(AbstractProvider): def __init__(self, provider_config: dict, provider_settings: dict) -> None: diff --git a/astrbot/core/provider/sources/gemini_embedding_source.py b/astrbot/core/provider/sources/gemini_embedding_source.py index baccf52a..562d1135 100644 --- a/astrbot/core/provider/sources/gemini_embedding_source.py +++ b/astrbot/core/provider/sources/gemini_embedding_source.py @@ -32,7 +32,6 @@ class GeminiEmbeddingProvider(EmbeddingProvider): self.model = provider_config.get( "embedding_model", "gemini-embedding-exp-03-07" ) - self.dimension = provider_config.get("embedding_dimensions", 768) async def get_embedding(self, text: str) -> list[float]: """ @@ -60,4 +59,4 @@ class GeminiEmbeddingProvider(EmbeddingProvider): def get_dim(self) -> int: """获取向量的维度""" - return self.dimension + return self.provider_config.get("embedding_dimensions", 768) diff --git a/astrbot/core/provider/sources/openai_embedding_source.py b/astrbot/core/provider/sources/openai_embedding_source.py index 79b2e83b..e6f692a3 100644 --- a/astrbot/core/provider/sources/openai_embedding_source.py +++ b/astrbot/core/provider/sources/openai_embedding_source.py @@ -22,7 +22,6 @@ class OpenAIEmbeddingProvider(EmbeddingProvider): timeout=int(provider_config.get("timeout", 20)), ) self.model = provider_config.get("embedding_model", "text-embedding-3-small") - self.dimension = provider_config.get("embedding_dimensions", 1024) async def get_embedding(self, text: str) -> list[float]: """ @@ -40,4 +39,4 @@ class OpenAIEmbeddingProvider(EmbeddingProvider): def get_dim(self) -> int: """获取向量的维度""" - return self.dimension + return self.provider_config.get("embedding_dimensions", 1024) diff --git a/astrbot/core/star/context.py b/astrbot/core/star/context.py index 31616e7d..0229f4db 100644 --- a/astrbot/core/star/context.py +++ b/astrbot/core/star/context.py @@ -19,6 +19,7 @@ from astrbot.core.platform import Platform from astrbot.core.platform.manager import PlatformManager from astrbot.core.platform_message_history_mgr import PlatformMessageHistoryManager from astrbot.core.astrbot_config_mgr import AstrBotConfigManager +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager from astrbot.core.persona_mgr import PersonaManager from .star import star_registry, StarMetadata, star_map from .star_handler import star_handlers_registry, StarHandlerMetadata, EventType @@ -55,6 +56,7 @@ class Context: message_history_manager: PlatformMessageHistoryManager, persona_manager: PersonaManager, astrbot_config_mgr: AstrBotConfigManager, + knowledge_base_manager: KnowledgeBaseManager, ): self._event_queue = event_queue """事件队列。消息平台通过事件队列传递消息事件。""" @@ -68,6 +70,7 @@ class Context: self.message_history_manager = message_history_manager self.persona_manager = persona_manager self.astrbot_config_mgr = astrbot_config_mgr + self.kb_manager = knowledge_base_manager def get_registered_star(self, star_name: str) -> StarMetadata | None: """根据插件名获取插件的 Metadata""" diff --git a/astrbot/dashboard/routes/__init__.py b/astrbot/dashboard/routes/__init__.py index ef2fa3e8..e1d58f62 100644 --- a/astrbot/dashboard/routes/__init__.py +++ b/astrbot/dashboard/routes/__init__.py @@ -11,6 +11,7 @@ from .conversation import ConversationRoute from .file import FileRoute from .session_management import SessionManagementRoute from .persona import PersonaRoute +from .knowledge_base import KnowledgeBaseRoute __all__ = [ "AuthRoute", @@ -26,4 +27,5 @@ __all__ = [ "FileRoute", "SessionManagementRoute", "PersonaRoute", + "KnowledgeBaseRoute", ] diff --git a/astrbot/dashboard/routes/config.py b/astrbot/dashboard/routes/config.py index 01b9b243..998240c9 100644 --- a/astrbot/dashboard/routes/config.py +++ b/astrbot/dashboard/routes/config.py @@ -1,4 +1,3 @@ -import typing import traceback import os import inspect @@ -45,9 +44,7 @@ def try_cast(value: str, type_: str): return None -def validate_config( - data, schema: dict, is_core: bool -) -> typing.Tuple[typing.List[str], typing.Dict]: +def validate_config(data, schema: dict, is_core: bool) -> tuple[list[str], dict]: errors = [] def validate(data: dict, metadata: dict = schema, path=""): @@ -178,6 +175,7 @@ class ConfigRoute(Route): "/config/provider/check_one": ("GET", self.check_one_provider_status), "/config/provider/list": ("GET", self.get_provider_config_list), "/config/provider/model_list": ("GET", self.get_provider_model_list), + "/config/provider/get_embedding_dim": ("POST", self.get_embedding_dim), } self.register_routes() @@ -601,6 +599,61 @@ class ConfigRoute(Route): logger.error(traceback.format_exc()) return Response().error(str(e)).__dict__ + async def get_embedding_dim(self): + """获取嵌入模型的维度""" + post_data = await request.json + provider_config = post_data.get("provider_config", None) + if not provider_config: + return Response().error("缺少参数 provider_config").__dict__ + + try: + # 动态导入 EmbeddingProvider + from astrbot.core.provider.provider import EmbeddingProvider + from astrbot.core.provider.register import provider_cls_map + + # 获取 provider 类型 + provider_type = provider_config.get("type", None) + if not provider_type: + return Response().error("provider_config 缺少 type 字段").__dict__ + + # 获取对应的 provider 类 + if provider_type not in provider_cls_map: + return ( + Response() + .error(f"未找到适用于 {provider_type} 的提供商适配器") + .__dict__ + ) + + provider_metadata = provider_cls_map[provider_type] + cls_type = provider_metadata.cls_type + + if not cls_type: + return Response().error(f"无法找到 {provider_type} 的类").__dict__ + + # 实例化 provider + inst = cls_type(provider_config, {}) + + # 检查是否是 EmbeddingProvider + if not isinstance(inst, EmbeddingProvider): + return Response().error("提供商不是 EmbeddingProvider 类型").__dict__ + + # 初始化 + if getattr(inst, "initialize", None): + await inst.initialize() + + # 获取嵌入向量维度 + vec = await inst.get_embedding("echo") + dim = len(vec) + + logger.info( + f"检测到 {provider_config.get('id', 'unknown')} 的嵌入向量维度为 {dim}" + ) + + return Response().ok({"embedding_dimensions": dim}).__dict__ + except Exception as e: + logger.error(traceback.format_exc()) + return Response().error(f"获取嵌入维度失败: {str(e)}").__dict__ + async def get_platform_list(self): """获取所有平台的列表""" platform_list = [] @@ -797,7 +850,7 @@ class ConfigRoute(Route): logger.warning( f"Failed to import required modules for platform {platform.name}: {e}" ) - except (OSError, IOError) as e: + except OSError as e: logger.warning(f"File system error for platform {platform.name} logo: {e}") except Exception as e: logger.warning( diff --git a/astrbot/dashboard/routes/knowledge_base.py b/astrbot/dashboard/routes/knowledge_base.py new file mode 100644 index 00000000..d8d0434d --- /dev/null +++ b/astrbot/dashboard/routes/knowledge_base.py @@ -0,0 +1,1065 @@ +"""知识库管理 API 路由""" + +import uuid +import aiofiles +import os +import traceback +import asyncio +from quart import request +from astrbot.core import logger +from astrbot.core.core_lifecycle import AstrBotCoreLifecycle +from .route import Route, Response, RouteContext +from ..utils import generate_tsne_visualization +from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider + + +class KnowledgeBaseRoute(Route): + """知识库管理路由 + + 提供知识库、文档、检索、会话配置等 API 接口 + """ + + def __init__( + self, + context: RouteContext, + core_lifecycle: AstrBotCoreLifecycle, + ) -> None: + super().__init__(context) + self.core_lifecycle = core_lifecycle + self.kb_manager = None # 延迟初始化 + self.kb_db = None + self.session_config_db = None # 会话配置数据库 + self.retrieval_manager = None + self.upload_progress = {} # 存储上传进度 {task_id: {status, file_index, file_total, stage, current, total}} + self.upload_tasks = {} # 存储后台上传任务 {task_id: {"status", "result", "error"}} + + # 注册路由 + self.routes = { + # 知识库管理 + "/kb/list": ("GET", self.list_kbs), + "/kb/create": ("POST", self.create_kb), + "/kb/get": ("GET", self.get_kb), + "/kb/update": ("POST", self.update_kb), + "/kb/delete": ("POST", self.delete_kb), + "/kb/stats": ("GET", self.get_kb_stats), + # 文档管理 + "/kb/document/list": ("GET", self.list_documents), + "/kb/document/upload": ("POST", self.upload_document), + "/kb/document/upload/progress": ("GET", self.get_upload_progress), + "/kb/document/get": ("GET", self.get_document), + "/kb/document/delete": ("POST", self.delete_document), + # # 块管理 + "/kb/chunk/list": ("GET", self.list_chunks), + "/kb/chunk/delete": ("POST", self.delete_chunk), + # # 多媒体管理 + # "/kb/media/list": ("GET", self.list_media), + # "/kb/media/delete": ("POST", self.delete_media), + # 检索 + "/kb/retrieve": ("POST", self.retrieve), + # 会话知识库配置 + "/kb/session/config/get": ("GET", self.get_session_kb_config), + "/kb/session/config/set": ("POST", self.set_session_kb_config), + "/kb/session/config/delete": ("POST", self.delete_session_kb_config), + } + self.register_routes() + + def _get_kb_manager(self): + return self.core_lifecycle.kb_manager + + async def _background_upload_task( + self, + task_id: str, + kb_helper, + files_to_upload: list, + chunk_size: int, + chunk_overlap: int, + batch_size: int, + tasks_limit: int, + max_retries: int, + ): + """后台上传任务""" + try: + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "processing", + "result": None, + "error": None, + } + self.upload_progress[task_id] = { + "status": "processing", + "file_index": 0, + "file_total": len(files_to_upload), + "stage": "waiting", + "current": 0, + "total": 100, + } + + uploaded_docs = [] + failed_docs = [] + + for file_idx, file_info in enumerate(files_to_upload): + try: + # 更新整体进度 + self.upload_progress[task_id].update( + { + "status": "processing", + "file_index": file_idx, + "file_name": file_info["file_name"], + "stage": "parsing", + "current": 0, + "total": 100, + } + ) + + # 创建进度回调函数 + async def progress_callback(stage, current, total): + if task_id in self.upload_progress: + self.upload_progress[task_id].update( + { + "status": "processing", + "file_index": file_idx, + "file_name": file_info["file_name"], + "stage": stage, + "current": current, + "total": total, + } + ) + + doc = await kb_helper.upload_document( + file_name=file_info["file_name"], + file_content=file_info["file_content"], + file_type=file_info["file_type"], + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + progress_callback=progress_callback, + ) + + uploaded_docs.append(doc.model_dump()) + except Exception as e: + logger.error(f"上传文档 {file_info['file_name']} 失败: {e}") + failed_docs.append( + {"file_name": file_info["file_name"], "error": str(e)} + ) + + # 更新任务完成状态 + result = { + "task_id": task_id, + "uploaded": uploaded_docs, + "failed": failed_docs, + "total": len(files_to_upload), + "success_count": len(uploaded_docs), + "failed_count": len(failed_docs), + } + + self.upload_tasks[task_id] = { + "status": "completed", + "result": result, + "error": None, + } + self.upload_progress[task_id]["status"] = "completed" + + except Exception as e: + logger.error(f"后台上传任务 {task_id} 失败: {e}") + logger.error(traceback.format_exc()) + self.upload_tasks[task_id] = { + "status": "failed", + "result": None, + "error": str(e), + } + if task_id in self.upload_progress: + self.upload_progress[task_id]["status"] = "failed" + + async def list_kbs(self): + """获取知识库列表 + + Query 参数: + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + - refresh_stats: 是否刷新统计信息 (默认 false，首次加载时可设为 true) + """ + try: + kb_manager = self._get_kb_manager() + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 20, type=int) + + kbs = await kb_manager.list_kbs() + + # 转换为字典列表 + kb_list = [] + for kb in kbs: + kb_list.append(kb.model_dump()) + + return ( + Response() + .ok({"items": kb_list, "page": page, "page_size": page_size}) + .__dict__ + ) + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库列表失败: {str(e)}").__dict__ + + async def create_kb(self): + """创建知识库 + + Body: + - kb_name: 知识库名称 (必填) + - description: 描述 (可选) + - emoji: 图标 (可选) + - embedding_provider_id: 嵌入模型提供商ID (可选) + - rerank_provider_id: 重排序模型提供商ID (可选) + - chunk_size: 分块大小 (可选, 默认512) + - chunk_overlap: 块重叠大小 (可选, 默认50) + - top_k_dense: 密集检索数量 (可选, 默认50) + - top_k_sparse: 稀疏检索数量 (可选, 默认50) + - top_m_final: 最终返回数量 (可选, 默认5) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + kb_name = data.get("kb_name") + if not kb_name: + return Response().error("知识库名称不能为空").__dict__ + + description = data.get("description") + emoji = data.get("emoji") + embedding_provider_id = data.get("embedding_provider_id") + rerank_provider_id = data.get("rerank_provider_id") + chunk_size = data.get("chunk_size") + chunk_overlap = data.get("chunk_overlap") + top_k_dense = data.get("top_k_dense") + top_k_sparse = data.get("top_k_sparse") + top_m_final = data.get("top_m_final") + + # pre-check embedding dim + if not embedding_provider_id: + return Response().error("缺少参数 embedding_provider_id").__dict__ + prv = await kb_manager.provider_manager.get_provider_by_id( + embedding_provider_id + ) # type: ignore + if not prv or not isinstance(prv, EmbeddingProvider): + return ( + Response().error(f"嵌入模型不存在或类型错误({type(prv)})").__dict__ + ) + try: + vec = await prv.get_embedding("astrbot") + if len(vec) != prv.get_dim(): + raise ValueError( + f"嵌入向量维度不匹配，实际是 {len(vec)}，然而配置是 {prv.get_dim()}" + ) + except Exception as e: + return Response().error(f"测试嵌入模型失败: {str(e)}").__dict__ + # pre-check rerank + if rerank_provider_id: + rerank_prv: RerankProvider = ( + await kb_manager.provider_manager.get_provider_by_id( + rerank_provider_id + ) + ) # type: ignore + if not rerank_prv: + return Response().error("重排序模型不存在").__dict__ + # 检查重排序模型可用性 + try: + res = await rerank_prv.rerank( + query="astrbot", documents=["astrbot knowledge base"] + ) + if not res: + raise ValueError("重排序模型返回结果异常") + except Exception as e: + return ( + Response() + .error(f"测试重排序模型失败: {str(e)}，请检查控制台日志输出。") + .__dict__ + ) + + kb_helper = await kb_manager.create_kb( + kb_name=kb_name, + description=description, + emoji=emoji, + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + top_k_dense=top_k_dense, + top_k_sparse=top_k_sparse, + top_m_final=top_m_final, + ) + kb = kb_helper.kb + + return Response().ok(kb.model_dump(), "创建知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"创建知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"创建知识库失败: {str(e)}").__dict__ + + async def get_kb(self): + """获取知识库详情 + + Query 参数: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + kb = kb_helper.kb + + return Response().ok(kb.model_dump()).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库详情失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库详情失败: {str(e)}").__dict__ + + async def update_kb(self): + """更新知识库 + + Body: + - kb_id: 知识库 ID (必填) + - kb_name: 新的知识库名称 (可选) + - description: 新的描述 (可选) + - emoji: 新的图标 (可选) + - embedding_provider_id: 新的嵌入模型提供商ID (可选) + - rerank_provider_id: 新的重排序模型提供商ID (可选) + - chunk_size: 分块大小 (可选) + - chunk_overlap: 块重叠大小 (可选) + - top_k_dense: 密集检索数量 (可选) + - top_k_sparse: 稀疏检索数量 (可选) + - top_m_final: 最终返回数量 (可选) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_name = data.get("kb_name") + description = data.get("description") + emoji = data.get("emoji") + embedding_provider_id = data.get("embedding_provider_id") + rerank_provider_id = data.get("rerank_provider_id") + chunk_size = data.get("chunk_size") + chunk_overlap = data.get("chunk_overlap") + top_k_dense = data.get("top_k_dense") + top_k_sparse = data.get("top_k_sparse") + top_m_final = data.get("top_m_final") + + # 检查是否至少提供了一个更新字段 + if all( + v is None + for v in [ + kb_name, + description, + emoji, + embedding_provider_id, + rerank_provider_id, + chunk_size, + chunk_overlap, + top_k_dense, + top_k_sparse, + top_m_final, + ] + ): + return Response().error("至少需要提供一个更新字段").__dict__ + + kb_helper = await kb_manager.update_kb( + kb_id=kb_id, + kb_name=kb_name, + description=description, + emoji=emoji, + embedding_provider_id=embedding_provider_id, + rerank_provider_id=rerank_provider_id, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + top_k_dense=top_k_dense, + top_k_sparse=top_k_sparse, + top_m_final=top_m_final, + ) + + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + kb = kb_helper.kb + return Response().ok(kb.model_dump(), "更新知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"更新知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"更新知识库失败: {str(e)}").__dict__ + + async def delete_kb(self): + """删除知识库 + + Body: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + success = await kb_manager.delete_kb(kb_id) + if not success: + return Response().error("知识库不存在").__dict__ + + return Response().ok(message="删除知识库成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除知识库失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除知识库失败: {str(e)}").__dict__ + + async def get_kb_stats(self): + """获取知识库统计信息 + + Query 参数: + - kb_id: 知识库 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + kb = kb_helper.kb + + stats = { + "kb_id": kb.kb_id, + "kb_name": kb.kb_name, + "doc_count": kb.doc_count, + "chunk_count": kb.chunk_count, + "created_at": kb.created_at.isoformat(), + "updated_at": kb.updated_at.isoformat(), + } + + return Response().ok(stats).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取知识库统计失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取知识库统计失败: {str(e)}").__dict__ + + # ===== 文档管理 API ===== + + async def list_documents(self): + """获取文档列表 + + Query 参数: + - kb_id: 知识库 ID (必填) + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 100, type=int) + + offset = (page - 1) * page_size + limit = page_size + + doc_list = await kb_helper.list_documents(offset=offset, limit=limit) + + doc_list = [doc.model_dump() for doc in doc_list] + + return ( + Response() + .ok({"items": doc_list, "page": page, "page_size": page_size}) + .__dict__ + ) + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取文档列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取文档列表失败: {str(e)}").__dict__ + + async def upload_document(self): + """上传文档 + + 支持两种方式: + 1. multipart/form-data 文件上传（支持多文件，最多10个） + 2. JSON 格式 base64 编码上传（支持多文件，最多10个） + + Form Data (multipart/form-data): + - kb_id: 知识库 ID (必填) + - file: 文件对象 (必填，可多个，字段名为 file, file1, file2, ... 或 files[]) + + JSON Body (application/json): + - kb_id: 知识库 ID (必填) + - files: 文件数组 (必填) + - file_name: 文件名 (必填) + - file_content: base64 编码的文件内容 (必填) + + 返回: + - task_id: 任务ID，用于查询上传进度和结果 + """ + try: + kb_manager = self._get_kb_manager() + + # 检查 Content-Type + content_type = request.content_type + kb_id = None + chunk_size = None + chunk_overlap = None + batch_size = 32 + tasks_limit = 3 + max_retries = 3 + files_to_upload = [] # 存储待上传的文件信息列表 + + if content_type and "multipart/form-data" not in content_type: + return ( + Response().error("Content-Type 须为 multipart/form-data").__dict__ + ) + form_data = await request.form + files = await request.files + + kb_id = form_data.get("kb_id") + chunk_size = int(form_data.get("chunk_size", 512)) + chunk_overlap = int(form_data.get("chunk_overlap", 50)) + batch_size = int(form_data.get("batch_size", 32)) + tasks_limit = int(form_data.get("tasks_limit", 3)) + max_retries = int(form_data.get("max_retries", 3)) + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + + # 收集所有文件 + file_list = [] + # 支持 file, file1, file2, ... 或 files[] 格式 + for key in files.keys(): + if key == "file" or key.startswith("file") or key == "files[]": + file_items = files.getlist(key) + file_list.extend(file_items) + + if not file_list: + return Response().error("缺少文件").__dict__ + + # 限制文件数量 + if len(file_list) > 10: + return Response().error("最多只能上传10个文件").__dict__ + + # 处理每个文件 + for file in file_list: + file_name = file.filename + + # 保存到临时文件 + temp_file_path = f"data/temp/{uuid.uuid4()}_{file_name}" + await file.save(temp_file_path) + + try: + # 异步读取文件内容 + async with aiofiles.open(temp_file_path, "rb") as f: + file_content = await f.read() + + # 提取文件类型 + file_type = ( + file_name.rsplit(".", 1)[-1].lower() if "." in file_name else "" + ) + + files_to_upload.append( + { + "file_name": file_name, + "file_content": file_content, + "file_type": file_type, + } + ) + finally: + # 清理临时文件 + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + + # 获取知识库 + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + # 生成任务ID + task_id = str(uuid.uuid4()) + + # 初始化任务状态 + self.upload_tasks[task_id] = { + "status": "pending", + "result": None, + "error": None, + } + + # 启动后台任务 + asyncio.create_task( + self._background_upload_task( + task_id=task_id, + kb_helper=kb_helper, + files_to_upload=files_to_upload, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + ) + ) + + return ( + Response() + .ok( + { + "task_id": task_id, + "file_count": len(files_to_upload), + "message": "task created, processing in background", + } + ) + .__dict__ + ) + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"上传文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"上传文档失败: {str(e)}").__dict__ + + async def get_upload_progress(self): + """获取上传进度和结果 + + Query 参数: + - task_id: 任务 ID (必填) + + 返回状态: + - pending: 任务待处理 + - processing: 任务处理中 + - completed: 任务完成 + - failed: 任务失败 + """ + try: + task_id = request.args.get("task_id") + if not task_id: + return Response().error("缺少参数 task_id").__dict__ + + # 检查任务是否存在 + if task_id not in self.upload_tasks: + return Response().error("找不到该任务").__dict__ + + task_info = self.upload_tasks[task_id] + status = task_info["status"] + + # 构建返回数据 + response_data = { + "task_id": task_id, + "status": status, + } + + # 如果任务正在处理，返回进度信息 + if status == "processing" and task_id in self.upload_progress: + response_data["progress"] = self.upload_progress[task_id] + + # 如果任务完成，返回结果 + if status == "completed": + response_data["result"] = task_info["result"] + # 清理已完成的任务 + # del self.upload_tasks[task_id] + # if task_id in self.upload_progress: + # del self.upload_progress[task_id] + + # 如果任务失败，返回错误信息 + if status == "failed": + response_data["error"] = task_info["error"] + + return Response().ok(response_data).__dict__ + + except Exception as e: + logger.error(f"获取上传进度失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取上传进度失败: {str(e)}").__dict__ + + async def get_document(self): + """获取文档详情 + + Query 参数: + - doc_id: 文档 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + doc_id = request.args.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + doc = await kb_helper.get_document(doc_id) + if not doc: + return Response().error("文档不存在").__dict__ + + return Response().ok(doc.model_dump()).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取文档详情失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取文档详情失败: {str(e)}").__dict__ + + async def delete_document(self): + """删除文档 + + Body: + - kb_id: 知识库 ID (必填) + - doc_id: 文档 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + doc_id = data.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_document(doc_id) + return Response().ok(message="删除文档成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文档失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文档失败: {str(e)}").__dict__ + + async def delete_chunk(self): + """删除文本块 + + Body: + - kb_id: 知识库 ID (必填) + - chunk_id: 块 ID (必填) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + kb_id = data.get("kb_id") + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + chunk_id = data.get("chunk_id") + if not chunk_id: + return Response().error("缺少参数 chunk_id").__dict__ + doc_id = data.get("doc_id") + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + + kb_helper = await kb_manager.get_kb(kb_id) + if not kb_helper: + return Response().error("知识库不存在").__dict__ + + await kb_helper.delete_chunk(chunk_id, doc_id) + return Response().ok(message="删除文本块成功").__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"删除文本块失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除文本块失败: {str(e)}").__dict__ + + async def list_chunks(self): + """获取块列表 + + Query 参数: + - kb_id: 知识库 ID (必填) + - page: 页码 (默认 1) + - page_size: 每页数量 (默认 20) + """ + try: + kb_manager = self._get_kb_manager() + kb_id = request.args.get("kb_id") + doc_id = request.args.get("doc_id") + page = request.args.get("page", 1, type=int) + page_size = request.args.get("page_size", 100, type=int) + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ + if not doc_id: + return Response().error("缺少参数 doc_id").__dict__ + kb_helper = await kb_manager.get_kb(kb_id) + offset = (page - 1) * page_size + limit = page_size + if not kb_helper: + return Response().error("知识库不存在").__dict__ + chunk_list = await kb_helper.get_chunks_by_doc_id( + doc_id=doc_id, offset=offset, limit=limit + ) + return ( + Response() + .ok( + data={ + "items": chunk_list, + "page": page, + "page_size": page_size, + "total": await kb_helper.get_chunk_count_by_doc_id(doc_id), + } + ) + .__dict__ + ) + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"获取块列表失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"获取块列表失败: {str(e)}").__dict__ + + # ===== 检索 API ===== + + async def retrieve(self): + """检索知识库 + + Body: + - query: 查询文本 (必填) + - kb_ids: 知识库 ID 列表 (必填) + - top_k: 返回结果数量 (可选, 默认 5) + - debug: 是否启用调试模式，返回 t-SNE 可视化图片 (可选, 默认 False) + """ + try: + kb_manager = self._get_kb_manager() + data = await request.json + + query = data.get("query") + kb_names = data.get("kb_names") + debug = data.get("debug", False) + + if not query: + return Response().error("缺少参数 query").__dict__ + if not kb_names or not isinstance(kb_names, list): + return Response().error("缺少参数 kb_names 或格式错误").__dict__ + + top_k = data.get("top_k", 5) + + results = await kb_manager.retrieve( + query=query, + kb_names=kb_names, + top_m_final=top_k, + ) + result_list = [] + if results: + result_list = results["results"] + + response_data = { + "results": result_list, + "total": len(result_list), + "query": query, + } + + # Debug 模式：生成 t-SNE 可视化 + if debug: + try: + img_base64 = await generate_tsne_visualization( + query, kb_names, kb_manager + ) + if img_base64: + response_data["visualization"] = img_base64 + except Exception as e: + logger.error(f"生成 t-SNE 可视化失败: {e}") + logger.error(traceback.format_exc()) + response_data["visualization_error"] = str(e) + + return Response().ok(response_data).__dict__ + + except ValueError as e: + return Response().error(str(e)).__dict__ + except Exception as e: + logger.error(f"检索失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"检索失败: {str(e)}").__dict__ + + # ===== 会话知识库配置 API ===== + + async def get_session_kb_config(self): + """获取会话的知识库配置 + + Query 参数: + - session_id: 会话 ID (必填) + + 返回: + - kb_ids: 知识库 ID 列表 + - top_k: 返回结果数量 + - enable_rerank: 是否启用重排序 + """ + try: + from astrbot.core import sp + + session_id = request.args.get("session_id") + + if not session_id: + return Response().error("缺少参数 session_id").__dict__ + + # 从 SharedPreferences 获取配置 + config = await sp.session_get(session_id, "kb_config", default={}) + + logger.debug(f"[KB配置] 读取到配置: session_id={session_id}") + + # 如果没有配置，返回默认值 + if not config: + config = {"kb_ids": [], "top_k": 5, "enable_rerank": True} + + return Response().ok(config).__dict__ + + except Exception as e: + logger.error(f"[KB配置] 获取配置时出错: {e}", exc_info=True) + return Response().error(f"获取会话知识库配置失败: {str(e)}").__dict__ + + async def set_session_kb_config(self): + """设置会话的知识库配置 + + Body: + - scope: 配置范围 (目前只支持 "session") + - scope_id: 会话 ID (必填) + - kb_ids: 知识库 ID 列表 (必填) + - top_k: 返回结果数量 (可选, 默认 5) + - enable_rerank: 是否启用重排序 (可选, 默认 true) + """ + try: + from astrbot.core import sp + + data = await request.json + + scope = data.get("scope") + scope_id = data.get("scope_id") + kb_ids = data.get("kb_ids", []) + top_k = data.get("top_k", 5) + enable_rerank = data.get("enable_rerank", True) + + # 验证参数 + if scope != "session": + return Response().error("目前仅支持 session 范围的配置").__dict__ + + if not scope_id: + return Response().error("缺少参数 scope_id").__dict__ + + if not isinstance(kb_ids, list): + return Response().error("kb_ids 必须是列表").__dict__ + + # 验证知识库是否存在 + kb_mgr = self._get_kb_manager() + invalid_ids = [] + valid_ids = [] + for kb_id in kb_ids: + kb_helper = await kb_mgr.get_kb(kb_id) + if kb_helper: + valid_ids.append(kb_id) + else: + invalid_ids.append(kb_id) + logger.warning(f"[KB配置] 知识库不存在: {kb_id}") + + if invalid_ids: + logger.warning(f"[KB配置] 以下知识库ID无效: {invalid_ids}") + + # 允许保存空列表，表示明确不使用任何知识库 + if kb_ids and not valid_ids: + # 只有当用户提供了 kb_ids 但全部无效时才报错 + return Response().error(f"所有提供的知识库ID都无效: {kb_ids}").__dict__ + + # 如果 kb_ids 为空列表，表示用户想清空配置 + if not kb_ids: + valid_ids = [] + + # 构建配置对象（只保存有效的ID） + config = { + "kb_ids": valid_ids, + "top_k": top_k, + "enable_rerank": enable_rerank, + } + + # 保存到 SharedPreferences + await sp.session_put(scope_id, "kb_config", config) + + # 立即验证是否保存成功 + verify_config = await sp.session_get(scope_id, "kb_config", default={}) + + if verify_config == config: + return ( + Response() + .ok( + {"valid_ids": valid_ids, "invalid_ids": invalid_ids}, + "保存知识库配置成功", + ) + .__dict__ + ) + else: + logger.error("[KB配置] 配置保存失败，验证不匹配") + return Response().error("配置保存失败").__dict__ + + except Exception as e: + logger.error(f"[KB配置] 设置配置时出错: {e}", exc_info=True) + return Response().error(f"设置会话知识库配置失败: {str(e)}").__dict__ + + async def delete_session_kb_config(self): + """删除会话的知识库配置 + + Body: + - scope: 配置范围 (目前只支持 "session") + - scope_id: 会话 ID (必填) + """ + try: + from astrbot.core import sp + + data = await request.json + + scope = data.get("scope") + scope_id = data.get("scope_id") + + # 验证参数 + if scope != "session": + return Response().error("目前仅支持 session 范围的配置").__dict__ + + if not scope_id: + return Response().error("缺少参数 scope_id").__dict__ + + # 从 SharedPreferences 删除配置 + await sp.session_remove(scope_id, "kb_config") + + return Response().ok(message="删除知识库配置成功").__dict__ + + except Exception as e: + logger.error(f"删除会话知识库配置失败: {e}") + logger.error(traceback.format_exc()) + return Response().error(f"删除会话知识库配置失败: {str(e)}").__dict__ diff --git a/astrbot/dashboard/server.py b/astrbot/dashboard/server.py index 9b9588d4..fca650fd 100644 --- a/astrbot/dashboard/server.py +++ b/astrbot/dashboard/server.py @@ -72,6 +72,7 @@ class AstrBotDashboard: ) self.persona_route = PersonaRoute(self.context, db, core_lifecycle) self.t2i_route = T2iRoute(self.context, core_lifecycle) + self.kb_route = KnowledgeBaseRoute(self.context, core_lifecycle) self.app.add_url_rule( "/api/plug/", diff --git a/astrbot/dashboard/utils.py b/astrbot/dashboard/utils.py new file mode 100644 index 00000000..4bdaf43c --- /dev/null +++ b/astrbot/dashboard/utils.py @@ -0,0 +1,161 @@ +import base64 +import os +import traceback +from io import BytesIO +from astrbot.api import logger +from astrbot.core.knowledge_base.kb_helper import KBHelper +from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager +from astrbot.core.db.vec_db.faiss_impl import FaissVecDB + + +async def generate_tsne_visualization( + query: str, kb_names: list[str], kb_manager: KnowledgeBaseManager +) -> str | None: + """生成 t-SNE 可视化图片 + + Args: + query: 查询文本 + kb_names: 知识库名称列表 + kb_manager: 知识库管理器 + + Returns: + 图片路径或 None + """ + try: + import faiss + import numpy as np + import matplotlib + + matplotlib.use("Agg") # 使用非交互式后端 + import matplotlib.pyplot as plt + from sklearn.manifold import TSNE + except ImportError as e: + raise Exception( + "缺少必要的库以生成 t-SNE 可视化。请安装 matplotlib 和 scikit-learn: {e}" + ) from e + + try: + # 获取第一个知识库的向量数据 + kb_helper: KBHelper | None = None + for kb_name in kb_names: + kb_helper = await kb_manager.get_kb_by_name(kb_name) + if kb_helper: + break + + if not kb_helper: + logger.warning("未找到知识库") + return None + + kb = kb_helper.kb + index_path = f"data/knowledge_base/{kb.kb_id}/index.faiss" + + # 读取 FAISS 索引 + if not os.path.exists(index_path): + logger.warning(f"FAISS 索引不存在: {index_path}") + return None + + index = faiss.read_index(index_path) + + if index.ntotal == 0: + logger.warning("索引为空") + return None + + # 提取所有向量 + logger.info(f"提取 {index.ntotal} 个向量用于可视化...") + if isinstance(index, faiss.IndexIDMap): + base_index = faiss.downcast_index(index.index) + if hasattr(base_index, "reconstruct_n"): + vectors = base_index.reconstruct_n(0, index.ntotal) + else: + vectors = np.zeros((index.ntotal, index.d), dtype=np.float32) + for i in range(index.ntotal): + base_index.reconstruct(i, vectors[i]) + elif hasattr(index, "reconstruct_n"): + vectors = index.reconstruct_n(0, index.ntotal) + else: + vectors = np.zeros((index.ntotal, index.d), dtype=np.float32) + for i in range(index.ntotal): + index.reconstruct(i, vectors[i]) + + # 获取查询向量 + vec_db: FaissVecDB = kb_helper.vec_db # type: ignore + embedding_provider = vec_db.embedding_provider + query_embedding = await embedding_provider.get_embedding(query) + query_vector = np.array([query_embedding], dtype=np.float32) + + # 合并所有向量和查询向量 + all_vectors = np.vstack([vectors, query_vector]) + + # t-SNE 降维 + logger.info("开始 t-SNE 降维...") + perplexity = min(30, all_vectors.shape[0] - 1) + tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity) + vectors_2d = tsne.fit_transform(all_vectors) + + # 分离知识库向量和查询向量 + kb_vectors_2d = vectors_2d[:-1] + query_vector_2d = vectors_2d[-1] + + # 可视化 + logger.info("生成可视化图表...") + plt.figure(figsize=(14, 10)) + + # 绘制知识库向量 + scatter = plt.scatter( + kb_vectors_2d[:, 0], + kb_vectors_2d[:, 1], + alpha=0.5, + s=40, + c=range(len(kb_vectors_2d)), + cmap="viridis", + label="Knowledge Base Vectors", + ) + + # 绘制查询向量（红色 X） + plt.scatter( + query_vector_2d[0], + query_vector_2d[1], + c="red", + s=300, + marker="X", + edgecolors="black", + linewidths=2, + label="Query", + zorder=5, + ) + + # 添加查询文本标注 + plt.annotate( + "Query", + (query_vector_2d[0], query_vector_2d[1]), + xytext=(10, 10), + textcoords="offset points", + fontsize=10, + bbox={"boxstyle": "round,pad=0.5", "fc": "yellow", "alpha": 0.7}, + arrowprops={"arrowstyle": "->", "connectionstyle": "arc3,rad=0"}, + ) + + plt.colorbar(scatter, label="Vector Index") + plt.title( + f"t-SNE Visualization: Query in Knowledge Base\n" + f"({index.ntotal} vectors, {index.d} dimensions, KB: {kb.kb_name})", + fontsize=14, + pad=20, + ) + plt.xlabel("t-SNE Dimension 1", fontsize=12) + plt.ylabel("t-SNE Dimension 2", fontsize=12) + plt.grid(True, alpha=0.3) + plt.legend(fontsize=10, loc="upper right") + + # base64 编码图片返回 + buffer = BytesIO() + plt.savefig(buffer, format="png", dpi=150, bbox_inches="tight") + plt.close() + buffer.seek(0) + img_base64 = base64.b64encode(buffer.read()).decode("utf-8") + return img_base64 + + except Exception as e: + logger.error(f"生成 t-SNE 可视化时出错: {e}") + logger.error(traceback.format_exc()) + return None diff --git a/dashboard/src/components/shared/AstrBotConfig.vue b/dashboard/src/components/shared/AstrBotConfig.vue index e45327eb..d6c6fee9 100644 --- a/dashboard/src/components/shared/AstrBotConfig.vue +++ b/dashboard/src/components/shared/AstrBotConfig.vue @@ -7,6 +7,8 @@ import ProviderSelector from './ProviderSelector.vue' import PersonaSelector from './PersonaSelector.vue' import KnowledgeBaseSelector from './KnowledgeBaseSelector.vue' import { useI18n } from '@/i18n/composables' +import axios from 'axios' +import { useToast } from '@/utils/toast' const props = defineProps({ metadata: { @@ -40,6 +42,7 @@ const currentEditingKey = ref('') const currentEditingLanguage = ref('json') const currentEditingTheme = ref('vs-light') let currentEditingKeyIterable = null +const loadingEmbeddingDim = ref(false) function openEditorDialog(key, value, theme, language) { currentEditingKey.value = key @@ -49,10 +52,34 @@ function openEditorDialog(key, value, theme, language) { dialog.value = true } + function saveEditedContent() { dialog.value = false } +async function getEmbeddingDimensions(providerConfig) { + if (loadingEmbeddingDim.value) return + + loadingEmbeddingDim.value = true + try { + const response = await axios.post('/api/config/provider/get_embedding_dim', { + provider_config: providerConfig + }) + + if (response.data.status != "error" && response.data.data?.embedding_dimensions) { + console.log(response.data.data.embedding_dimensions) + providerConfig.embedding_dimensions = response.data.data.embedding_dimensions + useToast().success("获取成功: " + response.data.data.embedding_dimensions) + } else { + useToast().error(response.data.message) + } + } catch (error) { + console.error('Error getting embedding dimensions:', error) + } finally { + loadingEmbeddingDim.value = false + } +} + function getValueBySelector(obj, selector) { const keys = selector.split('.') let current = obj @@ -184,6 +211,29 @@ function hasVisibleItemsAfter(items, currentIndex) { v-model="iterable[key]" /> + +
+ + + 自动检测 + +
+
diff --git a/dashboard/src/components/shared/KnowledgeBaseSelector.vue b/dashboard/src/components/shared/KnowledgeBaseSelector.vue index 3e0a8be1..a90e6f5e 100644 --- a/dashboard/src/components/shared/KnowledgeBaseSelector.vue +++ b/dashboard/src/components/shared/KnowledgeBaseSelector.vue @@ -1,11 +1,21 @@ + + +