Retrieving Relevant Information with Similarity Search in Go

Retrieving Relevant Information with Similarity Search

Welcome back! In the previous lesson, we explored how to generate embeddings for document chunks using Go-compatible APIs. Today, we will build on that knowledge by diving into vector databases and how they enable the efficient retrieval of relevant information through similarity search.

Vector databases are specialized storage systems designed to handle high-dimensional vector data, such as the embeddings we generated in the last lesson. They are crucial for performing similarity searches, which allow us to find document chunks that are semantically similar to a given query. In this lesson, we will use a custom in-memory vector store to understand the core concepts. While this implementation is designed for learning purposes, it follows idiomatic patterns that match production-ready vector stores like Weaviate, Pinecone, or Chroma—meaning you can easily swap the backend when you're ready to scale.

Preparing Documents with LangChain

Before we can perform a similarity search, we need to load and chunk our document using LangChain's document loaders and text splitters.

Here's how to do it in Go:

package main

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

    "github.com/tmc/langchaingo/documentloaders"
    "github.com/tmc/langchaingo/textsplitter"
)

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

    // Load the text file
    file, err := os.Open("data/alice_in_wonderland.txt")
    if err != nil {
        log.Fatalf("failed to open file: %v", err)
    }
    defer file.Close()

    loader := documentloaders.NewText(file)

    // Create a text splitter with specific chunk size and overlap
    splitter := textsplitter.NewRecursiveCharacter(
        textsplitter.WithChunkSize(500),
        textsplitter.WithChunkOverlap(100),
    )

    // Load and split the document in one step
    docs, err := loader.LoadAndSplit(ctx, splitter)
    if err != nil {
        log.Fatalf("failed to load and split document: %v", err)
    }

    if len(docs) == 0 {
        log.Fatalf("no documents loaded from source file")
    }

    fmt.Printf("Successfully loaded %d document chunks\n", len(docs))
}

This code demonstrates how to load a document and split it into chunks using LangChain's built-in tools. The LoadAndSplit method combines loading and chunking in a single operation, making our code more concise and maintainable.

Creating Embeddings and Vector Store

With our document chunks ready, the next step is to set up our embedding model and create a vector store. As you learned in the previous lesson, embeddings are numerical representations of text that capture semantic meaning.

We'll use a custom in-memory vector store that follows the same idiomatic patterns as production vector databases. This means the code you write here will be nearly identical when you switch to a real vector database backend.

package main

import (
    "context"
    "log"
    "os"

    "github.com/tmc/langchaingo/documentloaders"
    "github.com/tmc/langchaingo/embeddings"
    "github.com/tmc/langchaingo/llms/openai"
    "github.com/tmc/langchaingo/textsplitter"

    "codesignal/memstore"
)

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

    // Load and split document (same as before)
    file, err := os.Open("data/alice_in_wonderland.txt")
    if err != nil {
        log.Fatalf("failed to open file: %v", err)
    }
    defer file.Close()

    loader := documentloaders.NewText(file)
    splitter := textsplitter.NewRecursiveCharacter(
        textsplitter.WithChunkSize(500),
        textsplitter.WithChunkOverlap(100),
    )
    docs, err := loader.LoadAndSplit(ctx, splitter)
    if err != nil {
        log.Fatalf("failed to load and split document: %v", err)
    }

    // Initialize OpenAI LLM with embedding model
    llm, err := openai.New(
        openai.WithEmbeddingModel("text-embedding-3-small"),
    )
    if err != nil {
        log.Fatalf("failed to initialize OpenAI client: %v", err)
    }

    // Wrap the LLM in an Embedder
    embedder, err := embeddings.NewEmbedder(llm)
    if err != nil {
        log.Fatalf("failed to create embedder: %v", err)
    }

    // Create the in-memory vector store with the embedder
    store, err := memstore.New(
        memstore.WithEmbedder(embedder),
    )
    if err != nil {
        log.Fatalf("failed to create memstore: %v", err)
    }

    // Add all document chunks to the vector store
    _, err = store.AddDocuments(ctx, docs)
    if err != nil {
        log.Fatalf("failed to add documents to memstore: %v", err)
    }
}

Let's break down what's happening in this code:

  1. We initialize an OpenAI LLM client configured to use the text-embedding-3-small model.
  2. We wrap the LLM in an Embedder, which provides a consistent interface for generating embeddings.
  3. We create our custom in-memory vector store, passing the embedder as a configuration option.
  4. We add all document chunks to the store using AddDocuments, which automatically generates embeddings and stores them.

The memstore is a simple in-memory implementation perfect for learning and development. However, the API signatures (New, WithEmbedder, AddDocuments) follow the same patterns used by production vector stores in the LangChain ecosystem, making it trivial to swap backends later.

