springai使用chroma向量数据库

chushiyunen

191人浏览 · 2026-06-18 11:10:43

chushiyunen · 2026-06-18 11:10:43 发布

文章目录

开发环境用chromadb。

版本：
springai # 1.0.0-M6
chromadb # 0.5.23，要和springai匹配，所以不能选高版本

有两种方案：
1、python安装的也能用 # 这里用的这种

安装：
pip install chromadb==0.5.23 # 版本要和springai版本兼容

启动：
chroma run --host 0.0.0.0 --port 8000

验证：
http://localhost:8000/docs # 这个地址可以看到chromadb的接口即可

2、docker安装windows本地版
略

使用

maven依赖

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-chroma-store-spring-boot-starter</artifactId>
    <version>1.0.0-M6</version>
</dependency>

创建实体类

springai的Document已经支持chromadb，直接继承就行。

import org.springframework.ai.document.Document;
import java.util.Map;

public class KnowledgeDocument extends Document {

    /**
     * 业务创建时的构造方法（自动生成 ID）
     */
    public KnowledgeDocument(String content, String sourceFile) {
        super(content);
        this.getMetadata().put("source", sourceFile);
        this.getMetadata().put("created_at", System.currentTimeMillis());
    }

    /**
     * 内部全参构造方法（支持传入指定的 ID，用于从数据库/向量库还原对象）
     */
    private KnowledgeDocument(String id, String content, Map<String, Object> metadata) {
        super(id, content, metadata);
    }

    /**
     * 将 Spring AI 的 Document 转为我们的业务实体
     */
    public static KnowledgeDocument from(Document doc) {
        // 直接使用包含 ID 的构造方法，避免调用不存在的 setId 方法
        return new KnowledgeDocument(
                doc.getId(),
                doc.getText(),
                doc.getMetadata()
        );
    }
}

KnowledgeService接口类


public interface KnowledgeService {

    /**
     * 上传并解析文档入库
     */
    void uploadAndIndex(MultipartFile file) throws Exception;

    /**
     * 根据内容检索相关文档
     */
    List<KnowledgeDocument> search(String query, int topK);
}

KnowledgeServiceImpl接口实现类

import com.example.demo.entity.KnowledgeDocument;
import com.example.demo.service.KnowledgeService;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
public class KnowledgeServiceImpl implements KnowledgeService {

    private final VectorStore vectorStore;
    private final TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();

    @Override
    public void uploadAndIndex(MultipartFile file) throws Exception {
        // 1. 读取文件内容
        String content = new String(file.getBytes(), StandardCharsets.UTF_8);

        // 2. 创建原始 Document
        Document rawDoc = new KnowledgeDocument(content, file.getOriginalFilename());

        // 3. 文本分块 (Chunking)
        List<Document> chunks = tokenTextSplitter.apply(List.of(rawDoc));

        // 4. 存入 ChromaDB
        vectorStore.add(chunks);
    }

    @Override
    public List<KnowledgeDocument> search(String query, int topK) {
        SearchRequest request = SearchRequest.builder()
                .query(query)
                .topK(topK)
                .similarityThreshold(0.7)
                .build();

        List<Document> results = vectorStore.similaritySearch(request);
        // 将底层 Document 转换回我们的业务实体
        return results.stream()
                .map(KnowledgeDocument::from)
                .collect(Collectors.toList());
    }
}

controller


import com.example.demo.entity.JsonResult;
import com.example.demo.entity.KnowledgeDocument;
import com.example.demo.service.KnowledgeService;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.List;

/*
 * 知识库管理表
 */
@RestController()
@RequestMapping("/api")
@Slf4j
public class KnowledgeController {

    @Autowired
    private KnowledgeService knowledgeService;

    /**
     * 上传并解析文档入库
     */
    @PostMapping("/knowledge/upload")
    public JsonResult upload(@RequestParam("file") MultipartFile file) {
        String methodName = "知识库文档上传";
        JsonResult result = JsonResult.ok();
        try {
            log.info(methodName + "_操作开始, fileName={}", file.getOriginalFilename());

            // 参数校验
            if (file.isEmpty()) {
                throw new IllegalArgumentException("上传的文件不能为空");
            }

            knowledgeService.uploadAndIndex(file);

            result = JsonResult.ok("文档上传并解析成功");
            log.info(methodName + "_操作完成, fileName={}", file.getOriginalFilename());
            return result;
        } catch (IllegalArgumentException e) {
            log.error(methodName + "_操作失败, error=", e);
            return JsonResult.fail("-1", "操作失败," + e.getMessage());
        } catch (Exception e) {
            log.error(methodName + "_操作异常, error=", e);
            return JsonResult.fail("-1", "操作异常,请联系系统管理员");
        }
    }

    /**
     * 根据内容检索相关文档
     */
    @GetMapping("/knowledge/search")
    public JsonResult search(@RequestParam("query") String query,
                             @RequestParam(value = "topK", defaultValue = "3") Integer topK) {
        String methodName = "知识库内容检索";
        JsonResult result = JsonResult.ok();
        try {
            log.info(methodName + "_查询操作开始, query={}, topK={}", query, topK);

            // 参数校验
            if (StringUtils.isEmpty(query)) {
                throw new IllegalArgumentException("检索关键词不能为空");
            }

            List<KnowledgeDocument> documents = knowledgeService.search(query, topK);

            result = JsonResult.ok(documents);
            log.info(methodName + "_操作完成, 检索到 {} 条结果", documents.size());
            return result;
        } catch (IllegalArgumentException e) {
            log.error(methodName + "_操作失败, error=", e);
            return JsonResult.fail("-1", "操作失败," + e.getMessage());
        } catch (Exception e) {
            log.error(methodName + "_操作异常, error=", e);
            return JsonResult.fail("-1", "操作异常,请联系系统管理员");
        }
    }
}