Build an AI Text Classification API

// src/classify/engine.ts — Text classification with multi-label support and custom models
import { pool } from "../db";
import { Redis } from "ioredis";
import { createHash } from "node:crypto";

const redis = new Redis(process.env.REDIS_URL!);

interface ClassificationResult {
  labels: Array<{ label: string; confidence: number }>;
  sentiment: { score: number; label: "positive" | "negative" | "neutral" };
  spam: { isSpam: boolean; confidence: number };
  language: string;
  processingMs: number;
}

interface ClassificationModel {
  id: string;
  name: string;
  type: "category" | "sentiment" | "spam" | "custom";
  vocabulary: Map<string, Record<string, number>>;
  labels: string[];
  trainedAt: string;
  accuracy: number;
}

// Built-in category keywords (Naive Bayes-style)
const CATEGORY_KEYWORDS: Record<string, string[]> = {
  technology: ["api", "database", "code", "software", "deploy", "server", "cloud", "algorithm", "programming", "framework", "typescript", "python", "kubernetes", "docker", "ai", "machine learning"],
  business: ["revenue", "startup", "funding", "market", "growth", "investor", "profit", "strategy", "acquisition", "ipo", "valuation", "enterprise"],
  science: ["research", "study", "experiment", "hypothesis", "data", "journal", "peer-review", "discovery", "quantum", "genome", "climate"],
  design: ["ui", "ux", "figma", "layout", "typography", "color", "accessibility", "prototype", "wireframe", "responsive"],
  security: ["vulnerability", "exploit", "encryption", "firewall", "breach", "authentication", "zero-day", "malware", "phishing"],
  devops: ["ci/cd", "pipeline", "monitoring", "terraform", "ansible", "kubernetes", "docker", "deployment", "infrastructure"],
};

const SENTIMENT_WORDS = {
  positive: ["great", "excellent", "amazing", "love", "best", "awesome", "fantastic", "brilliant", "outstanding", "perfect", "wonderful", "impressive"],
  negative: ["bad", "terrible", "worst", "hate", "awful", "horrible", "poor", "broken", "failed", "disappointing", "frustrating", "useless"],
};

const SPAM_INDICATORS = ["buy now", "click here", "free money", "act now", "limited time", "congratulations you won", "nigerian prince", "earn $", "work from home", "100% free", "no obligation", "double your"];

