聊聊Spring AI的ETL Pipeline

序

本文主要研究一下Spring AI的ETL Pipeline

DocumentReader

org/springframework/ai/document/DocumentReader.java

public interface DocumentReader extends Supplier<List<Document>> {

    default List<Document> read() {
        return get();
    }

}

有TextReader、JsonReader、JsoupDocumentReader、MarkdownDocumentReader、PagePdfDocumentReader、ParagraphPdfDocumentReader、TikaDocumentReader这些实现

DocumentTransformer

org/springframework/ai/document/DocumentTransformer.java

public interface DocumentTransformer extends Function<List<Document>, List<Document>> {

    default List<Document> transform(List<Document> transform) {
        return apply(transform);
    }

}

有ContentFormatTransformer、KeywordMetadataEnricher、SummaryMetadataEnricher、TokenTextSplitter实现

DocumentWriter

org/springframework/ai/document/DocumentWriter.java

public interface DocumentWriter extends Consumer<List<Document>> {

    default void write(List<Document> documents) {
        accept(documents);
    }

}

有FileDocumentWriter、SimpleVectorStore、AzureVectorStore、CassandraVectorStore、ChromaVectorStore、CoherenceVectorStore、CosmosDBVectorStore、CouchbaseSearchVectorStore、ElasticsearchVectorStore、GemFireVectorStore、HanaCloudVectorStore、HanaCloudVectorStore、MariaDBVectorStore、MilvusVectorStore、MongoDBAtlasVectorStore、Neo4jVectorStore、OpenSearchVectorStore、OracleVectorStore、PgVectorStore、PineconeVectorStore、QdrantVectorStore、RedisVectorStore、TypesenseVectorStore、WeaviateVectorStore这些实现

示例

@BeforeEach
    void setUp() {
        DocumentReader markdownReader = new MarkdownDocumentReader(this.knowledgeBaseResource,
                MarkdownDocumentReaderConfig.defaultConfig());
        this.knowledgeBaseDocuments = markdownReader.read();
        this.pgVectorStore.add(this.knowledgeBaseDocuments);
    }

    @AfterEach
    void tearDown() {
        this.pgVectorStore.delete(this.knowledgeBaseDocuments.stream().map(Document::getId).toList());
    }

    @Test
    void ragBasic() {
        String question = "Where does the adventure of Anacletus and Birba take place?";

        RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
            .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
            .build();

        ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
            .build()
            .prompt(question)
            .advisors(ragAdvisor)
            .call()
            .chatResponse();

        assertThat(chatResponse).isNotNull();

        String response = chatResponse.getResult().getOutput().getText();
        System.out.println(response);
        assertThat(response).containsIgnoringCase("Highlands");

        evaluateRelevancy(question, chatResponse);
    }

这里通过MarkdownDocumentReader读取文档，然后写入到PgVectorStore，之后通过RetrievalAugmentationAdvisor去检索查询

小结

Spring AI提供了ETL(Extract, Transform, and Load) Pipeline用来处于文档，它由DocumentReader、DocumentTransformer、DocumentWriter三个组件组成。

doc

• Retrieval Augmented Generation (RAG) / ETL Pipeline