Generación mejorada por recuperación (RAG)

Firebase Genkit proporciona abstracciones que te ayudan a compilar flujos de generación con aumento de recuperación (RAG), así como complementos que ofrecen integraciones con herramientas relacionadas.

¿Qué es RAG?

La generación de recuperación mejorada es una técnica que se usa para incorporar fuentes de información en las respuestas de un LLM. Es importante poder hacerlo, ya que, si bien los LLMs suelen entrenarse con un cuerpo amplio de material, el uso práctico de los LLMs suele requerir conocimientos específicos del dominio (por ejemplo, es conveniente usar un LLM para responder las preguntas de los clientes sobre los productos de la empresa).

Una solución es ajustar el modelo con datos más específicos. Sin embargo, puede resultar caro en términos de costos de procesamiento y de esfuerzos necesarios para preparar los datos de entrenamiento adecuados.

Por el contrario, el modelo RAG incorpora fuentes de datos externas en una instrucción el tiempo que se pasa al modelo. Por ejemplo, imagina que la instrucción "¿Cuál es la relación de Bart con Lisa?" pudiera expandirse ("mejorarse") anteponiendo información relevante, dando como resultado el mensaje, "Los hijos de Homer y Marge se llaman Bart, Lisa y Maggie. ¿Cuál es la relación de Bart con a Lisa?"

Este enfoque tiene varias ventajas:

  • Puede ser más rentable porque no tienes que volver a entrenar el modelo.
  • Puedes actualizar tu fuente de datos de forma continua y el LLM puede aprovechar la información actualizada.
  • Ahora tienes la posibilidad de citar referencias en las respuestas de tu LLM.

Por otro lado, usar RAG naturalmente implica instrucciones más largas y ciertos cargos por servicios de la API de LLM por cada token de entrada que envías. En última instancia, debes evaluar las compensaciones de costos para tus aplicaciones.

RAG es un área muy amplia y se usan muchas técnicas diferentes para lograr la mejor calidad de RAG. El framework principal de Genkit ofrece tres abstracciones principales para ayudarte con la RAG:

  • Indexadores: agregan documentos a un "índice".
  • Incorporadores: transforman documentos en una representación vectorial.
  • Recuperadores: recuperan documentos de un "índice", según una consulta.

Estas definiciones son amplias a propósito porque Genkit no ofrece recomendaciones sobre qué es un "índice" o la forma exacta en que se recuperan los documentos. Genkit solo proporciona un formato Document, y el recuperador o el proveedor de implementación del indexador.

Indexadores

El índice es responsable de realizar un seguimiento de tus documentos para que puedas recuperar rápidamente documentos relevantes en función de una consulta específica. Esto es logra de la forma más habitual mediante una base de datos de vectores, que indexa tus documentos usando vectores multidimensionales llamados embeddings. Un embedding de texto (de manera opaca) representa los conceptos expresados en un fragmento de texto; estos se generan usando modelos de AA con propósitos especiales. Mediante la indexación de texto con su embedding, una base de datos de vectores puede agrupar texto relacionado de forma conceptual y recuperar documentos relacionadas con una nueva cadena de texto (la consulta).

Antes de que puedas recuperar documentos para generarlos, debes transferirlos al índice del documento. Un flujo de transferencia típico realiza lo siguiente:

  1. Divide documentos grandes en documentos más pequeños para que solo se utilicen las porciones relevantes para mejorar las instruciones: "fragmentación". Esto es necesario, ya que muchos LLMs tienen una ventana de contexto limitada, por lo que no es práctico incluir documentos completos con una instrucción.

    Genkit no proporciona bibliotecas de fragmentación integradas. Sin embargo, existen bibliotecas fuente abiertas que son compatibles con Genkit.

  2. Genera embeddings para cada fragmento Según la base de datos que utilices, podrías hacerlo de forma explícita con un modelo de generación de embeddings o podrías usar el generador de embeddings proporcionado por la base de datos.

  3. Agrega el bloque de texto y su índice a la base de datos.

Puedes ejecutar el flujo de transferencia con poca frecuencia o solo una vez si estás trabajando con una fuente de datos estable. Por otro lado, si trabajas con datos que cambian con frecuencia, puedes ejecutar continuamente el flujo de transferencia (por ejemplo, en un activador de Cloud Firestore, cuando se actualiza un documento).

Incorporadores

Un incorporador es una función que toma contenido (texto, imágenes, audio, etc.) y crea un vector numérico que codifica el significado semántico del contenido original. Como se mencionó anteriormente, los incorporadores se aprovechan como parte del proceso de indexación; sin embargo, también se pueden usar de forma independiente para crear incorporaciones sin un índice.

Recuperadores

Un recuperador es un concepto que encapsula la lógica relacionada con cualquier tipo de recuperación de datos. Los casos de recuperación más populares suelen incluir recuperación desde almacenes de vectores, sin embargo, en Genkit, un recuperador puede ser cualquier función que devuelva datos.

