Ingestion Pipelines & Freshness

import os

from langchain_classic.indexes import SQLRecordManager, index
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore

record_manager = SQLRecordManager(namespace="pgvector/docsassistant", db_url=os.environ["RECORD_MANAGER_DB_URL"])
record_manager.create_schema()


def sync(documents: list[Document], vector_store: VectorStore) -> dict:
    """Write new and changed chunks, delete stale chunks of the same sources, skip unchanged ones."""
    result = index(
        documents,
        record_manager,
        vector_store,
        cleanup="incremental",        # delete old chunks of every source present in this batch
        source_id_key="source",       # the metadata key that identifies the source document
        key_encoder="sha256",
    )
    return dict(result)               # {'num_added': .., 'num_updated': .., 'num_skipped': .., 'num_deleted': ..}


def delete_source(source: str, vector_store: VectorStore) -> int:
    """Remove every chunk of a deleted source document."""
    keys = record_manager.list_keys(group_ids=[source])
    if keys:
        vector_store.delete(ids=keys)
        record_manager.delete_keys(keys)
    return len(keys)

import json
import os

from confluent_kafka import Consumer
from langchain_core.documents import Document

from docsassistant.ingest.indexing import delete_source, sync
from docsassistant.ingest.pipeline import to_chunks, vector_store   # ladder + redact() + splitter

consumer = Consumer({
    "bootstrap.servers": os.environ["KAFKA_BOOTSTRAP_SERVERS"],
    "group.id": "docsassistant-reindex",
    "enable.auto.commit": False,        # commit only after the index write
    "auto.offset.reset": "earliest",
})
consumer.subscribe(["docs.public.documents"])

while True:
    msg = consumer.poll(1.0)
    if msg is None or msg.error():
        continue
    if msg.value() is None:              # tombstone after a delete: nothing left to do
        consumer.commit(message=msg)
        continue
    payload = json.loads(msg.value())["payload"]
    row = payload["after"] or payload["before"]
    source = f"cms/{row['id']}"
    if payload["op"] == "d":
        delete_source(source, vector_store)
    else:                                # c, u, r
        doc = Document(page_content=row["body"], metadata={"source": source, "acl": row["acl"]})
        sync(to_chunks(doc), vector_store)
    consumer.commit(message=msg)

package com.example.docsassistant.rag.ingest;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.IntStream;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.kafka.annotation.KafkaListener;
import org.springframework.kafka.support.Acknowledgment;
import org.springframework.stereotype.Component;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.json.JsonMapper;

@Component
public class CdcReindexListener {

    private final VectorStore vectorStore;
    private final TokenTextSplitter splitter = TokenTextSplitter.builder().withChunkSize(512).build();
    private final JsonMapper json = JsonMapper.builder().build();

    public CdcReindexListener(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    // requires spring.kafka.listener.ack-mode=manual
    @KafkaListener(topics = "docs.public.documents", groupId = "docsassistant-reindex")
    public void onChange(ConsumerRecord<String, String> record, Acknowledgment ack) {
        if (record.value() == null) {            // tombstone
            ack.acknowledge();
            return;
        }
        JsonNode payload = json.readTree(record.value()).path("payload");
        JsonNode row = payload.path("after").isNull() ? payload.path("before") : payload.path("after");
        String source = "cms/" + row.path("id").asString();

        var filter = new FilterExpressionBuilder().eq("source", source).build();
        vectorStore.delete(filter);                                  // remove the old version (or the deleted doc)
        if (!"d".equals(payload.path("op").asString())) {
            Document doc = new Document(row.path("body").asString(),
                    Map.of("source", source, "acl", row.path("acl").asString()));
            List<Document> chunks = splitter.apply(List.of(doc));
            List<Document> withIds = IntStream.range(0, chunks.size())
                    .mapToObj(i -> chunks.get(i).mutate()
                            .id(UUID.nameUUIDFromBytes((source + ":" + i).getBytes(StandardCharsets.UTF_8)).toString())
                            .build())
                    .toList();
            vectorStore.add(withIds);
        }
        ack.acknowledge();                                           // commit only after the write succeeded
    }
}

from langchain_core.documents import Document
from presidio_analyzer import AnalyzerEngine
from presidio_anonymizer import AnonymizerEngine
from presidio_anonymizer.entities import OperatorConfig

analyzer = AnalyzerEngine()            # needs a spaCy model, e.g. en_core_web_lg, installed
anonymizer = AnonymizerEngine()
ENTITIES = ["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER", "IBAN_CODE", "CREDIT_CARD", "IP_ADDRESS"]


def redact(doc: Document) -> Document:
    findings = analyzer.analyze(text=doc.page_content, entities=ENTITIES, language="en", score_threshold=0.6)
    result = anonymizer.anonymize(
        text=doc.page_content,
        analyzer_results=findings,
        operators={"DEFAULT": OperatorConfig("replace")},     # default replacement is "<ENTITY_TYPE>"
    )
    counts: dict[str, int] = {}
    for item in result.items:
        counts[item.entity_type] = counts.get(item.entity_type, 0) + 1
    return Document(page_content=result.text, metadata={**doc.metadata, "pii_redacted": counts})

package com.example.docsassistant.rag.ingest;

import java.util.List;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

@Service
public class EtlPipeline {

    private final VectorStore vectorStore;
    private final KeywordMetadataEnricher keywords;
    private final TokenTextSplitter splitter = TokenTextSplitter.builder().withChunkSize(512).build();

    public EtlPipeline(VectorStore vectorStore, ChatModel chatModel) {
        this.vectorStore = vectorStore;
        this.keywords = KeywordMetadataEnricher.builder(chatModel).keywordCount(5).build();
    }

    public int ingest(Resource resource, String source, String acl) {
        DocumentReader reader = new TikaDocumentReader(resource);                 // extract
        List<Document> docs = reader.get().stream()
                .map(d -> d.mutate().metadata("source", source).metadata("acl", acl).build())
                .toList();
        List<Document> chunks = keywords.apply(splitter.apply(docs));            // transform
        vectorStore.accept(chunks);                                               // load
        return chunks.size();
    }
}

package com.example.docsassistant.l4j.ingest;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;

public class Ingestion {

    private final EmbeddingStore<TextSegment> store;
    private final EmbeddingStoreIngestor ingestor;

    public Ingestion(EmbeddingStore<TextSegment> store, EmbeddingModel embeddingModel) {
        this.store = store;
        this.ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(DocumentSplitters.recursive(1500, 150))
                .textSegmentTransformer(segment -> TextSegment.from(
                        segment.metadata().getString("title") + "\n\n" + segment.text(),   // context header
                        segment.metadata()))
                .embeddingModel(embeddingModel)
                .embeddingStore(store)
                .build();
    }

    public IngestionResult replace(Document document, String source, String acl) {
        document.metadata().put("source", source).put("acl", acl);
        store.removeAll(metadataKey("source").isEqualTo(source));      // delete the previous version
        return ingestor.ingest(document);
    }

    public void delete(String source) {
        store.removeAll(metadataKey("source").isEqualTo(source));
    }
}