Advanced Indexing Patterns

import os
import uuid

from langchain.chat_models import init_chat_model
from langchain_classic.retrievers import MultiVectorRetriever
from langchain_core.documents import Document
from langchain_core.stores import InMemoryStore
from langchain_core.vectorstores import VectorStore
from pydantic import BaseModel, Field


class Questions(BaseModel):
    questions: list[str] = Field(min_length=3, max_length=5)


writer = init_chat_model(os.environ["INDEXING_MODEL"], temperature=0).with_structured_output(Questions)


def build_retriever(chunks: list[Document], vector_store: VectorStore) -> MultiVectorRetriever:
    retriever = MultiVectorRetriever(vectorstore=vector_store, docstore=InMemoryStore(), id_key="doc_id",
                                     search_kwargs={"k": 8})
    ids = [str(uuid.uuid5(uuid.NAMESPACE_URL, c.metadata["chunk_id"])) for c in chunks]
    derived = []
    for doc_id, chunk in zip(ids, chunks):
        qs = writer.invoke("Write the questions a reader could answer with this passage only:\n\n" + chunk.page_content)
        derived += [Document(page_content=q, metadata={"doc_id": doc_id, **chunk.metadata}) for q in qs.questions]
    retriever.vectorstore.add_documents(derived)
    retriever.docstore.mset(list(zip(ids, chunks)))
    return retriever

import os

from langchain.chat_models import init_chat_model
from langchain_core.documents import Document
from pydantic import BaseModel, Field


class Propositions(BaseModel):
    propositions: list[str] = Field(description="Atomic, self-contained facts; resolve every pronoun and reference")


extractor = init_chat_model(os.environ["INDEXING_MODEL"], temperature=0).with_structured_output(Propositions)


def to_propositions(chunk: Document) -> list[Document]:
    result = extractor.invoke(
        "Decompose the passage into propositions. Each proposition states exactly one fact, names its subject "
        "explicitly and can be understood without the passage.\n\nPassage:\n" + chunk.page_content)
    return [Document(page_content=p, metadata={**chunk.metadata, "parent_chunk_id": chunk.metadata["chunk_id"],
                                               "granularity": "proposition"})
            for p in result.propositions]

from collections import defaultdict
from collections.abc import Callable

from langchain_core.documents import Document


def auto_merge(children: list[Document], parent_size: Callable[[str], int],
               load_parent: Callable[[str], Document], ratio: float = 0.5) -> list[Document]:
    """Replace retrieved children by their parent when at least `ratio` of the parent's children were retrieved."""
    by_parent: dict[str, list[Document]] = defaultdict(list)
    for child in children:
        by_parent[child.metadata["parent_id"]].append(child)
    merged, emitted = [], set()
    for child in children:                                   # keep the retrieval order
        pid = child.metadata["parent_id"]
        if len(by_parent[pid]) / parent_size(pid) >= ratio:
            if pid not in emitted:
                merged.append(load_parent(pid))
                emitted.add(pid)
        else:
            merged.append(child)
    return merged

package com.example.docsassistant.rag.indexing;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.function.ToIntFunction;

import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;

public class AutoMergingPostProcessor implements DocumentPostProcessor {

    private final ToIntFunction<String> parentSize;
    private final Function<String, Document> loadParent;
    private final double ratio;

    public AutoMergingPostProcessor(ToIntFunction<String> parentSize, Function<String, Document> loadParent,
                                    double ratio) {
        this.parentSize = parentSize;
        this.loadParent = loadParent;
        this.ratio = ratio;
    }

    @Override
    public List<Document> process(Query query, List<Document> children) {
        Map<String, Integer> counts = new LinkedHashMap<>();
        for (Document child : children) {
            counts.merge((String) child.getMetadata().get("parent_id"), 1, Integer::sum);
        }
        List<Document> merged = new ArrayList<>();
        Set<String> emitted = new HashSet<>();
        for (Document child : children) {
            String parentId = (String) child.getMetadata().get("parent_id");
            if ((double) counts.get(parentId) / parentSize.applyAsInt(parentId) >= ratio) {
                if (emitted.add(parentId)) {
                    merged.add(loadParent.apply(parentId));
                }
            } else {
                merged.add(child);
            }
        }
        return merged;
    }
}

