Рецепт Spring AI: фильтрация результатов RAG с использованием метаданных по названию игры

Векторное сходство позволяет находить релевантный текст, но одного лишь уровня релевантности часто бывает недостаточно. Магазин может хранить тысячи фрагментов информации по разным темам: продукты, клиенты или, в данном случае, настольные игры. Когда несколько фрагментов одинаково хорошо соответствуют вопросу, как выбрать именно те из них, которые действительно релевантны? Необходимо добавить метаданные к каждому документу и использовать их для фильтрации во время поиска. Таким образом, модель будет рассматривать только ту часть данных, которую вы выбрали.

Настройка: три небольшие игры

Начальной точкой служит приложение Spring AI, реализующее технологию RAG без отдельной векторной базы данных; оно было модифицировано для обработки нескольких файлов. Демонстрационные файлы представляют собой краткие текстовые описания трех вымышленных игр: Camp Bowwow, Frog Panic и Taco Truck. Каждый из них находится в папке src/main/resources/documents/. Сначала рассмотрим файл Camp Bowwow в формате camp-bowwow.txt:

On each player's turn, they place a dog, food bowl, or toy into
one of the three play yards. The game ends when a player has
placed their last dog into a yard. Then each player who hasn't
played in that round gets to place one more item into a yard
until the last player so that all players have an equal number
of turns.

Текст «Frog Panic» находится в файле frog-panic.txt:

On their turn, a player starts by moving any flies on the board
to the same space as an adjacent frog (effectively the frog has
eaten that fly). Then the player can either place 1 fly adjacent
to a frog on the board as a distraction so that the frog won't
jump or they may place a frog of their player color onto any
empty space. If the frog is adjacent to another frog (even if it
is their own), the adjacent frog will jump two spaces in any
direction away from the frog that was just placed. The direction
of the jump is decided by the current player. If the frog jumps
off the board, then that frog is out of play.

The game ends when all players have placed all of their frogs
(not counting frogs that have jumped off the board). The player
with the most frogs remaining on the board wins.

А также информация о фургоне с тако в taco-truck.txt:

When it is a player's turn, they will draw a card from the
ingredients deck and add it to their hand. Then they may play
any ingredient card from their hand to add that ingredient to
one of up to three tacos in progress. Or, if they have a crunch
card in their hand, then they may play it to remove one
ingredient from any one opponent's taco. The first player to
completely build 3 tacos wins.

Теперь конфигурация загрузки должна загружать все три файла:

@Configuration
public class RagIngestionConfig {

  private static final Logger logger =
      LoggerFactory.getLogger(RagIngestionConfig.class);

  @Value("${rag.documents}")
  Resource[] documentResources;

  @Bean
  @Order(-1)
  ApplicationRunner load(VectorStore vectorStore) {
    return args -> {
      for (Resource documentResource : documentResources) {
        var filename = documentResource.getFilename();
        logger.info("Loading document from {}.", filename);
        var reader = new TikaDocumentReader(documentResource);
        var splitter = TokenTextSplitter.builder().build();
        vectorStore.accept(
            splitter.apply(
                reader.get()));
      }
      logger.info("Document loading complete.");
    };
  }
}

По сравнению с предыдущей версией различия незначительны. Вместо одного пути типа String свойство rag.documents передаёт массив объектов Resource. Метод load() проходит по этому массиву, разделяет каждый объект на фрагменты и записывает их в векторный хранилище.

Запустите приложение сейчас и задайте вопрос вроде «Что я могу сделать на своей очереди?», и появится задача. Все три кратких описания относятся к ходу игрока, поэтому фрагменты из любого из них могут показаться релевантными для векторного хранилища.

Что происходит без фильтра

Ответ может описывать одну игру или, что более вероятно, смесь правил из нескольких игр. Один из примеров выглядел так:

How can I help?

