[Spring AI] Spring AI로 RAG 구현하기 — pgvector, RetrievalAugmentationAdvisor, VectorStore

1. 전체 구성

1) 사용 기술

• Spring AI : RAG 구현을 위한 프레임워크
• OpenAI : 임베딩 모델(text-embedding-3-small), Chat 모델(gpt-4o-mini)
• Ollama : 로컬 Chat 모델(llama3.1:8b)
• pgvector : Vector DB

 

2) 전체 흐름

[사전 준비]
문서 텍스트
    → OpenAI 임베딩 모델이 벡터로 변환
    → pgvector에 저장

[질문 시]
사용자 질문
    → 임베딩 모델이 벡터로 변환
    → pgvector에서 코사인 유사도로 유사한 문서 검색
    → 임계치 이상 문서를 프롬프트 context에 삽입
    → LLM이 context를 참고해서 답변 생성

 

2. 의존성 및 설정

1) build.gradle

dependencies {
    implementation 'org.springframework.boot:spring-boot-starter-webflux'
    implementation 'org.springframework.ai:spring-ai-starter-model-openai'
    implementation 'org.springframework.ai:spring-ai-starter-model-ollama'
    implementation 'org.springframework.ai:spring-ai-starter-vector-store-pgvector'
    implementation 'org.springframework.ai:spring-ai-rag'
    compileOnly 'org.projectlombok:lombok'
    annotationProcessor 'org.projectlombok:lombok'
    developmentOnly 'org.springframework.boot:spring-boot-devtools'
}

 

2) application.properties

# OpenAI
spring.ai.openai.api-key=${OPENAI_API_KEY}
spring.ai.openai.chat.model=gpt-4o-mini
spring.ai.openai.embedding.model=text-embedding-3-small

# Ollama
spring.ai.ollama.base-url=http://localhost:11434
spring.ai.ollama.chat.options.model=llama3.1:8b
spring.ai.ollama.chat.options.temperature=0.2

# pgvector
spring.datasource.url=jdbc:postgresql://localhost:5432/postgres
spring.datasource.username=postgres
spring.datasource.password=postgres

spring.ai.vectorstore.pgvector.index-type=hnsw
spring.ai.vectorstore.pgvector.distance-type=cosine-distance
spring.ai.vectorstore.pgvector.dimensions=1536
spring.ai.vectorstore.pgvector.metadata-fields=type,docId,topic

# ollama embedding 비활성화 (pgvector는 OpenAI embedding 사용)
spring.ai.ollama.embedding.enabled=false

• ollama와 openai를 함께 쓸 경우 EmbeddingModel Bean이 두 개 등록되어 충돌이 발생한다
• pgvector는 OpenAI EmbeddingModel을 사용하도록 ollama embedding을 비활성화해야 한다

 

3. ChatClient 설정—ChatConfig.java

@Configuration
public class ChatConfig {

    @Bean
    ChatMemory chatMemory() {
        return MessageWindowChatMemory.builder()
                .maxMessages(20).build();
    }

    @Bean("openAiChatClient")
    @Primary
    ChatClient openAiChatClient(OpenAiChatModel chatModel,
            ChatMemory chatMemory, WeatherTool weatherTool) {
        return ChatClient.builder(chatModel)
                .defaultAdvisors(MessageChatMemoryAdvisor
                        .builder(chatMemory).build())
                .defaultTools(weatherTool)
                .build();
    }

    @Bean("ollamaChatClient")
    ChatClient ollamaChatClient(OllamaChatModel chatModel,
            ChatMemory chatMemory, WeatherTool weatherTool) {
        return ChatClient.builder(chatModel)
                .defaultAdvisors(MessageChatMemoryAdvisor
                        .builder(chatMemory).build())
                .defaultTools(weatherTool)
                .build();
    }

    // ollama와 openai EmbeddingModel 충돌 방지
    @Bean
    @Primary
    public EmbeddingModel embeddingModel(OpenAiEmbeddingModel embeddingModel) {
        return embeddingModel;
    }
}

• @RequiredArgsConstructor(Lombok)와 @Qualifier를 함께 쓰면 @Qualifier가 무시된다
• 생성자를 직접 작성해서 @Qualifier를 명시해야 한다

 

4. 문서 저장 — VectorStoreBootstrap.java

