freezmo/lib/chat-context.ts

"use client";

import type { ChatMessage, ChatScope, ChatSource } from "./chat";
import { estimateTokens } from "./chat-trim";
import type { Deck } from "./schemas";

/**
 * Client-side, dependency-free retrieval over the user's own notes.
 *
 * The server stays stateless: the chat route never reads the notes itself.
 * This module chunks the scoped decks (~800 tokens per chunk), ranks the
 * chunks against the question plus the recent turns with a small BM25 scorer
 * and returns the top chunks under a token budget, each labelled with its
 * deck and section. Only what comes back from here is sent to the model.
 */

export type ContextChunk = ChatSource & { text: string; score: number };

const CHUNK_CHARS = 3200; // ≈ 800 tokens
const CHUNK_OVERLAP = 320;
const DEFAULT_BUDGET_TOKENS = 7000;
const MAX_CHUNKS = 24;
const K1 = 1.2;
const B = 0.75;

export function deckHasNotes(deck: Deck): boolean {
  return Boolean(deck.sourceText?.trim());
}

/**
 * Older decks (or imported ones) may have no notes text at all. Their
 * generated content is the fallback — the UI says so explicitly.
 */
export function generatedDeckText(deck: Deck): string {
  const lines: string[] = [];
  if (deck.title) lines.push(`Deck: ${deck.title}`);
  if (deck.subject) lines.push(`Subject: ${deck.subject}`);
  if (deck.summary) lines.push(`Summary: ${deck.summary}`);
  const topics = new Set<string>();
  for (const activity of deck.activities) {
    if (activity.topic) topics.add(activity.topic);
  }
  if (topics.size) lines.push(`Topics: ${[...topics].join(", ")}`);
  lines.push("Generated study content:");
  for (const activity of deck.activities) {
    lines.push(...activityToLines(activity));
  }
  return lines.join("\n");
}

function activityToLines(activity: Deck["activities"][number]): string[] {
  const out: string[] = [];
  switch (activity.type) {
    case "multiple_choice":
      out.push(`- Q: ${activity.question} → ${activity.options[activity.correctIndex] ?? ""}`);
      break;
    case "true_false":
      out.push(`- ${activity.statement} → ${activity.answer ? "true" : "false"}`);
      break;
    case "fill_blank":
      out.push(`- ${activity.sentence} → ${activity.answer}`);
      break;
    case "matching":
      out.push(`- ${activity.instruction}`);
      for (const pair of activity.pairs) out.push(`  · ${pair.term} → ${pair.definition}`);
      break;
    case "ordering":
      out.push(`- ${activity.instruction}: ${activity.items.join(" → ")}`);
      break;
    case "short_answer":
      out.push(`- Q: ${activity.question} → ${activity.referenceAnswer}`);
      break;
    case "flashcard":
      out.push(`- ${activity.front} → ${activity.back}`);
      break;
  }
  out.push(`  (${activity.explanation})`);
  return out;
}

/**
 * Splits a deck's text into chunks. `### name` headings (added by the document
 * extractor) become section labels; anything else is numbered.
 */
function chunkText(text: string): { section: string; text: string }[] {
  const sections: { section: string; text: string }[] = [];
  let currentSection = "Notes";
  let current: string[] = [];
  const flush = () => {
    const joined = current.join("\n").trim();
    if (joined) sections.push({ section: currentSection, text: joined });
    current = [];
  };
  for (const line of text.split("\n")) {
    const heading = /^#{1,4}\s+(.+)$/.exec(line.trim());
    if (heading) {
      flush();
      currentSection = heading[1].trim().slice(0, 120) || "Notes";
      continue;
    }
    current.push(line);
  }
  flush();

  const chunks: { section: string; text: string }[] = [];
  for (const section of sections) {
    if (section.text.length <= CHUNK_CHARS) {
      chunks.push(section);
      continue;
    }
    let start = 0;
    let part = 1;
    while (start < section.text.length) {
      const end = Math.min(section.text.length, start + CHUNK_CHARS);
      const slice = section.text.slice(start, end).trim();
      if (slice) {
        chunks.push({
          section: section.section === "Notes" ? `Part ${part}` : `${section.section} · ${part}`,
          text: slice,
        });
      }
      part += 1;
      if (end >= section.text.length) break;
      start = end - CHUNK_OVERLAP;
    }
  }
  return chunks;
}

function tokenise(text: string): string[] {
  return text
    .toLowerCase()
    .normalize("NFKD")
    .replace(/[\u0300-\u036f]/g, "")
    .replace(/[^\p{L}\p{N}\s]/gu, " ")
    .split(/\s+/)
    .filter((word) => word.length > 2 && !STOPWORDS.has(word));
}

