Рецепт 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 у 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.)

Те саме запитання, інший фільтр, інший набір документів для отримання відповіді.

Понад просту схожість

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

Тут фільтром є проста назва ігри, але у реальних застосуваннях документи можна позначати такими атрибутами, як клієнт, продукт, тип документа, категорія, мова, дата, користувач, рівень безпеки чи будь-які інші характеристики, які допомагають звузити пошук. Поєднання семантичної схожості з фільтрацією за метаданами робить пошук більш точним та забезпечує моделі скорочений, але більш релевантний контекст.

Кілька практичних порад перед використанням цього шаблону в інших ситуаціях. Зберігайте консистентні значення метаданих, взяті з невеликого відомого набору, адже крок класифікації працює ефективно лише тоді, коли може обирати з фіксованого списку. Заздалегідь вирішіть, що має відбуватися, коли запит не знаходить відповідностей. Також вважайте метадані, пов’язані з безпекою, такі як ідентифікатор користувача чи обсяг дозволів, обов’язковою вимогою, яку має дотримуватися ваш код, а не чимось, що має бути враховане моделлю.