Build a Graph RAG system with LangChain and GraphRetriever (TypeScript)

import * as fs from "fs";
import { DataAPIClient, SomeDoc } from "@datastax/astra-db-ts";
import { OpenAIEmbeddings } from "@langchain/openai";

const endpoint = process.env.API_ENDPOINT; (1)
const applicationToken = process.env.APPLICATION_TOKEN;
const openaiApiKey = process.env.OPENAI_API_KEY;

if (!endpoint || !applicationToken || !openaiApiKey) {
  throw new Error(
    "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.",
  );
}

// Initialize the Astra DB Data API Client
const client = new DataAPIClient();
const database = client.db(endpoint, {
  token: applicationToken,
});

const dataFilePath = "PATH_TO_DATA_FILE"; (2)

// Read the JSON file and parse it into a JSON array
const rawData = fs.readFileSync(dataFilePath, "utf8");
const jsonData = JSON.parse(rawData);

// Helper to extract all hyperlinks from HTML anchor tags
function extractHyperlinks(html: string): string[] {
  const hrefRegex = /<a\s+(?:[^>]*?\s+)?href="([^"]*)"/gi;
  const links: string[] = [];
  let match;
  while ((match = hrefRegex.exec(html)) !== null) {
    links.push(match[1]);
  }
  return links;
}

(async function () {
  // Convert the JSON array into documents with extracted hyperlinks
  const documents: SomeDoc[] = jsonData.map(
    (data: { html_doc: string; url: string }) => {
      const links = extractHyperlinks(data.html_doc);
      return {
        content: data.html_doc,
        url: data.url,
        hyperlink: links,
      };
    },
  );

  // Generate embeddings for the documents using LangChain OpenAI Embeddings
  console.log("Generating embeddings...");
  const embedder = new OpenAIEmbeddings({ apiKey: openaiApiKey });
  const contents = documents.map((doc) => doc.content);
  const vectors = await embedder.embedDocuments(contents);
  const documentsWithEmbeddings = documents.map((document, i) => ({
    content: document.content,
    metadata: {
      url: document.url,
      hyperlink: document.hyperlink,
    },
    $vector: vectors[i],
  }));

  // Create a new collection
  // (If a collection with the same name already exists,
  //   that collection is used instead)
  console.log("Creating collection...");
  const collection = await database.createCollection("graph_rag_tutorial", {
    vector: {
      dimension: 1536,
      metric: "cosine",
    },
  });

  // In case a collection with this name already existed,
  // delete any documents in the collection
  await collection.deleteMany({});

  // Insert the documents into the collection
  console.log("Inserting documents...");
  await collection.insertMany(documentsWithEmbeddings);
})();

import { Collection } from "@datastax/astra-db-ts";
import { Document } from "@langchain/core/documents";
import {
  BaseRetriever,
  type BaseRetrieverInput,
} from "@langchain/core/retrievers";
import { OpenAIEmbeddings } from "@langchain/openai";

interface GraphRetrieverInput extends BaseRetrieverInput {
  collection: Collection;
  embedder: OpenAIEmbeddings;
  startK: number;
  selectK: number;
  maxDepth: number;
}

// This retriever first uses vector search to find relevant documents,
// then uses graph traversal to explore their connections.
export class GraphRetriever extends BaseRetriever {
  lc_namespace = ["langchain", "retrievers"];

  private collection: Collection;
  private embedder: OpenAIEmbeddings;
  private startK: number;
  private selectK: number;
  private maxDepth: number;

  constructor(fields: GraphRetrieverInput) {
    super(fields);
    this.collection = fields.collection;
    this.embedder = fields.embedder;
    this.startK = fields.startK;
    this.selectK = fields.selectK;
    this.maxDepth = fields.maxDepth;
  }

