Files
yunke/yunke-server/scripts/seed_chats.py
T
浅念 fd86adfb63 chore: seed chat messages for listed customers (seed_chats)
Iterates smoke.db customers, scrapes each conversation's recent messages
so timelines aren't empty. Bounded per-customer.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-05 09:54:40 +08:00

71 lines
2.4 KiB
Python

"""Seed chat messages for customers already in the list (so timelines aren't empty).
Iterates every Customer row in smoke.db, scrapes that conversation's recent
messages (WeChat msg/chats or WhatsApp chatRecords), and upserts them. Bounded
per-customer so it stays quick.
"""
import asyncio
import sqlite3
import time
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from app.db import models # noqa: F401
from app.scraper.wechat import scrape_wechat_messages
from app.scraper.whatsapp import scrape_whatsapp_messages
from app.storage.repos import upsert_messages
from scripts.smoke_wechat import _push_settings_into_environ, build_client_from_env
WINDOW_MS = 90 * 86_400_000 # last 90 days
MAX_PER_CUSTOMER = 60
async def main() -> None:
_push_settings_into_environ()
client = build_client_from_env()
con = sqlite3.connect("smoke.db")
rows = con.execute(
"SELECT channel, owner_account_id, friend_id FROM customer "
"WHERE owner_account_id IS NOT NULL AND friend_id IS NOT NULL"
).fetchall()
con.close()
engine = create_async_engine("sqlite+aiosqlite:///./smoke.db")
sm = async_sessionmaker(engine, expire_on_commit=False)
now = int(time.time() * 1000)
total = 0
customers_with_msgs = 0
for channel, owner, friend in rows:
recs: list[dict] = []
try:
if channel == "wechat":
agen = scrape_wechat_messages(client, owner, friend, now - WINDOW_MS, now)
else:
agen = scrape_whatsapp_messages(client, owner, friend, now - WINDOW_MS, now)
async for m in agen:
recs.append(m)
if len(recs) >= MAX_PER_CUSTOMER:
break
except Exception as e: # noqa: BLE001
print(f" [{channel}] {friend}: {type(e).__name__}: {e!r}")
continue
if recs:
async with sm() as session:
n = await upsert_messages(session, recs)
await session.commit()
total += n
customers_with_msgs += 1
print(f" [{channel}] {friend}: {n} messages")
else:
print(f" [{channel}] {friend}: (no messages in window)")
print(f"\nDONE: {total} messages across {customers_with_msgs} customers")
await client.aclose()
await engine.dispose()
if __name__ == "__main__":
asyncio.run(main())