Para crear un recuperador, puedes usar una de las implementaciones proporcionadas o crear una propia.

Indexadores, incorporadores y recuperadores compatibles

Genkit proporciona compatibilidad con indexadores y recuperadores a través de su sistema de complementos. Se admiten oficialmente los siguientes complementos:

Además, Genkit admite los siguientes almacenes de vectores mediante plantillas de código predefinidas, que puedes personalizar para la configuración de tu base de datos y esquema:

La compatibilidad con el modelo de embedding se proporciona a través de los siguientes complementos:

Complemento Modelos
IA generativa de Google Embedding de texto de Gecko
Google Vertex AI Embedding de texto de Gecko

Define un flujo de RAG

En los siguientes ejemplos, se muestra cómo podrías transferir una colección de documentos PDF del menú de un restaurante a una base de datos de vectores y recuperarlos para usarlos en un flujo que determina qué alimentos hay disponibles.

Instala dependencias para procesar archivos PDF

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

Agrega un almacén de vectores local a tu configuración

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,
      },
    ]),
  ],
});

Define un indexador

En el siguiente ejemplo, se muestra cómo crear un indexador para transferir una colección de documentos PDF y almacenarlos en una base de datos local de vectores.

Usa el recuperador de similitud vectorial local basado en archivos que Genkit proporciona listo para usar para realizar pruebas y crear prototipos (no uses en producción)

Crea el indexador

export const menuPdfIndexer = devLocalIndexerRef('menuQA');

Crea la configuración de fragmentación

En este ejemplo, se usa la biblioteca llm-chunk, que proporciona un separador de texto simple para dividir los documentos en segmentos que se pueden vectorizar.

La siguiente definición configura la función de fragmentación para garantizar un segmento de documento de entre 1,000 y 2,000 caracteres, dividido al final de una oración, con una superposición entre los fragmentos de 100 caracteres.

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

Puedes encontrar más opciones de fragmentación para esta biblioteca en la documentación de llm-chunk.

Define el flujo del 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,
    });
  }
);

Ejecuta el flujo del indexador

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

Después de ejecutar el flujo indexMenu, se iniciará la base de datos de vectores con documentos y lista para usar en flujos de Genkit, con pasos de recuperación.

Define un flujo con recuperación

En el siguiente ejemplo, se muestra cómo usar un recuperador en un flujo RAG. Como en el ejemplo del indexador, se usa el recuperador de vectores basado en archivos de Genkit, que no debes usar en producción.

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

Escribe tus propios indexadores y recuperadores

También es posible crear tu propio recuperador. Esto es útil si tus documentos se administran en un almacén de documentos que no es compatible con Genkit (p. ej., MySQL, Google Drive, etcétera). El SDK de Genkit proporciona métodos flexibles que te permiten proporcionar un código personalizado para recuperar documentos. También puedes definir recuperadores personalizados que se complementan con los recuperados existentes en Genkit y aplicar técnicas avanzadas de RAG (como una reclasificación o instrucción) por encima.

Recuperadores simples

Los buscadores simples te permiten convertir fácilmente el código existente en buscadores:

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

Recuperadores 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 y rerank son algo que debes implementar por tu cuenta, el framework no lo proporciona).

Luego, puedes cambiar el recuperador:

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

Reclasificadores y recuperación en dos etapas

Un modelo de reclasificación, también conocido como codificador cruzado, es un tipo de modelo que, dada una consulta y un documento, mostrará una puntuación de similitud. Usamos esta puntuación para reordenar los documentos según su relevancia para nuestra búsqueda. Las APIs de reordenación toman una lista de documentos (por ejemplo, el resultado de un recuperador) y los vuelven a ordenar en función de su relevancia para la consulta. Este paso puede ser útil para ajustar los resultados y garantizar que se use la información más pertinente en la instrucción proporcionada a un modelo generativo.

Ejemplo de reclasificador

Un reclasificador en Genkit se define en una sintaxis similar a la de los recuperadores y los indexadores. Este es un ejemplo de cómo usar un reranker en Genkit. Este flujo vuelve a clasificar un conjunto de documentos según su relevancia para la consulta proporcionada con un reclasificador predefinido de 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,
    }));
  }
);

Este reclasificador usa el complemento genkit de Vertex AI con semantic-ranker-512 para asignar puntuaciones y clasificar documentos. Cuanto más alta sea la puntuación, más relevante será el documento para la consulta.

Reclasificadores personalizados

También puedes definir reclasificadores personalizados para que se adapten a tu caso de uso específico. Esto es útil cuando necesitas volver a clasificar los documentos con tu propia lógica personalizada o un modelo personalizado. Este es un ejemplo sencillo de cómo definir un reasignador de clasificación 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);
  }
);

Una vez definido, este reajustador personalizado se puede usar como cualquier otro reajustador en tus flujos de RAG, lo que te brinda flexibilidad para implementar estrategias avanzadas de reajustamiento.