Geração aumentada de recuperação (RAG, na sigla em inglês)

O Firebase Genkit fornece abstrações que ajudam a criar fluxos de geração de recuperação aumentada (RAG, na sigla em inglês), bem como plug-ins que fornecem integrações com ferramentas relacionadas.

O que é RAG?

A geração aumentada de recuperação é uma técnica usada para incorporar fontes externas de informações nas respostas de um LLM. É importante conseguir fazer isso porque, embora os LLMs sejam treinados normalmente com um amplo conjunto de materiais, o uso prático deles geralmente requer conhecimento de domínio específico (por exemplo, você pode querer usar um LLM para responder a perguntas de clientes sobre os produtos da sua empresa).

Uma solução é ajustar o modelo usando dados mais específicos. No entanto, isso pode ser caro, tanto em termos de custo de computação quanto em termos do esforço necessário para preparar os dados de treinamento adequados.

Por outro lado, o RAG funciona incorporando fontes de dados externas em um comando no momento em que ele é passado para o modelo. Por exemplo, imagine que a instrução "Qual é a relação de Bart com Lisa?" pode ser expandida ("aumentada") com a inclusão de algumas informações relevantes no início da mensagem: "Os filhos de Homer e Marge são chamados de Bart, Lisa e Maggie. Qual é a relação de Bart com Lisa?"

Essa abordagem tem várias vantagens:

  • Ela pode ser mais econômica porque não é necessário treinar novamente o modelo.
  • Você pode atualizar continuamente sua fonte de dados, e o LLM pode usar imediatamente as informações atualizadas.
  • Agora é possível citar referências nas respostas do LLM.

Por outro lado, o uso de RAG significa naturalmente prompts mais longos, e alguns serviços da API LLM cobram por cada token de entrada enviado. Em última análise, você precisa avaliar as compensações de custos para seus aplicativos.

A RAG é uma área muito ampla, e muitas técnicas diferentes são usadas para conseguir a melhor qualidade de RAG. O framework principal do Genkit oferece duas abstrações principais para ajudar você a fazer o RAG:

  • Indexadores: adicionam documentos a um "índice".
  • Embedders: transforma documentos em uma representação vetorial.
  • Recuperadores: recuperam documentos de um "índice", dada uma consulta.

Essas definições são amplas de propósito porque o Genkit não opina sobre o que é um "índice" ou como exatamente os documentos são recuperados dele. O Genkit fornece apenas um formato Document e todo o restante é definido pelo provedor de implementação do indexador ou do recuperador.

Indexadores

O índice é responsável por acompanhar os documentos para que você recupere rapidamente os documentos relevantes a partir de uma consulta específica. Isso geralmente é feito usando um banco de dados de vetores, que indexa seus documentos usando vetores multidimensionais chamados embeddings. Um embedding de texto (opaca) representa os conceitos expressos por um trecho de texto. Eles são gerados usando modelos de ML para fins especiais. Ao indexar o texto usando a incorporação, um banco de dados vetorial é capaz de agrupar textos conceitualmente relacionados e recuperar documentos relacionados a uma nova string de texto (a consulta).

Antes de recuperar documentos para fins de geração, eles precisam ser ingeridos no índice de documentos. Um fluxo de ingestão típico faz o seguinte:

  1. Divida documentos grandes em documentos menores para que apenas as partes relevantes sejam usadas para aumentar os comandos, "dividindo". Isso é necessário porque muitos LLMs têm uma janela de contexto limitada, o que torna inviável incluir documentos inteiros com um comando.

    O Genkit não fornece bibliotecas de agrupamento integradas. No entanto, há bibliotecas de código aberto disponíveis compatíveis com ele.

  2. Gerar embeddings para cada bloco. Dependendo do banco de dados que você estiver usando, é possível fazer isso explicitamente com um modelo de geração de embedding ou usar o gerador de embeddings fornecido pelo banco de dados.

  3. Adicione o bloco de texto e seu índice ao banco de dados.

É possível executar o fluxo de ingestão com pouca frequência ou apenas uma vez se você estiver trabalhando com uma fonte de dados estável. Por outro lado, se você estiver trabalhando com dados que mudam com frequência, poderá executar continuamente o fluxo de ingestão (por exemplo, em um gatilho do Cloud Firestore, sempre que um documento for atualizado).

Incorporadores