package com.example.docsassistant.l4j.indexing;

import java.util.ArrayList;
import java.util.List;
import java.util.function.BiFunction;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.splitter.DocumentBySentenceSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;

public final class SentenceWindow {

    private SentenceWindow() {
    }

    /** One segment per sentence, numbered within its document so that neighbours can be found later. */
    public static List<TextSegment> sentences(Document document, String source) {
        String[] sentences = new DocumentBySentenceSplitter(300, 0).split(document.text());   // sentence detection only
        List<TextSegment> numbered = new ArrayList<>();
        for (int i = 0; i < sentences.length; i++) {
            numbered.add(TextSegment.from(sentences[i],
                    document.metadata().copy().put("source", source).put("sentence", i)));
        }
        return numbered;
    }

    /** Wraps a sentence retriever: each hit is replaced by the window of w sentences around it. */
    public static ContentRetriever windowed(ContentRetriever sentenceRetriever, int w,
                                            BiFunction<String, int[], String> loadRange) {
        return (Query query) -> sentenceRetriever.retrieve(query).stream()
                .map(hit -> {
                    var meta = hit.textSegment().metadata();
                    int i = meta.getInteger("sentence");
                    String window = loadRange.apply(meta.getString("source"), new int[] {Math.max(0, i - w), i + w});
                    return Content.from(TextSegment.from(window, meta));
                })
                .toList();
    }
}

import os
from collections import defaultdict

from langchain.chat_models import init_chat_model
from langchain_core.documents import Document

summariser = init_chat_model(os.environ["INDEXING_MODEL"], temperature=0)
LEVELS = ["section", "page", "component"]              # leaf chunks carry all three keys (full paths) in metadata


def summarise(texts: list[str]) -> str:
    return summariser.invoke(
        "Summarise the key facts, names and versions in these passages for search. Be specific.\n\n"
        + "\n\n---\n\n".join(texts)).text


def build_tree(chunks: list[Document]) -> list[Document]:
    """Return summary documents for every level; index them together with the leaf chunks."""
    nodes, summaries = chunks, []
    for level in LEVELS:
        groups: dict[str, list[Document]] = defaultdict(list)
        for node in nodes:
            groups[node.metadata[level]].append(node)
        nodes = [Document(page_content=summarise([n.page_content for n in members]),
                          metadata={**{k: v for k, v in members[0].metadata.items() if k != "chunk_id"},
                                    "chunk_id": f"raptor:{level}:{key}", "raptor_level": level,
                                    "children": [n.metadata["chunk_id"] for n in members]})
                 for key, members in groups.items()]
        summaries += nodes
    return summaries

package com.example.docsassistant.rag.indexing;

import java.util.List;
import java.util.function.Function;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;

public class ContextualChunkTransformer implements DocumentTransformer {

    private static final String PROMPT = """
            <document>
            {document}
            </document>
            Here is a chunk from the document:
            <chunk>
            {chunk}
            </chunk>
            Write one or two sentences that situate this chunk within the document, to improve search retrieval
            of the chunk. Answer only with those sentences.""";

    private final ChatClient chatClient;
    private final Function<Document, String> fullDocument;   // loads the parent document's text

    public ContextualChunkTransformer(ChatClient.Builder builder, Function<Document, String> fullDocument) {
        this.chatClient = builder.build();
        this.fullDocument = fullDocument;
    }

    @Override
    public List<Document> apply(List<Document> chunks) {
        return chunks.stream().map(chunk -> {
            String context = chatClient.prompt()
                    .user(u -> u.text(PROMPT)
                            .param("document", fullDocument.apply(chunk))
                            .param("chunk", chunk.getText()))
                    .call()
                    .content();
            return chunk.mutate().text(context + "\n\n" + chunk.getText())
                    .metadata("context", context).build();
        }).toList();
    }
}