rag - 检索文档与知识库（Repository）

RAG 里最重要的一个工作，就是“检索”。就是在一个“知识库（Repository）”里找“文档（Document）”。

1、概念

检索的内容，可称为文档（Document）。提供文档索检服务的对象，就是知识库（Repository）了。知识库可分为：

• 只读知识库，只用于检索的
• 可存储知识库，可用于检索，同时提供存储管理

可存储知识库，还会用到两个重要的工具：

• DocumentLoader，文档加载器
• DocumentSplitter，文档分割器

比如有个 pdf 文档，需要通过 DocumentLoader 加载，转为 Document 列表，然后存入 Repository。有时候 Document 太大，可能还要用 DocumentSplitter 分割成多个小 Document。

2、DocumentLoader 及内置的适配

• 基础接口

public interface DocumentLoader {
    //附加元数据
    DocumentLoader additionalMetadata(String key, Object value);

    //附加元数据
    DocumentLoader additionalMetadata(Map<String, Object> metadata);

    //加载文档
    List<Document> load() throws IOException;
}

• 内置的适配（可以根据业务，按需定制）

3、DocumentSplitter 及内置的适配

• 基础接口

public interface DocumentSplitter {
    //分割
    default List<Document> split(String text) {
        return split(Arrays.asList(new Document(text)));
    }

   //分割
    List<Document> split(List<Document> documents);
}

• 内置的适配（可以根据业务，按需定制）

• 使用示例

private void load(RepositoryStorable repository, String file) throws IOException {
    //加载器
    HtmlSimpleLoader loader = new HtmlSimpleLoader(new File(file));
    
    //加载后再分割（按需）
    List<Document> documents = new SplitterPipeline() 
            .next(new RegexTextSplitter("\n\n"))
            .next(new TokenSizeTextSplitter(500))
            .split(loader.load());
    
    //入库
    repository.insert(documents); 
}

4、Repository 及内置的适配

基础接口

public interface Repository {
    //检索
    default List<Document> search(String query) throws IOException {
        return search(new QueryCondition(query));
    }

    //检索
    List<Document> search(QueryCondition condition) throws IOException;
}

public interface RepositoryStorable extends Repository {
    //插入
    void insert(List<Document> documents) throws IOException;

    //删除
    void delete(String... ids) throws IOException;

    //是否存在 
    boolean exists(String id) throws IOException;
}

• 内置的适配（可以根据业务，按需定制）

7、RAG 要素联合演示

知识库准备

//构建 embeddingModel
EmbeddingModel embeddingModel = EmbeddingModel.of(embedding_apiUrl)
        .apiKey(embedding_apiKey)
        .model(embedding_model)
        .build();

//构建 repository
InMemoryRepository repository = new InMemoryRepository(embeddingModel);

//加载文档并存储
PdfLoader loader = new PdfLoader(new File("ticket.pdf")).additionalMetadata("file", "ticket.pdf");

//再次（按需）组合切割
List<Document> documents = new SplitterPipeline() //2.分割文档（确保不超过 max-token-size）
                .next(new RegexTextSplitter("\n\n"))
                .next(new TokenSizeTextSplitter(500))
                .split(loader.load());

//存储仓库
repository.insert(documents);

检索应用

//构建 chatModel
ChatModel chatModel = ChatModel.of(chat_apiUrl)
                .provider(provider)
                .model(model)
                .build();

//用户输入消息
String message = "刘德华今年有几场演唱会？";

//1. 检索
List<Document> context = repository.search(message);

//2. 消息增强
ChatMessage chatMessage = ChatMessage.augment(message, context);

//3. 提交大模型
ChatResponse resp = chatModel.prompt(chatMessage).call();

//打印结果
System.out.println(resp.getMessage());