Geração aumentada de recuperação (RAG)

O Firebase Genkit fornece abstrações que ajudam a criar geração aumentada de recuperação (RAG), 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 de informação nas respostas de LLMs. É importante ser capaz de fazer porque, embora os LLMs sejam normalmente treinados com um amplo material, o uso prático de LLMs requer conhecimento de domínio específico (por exemplo, talvez você queira usar um LLM para responder perguntas sobre os produtos da sua empresa).

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

Em contraste, o RAG funciona incorporando fontes de dados externas em um comando na hora em que é passado para o modelo. Por exemplo, é possível imaginar o comando, "Qual é a relação de Bart com Lisa?" pode ser expandida ("aumentado") adicionando algumas informações relevantes, resultando no comando, "Os filhos de Homer e Marge se chamam Bart, Lisa e Maggie. Qual é a relação de Bart com Lisa?"

Essa abordagem tem várias vantagens:

  • Essa opção 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 fazer uso imediato das informações atualizadas.
  • Agora você pode citar referências nas respostas do seu LLM.

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

RAG é uma área muito ampla e existem muitas técnicas diferentes usadas para obter RAG da melhor qualidade. O framework Genkit principal oferece duas abstrações principais para ajudar você a fazer a RAG

  • Indexadores: adicionam documentos a um "índice".
  • Incorporadores: transformam documentos em uma representação vetorial.
  • Recuperadores: recuperam documentos de um "índice", de acordo com uma consulta.

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

Indexadores

O índice é responsável por monitorar seus documentos de maneira que você possa recuperar rapidamente documentos relevantes com base em uma consulta específica. Na maioria das vezes, isso é feito usando um banco de dados vetorial, que indexa seus documentos usando vetores multidimensionais chamados embeddings. Um embedding de texto (opacamente) representa os conceitos expressos por uma passagem de texto. Eles são gerados usando modelos de ML para fins especiais. Ao indexar texto usando sua incorporação, um banco de dados vetorial é capaz de agrupar texto conceitualmente relacionado e recuperar documentos relacionados a uma nova sequência de texto (a consulta).

Antes de poder recuperar documentos para fins de geração, é necessário ingeri-los em seu í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 seus comandos: "fragmentação". Isto é necessário porque muitos LLMs têm uma janela de contexto limitada, tornando impraticável incluir documentos inteiros com um comando.

    O Genkit não fornece bibliotecas de chunking integradas. No entanto, existem bibliotecas de código aberto disponíveis que são compatíveis com o Genkit.

  2. Gerar embeddings para cada bloco. Dependendo do banco de dados que você está usando, você pode fazer isso explicitamente com um modelo de geração de incorporação ou pode usar o gerador de incorporação fornecido pelo banco de dados.

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

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

Incorporadores

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

Recuperadores

Um recuperador é um conceito que encapsula a lógica relacionada a qualquer tipo de recuperação de documento. Os casos de recuperação mais populares normalmente 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 seu próprio.

Indexadores, recuperadores e incorporadores compatíveis

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

Além disso, o Genkit oferece suporte aos seguintes armazenamentos de vetores por meio de modelos de código predefinidos, que você pode personalizar para a configuração e o esquema do seu banco de dados:

O suporte a modelos de embedding é fornecido pelos seguintes plug-ins:

Plug-in Modelos
IA generativa do Google Embedding de texto da Gecko
Vertex AI do Google Embedding de texto da Gecko

Como definir um fluxo RAG

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

Instalar dependências para o processamento de 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 vetorial local.

Ele usa o recuperador de similaridade vetorial baseado em arquivo local que o Genkit fornece pronto para uso para testes e prototipagem simples (não use na 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 divisão para garantir um segmento de documento entre 1.000 e 2.000 caracteres, quebrado no final de uma frase, com uma sobreposição entre blocos de 100 caracteres.

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

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

Definir o 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 vetorial será carregado com documentos e estará pronto para ser usado em fluxos Genkit com etapas de recuperação.

Definir um fluxo com recuperação

O exemplo a seguir mostra como você pode usar um recuperador em um fluxo RAG. Assim como o exemplo do indexador, este exemplo usa o recuperador de vetores baseado em arquivo do Genkit, que você não deve usar na 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 { gemini15Flash } 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: gemini15Flash,
      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;
  }
);

Escrever seus próprios indexadores e recuperadores

Também é possível criar seu próprio recuperador. Isso é útil se seus documentos são gerenciados em um repositório de documentos que não é compatível com o Genkit (por exemplo, MySQL, Google Drive etc.). O Genkit SDK fornece métodos flexíveis que permitem fornecer código personalizado para buscar documentos. Também é possível definir retrievers que se baseiam em retrievers existentes no Genkit e aplicam recursos avançados Técnicas de RAG (como reclassificação ou extensões de comando) no início.

Recuperadores simples

Os recuperadores simples permitem que você converta facilmente um código existente em retrievers:

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ê teria que implementar por conta própria, não fornecidos pelo framework)

Depois, basta trocar o retriever:

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