> What can I do on my turn?

 - Draw a card from the ingredients deck and add it to your hand.
 - Then you may either:
    - Play any ingredient card from your hand to add that ingredient
      to one of up to three tacos in progress, or
    - Play a crunch card (if you have one) to remove one ingredient
      from any one opponent's taco.
 - On your turn you place a dog, a food bowl, or a toy into one of
   the three play yards.
 - The game ends when a player has placed their last dog; then players
   who haven't yet played in that round get one more placement so all
   players have equal turns.

Каждый факт связан с вопросом, однако смешивание игр делает ответ бесполезным. Представьте то же самое, но с сотнями игр в магазине. Именно здесь и пригождается фильтрация метаданных.

Маркировка документов заголовком

Для демонстрации каждому документу присваивается заголовок игры в качестве метаданных, что позволяет сузить поиск до одной конкретной игры по этому значению. Код обработки данных незначительно изменяется, чтобы заголовок хранился вместе с фрагментами данных:

@Bean
@Order(-1)
ApplicationRunner load(VectorStore vectorStore) {
  return args -> {
    for (Resource documentResource : documentResources) {
      var filename = documentResource.getFilename();

  logger.info("Loading document from {}.", filename);
      var reader = new TikaDocumentReader(documentResource);
      var splitter = TokenTextSplitter.builder().build();

      var titleTag =
          filename.substring(0, filename.lastIndexOf('.'));

      vectorStore.accept(
          splitter.apply(
              reader.get().stream()
                  .peek(document ->
                      document.getMetadata()
                          .put("title", titleTag))
                  .toList()));
    }
    logger.info("Document loading complete.");
  };
}

Прежде чем документ попадает в инструмент разделения, код формирует заголовок на основе его имени файла и сохраняет его в метаданных под ключом title. Для файла camp-bowwow.txt это означает:

title = camp-bowwow

При разделении данные метаданных также передаются, поэтому каждый полученный фрагмент сохраняется в хранилище вместе с ними. Эмбеддинги по-прежнему обеспечивают семантическое сходство, в то время как метаданные добавляют вторую ось для узкой настройки поиска. Вместо того чтобы запрашивать только фрагменты, похожие на вопрос, можно запрашивать фрагменты, похожие на вопрос и принадлежащие к определенной игре.

Определение той игры, о которой идет речь

Далее приложение должно определить, к какой игре относится запрос. Есть нюанс: хранимые значения берутся из имён файлов вроде camp-bowwow.txt, что отличается от того, как обычно пишут название игры. Языковая модель может преодолеть эту проблему, переведя название из вопроса в одно из хранимых значений:

@Service
public class TitleHelper {

  private final ChatClient chatClient;

  public TitleHelper(ChatModel chatModel) {
    this.chatClient = ChatClient.builder(chatModel).build();
  }

  public String determineGameTitle(String question) {
    var title = chatClient.prompt()
        .user(userSpec -> userSpec
            .text("""
                Your job is to try to determine the title of a game from
                the question asked.
                The game choices are:
                - camp-bowwow
                - frog-panic
                - taco-truck
                - unknown
                If the game's title isn't explicitly mentioned in the
                question, or you don't recognize the game's title, then
                say "unknown".
                The question is:
                {question}
                """)
            .param("question", question))
        .call()
        .content();
    return title.equals("unknown") ? null : title;
  }
}

Обратите внимание, что TitleHelper не переиспользует основной объект ChatClient приложения. Он намеренно берет ChatModel и создает собственный клиент. Для простой задачи классификации не требуется память, инструменты RAG или другие функции, настроенные в основном клиенте; достаточен обычный ChatClient.

Запрос направляется в модель с заранее определенным списком допустимых названий, соответствующих значениям метаданных. Если модель не может распознать ни одну из этих игр, она возвращает unknown, что функция determineGameTitle() преобразует в null.

Применение фильтра к процессу поиска

После готовности вспомогательного класса в цикле чата в объекте ApplicationRunner требуется лишь несколько изменений:

