Skip to content

optimize(rag): 优化意图候选裁剪逻辑 #104

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

Open
brainbrain29 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

optimize(rag): 优化意图候选裁剪逻辑 #104

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
154 changes: 78 additions & 76 deletions bootstrap/src/main/java/com/nageoffer/ai/ragent/rag/core/intent/IntentResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import java.util.concurrent.CompletableFuture;
import java.util.HashMap;
import java.util.concurrent.Executor;
import java.util.PriorityQueue;

import static com.nageoffer.ai.ragent.rag.constant.RAGConstant.INTENT_MIN_SCORE;
import static com.nageoffer.ai.ragent.rag.constant.RAGConstant.MAX_INTENT_COUNT;
Expand Down Expand Up @@ -63,8 +64,7 @@ public List<SubQuestionIntent> resolve(RewriteResult rewriteResult) {
return new SubQuestionIntent(q, List.of());
}
},
intentClassifyExecutor
))
intentClassifyExecutor))
.toList();
List<SubQuestionIntent> subIntents = tasks.stream()
.map(CompletableFuture::join)
Expand Down Expand Up @@ -97,109 +97,111 @@ private List<NodeScore> classifyIntents(String question) {
}

/**
* 限制总意图数量不超过 MAX_INTENT_COUNT
* <p>
* 策略：
* 1. 如果总数未超限，直接返回
* 2. 如果超限，每个子问题至少保留 1 个最高分意图
* 3. 剩余配额按分数从高到低分配给其他意图
* 限制总意图数量不超过 MAX_INTENT_COUNT。
*
* 每个非空子问题保留一个最高分意图；其余配额分配给全局分数最高的候选。
*/
private List<SubQuestionIntent> capTotalIntents(List<SubQuestionIntent> subIntents) {
int totalIntents = subIntents.stream()
.mapToInt(si -> si.nodeScores().size())
.sum();
List<RankedIntentCandidate> guaranteedIntents = new ArrayList<>();
PriorityQueue<RankedIntentCandidate> additionalIntents = new PriorityQueue<>(this::compareAdditionalPriority);

int totalIntents = 0;
int remaining = MAX_INTENT_COUNT;
int order = 0;

for (int i = 0; i < subIntents.size(); i++) {
RankedIntentCandidate topIntent = null;

for (NodeScore nodeScore : subIntents.get(i).nodeScores()) {
totalIntents++;

RankedIntentCandidate candidate = new RankedIntentCandidate(new IntentCandidate(i, nodeScore), order++);

if (topIntent == null || compareScore(candidate, topIntent) > 0) {
if (topIntent != null) {
retainAdditionalIntent(additionalIntents, topIntent, remaining);
}
topIntent = candidate;
} else {
retainAdditionalIntent(additionalIntents, candidate, remaining);
}
}

if (topIntent != null) {
guaranteedIntents.add(topIntent);
remaining = Math.max(0, MAX_INTENT_COUNT - guaranteedIntents.size());
trimAdditionalIntents(additionalIntents, remaining);
}
}

// 未超限，直接返回
if (totalIntents <= MAX_INTENT_COUNT) {
return subIntents;
}

// 步骤1：收集所有意图，按子问题索引分组
List<IntentCandidate> allCandidates = collectAllCandidates(subIntents);
List<IntentCandidate> guaranteed = guaranteedIntents.stream()
.map(RankedIntentCandidate::candidate)
.toList();

// 步骤2：每个子问题保留最高分意图
List<IntentCandidate> guaranteedIntents = selectTopIntentPerSubQuestion(allCandidates, subIntents.size());
List<IntentCandidate> additional = additionalIntents.stream()
.sorted((left, right) -> compareAdditionalPriority(right, left))
.map(RankedIntentCandidate::candidate)
.toList();

// 步骤3：计算剩余配额
int remaining = MAX_INTENT_COUNT - guaranteedIntents.size();
return rebuildSubIntents(subIntents, guaranteed, additional);
}

