Skip to content

[Feat] RAG 파이프라인에 Reranker 로직 구현 #200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeong1112 merged 8 commits into from
Mar 24, 2026
Merged

[Feat] RAG 파이프라인에 Reranker 로직 구현 #200

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ff484e1
[Feat] #199 Cohere 관련 환경변수 추가
jeong1112 Mar 23, 2026
284b23e
[Feat] #199 Cohere 호출용 요청, 응답 DTO 구현
jeong1112 Mar 23, 2026
23d110d
[Feat] #199 Cohere Rerank API 호출 로직 구현
jeong1112 Mar 23, 2026
4774af1
[Feat] #199 정책 문서 Retrieval 과정에서 Reranker 로직 추가
jeong1112 Mar 23, 2026
9098e47
[Refactor] #199 수정된 Member 정보에 맞게 Context 수정
jeong1112 Mar 23, 2026
ef3272b
[Chore] #199 로그 삭제
jeong1112 Mar 23, 2026
1d8087b
[Chore] #199 중복되는 컨텍스트 출력 제거
jeong1112 Mar 24, 2026
9c926fe
[Chore] #199 Cohere API Fallback 시 리스트 사이즈 조절
jeong1112 Mar 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 30 additions & 11 deletions src/main/java/org/sopt/kareer/global/external/ai/builder/context/MemberContextBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import org.sopt.kareer.domain.member.repository.MemberVisaRepository;
import org.springframework.stereotype.Component;

import java.time.LocalDate;
import java.util.List;