Um embedder é uma função que usa conteúdos (texto, imagens, áudio etc.) e cria um vetor numérico que codifica o significado semântico do conteúdo original. Como mencionado acima, os embedders são aproveitados como parte do processo de indexação, no entanto, também podem ser usados de maneira independente para criar embeddings sem um índice.

Recuperadores

Um retriever é um conceito que encapsula a lógica relacionada a qualquer tipo de recuperação de documentos. Os casos de recuperação mais conhecidos geralmente incluem a recuperação de armazenamentos de vetores. No entanto, no Genkit, um recuperador pode ser qualquer função que retorne dados.

Para criar um recuperador, você pode usar uma das implementações fornecidas ou criar uma própria.

Indexadores, recuperados e incorporados compatíveis

O Genkit oferece suporte ao indexador e recuperador por meio do sistema de plug-ins. Estes plug-ins são oficialmente compatíveis:

Além disso, o Genkit é compatível com os seguintes armazenamentos de vetores por meio de modelos de código predefinidos, que podem ser personalizados de acordo com a configuração e o esquema do banco de dados:

O suporte ao modelo de embedding é fornecido por meio dos seguintes plug-ins:

Plug-in Modelos
IA generativa do Google Incorporação de texto da lagartixa
Vertex AI do Google Incorporação de texto da lagartixa

Como definir um fluxo RAG

Os exemplos a seguir mostram como ingerir uma coleção de documentos PDF do cardápio do restaurante em um banco de dados de vetores e recuperá-los para uso em um fluxo que determina quais alimentos estão disponíveis.

Instale dependências para processar PDFs

npm install llm-chunk pdf-parse
npm i -D --save @types/pdf-parse

Adicionar um repositório de vetores local à configuração

import {
  devLocalIndexerRef,
  devLocalVectorstore,
} from '@genkit-ai/dev-local-vectorstore';
import { textEmbeddingGecko, vertexAI } from '@genkit-ai/vertexai';

configureGenkit({
  plugins: [
    // vertexAI provides the textEmbeddingGecko embedder
    vertexAI(),

    // the local vector store requires an embedder to translate from text to vector
    devLocalVectorstore([
      {
        indexName: 'menuQA',
        embedder: textEmbeddingGecko,
      },
    ]),
  ],
});

Definir um indexador

O exemplo a seguir mostra como criar um indexador para ingerir uma coleção de documentos PDF e armazená-los em um banco de dados de vetores local.

Ele usa o recuperador de similaridade de vetores baseado em arquivos que o Genkit oferece pronto para uso para testes e prototipagem simples (não use em produção).

Criar o indexador

import { devLocalIndexerRef } from '@genkit-ai/dev-local-vectorstore';

export const menuPdfIndexer = devLocalIndexerRef('menuQA');

Criar configuração de divisão

Este exemplo usa a biblioteca llm-chunk, que fornece um divisor de texto simples para dividir documentos em segmentos que podem ser vetorizados.

A definição a seguir configura a função de agrupamento para garantir um segmento de documento de 1.000 a 2.000 caracteres, divididos no final de uma frase, com uma sobreposição entre partes de 100 caracteres.

const chunkingConfig = {
  minLength: 1000,
  maxLength: 2000,
  splitter: 'sentence',
  overlap: 100,
  delimiters: '',
} as any;

Mais opções de agrupamento para essa biblioteca podem ser encontradas na documentação de llm-chunk.

Definir seu fluxo do indexador

import { index } from '@genkit-ai/ai';
import { Document } from '@genkit-ai/ai/retriever';
import { defineFlow, run } from '@genkit-ai/flow';
import { readFile } from 'fs/promises';
import { chunk } from 'llm-chunk';
import path from 'path';
import pdf from 'pdf-parse';
import * as z from 'zod';

export const indexMenu = defineFlow(
  {
    name: 'indexMenu',
    inputSchema: z.string().describe('PDF file path'),
    outputSchema: z.void(),
  },
  async (filePath: string) => {
    filePath = path.resolve(filePath);

    // Read the pdf.
    const pdfTxt = await run('extract-text', () =>
      extractTextFromPdf(filePath)
    );

    // Divide the pdf text into segments.
    const chunks = await run('chunk-it', async () =>
      chunk(pdfTxt, chunkingConfig)
    );

    // Convert chunks of text into documents to store in the index.
    const documents = chunks.map((text) => {
      return Document.fromText(text, { filePath });
    });

    // Add documents to the index.
    await index({
      indexer: menuPdfIndexer,
      documents,
    });
  }
);

