learning-platform/src/lib/extractionPipeline.js

import * as db from './db';
import { callLLM, cachedSystem } from './llm';
import { extractionLimiter } from './llmRetry';
import { EMIT_KNOWLEDGE_GRAPH_TOOL } from './llmTools';

const MAX_KNOWN_TOPICS_HINT = 200;

/**
 * Build the "already-extracted topics" hint included in every chunk prompt.
 * Passes both ID and label so the model can match concepts by name and reuse
 * the exact ID + label rather than inventing near-duplicate variants.
 */
export function buildKnownIdsHint(topics) {
  if (!topics || !topics.length) return '';
  const recent = topics.slice(-MAX_KNOWN_TOPICS_HINT);
  return [
    'Already-extracted topics (reuse their ID and label exactly if the same concept appears here):',
    ...recent.map((t) => `- ${t.id}: "${t.label}"`),
    '',
  ].join('\n');
}

const EXTRACTION_SYSTEM_PROMPT = `You are an AI knowledge extractor for Respellion, an IT company built on radical transparency.
You receive a source text. Extract every distinct concept, role, and process from it and emit them through the emit_knowledge_graph tool.

CRITICAL INSTRUCTIONS FOR COMPLETENESS:
- Extract up to 15 of the most important distinct roles, processes, and concepts described or mentioned in the source text.
- Do not exceed 15 topics per chunk to prevent the response from being truncated.
- Facts should be integrated into the descriptions of other topics — never extracted as standalone topics.
- Keep descriptions concise (max 3 sentences) so the response fits.

Topic IDs are lowercase kebab-case slugs specific to the topic (e.g. "software-engineer", "data-quality-review"). Do not use generic IDs like "role-1" or "concept-2".

Assign a learning_relevance to every topic:
- "core": fundamental company knowledge.
- "standard": normal learning topics.
- "peripheral": good to know, low priority.
- "exclude": pure operational reference (printer guides, wifi passwords) that should never be tested.

Relation types: related_to | depends_on | part_of | executed_by.
`;

/**
 * Sentence-aware chunker with overlap.
 *
 * Targets ~2000 input tokens per chunk (`MAX_CHUNK_CHARS / 4`). Splits on
 * sentence boundaries first, then falls back to paragraph boundaries, and
 * hard-splits inside an oversized sentence as a last resort. Adjacent chunks
 * share `overlapChars` of trailing text to preserve cross-boundary context
 * for the model.
 *
 * Exported for unit tests; callers in this module use it directly.
 *
 * @param {string} text
 * @param {{ maxChars?: number, overlapChars?: number }} [opts]
 * @returns {string[]}
 */
export const MAX_CHUNK_CHARS = 8000;
export const OVERLAP_CHARS = 800;

export function chunkText(text, { maxChars = MAX_CHUNK_CHARS, overlapChars = OVERLAP_CHARS } = {}) {
  if (typeof text !== 'string' || !text.trim()) return [];
  const trimmed = text.trim();
  if (trimmed.length <= maxChars) return [trimmed];

  const units = splitIntoChunkableUnits(trimmed, maxChars);
  if (units.length === 0) return [];

  const chunks = [];
  let buf = '';
  let bufLen = 0; // length of new (non-overlap) content added since last flush

  for (const unit of units) {
    const wouldOverflow = (buf ? buf.length + 1 + unit.length : unit.length) > maxChars;
    if (wouldOverflow && bufLen > 0) {
      chunks.push(buf.trim());
      const overlap = buf.length > overlapChars ? buf.slice(-overlapChars) : '';
      buf = overlap;
      bufLen = 0;
    }
    // If the overlap + unit still won't fit, drop the overlap so the unit fits cleanly.
    if (buf && (buf.length + 1 + unit.length) > maxChars) {
      buf = '';
    }
    buf = buf ? buf + ' ' + unit : unit;
    bufLen += unit.length + (bufLen > 0 ? 1 : 0);
  }
  if (bufLen > 0 && buf.trim()) chunks.push(buf.trim());
  return chunks;
}

function splitIntoChunkableUnits(text, maxChars) {
  const paragraphs = text.split(/\n\s*\n+/);
  const units = [];
  for (const para of paragraphs) {
    const trimmedPara = para.trim();
    if (!trimmedPara) continue;
    const sentences = trimmedPara.split(/(?<=[.!?])\s+/);
    for (const s of sentences) {
      const sentence = s.trim();
      if (!sentence) continue;
      if (sentence.length <= maxChars) {
        units.push(sentence);
      } else {
        for (let i = 0; i < sentence.length; i += maxChars) {
          units.push(sentence.slice(i, i + maxChars));
        }
        console.warn(`[chunkText] Hard-split a sentence of ${sentence.length} chars (exceeds maxChars=${maxChars}).`);
      }
    }
  }
  return units;
}

