Discord link archiver bot with Gemini tags/summaries
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444"""Database models and connection management for Archivist."""
import asyncioimport loggingfrom datetime import datetime, timedelta, timezonefrom typing import Optional, List, Dict, Any
import asyncpgfrom pgvector.asyncpg import register_vector
logger = logging.getLogger("archivist.db")
# The vector extension must exist before the pool opens, because each pooled# connection runs `register_vector` in its init callback and that requires the# `vector` type to be resolvable at codec-registration time.SCHEMA_EXTENSION = "CREATE EXTENSION IF NOT EXISTS vector;"
SCHEMA = """-- ── Curated tag pool ──────────────────────────────────────────────────────────CREATE TABLE IF NOT EXISTS tags ( id SERIAL PRIMARY KEY, name TEXT UNIQUE NOT NULL, created_at TIMESTAMPTZ DEFAULT NOW());
-- ── Pending tag proposals (LLM new_tags) ─────────────────────────────────────-- Tags move from here into `tags` once they've been used enough times, or via-- an explicit `!promote-tag <name>` command from the owner.CREATE TABLE IF NOT EXISTS tag_proposals ( id SERIAL PRIMARY KEY, name TEXT UNIQUE NOT NULL, use_count INT NOT NULL DEFAULT 1, created_at TIMESTAMPTZ DEFAULT NOW(), last_used_at TIMESTAMPTZ DEFAULT NOW());
-- ── Archived links ────────────────────────────────────────────────────────────CREATE TABLE IF NOT EXISTS archived_links ( id SERIAL PRIMARY KEY, url TEXT NOT NULL, title TEXT, genre TEXT, gloss TEXT, summary TEXT, commentary TEXT, original_message TEXT, archive_url TEXT, archive_service TEXT, snapshot_path TEXT, embedding vector(768),
-- Discord metadata guild_id BIGINT, channel_id BIGINT NOT NULL, message_id BIGINT NOT NULL, author_id BIGINT NOT NULL, bot_message_id BIGINT,
-- Processing flags summary_skipped BOOLEAN DEFAULT FALSE, tags_skipped BOOLEAN DEFAULT FALSE, archive_skipped BOOLEAN DEFAULT FALSE, privacy_mode BOOLEAN DEFAULT FALSE,
-- Timestamps created_at TIMESTAMPTZ DEFAULT NOW(), archived_at TIMESTAMPTZ);
-- Full-text search column (generated, so it stays in sync automatically).-- Uses ALTER TABLE / ADD COLUMN IF NOT EXISTS so it's idempotent on existing DBs.ALTER TABLE archived_links ADD COLUMN IF NOT EXISTS fts tsvector GENERATED ALWAYS AS ( to_tsvector('english', coalesce(title, '') || ' ' || coalesce(gloss, '') || ' ' || coalesce(summary, '') || ' ' || coalesce(commentary, '') ) ) STORED;
CREATE TABLE IF NOT EXISTS link_tags ( link_id INTEGER REFERENCES archived_links(id) ON DELETE CASCADE, tag_id INTEGER REFERENCES tags(id) ON DELETE CASCADE, PRIMARY KEY (link_id, tag_id));
-- ── Key-value metadata store ──────────────────────────────────────────────────-- Used for operational bookmarks like `last_digest_sent_at`.CREATE TABLE IF NOT EXISTS meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL, updated_at TIMESTAMPTZ DEFAULT NOW());
-- ── Indexes ───────────────────────────────────────────────────────────────────-- `archived_links_url_unique` doubles as the url-lookup index AND enforces-- dedupe. On a pre-existing DB with duplicate URLs this will fail loudly —-- intentional: dedupe them by hand before re-running.DROP INDEX IF EXISTS idx_links_url;CREATE UNIQUE INDEX IF NOT EXISTS archived_links_url_unique ON archived_links(url);CREATE INDEX IF NOT EXISTS idx_links_channel ON archived_links(channel_id);CREATE INDEX IF NOT EXISTS idx_links_author ON archived_links(author_id);CREATE INDEX IF NOT EXISTS idx_links_created ON archived_links(created_at DESC);CREATE INDEX IF NOT EXISTS idx_links_fts ON archived_links USING GIN(fts);"""
class Database: def __init__(self, config: Dict[str, Any]): self.config = config self.pool: Optional[asyncpg.Pool] = None
async def connect(self): """Create connection pool and initialize schema.
Two-phase bootstrap so the pgvector codec can be registered on every pooled connection: first ensure the extension exists via a disposable connection, then open the pool with an init callback that registers the `vector` type. """ conn_kwargs = dict( host=self.config["host"], port=self.config["port"], database=self.config["database"], user=self.config["user"], password=self.config["password"], )
# Phase 1: ensure the vector extension exists. bootstrap = await asyncpg.connect(**conn_kwargs) try: await bootstrap.execute(SCHEMA_EXTENSION) finally: await bootstrap.close()
# Phase 2: open the pool with the vector codec registered per-connection. async def _init_conn(conn: asyncpg.Connection) -> None: await register_vector(conn)
self.pool = await asyncpg.create_pool( **conn_kwargs, min_size=2, max_size=10, init=_init_conn, ) async with self.pool.acquire() as conn: await conn.execute(SCHEMA) logger.info("Database connected and schema initialized")
async def close(self): if self.pool: await self.pool.close()
async def seed_tags(self, tags: List[str]): """Insert seed tags, ignoring duplicates.""" if not tags: return async with self.pool.acquire() as conn: await conn.executemany( "INSERT INTO tags (name) VALUES ($1) ON CONFLICT (name) DO NOTHING", [(t.lower().strip(),) for t in tags], ) logger.info(f"Seeded {len(tags)} tags")
async def get_tag_pool(self, proposal_threshold: int = 3) -> List[str]: """Return the curated tag pool plus any proposals used often enough.
`proposal_threshold` controls how many times a proposed tag must appear before it surfaces in the pool fed back to the LLM. The default of 3 prevents one-off LLM hallucinations from polluting the vocabulary. """ async with self.pool.acquire() as conn: curated = await conn.fetch("SELECT name FROM tags ORDER BY name") promoted = await conn.fetch( "SELECT name FROM tag_proposals WHERE use_count >= $1 ORDER BY name", proposal_threshold, ) names = {r["name"] for r in curated} | {r["name"] for r in promoted} return sorted(names)
async def record_proposals(self, tag_names: List[str]) -> None: """Record LLM-proposed tags in the proposals table.
Each call increments `use_count` for existing entries so the most frequently suggested tags naturally bubble up to the curated pool. """ if not tag_names: return async with self.pool.acquire() as conn: for name in tag_names: clean = name.lower().strip() if not clean: continue await conn.execute( """ INSERT INTO tag_proposals (name, use_count, last_used_at) VALUES ($1, 1, NOW()) ON CONFLICT (name) DO UPDATE SET use_count = tag_proposals.use_count + 1, last_used_at = NOW() """, clean, )
async def promote_tag(self, name: str) -> bool: """Move a proposal into the curated `tags` table.
Returns True if the tag was promoted (it existed in proposals or is new), False if it was already in the curated pool. """ clean = name.lower().strip() if not clean: return False async with self.pool.acquire() as conn: already = await conn.fetchval( "SELECT id FROM tags WHERE name = $1", clean ) if already is not None: return False await conn.execute( "INSERT INTO tags (name) VALUES ($1) ON CONFLICT (name) DO NOTHING", clean, ) # Remove from proposals if present (keeps the tables consistent). await conn.execute( "DELETE FROM tag_proposals WHERE name = $1", clean ) return True
async def ensure_tags(self, tag_names: List[str]) -> Dict[str, int]: """Ensure tags exist, create if needed. Returns {name: id} mapping.""" result = {} async with self.pool.acquire() as conn: for name in tag_names: clean_name = name.lower().strip() if not clean_name: continue row = await conn.fetchrow( """ INSERT INTO tags (name) VALUES ($1) ON CONFLICT (name) DO UPDATE SET name = EXCLUDED.name RETURNING id """, clean_name, ) if row: result[clean_name] = row["id"] return result
async def save_link( self, url: str, title: Optional[str], genre: Optional[str], gloss: Optional[str], summary: Optional[str], commentary: Optional[str], original_message: Optional[str], archive_url: Optional[str], archive_service: Optional[str], snapshot_path: Optional[str], embedding: Optional[List[float]], guild_id: Optional[int], channel_id: int, message_id: int, author_id: int, bot_message_id: Optional[int], tags: Optional[List[str]], proposed_tags: Optional[List[str]] = None, summary_skipped: bool = False, tags_skipped: bool = False, archive_skipped: bool = False, privacy_mode: bool = False, ) -> Optional[int]: """Save an archived link and its curated tags.
`tags` are the curated tags (from the known pool) — they get saved to `link_tags`. `proposed_tags` are new LLM suggestions — they go to `tag_proposals` for counting; once a proposal's use_count reaches the threshold it surfaces in the tag pool fed back to the LLM.
Returns the link ID, or None if a concurrent task already saved this URL (the pre-flight dedupe in `bot.process_link` closes the gap to microseconds, but `ON CONFLICT` makes the write truly idempotent). """ async with self.pool.acquire() as conn: async with conn.transaction(): link_id = await conn.fetchval( """ INSERT INTO archived_links ( url, title, genre, gloss, summary, commentary, original_message, archive_url, archive_service, snapshot_path, embedding, guild_id, channel_id, message_id, author_id, bot_message_id, summary_skipped, tags_skipped, archive_skipped, privacy_mode, archived_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21 ) ON CONFLICT (url) DO NOTHING RETURNING id """, url, title, genre, gloss, summary, commentary, original_message, archive_url, archive_service, snapshot_path, embedding, guild_id, channel_id, message_id, author_id, bot_message_id, summary_skipped, tags_skipped, archive_skipped, privacy_mode, datetime.now(timezone.utc) if archive_url else None, )
if link_id is None: logger.info(f"Skipped duplicate save: {url}") return None
if tags: tag_mapping = await self.ensure_tags(tags) if tag_mapping: await conn.executemany( "INSERT INTO link_tags (link_id, tag_id) VALUES ($1, $2)" " ON CONFLICT DO NOTHING", [(link_id, tid) for tid in tag_mapping.values()], )
# Record proposals outside the transaction — a failure here is non-fatal. if proposed_tags: try: await self.record_proposals(proposed_tags) except Exception: logger.exception("Failed to record tag proposals (non-fatal)")
logger.info(f"Saved link {link_id}: {url}") return link_id
async def find_existing_link(self, url: str) -> Optional[Dict[str, Any]]: """Return the existing archive row for `url`, or None.
Used for the pre-flight dedupe check — we short-circuit before fetch / LLM / archive work if we've already archived this URL. """ async with self.pool.acquire() as conn: row = await conn.fetchrow( """ SELECT id, url, title, channel_id, guild_id, bot_message_id, created_at FROM archived_links WHERE url = $1 ORDER BY id ASC LIMIT 1 """, url, ) return dict(row) if row else None
async def search_links(self, query_embedding: List[float], limit: int = 5) -> List[Dict[str, Any]]: """Perform vector similarity search.""" async with self.pool.acquire() as conn: rows = await conn.fetch( """ SELECT url, title, gloss, summary, archive_url FROM archived_links ORDER BY embedding <=> $1 LIMIT $2 """, query_embedding, limit ) return [dict(r) for r in rows]
async def get_daily_links(self) -> List[Dict[str, Any]]: """Return links from the last 24 hours (fallback for callers that don't use the windowed version).""" since = datetime.now(timezone.utc) - timedelta(hours=24) return await self.get_links_since(since)
async def get_links_since(self, since: datetime) -> List[Dict[str, Any]]: """Return links archived after *since* with user/channel metadata.""" async with self.pool.acquire() as conn: rows = await conn.fetch( """ SELECT url, title, genre, gloss, summary, commentary, channel_id, author_id, created_at FROM archived_links WHERE created_at > $1 ORDER BY created_at ASC """, since, ) return [dict(r) for r in rows]
# ── Meta key-value store ─────────────────────────────────────────────────
async def get_meta(self, key: str) -> Optional[str]: """Return the stored string value for *key*, or None if not set.""" async with self.pool.acquire() as conn: row = await conn.fetchrow( "SELECT value FROM meta WHERE key = $1", key ) return row["value"] if row else None
async def set_meta(self, key: str, value: str) -> None: """Upsert a key-value pair in the meta table.""" async with self.pool.acquire() as conn: await conn.execute( """ INSERT INTO meta (key, value, updated_at) VALUES ($1, $2, NOW()) ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, updated_at = NOW() """, key, value, )
# ── Full-text search ─────────────────────────────────────────────────────
async def search_links_text( self, query: str, limit: int = 5 ) -> List[Dict[str, Any]]: """Postgres full-text search over title, gloss, summary, and commentary.
Returns results ranked by `ts_rank`, most relevant first. Returns an empty list (not an error) if the `fts` column doesn't exist yet on a pre-migration DB — callers can fall back to vector search in that case. """ async with self.pool.acquire() as conn: try: rows = await conn.fetch( """ SELECT url, title, gloss, summary, archive_url, ts_rank(fts, plainto_tsquery('english', $1)) AS rank FROM archived_links WHERE fts @@ plainto_tsquery('english', $1) ORDER BY rank DESC LIMIT $2 """, query, limit, ) except Exception: logger.exception("Full-text search failed") return [] return [dict(r) for r in rows]