justvitamin/scripts/tag-reviews.ts

#!/usr/bin/env bun
// Tag every raw review with the aspects it supports.
// Usage: bun scripts/tag-reviews.ts [SKU1 SKU2 ...]   (defaults to all SKUs)
// Output: data/intelligence/<SKU>/review-aspects-supporting.json

import { join } from 'path'
import { readFileSync, writeFileSync, readdirSync, existsSync } from 'fs'

const ROOT = join(import.meta.dir, '..')
const INTEL = join(ROOT, 'data', 'intelligence')
const API_KEY = process.env.GEMINI_API_KEY
const MODEL = 'gemini-2.5-flash'
const BATCH_SIZE = 40
const CONCURRENCY = 4

if (!API_KEY) { console.error('GEMINI_API_KEY not set'); process.exit(1) }

type Aspect = { rank: number; aspect: string; sentiment: string; exampleQuotes?: string[] }
type Review = { comment: string; stars: number }

const SYSTEM = `You label customer reviews against a fixed list of aspects.
For each review, return ALL aspect ranks it clearly supports (the review states or strongly implies that aspect).
Match paraphrases (e.g. "achy joints" supports "joint pain"). Do NOT match weak/tangential mentions.
A review may match 0, 1, or many aspects. A short generic review like "great product" matches none.
Output strict JSON only: an array where index i corresponds to review i in the batch.
Each element is an array of integer aspect ranks.`

async function tagBatch(aspects: Aspect[], batch: { idx: number; r: Review }[]): Promise<number[][]> {
  const aspectList = aspects.map(a => `${a.rank}. [${a.sentiment}] ${a.aspect}`).join('\n')
  const reviewList = batch.map((b, i) => `[${i}] (${b.r.stars}★) ${b.r.comment.replace(/\s+/g, ' ').slice(0, 500)}`).join('\n')
  const userMsg = `ASPECTS:\n${aspectList}\n\nREVIEWS (${batch.length}):\n${reviewList}\n\nReturn JSON array of length ${batch.length}, each element an array of matching aspect ranks. Example: [[1,3],[],[2]]`

  for (let attempt = 0; attempt < 3; attempt++) {
    try {
      const url = `https://generativelanguage.googleapis.com/v1beta/models/${MODEL}:generateContent?key=${API_KEY}`
      const res = await fetch(url, {
        method: 'POST',
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          systemInstruction: { parts: [{ text: SYSTEM }] },
          contents: [{ role: 'user', parts: [{ text: userMsg }] }],
          generationConfig: {
            temperature: 0,
            responseMimeType: 'application/json',
            maxOutputTokens: 8000
          }
        })
      })
      if (!res.ok) {
        const txt = await res.text()
        if (res.status === 429 || res.status >= 500) {
          await new Promise(r => setTimeout(r, 2000 * (attempt + 1)))
          continue
        }
        throw new Error(`API ${res.status}: ${txt.slice(0, 200)}`)
      }
      const json: any = await res.json()
      const text = json.candidates?.[0]?.content?.parts?.[0]?.text || ''
      const match = text.match(/\[[\s\S]*\]/)
      if (!match) throw new Error('No JSON array in response')
      const arr = JSON.parse(match[0])
      if (!Array.isArray(arr) || arr.length !== batch.length) {
        throw new Error(`Expected length ${batch.length}, got ${Array.isArray(arr) ? arr.length : 'non-array'}`)
      }
      return arr.map((x: any) => Array.isArray(x) ? x.filter((n: any) => Number.isInteger(n)) : [])
    } catch (e: any) {
      if (attempt === 2) throw e
      await new Promise(r => setTimeout(r, 1500 * (attempt + 1)))
    }
  }
  return batch.map(() => [])
}

async function tagSku(sku: string) {
  const dir = join(INTEL, sku)
  const aspectsFile = join(dir, 'review-aspects.json')
  const reviewsFile = join(dir, '_raw_reviews.json')
  if (!existsSync(aspectsFile) || !existsSync(reviewsFile)) {
    console.log(`[${sku}] skip — missing files`); return
  }
  const aspectsDoc = JSON.parse(readFileSync(aspectsFile, 'utf8'))
  const reviews: Review[] = JSON.parse(readFileSync(reviewsFile, 'utf8'))
  const aspects: Aspect[] = aspectsDoc.aspects

  console.log(`[${sku}] ${reviews.length} reviews × ${aspects.length} aspects`)

  // Build batches with original indices
  const indexed = reviews.map((r, idx) => ({ idx, r })).filter(x => x.r.comment && x.r.comment.trim().length > 0)
  const batches: { idx: number; r: Review }[][] = []
  for (let i = 0; i < indexed.length; i += BATCH_SIZE) batches.push(indexed.slice(i, i + BATCH_SIZE))

  const supportingByRank: Record<number, number[]> = {}
  for (const a of aspects) supportingByRank[a.rank] = []

  let done = 0
  // Run with limited concurrency
  for (let i = 0; i < batches.length; i += CONCURRENCY) {
    const slice = batches.slice(i, i + CONCURRENCY)
    const results = await Promise.all(slice.map(b => tagBatch(aspects, b).catch(e => {
      console.error(`  batch err: ${e.message}`); return b.map(() => [])
    })))
    results.forEach((labels, k) => {
      const batch = slice[k]
      labels.forEach((ranks, j) => {
        const reviewIdx = batch[j].idx
        for (const rank of ranks) {
          if (supportingByRank[rank] !== undefined) supportingByRank[rank].push(reviewIdx)
        }
      })
    })
    done += slice.reduce((s, b) => s + b.length, 0)
    process.stdout.write(`\r[${sku}] tagged ${done}/${indexed.length}`)
  }
  process.stdout.write('\n')

  const out = {
    sku,
    generatedAt: new Date().toISOString().slice(0, 10),
    model: MODEL,
    totalReviews: reviews.length,
    supportingReviews: aspects.map(a => ({
      rank: a.rank,
      aspect: a.aspect,
      matchCount: supportingByRank[a.rank].length,
      reviewIndices: supportingByRank[a.rank].sort((x, y) => x - y)
    }))
  }
  const outFile = join(dir, 'review-aspects-supporting.json')
  writeFileSync(outFile, JSON.stringify(out, null, 2))
  console.log(`[${sku}] wrote ${outFile}`)
  for (const s of out.supportingReviews) {
    const stated = aspects.find(a => a.rank === s.rank)
    console.log(`  rank ${s.rank}: matched ${s.matchCount} reviews (stated: ${(aspectsDoc.aspects.find((a: any) => a.rank === s.rank) as any)?.reviews ?? '?'})`)
  }
}

const targets = process.argv.slice(2).length
  ? process.argv.slice(2)
  : readdirSync(INTEL).filter(d => existsSync(join(INTEL, d, 'review-aspects.json')))

for (const sku of targets) {
  try { await tagSku(sku) } catch (e: any) { console.error(`[${sku}] FAILED: ${e.message}`) }
}