forked from retoor/devplacepy
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any
|
|
|
|
from devplacepy.config import DEFAULT_CORRECTION_PROMPT
|
|
from devplacepy.database import get_table
|
|
from devplacepy.services.devii.errors import ToolInputError
|
|
|
|
logger = logging.getLogger("devii.ai_correction")
|
|
|
|
|
|
class AiCorrectionController:
|
|
def __init__(self, owner_kind: str, owner_id: str) -> None:
|
|
self._owner_kind = owner_kind
|
|
self._owner_id = owner_id
|
|
|
|
async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
|
|
if name == "ai_correction_get":
|
|
return self._get()
|
|
if name == "ai_correction_set":
|
|
return self._set(arguments)
|
|
raise ToolInputError(f"Unknown AI correction tool: {name}")
|
|
|
|
def _require_user(self) -> None:
|
|
if self._owner_kind != "user":
|
|
raise ToolInputError(
|
|
"AI content correction is only available for signed-in users."
|
|
)
|
|
|
|
def _coerce_bool(self, raw: Any) -> bool:
|
|
if isinstance(raw, bool):
|
|
return raw
|
|
return str(raw).strip().lower() in ("true", "1", "yes", "on")
|
|
|
|
def _value(self, arguments: dict[str, Any]) -> bool:
|
|
return self._coerce_bool(arguments.get("enabled"))
|
|
|
|
def _user(self) -> dict[str, Any]:
|
|
user = get_table("users").find_one(uid=self._owner_id)
|
|
if not user:
|
|
raise ToolInputError("User not found.")
|
|
return user
|
|
|
|
def _get(self) -> str:
|
|
self._require_user()
|
|
user = self._user()
|
|
return json.dumps(
|
|
{
|
|
"status": "success",
|
|
"enabled": bool(user.get("ai_correction_enabled")),
|
|
"sync": bool(user.get("ai_correction_sync")),
|
|
"prompt": user.get("ai_correction_prompt") or DEFAULT_CORRECTION_PROMPT,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
def _set(self, arguments: dict[str, Any]) -> str:
|
|
self._require_user()
|
|
user = self._user()
|
|
enabled = 1 if self._value(arguments) else 0
|
|
if "sync" in arguments and arguments.get("sync") is not None:
|
|
sync = 1 if self._coerce_bool(arguments.get("sync")) else 0
|
|
else:
|
|
sync = 1 if user.get("ai_correction_sync") else 0
|
|
prompt_raw = arguments.get("prompt")
|
|
if prompt_raw is None:
|
|
prompt = user.get("ai_correction_prompt") or DEFAULT_CORRECTION_PROMPT
|
|
else:
|
|
prompt = str(prompt_raw).strip()[:20000] or DEFAULT_CORRECTION_PROMPT
|
|
get_table("users").update(
|
|
{
|
|
"uid": self._owner_id,
|
|
"ai_correction_enabled": enabled,
|
|
"ai_correction_sync": sync,
|
|
"ai_correction_prompt": prompt,
|
|
},
|
|
["uid"],
|
|
)
|
|
from devplacepy.utils import clear_user_cache
|
|
|
|
clear_user_cache(self._owner_id)
|
|
return json.dumps(
|
|
{
|
|
"status": "success",
|
|
"enabled": bool(enabled),
|
|
"sync": bool(sync),
|
|
"prompt": prompt,
|
|
},
|
|
ensure_ascii=False,
|
|
)
|