private void retainAdditionalIntent(PriorityQueue<RankedIntentCandidate> candidates,
RankedIntentCandidate candidate,
int limit) {
if (limit <= 0) {
return;
}

// 步骤4：从剩余候选中按分数选择
List<IntentCandidate> additionalIntents = selectAdditionalIntents(allCandidates, guaranteedIntents, remaining);
if (candidates.size() < limit) {
candidates.offer(candidate);
return;
}

// 步骤5：合并并重建结果
return rebuildSubIntents(subIntents, guaranteedIntents, additionalIntents);
if (compareAdditionalPriority(candidate, candidates.peek()) > 0) {
candidates.poll();
candidates.offer(candidate);
}
}

/**
* 收集所有意图候选，标记所属子问题索引
*/
private List<IntentCandidate> collectAllCandidates(List<SubQuestionIntent> subIntents) {
List<IntentCandidate> candidates = new ArrayList<>();
for (int i = 0; i < subIntents.size(); i++) {
List<NodeScore> nodeScores = subIntents.get(i).nodeScores();
if (CollUtil.isEmpty(nodeScores)) {
continue;
}
for (NodeScore ns : nodeScores) {
candidates.add(new IntentCandidate(i, ns));
}
private void trimAdditionalIntents(PriorityQueue<RankedIntentCandidate> candidates, int limit) {
while (candidates.size() > limit) {
candidates.poll();
}
// 按分数降序排序
candidates.sort((a, b) -> Double.compare(b.nodeScore().getScore(), a.nodeScore().getScore()));
return candidates;
}

/**
* 每个子问题选择最高分意图（保底策略）
* 分数高者优先；同分时保持原始遍历顺序优先。
*/
private List<IntentCandidate> selectTopIntentPerSubQuestion(List<IntentCandidate> allCandidates, int subQuestionCount) {
List<IntentCandidate> topIntents = new ArrayList<>();
boolean[] selected = new boolean[subQuestionCount];

for (IntentCandidate candidate : allCandidates) {
int index = candidate.subQuestionIndex();
if (!selected[index]) {
topIntents.add(candidate);
selected[index] = true;
}
// 所有子问题都有了保底意图，提前退出
if (topIntents.size() == subQuestionCount) {
break;
}
}
return topIntents;
private int compareScore(RankedIntentCandidate left, RankedIntentCandidate right) {
int scoreComparison = Double.compare(
left.candidate().nodeScore().getScore(),
right.candidate().nodeScore().getScore());
return scoreComparison != 0
? scoreComparison
: Integer.compare(right.order(), left.order());
}

/**
* 从剩余候选中选择额外意图
* 小顶堆比较器：堆顶始终是当前额外候选中最应被淘汰的条目。
*/
private List<IntentCandidate> selectAdditionalIntents(List<IntentCandidate> allCandidates,
List<IntentCandidate> guaranteedIntents,
int remaining) {
if (remaining <= 0) {
return List.of();
}
private int compareAdditionalPriority(RankedIntentCandidate left, RankedIntentCandidate right) {
return compareScore(left, right);
}

List<IntentCandidate> additional = new ArrayList<>();
for (IntentCandidate candidate : allCandidates) {
// 跳过已经被选为保底的意图
if (guaranteedIntents.contains(candidate)) {
continue;
}
additional.add(candidate);
if (additional.size() >= remaining) {
break;
}
}
return additional;
private record RankedIntentCandidate(IntentCandidate candidate, int order) {
}

/**
* 根据选中的意图重建 SubQuestionIntent 列表
*/
private List<SubQuestionIntent> rebuildSubIntents(List<SubQuestionIntent> originalSubIntents,
List<IntentCandidate> guaranteedIntents,
List<IntentCandidate> additionalIntents) {
List<IntentCandidate> guaranteedIntents,
List<IntentCandidate> additionalIntents) {
// 合并所有选中的意图
List<IntentCandidate> allSelected = new ArrayList<>(guaranteedIntents);
allSelected.addAll(additionalIntents);
Expand Down