  async _getRelevantDocuments(query: string): Promise<Document[]> {
    // Generate embedding vector for the search query
    const queryVector = await this.embedder.embedQuery(query);

    // 1. Native vector search for initial seed documents
    const seedDocsCursor = this.collection.find(
      {},
      {
        sort: { $vector: queryVector },
        limit: this.startK,
      },
    );

    const seedDocs: Document[] = [];
    for await (const doc of seedDocsCursor) {
      seedDocs.push(
        new Document({
          pageContent: doc.content || "",
          metadata: doc.metadata || {},
        }),
      );
    }

    const retrievedDocs = [...seedDocs];
    const retrievedUrls = new Set<string>(
      retrievedDocs.map((doc) => doc.metadata.url).filter(Boolean),
    );

    // 2. Graph Traversal
    let currentDepth = 0;
    let frontier = [...seedDocs];

    while (
      currentDepth < this.maxDepth &&
      retrievedDocs.length < this.selectK &&
      frontier.length > 0
    ) {
      const nextFrontier: Document[] = [];

      // Extract all hyperlinks from the current frontier
      const linksToFetch = new Set<string>();
      for (const doc of frontier) {
        const links: string[] = doc.metadata.hyperlink || [];
        for (const link of links) {
          if (!retrievedUrls.has(link)) {
            linksToFetch.add(link);
          }
        }
      }

      if (linksToFetch.size === 0) {
        break;
      }

      // Fetch documents matching these URLs from Astra DB
      const fetchedDocs: Document[] = [];
      const linksArray = Array.from(linksToFetch);

      // Query the collection directly to perform metadata lookup
      const cursor = this.collection.find({
        "metadata.url": { $in: linksArray },
      });

      for await (const doc of cursor) {
        fetchedDocs.push(
          new Document({
            pageContent: doc.content || "",
            metadata: doc.metadata || {},
          }),
        );
      }

      // Add to retrieved list and prepare next frontier
      for (const doc of fetchedDocs) {
        if (retrievedDocs.length >= this.selectK) {
          break;
        }

        const url = doc.metadata.url;
        if (url && !retrievedUrls.has(url)) {
          retrievedUrls.add(url);
          retrievedDocs.push(doc);
          nextFrontier.push(doc);
        }
      }

      frontier = nextFrontier;
      currentDepth++;
    }

    return retrievedDocs;
  }
}

import { DataAPIClient } from "@datastax/astra-db-ts";
import { Document } from "@langchain/core/documents";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import {
  RunnablePassthrough,
  RunnableSequence,
} from "@langchain/core/runnables";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { GraphRetriever } from "./GraphRetriever"; (1)

const endpoint = process.env.API_ENDPOINT; (2)
const applicationToken = process.env.APPLICATION_TOKEN;
const openaiApiKey = process.env.OPENAI_API_KEY;

if (!endpoint || !applicationToken || !openaiApiKey) {
  throw new Error(
    "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.",
  );
}

// Initialize the Astra DB Data API Client and get the collection
const client = new DataAPIClient();
const database = client.db(endpoint, {
  token: applicationToken,
});
const collection = database.collection("graph_rag_tutorial");

// Initialize LangChain Embeddings
const embedder = new OpenAIEmbeddings({ apiKey: openaiApiKey });

// Initialize the LLM
const llm = new ChatOpenAI({
  model: "gpt-4o",
  apiKey: openaiApiKey,
});

// Define the prompt template
const template = `Answer the question based only on the following context:

{context}

Question: {question}
`;
const prompt = ChatPromptTemplate.fromTemplate(template);

// Initialize GraphRetriever
// This retriever first uses vector search to find relevant documents,
// then uses graph traversal to explore their connections.
const retriever = new GraphRetriever({
  collection,
  embedder,
  // Number of documents to fetch via vector search for starting the traversal
  startK: 3,
  // Maximum total documents to retrieve during traversal
  selectK: 10,
  // Maximum traversal depth
  // A value of 0 only performs vector search, but does not do any graph traversal
  maxDepth: 1,
});

// Helper function to format the retrieved documents for LLM context
function formatDocs(docs: Document[]): string {
  return docs.map((doc) => doc.pageContent).join("\n\n");
}

// Build the RAG chain:
// 1. Take the input question, retrieve relevant documents from a vector store,
//   and format the documents into a string.
//   Create a dictionary with the formatted documents and the original question.
// 2. Create the prompt by injecting the dictionary from step 1 into the prompt template.
// 3. Pass the formatted prompt to the LLM.
// 4. Clean and return the LLM output.
const chain = RunnableSequence.from([
  {
    context: retriever.pipe(formatDocs),
    question: new RunnablePassthrough(),
  },
  prompt,
  llm,
  new StringOutputParser(),
]);

// Try these questions to explore the knowledge graph:
const QUESTION = "What is close to the Space Needle?";
// Alternative questions:
// - "What is in the Lower Queen Anne neighborhood?"
// - "What is in the same neighborhood as the Space Needle?"
// - "What connects the 1962 World's Fair to modern Seattle?"
// - "Where is Chihuly Garden and Glass?"

(async function () {
  console.log(`\nQuestion: ${QUESTION}\n`);
  try {
    const response = await chain.invoke(QUESTION);
    console.log("Answer:");
    console.log(response);
  } catch (error) {
    console.error("Error during RAG query:", error);
  }
})();