@Component
@RequiredArgsConstructor
public class VectorStoreBootstrap {

    private static final String INIT_MARKER_TEXT = "__VECTORSTORE_V1";
    private final VectorStore vectorStore;

    @PostConstruct
    public void bootstrap() {
        // 이미 초기화된 경우 스킵
        if (isAlreadyBootstrapped(vectorStore)) return;

        List<Document> documents = List.of(
            new Document(INIT_MARKER_TEXT,
                Map.of("type", "bootstrap", "version", "v1", "app", "myapp")),
            new Document("Spring AI rocks!!",
                Map.of("docId", "seed-1", "topic", "spring-ai")),
            new Document("세상은 크고 구원은 세상의 모퉁이에 있습니다.",
                Map.of("docId", "seed-2", "type", "quote")),
            new Document("과거를 향해 앞으로 걸어가고 미래를 향해 되돌아갑니다.",
                Map.of("docId", "seed-3", "type", "quote", "theme", "time"))
        );

        vectorStore.add(documents);
    }

    boolean isAlreadyBootstrapped(VectorStore vectorStore) {
        var results = vectorStore.similaritySearch(
            SearchRequest.builder()
                .query(INIT_MARKER_TEXT)
                .topK(1)
                .filterExpression("type == 'bootstrap' && app == 'myapp'")
                .build());
        return results != null && !results.isEmpty();
    }
}

• @PostConstruct : 앱 시작 시 자동 실행
• isAlreadyBootstrapped : 중복 저장 방지. 이미 초기화 마커가 있으면 스킵
• vectorStore.add() : Spring AI가 자동으로 임베딩 변환 후 pgvector에 저장
• metadata : filterExpression으로 필터링할 때 사용. metadata-fields에 등록된 필드만 필터링 가능

 

5. RAG 구현 — ChatService

1) RAG Advisor 초기화

private static final double SIM_THRESHOLD = 0.35;
private static final int TOP_K = 1;

// 사용자 질문을 벡터로 변환 후 VectorStore에서 유사한 문서를 찾아오는 검색기
this.retriever = VectorStoreDocumentRetriever.builder()
        .vectorStore(vectorStore)
        .similarityThreshold(SIM_THRESHOLD)
        .topK(TOP_K)
        .build();

// 검색된 문서를 프롬프트 context에 자동으로 삽입해주는 RAG Advisor
this.ragAdvisor = RetrievalAugmentationAdvisor.builder()
        .documentRetriever(this.retriever)
        .build();

• SIM_THRESHOLD : 유사도 임계치. 이 값 이상인 문서만 검색 결과로 반환된다
• TOP_K : 최대 몇 개의 문서를 가져올지. 1로 설정하면 가장 유사한 문서 1개만 반환
• VectorStoreDocumentRetriever : 질문을 벡터로 변환 후 VectorStore에서 유사한 문서를 찾아오는 검색기
• RetrievalAugmentationAdvisor : 검색된 문서를 프롬프트 context에 자동으로 삽입해주는 Advisor

 

2) 유사도 검색 — hasRelevantDocs

private boolean hasRelevantDocs(String userInput) {
    var results = vectorStore.similaritySearch(
        SearchRequest.builder()
            .query(userInput)
            .topK(TOP_K)
            .filterExpression("type == 'quote'") // 메타데이터 필터
            .similarityThreshold(SIM_THRESHOLD)
            .build()
    );

    // 유사도 점수 로깅
    results.forEach(doc ->
        System.out.println("score: " + doc.getScore() + " | " + doc.getText())
    );

    return results != null && !results.isEmpty();
}

• similaritySearch : 질문을 벡터로 변환 후 VectorStore에서 유사한 문서를 검색
• filterExpression : SQL WHERE절과 같은 역할. metadata-fields에 등록된 필드만 사용 가능
• doc.getScore() : 유사도 점수 확인. 임계치 조정이 필요할 때 활용
• RAG Advisor에도 threshold가 있지만 hasRelevantDocs에서 먼저 확인해 관련 문서가 없으면 불필요한 RAG 적용을 막는다

 

3) RAG 적용 — chatMemoryStream

