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

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"regexp"

	"github.com/datastax/astra-db-go/v2/astra"
	"github.com/datastax/astra-db-go/v2/astra/filter"
	"github.com/datastax/astra-db-go/v2/astra/options"
	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/llms/openai"
)

// Helper to extract href values from HTML anchor tags
var hrefRegex = regexp.MustCompile(`(?i)<a\s+(?:[^>]*?\s+)?href="([^"]*)"`)

// Helper to extract href values from an HTML string
func extractHyperlinks(html string) []string {
	matches := hrefRegex.FindAllStringSubmatch(html, -1)
	links := make([]string, 0, len(matches))
	for _, m := range matches {
		links = append(links, m[1])
	}
	return links
}

func main() {
	ctx := context.Background()

	endpoint := os.Getenv("API_ENDPOINT") (1)
	applicationToken := os.Getenv("APPLICATION_TOKEN")
	openaiAPIKey := os.Getenv("OPENAI_API_KEY")

	if endpoint == "" || applicationToken == "" || openaiAPIKey == "" {
		log.Fatal(
			"Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.",
		)
	}

	dataFilePath := "PATH_TO_DATA_FILE" (2)

	// Read the JSON file and parse it into a slice of objects
	raw, err := os.ReadFile(dataFilePath)
	if err != nil {
		log.Fatal(err)
	}
	var jsonData []struct {
		HTMLDoc string `json:"html_doc"`
		URL     string `json:"url"`
	}
	if err := json.Unmarshal(raw, &jsonData); err != nil {
		log.Fatal(err)
	}

	// Convert the slice into documents with extracted hyperlinks
	type docRecord struct {
		content   string
		url       string
		hyperlink []string
	}
	docs := make([]docRecord, len(jsonData))
	for i, d := range jsonData {
		docs[i] = docRecord{
			content:   d.HTMLDoc,
			url:       d.URL,
			hyperlink: extractHyperlinks(d.HTMLDoc),
		}
	}

	// Generate embeddings for the documents
	// using LangChain OpenAI Embeddings
	fmt.Println("Generating embeddings...")
	llm, err := openai.New(openai.WithToken(openaiAPIKey))
	if err != nil {
		log.Fatal(err)
	}
	embedder, err := embeddings.NewEmbedder(llm)
	if err != nil {
		log.Fatal(err)
	}

	contents := make([]string, len(docs))
	for i, d := range docs {
		contents[i] = d.content
	}
	vectors, err := embedder.EmbedDocuments(ctx, contents)
	if err != nil {
		log.Fatal(err)
	}

	// Build the documents with embeddings
	docsWithEmbeddings := make([]any, len(docs))
	for i, d := range docs {
		docsWithEmbeddings[i] = astra.NewDocument{
			"content": d.content,
			"metadata": map[string]any{
				"url":       d.url,
				"hyperlink": d.hyperlink,
			},
			"$vector": vectors[i],
		}
	}

	// Initialize the Astra DB Data API client
	client := astra.NewClient()
	database := client.Database(endpoint, options.API().SetToken(applicationToken))

	// Create a new collection
	// (If a collection with the same name already exists,
	//   that collection is used instead)
	fmt.Println("Creating collection...")
	collection, err := database.CreateCollection(ctx, "graph_rag_tutorial",
		options.CreateCollection().UpdateVector(
			options.Vector().
				SetDimension(1536).
				SetMetric(options.MetricCosine),
		),
	)
	if err != nil {
		log.Fatal(err)
	}

	// In case a collection with this name already existed,
	// delete any documents in the collection
	if _, err := collection.DeleteMany(ctx, filter.F{}); err != nil {
		log.Fatal(err)
	}

	// Insert the documents into the collection
	fmt.Println("Inserting documents...")
	if _, err := collection.InsertMany(
		ctx,
		docsWithEmbeddings,
	); err != nil {
		log.Fatal(err)
	}
}

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/datastax/astra-db-go/v2/astra"
	"github.com/datastax/astra-db-go/v2/astra/options"
	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/schema"

	"tutorial-graph-rag/graphretriever" (1)
)

func main() {
	ctx := context.Background()

	endpoint := os.Getenv("API_ENDPOINT") (2)
	applicationToken := os.Getenv("APPLICATION_TOKEN")
	openaiAPIKey := os.Getenv("OPENAI_API_KEY")

	if endpoint == "" || applicationToken == "" || openaiAPIKey == "" {
		log.Fatal(
			"Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined.",
		)
	}

	// Initialize the Astra DB Data API client and get the collection
	client := astra.NewClient()
	database := client.Database(endpoint, options.API().SetToken(applicationToken))
	collection := database.Collection("graph_rag_tutorial")

	// Initialize LangChain OpenAI LLM
	llm, err := openai.New(
		openai.WithToken(openaiAPIKey),
		openai.WithModel("gpt-4o"),
	)
	if err != nil {
		log.Fatal(err)
	}

	// Initialize LangChain Embeddings
	embedder, err := embeddings.NewEmbedder(llm)
	if err != nil {
		log.Fatal(err)
	}

	// Initialize GraphRetriever.
	// This retriever first uses vector search to find relevant documents,
	// then uses graph traversal to explore their connections.
	retriever := &graphretriever.GraphRetriever{
		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,
	}

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

%s

Question: %s`

	// 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?"

	fmt.Printf("\nQuestion: %s\n\n", question)

	// Build the RAG chain:
	// 1. Retrieve relevant documents via vector search and graph traversal
	docs, err := retriever.GetRelevantDocuments(ctx, question)
	if err != nil {
		log.Fatalf("Error during retrieval: %v", err)
	}

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

	// 3. Inject the context and question into the prompt template
	prompt := fmt.Sprintf(promptTemplate, ctx2, question)

	// 4. Pass the formatted prompt to the LLM and print the response
	response, err := llm.Call(ctx, prompt)
	if err != nil {
		log.Fatalf("Error during RAG query: %v", err)
	}

	fmt.Println("Answer:")
	fmt.Println(response)
}

// Helper to concatenate content for use as LLM context
func formatDocs(docs []schema.Document) string {
	parts := make([]string, len(docs))
	for i, doc := range docs {
		parts[i] = doc.PageContent
	}
	return strings.Join(parts, "\n\n")
}