diff --git a/astrbot/core/db/plugin/sqlite_impl.py b/astrbot/core/db/plugin/sqlite_impl.py deleted file mode 100644 index 53cfb828..00000000 --- a/astrbot/core/db/plugin/sqlite_impl.py +++ /dev/null @@ -1,113 +0,0 @@ -import json -import aiosqlite -import os -from typing import Any -from .plugin_storage import PluginStorage -from astrbot.core.utils.astrbot_path import get_astrbot_data_path - -DBPATH = os.path.join(get_astrbot_data_path(), "plugin_data", "sqlite", "plugin_data.db") - - -class SQLitePluginStorage(PluginStorage): - """插件数据的 SQLite 存储实现类。 - - 该类提供异步方式将插件数据存储到 SQLite 数据库中,支持数据的增删改查操作。 - 所有数据以 (plugin, key) 作为复合主键进行索引。 - """ - - _instance = None # Standalone instance of the class - _db_conn = None - db_path = None - - def __new__(cls): - """ - 创建或获取 SQLitePluginStorage 的单例实例。 - 如果实例已存在,则返回现有实例;否则创建一个新实例。 - 数据在 `data/plugin_data/sqlite/plugin_data.db` 下。 - """ - os.makedirs(os.path.dirname(DBPATH), exist_ok=True) - if cls._instance is None: - cls._instance = super(SQLitePluginStorage, cls).__new__(cls) - cls._instance.db_path = DBPATH - return cls._instance - - async def _init_db(self): - """初始化数据库连接(只执行一次)""" - if SQLitePluginStorage._db_conn is None: - SQLitePluginStorage._db_conn = await aiosqlite.connect(self.db_path) - await self._setup_db() - - async def _setup_db(self): - """ - 异步初始化数据库。 - - 创建插件数据表,如果表不存在则创建,表结构包含 plugin、key 和 value 字段, - 其中 plugin 和 key 组合作为主键。 - """ - await self._db_conn.execute(""" - CREATE TABLE IF NOT EXISTS plugin_data ( - plugin TEXT, - key TEXT, - value TEXT, - PRIMARY KEY (plugin, key) - ) - """) - await self._db_conn.commit() - - async def set(self, plugin: str, key: str, value: Any): - """ - 异步存储数据。 - - 将指定插件的键值对存入数据库,如果键已存在则更新值。 - 值会被序列化为 JSON 字符串后存储。 - - Args: - plugin: 插件标识符 - key: 数据键名 - value: 要存储的数据值(任意类型,将被 JSON 序列化) - """ - await self._init_db() - await self._db_conn.execute( - "INSERT INTO plugin_data (plugin, key, value) VALUES (?, ?, ?) " - "ON CONFLICT(plugin, key) DO UPDATE SET value = excluded.value", - (plugin, key, json.dumps(value)), - ) - await self._db_conn.commit() - - async def get(self, plugin: str, key: str) -> Any: - """ - 异步获取数据。 - - 从数据库中获取指定插件和键名对应的值, - 返回的值会从 JSON 字符串反序列化为原始数据类型。 - - Args: - plugin: 插件标识符 - key: 数据键名 - - Returns: - Any: 存储的数据值,如果未找到则返回 None - """ - await self._init_db() - async with self._db_conn.execute( - "SELECT value FROM plugin_data WHERE plugin = ? AND key = ?", - (plugin, key), - ) as cursor: - row = await cursor.fetchone() - return json.loads(row[0]) if row else None - - async def delete(self, plugin: str, key: str): - """ - 异步删除数据。 - - 从数据库中删除指定插件和键名对应的数据项。 - - Args: - plugin: 插件标识符 - key: 要删除的数据键名 - """ - await self._init_db() - await self._db_conn.execute( - "DELETE FROM plugin_data WHERE plugin = ? AND key = ?", (plugin, key) - ) - await self._db_conn.commit() diff --git a/astrbot/core/rag/embedding/openai_source.py b/astrbot/core/rag/embedding/openai_source.py deleted file mode 100644 index dc09d84d..00000000 --- a/astrbot/core/rag/embedding/openai_source.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import List -from openai import AsyncOpenAI - - -class SimpleOpenAIEmbedding: - def __init__( - self, - model, - api_key, - api_base=None, - ) -> None: - self.client = AsyncOpenAI(api_key=api_key, base_url=api_base) - self.model = model - - async def get_embedding(self, text) -> List[float]: - """ - 获取文本的嵌入 - """ - embedding = await self.client.embeddings.create(input=text, model=self.model) - return embedding.data[0].embedding diff --git a/astrbot/core/rag/knowledge_db_mgr.py b/astrbot/core/rag/knowledge_db_mgr.py deleted file mode 100644 index f1c1f386..00000000 --- a/astrbot/core/rag/knowledge_db_mgr.py +++ /dev/null @@ -1,95 +0,0 @@ -import os -from typing import List, Dict -from astrbot.core import logger -from .store import Store -from astrbot.core.config import AstrBotConfig -from astrbot.core.utils.astrbot_path import get_astrbot_data_path - - -class KnowledgeDBManager: - def __init__(self, astrbot_config: AstrBotConfig) -> None: - self.db_path = os.path.join(get_astrbot_data_path(), "knowledge_db") - self.config = astrbot_config.get("knowledge_db", {}) - self.astrbot_config = astrbot_config - if not os.path.exists(self.db_path): - os.makedirs(self.db_path) - self.store_insts: Dict[str, Store] = {} - for name, cfg in self.config.items(): - if cfg["strategy"] == "embedding": - logger.info(f"加载 Chroma Vector Store:{name}") - try: - from .store.chroma_db import ChromaVectorStore - except ImportError as ie: - logger.error(f"{ie} 可能未安装 chromadb 库。") - continue - self.store_insts[name] = ChromaVectorStore( - name, cfg["embedding_config"] - ) - else: - logger.error(f"不支持的策略:{cfg['strategy']}") - - async def list_knowledge_db(self) -> List[str]: - return [ - f - for f in os.listdir(self.db_path) - if os.path.isfile(os.path.join(self.db_path, f)) - ] - - async def create_knowledge_db(self, name: str, config: Dict): - """ - config 格式: - ``` - { - "strategy": "embedding", # 目前只支持 embedding - "chunk_method": { - "strategy": "fixed", - "chunk_size": 100, - "overlap_size": 10 - }, - "embedding_config": { - "strategy": "openai", - "base_url": "", - "model": "", - "api_key": "" - } - } - ``` - """ - if name in self.config: - raise ValueError(f"知识库已存在:{name}") - - self.config[name] = config - self.astrbot_config["knowledge_db"] = self.config - self.astrbot_config.save_config() - - async def insert_record(self, name: str, text: str): - if name not in self.store_insts: - raise ValueError(f"未找到知识库:{name}") - - ret = [] - match self.config[name]["chunk_method"]["strategy"]: - case "fixed": - chunk_size = self.config[name]["chunk_method"]["chunk_size"] - chunk_overlap = self.config[name]["chunk_method"]["overlap_size"] - ret = self._fixed_chunk(text, chunk_size, chunk_overlap) - case _: - pass - - for chunk in ret: - await self.store_insts[name].save(chunk) - - async def retrive_records(self, name: str, query: str, top_n: int = 3) -> List[str]: - if name not in self.store_insts: - raise ValueError(f"未找到知识库:{name}") - - inst = self.store_insts[name] - return await inst.query(query, top_n) - - def _fixed_chunk(self, text: str, chunk_size: int, chunk_overlap: int) -> List[str]: - chunks = [] - start = 0 - while start < len(text): - end = start + chunk_size - chunks.append(text[start:end]) - start += chunk_size - chunk_overlap - return chunks diff --git a/astrbot/core/rag/store/__init__.py b/astrbot/core/rag/store/__init__.py deleted file mode 100644 index 0e74c5a0..00000000 --- a/astrbot/core/rag/store/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from typing import List - - -class Store: - async def save(self, text: str): - pass - - async def query(self, query: str, top_n: int = 3) -> List[str]: - pass diff --git a/astrbot/core/rag/store/chroma_db.py b/astrbot/core/rag/store/chroma_db.py deleted file mode 100644 index d4cfae94..00000000 --- a/astrbot/core/rag/store/chroma_db.py +++ /dev/null @@ -1,44 +0,0 @@ -import chromadb -import uuid -from typing import List, Dict -from astrbot.api import logger -from ..embedding.openai_source import SimpleOpenAIEmbedding -from . import Store -from astrbot.core.utils.astrbot_path import get_astrbot_data_path - - -class ChromaVectorStore(Store): - def __init__(self, name: str, embedding_cfg: Dict) -> None: - import os - self.chroma_client = chromadb.PersistentClient( - path=os.path.join(get_astrbot_data_path(), "long_term_memory_chroma.db") - ) - self.collection = self.chroma_client.get_or_create_collection(name=name) - self.embedding = None - if embedding_cfg["strategy"] == "openai": - self.embedding = SimpleOpenAIEmbedding( - model=embedding_cfg["model"], - api_key=embedding_cfg["api_key"], - api_base=embedding_cfg.get("base_url", None), - ) - - async def save(self, text: str, metadata: Dict = None): - logger.debug(f"Saving text: {text}") - embedding = await self.embedding.get_embedding(text) - - self.collection.upsert( - documents=text, - metadatas=metadata, - ids=str(uuid.uuid4()), - embeddings=embedding, - ) - - async def query( - self, query: str, top_n=3, metadata_filter: Dict = None - ) -> List[str]: - embedding = await self.embedding.get_embedding(query) - - results = self.collection.query( - query_embeddings=embedding, n_results=top_n, where=metadata_filter - ) - return results["documents"][0]