Losita/mai-bot
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ruff

Browse Source
This commit is contained in:
春河晴
2025-06-10 16:13:31 +09:00
parent 440e8bf7f3
commit 8d9a88a903
70 changed files with 1598 additions and 1642 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
src/chat/actions/plugin_api/database_api.py
View File

@@ -8,13 +8,16 @@ from peewee import Model, DoesNotExist
logger = get_logger("database_api")
class DatabaseAPI:
"""数据库API模块
提供了数据库操作相关的功能
"""
async def store_action_info(self, action_build_into_prompt: bool = False, action_prompt_display: str = "", action_done: bool = True) -> None:
async def store_action_info(
self, action_build_into_prompt: bool = False, action_prompt_display: str = "", action_done: bool = True
) -> None:
"""存储action执行信息到数据库
Args:
@@ -44,13 +47,13 @@ class DatabaseAPI:
chat_info_platform=chat_stream.platform,
user_id=chat_stream.user_info.user_id if chat_stream.user_info else "",
user_nickname=chat_stream.user_info.user_nickname if chat_stream.user_info else "",
user_cardname=chat_stream.user_info.user_cardname if chat_stream.user_info else ""
user_cardname=chat_stream.user_info.user_cardname if chat_stream.user_info else "",
)
logger.debug(f"{self.log_prefix} 已存储action信息: {action_prompt_display}")
except Exception as e:
logger.error(f"{self.log_prefix} 存储action信息时出错: {e}")
traceback.print_exc()
async def db_query(
self,
model_class: Type[Model],
@@ -59,12 +62,12 @@ class DatabaseAPI:
data: Dict[str, Any] = None,
limit: int = None,
order_by: List[str] = None,
single_result: bool = False
single_result: bool = False,
) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
"""执行数据库查询操作
这个方法提供了一个通用接口来执行数据库操作，包括查询、创建、更新和删除记录。
Args:
model_class: Peewee 模型类，例如 ActionRecords, Messages 等
query_type: 查询类型，可选值: "get", "create", "update", "delete", "count"
@@ -73,7 +76,7 @@ class DatabaseAPI:
limit: 限制结果数量
order_by: 排序字段列表，使用字段名，前缀'-'表示降序
single_result: 是否只返回单个结果
Returns:
根据查询类型返回不同的结果:
- "get": 返回查询结果列表或单个结果（如果 single_result=True）
@@ -81,24 +84,24 @@ class DatabaseAPI:
- "update": 返回受影响的行数
- "delete": 返回受影响的行数
- "count": 返回记录数量
示例:
# 查询最近10条消息
messages = await self.db_query(
Messages,
Messages,
query_type="get",
filters={"chat_id": chat_stream.stream_id},
limit=10,
order_by=["-time"]
)
# 创建一条记录
new_record = await self.db_query(
ActionRecords,
query_type="create",
data={"action_id": "123", "time": time.time(), "action_name": "TestAction"}
)
# 更新记录
updated_count = await self.db_query(
ActionRecords,
@@ -106,14 +109,14 @@ class DatabaseAPI:
filters={"action_id": "123"},
data={"action_done": True}
)
# 删除记录
deleted_count = await self.db_query(
ActionRecords,
query_type="delete",
filters={"action_id": "123"}
)
# 计数
count = await self.db_query(
Messages,
@@ -125,12 +128,12 @@ class DatabaseAPI:
# 构建基本查询
if query_type in ["get", "update", "delete", "count"]:
query = model_class.select()
# 应用过滤条件
if filters:
for field, value in filters.items():
query = query.where(getattr(model_class, field) == value)
# 执行查询
if query_type == "get":
# 应用排序
@@ -140,56 +143,56 @@ class DatabaseAPI:
query = query.order_by(getattr(model_class, field[1:]).desc())
else:
query = query.order_by(getattr(model_class, field))
# 应用限制
if limit:
query = query.limit(limit)
# 执行查询
results = list(query.dicts())
# 返回结果
if single_result:
return results[0] if results else None
return results
elif query_type == "create":
if not data:
raise ValueError("创建记录需要提供data参数")
# 创建记录
record = model_class.create(**data)
# 返回创建的记录
return model_class.select().where(model_class.id == record.id).dicts().get()
elif query_type == "update":
if not data:
raise ValueError("更新记录需要提供data参数")
# 更新记录
return query.update(**data).execute()
elif query_type == "delete":
# 删除记录
return query.delete().execute()
elif query_type == "count":
# 计数
return query.count()
else:
raise ValueError(f"不支持的查询类型: {query_type}")
except DoesNotExist:
# 记录不存在
if query_type == "get" and single_result:
return None
return []
except Exception as e:
logger.error(f"{self.log_prefix} 数据库操作出错: {e}")
traceback.print_exc()
# 根据查询类型返回合适的默认值
if query_type == "get":
return None if single_result else []
@@ -198,21 +201,18 @@ class DatabaseAPI:
raise "unknown query type"
async def db_raw_query(
self,
sql: str,
params: List[Any] = None,
fetch_results: bool = True
self, sql: str, params: List[Any] = None, fetch_results: bool = True
) -> Union[List[Dict[str, Any]], int, None]:
"""执行原始SQL查询
警告: 使用此方法需要小心，确保SQL语句已正确构造以避免SQL注入风险。
Args:
sql: 原始SQL查询字符串
params: 查询参数列表，用于替换SQL中的占位符
fetch_results: 是否获取查询结果，对于SELECT查询设为True，对于
UPDATE/INSERT/DELETE等操作设为False
Returns:
如果fetch_results为True，返回查询结果列表；
如果fetch_results为False，返回受影响的行数；
@@ -220,55 +220,51 @@ class DatabaseAPI:
"""
try:
cursor = db.execute_sql(sql, params or [])
if fetch_results:
# 获取列名
columns = [col[0] for col in cursor.description]
# 构建结果字典列表
results = []
for row in cursor.fetchall():
results.append(dict(zip(columns, row)))
return results
else:
# 返回受影响的行数
return cursor.rowcount
except Exception as e:
logger.error(f"{self.log_prefix} 执行原始SQL查询出错: {e}")
traceback.print_exc()
return None
async def db_save(
self,
model_class: Type[Model],
data: Dict[str, Any],
key_field: str = None,
key_value: Any = None
self, model_class: Type[Model], data: Dict[str, Any], key_field: str = None, key_value: Any = None
) -> Union[Dict[str, Any], None]:
"""保存数据到数据库（创建或更新）
如果提供了key_field和key_value，会先尝试查找匹配的记录进行更新；
如果没有找到匹配记录，或未提供key_field和key_value，则创建新记录。
Args:
model_class: Peewee模型类，如ActionRecords, Messages等
data: 要保存的数据字典
key_field: 用于查找现有记录的字段名，例如"action_id"
key_value: 用于查找现有记录的字段值
Returns:
Dict[str, Any]: 保存后的记录数据
None: 如果操作失败
示例:
# 创建或更新一条记录
record = await self.db_save(
ActionRecords,
{
"action_id": "123",
"time": time.time(),
"action_id": "123",
"time": time.time(),
"action_name": "TestAction",
"action_done": True
},
@@ -280,58 +276,50 @@ class DatabaseAPI:
# 如果提供了key_field和key_value，尝试更新现有记录
if key_field and key_value is not None:
# 查找现有记录
existing_records = list(model_class.select().where(
getattr(model_class, key_field) == key_value
).limit(1))
existing_records = list(
model_class.select().where(getattr(model_class, key_field) == key_value).limit(1)
)
if existing_records:
# 更新现有记录
existing_record = existing_records[0]
for field, value in data.items():
setattr(existing_record, field, value)
existing_record.save()
# 返回更新后的记录
updated_record = model_class.select().where(
model_class.id == existing_record.id
).dicts().get()
updated_record = model_class.select().where(model_class.id == existing_record.id).dicts().get()
return updated_record
# 如果没有找到现有记录或未提供key_field和key_value，创建新记录
new_record = model_class.create(**data)
# 返回创建的记录
created_record = model_class.select().where(
model_class.id == new_record.id
).dicts().get()
created_record = model_class.select().where(model_class.id == new_record.id).dicts().get()
return created_record
except Exception as e:
logger.error(f"{self.log_prefix} 保存数据库记录出错: {e}")
traceback.print_exc()
return None
async def db_get(
self,
model_class: Type[Model],
filters: Dict[str, Any] = None,
order_by: str = None,
limit: int = None
self, model_class: Type[Model], filters: Dict[str, Any] = None, order_by: str = None, limit: int = None
) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
"""从数据库获取记录
这是db_query方法的简化版本，专注于数据检索操作。
Args:
model_class: Peewee模型类
filters: 过滤条件，字段名和值的字典
order_by: 排序字段，前缀'-'表示降序，例如'-time'表示按时间降序
limit: 结果数量限制，如果为1则返回单个记录而不是列表
Returns:
如果limit=1，返回单个记录字典或None；
否则返回记录字典列表或空列表。
示例:
# 获取单个记录
record = await self.db_get(
@@ -339,7 +327,7 @@ class DatabaseAPI:
filters={"action_id": "123"},
limit=1
)
# 获取最近10条记录
records = await self.db_get(
Messages,
@@ -351,32 +339,32 @@ class DatabaseAPI:
try:
# 构建查询
query = model_class.select()
# 应用过滤条件
if filters:
for field, value in filters.items():
query = query.where(getattr(model_class, field) == value)
# 应用排序
if order_by:
if order_by.startswith("-"):
query = query.order_by(getattr(model_class, order_by[1:]).desc())
else:
query = query.order_by(getattr(model_class, order_by))
# 应用限制
if limit:
query = query.limit(limit)
# 执行查询
results = list(query.dicts())
# 返回结果
if limit == 1:
return results[0] if results else None
return results
except Exception as e:
logger.error(f"{self.log_prefix} 获取数据库记录出错: {e}")
traceback.print_exc()
return None if limit == 1 else []
return None if limit == 1 else []