How to Do Citations with Ragie

Matt Kauffman

Principal Engineer

In this cookbook, we’ll walk you through how to provide sources for generated responses using Ragie and OpenAI GPT4-o. This feature is very important for applications that require source-backed information.

‍Prerequisites

Ragie API Key: To access Ragie for document retrieval.
OpenAI API Key: To access OpenAI for response generation.
Zod Library: Used for schema validation to ensure response formatting.

‍

Setup

First, import the necessary libraries and initialize Ragie and OpenAI using your API keys.


import OpenAI from "openai";
import { zodResponseFormat } from "openai/helpers/zod";
import { z } from "zod";
import { Ragie } from "ragie";

const openai = new OpenAI({
  apiKey: process.env["OPENAI_API_KEY"],
});

const ragie = new Ragie({
  auth: process.env["RAGIE_API_KEY"],
});

‍

Schema Definition

Define a schema for the response structure. This schema will help format responses to include:

Message: The generated response text.
Used Document IDs: IDs of documents used to create the response, which we’ll use to generate citations.
‍


export const citedResponseSchema = z.object({
  usedDocumentIds: z
    .array(
      z
        .string()
        .describe(
          "The document_id of a chunk that was used as part of the generated response"
        )
    )
    .describe(
      "The document_ids of the chunks used to generate the response. document_ids of chunks not used in the response are not included in this list"
    ),
  message: z.string().describe("The response message"),
});

‍

System Message Creation

Create a function to format system messages that will guide the model. This message includes chunks retrieved from Ragie, instructing the model to incorporate this information into the generated response.


const createSystemMessage = (chunks: string[]) => `
  <YOUR SYSTEM MESSAGE>

  Here are relevant chunks that you can use to respond to the user. Remember to incorporate this information into your responses.
  
  ${chunks}
`;

‍

Retrieve Relevant Chunks from Ragie

Define a function that:

Retrieves relevant document chunks based on a query.
Generates a response with the retrieved chunks incorporated.
Includes citations from the used document IDs.


export async function generateCompletionWithCitations(query: string) {

  const chunks = await ragie.retrievals.retrieve({
    query,
    rerank: true,
    topK: 6,
    maxChunksPerDocument: 4,
  });

  const chunkText = chunks.scoredChunks.map((chunk) => JSON.stringify(chunk));

‍

This function calls ragie.retrievals.retrieve, which:

Retrieves the top-ranked chunks using rerank to enhance relevance.
Limits the result to top 6 results (topK: 6) and 4 chunks per document
(maxChunksPerDocument: 4), feel free to customize this based on your use-case.

‍

Generate the Completion with Citations

Once we retrieve the chunks, we pass them to the GPT-4o to generate a response using the citedResponseSchema format. This allows the model to reference specific document chunks as citations.


const completion = await openai.beta.chat.completions.parse({
    model: "gpt-4o-2024-08-06",
    temperature: 0.3,
    messages: [
      { role: "system", content: createSystemMessage(chunkText) },
      {
        role: "user",
        content: query,
      },
    ],
    response_format: zodResponseFormat(citedResponseSchema, "citedResponse"),
  });

Explanation:

Temperature: Set to 0.3 for controlled, factual responses.
Response Format: ZodResponseFormat will be used to validate the response format based on the citedResponseSchema. This provides OpenAI structured output with the expected schema for its output.

‍

Parse and Format Output with Citations

The parseOutput function processes the model’s response and retrieves the cited document’s metadata. It formats the output to include the citation as a hyperlink, referencing the document source URL. This assumes that your documents were loaded into Ragie with a `source_url` metadata filter. Adjust as needed to suit the metadata schema of your documents.


const parseOutput = async (output: { message: string; usedDocumentIds: string[] } | null) => {

    if (!output) return "";
  
    const citedDocs = output.usedDocumentIds.map((docId) => {
      const doc = chunks.scoredChunks.find((chunk) => chunk.document_id === docId);

      return doc ? `<a href="${doc.metadata.source_url}">${doc.metadata.name}</a>` : "Unknown Source";
    });

    return `${output.message} Cited from: ${citedDocs.join(", ")}`;
  };
‍

  return parseOutput(completion.choices[0].message.parsed);
}

Explanation:

Document Metadata: Retrieves metadata for the document ID used in the response to display the citation.
Hyperlinked Citation: Outputs the response message with a hyperlink to the document’s source.

‍

Full Code Example

Here’s the complete code for generating a response with citations using Ragie and OpenAI:


import OpenAI from "openai";
import { zodResponseFormat } from "openai/helpers/zod";
import { z } from "zod";
import { Ragie } from "ragie";
‍
const openai = new OpenAI({
  apiKey: process.env["OPENAI_API_KEY"],
});

‍

const ragie = new Ragie({
  auth: process.env["RAGIE_API_KEY"],
});

‍

export const citedResponseSchema = z.object({
  usedDocumentIds: z
    .array(
      z
        .string()
        .describe(
          "The document_id of a chunk that was used as part of the generated response"
        )
    )
    .describe(
      "The document_ids of the chunks used to generate the response. document_ids of chunks not used in the response are not included in this list"
    ),
  message: z.string().describe("The response message"),
});

‍

const createSystemMessage = (chunks: string[]) => `
  <YOUR SYSTEM MESSAGE>

  Here are relevant chunks that you can use to respond to the user. Remember to incorporate this information into your responses.
  ${chunks}
`;

‍

export async function generateCompletionWithCitations(query: string) {
  const chunks = await ragie.retrievals.retrieve({
    query,
    rerank: true,
    topK: 6,
    maxChunksPerDocument: 4,
  });

  const chunkText = chunks.scoredChunks.map((chunk) => JSON.stringify(chunk));

  const completion = await openai.beta.chat.completions.parse({
    model: "gpt-4o-2024-08-06",
    temperature: 0.3,
    messages: [
      { role: "system", content: createSystemMessage(chunkText) },
      { role: "user", content: query },
    ],
    response_format: zodResponseFormat(citedResponseSchema, "citedResponse"),
  });

‍

  const parseOutput = async (output: { message: string; usedDocumentIds: string[] } | null) => {

    if (!output) return "";
    
    const citedDocs = output.usedDocumentIds.map((docId) => {
      const doc = chunks.scoredChunks.find((chunk) => chunk.document_id === docId);

      return doc ? `<a href="${doc.metadata.source_url}">${doc.metadata.name}</a>` : "Unknown Source";
    });

    return `${output.message} Cited from: ${citedDocs.join(", ")}`;
  };

‍
  return parseOutput(completion.choices[0].message.parsed);
}

‍

Usage Example

Suppose a user asks: “What are the latest sustainability goals outlined by Company X?” The system retrieves relevant document chunks from Company X's 10-K filing in Ragie, passes them to OpenAI to generate a response, and includes a citation to the source document.

The output might look like:

“Company X has outlined key sustainability targets, including a 25% reduction in emissions by 2030. Cited from Company X 2023 10-K Report.”

‍

Conclusion

This recipe demonstrates how to use Ragie with OpenAI to cite responses, adding credibility to generated content. Try integrating this set up in your own applications for source-backed responses.