mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-26 00:53:54 +08:00
feat: add FastAPI manager API compatibility baseline
This commit is contained in:
Tyke Chen
@@ -0,0 +1,665 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# All interpolated SQL fragments are selected from closed column/table allowlists.
|
||||
# ruff: noqa: S608
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import bindparam, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
AGENT_COLUMNS = (
|
||||
"id",
|
||||
"user_id",
|
||||
"agent_code",
|
||||
"agent_name",
|
||||
"asr_model_id",
|
||||
"vad_model_id",
|
||||
"llm_model_id",
|
||||
"slm_model_id",
|
||||
"vllm_model_id",
|
||||
"tts_model_id",
|
||||
"tts_voice_id",
|
||||
"tts_language",
|
||||
"tts_volume",
|
||||
"tts_rate",
|
||||
"tts_pitch",
|
||||
"mem_model_id",
|
||||
"intent_model_id",
|
||||
"chat_history_conf",
|
||||
"system_prompt",
|
||||
"summary_memory",
|
||||
"lang_code",
|
||||
"language",
|
||||
"sort",
|
||||
"creator",
|
||||
"created_at",
|
||||
"updater",
|
||||
"updated_at",
|
||||
)
|
||||
AGENT_MUTABLE_COLUMNS = frozenset(AGENT_COLUMNS) - {"id", "user_id", "creator", "created_at"}
|
||||
TEMPLATE_COLUMNS = (
|
||||
"id",
|
||||
"agent_code",
|
||||
"agent_name",
|
||||
"asr_model_id",
|
||||
"vad_model_id",
|
||||
"llm_model_id",
|
||||
"vllm_model_id",
|
||||
"tts_model_id",
|
||||
"tts_voice_id",
|
||||
"tts_language",
|
||||
"tts_volume",
|
||||
"tts_rate",
|
||||
"tts_pitch",
|
||||
"mem_model_id",
|
||||
"intent_model_id",
|
||||
"chat_history_conf",
|
||||
"system_prompt",
|
||||
"summary_memory",
|
||||
"lang_code",
|
||||
"language",
|
||||
"sort",
|
||||
"creator",
|
||||
"created_at",
|
||||
"updater",
|
||||
"updated_at",
|
||||
)
|
||||
|
||||
|
||||
class AgentRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
@property
|
||||
def is_sqlite(self) -> bool:
|
||||
bind = self.session.get_bind()
|
||||
return bool(bind is not None and bind.dialect.name == "sqlite")
|
||||
|
||||
async def get_agent(self, agent_id: str, *, for_update: bool = False) -> dict[str, Any] | None:
|
||||
suffix = "" if self.is_sqlite or not for_update else " FOR UPDATE"
|
||||
return await self.fetch_one(
|
||||
f"SELECT {', '.join(AGENT_COLUMNS)} FROM ai_agent WHERE id=:id{suffix}", {"id": agent_id}
|
||||
)
|
||||
|
||||
async def check_agent_owner(self, agent_id: str, user_id: int, *, super_admin: bool) -> bool:
|
||||
if super_admin:
|
||||
return bool(await self.scalar("SELECT 1 FROM ai_agent WHERE id=:id LIMIT 1", {"id": agent_id}))
|
||||
return bool(
|
||||
await self.scalar(
|
||||
"SELECT 1 FROM ai_agent WHERE id=:id AND user_id=:user_id LIMIT 1",
|
||||
{"id": agent_id, "user_id": user_id},
|
||||
)
|
||||
)
|
||||
|
||||
async def list_user_agents(self, user_id: int, keyword: str | None) -> list[dict[str, Any]]:
|
||||
params: dict[str, Any] = {"user_id": user_id}
|
||||
where = "a.user_id=:user_id"
|
||||
if keyword is not None and keyword.strip():
|
||||
params["keyword"] = f"%{keyword}%"
|
||||
where += (
|
||||
" AND (a.agent_name LIKE :keyword"
|
||||
" OR EXISTS (SELECT 1 FROM ai_device d0 WHERE d0.agent_id=a.id"
|
||||
" AND d0.user_id=:user_id AND d0.mac_address LIKE :keyword)"
|
||||
" OR EXISTS (SELECT 1 FROM ai_agent_tag_relation tr0"
|
||||
" JOIN ai_agent_tag t0 ON t0.id=tr0.tag_id"
|
||||
" WHERE tr0.agent_id=a.id AND t0.deleted=0 AND t0.tag_name LIKE :keyword))"
|
||||
)
|
||||
return await self.fetch_all(
|
||||
"SELECT a.*, mt.model_name AS tts_model_name, ml.model_name AS llm_model_name,"
|
||||
" mv.model_name AS vllm_model_name, COALESCE(tv.name, vc.name) AS tts_voice_name,"
|
||||
" (SELECT MAX(d.last_connected_at) FROM ai_device d WHERE d.agent_id=a.id) AS last_connected_at,"
|
||||
" (SELECT COUNT(*) FROM ai_device d WHERE d.agent_id=a.id) AS device_count"
|
||||
" FROM ai_agent a"
|
||||
" LEFT JOIN ai_model_config mt ON mt.id=a.tts_model_id"
|
||||
" LEFT JOIN ai_model_config ml ON ml.id=a.llm_model_id"
|
||||
" LEFT JOIN ai_model_config mv ON mv.id=a.vllm_model_id"
|
||||
" LEFT JOIN ai_tts_voice tv ON tv.id=a.tts_voice_id"
|
||||
" LEFT JOIN ai_voice_clone vc ON vc.id=a.tts_voice_id"
|
||||
f" WHERE {where} ORDER BY a.created_at DESC",
|
||||
params,
|
||||
)
|
||||
|
||||
async def list_admin_agents(
|
||||
self, page: int, limit: int, order_field: str, ascending: bool
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
allowed = {"agent_name", "created_at", "updated_at", "sort", "id"}
|
||||
selected = order_field if order_field in allowed else "agent_name"
|
||||
direction = "ASC" if ascending else "DESC"
|
||||
total = int(await self.scalar("SELECT COUNT(*) FROM ai_agent") or 0)
|
||||
query = (
|
||||
f"SELECT {', '.join(AGENT_COLUMNS)} FROM ai_agent "
|
||||
f"ORDER BY {selected} {direction} LIMIT :limit OFFSET :offset"
|
||||
)
|
||||
rows = await self.fetch_all(
|
||||
query,
|
||||
{"limit": limit, "offset": (page - 1) * limit},
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def insert_agent(self, values: Mapping[str, Any]) -> int:
|
||||
columns = [column for column in AGENT_COLUMNS if column in values]
|
||||
placeholders = ", ".join(f":{column}" for column in columns)
|
||||
return await self.execute(
|
||||
f"INSERT INTO ai_agent ({', '.join(columns)}) VALUES ({placeholders})",
|
||||
{column: values[column] for column in columns},
|
||||
)
|
||||
|
||||
async def update_agent(self, agent_id: str, values: Mapping[str, Any], *, include_null: bool = False) -> int:
|
||||
selected = {
|
||||
key: value
|
||||
for key, value in values.items()
|
||||
if key in AGENT_MUTABLE_COLUMNS and (include_null or value is not None)
|
||||
}
|
||||
if not selected:
|
||||
return 0
|
||||
assignments = ", ".join(f"{column}=:{column}" for column in selected)
|
||||
return await self.execute(
|
||||
f"UPDATE ai_agent SET {assignments} WHERE id=:agent_id",
|
||||
{**selected, "agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def get_agent_plugins(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT m.id,m.agent_id,m.plugin_id,m.param_info,p.provider_code"
|
||||
" FROM ai_agent_plugin_mapping m LEFT JOIN ai_model_provider p ON p.id=m.plugin_id"
|
||||
" WHERE m.agent_id=:agent_id ORDER BY m.id ASC",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def replace_plugins(self, agent_id: str, plugins: Sequence[Mapping[str, Any]]) -> None:
|
||||
existing = await self.fetch_all(
|
||||
"SELECT id,plugin_id FROM ai_agent_plugin_mapping WHERE agent_id=:agent_id",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
by_plugin = {str(row["plugin_id"]): int(row["id"]) for row in existing}
|
||||
incoming = {str(item.get("plugin_id") or "") for item in plugins}
|
||||
remove_ids = [int(row["id"]) for row in existing if str(row["plugin_id"]) not in incoming]
|
||||
if remove_ids:
|
||||
statement = text("DELETE FROM ai_agent_plugin_mapping WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
await self.execute(statement, {"ids": remove_ids})
|
||||
for item in plugins:
|
||||
plugin_id = str(item.get("plugin_id") or "")
|
||||
params = {"agent_id": agent_id, "plugin_id": plugin_id, "param_info": item.get("param_info") or "{}"}
|
||||
if plugin_id in by_plugin:
|
||||
await self.execute(
|
||||
"UPDATE ai_agent_plugin_mapping SET param_info=:param_info WHERE id=:id",
|
||||
{"id": by_plugin[plugin_id], **params},
|
||||
)
|
||||
else:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_agent_plugin_mapping (id,agent_id,plugin_id,param_info)"
|
||||
" VALUES (:id,:agent_id,:plugin_id,:param_info)",
|
||||
{"id": int(item["id"]), **params},
|
||||
)
|
||||
|
||||
async def delete_plugins(self, agent_id: str) -> int:
|
||||
return await self.execute("DELETE FROM ai_agent_plugin_mapping WHERE agent_id=:id", {"id": agent_id})
|
||||
|
||||
async def get_context_provider(self, agent_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id,agent_id,context_providers,creator,created_at,updater,updated_at"
|
||||
" FROM ai_agent_context_provider WHERE agent_id=:id LIMIT 1",
|
||||
{"id": agent_id},
|
||||
)
|
||||
|
||||
async def upsert_context_provider(self, agent_id: str, encoded: str, new_id: str) -> None:
|
||||
existing = await self.get_context_provider(agent_id)
|
||||
if existing:
|
||||
await self.execute(
|
||||
"UPDATE ai_agent_context_provider SET context_providers=:value WHERE id=:id",
|
||||
{"value": encoded, "id": existing["id"]},
|
||||
)
|
||||
else:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_agent_context_provider (id,agent_id,context_providers)"
|
||||
" VALUES (:id,:agent_id,:value)",
|
||||
{"id": new_id, "agent_id": agent_id, "value": encoded},
|
||||
)
|
||||
|
||||
async def get_correct_word_ids(self, agent_id: str) -> list[str]:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT file_id FROM ai_agent_correct_word_mapping WHERE agent_id=:id", {"id": agent_id}
|
||||
)
|
||||
return [str(row["file_id"]) for row in rows]
|
||||
|
||||
async def replace_correct_words(
|
||||
self, agent_id: str, file_ids: Sequence[str], user_id: int, now: datetime, ids: Sequence[str]
|
||||
) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_mapping WHERE agent_id=:id", {"id": agent_id})
|
||||
await self.execute_many(
|
||||
"INSERT INTO ai_agent_correct_word_mapping"
|
||||
" (id,agent_id,file_id,creator,created_at,updater,updated_at)"
|
||||
" VALUES (:id,:agent_id,:file_id,:user_id,:now,:user_id,:now)",
|
||||
[
|
||||
{"id": mapping_id, "agent_id": agent_id, "file_id": file_id, "user_id": user_id, "now": now}
|
||||
for mapping_id, file_id in zip(ids, file_ids, strict=True)
|
||||
],
|
||||
)
|
||||
|
||||
async def get_agent_tags(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT t.id,t.tag_name,t.sort,r.sort AS relation_sort"
|
||||
" FROM ai_agent_tag t JOIN ai_agent_tag_relation r ON t.id=r.tag_id"
|
||||
" WHERE r.agent_id=:id AND t.deleted=0 ORDER BY r.sort ASC,r.created_at ASC",
|
||||
{"id": agent_id},
|
||||
)
|
||||
|
||||
async def get_tags_for_agents(self, agent_ids: Sequence[str]) -> list[dict[str, Any]]:
|
||||
if not agent_ids:
|
||||
return []
|
||||
statement = text(
|
||||
"SELECT t.id,t.tag_name,r.agent_id,r.sort AS relation_sort"
|
||||
" FROM ai_agent_tag t JOIN ai_agent_tag_relation r ON t.id=r.tag_id"
|
||||
" WHERE r.agent_id IN :ids AND t.deleted=0 ORDER BY r.sort ASC,r.created_at ASC"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
return await self.fetch_all(statement, {"ids": list(agent_ids)})
|
||||
|
||||
async def list_tags(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all("SELECT id,tag_name,sort FROM ai_agent_tag WHERE deleted=0 ORDER BY sort ASC")
|
||||
|
||||
async def get_tag(self, tag_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_agent_tag WHERE id=:id", {"id": tag_id})
|
||||
|
||||
async def find_active_tag_by_name(self, tag_name: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_agent_tag WHERE tag_name=:name AND deleted=0 LIMIT 1", {"name": tag_name}
|
||||
)
|
||||
|
||||
async def find_any_tag_by_name(self, tag_name: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_agent_tag WHERE tag_name=:name LIMIT 1", {"name": tag_name})
|
||||
|
||||
async def insert_tag(self, values: Mapping[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"INSERT INTO ai_agent_tag"
|
||||
" (id,tag_name,sort,deleted,creator,created_at,updater,updated_at)"
|
||||
" VALUES (:id,:tag_name,:sort,:deleted,:creator,:created_at,:updater,:updated_at)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def soft_delete_tag(self, tag_id: str, now: datetime) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_agent_tag SET deleted=1,updated_at=:now WHERE id=:id",
|
||||
{"id": tag_id, "now": now},
|
||||
)
|
||||
|
||||
async def replace_tag_relations(self, agent_id: str, relations: Sequence[Mapping[str, Any]]) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_tag_relation WHERE agent_id=:id", {"id": agent_id})
|
||||
await self.execute_many(
|
||||
"INSERT INTO ai_agent_tag_relation"
|
||||
" (id,agent_id,tag_id,sort,creator,created_at,updater,updated_at)"
|
||||
" VALUES (:id,:agent_id,:tag_id,:sort,:creator,:created_at,:updater,:updated_at)",
|
||||
relations,
|
||||
)
|
||||
|
||||
async def get_model_config(self, model_id: str | None) -> dict[str, Any] | None:
|
||||
if not model_id:
|
||||
return None
|
||||
return await self.fetch_one("SELECT * FROM ai_model_config WHERE id=:id", {"id": model_id})
|
||||
|
||||
async def get_default_llm_config(self) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_model_config WHERE model_type='LLM' AND is_enabled=1"
|
||||
" ORDER BY is_default DESC,sort ASC LIMIT 1"
|
||||
)
|
||||
|
||||
async def get_model_provider(self, provider_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_model_provider WHERE id=:id", {"id": provider_id})
|
||||
|
||||
async def get_timbre(self, timbre_id: str | None) -> dict[str, Any] | None:
|
||||
if not timbre_id:
|
||||
return None
|
||||
row = await self.fetch_one("SELECT * FROM ai_tts_voice WHERE id=:id", {"id": timbre_id})
|
||||
if row is None:
|
||||
row = await self.fetch_one("SELECT * FROM ai_voice_clone WHERE id=:id", {"id": timbre_id})
|
||||
return row
|
||||
|
||||
async def find_timbre_by_voice_code(self, model_id: str, voice_code: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_tts_voice WHERE tts_model_id=:model_id AND tts_voice=:voice LIMIT 1",
|
||||
{"model_id": model_id, "voice": voice_code},
|
||||
)
|
||||
|
||||
async def get_device_by_mac(self, mac_address: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_device WHERE mac_address=:mac ORDER BY id DESC LIMIT 1", {"mac": mac_address}
|
||||
)
|
||||
|
||||
async def get_agent_by_device_mac(self, mac_address: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {', '.join('a.' + column for column in AGENT_COLUMNS)}"
|
||||
" FROM ai_device d LEFT JOIN ai_agent a ON d.agent_id=a.id"
|
||||
" WHERE d.mac_address=:mac ORDER BY d.id DESC LIMIT 1",
|
||||
{"mac": mac_address},
|
||||
)
|
||||
|
||||
async def update_device_connection(self, device_id: str, now: datetime) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_device SET last_connected_at=:now WHERE id=:id", {"id": device_id, "now": now}
|
||||
)
|
||||
|
||||
async def insert_chat_audio(self, audio_id: str, audio: bytes) -> int:
|
||||
return await self.execute(
|
||||
"INSERT INTO ai_agent_chat_audio (id,audio) VALUES (:id,:audio)", {"id": audio_id, "audio": audio}
|
||||
)
|
||||
|
||||
async def get_chat_audio(self, audio_id: str) -> bytes | None:
|
||||
value = await self.scalar("SELECT audio FROM ai_agent_chat_audio WHERE id=:id", {"id": audio_id})
|
||||
return bytes(value) if value is not None else None
|
||||
|
||||
async def insert_chat_history(self, values: Mapping[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"INSERT INTO ai_agent_chat_history"
|
||||
" (mac_address,agent_id,session_id,chat_type,content,audio_id,created_at)"
|
||||
" VALUES (:mac_address,:agent_id,:session_id,:chat_type,:content,:audio_id,:created_at)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def get_session_agent_id(self, session_id: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT agent_id FROM ai_agent_chat_history WHERE session_id=:id LIMIT 1", {"id": session_id}
|
||||
)
|
||||
return str(value) if value is not None else None
|
||||
|
||||
async def get_audio_agent_id(self, audio_id: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT agent_id FROM ai_agent_chat_history WHERE audio_id=:id LIMIT 1", {"id": audio_id}
|
||||
)
|
||||
return str(value) if value is not None else None
|
||||
|
||||
async def is_audio_owned(self, audio_id: str, agent_id: str) -> bool:
|
||||
count = await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_agent_chat_history WHERE audio_id=:audio_id AND agent_id=:agent_id",
|
||||
{"audio_id": audio_id, "agent_id": agent_id},
|
||||
)
|
||||
return int(count or 0) == 1
|
||||
|
||||
async def get_audio_content(self, audio_id: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT content FROM ai_agent_chat_history WHERE audio_id=:id LIMIT 1", {"id": audio_id}
|
||||
)
|
||||
return str(value) if value is not None else None
|
||||
|
||||
async def get_chat_history(self, agent_id: str, session_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT created_at,chat_type,content,audio_id,mac_address"
|
||||
" FROM ai_agent_chat_history WHERE agent_id=:agent_id AND session_id=:session_id"
|
||||
" ORDER BY created_at ASC",
|
||||
{"agent_id": agent_id, "session_id": session_id},
|
||||
)
|
||||
|
||||
async def get_recent_user_history(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT content,audio_id FROM ai_agent_chat_history"
|
||||
" WHERE agent_id=:id AND chat_type=1 AND audio_id IS NOT NULL ORDER BY id DESC LIMIT 50",
|
||||
{"id": agent_id},
|
||||
)
|
||||
|
||||
async def list_sessions(self, agent_id: str, page: int, limit: int) -> tuple[list[dict[str, Any]], int]:
|
||||
total = int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM (SELECT session_id FROM ai_agent_chat_history"
|
||||
" WHERE agent_id=:id GROUP BY session_id) sessions",
|
||||
{"id": agent_id},
|
||||
)
|
||||
or 0
|
||||
)
|
||||
rows = await self.fetch_all(
|
||||
"SELECT h.session_id,MAX(h.created_at) AS created_at,COUNT(*) AS chat_count,"
|
||||
" (SELECT t.title FROM ai_agent_chat_title t WHERE t.session_id=h.session_id LIMIT 1) AS title"
|
||||
" FROM ai_agent_chat_history h WHERE h.agent_id=:id GROUP BY h.session_id"
|
||||
" ORDER BY created_at DESC LIMIT :limit OFFSET :offset",
|
||||
{"id": agent_id, "limit": limit, "offset": (page - 1) * limit},
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def upsert_chat_title(self, session_id: str, title: str, now: datetime, title_id: str) -> None:
|
||||
existing = await self.fetch_one(
|
||||
"SELECT id FROM ai_agent_chat_title WHERE session_id=:session_id LIMIT 1", {"session_id": session_id}
|
||||
)
|
||||
if existing:
|
||||
await self.execute(
|
||||
"UPDATE ai_agent_chat_title SET title=:title,updated_at=:now WHERE id=:id",
|
||||
{"id": existing["id"], "title": title, "now": now},
|
||||
)
|
||||
else:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_agent_chat_title (id,session_id,title,created_at,updated_at)"
|
||||
" VALUES (:id,:session_id,:title,:now,:now)",
|
||||
{"id": title_id, "session_id": session_id, "title": title, "now": now},
|
||||
)
|
||||
|
||||
async def delete_chat_history(self, agent_id: str, *, delete_audio: bool, delete_text: bool) -> None:
|
||||
if delete_audio:
|
||||
ids = await self.fetch_all(
|
||||
"SELECT DISTINCT audio_id FROM ai_agent_chat_history WHERE agent_id=:id AND audio_id IS NOT NULL",
|
||||
{"id": agent_id},
|
||||
)
|
||||
audio_ids = [str(row["audio_id"]) for row in ids]
|
||||
for offset in range(0, len(audio_ids), 1000):
|
||||
batch = audio_ids[offset : offset + 1000]
|
||||
statement = text("DELETE FROM ai_agent_chat_audio WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
await self.execute(statement, {"ids": batch})
|
||||
if delete_audio and not delete_text:
|
||||
await self.execute("UPDATE ai_agent_chat_history SET audio_id=NULL WHERE agent_id=:id", {"id": agent_id})
|
||||
if delete_text:
|
||||
await self.execute("DELETE FROM ai_agent_chat_history WHERE agent_id=:id", {"id": agent_id})
|
||||
|
||||
async def delete_agent_cascade(self, agent_id: str) -> None:
|
||||
devices = await self.fetch_all("SELECT mac_address FROM ai_device WHERE agent_id=:id", {"id": agent_id})
|
||||
macs = [str(row["mac_address"]) for row in devices if row.get("mac_address") is not None]
|
||||
await self.execute("DELETE FROM ai_device WHERE agent_id=:id", {"id": agent_id})
|
||||
if macs:
|
||||
statement = text(
|
||||
"DELETE FROM ai_device_address_book WHERE mac_address IN :macs OR target_mac IN :targets"
|
||||
).bindparams(bindparam("macs", expanding=True), bindparam("targets", expanding=True))
|
||||
await self.execute(statement, {"macs": macs, "targets": macs})
|
||||
await self.delete_chat_history(agent_id, delete_audio=True, delete_text=True)
|
||||
for table in (
|
||||
"ai_agent_plugin_mapping",
|
||||
"ai_agent_context_provider",
|
||||
"ai_agent_correct_word_mapping",
|
||||
"ai_agent_tag_relation",
|
||||
"ai_agent_snapshot",
|
||||
):
|
||||
await self.execute(f"DELETE FROM {table} WHERE agent_id=:id", {"id": agent_id})
|
||||
await self.execute("DELETE FROM ai_agent WHERE id=:id", {"id": agent_id})
|
||||
|
||||
async def list_templates(
|
||||
self, *, name: str | None = None, page: int | None = None, limit: int | None = None
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
params: dict[str, Any] = {}
|
||||
where = ""
|
||||
if name:
|
||||
where = " WHERE agent_name LIKE :name"
|
||||
params["name"] = f"%{name}%"
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_agent_template{where}", params) or 0)
|
||||
paging = ""
|
||||
if page is not None and limit is not None:
|
||||
params.update(limit=limit, offset=(page - 1) * limit)
|
||||
paging = " LIMIT :limit OFFSET :offset"
|
||||
rows = await self.fetch_all(f"SELECT * FROM ai_agent_template{where} ORDER BY sort ASC{paging}", params)
|
||||
return rows, total
|
||||
|
||||
async def get_template(self, template_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_agent_template WHERE id=:id", {"id": template_id})
|
||||
|
||||
async def get_default_template(self) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_agent_template ORDER BY sort ASC LIMIT 1")
|
||||
|
||||
async def next_template_sort(self) -> int:
|
||||
rows = await self.fetch_all("SELECT sort FROM ai_agent_template WHERE sort IS NOT NULL ORDER BY sort ASC")
|
||||
expected = 1
|
||||
for row in rows:
|
||||
value = int(row["sort"])
|
||||
if value > expected:
|
||||
return expected
|
||||
expected = value + 1
|
||||
return expected
|
||||
|
||||
async def insert_template(self, values: Mapping[str, Any]) -> int:
|
||||
columns = [column for column in TEMPLATE_COLUMNS if column in values]
|
||||
return await self.execute(
|
||||
f"INSERT INTO ai_agent_template ({', '.join(columns)})"
|
||||
f" VALUES ({', '.join(':' + column for column in columns)})",
|
||||
{column: values[column] for column in columns},
|
||||
)
|
||||
|
||||
async def update_template(self, template_id: str, values: Mapping[str, Any]) -> int:
|
||||
selected = {
|
||||
key: value for key, value in values.items() if key in TEMPLATE_COLUMNS and key != "id" and value is not None
|
||||
}
|
||||
if not selected:
|
||||
return 0
|
||||
return await self.execute(
|
||||
f"UPDATE ai_agent_template SET {', '.join(key + '=:' + key for key in selected)} WHERE id=:id",
|
||||
{**selected, "id": template_id},
|
||||
)
|
||||
|
||||
async def delete_template(self, template_id: str) -> int:
|
||||
return await self.execute("DELETE FROM ai_agent_template WHERE id=:id", {"id": template_id})
|
||||
|
||||
async def reorder_templates(self, deleted_sort: int) -> int:
|
||||
return await self.execute("UPDATE ai_agent_template SET sort=sort-1 WHERE sort>:sort", {"sort": deleted_sort})
|
||||
|
||||
async def delete_templates(self, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
statement = text("DELETE FROM ai_agent_template WHERE id IN :ids").bindparams(bindparam("ids", expanding=True))
|
||||
return await self.execute(statement, {"ids": list(ids)})
|
||||
|
||||
async def list_voiceprints(self, agent_id: str, user_id: int) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id,audio_id,source_name,introduce,create_date"
|
||||
" FROM ai_agent_voice_print WHERE agent_id=:agent_id AND creator=:user_id",
|
||||
{"agent_id": agent_id, "user_id": user_id},
|
||||
)
|
||||
|
||||
async def list_voiceprint_ids(self, agent_id: str) -> list[str]:
|
||||
rows = await self.fetch_all("SELECT id FROM ai_agent_voice_print WHERE agent_id=:id", {"id": agent_id})
|
||||
return [str(row["id"]) for row in rows]
|
||||
|
||||
async def get_voiceprint(self, voiceprint_id: str, user_id: int | None = None) -> dict[str, Any] | None:
|
||||
where = "id=:id"
|
||||
params: dict[str, Any] = {"id": voiceprint_id}
|
||||
if user_id is not None:
|
||||
where += " AND creator=:user_id"
|
||||
params["user_id"] = user_id
|
||||
return await self.fetch_one(f"SELECT * FROM ai_agent_voice_print WHERE {where} LIMIT 1", params)
|
||||
|
||||
async def insert_voiceprint(self, values: Mapping[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"INSERT INTO ai_agent_voice_print"
|
||||
" (id,agent_id,audio_id,source_name,introduce,creator,create_date,updater,update_date)"
|
||||
" VALUES (:id,:agent_id,:audio_id,:source_name,:introduce,:creator,:create_date,:updater,:update_date)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update_voiceprint(self, voiceprint_id: str, user_id: int, values: Mapping[str, Any]) -> int:
|
||||
allowed = {"audio_id", "source_name", "introduce", "updater", "update_date"}
|
||||
selected = {key: value for key, value in values.items() if key in allowed and value is not None}
|
||||
if not selected:
|
||||
return 0
|
||||
return await self.execute(
|
||||
f"UPDATE ai_agent_voice_print SET {', '.join(key + '=:' + key for key in selected)}"
|
||||
" WHERE id=:id AND creator=:user_id",
|
||||
{**selected, "id": voiceprint_id, "user_id": user_id},
|
||||
)
|
||||
|
||||
async def delete_voiceprint(self, voiceprint_id: str, user_id: int) -> int:
|
||||
return await self.execute(
|
||||
"DELETE FROM ai_agent_voice_print WHERE id=:id AND creator=:user_id",
|
||||
{"id": voiceprint_id, "user_id": user_id},
|
||||
)
|
||||
|
||||
async def snapshot_max_version(self, agent_id: str) -> int:
|
||||
return int(
|
||||
await self.scalar(
|
||||
"SELECT COALESCE(MAX(version_no),0) FROM ai_agent_snapshot WHERE agent_id=:id", {"id": agent_id}
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
async def latest_snapshot(self, agent_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_agent_snapshot WHERE agent_id=:id ORDER BY version_no DESC LIMIT 1", {"id": agent_id}
|
||||
)
|
||||
|
||||
async def next_snapshot(self, agent_id: str, version_no: int) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_agent_snapshot WHERE agent_id=:id AND version_no>:version"
|
||||
" ORDER BY version_no ASC LIMIT 1",
|
||||
{"id": agent_id, "version": version_no},
|
||||
)
|
||||
|
||||
async def get_snapshot(self, snapshot_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one("SELECT * FROM ai_agent_snapshot WHERE id=:id", {"id": snapshot_id})
|
||||
|
||||
async def list_snapshots(
|
||||
self, agent_id: str, page: int, limit: int, max_version_no: int | None
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
params: dict[str, Any] = {"id": agent_id, "limit": limit, "offset": (page - 1) * limit}
|
||||
extra = ""
|
||||
if max_version_no is not None:
|
||||
extra = " AND version_no<=:max_version"
|
||||
params["max_version"] = max_version_no
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_agent_snapshot WHERE agent_id=:id{extra}", params) or 0)
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_agent_snapshot WHERE agent_id=:id{extra}"
|
||||
" ORDER BY version_no DESC LIMIT :limit OFFSET :offset",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def insert_snapshot_next_version(self, values: Mapping[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"INSERT INTO ai_agent_snapshot"
|
||||
" (id,agent_id,user_id,version_no,snapshot_data,changed_fields,source,"
|
||||
" restore_from_snapshot_id,restore_from_version_no,creator,created_at,redaction_version)"
|
||||
" SELECT :id,:agent_id,:user_id,COALESCE(MAX(version_no),0)+1,:snapshot_data,:changed_fields,:source,"
|
||||
" :restore_from_snapshot_id,:restore_from_version_no,:creator,:created_at,:redaction_version"
|
||||
" FROM ai_agent_snapshot WHERE agent_id=:agent_id",
|
||||
values,
|
||||
)
|
||||
|
||||
async def prune_snapshots(self, agent_id: str, keep: int) -> int:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT id FROM ai_agent_snapshot WHERE agent_id=:id ORDER BY version_no DESC LIMIT :keep",
|
||||
{"id": agent_id, "keep": keep},
|
||||
)
|
||||
retained = [str(row["id"]) for row in rows]
|
||||
if not retained:
|
||||
return 0
|
||||
statement = text("DELETE FROM ai_agent_snapshot WHERE agent_id=:agent_id AND id NOT IN :retained").bindparams(
|
||||
bindparam("retained", expanding=True)
|
||||
)
|
||||
return await self.execute(statement, {"agent_id": agent_id, "retained": retained})
|
||||
|
||||
async def delete_snapshot(self, snapshot_id: str) -> int:
|
||||
return await self.execute("DELETE FROM ai_agent_snapshot WHERE id=:id", {"id": snapshot_id})
|
||||
|
||||
async def list_legacy_snapshots(self, after_id: str | None, limit: int, version: int) -> list[dict[str, Any]]:
|
||||
params: dict[str, Any] = {"version": version, "limit": limit}
|
||||
extra = ""
|
||||
if after_id is not None:
|
||||
extra = " AND id>:after_id"
|
||||
params["after_id"] = after_id
|
||||
return await self.fetch_all(
|
||||
f"SELECT id,snapshot_data,redaction_version FROM ai_agent_snapshot"
|
||||
f" WHERE redaction_version<:version{extra} ORDER BY id ASC LIMIT :limit",
|
||||
params,
|
||||
)
|
||||
|
||||
async def update_redacted_snapshot(self, snapshot_id: str, snapshot_data: str, version: int) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_agent_snapshot SET snapshot_data=:data,redaction_version=:version"
|
||||
" WHERE id=:id AND redaction_version<:version",
|
||||
{"id": snapshot_id, "data": snapshot_data, "version": version},
|
||||
)
|
||||
@@ -0,0 +1,105 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class ConfigRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def list_params(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT param_code, param_value, value_type FROM sys_params WHERE param_type = 1"
|
||||
)
|
||||
|
||||
async def get_param_value(self, code: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT param_value FROM sys_params WHERE param_code = :code LIMIT 1",
|
||||
{"code": code},
|
||||
)
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def get_default_template(self) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, agent_code, agent_name, asr_model_id, vad_model_id, llm_model_id, "
|
||||
"vllm_model_id, tts_model_id, tts_voice_id, tts_language, tts_volume, tts_rate, tts_pitch, "
|
||||
"mem_model_id, intent_model_id, chat_history_conf, system_prompt, summary_memory, lang_code, "
|
||||
"language, sort FROM ai_agent_template ORDER BY sort ASC LIMIT 1"
|
||||
)
|
||||
|
||||
async def get_device_by_mac(self, mac_address: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, user_id, mac_address, board, agent_id, app_version, auto_update "
|
||||
"FROM ai_device WHERE mac_address = :mac_address LIMIT 1",
|
||||
{"mac_address": mac_address},
|
||||
)
|
||||
|
||||
async def get_agent(self, agent_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, user_id, agent_code, agent_name, asr_model_id, vad_model_id, llm_model_id, slm_model_id, "
|
||||
"vllm_model_id, tts_model_id, tts_voice_id, tts_language, tts_volume, tts_rate, tts_pitch, "
|
||||
"mem_model_id, intent_model_id, chat_history_conf, system_prompt, summary_memory, lang_code, language "
|
||||
"FROM ai_agent WHERE id = :id LIMIT 1",
|
||||
{"id": agent_id},
|
||||
)
|
||||
|
||||
async def get_model(self, model_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, model_type, model_code, model_name, is_default, is_enabled, config_json, doc_link, "
|
||||
"remark, sort, creator, create_date, updater, update_date "
|
||||
"FROM ai_model_config WHERE id = :id LIMIT 1",
|
||||
{"id": model_id},
|
||||
)
|
||||
|
||||
async def get_timbre(self, timbre_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, languages, name, remark, reference_audio, reference_text, sort, tts_model_id, "
|
||||
"tts_voice, voice_demo FROM ai_tts_voice WHERE id = :id LIMIT 1",
|
||||
{"id": timbre_id},
|
||||
)
|
||||
|
||||
async def get_voice_clone(self, clone_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, name, model_id, voice_id, languages, user_id, train_status, train_error "
|
||||
"FROM ai_voice_clone WHERE id = :id LIMIT 1",
|
||||
{"id": clone_id},
|
||||
)
|
||||
|
||||
async def get_plugin_mappings(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT m.id, m.agent_id, m.plugin_id, m.param_info, "
|
||||
"(SELECT p.provider_code FROM ai_model_provider p WHERE p.id = m.plugin_id LIMIT 1) AS provider_code "
|
||||
"FROM ai_agent_plugin_mapping m WHERE m.agent_id = :agent_id",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def get_dataset(self, dataset_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, dataset_id, rag_model_id, name, description, status "
|
||||
"FROM ai_rag_dataset WHERE id = :id LIMIT 1",
|
||||
{"id": dataset_id},
|
||||
)
|
||||
|
||||
async def get_context_providers(self, agent_id: str) -> Any:
|
||||
return await self.scalar(
|
||||
"SELECT context_providers FROM ai_agent_context_provider WHERE agent_id = :agent_id LIMIT 1",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def get_voiceprints(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, agent_id, source_name, introduce, create_date "
|
||||
"FROM ai_agent_voice_print WHERE agent_id = :agent_id ORDER BY create_date ASC",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def get_correct_word_items(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT i.source_word, i.target_word FROM ai_agent_correct_word_mapping m "
|
||||
"JOIN ai_agent_correct_word_item i ON i.file_id = m.file_id WHERE m.agent_id = :agent_id",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
@@ -0,0 +1,115 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class CorrectWordRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def name_exists(self, user_id: int, file_name: str, exclude_id: str | None = None) -> bool:
|
||||
return bool(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_agent_correct_word_file WHERE creator=:creator AND file_name=:file_name "
|
||||
"AND (:exclude_id IS NULL OR id<>:exclude_id)",
|
||||
{"creator": user_id, "file_name": file_name, "exclude_id": exclude_id},
|
||||
)
|
||||
)
|
||||
|
||||
async def insert_file(self, values: dict[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_agent_correct_word_file "
|
||||
"(id, file_name, word_count, content, creator, created_at) "
|
||||
"VALUES (:id, :file_name, :word_count, :content, :creator, :now)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def insert_items(self, values: Sequence[dict[str, Any]]) -> None:
|
||||
await self.execute_many(
|
||||
"INSERT INTO ai_agent_correct_word_item (id, file_id, source_word, target_word) "
|
||||
"VALUES (:id, :file_id, :source_word, :target_word)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def get_file(self, file_id: str, *, for_update: bool = False) -> dict[str, Any] | None:
|
||||
suffix = " FOR UPDATE" if for_update and self.session.get_bind().dialect.name != "sqlite" else ""
|
||||
return await self.fetch_one(
|
||||
f"SELECT * FROM ai_agent_correct_word_file WHERE id=:id{suffix}", # noqa: S608
|
||||
{"id": file_id},
|
||||
)
|
||||
|
||||
async def update_file(self, values: dict[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_agent_correct_word_file SET file_name=:file_name, word_count=:word_count, "
|
||||
"content=:content, updater=:updater, updated_at=:now WHERE id=:id",
|
||||
values,
|
||||
)
|
||||
|
||||
async def list_files(
|
||||
self, user_id: int, *, offset: int | None = None, limit: int | None = None
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
total = int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_agent_correct_word_file WHERE creator=:creator", {"creator": user_id}
|
||||
)
|
||||
or 0
|
||||
)
|
||||
if offset is None or limit is None:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT * FROM ai_agent_correct_word_file WHERE creator=:creator ORDER BY created_at DESC",
|
||||
{"creator": user_id},
|
||||
)
|
||||
else:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT * FROM ai_agent_correct_word_file WHERE creator=:creator ORDER BY created_at DESC "
|
||||
"LIMIT :offset, :limit",
|
||||
{"creator": user_id, "offset": offset, "limit": limit},
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def delete_file_graph(self, file_id: str) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_mapping WHERE file_id=:id", {"id": file_id})
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_item WHERE file_id=:id", {"id": file_id})
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_file WHERE id=:id", {"id": file_id})
|
||||
|
||||
async def delete_items(self, file_id: str) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_item WHERE file_id=:id", {"id": file_id})
|
||||
|
||||
async def items_for_agent(self, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT i.source_word, i.target_word FROM ai_agent_correct_word_item i "
|
||||
"JOIN ai_agent_correct_word_mapping m ON m.file_id=i.file_id WHERE m.agent_id=:agent_id",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def file_ids_for_agent(self, agent_id: str) -> list[str]:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT file_id FROM ai_agent_correct_word_mapping WHERE agent_id=:agent_id", {"agent_id": agent_id}
|
||||
)
|
||||
return [str(row["file_id"]) for row in rows]
|
||||
|
||||
async def replace_agent_mappings(
|
||||
self, agent_id: str, file_ids: Sequence[str], user_id: int, now: Any
|
||||
) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_correct_word_mapping WHERE agent_id=:agent_id", {"agent_id": agent_id})
|
||||
await self.execute_many(
|
||||
"INSERT INTO ai_agent_correct_word_mapping "
|
||||
"(id, agent_id, file_id, creator, created_at, updater, updated_at) "
|
||||
"VALUES (:id, :agent_id, :file_id, :user_id, :now, :user_id, :now)",
|
||||
[
|
||||
{
|
||||
"id": uuid.uuid4().hex,
|
||||
"agent_id": agent_id,
|
||||
"file_id": file_id,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
}
|
||||
for file_id in file_ids
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,294 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# Every interpolated SQL fragment below is a module constant or a service-side allowlist.
|
||||
# ruff: noqa: S608
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
DEVICE_COLUMNS = (
|
||||
"id, user_id, mac_address, last_connected_at, auto_update, board, alias, "
|
||||
"agent_id, app_version, sort, updater, update_date, creator, create_date"
|
||||
)
|
||||
OTA_COLUMNS = (
|
||||
"id, firmware_name, type, version, size, remark, firmware_path, sort, "
|
||||
"updater, update_date, creator, create_date"
|
||||
)
|
||||
ADDRESS_BOOK_COLUMNS = (
|
||||
"mac_address, target_mac, alias, has_permission, creator, create_date, updater, update_date"
|
||||
)
|
||||
|
||||
|
||||
class DeviceRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def get_device(self, device_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {DEVICE_COLUMNS} FROM ai_device WHERE id = :id LIMIT 1",
|
||||
{"id": device_id},
|
||||
)
|
||||
|
||||
async def get_device_by_mac(self, mac_address: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {DEVICE_COLUMNS} FROM ai_device WHERE mac_address = :mac_address LIMIT 1",
|
||||
{"mac_address": mac_address},
|
||||
)
|
||||
|
||||
async def get_user_devices(self, user_id: int, agent_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
f"SELECT {DEVICE_COLUMNS} FROM ai_device WHERE user_id = :user_id AND agent_id = :agent_id",
|
||||
{"user_id": user_id, "agent_id": agent_id},
|
||||
)
|
||||
|
||||
async def insert_device(self, values: Mapping[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_device "
|
||||
"(id, user_id, mac_address, last_connected_at, auto_update, board, alias, agent_id, app_version, "
|
||||
"sort, updater, update_date, creator, create_date) "
|
||||
"VALUES (:id, :user_id, :mac_address, :last_connected_at, :auto_update, :board, :alias, :agent_id, "
|
||||
":app_version, :sort, :updater, :update_date, :creator, :create_date)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update_device_info(
|
||||
self,
|
||||
device_id: str,
|
||||
*,
|
||||
auto_update: int | None,
|
||||
alias: str | None,
|
||||
updater: int,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
assignments = ["updater = :updater", "update_date = :now"]
|
||||
params: dict[str, Any] = {"id": device_id, "updater": updater, "now": now}
|
||||
if auto_update is not None:
|
||||
assignments.append("auto_update = :auto_update")
|
||||
params["auto_update"] = auto_update
|
||||
if alias is not None:
|
||||
assignments.append("alias = :alias")
|
||||
params["alias"] = alias
|
||||
return await self.execute(
|
||||
f"UPDATE ai_device SET {', '.join(assignments)} WHERE id = :id",
|
||||
params,
|
||||
)
|
||||
|
||||
async def update_connection(
|
||||
self,
|
||||
device_id: str,
|
||||
*,
|
||||
app_version: str | None,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
if app_version is None or not app_version.strip():
|
||||
return await self.execute(
|
||||
"UPDATE ai_device SET last_connected_at = :now WHERE id = :id",
|
||||
{"id": device_id, "now": now},
|
||||
)
|
||||
return await self.execute(
|
||||
"UPDATE ai_device SET last_connected_at = :now, app_version = :app_version WHERE id = :id",
|
||||
{"id": device_id, "now": now, "app_version": app_version},
|
||||
)
|
||||
|
||||
async def delete_device_for_user(self, device_id: str, user_id: int) -> int:
|
||||
return await self.execute(
|
||||
"DELETE FROM ai_device WHERE id = :id AND user_id = :user_id",
|
||||
{"id": device_id, "user_id": user_id},
|
||||
)
|
||||
|
||||
async def get_address_book(self, mac_address: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
f"SELECT {ADDRESS_BOOK_COLUMNS} FROM ai_device_address_book "
|
||||
"WHERE mac_address = :mac_address ORDER BY update_date DESC",
|
||||
{"mac_address": mac_address},
|
||||
)
|
||||
|
||||
async def get_all_address_book(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(f"SELECT {ADDRESS_BOOK_COLUMNS} FROM ai_device_address_book")
|
||||
|
||||
async def get_address_book_record(self, mac_address: str, target_mac: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {ADDRESS_BOOK_COLUMNS} FROM ai_device_address_book "
|
||||
"WHERE mac_address = :mac_address AND target_mac = :target_mac LIMIT 1",
|
||||
{"mac_address": mac_address, "target_mac": target_mac},
|
||||
)
|
||||
|
||||
async def get_aliases(self, mac_address: str) -> list[str]:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT alias FROM ai_device_address_book WHERE mac_address = :mac_address",
|
||||
{"mac_address": mac_address},
|
||||
)
|
||||
return [str(row["alias"]) for row in rows if row.get("alias") not in (None, "")]
|
||||
|
||||
async def insert_address_book(
|
||||
self,
|
||||
*,
|
||||
mac_address: str,
|
||||
target_mac: str,
|
||||
alias: str | None,
|
||||
has_permission: bool | None,
|
||||
actor: int,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_device_address_book "
|
||||
"(mac_address, target_mac, alias, has_permission, creator, create_date, updater, update_date) "
|
||||
"VALUES (:mac_address, :target_mac, :alias, :has_permission, :actor, :now, :actor, :now)",
|
||||
{
|
||||
"mac_address": mac_address,
|
||||
"target_mac": target_mac,
|
||||
"alias": alias,
|
||||
"has_permission": has_permission,
|
||||
"actor": actor,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_address_alias(
|
||||
self,
|
||||
mac_address: str,
|
||||
target_mac: str,
|
||||
alias: str | None,
|
||||
*,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_device_address_book SET alias = :alias, update_date = :now "
|
||||
"WHERE mac_address = :mac_address AND target_mac = :target_mac",
|
||||
{
|
||||
"mac_address": mac_address,
|
||||
"target_mac": target_mac,
|
||||
"alias": alias,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_address_permission(
|
||||
self,
|
||||
mac_address: str,
|
||||
target_mac: str,
|
||||
has_permission: bool,
|
||||
*,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_device_address_book "
|
||||
"SET has_permission = :has_permission, update_date = :now "
|
||||
"WHERE mac_address = :mac_address AND target_mac = :target_mac",
|
||||
{
|
||||
"mac_address": mac_address,
|
||||
"target_mac": target_mac,
|
||||
"has_permission": has_permission,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_address_books_for_macs(self, mac_addresses: Sequence[str]) -> int:
|
||||
if not mac_addresses:
|
||||
return 0
|
||||
placeholders = ", ".join(f":mac_{index}" for index in range(len(mac_addresses)))
|
||||
params = {f"mac_{index}": mac for index, mac in enumerate(mac_addresses)}
|
||||
return await self.execute(
|
||||
f"DELETE FROM ai_device_address_book WHERE mac_address IN ({placeholders}) "
|
||||
f"OR target_mac IN ({placeholders})",
|
||||
params,
|
||||
)
|
||||
|
||||
async def count_ota(self, firmware_name: str | None = None) -> int:
|
||||
where = ""
|
||||
params: dict[str, Any] = {}
|
||||
if firmware_name is not None and firmware_name.strip():
|
||||
where = " WHERE firmware_name LIKE :firmware_name"
|
||||
params["firmware_name"] = f"%{firmware_name}%"
|
||||
return int(await self.scalar(f"SELECT COUNT(*) FROM ai_ota{where}", params) or 0)
|
||||
|
||||
async def list_ota(
|
||||
self,
|
||||
*,
|
||||
page: int,
|
||||
limit: int,
|
||||
firmware_name: str | None,
|
||||
order_fields: Sequence[str],
|
||||
ascending: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
where = ""
|
||||
params: dict[str, Any] = {"limit": limit, "offset": max(page - 1, 0) * limit}
|
||||
if firmware_name is not None and firmware_name.strip():
|
||||
where = " WHERE firmware_name LIKE :firmware_name"
|
||||
params["firmware_name"] = f"%{firmware_name}%"
|
||||
direction = "ASC" if ascending else "DESC"
|
||||
order_by = ", ".join(f"{field} {direction}" for field in order_fields)
|
||||
return await self.fetch_all(
|
||||
f"SELECT {OTA_COLUMNS} FROM ai_ota{where} ORDER BY {order_by} LIMIT :limit OFFSET :offset",
|
||||
params,
|
||||
)
|
||||
|
||||
async def get_ota(self, ota_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {OTA_COLUMNS} FROM ai_ota WHERE id = :id LIMIT 1",
|
||||
{"id": ota_id},
|
||||
)
|
||||
|
||||
async def get_first_ota_by_type(self, ota_type: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {OTA_COLUMNS} FROM ai_ota WHERE type = :type LIMIT 1",
|
||||
{"type": ota_type},
|
||||
)
|
||||
|
||||
async def get_latest_ota(self, ota_type: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
f"SELECT {OTA_COLUMNS} FROM ai_ota WHERE type = :type ORDER BY update_date DESC LIMIT 1",
|
||||
{"type": ota_type},
|
||||
)
|
||||
|
||||
async def count_duplicate_ota(self, *, ota_id: str, ota_type: str | None, version: str | None) -> int:
|
||||
return int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_ota WHERE type = :type AND version = :version AND id <> :id",
|
||||
{"id": ota_id, "type": ota_type, "version": version},
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
async def insert_ota(self, values: Mapping[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_ota "
|
||||
"(id, firmware_name, type, version, size, remark, firmware_path, sort, updater, update_date, creator, "
|
||||
"create_date) VALUES (:id, :firmware_name, :type, :version, :size, :remark, :firmware_path, :sort, "
|
||||
":updater, :update_date, :creator, :create_date)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update_ota(self, ota_id: str, values: Mapping[str, Any]) -> int:
|
||||
allowed = {
|
||||
"firmware_name",
|
||||
"type",
|
||||
"version",
|
||||
"size",
|
||||
"remark",
|
||||
"firmware_path",
|
||||
"sort",
|
||||
"updater",
|
||||
"update_date",
|
||||
"creator",
|
||||
"create_date",
|
||||
}
|
||||
selected = {key: value for key, value in values.items() if key in allowed and value is not None}
|
||||
if not selected:
|
||||
return 0
|
||||
assignments = ", ".join(f"{key} = :{key}" for key in selected)
|
||||
return await self.execute(
|
||||
f"UPDATE ai_ota SET {assignments} WHERE id = :id",
|
||||
{"id": ota_id, **selected},
|
||||
)
|
||||
|
||||
async def delete_ota(self, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
placeholders = ", ".join(f":id_{index}" for index in range(len(ids)))
|
||||
params = {f"id_{index}": value for index, value in enumerate(ids)}
|
||||
return await self.execute(f"DELETE FROM ai_ota WHERE id IN ({placeholders})", params)
|
||||
@@ -0,0 +1,344 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from sqlalchemy import bindparam, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class KnowledgeRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def dataset_page(
|
||||
self, user_id: int, name: str | None, offset: int, limit: int
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
where = (
|
||||
"WHERE creator=:creator AND (:name IS NULL OR :name='' OR name LIKE CONCAT('%', :name, '%'))"
|
||||
)
|
||||
params = {"creator": user_id, "name": name, "offset": offset, "limit": limit}
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_rag_dataset {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_rag_dataset {where} ORDER BY created_at DESC LIMIT :offset, :limit", # noqa: S608
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def get_dataset(self, identifier: str, *, for_update: bool = False) -> dict[str, Any] | None:
|
||||
suffix = " FOR UPDATE" if for_update and self.session.get_bind().dialect.name != "sqlite" else ""
|
||||
return await self.fetch_one(
|
||||
f"SELECT * FROM ai_rag_dataset WHERE dataset_id=:id OR id=:id LIMIT 1{suffix}", # noqa: S608
|
||||
{"id": identifier},
|
||||
)
|
||||
|
||||
async def datasets_by_ids(self, identifiers: Sequence[str]) -> list[dict[str, Any]]:
|
||||
if not identifiers:
|
||||
return []
|
||||
statement = text("SELECT * FROM ai_rag_dataset WHERE dataset_id IN :ids OR id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
result = await self.session.execute(statement, {"ids": list(identifiers)})
|
||||
return [dict(row) for row in result.mappings().all()]
|
||||
|
||||
async def duplicate_dataset_name(self, user_id: int, name: str, exclude_id: str | None = None) -> bool:
|
||||
return bool(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_rag_dataset WHERE creator=:creator AND name=:name "
|
||||
"AND (:exclude_id IS NULL OR id<>:exclude_id)",
|
||||
{"creator": user_id, "name": name, "exclude_id": exclude_id},
|
||||
)
|
||||
)
|
||||
|
||||
async def dataset_id_conflict(self, dataset_id: str, exclude_id: str) -> bool:
|
||||
return bool(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_rag_dataset WHERE dataset_id=:dataset_id AND id<>:exclude_id",
|
||||
{"dataset_id": dataset_id, "exclude_id": exclude_id},
|
||||
)
|
||||
)
|
||||
|
||||
async def insert_dataset(self, values: dict[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_rag_dataset "
|
||||
"(id,dataset_id,rag_model_id,tenant_id,name,avatar,description,embedding_model,permission,chunk_method,"
|
||||
"parser_config,chunk_count,document_count,token_num,status,creator,created_at,updater,updated_at) VALUES "
|
||||
"(:id,:dataset_id,:rag_model_id,:tenant_id,:name,:avatar,:description,:embedding_model,:permission,"
|
||||
":chunk_method,:parser_config,:chunk_count,:document_count,:token_num,:status,:creator,:created_at,"
|
||||
":updater,:updated_at)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update_dataset(self, values: dict[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_rag_dataset SET dataset_id=COALESCE(:dataset_id,dataset_id),"
|
||||
"rag_model_id=COALESCE(:rag_model_id,rag_model_id),name=COALESCE(:name,name),"
|
||||
"avatar=COALESCE(:avatar,avatar),description=COALESCE(:description,description),"
|
||||
"embedding_model=COALESCE(:embedding_model,embedding_model),"
|
||||
"permission=COALESCE(:permission,permission),chunk_method=COALESCE(:chunk_method,chunk_method),"
|
||||
"parser_config=COALESCE(:parser_config,parser_config),chunk_count=COALESCE(:chunk_count,chunk_count),"
|
||||
"token_num=COALESCE(:token_num,token_num),status=COALESCE(:status,status),"
|
||||
"creator=COALESCE(:creator,creator),created_at=COALESCE(:created_at,created_at),updater=:updater,"
|
||||
"updated_at=:updated_at WHERE id=:id",
|
||||
values,
|
||||
)
|
||||
|
||||
async def delete_dataset_local(self, row: dict[str, Any]) -> None:
|
||||
await self.execute("DELETE FROM ai_agent_plugin_mapping WHERE plugin_id=:id", {"id": row["id"]})
|
||||
await self.execute("DELETE FROM ai_rag_dataset WHERE id=:id", {"id": row["id"]})
|
||||
|
||||
async def rag_models(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, model_name, config_json FROM ai_model_config WHERE model_type='RAG' AND is_enabled=1 "
|
||||
"ORDER BY is_default DESC, create_date DESC"
|
||||
)
|
||||
|
||||
async def rag_config(self, model_id: str) -> dict[str, Any]:
|
||||
row = await self.fetch_one("SELECT config_json FROM ai_model_config WHERE id=:id", {"id": model_id})
|
||||
if row is None or row.get("config_json") is None:
|
||||
from app.core.errors import AppError
|
||||
|
||||
raise AppError(10164)
|
||||
raw = row["config_json"]
|
||||
if isinstance(raw, bytes):
|
||||
raw = raw.decode("utf-8")
|
||||
config = dict(raw) if isinstance(raw, dict) else dict(json.loads(str(raw)))
|
||||
config.setdefault("type", "ragflow")
|
||||
return config
|
||||
|
||||
async def documents_page(
|
||||
self,
|
||||
dataset_id: str,
|
||||
*,
|
||||
name: str | None,
|
||||
status: str | None,
|
||||
offset: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
where = (
|
||||
"WHERE dataset_id=:dataset_id "
|
||||
"AND (:name IS NULL OR :name='' OR name LIKE CONCAT('%', :name, '%')) "
|
||||
"AND (:status IS NULL OR :status='' OR status=:status)"
|
||||
)
|
||||
params = {"dataset_id": dataset_id, "name": name, "status": status, "offset": offset, "limit": limit}
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_rag_knowledge_document {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_rag_knowledge_document {where} " # noqa: S608
|
||||
"ORDER BY created_at DESC LIMIT :offset, :limit",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def all_documents(self, dataset_id: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT * FROM ai_rag_knowledge_document WHERE dataset_id=:dataset_id", {"dataset_id": dataset_id}
|
||||
)
|
||||
|
||||
async def documents_by_remote_ids(self, dataset_id: str, ids: Sequence[str]) -> list[dict[str, Any]]:
|
||||
if not ids:
|
||||
return []
|
||||
statement = text(
|
||||
"SELECT * FROM ai_rag_knowledge_document WHERE dataset_id=:dataset_id AND document_id IN :ids"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
result = await self.session.execute(statement, {"dataset_id": dataset_id, "ids": list(ids)})
|
||||
return [dict(row) for row in result.mappings().all()]
|
||||
|
||||
async def upsert_document(self, dataset_id: str, remote: dict[str, Any], *, creator: int | None = None) -> bool:
|
||||
document_id = str(remote.get("id") or remote.get("document_id") or "")
|
||||
existing = await self.fetch_one(
|
||||
"SELECT id,created_at FROM ai_rag_knowledge_document WHERE document_id=:id", {"id": document_id}
|
||||
)
|
||||
name = remote.get("name")
|
||||
size = remote.get("size")
|
||||
if size is None:
|
||||
size = remote.get("file_size")
|
||||
meta_fields = remote.get("meta_fields")
|
||||
if meta_fields is None:
|
||||
meta_fields = remote.get("meta")
|
||||
error = remote.get("progress_msg")
|
||||
if error is None:
|
||||
error = remote.get("error")
|
||||
synced_at = _shanghai_now_naive()
|
||||
created_at = remote.get("created_at")
|
||||
if not isinstance(created_at, datetime):
|
||||
created_at = _millis_date(remote.get("create_time"))
|
||||
updated_at = remote.get("updated_at")
|
||||
if not isinstance(updated_at, datetime):
|
||||
updated_at = _millis_date(remote.get("update_time"))
|
||||
values = {
|
||||
"id": existing["id"] if existing else uuid.uuid4().hex,
|
||||
"dataset_id": remote.get("dataset_id") or dataset_id,
|
||||
"document_id": document_id,
|
||||
"name": name,
|
||||
"size": size,
|
||||
"type": _file_type(str(name or "")),
|
||||
"chunk_method": remote.get("chunk_method"),
|
||||
"parser_config": _json_dump(remote.get("parser_config")),
|
||||
"status": _remote_status(remote.get("status")),
|
||||
"run": remote.get("run"),
|
||||
"progress": remote.get("progress"),
|
||||
"thumbnail": remote.get("thumbnail"),
|
||||
"process_duration": remote.get("process_duration"),
|
||||
"meta_fields": _json_dump(meta_fields),
|
||||
"source_type": remote.get("source_type"),
|
||||
"error": error,
|
||||
"chunk_count": remote.get("chunk_count") or 0,
|
||||
"token_count": remote.get("token_count") or 0,
|
||||
"enabled": 1,
|
||||
"creator": creator,
|
||||
"created_at": existing.get("created_at") if existing else (created_at or synced_at),
|
||||
"updated_at": updated_at or synced_at,
|
||||
"synced_at": synced_at,
|
||||
}
|
||||
if existing:
|
||||
await self.execute(
|
||||
"UPDATE ai_rag_knowledge_document SET dataset_id=:dataset_id,document_id=:document_id,name=:name,"
|
||||
"size=:size,type=:type,chunk_method=:chunk_method,parser_config=:parser_config,status=:status,run=:run,"
|
||||
"progress=:progress,thumbnail=:thumbnail,process_duration=:process_duration,meta_fields=:meta_fields,"
|
||||
"source_type=:source_type,error=:error,chunk_count=:chunk_count,token_count=:token_count,enabled=:enabled,"
|
||||
"updated_at=:updated_at,last_sync_at=:synced_at WHERE id=:id",
|
||||
values,
|
||||
)
|
||||
return False
|
||||
await self.execute(
|
||||
"INSERT INTO ai_rag_knowledge_document "
|
||||
"(id,dataset_id,document_id,name,size,type,chunk_method,parser_config,status,run,progress,thumbnail,"
|
||||
"process_duration,meta_fields,source_type,error,chunk_count,token_count,enabled,creator,created_at,updated_at,"
|
||||
"last_sync_at) VALUES (:id,:dataset_id,:document_id,:name,:size,:type,:chunk_method,:parser_config,:status,"
|
||||
":run,:progress,:thumbnail,:process_duration,:meta_fields,:source_type,:error,:chunk_count,:token_count,"
|
||||
":enabled,:creator,COALESCE(:created_at,:synced_at),COALESCE(:updated_at,:synced_at),:synced_at)",
|
||||
values,
|
||||
)
|
||||
return True
|
||||
|
||||
async def update_stats(self, dataset_id: str, docs: int, chunks: int, tokens: int) -> None:
|
||||
await self.execute(
|
||||
"UPDATE ai_rag_dataset SET document_count=document_count+:docs,chunk_count=chunk_count+:chunks,"
|
||||
"token_num=token_num+:tokens,updated_at=:now WHERE dataset_id=:dataset_id",
|
||||
{
|
||||
"dataset_id": dataset_id,
|
||||
"docs": docs,
|
||||
"chunks": chunks,
|
||||
"tokens": tokens,
|
||||
"now": _shanghai_now_naive(),
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_document_shadows(self, dataset_id: str, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
statement = text(
|
||||
"DELETE FROM ai_rag_knowledge_document WHERE dataset_id=:dataset_id AND document_id IN :ids"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
result = await self.session.execute(statement, {"dataset_id": dataset_id, "ids": list(ids)})
|
||||
return int(getattr(result, "rowcount", 0) or 0)
|
||||
|
||||
async def mark_documents_running(self, dataset_id: str, ids: Sequence[str], now: datetime) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
statement = text(
|
||||
"UPDATE ai_rag_knowledge_document SET run='RUNNING',status='1',updated_at=:now "
|
||||
"WHERE dataset_id=:dataset_id AND document_id IN :ids"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
result = await self.session.execute(statement, {"dataset_id": dataset_id, "ids": list(ids), "now": now})
|
||||
return int(getattr(result, "rowcount", 0) or 0)
|
||||
|
||||
async def mark_document_remote_deleted(self, document_id: str, now: datetime) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_rag_knowledge_document SET run='CANCEL',error=:error,updated_at=:now,last_sync_at=:now "
|
||||
"WHERE document_id=:document_id",
|
||||
{
|
||||
"document_id": document_id,
|
||||
"error": "文档在远程服务中已被删除",
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def sync_running_document(
|
||||
self,
|
||||
dataset_id: str,
|
||||
document_id: str,
|
||||
remote: dict[str, Any],
|
||||
now: datetime,
|
||||
) -> int:
|
||||
"""Update exactly the columns touched by Java's status-sync helper."""
|
||||
updated_at = _millis_date(remote.get("update_time")) or now
|
||||
meta_fields = remote.get("meta_fields")
|
||||
assignments = (
|
||||
"status=:status,run=:run,progress=:progress,chunk_count=:chunk_count,token_count=:token_count,"
|
||||
"error=:error,process_duration=:process_duration,thumbnail=:thumbnail,updated_at=:updated_at,"
|
||||
"last_sync_at=:now"
|
||||
)
|
||||
if meta_fields is not None:
|
||||
assignments += ",meta_fields=:meta_fields"
|
||||
return await self.execute(
|
||||
f"UPDATE ai_rag_knowledge_document SET {assignments} " # noqa: S608
|
||||
"WHERE document_id=:document_id AND dataset_id=:dataset_id",
|
||||
{
|
||||
"dataset_id": dataset_id,
|
||||
"document_id": document_id,
|
||||
"status": remote.get("status"),
|
||||
"run": remote.get("run"),
|
||||
"progress": remote.get("progress"),
|
||||
"chunk_count": remote.get("chunk_count"),
|
||||
"token_count": remote.get("token_count"),
|
||||
"error": remote.get("progress_msg") if remote.get("progress_msg") is not None else remote.get("error"),
|
||||
"process_duration": remote.get("process_duration"),
|
||||
"thumbnail": remote.get("thumbnail"),
|
||||
"meta_fields": _json_dump(meta_fields),
|
||||
"updated_at": updated_at,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def running_documents(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all("SELECT * FROM ai_rag_knowledge_document WHERE run='RUNNING' AND status='1'")
|
||||
|
||||
|
||||
def _json_dump(value: Any) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
|
||||
def _millis_date(value: Any) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
timezone = ZoneInfo(get_settings().timezone)
|
||||
return datetime.fromtimestamp(float(value) / 1000, timezone).replace(tzinfo=None)
|
||||
except (TypeError, ValueError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
def _file_type(name: str) -> str:
|
||||
last_dot = name.rfind(".")
|
||||
if last_dot <= 0 or last_dot == len(name) - 1:
|
||||
return "unknown"
|
||||
extension = name.rsplit(".", 1)[1].lower()
|
||||
if extension in {"pdf", "doc", "docx", "txt", "md", "mdx"}:
|
||||
return "document"
|
||||
if extension in {"csv", "xls", "xlsx"}:
|
||||
return "spreadsheet"
|
||||
if extension in {"ppt", "pptx"}:
|
||||
return "presentation"
|
||||
return extension
|
||||
|
||||
|
||||
def _remote_status(value: Any) -> str:
|
||||
if value is None or (isinstance(value, str) and not value.strip()):
|
||||
return "1"
|
||||
return str(value)
|
||||
|
||||
|
||||
def _shanghai_now_naive() -> datetime:
|
||||
return datetime.now(ZoneInfo(get_settings().timezone)).replace(tzinfo=None)
|
||||
@@ -0,0 +1,243 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import bindparam, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class ModelRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def list_model_names(self, model_type: str, model_name: str | None) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, model_name FROM ai_model_config "
|
||||
"WHERE model_type = :model_type AND is_enabled = 1 "
|
||||
"AND (:model_name IS NULL OR :model_name = '' OR model_name LIKE CONCAT('%', :model_name, '%')) "
|
||||
"ORDER BY sort ASC",
|
||||
{"model_type": model_type, "model_name": model_name},
|
||||
)
|
||||
|
||||
async def list_llm_names(self, model_name: str | None) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, model_name, config_json FROM ai_model_config "
|
||||
"WHERE model_type = 'llm' AND is_enabled = 1 "
|
||||
"AND (:model_name IS NULL OR :model_name = '' OR model_name LIKE CONCAT('%', :model_name, '%'))",
|
||||
{"model_name": model_name},
|
||||
)
|
||||
|
||||
async def list_providers_by_type(self, model_type: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT * FROM ai_model_provider WHERE model_type = :model_type ORDER BY sort ASC",
|
||||
{"model_type": model_type or ""},
|
||||
)
|
||||
|
||||
async def list_providers(
|
||||
self,
|
||||
*,
|
||||
model_type: str | None,
|
||||
name: str | None,
|
||||
offset: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
where = (
|
||||
"WHERE (:model_type IS NULL OR :model_type = '' OR model_type = :model_type) "
|
||||
"AND (:name IS NULL OR :name = '' OR name LIKE CONCAT('%', :name, '%') "
|
||||
"OR provider_code LIKE CONCAT('%', :name, '%'))"
|
||||
)
|
||||
params = {"model_type": model_type, "name": name, "offset": offset, "limit": limit}
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_model_provider {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_model_provider {where} " # noqa: S608
|
||||
"ORDER BY model_type ASC, sort ASC LIMIT :offset, :limit",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def list_model_configs(
|
||||
self,
|
||||
*,
|
||||
model_type: str,
|
||||
model_name: str | None,
|
||||
offset: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
where = (
|
||||
"WHERE model_type = :model_type AND "
|
||||
"(:model_name IS NULL OR :model_name = '' OR model_name LIKE CONCAT('%', :model_name, '%'))"
|
||||
)
|
||||
params = {"model_type": model_type, "model_name": model_name, "offset": offset, "limit": limit}
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_model_config {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_model_config {where} " # noqa: S608
|
||||
"ORDER BY is_enabled DESC, sort ASC LIMIT :offset, :limit",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def get_provider(self, model_type: str, provider_code: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT * FROM ai_model_provider WHERE model_type = :model_type AND provider_code = :provider_code LIMIT 1",
|
||||
{"model_type": model_type or "", "provider_code": provider_code or ""},
|
||||
)
|
||||
|
||||
async def get_model(self, model_id: str, *, for_update: bool = False) -> dict[str, Any] | None:
|
||||
suffix = " FOR UPDATE" if for_update and self.session.get_bind().dialect.name != "sqlite" else ""
|
||||
return await self.fetch_one(
|
||||
f"SELECT * FROM ai_model_config WHERE id = :id LIMIT 1{suffix}", # noqa: S608
|
||||
{"id": model_id},
|
||||
)
|
||||
|
||||
async def insert_model(self, values: dict[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_model_config "
|
||||
"(id, model_type, model_code, model_name, is_default, is_enabled, config_json, doc_link, remark, sort) "
|
||||
"VALUES (:id, :model_type, :model_code, :model_name, :is_default, COALESCE(:is_enabled, 0), "
|
||||
":config_json, :doc_link, :remark, COALESCE(:sort, 0))",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update_model(self, values: dict[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_model_config SET model_type=:model_type, model_code=:model_code, "
|
||||
"model_name=COALESCE(:model_name, model_name), is_default=:is_default, "
|
||||
"is_enabled=COALESCE(:is_enabled, is_enabled), config_json=:config_json, doc_link=:doc_link, "
|
||||
"remark=COALESCE(:remark, remark), sort=COALESCE(:sort, sort) WHERE id=:id",
|
||||
values,
|
||||
)
|
||||
|
||||
async def delete_model(self, model_id: str) -> int:
|
||||
return await self.execute("DELETE FROM ai_model_config WHERE id = :id", {"id": model_id})
|
||||
|
||||
async def model_agent_references(self, model_id: str) -> list[str]:
|
||||
rows = await self.fetch_all(
|
||||
"SELECT agent_name FROM ai_agent WHERE vad_model_id=:id OR asr_model_id=:id OR llm_model_id=:id "
|
||||
"OR tts_model_id=:id OR mem_model_id=:id OR vllm_model_id=:id OR intent_model_id=:id",
|
||||
{"id": model_id},
|
||||
)
|
||||
return [str(row.get("agent_name") or "") for row in rows]
|
||||
|
||||
async def intent_reference_count(self, model_id: str) -> int:
|
||||
return int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_model_config WHERE model_type='Intent' AND CAST(config_json AS CHAR) LIKE "
|
||||
"CONCAT('%', :id, '%')",
|
||||
{"id": model_id},
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
async def set_models_default(self, model_type: str, value: int) -> None:
|
||||
await self.execute(
|
||||
"UPDATE ai_model_config SET is_default=:value WHERE model_type=:model_type",
|
||||
{"value": value, "model_type": model_type},
|
||||
)
|
||||
|
||||
async def set_model_enabled(self, model_id: str, status: int) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_model_config SET is_enabled=:status WHERE id=:id",
|
||||
{"status": status, "id": model_id},
|
||||
)
|
||||
|
||||
async def update_default_template_models(self, model_type: str, model_id: str) -> None:
|
||||
columns = {
|
||||
"ASR": ("asr_model_id",),
|
||||
"VAD": ("vad_model_id",),
|
||||
"LLM": ("llm_model_id",),
|
||||
"TTS": ("tts_model_id", "tts_voice_id"),
|
||||
"VLLM": ("vllm_model_id",),
|
||||
"MEMORY": ("mem_model_id",),
|
||||
"INTENT": ("intent_model_id",),
|
||||
}.get(model_type.upper())
|
||||
if not columns:
|
||||
return
|
||||
if columns == ("tts_model_id", "tts_voice_id"):
|
||||
await self.execute(
|
||||
"UPDATE ai_agent_template SET tts_model_id=:id, tts_voice_id=NULL WHERE sort >= 0",
|
||||
{"id": model_id},
|
||||
)
|
||||
else:
|
||||
column = columns[0]
|
||||
await self.session.execute(
|
||||
text(f"UPDATE ai_agent_template SET {column}=:id WHERE sort >= 0"), # noqa: S608
|
||||
{"id": model_id},
|
||||
)
|
||||
|
||||
async def insert_provider(self, values: dict[str, Any]) -> None:
|
||||
if self.session.get_bind().dialect.name == "sqlite":
|
||||
statement = (
|
||||
"INSERT INTO ai_model_provider "
|
||||
"(id, model_type, provider_code, name, fields, sort, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :model_type, :provider_code, :name, :fields, :sort, :creator, :now, :updater, :now)"
|
||||
)
|
||||
else:
|
||||
statement = (
|
||||
"INSERT INTO ai_model_provider "
|
||||
"(id, model_type, provider_code, name, fields, sort, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :model_type, :provider_code, :name, CAST(:fields AS JSON), :sort, :creator, :now, "
|
||||
":updater, :now)"
|
||||
)
|
||||
await self.execute(statement, values)
|
||||
|
||||
async def update_provider(self, values: dict[str, Any]) -> int:
|
||||
if self.session.get_bind().dialect.name == "sqlite":
|
||||
statement = (
|
||||
"UPDATE ai_model_provider SET model_type=:model_type, provider_code=:provider_code, name=:name, "
|
||||
"fields=:fields, sort=:sort, updater=:updater, update_date=:now WHERE id=:id"
|
||||
)
|
||||
else:
|
||||
statement = (
|
||||
"UPDATE ai_model_provider SET model_type=:model_type, provider_code=:provider_code, name=:name, "
|
||||
"fields=CAST(:fields AS JSON), sort=:sort, updater=:updater, update_date=:now WHERE id=:id"
|
||||
)
|
||||
return await self.execute(statement, values)
|
||||
|
||||
async def delete_providers(self, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
statement = text("DELETE FROM ai_model_provider WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
result = await self.session.execute(statement, {"ids": list(ids)})
|
||||
return int(getattr(result, "rowcount", 0) or 0)
|
||||
|
||||
async def list_plugins_for_user(self, user_id: int) -> list[dict[str, Any]]:
|
||||
providers = await self.fetch_all("SELECT * FROM ai_model_provider WHERE model_type='Plugin'")
|
||||
datasets = await self.fetch_all(
|
||||
"SELECT id, name, created_at, updated_at FROM ai_rag_dataset WHERE creator=:creator AND status=1",
|
||||
{"creator": user_id},
|
||||
)
|
||||
providers.extend(
|
||||
{
|
||||
"id": row["id"],
|
||||
"model_type": "Rag",
|
||||
"name": f"[知识库]{row['name']}",
|
||||
"provider_code": "ragflow",
|
||||
"fields": "[]",
|
||||
"sort": 0,
|
||||
"create_date": row.get("created_at"),
|
||||
"update_date": row.get("updated_at"),
|
||||
"creator": 0,
|
||||
"updater": 0,
|
||||
}
|
||||
for row in datasets
|
||||
)
|
||||
return providers
|
||||
|
||||
|
||||
def parse_json_object(value: Any) -> dict[str, Any] | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, dict):
|
||||
return dict(value)
|
||||
if isinstance(value, bytes):
|
||||
value = value.decode("utf-8")
|
||||
if isinstance(value, str):
|
||||
parsed = json.loads(value)
|
||||
return dict(parsed) if isinstance(parsed, dict) else None
|
||||
return None
|
||||
@@ -0,0 +1,153 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class SecurityRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def get_param_value(self, code: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT param_value FROM sys_params WHERE param_code = :code LIMIT 1",
|
||||
{"code": code},
|
||||
)
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def get_mobile_area_items(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT d.dict_label AS name, d.dict_value AS `key` "
|
||||
"FROM sys_dict_data d "
|
||||
"LEFT JOIN sys_dict_type t ON d.dict_type_id = t.id "
|
||||
"WHERE t.dict_type = :dict_type ORDER BY d.sort ASC",
|
||||
{"dict_type": "MOBILE_AREA"},
|
||||
)
|
||||
|
||||
async def get_user_by_username(self, username: str | None) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, username, password, super_admin, status, creator, create_date, updater, update_date "
|
||||
"FROM sys_user WHERE username = :username LIMIT 1",
|
||||
{"username": username},
|
||||
)
|
||||
|
||||
async def get_user_by_id(self, user_id: int) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, username, password, super_admin, status, creator, create_date, updater, update_date "
|
||||
"FROM sys_user WHERE id = :id LIMIT 1",
|
||||
{"id": user_id},
|
||||
)
|
||||
|
||||
async def count_users(self) -> int:
|
||||
return int(await self.scalar("SELECT COUNT(*) FROM sys_user") or 0)
|
||||
|
||||
async def insert_user(
|
||||
self,
|
||||
*,
|
||||
user_id: int,
|
||||
username: str | None,
|
||||
password: str,
|
||||
super_admin: int,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO sys_user "
|
||||
"(id, username, password, super_admin, status, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :username, :password, :super_admin, 1, NULL, :now, NULL, :now)",
|
||||
{
|
||||
"id": user_id,
|
||||
"username": username,
|
||||
"password": password,
|
||||
"super_admin": super_admin,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def get_token_by_user_id(self, user_id: int, *, for_update: bool = False) -> dict[str, Any] | None:
|
||||
sql = (
|
||||
"SELECT id, user_id, token, expire_date, update_date, create_date "
|
||||
"FROM sys_user_token WHERE user_id = :user_id LIMIT 1 FOR UPDATE"
|
||||
if for_update and self._supports_for_update()
|
||||
else "SELECT id, user_id, token, expire_date, update_date, create_date "
|
||||
"FROM sys_user_token WHERE user_id = :user_id LIMIT 1"
|
||||
)
|
||||
return await self.fetch_one(
|
||||
sql,
|
||||
{"user_id": user_id},
|
||||
)
|
||||
|
||||
async def insert_token(
|
||||
self,
|
||||
*,
|
||||
token_id: int,
|
||||
user_id: int,
|
||||
token: str,
|
||||
now: datetime,
|
||||
expire_date: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO sys_user_token (id, user_id, token, expire_date, update_date, create_date) "
|
||||
"VALUES (:id, :user_id, :token, :expire_date, :now, :now)",
|
||||
{
|
||||
"id": token_id,
|
||||
"user_id": user_id,
|
||||
"token": token,
|
||||
"expire_date": expire_date,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_token(self, *, token_id: int, token: str, now: datetime, expire_date: datetime) -> None:
|
||||
await self.execute(
|
||||
"UPDATE sys_user_token SET token = :token, expire_date = :expire_date, update_date = :now "
|
||||
"WHERE id = :id",
|
||||
{"id": token_id, "token": token, "expire_date": expire_date, "now": now},
|
||||
)
|
||||
|
||||
async def update_password(
|
||||
self,
|
||||
user_id: int,
|
||||
password_hash: str,
|
||||
now: datetime,
|
||||
*,
|
||||
preserve_audit_fields: bool = False,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_user SET password = :password, "
|
||||
"update_date = CASE WHEN :preserve_audit = 1 THEN update_date ELSE :now END WHERE id = :id",
|
||||
{
|
||||
"id": user_id,
|
||||
"password": password_hash,
|
||||
"now": now,
|
||||
"preserve_audit": int(preserve_audit_fields),
|
||||
},
|
||||
)
|
||||
|
||||
async def expire_user_token(self, user_id: int, expire_date: datetime) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_user_token SET expire_date = :expire_date WHERE user_id = :user_id",
|
||||
{"user_id": user_id, "expire_date": expire_date},
|
||||
)
|
||||
|
||||
def _supports_for_update(self) -> bool:
|
||||
bind = self.session.get_bind()
|
||||
return bind.dialect.name != "sqlite"
|
||||
|
||||
|
||||
async def raw_user_token(session: AsyncSession, token: str) -> dict[str, Any] | None:
|
||||
result = await session.execute(
|
||||
text(
|
||||
"SELECT t.id AS token_id, t.user_id, t.token, t.expire_date, "
|
||||
"u.username, u.super_admin, u.status "
|
||||
"FROM sys_user_token t JOIN sys_user u ON u.id = t.user_id "
|
||||
"WHERE t.token = :token LIMIT 1"
|
||||
),
|
||||
{"token": token},
|
||||
)
|
||||
row = result.mappings().first()
|
||||
return dict(row) if row is not None else None
|
||||
@@ -0,0 +1,494 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import bindparam, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class SysRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def page_users(self, *, mobile: str | None, page: int, limit: int) -> tuple[list[dict[str, Any]], int]:
|
||||
pattern = f"%{mobile}%" if mobile else None
|
||||
params = {"mobile": pattern, "offset": (page - 1) * limit, "limit": limit}
|
||||
total = int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM sys_user WHERE (:mobile IS NULL OR username LIKE :mobile)",
|
||||
params,
|
||||
)
|
||||
or 0
|
||||
)
|
||||
rows = await self.fetch_all(
|
||||
"SELECT u.id, u.username, u.status, u.create_date, "
|
||||
"(SELECT COUNT(*) FROM ai_device d WHERE d.user_id = u.id) AS device_count "
|
||||
"FROM sys_user u WHERE (:mobile IS NULL OR u.username LIKE :mobile) "
|
||||
"ORDER BY u.id ASC LIMIT :limit OFFSET :offset",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def reset_user_password(
|
||||
self,
|
||||
user_id: int,
|
||||
password_hash: str,
|
||||
updater: int,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_user SET password = :password, updater = :updater, update_date = :now WHERE id = :id",
|
||||
{"id": user_id, "password": password_hash, "updater": updater, "now": now},
|
||||
)
|
||||
|
||||
async def change_user_status(self, status: int, user_ids: list[int], updater: int, now: datetime) -> int:
|
||||
statement = text(
|
||||
"UPDATE sys_user SET status = :status, updater = :updater, update_date = :now WHERE id IN :ids"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
return await self.execute(
|
||||
statement,
|
||||
{"status": status, "updater": updater, "now": now, "ids": user_ids},
|
||||
)
|
||||
|
||||
async def delete_user_cascade(self, user_id: int) -> None:
|
||||
agent_rows = await self.fetch_all("SELECT id FROM ai_agent WHERE user_id = :user_id", {"user_id": user_id})
|
||||
agent_ids = [str(row["id"]) for row in agent_rows]
|
||||
await self.execute("DELETE FROM sys_user WHERE id = :id", {"id": user_id})
|
||||
await self.execute("DELETE FROM ai_device WHERE user_id = :id", {"id": user_id})
|
||||
for agent_id in agent_ids:
|
||||
audio_rows = await self.fetch_all(
|
||||
"SELECT DISTINCT audio_id FROM ai_agent_chat_history "
|
||||
"WHERE agent_id = :agent_id AND audio_id IS NOT NULL",
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
audio_ids = [str(row["audio_id"]) for row in audio_rows]
|
||||
if audio_ids:
|
||||
statement = text("DELETE FROM ai_agent_chat_audio WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
await self.execute(statement, {"ids": audio_ids})
|
||||
for table in (
|
||||
"ai_agent_chat_history",
|
||||
"ai_agent_plugin_mapping",
|
||||
"ai_agent_context_provider",
|
||||
"ai_agent_correct_word_mapping",
|
||||
"ai_agent_tag_relation",
|
||||
"ai_agent_snapshot",
|
||||
):
|
||||
# Table names are a closed list mirroring AgentServiceImpl.deleteAgent.
|
||||
await self.execute(
|
||||
f"DELETE FROM {table} WHERE agent_id = :agent_id", # noqa: S608 - closed table list above
|
||||
{"agent_id": agent_id},
|
||||
)
|
||||
await self.execute("DELETE FROM ai_device WHERE agent_id = :agent_id", {"agent_id": agent_id})
|
||||
await self.execute("DELETE FROM ai_agent WHERE id = :agent_id", {"agent_id": agent_id})
|
||||
|
||||
async def page_devices(
|
||||
self,
|
||||
*,
|
||||
keywords: str | None,
|
||||
page: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
pattern = f"%{keywords}%" if keywords else None
|
||||
params = {"keywords": pattern, "offset": (page - 1) * limit, "limit": limit}
|
||||
total = int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_device WHERE (:keywords IS NULL OR alias LIKE :keywords)",
|
||||
params,
|
||||
)
|
||||
or 0
|
||||
)
|
||||
rows = await self.fetch_all(
|
||||
"SELECT d.id, d.user_id, d.mac_address, d.last_connected_at, d.auto_update, d.board, d.alias, "
|
||||
"d.agent_id, d.app_version, d.sort, d.create_date, d.update_date, u.username AS bind_user_name "
|
||||
"FROM ai_device d LEFT JOIN sys_user u ON u.id = d.user_id "
|
||||
"WHERE (:keywords IS NULL OR d.alias LIKE :keywords) "
|
||||
"ORDER BY d.mac_address ASC LIMIT :limit OFFSET :offset",
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def page_params(
|
||||
self,
|
||||
*,
|
||||
param_code: str | None,
|
||||
page: int,
|
||||
limit: int,
|
||||
order_field: str | None,
|
||||
order: str | None,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
pattern = f"%{param_code}%" if param_code else None
|
||||
params = {"pattern": pattern, "offset": (page - 1) * limit, "limit": limit}
|
||||
where = "param_type = 1 AND (:pattern IS NULL OR param_code LIKE :pattern OR remark LIKE :pattern)"
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM sys_params WHERE {where}", params) or 0) # noqa: S608
|
||||
allowed = {
|
||||
"id": "id",
|
||||
"paramCode": "param_code",
|
||||
"paramValue": "param_value",
|
||||
"valueType": "value_type",
|
||||
"createDate": "create_date",
|
||||
"updateDate": "update_date",
|
||||
}
|
||||
order_column = allowed.get(order_field or "")
|
||||
order_clause = ""
|
||||
if order_column is not None:
|
||||
direction = "ASC" if (order or "").lower() == "asc" else "DESC"
|
||||
order_clause = f" ORDER BY {order_column} {direction}"
|
||||
sql = (
|
||||
"SELECT id, param_code, param_value, value_type, remark, create_date, update_date " # noqa: S608
|
||||
f"FROM sys_params WHERE {where}{order_clause} LIMIT :limit OFFSET :offset"
|
||||
)
|
||||
return await self.fetch_all(sql, params), total # noqa: S608
|
||||
|
||||
async def list_config_params(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, param_code, param_value, value_type, remark, create_date, update_date "
|
||||
"FROM sys_params WHERE param_type = 1"
|
||||
)
|
||||
|
||||
async def get_param(self, param_id: int) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, param_code, param_value, value_type, remark, create_date, update_date "
|
||||
"FROM sys_params WHERE id = :id",
|
||||
{"id": param_id},
|
||||
)
|
||||
|
||||
async def get_param_value(self, code: str) -> str | None:
|
||||
value = await self.scalar("SELECT param_value FROM sys_params WHERE param_code = :code", {"code": code})
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def insert_param(
|
||||
self,
|
||||
*,
|
||||
param_id: int,
|
||||
param_code: str,
|
||||
param_value: str,
|
||||
value_type: str,
|
||||
remark: str | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO sys_params "
|
||||
"(id, param_code, param_value, value_type, param_type, remark, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :code, :value, :value_type, 1, :remark, :user_id, :now, :user_id, :now)",
|
||||
{
|
||||
"id": param_id,
|
||||
"code": param_code,
|
||||
"value": param_value,
|
||||
"value_type": value_type,
|
||||
"remark": remark,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_param(
|
||||
self,
|
||||
*,
|
||||
param_id: int,
|
||||
param_code: str,
|
||||
param_value: str,
|
||||
value_type: str,
|
||||
remark: str | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_params SET param_code = :code, param_value = :value, value_type = :value_type, "
|
||||
"remark = CASE WHEN :has_remark = 1 THEN :remark ELSE remark END, updater = :user_id, update_date = :now "
|
||||
"WHERE id = :id",
|
||||
{
|
||||
"id": param_id,
|
||||
"code": param_code,
|
||||
"value": param_value,
|
||||
"value_type": value_type,
|
||||
"has_remark": int(remark is not None),
|
||||
"remark": remark,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_param_value_by_code(self, code: str, value: str, user_id: int, now: datetime) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_params SET param_value = :value, updater = :user_id, update_date = :now "
|
||||
"WHERE param_code = :code",
|
||||
{"code": code, "value": value, "user_id": user_id, "now": now},
|
||||
)
|
||||
|
||||
async def param_codes_for_ids(self, ids: list[int]) -> list[str]:
|
||||
statement = text("SELECT param_code FROM sys_params WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
rows = await self.fetch_all(statement, {"ids": ids})
|
||||
return [str(row["param_code"]) for row in rows]
|
||||
|
||||
async def delete_params(self, ids: list[int]) -> int:
|
||||
statement = text("DELETE FROM sys_params WHERE id IN :ids").bindparams(bindparam("ids", expanding=True))
|
||||
return await self.execute(statement, {"ids": ids})
|
||||
|
||||
async def delete_plugin_mapping_by_plugin_id(self, plugin_id: str) -> int:
|
||||
return await self.execute(
|
||||
"DELETE FROM ai_agent_plugin_mapping WHERE plugin_id = :plugin_id",
|
||||
{"plugin_id": plugin_id},
|
||||
)
|
||||
|
||||
async def page_dict_types(
|
||||
self,
|
||||
*,
|
||||
dict_type: str | None,
|
||||
dict_name: str | None,
|
||||
page: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
params = {
|
||||
"dict_type": f"%{dict_type}%" if dict_type else None,
|
||||
"dict_name": f"%{dict_name}%" if dict_name else None,
|
||||
"offset": (page - 1) * limit,
|
||||
"limit": limit,
|
||||
}
|
||||
where = (
|
||||
"(:dict_type IS NULL OR t.dict_type LIKE :dict_type) "
|
||||
"AND (:dict_name IS NULL OR t.dict_name LIKE :dict_name)"
|
||||
)
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM sys_dict_type t WHERE {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
"SELECT t.id, t.dict_type, t.dict_name, t.remark, t.sort, t.creator, t.create_date, t.updater, " # noqa: S608
|
||||
"t.update_date, creator.username AS creator_name, updater.username AS updater_name "
|
||||
"FROM sys_dict_type t LEFT JOIN sys_user creator ON creator.id = t.creator "
|
||||
"LEFT JOIN sys_user updater ON updater.id = t.updater "
|
||||
f"WHERE {where} ORDER BY t.sort ASC LIMIT :limit OFFSET :offset", # noqa: S608
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def get_dict_type(self, type_id: int) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, dict_type, dict_name, remark, sort, creator, create_date, updater, update_date "
|
||||
"FROM sys_dict_type WHERE id = :id",
|
||||
{"id": type_id},
|
||||
)
|
||||
|
||||
async def dict_type_exists(self, dict_type: str | None, *, exclude_id: int | None = None) -> bool:
|
||||
count = await self.scalar(
|
||||
"SELECT COUNT(*) FROM sys_dict_type WHERE dict_type = :dict_type "
|
||||
"AND (:exclude_id IS NULL OR id <> :exclude_id)",
|
||||
{"dict_type": dict_type, "exclude_id": exclude_id},
|
||||
)
|
||||
return int(count or 0) > 0
|
||||
|
||||
async def insert_dict_type(
|
||||
self,
|
||||
*,
|
||||
type_id: int,
|
||||
dict_type: str | None,
|
||||
dict_name: str | None,
|
||||
remark: str | None,
|
||||
sort: int | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO sys_dict_type "
|
||||
"(id, dict_type, dict_name, remark, sort, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :dict_type, :dict_name, :remark, :sort, :user_id, :now, :user_id, :now)",
|
||||
{
|
||||
"id": type_id,
|
||||
"dict_type": dict_type,
|
||||
"dict_name": dict_name,
|
||||
"remark": remark,
|
||||
"sort": sort,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_dict_type(
|
||||
self,
|
||||
*,
|
||||
type_id: int | None,
|
||||
dict_type: str | None,
|
||||
dict_name: str | None,
|
||||
remark: str | None,
|
||||
sort: int | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_dict_type SET "
|
||||
"dict_type = CASE WHEN :has_dict_type = 1 THEN :dict_type ELSE dict_type END, "
|
||||
"dict_name = CASE WHEN :has_dict_name = 1 THEN :dict_name ELSE dict_name END, "
|
||||
"remark = CASE WHEN :has_remark = 1 THEN :remark ELSE remark END, "
|
||||
"sort = CASE WHEN :has_sort = 1 THEN :sort ELSE sort END, updater = :user_id, update_date = :now "
|
||||
"WHERE id = :id",
|
||||
{
|
||||
"id": type_id,
|
||||
"has_dict_type": int(dict_type is not None),
|
||||
"dict_type": dict_type,
|
||||
"has_dict_name": int(dict_name is not None),
|
||||
"dict_name": dict_name,
|
||||
"has_remark": int(remark is not None),
|
||||
"remark": remark,
|
||||
"has_sort": int(sort is not None),
|
||||
"sort": sort,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def delete_dict_types(self, ids: list[int]) -> None:
|
||||
statement_data = text("DELETE FROM sys_dict_data WHERE dict_type_id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
statement_types = text("DELETE FROM sys_dict_type WHERE id IN :ids").bindparams(
|
||||
bindparam("ids", expanding=True)
|
||||
)
|
||||
await self.execute(statement_data, {"ids": ids})
|
||||
await self.execute(statement_types, {"ids": ids})
|
||||
|
||||
async def page_dict_data(
|
||||
self,
|
||||
*,
|
||||
dict_type_id: int | None,
|
||||
dict_label: str | None,
|
||||
dict_value: str | None,
|
||||
page: int,
|
||||
limit: int,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
params = {
|
||||
"type_id": dict_type_id,
|
||||
"dict_label": f"%{dict_label}%" if dict_label else None,
|
||||
"dict_value": f"%{dict_value}%" if dict_value else None,
|
||||
"offset": (page - 1) * limit,
|
||||
"limit": limit,
|
||||
}
|
||||
where = (
|
||||
"d.dict_type_id = :type_id AND (:dict_label IS NULL OR d.dict_label LIKE :dict_label) "
|
||||
"AND (:dict_value IS NULL OR d.dict_value LIKE :dict_value)"
|
||||
)
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM sys_dict_data d WHERE {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
"SELECT d.id, d.dict_type_id, d.dict_label, d.dict_value, d.remark, d.sort, d.creator, " # noqa: S608
|
||||
"d.create_date, d.updater, d.update_date, creator.username AS creator_name, "
|
||||
"updater.username AS updater_name FROM sys_dict_data d "
|
||||
"LEFT JOIN sys_user creator ON creator.id = d.creator "
|
||||
"LEFT JOIN sys_user updater ON updater.id = d.updater "
|
||||
f"WHERE {where} ORDER BY d.sort ASC LIMIT :limit OFFSET :offset", # noqa: S608
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def get_dict_data(self, data_id: int) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, dict_type_id, dict_label, dict_value, remark, sort, creator, create_date, updater, update_date "
|
||||
"FROM sys_dict_data WHERE id = :id",
|
||||
{"id": data_id},
|
||||
)
|
||||
|
||||
async def dict_data_label_exists(
|
||||
self,
|
||||
dict_type_id: int | None,
|
||||
compared_label: str | None,
|
||||
*,
|
||||
exclude_id: int | None = None,
|
||||
) -> bool:
|
||||
count = await self.scalar(
|
||||
"SELECT COUNT(*) FROM sys_dict_data WHERE dict_type_id = :type_id AND dict_label = :label "
|
||||
"AND (:exclude_id IS NULL OR id <> :exclude_id)",
|
||||
{"type_id": dict_type_id, "label": compared_label, "exclude_id": exclude_id},
|
||||
)
|
||||
return int(count or 0) > 0
|
||||
|
||||
async def dict_type_code(self, type_id: int | None) -> str | None:
|
||||
value = await self.scalar("SELECT dict_type FROM sys_dict_type WHERE id = :id", {"id": type_id})
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def insert_dict_data(
|
||||
self,
|
||||
*,
|
||||
data_id: int,
|
||||
dict_type_id: int | None,
|
||||
dict_label: str | None,
|
||||
dict_value: str | None,
|
||||
remark: str | None,
|
||||
sort: int | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO sys_dict_data "
|
||||
"(id, dict_type_id, dict_label, dict_value, remark, sort, creator, create_date, updater, update_date) "
|
||||
"VALUES (:id, :type_id, :label, :value, :remark, :sort, :user_id, :now, :user_id, :now)",
|
||||
{
|
||||
"id": data_id,
|
||||
"type_id": dict_type_id,
|
||||
"label": dict_label,
|
||||
"value": dict_value,
|
||||
"remark": remark,
|
||||
"sort": sort,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def update_dict_data(
|
||||
self,
|
||||
*,
|
||||
data_id: int | None,
|
||||
dict_type_id: int | None,
|
||||
dict_label: str | None,
|
||||
dict_value: str | None,
|
||||
remark: str | None,
|
||||
sort: int | None,
|
||||
user_id: int,
|
||||
now: datetime,
|
||||
) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE sys_dict_data SET "
|
||||
"dict_type_id = CASE WHEN :has_type_id = 1 THEN :type_id ELSE dict_type_id END, "
|
||||
"dict_label = CASE WHEN :has_label = 1 THEN :label ELSE dict_label END, "
|
||||
"dict_value = CASE WHEN :has_value = 1 THEN :value ELSE dict_value END, "
|
||||
"remark = CASE WHEN :has_remark = 1 THEN :remark ELSE remark END, "
|
||||
"sort = CASE WHEN :has_sort = 1 THEN :sort ELSE sort END, updater = :user_id, update_date = :now "
|
||||
"WHERE id = :id",
|
||||
{
|
||||
"id": data_id,
|
||||
"has_type_id": int(dict_type_id is not None),
|
||||
"type_id": dict_type_id,
|
||||
"has_label": int(dict_label is not None),
|
||||
"label": dict_label,
|
||||
"has_value": int(dict_value is not None),
|
||||
"value": dict_value,
|
||||
"has_remark": int(remark is not None),
|
||||
"remark": remark,
|
||||
"has_sort": int(sort is not None),
|
||||
"sort": sort,
|
||||
"user_id": user_id,
|
||||
"now": now,
|
||||
},
|
||||
)
|
||||
|
||||
async def dict_type_codes_for_data_ids(self, ids: list[int]) -> list[str]:
|
||||
statement = text(
|
||||
"SELECT DISTINCT t.dict_type FROM sys_dict_type t JOIN sys_dict_data d ON d.dict_type_id = t.id "
|
||||
"WHERE d.id IN :ids"
|
||||
).bindparams(bindparam("ids", expanding=True))
|
||||
rows = await self.fetch_all(statement, {"ids": ids})
|
||||
return [str(row["dict_type"]) for row in rows]
|
||||
|
||||
async def delete_dict_data(self, ids: list[int]) -> int:
|
||||
statement = text("DELETE FROM sys_dict_data WHERE id IN :ids").bindparams(bindparam("ids", expanding=True))
|
||||
return await self.execute(statement, {"ids": ids})
|
||||
|
||||
async def dict_items(self, dict_type: str) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT d.dict_label AS name, d.dict_value AS `key` FROM sys_dict_data d "
|
||||
"LEFT JOIN sys_dict_type t ON d.dict_type_id = t.id "
|
||||
"WHERE t.dict_type = :dict_type ORDER BY d.sort ASC",
|
||||
{"dict_type": dict_type},
|
||||
)
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import bindparam, text
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
|
||||
class TimbreRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def page(
|
||||
self, *, tts_model_id: str, name: str | None, offset: int, limit: int
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
where = (
|
||||
"WHERE tts_model_id=:tts_model_id AND "
|
||||
"(:name IS NULL OR :name='' OR name LIKE CONCAT('%', :name, '%'))"
|
||||
)
|
||||
params = {"tts_model_id": tts_model_id, "name": name, "offset": offset, "limit": limit}
|
||||
total = int(await self.scalar(f"SELECT COUNT(*) FROM ai_tts_voice {where}", params) or 0) # noqa: S608
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT * FROM ai_tts_voice {where} LIMIT :offset, :limit", # noqa: S608
|
||||
params,
|
||||
)
|
||||
return rows, total
|
||||
|
||||
async def insert(self, values: dict[str, Any]) -> None:
|
||||
await self.execute(
|
||||
"INSERT INTO ai_tts_voice "
|
||||
"(id, languages, name, remark, reference_audio, reference_text, sort, tts_model_id, tts_voice, "
|
||||
"voice_demo, creator, create_date) VALUES (:id, :languages, :name, :remark, :reference_audio, "
|
||||
":reference_text, :sort, :tts_model_id, :tts_voice, :voice_demo, :creator, :now)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def update(self, values: dict[str, Any]) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_tts_voice SET languages=:languages, name=:name, remark=COALESCE(:remark, remark), "
|
||||
"reference_audio=COALESCE(:reference_audio, reference_audio), "
|
||||
"reference_text=COALESCE(:reference_text, reference_text), sort=:sort, "
|
||||
"tts_model_id=:tts_model_id, tts_voice=:tts_voice, "
|
||||
"voice_demo=COALESCE(:voice_demo, voice_demo), updater=:updater, "
|
||||
"update_date=:now WHERE id=:id",
|
||||
values,
|
||||
)
|
||||
|
||||
async def delete(self, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
statement = text("DELETE FROM ai_tts_voice WHERE id IN :ids").bindparams(bindparam("ids", expanding=True))
|
||||
result = await self.session.execute(statement, {"ids": list(ids)})
|
||||
return int(getattr(result, "rowcount", 0) or 0)
|
||||
|
||||
async def voices(
|
||||
self, model_id: str, name: str | None, user_id: int
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
normal = await self.fetch_all(
|
||||
"SELECT id, name, voice_demo, languages FROM ai_tts_voice WHERE tts_model_id=:model_id "
|
||||
"AND (:name IS NULL OR :name='' OR name LIKE CONCAT('%', :name, '%'))",
|
||||
{"model_id": model_id or "", "name": name},
|
||||
)
|
||||
clones = await self.fetch_all(
|
||||
"SELECT id, name, voice_id AS voice_demo, languages FROM ai_voice_clone "
|
||||
"WHERE model_id=:model_id AND user_id=:user_id AND train_status=2",
|
||||
{"model_id": model_id, "user_id": user_id},
|
||||
)
|
||||
return normal, clones
|
||||
@@ -0,0 +1,170 @@
|
||||
from __future__ import annotations
|
||||
|
||||
# Every interpolated SQL fragment below is a module constant or a service-side allowlist.
|
||||
# ruff: noqa: S608
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.database import Repository
|
||||
|
||||
VOICE_COLUMNS = (
|
||||
"id, name, model_id, voice_id, languages, user_id, voice, train_status, train_error, creator, create_date"
|
||||
)
|
||||
|
||||
|
||||
class VoiceCloneRepository(Repository):
|
||||
def __init__(self, session: AsyncSession):
|
||||
super().__init__(session)
|
||||
|
||||
async def count(self, *, name: str | None, user_id: str | None) -> int:
|
||||
where, params = self._filters(name=name, user_id=user_id)
|
||||
return int(await self.scalar(f"SELECT COUNT(*) FROM ai_voice_clone{where}", params) or 0)
|
||||
|
||||
async def page(
|
||||
self,
|
||||
*,
|
||||
page: int,
|
||||
limit: int,
|
||||
name: str | None,
|
||||
user_id: str | None,
|
||||
order_fields: Sequence[str],
|
||||
ascending: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
where, params = self._filters(name=name, user_id=user_id)
|
||||
params.update(limit=limit, offset=max(page - 1, 0) * limit)
|
||||
direction = "ASC" if ascending else "DESC"
|
||||
order_by = ", ".join(f"{field} {direction}" for field in order_fields)
|
||||
return await self.fetch_all(
|
||||
f"SELECT {VOICE_COLUMNS} FROM ai_voice_clone{where} "
|
||||
f"ORDER BY {order_by} LIMIT :limit OFFSET :offset",
|
||||
params,
|
||||
)
|
||||
|
||||
async def get(self, voice_id: str | None) -> dict[str, Any] | None:
|
||||
if voice_id is None:
|
||||
return None
|
||||
return await self.fetch_one(
|
||||
f"SELECT {VOICE_COLUMNS} FROM ai_voice_clone WHERE id = :id LIMIT 1",
|
||||
{"id": voice_id},
|
||||
)
|
||||
|
||||
async def list_by_user(self, user_id: int) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
f"SELECT {VOICE_COLUMNS} FROM ai_voice_clone "
|
||||
"WHERE user_id = :user_id ORDER BY create_date DESC",
|
||||
{"user_id": user_id},
|
||||
)
|
||||
|
||||
async def voice_id_count(self, *, model_id: str, voice_id: str) -> int:
|
||||
return int(
|
||||
await self.scalar(
|
||||
"SELECT COUNT(*) FROM ai_voice_clone WHERE voice_id = :voice_id AND model_id = :model_id",
|
||||
{"model_id": model_id, "voice_id": voice_id},
|
||||
)
|
||||
or 0
|
||||
)
|
||||
|
||||
async def insert_many(self, values: Sequence[Mapping[str, Any]]) -> int:
|
||||
return await self.execute_many(
|
||||
"INSERT INTO ai_voice_clone "
|
||||
"(id, name, model_id, voice_id, languages, user_id, voice, train_status, train_error, creator, "
|
||||
"create_date) VALUES (:id, :name, :model_id, :voice_id, :languages, :user_id, :voice, :train_status, "
|
||||
":train_error, :creator, :create_date)",
|
||||
values,
|
||||
)
|
||||
|
||||
async def delete_many(self, ids: Sequence[str]) -> int:
|
||||
if not ids:
|
||||
return 0
|
||||
placeholders = ", ".join(f":id_{index}" for index in range(len(ids)))
|
||||
params = {f"id_{index}": value for index, value in enumerate(ids)}
|
||||
return await self.execute(f"DELETE FROM ai_voice_clone WHERE id IN ({placeholders})", params)
|
||||
|
||||
async def update_voice(self, voice_id: str, data: bytes) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_voice_clone SET voice = :voice, train_status = 0 WHERE id = :id",
|
||||
{"id": voice_id, "voice": data},
|
||||
)
|
||||
|
||||
async def update_name(self, voice_id: str, name: str) -> int:
|
||||
return await self.execute(
|
||||
"UPDATE ai_voice_clone SET name = :name WHERE id = :id",
|
||||
{"id": voice_id, "name": name},
|
||||
)
|
||||
|
||||
async def update_training(
|
||||
self,
|
||||
voice_id: str,
|
||||
*,
|
||||
train_status: int,
|
||||
train_error: str | None,
|
||||
speaker_id: str | None = None,
|
||||
) -> int:
|
||||
if speaker_id is None:
|
||||
return await self.execute(
|
||||
"UPDATE ai_voice_clone SET train_status = :train_status, train_error = :train_error WHERE id = :id",
|
||||
{"id": voice_id, "train_status": train_status, "train_error": train_error},
|
||||
)
|
||||
return await self.execute(
|
||||
"UPDATE ai_voice_clone SET train_status = :train_status, train_error = :train_error, "
|
||||
"voice_id = :speaker_id WHERE id = :id",
|
||||
{
|
||||
"id": voice_id,
|
||||
"train_status": train_status,
|
||||
"train_error": train_error,
|
||||
"speaker_id": speaker_id,
|
||||
},
|
||||
)
|
||||
|
||||
async def get_model_config(self, model_id: str) -> dict[str, Any] | None:
|
||||
return await self.fetch_one(
|
||||
"SELECT id, model_name, config_json FROM ai_model_config WHERE id = :id LIMIT 1",
|
||||
{"id": model_id},
|
||||
)
|
||||
|
||||
async def get_model_name(self, model_id: str) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT model_name FROM ai_model_config WHERE id = :id LIMIT 1",
|
||||
{"id": model_id},
|
||||
)
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def get_usernames(self, user_ids: Sequence[int]) -> dict[int, str]:
|
||||
if not user_ids:
|
||||
return {}
|
||||
unique_ids = list(dict.fromkeys(user_ids))
|
||||
placeholders = ", ".join(f":user_{index}" for index in range(len(unique_ids)))
|
||||
params = {f"user_{index}": value for index, value in enumerate(unique_ids)}
|
||||
rows = await self.fetch_all(
|
||||
f"SELECT id, username FROM sys_user WHERE id IN ({placeholders})",
|
||||
params,
|
||||
)
|
||||
return {int(row["id"]): str(row["username"]) for row in rows}
|
||||
|
||||
async def get_username(self, user_id: int) -> str | None:
|
||||
value = await self.scalar(
|
||||
"SELECT username FROM sys_user WHERE id = :id LIMIT 1",
|
||||
{"id": user_id},
|
||||
)
|
||||
return None if value is None else str(value)
|
||||
|
||||
async def get_tts_platforms(self) -> list[dict[str, Any]]:
|
||||
return await self.fetch_all(
|
||||
"SELECT id, model_name AS modelName FROM ai_model_config "
|
||||
"WHERE model_type = 'TTS' AND JSON_EXTRACT(config_json, '$.type') = 'huoshan_double_stream'"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _filters(*, name: str | None, user_id: str | None) -> tuple[str, dict[str, Any]]:
|
||||
clauses: list[str] = []
|
||||
params: dict[str, Any] = {}
|
||||
if user_id is not None and user_id.strip():
|
||||
clauses.append("user_id = :user_id")
|
||||
params["user_id"] = user_id
|
||||
if name is not None and name.strip():
|
||||
clauses.append("(name LIKE :name OR voice_id = :exact_name)")
|
||||
params["name"] = f"%{name}%"
|
||||
params["exact_name"] = name
|
||||
return (" WHERE " + " AND ".join(clauses) if clauses else "", params)
|
||||
Reference in New Issue
Block a user