Files
copykar/db/repository.py

1052 lines
45 KiB
Python

import json
import asyncpg
from datetime import datetime, timezone
from typing import List, Optional, Dict, Any, Tuple
from db.models import SourceChannel, SourceWebsite, TargetChannel, Post, Setting, AILog, AIProviderProfile, ChannelCategory
from db.database import get_db_pool
def _parse_post_row(row: asyncpg.Record) -> Post:
data = dict(row)
if isinstance(data.get("published_to"), str):
try:
data["published_to"] = json.loads(data["published_to"])
except Exception:
data["published_to"] = []
elif data.get("published_to") is None:
data["published_to"] = []
return Post(**data)
def _parse_source_website_row(row: asyncpg.Record) -> SourceWebsite:
data = dict(row)
if isinstance(data.get("api_config"), str):
try:
data["api_config"] = json.loads(data["api_config"])
except Exception:
data["api_config"] = {}
elif data.get("api_config") is None:
data["api_config"] = {}
return SourceWebsite(**data)
class Repository:
def __init__(self, dsn: Optional[str] = None):
self.dsn = dsn
async def _get_pool(self) -> asyncpg.Pool:
if self.dsn:
return await get_db_pool(self.dsn)
return await get_db_pool()
# --- Source Channels ---
async def add_source(self, channel_id: int, title: Optional[str] = None, username: Optional[str] = None) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
INSERT INTO sources (channel_id, title, username)
VALUES ($1, $2, $3)
ON CONFLICT(channel_id) DO UPDATE SET
title = EXCLUDED.title,
username = EXCLUDED.username,
is_active = TRUE
RETURNING id;
""",
channel_id, title, username,
)
return row["id"]
async def get_active_sources(self) -> List[SourceChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE ORDER BY id ASC;")
return [SourceChannel(**dict(r)) for r in rows]
async def get_source_by_channel_id(self, channel_id: int) -> Optional[SourceChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM sources WHERE channel_id = $1;", channel_id)
return SourceChannel(**dict(row)) if row else None
async def get_source_by_id(self, source_id: int) -> Optional[SourceChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM sources WHERE id = $1;", source_id)
return SourceChannel(**dict(row)) if row else None
async def update_source_context_count(self, source_id: int, count: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE sources SET context_message_count = $1 WHERE id = $2;", max(0, count), source_id)
async def delete_source(self, source_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE sources SET is_active = FALSE WHERE id = $1;", source_id)
# --- Source Websites ---
async def add_source_website(
self,
name: str,
url: str,
category_id: Optional[int] = None,
check_interval_min: int = 30,
auto_reanalyze_hours: int = 24,
api_config: Optional[Dict[str, Any]] = None,
custom_instructions: str = ""
) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
cfg_json = json.dumps(api_config or {})
row = await conn.fetchrow(
"""
INSERT INTO source_websites (name, url, category_id, check_interval_min, auto_reanalyze_hours, api_config, custom_instructions)
VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7)
ON CONFLICT(url) DO UPDATE SET
name = EXCLUDED.name,
category_id = COALESCE(EXCLUDED.category_id, source_websites.category_id),
custom_instructions = COALESCE(NULLIF(EXCLUDED.custom_instructions, ''), source_websites.custom_instructions),
is_active = TRUE
RETURNING id;
""",
name, url, category_id, check_interval_min, auto_reanalyze_hours, cfg_json, custom_instructions
)
return row["id"]
async def get_active_source_websites(self) -> List[SourceWebsite]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch("SELECT * FROM source_websites WHERE is_active = TRUE ORDER BY id ASC;")
return [_parse_source_website_row(r) for r in rows]
async def get_source_websites(self) -> List[SourceWebsite]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch("SELECT * FROM source_websites WHERE is_active = TRUE ORDER BY id ASC;")
return [_parse_source_website_row(r) for r in rows]
async def get_source_website_by_id(self, site_id: int) -> Optional[SourceWebsite]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM source_websites WHERE id = $1;", site_id)
return _parse_source_website_row(row) if row else None
async def get_source_websites_by_category(self, category_id: Optional[int]) -> List[SourceWebsite]:
pool = await self._get_pool()
async with pool.acquire() as conn:
if category_id is None:
rows = await conn.fetch("SELECT * FROM source_websites WHERE is_active = TRUE AND category_id IS NULL ORDER BY id ASC;")
else:
rows = await conn.fetch("SELECT * FROM source_websites WHERE is_active = TRUE AND category_id = $1 ORDER BY id ASC;", category_id)
return [_parse_source_website_row(r) for r in rows]
async def update_source_website_category(self, site_id: int, category_id: Optional[int]) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE source_websites SET category_id = $1 WHERE id = $2;", category_id, site_id)
async def update_source_website_api_config(self, site_id: int, api_config: Dict[str, Any]) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"""
UPDATE source_websites
SET api_config = $1::jsonb,
last_reanalyzed_at = CURRENT_TIMESTAMP,
last_error = NULL,
last_error_at = NULL
WHERE id = $2;
""",
json.dumps(api_config), site_id
)
async def update_source_website_fetch_status(self, site_id: int, error: Optional[str] = None) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
if error:
await conn.execute(
"""
UPDATE source_websites
SET last_error = $1, last_error_at = CURRENT_TIMESTAMP
WHERE id = $2;
""",
error, site_id
)
else:
await conn.execute(
"""
UPDATE source_websites
SET last_fetched_at = CURRENT_TIMESTAMP, last_error = NULL, last_error_at = NULL
WHERE id = $1;
""",
site_id
)
async def update_source_website_interval(self, site_id: int, interval_min: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE source_websites SET check_interval_min = $1 WHERE id = $2;", max(1, interval_min), site_id)
async def update_source_website_reanalyze_hours(self, site_id: int, hours: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE source_websites SET auto_reanalyze_hours = $1 WHERE id = $2;", max(0, hours), site_id)
async def update_source_website_custom_instructions(self, site_id: int, instructions: str) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE source_websites SET custom_instructions = $1 WHERE id = $2;", instructions.strip(), site_id)
async def delete_source_website(self, site_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE source_websites SET is_active = FALSE WHERE id = $1;", site_id)
async def update_target_context_count(self, target_id: int, count: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE targets SET context_message_count = $1 WHERE id = $2;", max(0, count), target_id)
async def get_recent_target_posts(
self,
target_id: int,
limit: int = 10,
exclude_post_id: Optional[int] = None
) -> List[Post]:
if limit <= 0:
return []
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT * FROM posts
WHERE is_deleted = FALSE
AND ($2::BIGINT IS NULL OR id != $2)
AND (
target_channel_id = $1
OR (
published_to IS NOT NULL
AND jsonb_typeof(published_to) = 'array'
AND EXISTS (
SELECT 1 FROM jsonb_array_elements(published_to) elem
WHERE (elem->>'target_id')::bigint = $1
)
)
)
ORDER BY COALESCE(published_at, created_at) DESC, id DESC
LIMIT $3;
""",
target_id, exclude_post_id, limit
)
posts = [_parse_post_row(r) for r in rows]
posts.reverse()
return posts
async def delete_target(self, target_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE targets SET is_active = FALSE WHERE id = $1;", target_id)
# --- Target Channels ---
async def add_target(
self,
channel_id: int,
title: Optional[str] = None,
username: Optional[str] = None,
post_interval_min: int = 30,
personality: str = "",
custom_footer: str = ""
) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
INSERT INTO targets (channel_id, title, username, post_interval_min, personality, custom_footer)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT(channel_id) DO UPDATE SET
title = EXCLUDED.title,
username = EXCLUDED.username,
post_interval_min = EXCLUDED.post_interval_min,
is_active = TRUE
RETURNING id;
""",
channel_id, title, username, post_interval_min, personality, custom_footer,
)
return row["id"]
async def update_target_personality(self, target_id: int, personality: str, custom_footer: Optional[str] = None) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
if custom_footer is not None:
await conn.execute(
"UPDATE targets SET personality = $1, custom_footer = $2 WHERE id = $3;",
personality, custom_footer, target_id
)
else:
await conn.execute(
"UPDATE targets SET personality = $1 WHERE id = $2;",
personality, target_id
)
async def update_target_footer(self, target_id: int, custom_footer: str) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET custom_footer = $1 WHERE id = $2;",
custom_footer, target_id
)
async def update_target_language(self, target_id: int, language: str) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET language = $1 WHERE id = $2;",
language, target_id
)
async def update_target_custom_prompt(self, target_id: int, custom_prompt: str) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET custom_prompt = $1 WHERE id = $2;",
custom_prompt, target_id
)
async def update_target_dispatch_order(self, target_id: int, dispatch_order: str) -> None:
order = "random" if str(dispatch_order).strip().lower() == "random" else "order"
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET dispatch_order = $1 WHERE id = $2;",
order, target_id
)
async def update_target_schedule(
self,
target_id: int,
post_interval_min: Optional[int] = None,
sleep_start_hour: Optional[int] = None,
sleep_end_hour: Optional[int] = None,
is_sleep_enabled: Optional[bool] = None
) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
target = await self.get_target_by_id(target_id)
if not target:
return
new_interval = post_interval_min if post_interval_min is not None else target.post_interval_min
new_start = sleep_start_hour if sleep_start_hour is not None else target.sleep_start_hour
new_end = sleep_end_hour if sleep_end_hour is not None else target.sleep_end_hour
new_enabled = is_sleep_enabled if is_sleep_enabled is not None else target.is_sleep_enabled
await conn.execute(
"""
UPDATE targets
SET post_interval_min = $1, sleep_start_hour = $2, sleep_end_hour = $3, is_sleep_enabled = $4
WHERE id = $5;
""",
new_interval, new_start, new_end, new_enabled, target_id
)
async def get_active_targets(self) -> List[TargetChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE ORDER BY id ASC;")
return [TargetChannel(**dict(r)) for r in rows]
async def get_target_by_id(self, target_id: int) -> Optional[TargetChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM targets WHERE id = $1;", target_id)
return TargetChannel(**dict(row)) if row else None
async def set_target_auto_sources(self, target_id: int, source_channel_ids: List[int]) -> None:
"""Replace the set of source channels auto-routed into this target."""
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET auto_source_ids = $1::bigint[] WHERE id = $2;",
sorted(set(source_channel_ids)), target_id,
)
async def toggle_target_auto_source(self, target_id: int, source_channel_id: int) -> bool:
"""Add or remove one source from a target's auto-route list. Returns the new state."""
target = await self.get_target_by_id(target_id)
if not target:
return False
current = set(target.auto_source_ids or [])
enabled = source_channel_id not in current
if enabled:
current.add(source_channel_id)
else:
current.discard(source_channel_id)
await self.set_target_auto_sources(target_id, list(current))
return enabled
async def get_targets_auto_routed_from(self, source_channel_id: int) -> List[TargetChannel]:
"""Active targets that have subscribed to this source channel."""
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT * FROM targets
WHERE is_active = TRUE AND $1 = ANY(auto_source_ids)
ORDER BY id ASC;
""",
source_channel_id,
)
return [TargetChannel(**dict(r)) for r in rows]
async def update_target_last_post(self, target_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE targets SET last_post_time = CURRENT_TIMESTAMP WHERE id = $1;",
target_id,
)
# --- Posts & Multi-Channel Dispatch ---
async def create_raw_post(
self,
source_channel_id: int,
source_message_id: int,
raw_text: Optional[str],
media_path: Optional[str] = None,
media_type: Optional[str] = None,
content_hash: Optional[str] = None,
tags: Optional[List[str]] = None,
subject: Optional[str] = None,
is_duplicate: bool = False,
duplicate_of_id: Optional[int] = None,
similarity_reason: Optional[str] = None,
source_created_at: Optional[datetime] = None,
) -> Optional[int]:
pool = await self._get_pool()
async with pool.acquire() as conn:
try:
row = await conn.fetchrow(
"""
INSERT INTO posts (
source_channel_id, source_message_id, raw_text, media_path,
media_type, content_hash, tags, subject, is_duplicate, duplicate_of_id, similarity_reason, source_created_at, status
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, 'pending_review')
RETURNING id;
""",
source_channel_id,
source_message_id,
raw_text,
media_path,
media_type,
content_hash,
tags or [],
subject,
is_duplicate,
duplicate_of_id,
similarity_reason,
source_created_at,
)
return row["id"] if row else None
except asyncpg.UniqueViolationError:
return None
async def get_post_by_id(self, post_id: int) -> Optional[Post]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM posts WHERE id = $1;", post_id)
return _parse_post_row(row) if row else None
async def get_recent_source_posts(
self,
source_channel_id: int,
limit: int = 5,
exclude_post_id: Optional[int] = None
) -> List[Post]:
"""Fetch the most recent posts from this source channel for narrative context."""
if limit <= 0:
return []
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT * FROM posts
WHERE source_channel_id = $1
AND is_deleted = FALSE
AND ($2::bigint IS NULL OR id <> $2)
ORDER BY COALESCE(source_created_at, created_at) DESC, id DESC
LIMIT $3;
""",
source_channel_id, exclude_post_id, limit
)
# Return in chronological order so the AI sees the natural progression (earliest to latest)
posts = [_parse_post_row(r) for r in rows]
posts.reverse()
return posts
async def find_candidate_posts_by_tags(
self,
tags: List[str],
exclude_post_id: Optional[int] = None,
limit: int = 10
) -> List[Post]:
"""Find recent posts that share at least one tag."""
if not tags:
return []
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT * FROM posts
WHERE tags && $1::text[]
AND is_deleted = FALSE
AND ($2::bigint IS NULL OR id <> $2)
ORDER BY id DESC
LIMIT $3;
""",
tags, exclude_post_id, limit
)
return [_parse_post_row(r) for r in rows]
async def find_duplicate_post(self, content_hash: Optional[str], exclude_post_id: Optional[int] = None) -> Optional[Post]:
"""Return the earliest post already carrying this content hash, if any."""
if not content_hash:
return None
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT * FROM posts
WHERE content_hash = $1
AND is_deleted = FALSE
AND ($2::bigint IS NULL OR id <> $2)
ORDER BY id ASC
LIMIT 1;
""",
content_hash, exclude_post_id,
)
return _parse_post_row(row) if row else None
async def count_posts_by_status(self, status: str) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval("SELECT COUNT(*) FROM posts WHERE status = $1;", status) or 0
async def count_posts_from_source(self, source_channel_id: int) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval(
"SELECT COUNT(*) FROM posts WHERE source_channel_id = $1 AND is_deleted = FALSE;",
source_channel_id,
) or 0
async def get_posts_by_status(self, status: str, limit: int = 20) -> List[Post]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"SELECT * FROM posts WHERE status = $1 ORDER BY id ASC LIMIT $2;",
status, limit,
)
return [_parse_post_row(r) for r in rows]
async def get_unreviewed_posts(self, limit: int = 50) -> List[Post]:
"""Pending posts that never made it onto a review card in the admin channel."""
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT * FROM posts
WHERE status = 'pending_review'
AND is_deleted = FALSE
AND review_message_id IS NULL
ORDER BY id ASC
LIMIT $1;
""",
limit,
)
return [_parse_post_row(r) for r in rows]
async def update_review_message_id(self, post_id: int, review_message_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE posts SET review_message_id = $1 WHERE id = $2;",
review_message_id, post_id,
)
async def _upsert_published_entry(
self,
post_id: int,
target_id: int,
target_title: str,
published_at: Optional[datetime],
mark_published: bool,
) -> None:
"""Add or stamp this target's entry in posts.published_to.
The list is read, edited and written back as a single parameterised jsonb value so
that titles containing quotes or backslashes cannot corrupt the document.
"""
pool = await self._get_pool()
async with pool.acquire() as conn:
async with conn.transaction():
row = await conn.fetchrow("SELECT published_to FROM posts WHERE id = $1 FOR UPDATE;", post_id)
if not row:
return
entries = row["published_to"]
if isinstance(entries, str):
try:
entries = json.loads(entries)
except Exception:
entries = []
if not isinstance(entries, list):
entries = []
stamp = published_at.isoformat() if published_at else None
for entry in entries:
if isinstance(entry, dict) and entry.get("target_id") == target_id:
entry["target_title"] = target_title
if stamp:
entry["published_at"] = stamp
break
else:
entries.append({
"target_id": target_id,
"target_title": target_title,
"published_at": stamp,
})
if mark_published:
await conn.execute(
"""
UPDATE posts
SET published_to = $1::jsonb,
status = 'published',
published_at = COALESCE(published_at, CURRENT_TIMESTAMP)
WHERE id = $2;
""",
json.dumps(entries, ensure_ascii=False), post_id,
)
else:
await conn.execute(
"UPDATE posts SET published_to = $1::jsonb WHERE id = $2;",
json.dumps(entries, ensure_ascii=False), post_id,
)
async def record_post_queued_to_target(self, post_id: int, target_id: int, target_title: str) -> None:
"""Note that a post is waiting in a target's delivery queue. Status is left untouched."""
await self._upsert_published_entry(post_id, target_id, target_title, published_at=None, mark_published=False)
async def record_post_published_to_target(self, post_id: int, target_id: int, target_title: str) -> None:
"""Stamp the post as actually delivered to a target channel."""
await self._upsert_published_entry(
post_id, target_id, target_title,
published_at=datetime.now(timezone.utc), mark_published=True,
)
# --- Error tracking ---
async def get_open_error_summary(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Unresolved errors grouped by service + exception type, newest group first."""
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT service_name, error_type, COUNT(*) AS occurrences,
MAX(created_at) AS last_seen,
(ARRAY_AGG(error_message ORDER BY created_at DESC))[1] AS last_message
FROM error_logs
WHERE resolved = FALSE
GROUP BY service_name, error_type
ORDER BY MAX(created_at) DESC
LIMIT $1;
""",
limit,
)
return [dict(r) for r in rows]
async def count_open_errors(self) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval("SELECT COUNT(*) FROM error_logs WHERE resolved = FALSE;") or 0
async def resolve_errors(
self,
service_name: Optional[str] = None,
error_type: Optional[str] = None,
note: Optional[str] = None,
) -> int:
"""Mark matching unresolved errors as fixed. Both filters None resolves everything."""
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval(
"""
WITH updated AS (
UPDATE error_logs
SET resolved = TRUE,
resolved_at = CURRENT_TIMESTAMP,
resolved_note = COALESCE($3, resolved_note)
WHERE resolved = FALSE
AND ($1::text IS NULL OR service_name = $1)
AND ($2::text IS NULL OR error_type = $2)
RETURNING 1
)
SELECT COUNT(*) FROM updated;
""",
service_name, error_type, note,
) or 0
async def reject_post(self, post_id: int, rejection_reason: Optional[str] = None) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
if rejection_reason:
await conn.execute(
"UPDATE posts SET status = 'rejected', rejection_reason = $1 WHERE id = $2;",
rejection_reason, post_id
)
else:
await conn.execute("UPDATE posts SET status = 'rejected' WHERE id = $1;", post_id)
async def soft_delete_post(self, post_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE posts SET is_deleted = TRUE, status = 'deleted' WHERE id = $1;", post_id)
# --- System Settings ---
async def get_setting(self, key: str, default: Optional[str] = None) -> Optional[str]:
pool = await self._get_pool()
async with pool.acquire() as conn:
val = await conn.fetchval("SELECT value FROM settings WHERE key = $1;", key)
return val if val is not None else default
async def set_setting(self, key: str, value: str, description: Optional[str] = None) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO settings (key, value, description)
VALUES ($1, $2, $3)
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, description = COALESCE(EXCLUDED.description, settings.description);
""",
key, value, description
)
async def get_all_settings(self) -> Dict[str, str]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch("SELECT key, value FROM settings;")
return {r["key"]: r["value"] for r in rows}
# --- AI Logs ---
async def record_ai_log(
self,
action_name: str,
provider: str,
model: str,
prompt: str,
system_prompt: Optional[str] = None,
response_text: Optional[str] = None,
duration_sec: float = 0.0,
status: str = "success",
error_message: Optional[str] = None,
) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval(
"""
INSERT INTO ai_logs (action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
RETURNING id;
""",
action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message
)
async def get_recent_ai_logs(self, limit: int = 10) -> List[AILog]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT id, action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message,
to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at
FROM ai_logs
ORDER BY id DESC
LIMIT $1;
""",
limit
)
return [AILog(**dict(r)) for r in rows]
async def get_ai_log_by_id(self, log_id: int) -> Optional[AILog]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT id, action_name, provider, model, prompt, system_prompt, response_text, duration_sec, status, error_message,
to_char(created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at
FROM ai_logs
WHERE id = $1;
""",
log_id
)
return AILog(**dict(row)) if row else None
# --- AI Providers Management ---
async def ensure_default_providers(self) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
count = await conn.fetchval("SELECT COUNT(*) FROM ai_providers;")
if count == 0:
await conn.execute(
"""
INSERT INTO ai_providers (name, provider_type, model, base_url, api_key, reasoning_effort, is_active)
VALUES
('AGY (سرور داخلی)', 'agy', 'antigravity', 'http://host.docker.internal:8088/v1', '', '', TRUE),
('OpenRouter / OpenAI', 'openai', 'google/gemini-3.5-flash', 'https://openrouter.ai/api/v1', '', '', FALSE),
('Google Gemini Direct', 'gemini', 'gemini-1.5-flash', 'https://generativelanguage.googleapis.com/v1beta', '', '', FALSE);
"""
)
async def add_provider_profile(
self,
name: str,
provider_type: str,
model: str,
base_url: str = "",
api_key: str = "",
reasoning_effort: str = "",
is_active: bool = False,
supports_vision: bool = False
) -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
if is_active:
await conn.execute("UPDATE ai_providers SET is_active = FALSE;")
return await conn.fetchval(
"""
INSERT INTO ai_providers (name, provider_type, model, base_url, api_key, reasoning_effort, is_active, supports_vision)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
RETURNING id;
""",
name, provider_type, model, base_url, api_key, reasoning_effort, is_active, supports_vision
)
async def get_provider_profiles(self) -> List[AIProviderProfile]:
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active,
p.fallback_provider_id, p.supports_vision,
fb.name as fallback_provider_name,
to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at
FROM ai_providers p
LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id
ORDER BY p.id ASC;
"""
)
return [AIProviderProfile(**dict(r)) for r in rows]
async def get_provider_profile_by_id(self, profile_id: int) -> Optional[AIProviderProfile]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active,
p.fallback_provider_id, p.supports_vision,
fb.name as fallback_provider_name,
to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at
FROM ai_providers p
LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id
WHERE p.id = $1;
""",
profile_id
)
return AIProviderProfile(**dict(row)) if row else None
async def get_active_provider_profile(self) -> Optional[AIProviderProfile]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
SELECT p.id, p.name, p.provider_type, p.model, p.base_url, p.api_key, p.reasoning_effort, p.is_active,
p.fallback_provider_id, p.supports_vision,
fb.name as fallback_provider_name,
to_char(p.created_at, 'YYYY-MM-DD HH24:MI:SS') as created_at
FROM ai_providers p
LEFT JOIN ai_providers fb ON p.fallback_provider_id = fb.id
WHERE p.is_active = TRUE
ORDER BY p.id ASC
LIMIT 1;
"""
)
return AIProviderProfile(**dict(row)) if row else None
async def set_active_provider_profile(self, profile_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
async with conn.transaction():
await conn.execute("UPDATE ai_providers SET is_active = FALSE;")
await conn.execute("UPDATE ai_providers SET is_active = TRUE WHERE id = $1;", profile_id)
async def update_provider_fallback(self, profile_id: int, fallback_provider_id: Optional[int]) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE ai_providers SET fallback_provider_id = $1 WHERE id = $2;",
fallback_provider_id, profile_id
)
async def update_provider_vision(self, profile_id: int, supports_vision: bool) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE ai_providers SET supports_vision = $1 WHERE id = $2;",
supports_vision, profile_id
)
async def update_provider_profile(
self,
profile_id: int,
name: Optional[str] = None,
model: Optional[str] = None,
base_url: Optional[str] = None,
api_key: Optional[str] = None,
reasoning_effort: Optional[str] = None,
) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"""
UPDATE ai_providers
SET name = COALESCE($2, name),
model = COALESCE($3, model),
base_url = COALESCE($4, base_url),
api_key = COALESCE($5, api_key),
reasoning_effort = COALESCE($6, reasoning_effort)
WHERE id = $1;
""",
profile_id, name, model, base_url, api_key, reasoning_effort
)
async def update_provider_reasoning_effort(self, profile_id: int, reasoning_effort: str) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"UPDATE ai_providers SET reasoning_effort = $1 WHERE id = $2;",
reasoning_effort, profile_id
)
async def delete_provider_profile(self, profile_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM ai_providers WHERE id = $1;", profile_id)
# --- Channel Categories ---
async def create_category(self, name: str, cat_type: str = "both", description: str = "") -> int:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow(
"""
INSERT INTO channel_categories (name, type, description)
VALUES ($1, $2, $3)
RETURNING id;
""",
name, cat_type, description
)
return row["id"]
async def get_categories(self, cat_type: Optional[str] = None) -> List[ChannelCategory]:
pool = await self._get_pool()
async with pool.acquire() as conn:
if cat_type:
rows = await conn.fetch(
"SELECT * FROM channel_categories WHERE type = $1 OR type = 'both' ORDER BY id ASC;",
cat_type
)
else:
rows = await conn.fetch("SELECT * FROM channel_categories ORDER BY id ASC;")
return [ChannelCategory(**dict(r)) for r in rows]
async def get_category_by_id(self, cat_id: int) -> Optional[ChannelCategory]:
pool = await self._get_pool()
async with pool.acquire() as conn:
row = await conn.fetchrow("SELECT * FROM channel_categories WHERE id = $1;", cat_id)
return ChannelCategory(**dict(row)) if row else None
async def update_category(self, cat_id: int, name: Optional[str] = None, description: Optional[str] = None, cat_type: Optional[str] = None) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute(
"""
UPDATE channel_categories
SET name = COALESCE($2, name),
description = COALESCE($3, description),
type = COALESCE($4, type)
WHERE id = $1;
""",
cat_id, name, description, cat_type
)
async def delete_category(self, cat_id: int) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE sources SET category_id = NULL WHERE category_id = $1;", cat_id)
await conn.execute("UPDATE targets SET category_id = NULL WHERE category_id = $1;", cat_id)
await conn.execute("DELETE FROM channel_categories WHERE id = $1;", cat_id)
async def set_source_category(self, source_id: int, category_id: Optional[int]) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE sources SET category_id = $1 WHERE id = $2;", category_id, source_id)
async def set_target_category(self, target_id: int, category_id: Optional[int]) -> None:
pool = await self._get_pool()
async with pool.acquire() as conn:
await conn.execute("UPDATE targets SET category_id = $1 WHERE id = $2;", category_id, target_id)
async def get_sources_by_category(self, category_id: Optional[int]) -> List[SourceChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
if category_id is None:
rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE AND category_id IS NULL ORDER BY id ASC;")
else:
rows = await conn.fetch("SELECT * FROM sources WHERE is_active = TRUE AND category_id = $1 ORDER BY id ASC;", category_id)
return [SourceChannel(**dict(r)) for r in rows]
async def get_targets_by_category(self, category_id: Optional[int]) -> List[TargetChannel]:
pool = await self._get_pool()
async with pool.acquire() as conn:
if category_id is None:
rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE AND category_id IS NULL ORDER BY id ASC;")
else:
rows = await conn.fetch("SELECT * FROM targets WHERE is_active = TRUE AND category_id = $1 ORDER BY id ASC;", category_id)
return [TargetChannel(**dict(r)) for r in rows]
async def get_category_channel_counts(self, category_id: int) -> Dict[str, int]:
pool = await self._get_pool()
async with pool.acquire() as conn:
src_count = await conn.fetchval("SELECT COUNT(*) FROM sources WHERE is_active = TRUE AND category_id = $1;", category_id) or 0
trg_count = await conn.fetchval("SELECT COUNT(*) FROM targets WHERE is_active = TRUE AND category_id = $1;", category_id) or 0
return {"sources": src_count, "targets": trg_count}
# --- System Global Operational State ---
async def is_system_paused(self) -> bool:
val = await self.get_setting("system_is_paused", "false")
return str(val).strip().lower() in ("true", "1", "yes")
async def set_system_paused(self, paused: bool) -> None:
await self.set_setting(
"system_is_paused",
"true" if paused else "false",
description="Global emergency operational pause"
)