homelab-brain/chatlogger/query.py

"""CLI access to the chatlogger SQLite database.

Usage:
    chatlogger recent [--chat <id>] [-n 50]
    chatlogger search "query string" [--chat <id>] [-n 30]
    chatlogger stats
    chatlogger chats
    chatlogger since <minutes> [--chat <id>]
    chatlogger export [--chat <id>] [--since-ts N]   # JSONL to stdout
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

from store import Store, format_messages

DEFAULT_DB = os.environ.get("ARAKAWA_BOT_DB", "/opt/chatlogger/chats.db")
TZ_OFFSET_HOURS = int(os.environ.get("ARAKAWA_BOT_TZ_OFFSET_HOURS", "7"))


def cmd_recent(args, store: Store) -> int:
    msgs = store.recent(chat_id=args.chat, limit=args.n)
    if not msgs:
        print("(no messages)")
        return 0
    print(format_messages(msgs, tz_offset_hours=TZ_OFFSET_HOURS))
    return 0


def cmd_search(args, store: Store) -> int:
    msgs = store.search(query=args.query, chat_id=args.chat, limit=args.n)
    if not msgs:
        print(f"(no matches for {args.query!r})")
        return 0
    print(format_messages(msgs, tz_offset_hours=TZ_OFFSET_HOURS))
    return 0


def cmd_stats(args, store: Store) -> int:
    s = store.stats()
    print(f"Total messages: {s['total_messages']}")
    print(f"Total edits:    {s['total_edits']}")
    print(f"DB size:        {s['db_size_bytes'] / 1024:.1f} KB")
    print()
    print("Per chat:")
    for c in s["chats"]:
        last = time.strftime("%Y-%m-%d %H:%M", time.gmtime(c["last_ts"] + TZ_OFFSET_HOURS * 3600))
        print(f"  {c['chat_id']:>20}  {c['n']:>6}  last={last}  title={c['title'] or '?'}")
    return 0


def cmd_chats(args, store: Store) -> int:
    rows = store.chats()
    if not rows:
        print("(no chats)")
        return 0
    for c in rows:
        last = time.strftime("%Y-%m-%d %H:%M", time.gmtime(c["last_ts"] + TZ_OFFSET_HOURS * 3600))
        print(f"{c['chat_id']:>20}  {c['type']:<12}  count={c['n']:<5}  last={last}  {c['title'] or ''}")
    return 0


def cmd_since(args, store: Store) -> int:
    since_ts = int(time.time()) - args.minutes * 60
    msgs = store.since(chat_id=args.chat, since_ts=since_ts, limit=10_000)
    if not msgs:
        print("(no messages in window)")
        return 0
    print(format_messages(msgs, tz_offset_hours=TZ_OFFSET_HOURS))
    return 0


def cmd_export(args, store: Store) -> int:
    if args.since_ts:
        msgs = store.since(chat_id=args.chat, since_ts=args.since_ts, limit=1_000_000)
    else:
        msgs = store.recent(chat_id=args.chat, limit=1_000_000)
    for m in msgs:
        print(json.dumps({
            "id": m.id,
            "chat_id": m.chat_id,
            "chat_title": m.chat_title,
            "message_id": m.message_id,
            "user_id": m.user_id,
            "user_display": m.user_display,
            "text": m.text,
            "media_type": m.media_type,
            "media_caption": m.media_caption,
            "reply_to_message_id": m.reply_to_message_id,
            "timestamp_utc": m.timestamp_utc,
            "edit_timestamp_utc": m.edit_timestamp_utc,
        }, ensure_ascii=False))
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="chatlogger", description="Arakawa Concierge chat log query tool")
    p.add_argument("--db", default=DEFAULT_DB, help=f"SQLite path (default: {DEFAULT_DB})")
    sub = p.add_subparsers(dest="cmd", required=True)

    r = sub.add_parser("recent", help="Last N messages (chronological)")
    r.add_argument("--chat", type=int, default=None)
    r.add_argument("-n", type=int, default=50)
    r.set_defaults(func=cmd_recent)

    s = sub.add_parser("search", help="FTS5 search")
    s.add_argument("query")
    s.add_argument("--chat", type=int, default=None)
    s.add_argument("-n", type=int, default=30)
    s.set_defaults(func=cmd_search)

    st = sub.add_parser("stats", help="DB statistics")
    st.set_defaults(func=cmd_stats)

    ch = sub.add_parser("chats", help="List known chats")
    ch.set_defaults(func=cmd_chats)

    si = sub.add_parser("since", help="Messages from the last N minutes")
    si.add_argument("minutes", type=int)
    si.add_argument("--chat", type=int, default=None)
    si.set_defaults(func=cmd_since)

    ex = sub.add_parser("export", help="Export as JSONL")
    ex.add_argument("--chat", type=int, default=None)
    ex.add_argument("--since-ts", type=int, default=None)
    ex.set_defaults(func=cmd_export)

    return p


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    store = Store(args.db)
    return args.func(args, store)


if __name__ == "__main__":
    sys.exit(main())