448 lines
19 KiB
Python
448 lines
19 KiB
Python
from contextlib import closing
|
|
import datetime
|
|
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 user_training_notes (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
telegram_user_id INTEGER NOT NULL,
|
|
note_text TEXT NOT NULL,
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE (telegram_user_id, note_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 generated_ideas (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
telegram_user_id INTEGER NOT NULL,
|
|
title TEXT NOT NULL,
|
|
angle TEXT NOT NULL,
|
|
why TEXT NOT NULL,
|
|
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
|
|
)
|
|
"""
|
|
)
|
|
api_log_columns = {
|
|
row["name"] for row in connection.execute("PRAGMA table_info(api_call_logs)").fetchall()
|
|
}
|
|
if "operation_id" not in api_log_columns:
|
|
connection.execute("ALTER TABLE api_call_logs ADD COLUMN operation_id TEXT NOT NULL DEFAULT ''")
|
|
if "duration_ms" not in api_log_columns:
|
|
connection.execute("ALTER TABLE api_call_logs ADD COLUMN duration_ms INTEGER")
|
|
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)"
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS subscriptions (
|
|
telegram_user_id INTEGER PRIMARY KEY,
|
|
trial_started_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
monthly_days INTEGER NOT NULL DEFAULT 30,
|
|
last_paid_at TEXT,
|
|
is_active INTEGER NOT NULL DEFAULT 1,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
"""
|
|
)
|
|
|
|
|
|
# Периоды: триал 3 дня, подписка 30 дней (100 звёзд)
|
|
TRIAL_DAYS = 3
|
|
SUBSCRIPTION_MONTH_DAYS = 30
|
|
SUBSCRIPTION_PRICE_STARS = 100
|
|
|
|
|
|
def _days_until(sub: sqlite3.Row, now: datetime.datetime | None = None) -> int:
|
|
"""Сколько дней осталось от оплаты (триал или подписка). Отрицательно = истекло/отключено."""
|
|
if not sub:
|
|
return -1
|
|
if sub["is_active"] == 0:
|
|
return -1
|
|
# SQLite CURRENT_TIMESTAMP возвращает UTC-нейтральный datetime.
|
|
def utc(value: str) -> datetime.datetime:
|
|
return datetime.datetime.fromisoformat(value.replace("Z", "+00:00"))
|
|
|
|
if sub["last_paid_at"] is None:
|
|
# Триал: trial_started_at + 3 дня
|
|
trial_start = utc(sub["trial_started_at"])
|
|
if trial_start.tzinfo is not None:
|
|
trial_start = trial_start.astimezone(datetime.timezone.utc).replace(tzinfo=None)
|
|
expires = trial_start + datetime.timedelta(days=TRIAL_DAYS)
|
|
else:
|
|
expires = utc(sub["last_paid_at"])
|
|
if expires.tzinfo is not None:
|
|
expires = expires.astimezone(datetime.timezone.utc).replace(tzinfo=None)
|
|
now = now or datetime.datetime.now(datetime.timezone.utc)
|
|
now = now.replace(tzinfo=None)
|
|
delta = expires - now
|
|
# Полные дни, пока expires > now. При истечении (expires <= now) считаем -1.
|
|
return delta.days if expires > now else -1
|
|
|
|
|
|
def subscription_row(user_id: int, db_path: str | None = None) -> sqlite3.Row | None:
|
|
with closing(connect(db_path)) as connection:
|
|
return connection.execute(
|
|
"SELECT * FROM subscriptions WHERE telegram_user_id = ?", (user_id,)
|
|
).fetchone()
|
|
|
|
|
|
def ensure_subscription(user_id: int, db_path: str | None = None) -> None:
|
|
"""Создаёт запись подписки с триалом, если её нет."""
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO subscriptions (telegram_user_id) VALUES (?)", (user_id,)
|
|
)
|
|
|
|
|
|
def activate_subscription(user_id: int, months: int = 1, db_path: str | None = None) -> int:
|
|
"""Активирует/продлевает подписку на заданное число месяцев. Возвращает дни нового срока."""
|
|
with closing(connect(db_path)) as connection, connection:
|
|
sub = subscription_row(user_id, db_path)
|
|
if sub and sub["last_paid_at"]:
|
|
# Продление: отсчитываем от момента оплаты
|
|
now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
|
elif sub:
|
|
# Первый платёж после триала: отсчитываем ровно 30 дней от момента оплаты
|
|
now = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
|
|
else:
|
|
raise RuntimeError("Subscription record not found")
|
|
new_expires = now + datetime.timedelta(days=SUBSCRIPTION_MONTH_DAYS * months)
|
|
connection.execute(
|
|
"""
|
|
UPDATE subscriptions
|
|
SET last_paid_at = ?,
|
|
is_active = 1,
|
|
monthly_days = ?,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE telegram_user_id = ?
|
|
""",
|
|
(new_expires.isoformat(), SUBSCRIPTION_MONTH_DAYS * months, user_id),
|
|
)
|
|
return (new_expires - now).days
|
|
|
|
|
|
def deactivate_subscription(user_id: int, db_path: str | None = None) -> None:
|
|
"""Отключает подписку (период догорает до конца)."""
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute(
|
|
"UPDATE subscriptions SET is_active = 0 WHERE telegram_user_id = ?", (user_id,)
|
|
)
|
|
|
|
|
|
def is_subscribed(user_id: int, db_path: str | None = None) -> bool:
|
|
"""Правда ли подписка активна (срок не истёк)."""
|
|
sub = subscription_row(user_id, db_path)
|
|
if not sub:
|
|
return False
|
|
if not sub["is_active"]:
|
|
return False
|
|
return _days_until(sub) > 0
|
|
|
|
|
|
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 add_training_note(user_id: int, note: str, db_path: str | None = None) -> list[str]:
|
|
text = note.strip()
|
|
if not text:
|
|
return get_training_notes(user_id, db_path)
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,))
|
|
connection.execute(
|
|
"INSERT OR IGNORE INTO user_training_notes (telegram_user_id, note_text) VALUES (?, ?)",
|
|
(user_id, text),
|
|
)
|
|
rows = connection.execute(
|
|
"SELECT note_text FROM user_training_notes WHERE telegram_user_id = ? ORDER BY id",
|
|
(user_id,),
|
|
).fetchall()
|
|
return [str(row["note_text"]) for row in rows]
|
|
|
|
|
|
def get_training_notes(user_id: int, db_path: str | None = None) -> list[str]:
|
|
with closing(connect(db_path)) as connection:
|
|
rows = connection.execute(
|
|
"SELECT note_text FROM user_training_notes WHERE telegram_user_id = ? ORDER BY id",
|
|
(user_id,),
|
|
).fetchall()
|
|
return [str(row["note_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 = "",
|
|
operation_id: str | None = None,
|
|
duration_ms: int | None = None,
|
|
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, operation_id, duration_ms
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(user_id, stage[:100], model[:200], prompt_sha256[:64], request_json,
|
|
response_json, status_code, error[:2000], (operation_id or "")[:32], duration_ms),
|
|
)
|
|
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, operation_id, duration_ms, 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, operation_id, duration_ms, 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:
|
|
clear_user_data(user_id, db_path)
|
|
|
|
|
|
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 save_idea(user_id: int, idea: dict[str, str], db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("INSERT OR IGNORE INTO users (telegram_user_id) VALUES (?)", (user_id,))
|
|
connection.execute(
|
|
"INSERT INTO generated_ideas (telegram_user_id, title, angle, why) VALUES (?, ?, ?, ?)",
|
|
(user_id, idea["title"], idea["angle"], idea["why"]),
|
|
)
|
|
|
|
|
|
def get_idea_history(user_id: int, limit: int = 30, db_path: str | None = None) -> list[dict[str, str]]:
|
|
with closing(connect(db_path)) as connection:
|
|
rows = connection.execute(
|
|
"SELECT title, angle, why FROM generated_ideas WHERE telegram_user_id = ? ORDER BY id DESC LIMIT ?",
|
|
(user_id, max(1, min(limit, 100))),
|
|
).fetchall()
|
|
return [{"title": str(row["title"]), "angle": str(row["angle"]), "why": str(row["why"])} for row in rows]
|
|
|
|
|
|
def clear_user_data(user_id: int, db_path: str | None = None) -> None:
|
|
with closing(connect(db_path)) as connection, connection:
|
|
connection.execute("DELETE FROM user_training_notes WHERE telegram_user_id = ?", (user_id,))
|
|
connection.execute("DELETE FROM user_post_samples WHERE telegram_user_id = ?", (user_id,))
|
|
connection.execute("DELETE FROM generated_posts WHERE telegram_user_id = ?", (user_id,))
|
|
connection.execute("DELETE FROM generated_ideas WHERE telegram_user_id = ?", (user_id,))
|
|
connection.execute("DELETE FROM users WHERE telegram_user_id = ?", (user_id,))
|
|
|
|
|
|
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),
|
|
)
|