const STOPWORDS = new Set([
  "the", "and", "for", "are", "was", "were", "with", "that", "this", "from",
  "have", "has", "had", "not", "you", "your", "what", "when", "where", "which",
  "how", "why", "who", "can", "could", "would", "should", "about", "into",
  "than", "then", "them", "they", "there", "here", "does", "did", "doing",
  "jak", "kdy", "kde", "kter", "která", "které", "který", "což", "podle",
  "jako", "jsou", "byl", "byla", "bylo", "být", "mám", "máte", "může",
  "proč", "proto", "tak", "také", "nebo", "ale", "ani", "už", "ještě",
]);

/**
 * BM25 over the chunks. The query is the new question plus the last few turns,
 * so follow-ups like "and why?" still retrieve the right material.
 */
function rank(chunks: { section: string; text: string }[], query: string): number[] {
  const terms = [...new Set(tokenise(query))];
  if (terms.length === 0) return chunks.map(() => 0);
  const docs = chunks.map((chunk) => tokenise(chunk.text));
  const lengths = docs.map((doc) => doc.length);
  const avg = lengths.reduce((sum, value) => sum + value, 0) / Math.max(1, lengths.length);
  const df = new Map<string, number>();
  for (const doc of docs) {
    for (const term of new Set(doc)) df.set(term, (df.get(term) ?? 0) + 1);
  }
  const N = docs.length;
  return docs.map((doc, index) => {
    const counts = new Map<string, number>();
    for (const term of doc) counts.set(term, (counts.get(term) ?? 0) + 1);
    let score = 0;
    for (const term of terms) {
      const freq = counts.get(term);
      if (!freq) continue;
      const docFreq = df.get(term) ?? 0;
      const idf = Math.log(1 + (N - docFreq + 0.5) / (docFreq + 0.5));
      const norm = K1 * (1 - B + (B * lengths[index]) / Math.max(1, avg));
      score += (idf * (freq * (K1 + 1))) / (freq + norm);
    }
    return score;
  });
}

export type ChatContextResult = {
  chunks: ContextChunk[];
  /** Decks in scope whose notes text is missing (fallback content was used). */
  missingNotes: string[];
  scopedDecks: Deck[];
};

/**
 * Builds the context for one request: scoped decks → chunks → BM25 top-k under
 * the token budget. `extraChunks` (e.g. a study result handed over from the
 * results screen) are always kept and come first.
 */
export function buildChatContext(input: {
  decks: Deck[];
  scope: ChatScope;
  question: string;
  history: ChatMessage[];
  useNotes: boolean;
  extraChunks?: { deckId: string; deckTitle: string; section: string; text: string }[];
  budgetTokens?: number;
}): ChatContextResult {
  const scopedDecks =
    input.scope.kind === "none"
      ? []
      : input.scope.kind === "all"
        ? input.decks.filter((deck) => deck.activities.length > 0)
        : input.decks.filter((deck) => input.scope.kind === "decks" && input.scope.deckIds.includes(deck.id));

  const budget = input.budgetTokens ?? DEFAULT_BUDGET_TOKENS;
  const chunks: ContextChunk[] = [];
  const missingNotes: string[] = [];
  let used = 0;

  const push = (chunk: ContextChunk) => {
    const cost = estimateTokens(chunk.text);
    if (used + cost > budget) return false;
    chunks.push(chunk);
    used += cost;
    return true;
  };

  for (const extra of input.extraChunks ?? []) {
    push({ ...extra, score: Number.POSITIVE_INFINITY });
  }

  const recent = input.history
    .slice(-4)
    .map((message) => message.content)
    .join("\n");
  const query = `${input.question}\n${recent}`;

  for (const deck of scopedDecks) {
    const hasNotes = deckHasNotes(deck);
    if (!hasNotes) missingNotes.push(deck.id);
    const text = hasNotes && input.useNotes ? deck.sourceText : generatedDeckText(deck);
    const parts = chunkText(text);
    const scores = rank(parts, query);
    const ranked = parts
      .map((part, index) => ({ part, score: scores[index] }))
      .sort((a, b) => b.score - a.score);
    const top = ranked.filter((entry) => entry.score > 0).slice(0, 6);
    const ordered = top.length > 0 ? top : ranked.slice(0, 2);
    for (const entry of ordered) {
      if (chunks.length >= MAX_CHUNKS) break;
      if (!push({
        deckId: deck.id,
        deckTitle: deck.title,
        section: entry.part.section,
        text: entry.part.text,
        score: entry.score,
      })) {
        break;
      }
    }
  }

  return { chunks, missingNotes, scopedDecks };
}