Spring AI: Кращі результати RAG завдяки гібридному пошуку з використанням Lucene та RRF
Об’єднайте векторний пошук із Lucene BM25 у ретрівері Spring AI та поєднайте рейтинги за допомогою Reciprocal Rank Fusion, щоб RAG знаходив як семантично схожі, так і точні терміни.
Джерела: Spring AI Recipe: Better RAG Results with Hybrid Search
Векторний пошук добре розуміє значення запиту. Лексичний пошук ефективний у пошуку слів, які були фактично введені. Якщо поєднати їх у гібридний пошук, то застосунок типу RAG отримує переваги обох підходів.
Пошук за схожістю за допомогою ембеддингів добре підходить для знаходження документів, семантично близьких до запиту, але може бути менш надійним, коли запит містить важливу специфічну термінологію. BM25 (Best Match 25) — це алгоритм лексичного ранжування, який робить протилежне: він знаходить відповідні терміни, враховуючи частоту їхнього зустрічання, рідкісність та довжину документа. Їхнє сумісне використання допомагає компенсувати слабкі сторони кожного з підходів.
Тут пропонується гібридний DocumentRetriever для Spring AI, який здійснює запити до сховища векторів та індексу BM25 та об’єднує результати в один відсортований список. BM25 походить з Apache Lucene — зрілої та швидкої бібліотеки пошуку, яка вже довгий час є основою пошуку в Java. Перша зміна стосується процесу надходження даних, щоб кожен фрагмент записувався як у сховище векторів, так і в Lucene.
Запис фрагментів у сховище векторів та в Lucene
Почніть з додавання Lucene до процесу компіляції:
implementation 'org.apache.lucene:lucene-core:10.5.1'
Далі йде компонент, який приймає об’єкти Document від Spring AI та індексує їх у Lucene. Це LuceneDocumentWriter:
public class LuceneDocumentWriter implements AutoCloseable {
private static final String ID_FIELD = "id";
private static final String CONTENT_FIELD = "content";
private static final String METADATA_FIELD = "metadata";
private final Analyzer analyzer;
private final Directory directory;
private final IndexWriter writer;
private final ObjectMapper objectMapper;
public LuceneDocumentWriter(Path indexPath) throws IOException {
this.analyzer = new StandardAnalyzer();
this.directory = FSDirectory.open(indexPath);
var config = new IndexWriterConfig(analyzer);
this.writer = new IndexWriter(directory, config);
this.objectMapper = new ObjectMapper();
}
public void add(List<Document> documents) throws IOException {
for (Document document : documents) {
var luceneDocument = new org.apache.lucene.document.Document();
luceneDocument.add(
new StringField(ID_FIELD, document.getId(), Field.Store.YES));
luceneDocument.add(
new TextField(CONTENT_FIELD, document.getText(), Field.Store.YES));
luceneDocument.add(
new StoredField(METADATA_FIELD,
objectMapper.writeValueAsString(document.getMetadata())));
writer.addDocument(luceneDocument);
}
writer.commit();
}
@Override
public void close() throws IOException {
writer.close();
directory.close();
analyzer.close();
}
}
У поточних версіях Lucene за замовчуванням використовується алгоритм BM25 для оцінки схожості, тому для його увімкнення нічого особливого не потрібно. Якщо ви хочете прямо вказати цей вибір, додайте цей рядок після створення IndexWriterConfig:
config.setSimilarity(new BM25Similarity());
Код імпорту використовує функцію add() для додавання документів до Lucene. Його параметр List<Document> містить документи Spring AI (org.springframework.ai.document.Document), тоді як Lucene індексує свій власний тип документів (org.apache.lucene.document.Document). Тож основна частина методу копіює ID, текст та метадані кожного документа Spring AI у поля нового документа Lucene, який потім додається до індексу.
Коли об’єкт запису готовий, оголосіть його як bean, вставте його до ApplicationRunner, який виконує процес імпорту, та збережіть кожну партію даних у обох сховищах:
@Bean
LuceneDocumentWriter luceneDocumentWriter(
@Value("${lucene.index.path}") Path luceneIndexPath) throws IOException {
return new LuceneDocumentWriter(luceneIndexPath);
}
@Bean
ApplicationRunner loadSpringBootDocumentation(
VectorStore vectorStore, LuceneDocumentWriter luceneDocumentWriter) {
return args -> {
// ...code snipped out for brevity
vectorStore.add(chunks);
luceneDocumentWriter.add(chunks);
// ...code snipped out for brevity
};
}
Шлях індексу вказується за допомогою @Value, тож встановіть його у файлі application.properties:
lucene.index.path=/tmp/lucene/boot-docs
Будь-який шлях підійде, але запам’ятайте його, оскільки основний додаток RAG має відкривати той самий індекс Lucene.
Якщо ви вже працювали з попереднім кроком HyDE у тому ж проекті, очистіть наявні дані векторного сховища перед тим, як знову запустити процес імпорту, інакше будуть зберігатися дублікати фрагментів. Після завершення імпорту у вас знову має бути 300 фрагментів у 89 документах, які тепер будуть доступні як через векторне сховище, так і через Lucene.
Запити до обох джерел
Основному додатку RAG потрібні Lucene Core та парсер запитів Lucene:
implementation 'org.apache.lucene:lucene-core:10.5.1'
implementation 'org.apache.lucene:lucene-queryparser:10.5.1'
Тепер створіть компонент, який шукає в індексі Lucene та повертає знахідки у вигляді об’єктів Document від Spring AI:
public class LuceneSearch implements AutoCloseable {
private static final String ID_FIELD = "id";
private static final String CONTENT_FIELD = "content";
private static final String METADATA_FIELD = "metadata";
private final Analyzer analyzer;
private final Directory directory;
private final DirectoryReader reader;
private final IndexSearcher searcher;
private final ObjectMapper objectMapper;
public LuceneSearch(Path indexPath) throws IOException {
this.analyzer = new StandardAnalyzer();
this.directory = FSDirectory.open(indexPath);
this.reader = DirectoryReader.open(directory);
this.searcher = new IndexSearcher(reader);
this.objectMapper = new ObjectMapper();
}
public List<Document> search(String queryText, int topK) {
try {
var parser = new QueryParser(CONTENT_FIELD, analyzer);
var query = parser.parse(
QueryParser.escape(queryText));
var hits = searcher.search(query, topK);
var storedFields = searcher.storedFields();
var results = new ArrayList<Document>();
for (var hit : hits.scoreDocs) {
var luceneDocument =
storedFields.document(hit.doc);
var metadata = objectMapper.readValue(
luceneDocument.get(METADATA_FIELD),
new TypeReference<Map<String, Object>>() {
});
results.add(new Document(
luceneDocument.get(ID_FIELD),
luceneDocument.get(CONTENT_FIELD),
metadata));
}
return results;
}
catch (IOException | ParseException e) {
throw new IllegalStateException("Failed to search Lucene index", e);
}
}
@Override
public void close() throws IOException {
reader.close();
directory.close();
analyzer.close();
}
}
Не заглиблюючись у API пошуку Lucene, функція search() бере текст запиту та запитує IndexSearcher про документи з найкращою відповідністю, причому параметр topK обмежує кількість отриманих результатів. Як і під час індексування, за замовчуванням використовується алгоритм BM25 для оцінки схожості, і якщо ви хочете це побачити, можете явно налаштувати пошуковик:
searcher.setSimilarity(new BM25Similarity());
Решта функції search() перетворює кожен документ Lucene Document назад у документ Spring AI Document, щоб він міг брати участь у процесі RAG.
Зараз існує два незалежні способи отримання документів. Пошук за векторною схожістю через VectorStore ефективно знаходить семантично пов’язаний контент. Пошук BM25 через LuceneSearch ефективно знаходить важливі терміни у запиті. Кожен з цих методів може сам по собі дати корисні результати, але їх поєднання забезпечує як семантичну, так і лексичну точність. HybridDocumentRetriever здійснює саме це поєднання.
Гібридний збирач документів
Ось його суть:
public class HybridDocumentRetriever implements DocumentRetriever {
private static final Logger logger =
LoggerFactory.getLogger(HybridDocumentRetriever.class);
private final VectorStore vectorStore;
private final LuceneSearch luceneSearch;
public HybridDocumentRetriever(
VectorStore vectorStore, LuceneSearch luceneSearch) {
this.vectorStore = vectorStore;
this.luceneSearch = luceneSearch;
}
@Override
public List<Document> retrieve(Query query) {
var vectorResults = vectorStore.similaritySearch(
SearchRequest.builder()
.query(query.text())
.topK(10)
.build());
var bm25Results = luceneSearch.search(query.text(), 10);
var fusedResults =
reciprocalRankFusion(vectorResults, bm25Results, 10);
logResults("VECTOR", vectorResults);
logResults("BM25", bm25Results);
logResults("FUSED", fusedResults);
return fusedResults;
}
// ...
}
HybridDocumentRetriever реалізує клас DocumentRetriever від Spring AI — компонент, який завантажує документи для RetrievalAugmentationAdvisor. Раніше використовувався VectorStoreDocumentRetriever, який запозичує дані лише з векторного сховища. Тут його замінено на власну реалізацію.
Його метод retrieve() виконує два пошуки для одного запиту. Спочатку він викликає similaritySearch() для 10 найкращих результатів за критерієм схожості. Потім запитує Lucene про 10 найкращих результатів за алгоритмом BM25. Результатом є два ранжовані списки: деякі документи з’являються в обох, а інші — лише в одному. Якби списки було просто об’єднано, інформація про ранжування була б втрачена, тому при злитті необхідно дотримуватися позиції кожного документа в кожному списку. Саме це робить метод Reciprocal Rank Fusion.
Злиття за допомогою Reciprocal Rank Fusion
Метод фузії виглядає так:
private List<Document> reciprocalRankFusion(
List<Document> vectorResults,
List<Document> bm25Results,
int limit) {
Map<String, Double> scores = new HashMap<>();
Map<String, Document> documents = new HashMap<>();
addRrfScores(vectorResults, scores, documents);
addRrfScores(bm25Results, scores, documents);
return scores.entrySet().stream()
.sorted(Map.Entry.<String, Double>comparingByValue().reversed())
.limit(limit)
.map(entry -> documents.get(entry.getKey()))
.toList();
}
private void addRrfScores(
List<Document> results,
Map<String, Double> scores,
Map<String, Document> documents) {
final int k = 60;
for (int i = 0; i < results.size(); i++) {
var document = results.get(i);
var rank = i + 1;
documents.putIfAbsent(document.getId(), document);
scores.merge(
document.getId(),
1.0 / (k + rank),
Double::sum);
}
}
Функція взаємного об’єднання рангів (RRF) перетворює кілька наборів результатів із рангуванням на одне єдине рангування. Замість спроби порівнювати первинні оцінки з двох дуже різних алгоритмів пошуку, вона оцінює документи за їхнім місцем у кожному списку. Для кожного списку документ отримує:
1 / (k + rank)
Тут k — це константа, яка зазвичай дорівнює 60 для RRF, тоді як rank позначає місце документа в цьому списку. Документ, який є у обох списках, отримує суму обох оцінок. Результат є інтуїтивно зрозумілим: документи, які посідають високі позиції в одному з пошуків, мають хороші результати, а ті, що посідають високі позиції в обох, зазвичай піднімаються на вершину, що саме і потрібно під час поєднання семантичного та лексичного пошуку.
У кінці функції retrieve() фіксуються всі три рейтинги, щоб ефект був видимим:
private void logResults(String label, List<Document> documents) {
logger.debug("{}", label);
for (int i = 0; i < documents.size(); i++) {
var document = documents.get(i);
var text = document.getText()
.replaceAll("\\s+", " ")
.trim();
var snippet = text.substring(
0, Math.min(text.length(), 80));
logger.debug(
"{}. {} - {}",
i + 1,
document.getId(),
snippet);
}
}
Цей логування не є необхідним для самого гібридного пошуку, але під час експериментування воно дозволяє порівнювати векторні, BM25 та злиті способи сортування одночасно.
Підключення гібридного механізму пошуку
Зареєструйте LuceneSearch та HybridDocumentRetriever як beanи:
@Bean
LuceneSearch luceneSearch(
@Value("${lucene.index.path}") Path luceneIndexPath) throws IOException {
return new LuceneSearch(luceneIndexPath);
}
@Bean
HybridDocumentRetriever hybridDocumentRetriever(
VectorStore vectorStore, LuceneSearch luceneSearch) {
return new HybridDocumentRetriever(vectorStore, luceneSearch);
}
Як і під час завантаження даних, LuceneSearch потребує адреси індексу Lucene:
lucene.index.path=/tmp/lucene/boot-docs
Переконайтеся, що цей шлях збігається з тим, який використовувала програма імпорту. Потім змініть ChatClientBuilderCustomizer, який налаштовує RetrievalAugmentationAdvisor, щоб він використовував HybridDocumentRetriever замість VectorStoreDocumentRetriever:
@Bean
ChatClientBuilderCustomizer addRagAdvisor(
QueryTransformer[] queryTransformers,
HybridDocumentRetriever hybridDocumentRetriever) {
return builder -> {
var ragAdvisor = RetrievalAugmentationAdvisor.builder()
.queryTransformers(queryTransformers)
.documentRetriever(hybridDocumentRetriever)
.documentJoiner(documentsForQuery ->
documentsForQuery.values().stream()
.flatMap(List::stream)
.flatMap(List::stream)
.toList())
.build();
builder.defaultAdvisors(ragAdvisor);
};
}
Окрім передачі засобу отримання документів через documentRetriever(), налаштовувач також передає documentJoiner() об’єкт DocumentJoiner, створений за допомогою лямбда-функції. Завданням цього засобу є об’єднання кількох списків отриманих документів у одну колекцію. Стандартний ConcatenationDocumentJoiner знову сортує документи за їхнім рейтингом, що скасовує порядок, який щойно було створено за допомогою RRF. Тому код надає засіб об’єднання, який фактично нічого не робить та залишає об’єднаний рейтинг без змін.
Ви можете вважати, що компонент joiner — це природне місце для RRF, і концептуально це так. Такий підхід також передбачає розділення векторного пошуку та пошуку за допомогою Lucene на два окремі механізми отримання даних. Проблема полягає у тому, що конструктор advisor приймає лише один DocumentRetriever, тому ці два механізми не можна налаштувати окремо та передати обидва набори результатів до joiner. Простішим рішенням є виконання обох пошуків та їх об’єднання всередині HybridDocumentRetriever, а потім використання joiner для збереження порядку об’єднаних результатів. Якщо Spring AI у майбутньому підтримає кілька механізмів отримання даних, цю конструкцію варто переглянути.
Це все щодо підключень. RetrievalAugmentationAdvisor не повинен знати про наявність двох стратегій пошуку. Для нього HybridDocumentRetriever — це просто компонент, який надає документи, додані до контексту запиту.
Бачення різниці
Увімкніть журналування для відладки для HybridDocumentRetriever, щоб ви могли порівняти три рейтинги:
logging.level.com.example.essentialrag.HybridDocumentRetriever=debug
Запустіть додаток та поставте запитання щодо використання @ConfigurationProperties для мапування зовнішньої конфігурації на об’єкт Java з перевіркою. Результати дебагування покажуть три набори даних із позначками VECTOR, BM25 та FUSED. Ранжування за векторним підходом починається ось так:
VECTOR
1. d91bc5c8-... - If the annotation does not provide any prefix...
2. 049f6484-... - For YAML, you can use the following syntax...
3. 152dde59-... - In this setup, the presence of a single parameterized constructor...
...
BM25 створює інший порядок:
BM25
1. 4c7128dc-... - There are spring-boot-starter-data-couchbase...
2. a3dbefae-... - fun main(args: Array<String>)...
3. 1ab4c223-... - Configuring System Environment Properties...
4. d91bc5c8-... - If the annotation does not provide any prefix...
...
А результат RRF поєднує їх:
FUSED
1. d91bc5c8-... - If the annotation does not provide any prefix...
2. 1ab4c223-... - Configuring System Environment Properties...
3. 049f6484-... - For YAML, you can use the following syntax...
4. 152dde59-... - In this setup, the presence of a single parameterized constructor...
5. 091acd83-... - The default unit is bytes...
...
Цей контраст показує, чому гібридний пошук є ефективним. Векторний пошук добре справляється з концептуальними аспектами запиту: зовнішня конфігурація, прив’язка, об’єкти Java та перевірка. BM25 додає лексичне порівняння та надає більше ваги унікальним термінам, таким як ConfigurationProperties. Потім RRF об’єднує ці два підходи. Документи, які отримують високі оцінки в обох методах, піднімаються вгору, тоді як сильні результати, знайдені лише однією стратегією, все одно можуть потрапити до кінцевого списку.
Тому цей потік не мусить вибирати між розумінням того, що має на увазі користувач, та порівнянням слів, які він використав. За допомогою гібридного пошуку можна робити і те, і інше.
Пов’язана література
- Гібридний пошук у RAG: практичне застосування BM25, векторів та методу взаємного рангування — Чому RAG, що використовує лише вектори, втрачає точні ідентифікатори, як метод RRF поєднує BM25 з ембеддингами, де гібридний пошук дає гірші результати та чому фільтри метаданих мають використовуватися перед об’єднанням даних.
- Міркування над ієрархічним деревом документів замість пошуку в векторах — PageIndex створює ієрархічний індекс, який дозволяє ШІ навігувати в ньому, уникаючи розділення документів на частини та використання ембеддингів. Як це працює, результати тестування на FinanceBench та коли варто його використовувати.