Build a Graph RAG system with LangChain and GraphRetriever (C#)

using System.Text.Json;
using System.Text.Json.Serialization;
using System.Text.RegularExpressions;
using DataStax.AstraDB.DataApi;
using DataStax.AstraDB.DataApi.Core;
using LangChain.Providers;
using LangChain.Providers.OpenAI;
using AstraDocument = DataStax.AstraDB.DataApi.Collections.Document;

namespace TutorialGraphRag;

public class DataItem
{
  [JsonPropertyName("html_doc")]
  public string HtmlDoc { get; set; } = string.Empty;

  [JsonPropertyName("url")]
  public string Url { get; set; } = string.Empty;
}

public class TutorialGraphRagCreateStore
{
  // Helper to extract all hyperlinks from HTML anchor tags
  private static readonly Regex HrefRegex = new(
    @"<a\s+(?:[^>]*?\s+)?href=""([^""]*)""",
    RegexOptions.IgnoreCase | RegexOptions.Compiled
  );

  private static List<string> ExtractHyperlinks(string html)
  {
    var matches = HrefRegex.Matches(html);
    var links = new List<string>();
    foreach (Match match in matches)
    {
      if (match.Groups.Count > 1)
      {
        links.Add(match.Groups[1].Value);
      }
    }
    return links;
  }

  public static async Task Main(string[] args)
  {
    string? endpoint = Environment.GetEnvironmentVariable("API_ENDPOINT"); (1)
    string? applicationToken = Environment.GetEnvironmentVariable(
      "APPLICATION_TOKEN"
    );
    string? openaiApiKey = Environment.GetEnvironmentVariable(
      "OPENAI_API_KEY"
    );

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

    // Initialize the Astra DB Data API Client
    var client = new DataAPIClient();
    var database = client.GetDatabase(endpoint, applicationToken);

    const string dataFilePath = "PATH_TO_DATA_FILE"; (2)

    // Read the JSON file and parse it into a DataItem list
    string rawData = await File.ReadAllTextAsync(dataFilePath);
    var jsonData =
      JsonSerializer.Deserialize<List<DataItem>>(rawData)
      ?? new List<DataItem>();

    // Convert the JSON data into documents with extracted hyperlinks
    var documents =
      new List<(string Content, string Url, List<string> Hyperlinks)>();
    foreach (var item in jsonData)
    {
      var links = ExtractHyperlinks(item.HtmlDoc);
      documents.Add((item.HtmlDoc, item.Url, links));
    }

    // Generate embeddings for the documents using LangChain OpenAI Embeddings
    Console.WriteLine("Generating embeddings...");
    var provider = new OpenAiProvider(openaiApiKey);
    var embedder = new OpenAiEmbeddingModel(
      provider,
      "text-embedding-3-small"
    );

    var contents = documents.ConvertAll(d => d.Content);
    var embeddingResponse = await embedder.CreateEmbeddingsAsync(
      new EmbeddingRequest { Strings = contents }
    );
    var vectors = embeddingResponse.Values;

    var documentsWithEmbeddings = new List<AstraDocument>();
    for (int i = 0; i < documents.Count; i++)
    {
      var doc = new AstraDocument
      {
        ["content"] = documents[i].Content,
        ["metadata"] = new Dictionary<string, object>
        {
          ["url"] = documents[i].Url,
          ["hyperlink"] = documents[i].Hyperlinks,
        },
        ["$vector"] = vectors[i],
      };
      documentsWithEmbeddings.Add(doc);
    }

    // Create a new collection
    // (If a collection with the same name already exists,
    //   that collection is used instead)
    Console.WriteLine("Creating collection...");
    var definition = new CollectionDefinition
    {
      Vector = new VectorOptions
      {
        Dimension = 1536,
        Metric = SimilarityMetric.Cosine,
      },
    };
    var collection = await database.CreateCollectionAsync(
      "graph_rag_tutorial",
      definition
    );

    // In case a collection with this name already existed,
    // delete any documents in the collection
    await collection.DeleteManyAsync(
      Builders<AstraDocument>.CollectionFilter.Empty()
    );

    // Insert the documents into the collection
    Console.WriteLine("Inserting documents...");
    await collection.InsertManyAsync(documentsWithEmbeddings);
  }
}

using DataStax.AstraDB.DataApi;
using LangChain.Providers;
using LangChain.Providers.OpenAI;
using LangChainDocument = LangChain.DocumentLoaders.Document;

namespace TutorialGraphRag;

public class TutorialGraphRagRetrieval
{
  // Helper function to format the retrieved documents for LLM context
  private static string FormatDocs(List<LangChainDocument> docs)
  {
    return string.Join("\n\n", docs.Select(doc => doc.PageContent));
  }

  public static async Task Main(string[] args)
  {
    string? endpoint = Environment.GetEnvironmentVariable("API_ENDPOINT"); (1)
    string? applicationToken = Environment.GetEnvironmentVariable(
      "APPLICATION_TOKEN"
    );
    string? openaiApiKey = Environment.GetEnvironmentVariable(
      "OPENAI_API_KEY"
    );

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

    // Initialize the Astra DB Data API Client and get the collection
    var client = new DataAPIClient();
    var database = client.GetDatabase(endpoint, applicationToken);
    var collection = database.GetCollection("graph_rag_tutorial");

    // Initialize LangChain OpenAI Provider, Embeddings, and LLM
    var provider = new OpenAiProvider(openaiApiKey);
    var embedder = new OpenAiEmbeddingModel(
      provider,
      "text-embedding-3-small"
    );
    var llm = new OpenAiChatModel(provider, "gpt-4o");

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

      {0}

      Question: {1}
      """;

    // Initialize GraphRetriever
    // This retriever first uses vector search to find relevant documents,
    // then uses graph traversal to explore their connections.
    var retriever = new GraphRetriever( (2)
      new GraphRetrieverOptions
      {
        Collection = collection,
        Embedder = 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,
      }
    );

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

    Console.WriteLine($"\nQuestion: {question}\n");

    try
    {
      // Build the RAG chain:
      // 1. Retrieve relevant documents via vector search and graph traversal
      var docs = await retriever.GetRelevantDocumentsAsync(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
      var response = await llm.GenerateAsync(prompt);

      Console.WriteLine("Answer:");
      Console.WriteLine(response.LastMessageContent);
    }
    catch (Exception error)
    {
      Console.Error.WriteLine($"Error during RAG query: {error}");
    }
  }
}