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 fazer isso 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 usar imediatamente as 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 três 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, é possível fazer isso explicitamente com um modelo de geração de incorporação ou 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 sua própria.

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 processar PDFs

npm install llm-chunk pdf-parse @genkit-ai/dev-local-vectorstore
npm i -D --save @types/pdf-parse

Adicionar uma loja de vetores local à configuração

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

const ai = genkit({
  plugins: [
    // vertexAI provides the textEmbedding004 embedder
    vertexAI(),

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

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 em testes e prototipagem simples (não use na produção)

Criar o indexador

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 em blocos para garantir um segmento de documento de 1.000 a 2.000 caracteres, dividido no final de uma frase, com uma sobreposição entre pedaços de 100 caracteres.

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

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

Definir o fluxo do indexador

import { Document } from 'genkit/retriever';
import { chunk } from 'llm-chunk';
import { readFile } from 'fs/promises';
import path from 'path';
import pdf from 'pdf-parse';

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

export const indexMenu = ai.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 ai.index({
      indexer: menuPdfIndexer,
      documents,
    });
  }
);

Executar o fluxo do indexador

genkit flow:run indexMenu "'menu.pdf'"

Depois de executar o fluxo indexMenu, o banco de dados vetorial será propagado com documentos e 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 { devLocalRetrieverRef } from '@genkit-ai/dev-local-vectorstore';

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

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

    // generate a response
   const { text } = await ai.generate({
      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}`,
      docs,
    });

    return text;
  }
);

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. Você também pode definir recuperadores personalizados que se baseiam em recuperadores existentes no Genkit e aplicar técnicas avançadas de RAG (como reclassificação ou extensão de comando) na parte de cima.

Recuperadores simples

Os retrievers simples permitem converter facilmente o código atual em retrievers:

import { z } from "genkit";
import { searchEmails } from "./db";

ai.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;
  }
);

Retrievers personalizados

import {
  CommonRetrieverOptionsSchema,
} from 'genkit/retriever';
import { z } from 'genkit';

export const menuRetriever = devLocalRetrieverRef('menuQA');

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

const advancedMenuRetriever = ai.defineRetriever(
  {
    name: `custom/advancedMenuRetriever`,
    configSchema: advancedMenuRetrieverOptionsSchema,
  },
  async (input, options) => {
    const extendedPrompt = await extendPrompt(input);
    const docs = await ai.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ê precisa implementar, não fornecido pelo framework.

E você pode trocar o retriever:

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

Rerankers e recuperação em duas etapas

Um modelo de reclassificação, também conhecido como codificador cruzado, é um tipo de modelo que, dada uma consulta e um documento, gera uma pontuação de similaridade. Usamos essa pontuação para reordenar os documentos por relevância para nossa consulta. As APIs de reclassificação usam uma lista de documentos (por exemplo, a saída de um recuperador) e reordenam os documentos com base na relevância deles para a consulta. Essa etapa pode ser útil para ajustar os resultados e garantir que as informações mais pertinentes sejam usadas na instrução fornecida a um modelo generativo.

Exemplo de Reranker

Um reclassificador no Genkit é definido em uma sintaxe semelhante a dos extratores e indexadores. Confira um exemplo de uso de um reranker no Genkit. Esse fluxo classifica novamente um conjunto de documentos com base na relevância deles para a consulta fornecida usando um reranker predefinido da Vertex AI.

const FAKE_DOCUMENT_CONTENT = [
  'pythagorean theorem',
  'e=mc^2',
  'pi',
  'dinosaurs',
  'quantum mechanics',
  'pizza',
  'harry potter',
];

export const rerankFlow = ai.defineFlow(
  {
    name: 'rerankFlow',
    inputSchema: z.object({ query: z.string() }),
    outputSchema: z.array(
      z.object({
        text: z.string(),
        score: z.number(),
      })
    ),
  },
  async ({ query }) => {
    const documents = FAKE_DOCUMENT_CONTENT.map((text) =>
       ({ content: text })
    );

    const rerankedDocuments = await ai.rerank({
      reranker: 'vertexai/semantic-ranker-512',
      query:  ({ content: query }),
      documents,
    });

    return rerankedDocuments.map((doc) => ({
      text: doc.content,
      score: doc.metadata.score,
    }));
  }
);

Essa ferramenta usa o plug-in do genkit da Vertex AI com semantic-ranker-512 para avaliar e classificar documentos. Quanto maior a pontuação, mais relevante é o documento para a consulta.

Reclassificadores personalizados

Também é possível definir reclassificadores personalizados para atender ao seu caso de uso específico. Isso é útil quando você precisa classificar documentos novamente usando sua própria lógica ou um modelo personalizado. Confira um exemplo simples de como definir um reclassificador personalizado:

export const customReranker = ai.defineReranker(
  {
    name: 'custom/reranker',
    configSchema: z.object({
      k: z.number().optional(),
    }),
  },
  async (query, documents, options) => {
    // Your custom reranking logic here
    const rerankedDocs = documents.map((doc) => {
      const score = Math.random(); // Assign random scores for demonstration
      return {
        ...doc,
        metadata: { ...doc.metadata, score },
      };
    });

    return rerankedDocs.sort((a, b) => b.metadata.score - a.metadata.score).slice(0, options.k || 3);
  }
);

Depois de definido, esse reclassificador personalizado pode ser usado como qualquer outro reclassificador nos fluxos de RAG, dando flexibilidade para implementar estratégias avançadas de reclassificação.