learning-platform/app/services/ingestion/src/pipeline/embed.ts

import { v4 as uuid } from 'uuid';
import { openai, EMBEDDING_MODEL, EMBEDDING_BATCH_SIZE } from '../lib/openai.js';
import { qdrant, QDRANT_COLLECTIONS } from '../lib/qdrant.js';
import { updateTopicQdrantIds } from './write.js';
import type { Chunk, WrittenTopic, SourceChunkPayload, TopicSummaryPayload } from '../types.js';

async function embedTexts(texts: string[]): Promise<number[][]> {
  const vectors: number[][] = [];

  for (let i = 0; i < texts.length; i += EMBEDDING_BATCH_SIZE) {
    const batch = texts.slice(i, i + EMBEDDING_BATCH_SIZE);
    const response = await openai.embeddings.create({
      model: EMBEDDING_MODEL,
      input: batch,
    });
    for (const item of response.data) {
      vectors.push(item.embedding);
    }
  }

  return vectors;
}

export async function embedAndStore(
  chunks: Chunk[],
  writtenTopics: WrittenTopic[],
  onProgress: (embedded: number) => void,
): Promise<void> {
  // Build chunk → topic mapping.
  // The AI labels chunks as [CHUNK-<uuid>] so sourceChunkIds may carry that prefix;
  // strip it so lookups match the bare UUID used as the Qdrant point ID.
  const normalise = (id: string): string => id.replace(/^CHUNK-/i, '');

  const chunkTopicMap = new Map<string, string>();
  const chunkThemeMap = new Map<string, string>();
  for (const topic of writtenTopics) {
    for (const rawId of topic.sourceChunkIds) {
      const chunkId = normalise(rawId);
      chunkTopicMap.set(chunkId, topic.id);
      chunkThemeMap.set(chunkId, topic.themeId);
    }
  }

  // -------------------------------------------------------------------------
  // Source chunks
  // -------------------------------------------------------------------------
  const chunkTexts = chunks.map(c => c.text);
  const chunkVectors = await embedTexts(chunkTexts);

  const sourcePoints = chunks.map((c, i) => {
    const vector = chunkVectors[i];
    if (!vector) throw new Error(`Missing embedding for chunk index ${i}`);

    const payload: SourceChunkPayload = {
      source_document_id: c.documentId,
      chunk_index: c.index,
      text: c.text,
      theme_id: chunkThemeMap.get(c.id) ?? null,
      topic_id: chunkTopicMap.get(c.id) ?? null,
      format: c.format,
    };

    return { id: c.id, vector, payload };
  });

  // Upsert in batches
  for (let i = 0; i < sourcePoints.length; i += EMBEDDING_BATCH_SIZE) {
    const batch = sourcePoints.slice(i, i + EMBEDDING_BATCH_SIZE);
    await qdrant.upsert(QDRANT_COLLECTIONS.SOURCE_CHUNKS, { points: batch });
    onProgress(i + batch.length);
  }

  // -------------------------------------------------------------------------
  // Topic summaries
  // -------------------------------------------------------------------------
  const topicTexts = writtenTopics.map(t => t.body);
  const topicVectors = await embedTexts(topicTexts);

  const summaryPoints = writtenTopics.map((topic, i) => {
    const vector = topicVectors[i];
    if (!vector) throw new Error(`Missing embedding for topic index ${i}`);

    const payload: TopicSummaryPayload = {
      topic_id: topic.id,
      theme_id: topic.themeId,
      title: topic.title,
      text: topic.body,
    };

    return { id: uuid(), vector, payload };
  });

  for (let i = 0; i < summaryPoints.length; i += EMBEDDING_BATCH_SIZE) {
    const batch = summaryPoints.slice(i, i + EMBEDDING_BATCH_SIZE);
    await qdrant.upsert(QDRANT_COLLECTIONS.TOPIC_SUMMARIES, { points: batch });
  }

  // -------------------------------------------------------------------------
  // Update topics.qdrant_chunk_ids in PocketBase
  // -------------------------------------------------------------------------
  for (const topic of writtenTopics) {
    const qdrantIds = topic.sourceChunkIds
      .map(id => normalise(id))
      .filter(id => chunkTopicMap.get(id) === topic.id);
    if (qdrantIds.length > 0) {
      await updateTopicQdrantIds(topic.id, qdrantIds);
    }
  }
}