// Classify text
export async function classify(text: string, options?: { models?: string[]; threshold?: number }): Promise<ClassificationResult> {
  const start = Date.now();
  const threshold = options?.threshold || 0.3;

  // Check cache
  const cacheKey = `classify:${createHash("md5").update(text.slice(0, 500)).digest("hex")}`;
  const cached = await redis.get(cacheKey);
  if (cached) return { ...JSON.parse(cached), processingMs: 1 };

  const lower = text.toLowerCase();
  const words = tokenize(lower);

  // Category classification
  const categoryScores: Record<string, number> = {};
  for (const [category, keywords] of Object.entries(CATEGORY_KEYWORDS)) {
    const matchCount = keywords.filter((kw) => lower.includes(kw)).length;
    const score = matchCount / keywords.length;
    if (score > 0) categoryScores[category] = score;
  }

  const labels = Object.entries(categoryScores)
    .filter(([, score]) => score >= threshold)
    .sort((a, b) => b[1] - a[1])
    .map(([label, confidence]) => ({ label, confidence: Math.round(confidence * 100) / 100 }));

  // Sentiment analysis
  let sentimentScore = 0;
  for (const word of words) {
    if (SENTIMENT_WORDS.positive.includes(word)) sentimentScore += 0.15;
    if (SENTIMENT_WORDS.negative.includes(word)) sentimentScore -= 0.15;
  }
  sentimentScore = Math.max(-1, Math.min(1, sentimentScore));
  const sentimentLabel = sentimentScore > 0.1 ? "positive" : sentimentScore < -0.1 ? "negative" : "neutral";

  // Spam detection
  const spamMatchCount = SPAM_INDICATORS.filter((indicator) => lower.includes(indicator)).length;
  const hasExcessiveCaps = (text.match(/[A-Z]/g) || []).length / text.length > 0.5;
  const hasExcessiveLinks = (text.match(/https?:\/\//g) || []).length > 3;
  const spamConfidence = Math.min(1, spamMatchCount * 0.3 + (hasExcessiveCaps ? 0.2 : 0) + (hasExcessiveLinks ? 0.2 : 0));

  // Language detection (simplified)
  const language = detectLanguage(text);

  const result: ClassificationResult = {
    labels: labels.length > 0 ? labels : [{ label: "general", confidence: 0.5 }],
    sentiment: { score: sentimentScore, label: sentimentLabel },
    spam: { isSpam: spamConfidence > 0.5, confidence: spamConfidence },
    language,
    processingMs: Date.now() - start,
  };

  // Cache for 1 hour
  await redis.setex(cacheKey, 3600, JSON.stringify(result));

  // Track analytics
  await redis.hincrby("classify:stats", "total", 1);
  for (const label of labels) await redis.hincrby("classify:stats", `cat:${label.label}`, 1);
  await redis.hincrby("classify:stats", `sentiment:${sentimentLabel}`, 1);
  if (result.spam.isSpam) await redis.hincrby("classify:stats", "spam", 1);

  return result;
}

// Batch classification
export async function batchClassify(texts: Array<{ id: string; text: string }>): Promise<Array<{ id: string; result: ClassificationResult }>> {
  const results = [];
  for (const item of texts) {
    const result = await classify(item.text);
    results.push({ id: item.id, result });
  }
  return results;
}

// Train custom model from labeled examples
export async function trainCustomModel(modelName: string, examples: Array<{ text: string; labels: string[] }>): Promise<{ accuracy: number; modelId: string }> {
  const vocabulary = new Map<string, Record<string, number>>();
  const labelCounts: Record<string, number> = {};

  // Build word-label frequency table
  for (const example of examples) {
    const words = tokenize(example.text.toLowerCase());
    for (const label of example.labels) {
      labelCounts[label] = (labelCounts[label] || 0) + 1;
      for (const word of words) {
        if (!vocabulary.has(word)) vocabulary.set(word, {});
        const wordLabels = vocabulary.get(word)!;
        wordLabels[label] = (wordLabels[label] || 0) + 1;
      }
    }
  }

  // Store model
  const modelId = `model-${createHash("md5").update(modelName).digest("hex").slice(0, 8)}`;
  await pool.query(
    `INSERT INTO classification_models (id, name, vocabulary, label_counts, trained_at, example_count)
     VALUES ($1, $2, $3, $4, NOW(), $5)
     ON CONFLICT (name) DO UPDATE SET vocabulary = $3, label_counts = $4, trained_at = NOW(), example_count = $5`,
    [modelId, modelName, JSON.stringify(Object.fromEntries(vocabulary)), JSON.stringify(labelCounts), examples.length]
  );

  // Simple accuracy estimate (leave-one-out on small sample)
  let correct = 0;
  const testSize = Math.min(examples.length, 50);
  for (let i = 0; i < testSize; i++) {
    const prediction = await classify(examples[i].text);
    if (examples[i].labels.some((l) => prediction.labels.some((p) => p.label === l))) correct++;
  }

  return { accuracy: correct / testSize, modelId };
}

function tokenize(text: string): string[] {
  return text.toLowerCase().replace(/[^a-z0-9\s/\-]/g, "").split(/\s+/).filter((w) => w.length > 2);
}

function detectLanguage(text: string): string {
  const patterns: Record<string, RegExp> = {
    en: /\b(the|is|are|was|were|have|has|been|will|would|could|should)\b/gi,
    es: /\b(el|la|los|las|es|son|fue|ser|estar|tiene|hacer)\b/gi,
    fr: /\b(le|la|les|est|sont|être|avoir|fait|dans|pour)\b/gi,
    de: /\b(der|die|das|ist|sind|hat|haben|werden|nicht|ein)\b/gi,
  };
  let bestLang = "en", bestScore = 0;
  for (const [lang, pattern] of Object.entries(patterns)) {
    const matches = (text.match(pattern) || []).length;
    if (matches > bestScore) { bestScore = matches; bestLang = lang; }
  }
  return bestLang;
}

// Analytics
export async function getClassificationStats(): Promise<Record<string, number>> {
  return redis.hgetall("classify:stats");
}