@Bean
ApplicationRunner go(
    ChatClient chatClient,
    TitleHelper titleHelper) {

  return args -> {
    System.out.println("How can I help?\n");

    try (Scanner scanner = new Scanner(System.in)) {
      while (true) {
        System.out.print("> ");
        if (!scanner.hasNextLine()) break;
        var input = scanner.nextLine();
        if (input.isBlank()) continue;

        var requestSpec = chatClient.prompt(input);

        var gameTitle =
            titleHelper.determineGameTitle(input);

        var answer = requestSpec
            .advisors(spec -> {
              spec.param(
                  ChatMemory.CONVERSATION_ID,
                  "DEMO");
              if (gameTitle != null) {
                spec.param(
                    QuestionAnswerAdvisor.FILTER_EXPRESSION,
                    String.format(
                        "title == '%s'",
                        gameTitle));
              }
            })
            .call()
            .content();

        System.out.println("\n - " + answer);
      }
    }
  };
}

Изменяются три вещи:

• TitleHelper вставляется в механизм выполнения.
• Он используется для определения того, к какой игре относится вопрос.
• При нахождении названия фильтрационное выражение передается в QuestionAnswerAdvisor, поэтому учитываются только документы, метаданные title которых совпадают.

Само выражение довольно краткое:

title == 'camp-bowwow'

Важно то, где происходит фильтрация. Речь не идет о запросе к ЯИИ игнорировать несвязанные документы после их загрузки. Фильтр ограничивает сам поиск по векторам, поэтому результаты возвращают только фрагменты выбранной игры, а модель с самого начала получает сфокусированный контекст.

Если игра не обнаружена, фильтр просто не применяется, и поиск охватывает всё. Такой вариант резерва стоит сохранить. Неверный или отсутствующий прогноз не должен нарушать работу процесса, а поведение без фильтрации остается прежним.

Попробовать сами

Запустите приложение и спросите, что делает игрок на своей ходу, на этот раз укажав название игры:

How can I help?

> I'm playing Camp Bowwow. What do I do when it's my turn?

 - On your turn you place one item - either a dog, a food bowl,
   or a toy - into one of the three play yards.

(For reference: the game ends when a player has placed their last
dog; then any players who haven't yet played in that round each get
one more item turn in order until the last player, so all players
have equal turns.)

Поиск ограничен блоками с title, равным camp-bowwow, поэтому ответ не может отклоняться от правил Taco Truck или Frog Panic. Попробуйте другой вариант:

> I'm playing Frog Panic. What do I do when it's my turn?

- On your turn:
1. First, move any flies on the board onto the same space as an
   adjacent frog (the frog eats that fly).
2. Then choose one action:
   - Place 1 fly adjacent to a frog on the board as a distraction
     (so that frog won't jump), or
   - Place a frog of your color onto any empty space.
If the frog you place is adjacent to another frog (even your own),
that adjacent frog jumps two spaces away in any direction you choose.
If a frog jumps off the board it is out of play.

(For reference: the game ends when all players have placed all their
frogs; the player with the most frogs remaining on the board wins.)

Тот же вопрос, другой фильтр, другой набор документов, лежащих в основе ответа.

За пределами сходства

Векторный поиск отлично справляется с определением тех документов, которые наиболее похожи на этот запрос. Однако он сам по себе не может учесть всю информацию о месте поиска. Метаданные — это способ внедрения таких дополнительных знаний в процесс поиска.

Здесь фильтр представляет собой простое название игры, но в реальных приложениях документы могут маркироваться информацией о клиенте, продукте, типе документа, категории, языке, дате, арендаторе, уровне безопасности или любым другим атрибутом, который позволяет сузить круг поиска. Сочетание семантического сходства с фильтрацией по метаданным делает поиск более точным и предоставляет модели более компактный и релевантный контекст.

Несколько практических рекомендаций перед использованием этого шаблона в других случаях. Сохраняйте консистентность значений метаданных и берите их из небольшого известного набора, поскольку шаг классификации работает эффективно только тогда, когда у него есть возможность выбирать из фиксированного списка. Заранее определите, что произойдет, если вопрос не совпадает ни с чем. Кроме того, считайте метаданные, связанные с безопасностью, такие как информация о арендаторе или диапазон разрешений, строгим требованием, применяемым вашим кодом, а не чем-то, что модель должна соблюдать.