import static org.sopt.kareer.domain.member.exception.MemberErrorCode.MEMBER_NOT_FOUND;
Expand All @@ -27,17 +28,23 @@ public MemberAndContext load(Long memberId) {

StringBuilder sb = new StringBuilder();
sb.append("User Profile\n");
sb.append("- name: ").append(nullToEmpty(member.getName())).append("\n");
sb.append("- country: ").append(member.getCountry() != null ? member.getCountry().name() : "").append("\n");
sb.append("- primaryMajor: ").append(nullToEmpty(member.getPrimaryMajor())).append("\n");
sb.append("- secondaryMajor: ").append(nullToEmpty(member.getSecondaryMajor())).append("\n");
sb.append("- targetJob: ").append(nullToEmpty(member.getTargetJob())).append("\n");
sb.append("- languageLevel: ").append(member.getLanguageLevel() != null ? member.getLanguageLevel().name() : "").append("\n");
sb.append("- degree: ").append(member.getDegree() != null ? member.getDegree().name() : "").append("\n");
sb.append("- graduationDate: ").append(member.getGraduationDate() != null ? member.getGraduationDate() : "").append("\n");
sb.append("- expectedGraduationDate: ").append(member.getExpectedGraduationDate() != null ? member.getExpectedGraduationDate() : "").append("\n");
sb.append("- targetJobSkill: ").append(nullToEmpty(member.getTargetJobSkill())).append("\n");
sb.append("- personalBackground: ").append(nullToEmpty(member.getPersonalBackground())).append("\n");
appendLine(sb, "name", member.getName());
appendLine(sb, "email", member.getEmail());
appendLine(sb, "birthDate", member.getBirthDate());
appendLine(sb, "country", member.getCountry() != null ? member.getCountry().name() : "");
appendLine(sb, "university", member.getUniversity());
appendLine(sb, "primaryMajor", member.getPrimaryMajor());
appendLine(sb, "secondaryMajor", member.getSecondaryMajor());
appendLine(sb, "targetJob", member.getTargetJob());
appendLine(sb, "targetJobSkill", member.getTargetJobSkill());
appendLine(sb, "fieldsOfInterest", member.getFieldsOfInterest());
appendLine(sb, "preparationStatus", member.getPreparationStatus());
appendLine(sb, "personalBackground", member.getPersonalBackground());
appendLine(sb, "languageLevel", member.getLanguageLevel() != null ? member.getLanguageLevel().name() : "");
appendLine(sb, "englishLevel", member.getEnglishLevel() != null ? member.getEnglishLevel().name() : "");
appendLine(sb, "degree", member.getDegree() != null ? member.getDegree().name() : "");
appendLine(sb, "graduationDate", member.getGraduationDate());
appendLine(sb, "expectedGraduationDate", member.getExpectedGraduationDate());

sb.append("Visa Info\n");
for (MemberVisa v : visas) {
Expand All @@ -56,5 +63,17 @@ private String nullToEmpty(String s) {
return s == null ? "" : s;
}

private void appendLine(StringBuilder sb, String key, String value) {
sb.append("- ").append(key).append(": ").append(nullToEmpty(value)).append("\n");
}

private void appendLine(StringBuilder sb, String key, LocalDate value) {
sb.append("- ").append(key).append(": ").append(toText(value)).append("\n");
}

private String toText(LocalDate value) {
return value == null ? "" : value.toString();
}

public record MemberAndContext(Member member, String contextText) {}
}
13 changes: 11 additions & 2 deletions src/main/java/org/sopt/kareer/global/external/ai/service/PolicyDocumentRetriever.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import org.sopt.kareer.domain.member.entity.MemberVisa;
import org.sopt.kareer.global.external.ai.builder.query.PolicyQueryBuilder;
import org.sopt.kareer.global.external.ai.properties.RoadmapRagProperties;
import org.sopt.kareer.global.external.cohere.service.CohereRerankClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
Expand All @@ -18,16 +19,24 @@ public class PolicyDocumentRetriever {

private final PgVectorStore policyDocumentVectorStore;
private final RoadmapRagProperties props;
private final CohereRerankClient cohereRerankClient;

public List<Document> retrievePolicy(Member member, MemberVisa visa) {
String query = PolicyQueryBuilder.buildPolicyQuery(member, visa);

return policyDocumentVectorStore.similaritySearch(
List<Document> candidates = policyDocumentVectorStore.similaritySearch(
SearchRequest.builder()
.query(query)
.topK(props.policyTopK())
.topK(props.candidatePoolTopK())
.build()
);

List<Document> reranked = cohereRerankClient.rerank(query, candidates, props.policyTopK());
if(reranked.size() > props.policyTopK()) {
return reranked.subList(0, props.policyTopK());
}
return reranked;

}

}
13 changes: 13 additions & 0 deletions src/main/java/org/sopt/kareer/global/external/cohere/dto/request/CohereRerankRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package org.sopt.kareer.global.external.cohere.dto.request;

import com.fasterxml.jackson.annotation.JsonProperty;

import java.util.List;

public record CohereRerankRequest(
String model,
String query,
List<String> documents,
@JsonProperty("top_n")
Integer topN
) {}
16 changes: 16 additions & 0 deletions src/main/java/org/sopt/kareer/global/external/cohere/dto/response/CohereRerankResponse.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
package org.sopt.kareer.global.external.cohere.dto.response;

import com.fasterxml.jackson.annotation.JsonProperty;

import java.util.List;

public record CohereRerankResponse(
String id,
List<Result> results
) {
public record Result(
Integer index,
@JsonProperty("relevance_score")
Double relevanceScore
) {}
}
15 changes: 15 additions & 0 deletions src/main/java/org/sopt/kareer/global/external/cohere/properties/CohereProperties.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
package org.sopt.kareer.global.external.cohere.properties;

import org.springframework.boot.context.properties.ConfigurationProperties;

@ConfigurationProperties(prefix = "cohere")
public record CohereProperties(
String apiKey,
String baseUrl,
Rerank rerank
) {
public record Rerank(
String model,
int topN
) {}
}
74 changes: 74 additions & 0 deletions src/main/java/org/sopt/kareer/global/external/cohere/service/CohereRerankClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
package org.sopt.kareer.global.external.cohere.service;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.sopt.kareer.global.external.cohere.dto.request.CohereRerankRequest;
import org.sopt.kareer.global.external.cohere.dto.response.CohereRerankResponse;
import org.sopt.kareer.global.external.cohere.properties.CohereProperties;
import org.springframework.ai.document.Document;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestClient;

import java.util.ArrayList;
import java.util.List;

@Slf4j
@Component
@RequiredArgsConstructor
public class CohereRerankClient {

private final CohereProperties cohereProperties;

private RestClient restClient() {
return RestClient.builder()
.baseUrl(cohereProperties.baseUrl())
.defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + cohereProperties.apiKey())
.defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
.build();
}

public List<Document> rerank(String query, List<Document> documents, Integer topN) {
if (documents == null || documents.isEmpty()) {
return List.of();
}

List<String> serializedDocs = documents.stream()
.map(Document::getText)
.toList();

CohereRerankRequest request = new CohereRerankRequest(
cohereProperties.rerank().model(),
query,
serializedDocs,
topN != null
? Math.min(topN, documents.size())
: Math.min(cohereProperties.rerank().topN(), documents.size())
);

try {
CohereRerankResponse response = restClient()
.post()
.uri("/v2/rerank")
.body(request)
.retrieve()
.body(CohereRerankResponse.class);

if (response == null || response.results() == null || response.results().isEmpty()) {
log.warn("Cohere rerank response empty. query={}", query);
return documents;
}
Comment on lines +58 to +61
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

응답이 비어있을 때의 fallback 동작 확인 필요

응답이 null이거나 비어있을 때 원본 documents(candidatePoolTopK 크기)를 그대로 반환합니다. 이 경우에도 호출부에서 기대하는 topN 크기와 불일치가 발생합니다. PolicyDocumentRetriever에서 언급한 것처럼 일관된 크기 처리가 필요합니다.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In
`@src/main/java/org/sopt/kareer/global/external/cohere/service/CohereRerankClient.java`
around lines 58 - 61, The current fallback in CohereRerankClient returns the
original documents when response is empty which can mismatch the expected topN
size; update the fallback to return a list whose size matches the requested topN
(or candidatePoolTopK if topN not passed) by trimming documents if longer or
padding/duplicating with the best-available items if shorter, and ensure this
logic is used in the method that processes response and returns 'documents'
(refer to variables/methods: response, documents, CohereRerankClient,
PolicyDocumentRetriever, candidatePoolTopK, topN) so callers always receive a
consistently-sized result.


List<Document> reranked = new ArrayList<>();
for (CohereRerankResponse.Result result : response.results()) {
reranked.add(documents.get(result.index()));
}
return reranked;
Comment on lines +63 to +67
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Cohere 응답의 index 값 유효성 검증 필요

Cohere API 응답의 result.index()가 입력 documents 리스트의 범위를 벗어날 경우 IndexOutOfBoundsException이 발생할 수 있습니다. 외부 API 응답이므로 방어적 검증을 추가하는 것이 안전합니다.

🛡️ 인덱스 유효성 검증 추가 
             List<Document> reranked = new ArrayList<>();
             for (CohereRerankResponse.Result result : response.results()) {
+                if (result.index() < 0 || result.index() >= documents.size()) {
+                    log.warn("Cohere returned invalid index: {}. Skipping.", result.index());
+                    continue;
+                }
                 reranked.add(documents.get(result.index()));
             }
             return reranked;
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In
`@src/main/java/org/sopt/kareer/global/external/cohere/service/CohereRerankClient.java`
around lines 63 - 67, The loop in CohereRerankClient that builds 'reranked' uses
result.index() from response.results() without bounds checks, which can throw
IndexOutOfBoundsException if the Cohere response is invalid; update the code in
CohereRerankClient (where response.results() is iterated and
documents.get(result.index()) is called) to validate that result.index() is
within 0..documents.size()-1 before accessing documents, and on invalid indexes
either skip the entry and log a warning including the offending index and the
documents size or throw a controlled custom exception so the caller can handle
it.


} catch (Exception e) {
log.error("Cohere rerank failed. fallback to original order. query={}", query, e);
return documents;
}
}
}
7 changes: 7 additions & 0 deletions src/main/resources/application.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,13 @@ management:
enabled: true
health:
show-details: always

cohere:
api-key: ${COHERE_API_KEY}
base-url: https://api.cohere.com
rerank:
model: rerank-multilingual-v3.0
top-n: 10
---
spring:
config:
Expand Down
Loading