public Flux<String> chatMemoryStream(String userInput, String userId) {

    boolean useRag = hasRelevantDocs(userInput);

    var prompt = ollamaChatClient.prompt()
            .user(userInput)
            .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, userId));

    if (useRag) {
        // advisors()는 새 객체를 반환하므로 반드시 재할당해야 함
        prompt = prompt.advisors(ragAdvisor);
    }

    return prompt.stream()
            .content()
            .map(this::preserveSseLeadingSpace)
            .concatWith(Flux.just("[[END]]"));
}

• hasRelevantDocs로 관련 문서 존재 여부를 먼저 확인
• 관련 문서가 있을 때만 ragAdvisor를 추가해 context 기반 답변 생성
• advisors()는 새 객체를 반환하기 때문에 반드시 재할당해야 적용된다
• concatWith(Flux.just("[[END]]")) : 스트리밍 종료 시그널을 클라이언트에 전달

 

전체 코드

@Service
public class ChatService {

    private final ChatClient openAiChatClient;
    private final ChatClient ollamaChatClient;
    private final VectorStore vectorStore;
    private final VectorStoreDocumentRetriever retriever;
    private final RetrievalAugmentationAdvisor ragAdvisor;

    private static final double SIM_THRESHOLD = 0.35;
    private static final int TOP_K = 1;

    public ChatService(
            @Qualifier("openAiChatClient") ChatClient openAiChatClient,
            @Qualifier("ollamaChatClient") ChatClient ollamaChatClient,
            VectorStore vectorStore) {
        this.openAiChatClient = openAiChatClient;
        this.ollamaChatClient = ollamaChatClient;
        this.vectorStore = vectorStore;

        this.retriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(SIM_THRESHOLD)
                .topK(TOP_K)
                .build();

        this.ragAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(this.retriever)
                .build();
    }

    public Mono<String> chatMemory(String userInput, String userId) {
        return Mono.fromCallable(() -> ollamaChatClient.prompt()
                .user(userInput)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, userId))
                .call()
                .content());
    }

    public Flux<String> chatMemoryStream(String userInput, String userId) {
        boolean useRag = hasRelevantDocs(userInput);

        var prompt = ollamaChatClient.prompt()
                .user(userInput)
                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, userId));

        if (useRag) {
            prompt = prompt.advisors(ragAdvisor);
        }

        return prompt.stream()
                .content()
                .map(this::preserveSseLeadingSpace)
                .concatWith(Flux.just("[[END]]"));
    }

    private boolean hasRelevantDocs(String userInput) {
        var results = vectorStore.similaritySearch(
            SearchRequest.builder()
                .query(userInput)
                .topK(TOP_K)
                .filterExpression("type == 'poem'")
                .similarityThreshold(SIM_THRESHOLD)
                .build()
        );

        results.forEach(doc ->
            System.out.println("score: " + doc.getScore() + " | " + doc.getText())
        );

        return results != null && !results.isEmpty();
    }

    private String preserveSseLeadingSpace(String chunk) {
        if (chunk == null) return "";
        return chunk.startsWith(" ") ? " " + chunk : chunk;
    }
}

6. 주의사항

1) metadata-fields 설정과 테이블 재생성

• filterExpression은 유사도 검색 시 메타데이터 조건으로 검색 범위를 좁히는 필터로, SQL WHERE절과 같은 역할을 한다
• pgvector는 필터링을 위해 해당 필드를 별도 컬럼으로 인덱싱하기 때문에 metadata-fields에 사용할 필드를 미리 등록해야 한다
• 등록하지 않으면 filterExpression 자체가 동작하지 않는다

spring.ai.vectorstore.pgvector.metadata-fields=type,docId,topic,author,theme

• 설정 추가 후 vector_store 테이블을 DROP하고 재시작해야 반영된다

DROP TABLE vector_store;

 

2) 임계치(threshold) 튜닝

• 임계치가 너무 높으면 관련 문서가 있어도 검색되지 않는다
• 임계치가 너무 낮으면 관련 없는 문서가 context에 포함되어 답변 품질이 떨어진다
• doc.getScore()로 실제 유사도 점수를 확인하면서 조정하는 것이 좋다

 

3) 토큰 비용

• 검색된 문서가 프롬프트에 포함되기 때문에 문서 길이만큼 토큰이 추가로 소비된다
• OpenAI 같은 서드파티 API 사용 시 topK와 문서 크기 설정에 주의해야 한다