export async function processSourceText(textContent, sourceName, { signal } = {}) {
  const existing = await db.getSources();
  const alreadyDone = existing.find(
    s => s.name === sourceName && s.status === 'completed'
  );
  if (alreadyDone) {
    throw new Error(`"${sourceName}" has already been processed. Delete the existing source first if you want to re-analyse it.`);
  }

  const rec = await db.addSource({ name: sourceName, status: 'processing' });
  const sourceId = rec.id;

  try {
    const chunks = chunkText(textContent);
    console.log(`[Pipeline] Split "${sourceName}" into ${chunks.length} chunks for processing.`);

    // Persist initial progress so other sessions/reloads can see it
    await db.updateSourceProgress(sourceId, { current: 0, total: chunks.length, message: 'Starting extraction...' });

    const existingTopics = await db.getTopics();
    const knownTopics = existingTopics.map((t) => ({ id: t.id, label: t.label }));

    const allExtractedTopics = [];
    const allExtractedRelations = [];

    for (let i = 0; i < chunks.length; i++) {
      if (signal?.aborted) throw signal.reason ?? new DOMException('Aborted', 'AbortError');

      console.log(`[Pipeline] Processing chunk ${i + 1}/${chunks.length} (${chunks[i].length} chars)...`);

      // Update progress before each chunk
      await db.updateSourceProgress(sourceId, {
        current: i,
        total: chunks.length,
        message: `Extracting chunk ${i + 1} of ${chunks.length}...`,
      });

      const hint = buildKnownIdsHint(knownTopics);
      const result = await callLLM({
        task: 'extract.source',
        tier: 'standard',
        system: cachedSystem(EXTRACTION_SYSTEM_PROMPT),
        user: `${hint}Analyze this part of the document (${i + 1}/${chunks.length}):\n\n${chunks[i]}`,
        tools: [EMIT_KNOWLEDGE_GRAPH_TOOL],
        toolChoice: { type: 'tool', name: EMIT_KNOWLEDGE_GRAPH_TOOL.name },
        maxTokens: 8192,
        timeoutMs: 180_000,
        limiter: extractionLimiter,
        signal,
      });

      const extractedData = result.toolUses[0]?.input;
      if (!extractedData) throw new Error(`Extraction did not emit a tool result for chunk ${i + 1}.`);

      if (Array.isArray(extractedData.topics)) {
        allExtractedTopics.push(...extractedData.topics);
        for (const t of extractedData.topics) {
          if (t?.id && !knownTopics.some((k) => k.id === t.id)) {
            knownTopics.push({ id: t.id, label: t.label });
          }
        }
      }
      if (Array.isArray(extractedData.relations)) {
        allExtractedRelations.push(...extractedData.relations);
      }
    }

    await db.updateSourceProgress(sourceId, { current: chunks.length, total: chunks.length, message: 'Merging results...' });
    await mergeKnowledgeGraph(existingTopics, { topics: allExtractedTopics, relations: allExtractedRelations });
    await db.updateSourceStatus(sourceId, 'completed');

    return { success: true, data: { topics: allExtractedTopics, relations: allExtractedRelations } };

  } catch (error) {
    const isAbort = error?.name === 'AbortError';
    await db.updateSourceStatus(sourceId, isAbort ? 'cancelled' : 'failed', isAbort ? 'cancelled by user' : error.message);
    throw error;
  }
}

async function mergeKnowledgeGraph(existingTopics, newData) {
  const existingRelations = await db.getRelations();

  const topicsMap = new Map(existingTopics.map(t => [t.id, t]));

  if (newData.topics && Array.isArray(newData.topics)) {
    for (const t of newData.topics) {
      if (topicsMap.has(t.id)) {
        const existing = topicsMap.get(t.id);
        const merged = {
          ...existing,
          ...t,
          description: t.description || existing.description,
        };
        if (existing.relevance_locked) {
          merged.learning_relevance = existing.learning_relevance;
          merged.relevance_locked = true;
        }
        topicsMap.set(t.id, merged);
      } else {
        topicsMap.set(t.id, t);
      }
    }
  }

  const newRelations = [...existingRelations];
  if (newData.relations && Array.isArray(newData.relations)) {
    for (const r of newData.relations) {
      const isDup = newRelations.some(ex => ex.source === r.source && ex.target === r.target && ex.type === r.type);
      if (!isDup) {
        newRelations.push(r);
      }
    }
  }

  await db.saveTopics(Array.from(topicsMap.values()));
  await db.saveRelations(newRelations);
}