Function deduplicateChunks

lib/edge-runtime/utils.ts:183–263

(
  chunks: SimilaritySearchResult[],
  nTextChunksInclude: number,
)

}

export function deduplicateChunks(
  chunks: SimilaritySearchResult[],
  nTextChunksInclude: number,
): SimilaritySearchResult[] {
  /** Deduplicate chunks by combining chunks with the same page_url, page_title and
   * section_title.
   *
   * Each doc_chunk in the DB has multiple text chunks. If the doc_chunks with the
   * most similar embeddings are from the same document, we don't want to add the
   * same text twice.
   *
   * E.g. If one doc_chunk contains [a,b,c], the other [b,c,d], then we have a chunk
   * with [a,b,c,d] (no duplication of b and c). **/
  // Deep copy to avoid modifying the original array
  chunks = JSON.parse(JSON.stringify(chunks));
  const getChunkUniquenessString = (chunk: SimilaritySearchResult) =>
    `${chunk.page_url} ${chunk.page_title} ${chunk.section_title}`;

  const chunkOrdering = chunks.map((chunk) => getChunkUniquenessString(chunk));
  chunks.sort((a, b) =>
    `${chunkOrdering.findIndex((str) => str === getChunkUniquenessString(b))} ${
      b.chunk_idx
    }` >
    `${chunkOrdering.findIndex((str) => str === getChunkUniquenessString(a))} ${
      a.chunk_idx
    }`
      ? -1
      : 1,
  );
  const deduped: SimilaritySearchResult[] = [chunks[0]];

  // Don't need to dedupe the first chunk, hence start with i=1
  for (let i = 1; i < chunks.length; i++) {
    const chunk = chunks[i];
    // Check if there's another chunk with the same page_url, page_title and section_title
    const matchedChunk = deduped.find(
      (seenCh) =>
        seenCh.page_url === chunk.page_url &&
        seenCh.page_title === chunk.page_title &&
        seenCh.section_title === chunk.section_title,
    );
    if (matchedChunk) {
      if (
        chunk.chunk_idx > matchedChunk.chunk_idx &&
        // Below check stops gaps in the text chunks
        chunk.chunk_idx - matchedChunk.chunk_idx <=
          matchedChunk.text_chunks.length
      ) {
        chunk.text_chunks.forEach((textChunk) => {
          if (!matchedChunk.text_chunks.includes(textChunk)) {
            // push() adds to the end of the array
            matchedChunk.text_chunks.push(textChunk);
          }
        });
      } else if (
        chunk.chunk_idx < matchedChunk.chunk_idx &&
        matchedChunk.chunk_idx - chunk.chunk_idx <= chunk.text_chunks.length
      ) {