264 lines
10 KiB
Python
264 lines
10 KiB
Python
from contextlib import closing
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_DB_PATH = "/app/data/social_writer.sqlite3"
|
|
API_LOG_RETENTION_DAYS = 30
|
|
API_LOG_MAX_FIELD_BYTES = 100_000
|
|
|
|
|
|
def connect(db_path: str | None = None) -> sqlite3.Connection:
|
|
path = Path(db_path or os.environ.get("BOT_DB_PATH", DEFAULT_DB_PATH))
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(path, timeout=10)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.execute("PRAGMA busy_timeout = 10000")
|
|
connection.execute("PRAGMA journal_mode = WAL")
|
|
return connection
|
|
|
|
|
|
def init_db(db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
telegram_user_id INTEGER PRIMARY KEY,
|
|
tone_of_voice TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS user_post_samples (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
telegram_user_id INTEGER NOT NULL,
|
|
post_text TEXT NOT NULL,
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE (telegram_user_id, post_text)
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS generated_posts (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
telegram_user_id INTEGER NOT NULL,
|
|
post_text TEXT NOT NULL,
|
|
origin TEXT NOT NULL DEFAULT 'generated',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS api_call_logs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
telegram_user_id INTEGER,
|
|
stage TEXT NOT NULL,
|
|
model TEXT NOT NULL DEFAULT '',
|
|
prompt_sha256 TEXT NOT NULL DEFAULT '',
|
|
request_json TEXT NOT NULL,
|
|
response_json TEXT NOT NULL DEFAULT '',
|
|
status_code INTEGER,
|
|
error TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_api_call_logs_created_at ON api_call_logs(created_at)"
|
|
)
|
|
connection.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_api_call_logs_user_stage ON api_call_logs(telegram_user_id, stage)"
|
|
)
|
|
|
|
|
|
def add_post_samples(user_id: int, posts: list[str], db_path: str | None = None) -> list[str]:
|
|
samples = [post.strip() for post in posts if post and post.strip()]
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,))
|
|
connection.executemany(
|
|
"INSERT OR IGNORE INTO user_post_samples (telegram_user_id, post_text) VALUES (?, ?)",
|
|
[(user_id, post) for post in samples],
|
|
)
|
|
rows = connection.execute(
|
|
"SELECT post_text FROM user_post_samples WHERE telegram_user_id = ? ORDER BY id",
|
|
(user_id,),
|
|
).fetchall()
|
|
return [str(row["post_text"]) for row in rows]
|
|
|
|
|
|
def get_post_samples(user_id: int, db_path: str | None = None) -> list[str]:
|
|
with closing(connect(db_path)) as connection:
|
|
rows = connection.execute(
|
|
"SELECT post_text FROM user_post_samples WHERE telegram_user_id = ? ORDER BY id",
|
|
(user_id,),
|
|
).fetchall()
|
|
return [str(row["post_text"]) for row in rows]
|
|
|
|
|
|
def _bounded_json(value: str, max_bytes: int = API_LOG_MAX_FIELD_BYTES) -> str:
|
|
if len(value.encode("utf-8")) <= max_bytes:
|
|
return value
|
|
preview = value.encode("utf-8")[:max_bytes // 4].decode("utf-8", errors="ignore")
|
|
return json.dumps({"truncated": True, "preview": preview}, ensure_ascii=False)
|
|
|
|
|
|
def save_api_call_log(
|
|
*,
|
|
stage: str,
|
|
request_json: str,
|
|
response_json: str = "",
|
|
model: str = "",
|
|
prompt_sha256: str = "",
|
|
user_id: int | None = None,
|
|
status_code: int | None = None,
|
|
error: str = "",
|
|
db_path: str | None = None,
|
|
) -> int:
|
|
request_json = _bounded_json(request_json)
|
|
response_json = _bounded_json(response_json) if response_json else ""
|
|
with closing(connect(db_path)) as connection, connection:
|
|
cursor = connection.execute(
|
|
"""
|
|
INSERT INTO api_call_logs (
|
|
telegram_user_id, stage, model, prompt_sha256, request_json,
|
|
response_json, status_code, error
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(user_id, stage[:100], model[:200], prompt_sha256[:64], request_json,
|
|
response_json, status_code, error[:2000]),
|
|
)
|
|
connection.execute(
|
|
"DELETE FROM api_call_logs WHERE created_at < datetime('now', ?)",
|
|
(f"-{API_LOG_RETENTION_DAYS} days",),
|
|
)
|
|
if cursor.lastrowid is None:
|
|
raise RuntimeError("Could not save API log record")
|
|
return cursor.lastrowid
|
|
|
|
|
|
def list_api_call_logs(
|
|
limit: int = 50, user_id: int | None = None, db_path: str | None = None
|
|
) -> list[dict[str, object]]:
|
|
limit = max(1, min(limit, 200))
|
|
with closing(connect(db_path)) as connection:
|
|
if user_id is None:
|
|
rows = connection.execute(
|
|
"SELECT id, telegram_user_id, stage, model, prompt_sha256, status_code, error, created_at, "
|
|
"length(request_json) AS request_chars, length(response_json) AS response_chars "
|
|
"FROM api_call_logs ORDER BY id DESC LIMIT ?", (limit,),
|
|
).fetchall()
|
|
else:
|
|
rows = connection.execute(
|
|
"SELECT id, telegram_user_id, stage, model, prompt_sha256, status_code, error, created_at, "
|
|
"length(request_json) AS request_chars, length(response_json) AS response_chars "
|
|
"FROM api_call_logs WHERE telegram_user_id = ? ORDER BY id DESC LIMIT ?", (user_id, limit),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
def get_api_call_log(log_id: int, db_path: str | None = None) -> dict[str, object] | None:
|
|
with closing(connect(db_path)) as connection:
|
|
row = connection.execute("SELECT * FROM api_call_logs WHERE id = ?", (log_id,)).fetchone()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def clear_user_tone_data(user_id: int, db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("DELETE FROM user_post_samples WHERE telegram_user_id = ?", (user_id,))
|
|
connection.execute(
|
|
"UPDATE users SET tone_of_voice = '', updated_at = CURRENT_TIMESTAMP WHERE telegram_user_id = ?",
|
|
(user_id,),
|
|
)
|
|
|
|
|
|
def get_post_history(user_id: int, limit: int = 30, db_path: str | None = None) -> list[dict[str, object]]:
|
|
limit = max(1, min(limit, 100))
|
|
with closing(connect(db_path)) as connection:
|
|
rows = connection.execute(
|
|
"SELECT post_text, origin, created_at FROM generated_posts "
|
|
"WHERE telegram_user_id = ? ORDER BY id DESC LIMIT ?",
|
|
(user_id, limit),
|
|
).fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
def save_generated_posts(user_id: int, posts: list[str], origin: str, db_path: str | None = None) -> None:
|
|
cleaned = [post.strip() for post in posts if post and post.strip()]
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,))
|
|
connection.executemany(
|
|
"INSERT INTO generated_posts (telegram_user_id, post_text, origin) VALUES (?, ?, ?)",
|
|
[(user_id, post, origin[:30]) for post in cleaned],
|
|
)
|
|
|
|
|
|
def get_recent_post_history_text(user_id: int, limit: int = 30, char_limit: int = 12000, db_path: str | None = None) -> list[str]:
|
|
history = get_post_history(user_id, limit, db_path)
|
|
result: list[str] = []
|
|
total = 0
|
|
for row in history:
|
|
text = str(row["post_text"])
|
|
if total + len(text) > char_limit:
|
|
break
|
|
result.append(text)
|
|
total += len(text)
|
|
return result
|
|
|
|
|
|
def get_tone_of_voice(user_id: int, db_path: str | None = None) -> str:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,)
|
|
)
|
|
row = connection.execute(
|
|
"SELECT tone_of_voice FROM users WHERE telegram_user_id = ?", (user_id,)
|
|
).fetchone()
|
|
return str(row["tone_of_voice"] or "") if row else ""
|
|
|
|
|
|
def get_deploy_notice_version(db_path: str | None = None) -> str:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"CREATE TABLE IF NOT EXISTS app_state (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
|
|
)
|
|
row = connection.execute(
|
|
"SELECT value FROM app_state WHERE key = 'deploy_notice_version'"
|
|
).fetchone()
|
|
return str(row["value"]) if row else ""
|
|
|
|
|
|
def set_deploy_notice_version(version: str, db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"CREATE TABLE IF NOT EXISTS app_state (key TEXT PRIMARY KEY, value TEXT NOT NULL)"
|
|
)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO app_state (key, value) VALUES ('deploy_notice_version', ?)
|
|
ON CONFLICT(key) DO UPDATE SET value = excluded.value
|
|
""",
|
|
(version,),
|
|
)
|
|
|
|
|
|
def save_tone_of_voice(user_id: int, tone_of_voice: str, db_path: str | None = None) -> None:
|
|
tone = tone_of_voice.strip()[:1200]
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO users (telegram_user_id, tone_of_voice, updated_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
ON CONFLICT(telegram_user_id) DO UPDATE SET
|
|
tone_of_voice = excluded.tone_of_voice,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
""",
|
|
(user_id, tone),
|
|
)
|