In this blog post, we will explore how to build a RAG (Retrieval Augmented Generation) system using Ent, Atlas, and pgvector.

RAG is a technique that augments the power of generative models by incorporating a retrieval step. Instead of relying solely on the model’s internal knowledge, we can retrieve relevant documents or data from an external source and use that information to produce more accurate, context-aware responses. This approach is particularly useful when building applications such as question-answering systems, chatbots, or any scenario where up-to-date or domain-specific knowledge is needed.

Setting Up our Ent schema​

Let's begin our tutorial by initializing the Go module which we will be using for our project:

go mod init github.com/rotemtam/entrag # Feel free to replace the module path with your own

In this project we will use Ent, an entity framework for Go, to define our database schema. The database will store the documents we want to retrieve (chunked to a fixed size) and the vectors representing each chunk. Initialize the Ent project by running the following command:

go run -mod=mod entgo.io/ent/cmd/ent new Embedding Chunk

This command creates placeholders for our data models. Our project should look like this:

├── ent
│   ├── generate.go
│   └── schema
│       ├── chunk.go
│       └── embedding.go
├── go.mod
└── go.sum

Next, let's define the schema for the Chunk model. Open the ent/schema/chunk.go file and define the schema as follows:

package schema

import (
    "entgo.io/ent"
    "entgo.io/ent/schema/edge"
    "entgo.io/ent/schema/field"
)

// Chunk holds the schema definition for the Chunk entity.
type Chunk struct {
    ent.Schema
}

// Fields of the Chunk.
func (Chunk) Fields() []ent.Field {
    return []ent.Field{
        field.String("path"),
        field.Int("nchunk"),
        field.Text("data"),
    }
}

// Edges of the Chunk.
func (Chunk) Edges() []ent.Edge {
    return []ent.Edge{
        edge.To("embedding", Embedding.Type).StorageKey(edge.Column("chunk_id")).Unique(),
    }
}

This schema defines a Chunk entity with three fields: path, nchunk, and data. The path field stores the path of the document, nchunk stores the chunk number, and data stores the chunked text data. We also define an edge to the Embedding entity, which will store the vector representation of the chunk.

Before we proceed, let's install the pgvector package. pgvector is a PostgreSQL extension that provides support for vector operations and similarity search. We will need it to store and retrieve the vector representations of our chunks.

go get github.com/pgvector/pgvector-go

Next, let's define the schema for the Embedding model. Open the ent/schema/embedding.go file and define the schema as follows:

package schema

import (
    "entgo.io/ent"
    "entgo.io/ent/dialect"
    "entgo.io/ent/dialect/entsql"
    "entgo.io/ent/schema/edge"
    "entgo.io/ent/schema/field"
    "entgo.io/ent/schema/index"
    "github.com/pgvector/pgvector-go"
)

// Embedding holds the schema definition for the Embedding entity.
type Embedding struct {
    ent.Schema
}

// Fields of the Embedding.
func (Embedding) Fields() []ent.Field {
    return []ent.Field{
        field.Other("embedding", pgvector.Vector{}).
            SchemaType(map[string]string{
                dialect.Postgres: "vector(1536)",
            }),
    }
}

// Edges of the Embedding.
func (Embedding) Edges() []ent.Edge {
    return []ent.Edge{
        edge.From("chunk", Chunk.Type).Ref("embedding").Unique().Required(),
    }
}

func (Embedding) Indexes() []ent.Index {
    return []ent.Index{
        index.Fields("embedding").
            Annotations(
                entsql.IndexType("hnsw"),
                entsql.OpClass("vector_l2_ops"),
            ),
    }
}

This schema defines an Embedding entity with a single field embedding of type pgvector.Vector. The embedding field stores the vector representation of the chunk. We also define an edge to the Chunk entity and an index on the embedding field using the hnsw index type and vector_l2_ops operator class. This index will enable us to perform efficient similarity searches on the embeddings.

Finally, let's generate the Ent code by running the following commands:

go mod tidy
go generate ./...

Ent will generate the necessary code for our models based on the schema definitions.

Setting Up the database​

Next, let's set up the PostgreSQL database. We will use Docker to run a PostgreSQL instance locally. As we need the pgvector extension, we will use the pgvector/pgvector:pg17 Docker image, which comes with the extension pre-installed.

docker run --rm --name postgres -e POSTGRES_PASSWORD=pass -p 5432:5432 -d pgvector/pgvector:pg17

We will be using Atlas, a database schema-as-code tool that integrates with Ent, to manage our database schema. Install Atlas by running the following command:

curl -sSfL https://atlasgo.io/install.sh | sh

For other installation options, see the Atlas installation docs.

As we are going to managing extensions, we need an Atlas Pro account. You can sign up for a free trial by running:

atlas login

Now, let's create our base configuration base.pg.hcl which provides the vector extension for the public schema:

schema "public" {
}

extension "vector" {
  schema = schema.public
}

Now, let's create our Atlas configuration which composes the base.pg.hcl file with the Ent schema:

data "composite_schema" "schema" {
  schema {
    url = "file://base.pg.hcl"
  }
  schema "public" {
    url = "ent://ent/schema"
  }
}

env "local" {
  url = getenv("DB_URL")
  schema {
    src = data.composite_schema.schema.url
  }
  dev = "docker://pgvector/pg17/dev"
}