async function extractTextFromPdf(filePath: string) {
  const pdfFile = path.resolve(filePath);
  const dataBuffer = await readFile(pdfFile);
  const data = await pdf(dataBuffer);
  return data.text;
}

Executar o fluxo do indexador

genkit flow:run indexMenu "'../pdfs'"

Depois de executar o fluxo indexMenu, o banco de dados de vetores será alimentado com documentos e estará pronto para ser usado em fluxos do Genkit com etapas de recuperação.

Definir um fluxo com recuperação

O exemplo a seguir mostra como usar um retriever em um fluxo RAG. Assim como no exemplo do indexador, este exemplo usa o recuperador de vetor baseado em arquivo do Genkit, que não pode ser usado em produção.

import { generate } from '@genkit-ai/ai';
import { retrieve } from '@genkit-ai/ai/retriever';
import { devLocalRetrieverRef } from '@genkit-ai/dev-local-vectorstore';
import { defineFlow } from '@genkit-ai/flow';
import { geminiPro } from '@genkit-ai/vertexai';
import * as z from 'zod';

// Define the retriever reference
export const menuRetriever = devLocalRetrieverRef('menuQA');

export const menuQAFlow = defineFlow(
  { name: 'menuQA', inputSchema: z.string(), outputSchema: z.string() },
  async (input: string) => {
    // retrieve relevant documents
    const docs = await retrieve({
      retriever: menuRetriever,
      query: input,
      options: { k: 3 },
    });

    // generate a response
    const llmResponse = await generate({
      model: geminiPro,
      prompt: `
    You are acting as a helpful AI assistant that can answer 
    questions about the food available on the menu at Genkit Grub Pub.
    
    Use only the context provided to answer the question.
    If you don't know, do not make up an answer.
    Do not add or change items on the menu.

    Question: ${input}
    `,
      context: docs,
    });

    const output = llmResponse.text();
    return output;
  }
);

Crie seus próprios indexadores e recuperados

Também é possível criar seu próprio recuperador. Isso é útil quando os documentos são gerenciados em um repositório que não é compatível com o Genkit (por exemplo: MySQL, Google Drive etc.). O SDK do Genkit oferece métodos flexíveis que permitem fornecer código personalizado para buscar documentos. Também é possível definir coletores personalizados que se baseiam em recuperadores existentes no Genkit e aplicar técnicas avançadas de RAG (como extensões de comando ou reclassificação).

Recuperadores simples

Recuperadores simples permitem converter facilmente código existente em recuperador:

import {
  defineSimpleRetriever,
  retrieve
} from '@genkit-ai/ai/retriever';
import { searchEmails } from './db';
import { z } from 'zod';

defineSimpleRetriever({
  name: 'myDatabase',
  configSchema: z.object({
    limit: z.number().optional()
  }).optional(),
  // we'll extract "message" from the returned email item
  content: 'message',
  // and several keys to use as metadata
  metadata: ['from', 'to', 'subject'],
} async (query, config) => {
  const result = await searchEmails(query.text(), {limit: config.limit});
  return result.data.emails;
});

Recuperadores personalizados

import {
  CommonRetrieverOptionsSchema,
  defineRetriever,
  retrieve,
} from '@genkit-ai/ai/retriever';
import * as z from 'zod';

export const menuRetriever = devLocalRetrieverRef('menuQA');

const advancedMenuRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({
  preRerankK: z.number().max(1000),
});

const advancedMenuRetriever = defineRetriever(
  {
    name: `custom/advancedMenuRetriever`,
    configSchema: advancedMenuRetrieverOptionsSchema,
  },
  async (input, options) => {
    const extendedPrompt = await extendPrompt(input);
    const docs = await retrieve({
      retriever: menuRetriever,
      query: extendedPrompt,
      options: { k: options.preRerankK || 10 },
    });
    const rerankedDocs = await rerank(docs);
    return rerankedDocs.slice(0, options.k || 3);
  }
);

extendPrompt e rerank são algo que você precisaria implementar por conta própria, não fornecido pelo framework.

Depois, é só trocar seu recuperador:

const docs = await retrieve({
  retriever: advancedRetriever,
  query: input,
  options: { preRerankK: 7, k: 3 },
});