243 lines
9.3 KiB
Python
243 lines
9.3 KiB
Python
"""AI 角色 Prompt 配置服务。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Optional
|
|
|
|
from sqlalchemy import and_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ..config.ai_prompts import build_default_role_prompt_bundle
|
|
from ..models import Integration
|
|
|
|
from ..utils.timezone import now_cn_aware
|
|
INTEGRATION_TYPE = "ai_prompt_profiles"
|
|
ROLE_KEY_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{1,31}$")
|
|
MAX_ROLES = 12
|
|
|
|
|
|
def _safe_json(value: Any) -> dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
try:
|
|
return json.loads(json.dumps(value))
|
|
except Exception:
|
|
return dict(value)
|
|
if isinstance(value, str):
|
|
try:
|
|
parsed = json.loads(value)
|
|
if isinstance(parsed, dict):
|
|
return parsed
|
|
except Exception:
|
|
return {}
|
|
return {}
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return now_cn_aware().isoformat()
|
|
|
|
|
|
class PromptProfileService:
|
|
"""管理租户级 AI 角色 Prompt 配置。"""
|
|
|
|
async def get_profiles(self, db: AsyncSession, *, tenant_id: int) -> dict[str, Any]:
|
|
integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id)
|
|
normalized = self._normalize_bundle(bundle)
|
|
if normalized != bundle:
|
|
integration.config = json.loads(json.dumps(normalized, ensure_ascii=False))
|
|
integration.is_active = True
|
|
await db.commit()
|
|
return normalized
|
|
|
|
async def update_profiles(
|
|
self,
|
|
db: AsyncSession,
|
|
*,
|
|
tenant_id: int,
|
|
roles: list[dict[str, Any]],
|
|
active_role_key: Optional[str] = None,
|
|
) -> dict[str, Any]:
|
|
if not roles:
|
|
raise ValueError("至少需要保留一个角色")
|
|
if len(roles) > MAX_ROLES:
|
|
raise ValueError(f"角色数量不能超过 {MAX_ROLES} 个")
|
|
|
|
dedup: dict[str, dict[str, Any]] = {}
|
|
for idx, raw in enumerate(roles):
|
|
key = str(raw.get("role_key") or "").strip().lower()
|
|
if not key:
|
|
raise ValueError(f"第 {idx + 1} 个角色缺少 role_key")
|
|
if not ROLE_KEY_RE.match(key):
|
|
raise ValueError(f"角色标识 {key} 格式不正确，仅支持小写字母/数字/_/-")
|
|
if key in dedup:
|
|
raise ValueError(f"角色标识重复: {key}")
|
|
|
|
role_name = str(raw.get("role_name") or "").strip()
|
|
system_prompt = str(raw.get("system_prompt") or "").strip()
|
|
if not role_name:
|
|
raise ValueError(f"角色 {key} 缺少角色名称")
|
|
if len(role_name) > 32:
|
|
raise ValueError(f"角色 {key} 名称过长（最多 32 字）")
|
|
if len(system_prompt) < 20:
|
|
raise ValueError(f"角色 {key} 的提示词过短（至少 20 字）")
|
|
if len(system_prompt) > 8000:
|
|
raise ValueError(f"角色 {key} 的提示词过长（最多 8000 字）")
|
|
|
|
role = {
|
|
"role_key": key,
|
|
"role_name": role_name,
|
|
"role_description": str(raw.get("role_description") or "").strip()[:300],
|
|
"system_prompt": system_prompt,
|
|
"enabled": bool(raw.get("enabled", True)),
|
|
"updated_at": str(raw.get("updated_at") or _now_iso()),
|
|
}
|
|
dedup[key] = role
|
|
|
|
ordered_roles = list(dedup.values())
|
|
enabled_keys = [str(r["role_key"]) for r in ordered_roles if bool(r.get("enabled"))]
|
|
if not enabled_keys:
|
|
raise ValueError("至少需要启用一个角色")
|
|
|
|
requested_active = str(active_role_key or "").strip().lower()
|
|
if requested_active and requested_active in enabled_keys:
|
|
active = requested_active
|
|
else:
|
|
active = enabled_keys[0]
|
|
|
|
bundle = {
|
|
"version": 1,
|
|
"active_role_key": active,
|
|
"roles": ordered_roles,
|
|
"updated_at": _now_iso(),
|
|
}
|
|
|
|
integration, _ = await self._get_or_init_bundle(db, tenant_id=tenant_id)
|
|
integration.config = json.loads(json.dumps(bundle, ensure_ascii=False))
|
|
integration.is_active = True
|
|
await db.commit()
|
|
return self._normalize_bundle(bundle)
|
|
|
|
async def get_active_role(self, db: AsyncSession, *, tenant_id: int) -> dict[str, Any]:
|
|
data = await self.get_profiles(db, tenant_id=tenant_id)
|
|
active_key = str(data.get("active_role_key") or "")
|
|
roles = data.get("roles", [])
|
|
if not isinstance(roles, list):
|
|
roles = []
|
|
active_role = next((r for r in roles if str(r.get("role_key")) == active_key), None)
|
|
if active_role is None and roles:
|
|
active_role = roles[0]
|
|
return {
|
|
"active_role_key": active_key,
|
|
"role": active_role,
|
|
}
|
|
|
|
async def _get_or_init_bundle(
|
|
self,
|
|
db: AsyncSession,
|
|
*,
|
|
tenant_id: int,
|
|
) -> tuple[Integration, dict[str, Any]]:
|
|
result = await db.execute(
|
|
select(Integration).where(
|
|
and_(
|
|
Integration.tenant_id == int(tenant_id),
|
|
Integration.type == INTEGRATION_TYPE,
|
|
)
|
|
)
|
|
)
|
|
integration = result.scalar_one_or_none()
|
|
|
|
if integration is None:
|
|
bundle = self._normalize_bundle(build_default_role_prompt_bundle())
|
|
integration = Integration(
|
|
tenant_id=int(tenant_id),
|
|
type=INTEGRATION_TYPE,
|
|
config=json.loads(json.dumps(bundle, ensure_ascii=False)),
|
|
is_active=True,
|
|
)
|
|
db.add(integration)
|
|
await db.commit()
|
|
return integration, bundle
|
|
|
|
bundle = self._normalize_bundle(_safe_json(integration.config))
|
|
return integration, bundle
|
|
|
|
def _normalize_bundle(self, raw: dict[str, Any]) -> dict[str, Any]:
|
|
defaults = build_default_role_prompt_bundle()
|
|
default_roles_raw = defaults.get("roles", [])
|
|
default_roles: dict[str, dict[str, Any]] = {}
|
|
if isinstance(default_roles_raw, list):
|
|
for row in default_roles_raw:
|
|
if not isinstance(row, dict):
|
|
continue
|
|
key = str(row.get("role_key") or "").strip().lower()
|
|
if key:
|
|
default_roles[key] = {
|
|
"role_key": key,
|
|
"role_name": str(row.get("role_name") or key)[:32],
|
|
"role_description": str(row.get("role_description") or "")[:300],
|
|
"system_prompt": str(row.get("system_prompt") or "").strip(),
|
|
"enabled": bool(row.get("enabled", True)),
|
|
"updated_at": str(row.get("updated_at") or _now_iso()),
|
|
}
|
|
|
|
incoming_roles = raw.get("roles") if isinstance(raw.get("roles"), list) else []
|
|
roles: list[dict[str, Any]] = []
|
|
seen: set[str] = set()
|
|
|
|
for item in incoming_roles:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
key = str(item.get("role_key") or "").strip().lower()
|
|
if not key or key in seen or not ROLE_KEY_RE.match(key):
|
|
continue
|
|
seen.add(key)
|
|
|
|
fallback = default_roles.get(key, {})
|
|
prompt_text = str(item.get("system_prompt") or fallback.get("system_prompt") or "").strip()
|
|
if len(prompt_text) < 20:
|
|
prompt_text = str(fallback.get("system_prompt") or "").strip()
|
|
if len(prompt_text) < 20:
|
|
continue
|
|
|
|
roles.append(
|
|
{
|
|
"role_key": key,
|
|
"role_name": str(item.get("role_name") or fallback.get("role_name") or key).strip()[:32] or key,
|
|
"role_description": str(item.get("role_description") or fallback.get("role_description") or "").strip()[:300],
|
|
"system_prompt": prompt_text[:8000],
|
|
"enabled": bool(item.get("enabled", fallback.get("enabled", True))),
|
|
"updated_at": str(item.get("updated_at") or _now_iso()),
|
|
}
|
|
)
|
|
if len(roles) >= MAX_ROLES:
|
|
break
|
|
|
|
if not roles:
|
|
roles = list(default_roles.values())[:MAX_ROLES]
|
|
|
|
enabled_keys = [str(r["role_key"]) for r in roles if bool(r.get("enabled"))]
|
|
if not enabled_keys and roles:
|
|
roles[0]["enabled"] = True
|
|
enabled_keys = [str(roles[0]["role_key"])]
|
|
|
|
active_raw = str(raw.get("active_role_key") or "").strip().lower()
|
|
active_role_key = active_raw if active_raw in enabled_keys else (enabled_keys[0] if enabled_keys else str(roles[0]["role_key"]))
|
|
|
|
return {
|
|
"version": 1,
|
|
"active_role_key": active_role_key,
|
|
"roles": roles,
|
|
"updated_at": str(raw.get("updated_at") or _now_iso()),
|
|
"tips": {
|
|
"active_role_key": "当前生效角色，影响 AI 回复语气与任务分派风格。",
|
|
"system_prompt": "建议写清角色职责、回复风格、输出边界；支持后续扩展多个角色。",
|
|
"role_key": "角色标识用于系统路由，保存后建议保持稳定。",
|
|
},
|
|
}
|
|
|
|
|
|
prompt_profile_service = PromptProfileService()
|