douyin-archive/lib/hashtag-related.server.ts

import {
  HASHTAG_EXTRACT_PG,
  META_HASHTAG_PG,
  META_HASHTAGS,
} from "@/lib/hashtag";
import { prisma } from "@/lib/prisma";

export type RelatedHashtag = {
  name: string;
  /** 和目标 tag 一起出现的作品数 */
  together: number;
  /** 这个 tag 自己的作品数 */
  count: number;
};

export type RelatedHashtagsResult = {
  name: string;
  count: number;
  related: RelatedHashtag[];
};

const CACHE_MS = 60_000;
const cache = new Map<string, { at: number; value: RelatedHashtagsResult }>();

/**
 * 按「同一个作品描述里一起出现」计算相关 hashtag（共现图），每次实时从描述里解析，不需要额外的表。
 *
 * 打分取两个方向里较大的那个：
 * - 子话题：候选 tag 大多和目标一起出现（together / 候选数），乘 ln(1+together) 防止只出现一次的巧合排到前面；
 *   例：#明日方舟 → #明日方舟七周年、#特蕾西娅
 * - 上位话题：目标大多和候选一起出现（together / 目标数），乘 ln(1+候选数)，让热门的上位 tag 排前面；
 *   例：#圣聆初雪（只有 2 个作品）→ #明日方舟
 * cos、二次元这类「万金油」tag 与平台活动 tag 不参与推荐。
 */
export async function relatedHashtags(
  rawName: string,
  limit = 12,
): Promise<RelatedHashtagsResult> {
  const name = rawName.toLowerCase();
  const key = `${name}:${limit}`;
  const hit = cache.get(key);
  if (hit && Date.now() - hit.at < CACHE_MS) return hit.value;

  const rows = await prisma.$queryRaw<
    { name: string; together: number; count: number; target: number }[]
  >`
    WITH ht AS (
      SELECT DISTINCT w.aweme_id, lower(m[1]) AS tag
      FROM (
        SELECT aweme_id, "desc" FROM "Video"
        UNION ALL
        SELECT aweme_id, "desc" FROM "ImagePost"
      ) w,
      regexp_matches(w."desc", ${HASHTAG_EXTRACT_PG}, 'g') m
    ),
    target AS (SELECT aweme_id FROM ht WHERE tag = ${name}),
    target_count AS (SELECT count(*)::int AS c FROM target),
    co AS (
      SELECT ht.tag, count(*)::int AS together
      FROM ht JOIN target USING (aweme_id)
      WHERE ht.tag <> ${name}
      GROUP BY ht.tag
    ),
    counts AS (
      SELECT tag, count(*)::int AS c FROM ht
      WHERE tag IN (SELECT tag FROM co)
      GROUP BY tag
    )
    SELECT
      co.tag AS name,
      co.together,
      counts.c AS count,
      tc.c AS target
    FROM co
    JOIN counts USING (tag)
    CROSS JOIN target_count tc
    WHERE co.together >= CASE WHEN tc.c >= 4 THEN 2 ELSE 1 END
      AND co.tag <> ALL (${META_HASHTAGS}::text[])
      AND co.tag !~ ${META_HASHTAG_PG}
    ORDER BY GREATEST(
      co.together::float8 / counts.c * ln(1 + co.together),
      co.together::float8 / tc.c * ln(1 + counts.c)
    ) DESC, co.together DESC, co.tag
    LIMIT ${limit}
  `;

  let count = rows[0]?.target;
  if (count === undefined) {
    // 没有相关 tag 时单独数一下目标 tag 的作品数
    const [{ c }] = await prisma.$queryRaw<{ c: number }[]>`
      SELECT count(DISTINCT w.aweme_id)::int AS c
      FROM (
        SELECT aweme_id, "desc" FROM "Video"
        UNION ALL
        SELECT aweme_id, "desc" FROM "ImagePost"
      ) w,
      regexp_matches(w."desc", ${HASHTAG_EXTRACT_PG}, 'g') m
      WHERE lower(m[1]) = ${name}
    `;
    count = c;
  }

  const value: RelatedHashtagsResult = {
    name,
    count,
    related: rows.map(({ name, together, count }) => ({
      name,
      together,
      count,
    })),
  };
  cache.set(key, { at: Date.now(), value });
  return value;
}