Performing Similarity Search

Now that we have our vector store populated with embedded document chunks, we can perform a similarity search to retrieve relevant documents. Similarity search finds document chunks whose embeddings are closest to a given query's embedding, allowing us to extract information that is semantically similar to the query.

Here's the complete example with similarity search:

package main

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

    "github.com/tmc/langchaingo/documentloaders"
    "github.com/tmc/langchaingo/embeddings"
    "github.com/tmc/langchaingo/llms/openai"
    "github.com/tmc/langchaingo/textsplitter"

    "codesignal/memstore"
)

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

    // Load and prepare documents (abbreviated for clarity)
    file, err := os.Open("data/alice_in_wonderland.txt")
    if err != nil {
        log.Fatalf("failed to open file: %v", err)
    }
    defer file.Close()

    loader := documentloaders.NewText(file)
    splitter := textsplitter.NewRecursiveCharacter(
        textsplitter.WithChunkSize(500),
        textsplitter.WithChunkOverlap(100),
    )
    docs, err := loader.LoadAndSplit(ctx, splitter)
    if err != nil {
        log.Fatalf("failed to load and split document: %v", err)
    }

    // Initialize embedder and vector store
    llm, err := openai.New(
        openai.WithEmbeddingModel("text-embedding-3-small"),
    )
    if err != nil {
        log.Fatalf("failed to initialize OpenAI client: %v", err)
    }

    embedder, err := embeddings.NewEmbedder(llm)
    if err != nil {
        log.Fatalf("failed to create embedder: %v", err)
    }

    store, err := memstore.New(
        memstore.WithEmbedder(embedder),
    )
    if err != nil {
        log.Fatalf("failed to create memstore: %v", err)
    }

    _, err = store.AddDocuments(ctx, docs)
    if err != nil {
        log.Fatalf("failed to add documents to memstore: %v", err)
    }

    // Perform similarity search
    query := "Who is Alice?"
    k := 3 // Number of results to retrieve

    results, err := store.SimilaritySearch(ctx, query, k)
    if err != nil {
        log.Fatalf("similarity search failed: %v", err)
    }

    // Display results
    fmt.Printf("Query: %s\n\n", query)
    if len(results) == 0 {
        fmt.Println("No results found.")
        return
    }

    for i, doc := range results {
        fmt.Printf("Result %d:\n%s\n---\n\n", i+1, doc.PageContent)
    }
}

This code demonstrates how to perform a similarity search and display the results. The SimilaritySearch method takes three parameters:

  1. A context for managing the request lifecycle
  2. The query string to search for
  3. The number of top results to return (k)

The method returns a slice of documents ranked by similarity to the query. Each document contains the PageContent field with the actual text chunk. Your exact results will vary slightly depending on chunking boundaries and embedding behavior, but the program will print the query followed by the top k most similar chunks:

Query: Who is Alice?

Result 1:
Alice was beginning to get very tired of sitting by her sister on the bank, and of having nothing to do...
---

Result 2:
So she was considering in her own mind (as well as she could, for the hot day made her feel very sleepy and stupid)...
---

Result 3:
There was nothing so VERY remarkable in that; nor did Alice think it so VERY much out of the way to hear the Rabbit say to itself...
---

If no chunks are similar enough (or if the store is empty), you may see:

Query: Who is Alice?

No results found.

Swapping Vector Store Backends

One of the key advantages of using idiomatic patterns is that switching from our in-memory store to a production vector database is straightforward. For example, if you wanted to use Weaviate instead, you would simply replace:

store, err := memstore.New(
    memstore.WithEmbedder(embedder),
)

with:

store, err := weaviate.New(
    weaviate.WithScheme("http"),
    weaviate.WithHost("localhost:8080"),
    weaviate.WithEmbedder(embedder),
)

The rest of your code—including AddDocuments and SimilaritySearch—remains unchanged because all vector stores in the LangChain ecosystem implement the same interface.

Summary and Next Steps

In this lesson, you learned how to create an in-memory vector store and perform similarity search to retrieve relevant information from documents using LangChain in Go. We built on your knowledge of document loading, splitting, and embedding to enable efficient document retrieval.

Key takeaways:

  • Vector stores are specialized databases for high-dimensional vector data
  • The in-memory store we used follows idiomatic patterns that match production databases
  • SimilaritySearch finds semantically similar documents based on embedding distance
  • You can easily swap vector store backends thanks to consistent interface design

As you move on to the practice exercises, I encourage you to experiment with different documents and queries to solidify your understanding. This hands-on practice will prepare you for the next unit, where we will continue to build on these skills. Keep up the great work, and I look forward to seeing you in the next lesson!

Previous LessonNext Lesson: Asking Questions with Retrieved Context and Templates in Go
Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal