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

package com.example;

import com.datastax.astra.client.DataAPIClient;
import com.datastax.astra.client.collections.Collection;
import com.datastax.astra.client.collections.definition.CollectionDefinition;
import com.datastax.astra.client.collections.definition.documents.Document;
import com.datastax.astra.client.core.vector.SimilarityMetric;
import com.datastax.astra.client.databases.Database;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModelName;
import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class TutorialGraphRagCreateStore {

  // Helper to extract href values from HTML anchor tags
  private static final Pattern HREF_PATTERN =
      Pattern.compile("<a\\s+(?:[^>]*?\\s+)?href=\"([^\"]*)\"", Pattern.CASE_INSENSITIVE);

  // Helper to extract href values from an HTML string
  private static List<String> extractHyperlinks(String html) {
    List<String> links = new ArrayList<>();
    Matcher matcher = HREF_PATTERN.matcher(html);
    while (matcher.find()) {
      links.add(matcher.group(1));
    }
    return links;
  }

  public static void main(String[] args) throws Exception {
    String endpoint = System.getenv("API_ENDPOINT"); (1)
    String applicationToken = System.getenv("APPLICATION_TOKEN");
    String openaiApiKey = System.getenv("OPENAI_API_KEY");

    if (endpoint == null
        || endpoint.isBlank()
        || applicationToken == null
        || applicationToken.isBlank()
        || openaiApiKey == null
        || openaiApiKey.isBlank()) {
      throw new IllegalStateException(
          "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.");
    }

    // Initialize the Astra DB Data API Client
    DataAPIClient client = new DataAPIClient(applicationToken);
    Database database = client.getDatabase(endpoint);

    String dataFilePath = "modules/tutorials/attachments/graph_rag_dataset.json"; (2)

    // Read the JSON file and parse it into a list of data items
    ObjectMapper mapper = new ObjectMapper();
    List<Map<String, String>> jsonData =
        mapper.readValue(new File(dataFilePath), new TypeReference<List<Map<String, String>>>() {});

    // Build content strings and extract hyperlinks
    List<String> contents = new ArrayList<>();
    List<List<String>> allHyperlinks = new ArrayList<>();
    List<String> urls = new ArrayList<>();
    for (Map<String, String> item : jsonData) {
      String htmlDoc = item.get("html_doc");
      contents.add(htmlDoc);
      urls.add(item.get("url"));
      allHyperlinks.add(extractHyperlinks(htmlDoc));
    }

    // Generate embeddings for the documents using LangChain4j OpenAI Embeddings
    System.out.println("Generating embeddings...");
    OpenAiEmbeddingModel embeddingModel =
        OpenAiEmbeddingModel.builder()
            .apiKey(openaiApiKey)
            .modelName(OpenAiEmbeddingModelName.TEXT_EMBEDDING_3_SMALL)
            .build();

    List<TextSegment> segments =
        contents.stream().map(TextSegment::from).collect(Collectors.toList());
    List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

    // Build Astra DB documents with content, metadata, and vector
    List<Document> documentsWithEmbeddings = new ArrayList<>();
    for (int i = 0; i < jsonData.size(); i++) {
      Map<String, Object> metadata = new HashMap<>();
      metadata.put("url", urls.get(i));
      metadata.put("hyperlink", allHyperlinks.get(i));

      Document doc =
          new Document()
              .append("content", contents.get(i))
              .append("metadata", metadata)
              .vector(embeddings.get(i).vector());
      documentsWithEmbeddings.add(doc);
    }

    // Create a new collection
    // (If a collection with the same name already exists,
    //   that collection is used instead)
    System.out.println("Creating collection...");
    Collection<Document> collection =
        database.createCollection(
            "graph_rag_tutorial", new CollectionDefinition().vector(1536, SimilarityMetric.COSINE));

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

    // Insert the documents into the collection
    System.out.println("Inserting documents...");
    collection.insertMany(documentsWithEmbeddings);
  }
}

package com.example;

import com.datastax.astra.client.DataAPIClient;
import com.datastax.astra.client.collections.Collection;
import com.datastax.astra.client.collections.definition.documents.Document;
import com.datastax.astra.client.databases.Database;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiChatModelName;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModelName;
import java.util.List;

public class TutorialGraphRagRetrieval {

  // Helper function to format the retrieved documents for LLM context
  private static String formatDocs(List<Document> docs) {
    StringBuilder sb = new StringBuilder();
    for (int i = 0; i < docs.size(); i++) {
      if (i > 0) sb.append("\n\n");
      Object content = docs.get(i).get("content");
      sb.append(content != null ? content.toString() : "");
    }
    return sb.toString();
  }

  public static void main(String[] args) {
    String endpoint = System.getenv("API_ENDPOINT"); (1)
    String applicationToken = System.getenv("APPLICATION_TOKEN");
    String openaiApiKey = System.getenv("OPENAI_API_KEY");

    if (endpoint == null
        || endpoint.isBlank()
        || applicationToken == null
        || applicationToken.isBlank()
        || openaiApiKey == null
        || openaiApiKey.isBlank()) {
      throw new IllegalStateException(
          "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.");
    }

    // Initialize the Astra DB Data API Client and get the collection
    DataAPIClient client = new DataAPIClient(applicationToken);
    Database database = client.getDatabase(endpoint);
    Collection<Document> collection = database.getCollection("graph_rag_tutorial");

    // Initialize LangChain4j OpenAI Embeddings and LLM
    OpenAiEmbeddingModel embeddingModel =
        OpenAiEmbeddingModel.builder()
            .apiKey(openaiApiKey)
            .modelName(OpenAiEmbeddingModelName.TEXT_EMBEDDING_3_SMALL)
            .build();

    OpenAiChatModel llm =
        OpenAiChatModel.builder()
            .apiKey(openaiApiKey)
            .modelName(OpenAiChatModelName.GPT_4_O)
            .build();

    // Define the prompt template
    String template =
        "Answer the question based only on the following context:\n\n" + "%s\n\n" + "Question: %s";

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

    // Try these questions to explore the knowledge graph:
    String 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?"

    System.out.println("\nQuestion: " + question + "\n");

    try {
      // Build the RAG chain:
      // 1. Retrieve relevant documents via vector search and graph traversal
      List<Document> docs = retriever.getRelevantDocuments(question);

      // 2. Format the retrieved documents into a single context string
      String context = formatDocs(docs);

      // 3. Inject the context and question into the prompt template
      String prompt = String.format(template, context, question);

      // 4. Pass the formatted prompt to the LLM and print the response
      String response = llm.chat(prompt);

      System.out.println("Answer:");
      System.out.println(response);
    } catch (Exception e) {
      System.err.println("Error during RAG query: " + e);
    }
  }
}