Esta página fue traducida por PageTurner AI (beta). No está respaldada oficialmente por el proyecto. ¿Encontraste un error? Reportar problema →
En esta entrada de blog, exploraremos cómo construir un sistema RAG (Retrieval Augmented Generation) utilizando Ent, Atlas y pgvector.
RAG es una técnica que potencia los modelos generativos incorporando un paso de recuperación. En lugar de depender únicamente del conocimiento interno del modelo, podemos recuperar documentos o datos relevantes de una fuente externa y usar esa información para producir respuestas más precisas y contextualizadas. Este enfoque es especialmente útil al construir aplicaciones como sistemas de preguntas y respuestas, chatbots o cualquier escenario donde se necesite conocimiento actualizado o específico de un dominio.
Configuración de nuestro esquema de Ent
Comencemos el tutorial inicializando el módulo de Go que usaremos para nuestro proyecto:
go mod init github.com/rotemtam/entrag # Feel free to replace the module path with your own
En este proyecto usaremos Ent, un framework de entidades para Go, para definir nuestro esquema de base de datos. La base de datos almacenará los documentos que queremos recuperar (divididos en fragmentos de tamaño fijo) y los vectores que representan cada fragmento. Inicializa el proyecto de Ent ejecutando el siguiente comando:
go run -mod=mod entgo.io/ent/cmd/ent new Embedding Chunk
Este comando crea marcadores de posición para nuestros modelos de datos. Nuestro proyecto debería tener este aspecto:
├── ent
│ ├── generate.go
│ └── schema
│ ├── chunk.go
│ └── embedding.go
├── go.mod
└── go.sum
A continuación, definamos el esquema para el modelo Chunk. Abre el archivo ent/schema/chunk.go y define el esquema de la siguiente manera:
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(),
}
}
Este esquema define una entidad Chunk con tres campos: path, nchunk y data. El campo path almacena la ruta
del documento, nchunk almacena el número de fragmento y data almacena los datos de texto fragmentados. También definimos una relación
con la entidad Embedding, que almacenará la representación vectorial del fragmento.
Antes de continuar, instalemos el paquete pgvector. pgvector es una extensión de PostgreSQL que proporciona soporte para
operaciones vectoriales y búsqueda por similitud. La necesitaremos para almacenar y recuperar las representaciones vectoriales de nuestros fragmentos.
go get github.com/pgvector/pgvector-go
A continuación, definamos el esquema para el modelo Embedding. Abre el archivo ent/schema/embedding.go y define el esquema
de la siguiente manera:
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"),
),
}
}
Este esquema define una entidad Embedding con un único campo embedding de tipo pgvector.Vector. El campo embedding
almacena la representación vectorial del fragmento. También definimos una relación con la entidad Chunk y un índice en el
campo embedding usando el tipo de índice hnsw y la clase de operadores vector_l2_ops. Este índice nos permitirá realizar
búsquedas de similitud eficientes sobre los embeddings.
Finalmente, generemos el código de Ent ejecutando los siguientes comandos:
go mod tidy
go generate ./...
Ent generará el código necesario para nuestros modelos basándose en las definiciones del esquema.
Configuración de la base de datos
A continuación, configuremos la base de datos PostgreSQL. Usaremos Docker para ejecutar una instancia de PostgreSQL localmente. Como necesitamos la
extensión pgvector, usaremos la imagen de Docker pgvector/pgvector:pg17, que viene con la extensión
preinstalada.
docker run --rm --name postgres -e POSTGRES_PASSWORD=pass -p 5432:5432 -d pgvector/pgvector:pg17
Utilizaremos Atlas, una herramienta de esquema de base de datos como código que se integra con Ent, para gestionar nuestro esquema de base de datos. Instala Atlas ejecutando el siguiente comando:
curl -sSfL https://atlasgo.io/install.sh | sh
Para otras opciones de instalación, consulta la documentación de instalación de Atlas.
Como vamos a gestionar extensiones, necesitamos una cuenta Atlas Pro. Puedes registrarte para una prueba gratuita ejecutando:
atlas login
Si prefieres omitir Atlas, puedes aplicar el esquema directamente a la base de datos usando las sentencias de este archivo
Ahora, creemos nuestra configuración base base.pg.hcl que proporciona la extensión vector para el esquema público:
schema "public" {
}
extension "vector" {
schema = schema.public
}
Definamos nuestra configuración de Atlas que combina el archivo base.pg.hcl con el esquema de Ent:
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"
}
Esta configuración define un esquema compuesto que incluye el archivo base.pg.hcl y el esquema de Ent. También definimos un entorno llamado local que usa el esquema compuesto para desarrollo local. El campo dev especifica la URL de la Base de Datos de Desarrollo, que Atlas utiliza para normalizar esquemas y realizar cálculos.
Apliquemos el esquema a la base de datos ejecutando:
export DB_URL='postgresql://postgres:pass@localhost:5432/postgres?sslmode=disable'
atlas schema apply --env local
Atlas cargará el estado deseado de la base de datos desde nuestra configuración, lo comparará con el estado actual y creará un plan de migración:
Planning migration statements (5 in total):
-- create extension "vector":
-> CREATE EXTENSION "vector" WITH SCHEMA "public" VERSION "0.8.0";
-- create "chunks" table:
-> CREATE TABLE "public"."chunks" (
"id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY,
"path" character varying NOT NULL,
"nchunk" bigint NOT NULL,
"data" text NOT NULL,
PRIMARY KEY ("id")
);
-- create "embeddings" table:
-> CREATE TABLE "public"."embeddings" (
"id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY,
"embedding" public.vector(1536) NOT NULL,
"chunk_id" bigint NOT NULL,
PRIMARY KEY ("id"),
CONSTRAINT "embeddings_chunks_embedding" FOREIGN KEY ("chunk_id") REFERENCES "public"."chunks" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
);
-- create index "embedding_embedding" to table: "embeddings":
-> CREATE INDEX "embedding_embedding" ON "public"."embeddings" USING hnsw ("embedding" vector_l2_ops);
-- create index "embeddings_chunk_id_key" to table: "embeddings":
-> CREATE UNIQUE INDEX "embeddings_chunk_id_key" ON "public"."embeddings" ("chunk_id");
-------------------------------------------
Analyzing planned statements (5 in total):
-- non-optimal columns alignment:
-- L4: Table "chunks" has 8 redundant bytes of padding per row. To reduce disk space,
the optimal order of the columns is as follows: "id", "nchunk", "path",
"data" https://atlasgo.io/lint/analyzers#PG110
-- ok (370.25µs)
-------------------------
-- 114.306667ms
-- 5 schema changes
-- 1 diagnostic
-------------------------------------------
? Approve or abort the plan:
▸ Approve and apply
Abort
Además de planificar el cambio, Atlas proporciona diagnósticos y sugerencias para optimizar el esquema. En este caso sugiere reordenar las columnas en la tabla chunks para reducir espacio en disco. Como no nos preocupa el espacio en este tutorial, podemos continuar seleccionando Approve and apply.
Para verificar que el esquema se aplicó correctamente, podemos reejecutar atlas schema apply. Atlas mostrará:
Schema is synced, no changes to be made
Configuración inicial del CLI
Con el esquema de base de datos listo, creemos nuestra aplicación CLI. Para este tutorial usaremos la librería alecthomas/kong para construir una pequeña aplicación que cargue, indexe y consulte documentos.
Primero instala la librería kong:
go get github.com/alecthomas/kong
Crea un archivo cmd/entrag/main.go con este contenido:
package main
import (
"fmt"
"os"
"github.com/alecthomas/kong"
)
// CLI holds global options and subcommands.
type CLI struct {
// DBURL is read from the environment variable DB_URL.
DBURL string `kong:"env='DB_URL',help='Database URL for the application.'"`
OpenAIKey string `kong:"env='OPENAI_KEY',help='OpenAI API key for the application.'"`
// Subcommands
Load *LoadCmd `kong:"cmd,help='Load command that accepts a path.'"`
Index *IndexCmd `kong:"cmd,help='Create embeddings for any chunks that do not have one.'"`
Ask *AskCmd `kong:"cmd,help='Ask a question about the indexed documents'"`
}
func main() {
var cli CLI
app := kong.Parse(&cli,
kong.Name("entrag"),
kong.Description("Ask questions about markdown files."),
kong.UsageOnError(),
)
if err := app.Run(&cli); err != nil {
fmt.Fprintf(os.Stderr, "Error: %s\n", err)
os.Exit(1)
}
}
Crea otro archivo cmd/entrag/rag.go con este contenido:
package main
type (
// LoadCmd loads the markdown files into the database.
LoadCmd struct {
Path string `help:"path to dir with markdown files" type:"existingdir" required:""`
}
// IndexCmd creates the embedding index on the database.
IndexCmd struct {
}
// AskCmd is another leaf command.
AskCmd struct {
// Text is the positional argument for the ask command.
Text string `kong:"arg,required,help='Text for the ask command.'"`
}
)
Verifica que la aplicación CLI funciona ejecutando:
go run ./cmd/entrag --help
Si todo está configurado correctamente, verás la ayuda de la aplicación:
Usage: entrag <command> [flags]
Ask questions about markdown files.
Flags:
-h, --help Show context-sensitive help.
--dburl=STRING Database URL for the application ($DB_URL).
--open-ai-key=STRING OpenAI API key for the application ($OPENAI_KEY).
Commands:
load --path=STRING [flags]
Load command that accepts a path.
index [flags]
Create embeddings for any chunks that do not have one.
ask <text> [flags]
Ask a question about the indexed documents
Run "entrag <command> --help" for more information on a command.
Cargar documentos en la base de datos
Necesitamos archivos Markdown para cargar. Crea un directorio data y añade archivos Markdown. Para este ejemplo, descargué el repositorio ent/ent y usé el directorio docs como fuente.
Implementemos el comando LoadCmd para cargar los archivos. Añade este código en cmd/entrag/rag.go:
const (
tokenEncoding = "cl100k_base"
chunkSize = 1000
)
// Run is the method called when the "load" command is executed.
func (cmd *LoadCmd) Run(ctx *CLI) error {
client, err := ctx.entClient()
if err != nil {
return fmt.Errorf("failed opening connection to postgres: %w", err)
}
tokTotal := 0
return filepath.WalkDir(ctx.Load.Path, func(path string, d fs.DirEntry, err error) error {
if filepath.Ext(path) == ".mdx" || filepath.Ext(path) == ".md" {
chunks := breakToChunks(path)
for i, chunk := range chunks {
tokTotal += len(chunk)
client.Chunk.Create().
SetData(chunk).
SetPath(path).
SetNchunk(i).
SaveX(context.Background())
}
}
return nil
})
}
func (c *CLI) entClient() (*ent.Client, error) {
return ent.Open("postgres", c.DBURL)
}
Este código define el método Run para LoadCmd. Lee los archivos Markdown de la ruta especificada, los divide en fragmentos de 1000 tokens y los guarda en la base de datos. Usamos entClient para crear un cliente Ent con la URL de base de datos especificada en las opciones CLI.
Para la implementación de breakToChunks, consulta el código completo
en el repositorio entrag, que se basa casi por completo en
la introducción a RAG en Go de Eli Bendersky.
Finalmente, ejecutemos el comando load para cargar los archivos markdown en la base de datos:
go run ./cmd/entrag load --path=data
Cuando termine el comando, deberías ver los fragmentos cargados en la base de datos. Para verificar, ejecuta:
docker exec -it postgres psql -U postgres -d postgres -c "SELECT COUNT(*) FROM chunks;"
Deberías ver algo similar a:
count
-------
276
(1 row)
Indexando los embeddings
Ahora que hemos cargado los documentos en la base de datos, necesitamos crear embeddings para cada fragmento. Usaremos
la API de OpenAI para generar los embeddings. Primero instala el paquete openai:
go get github.com/sashabaranov/go-openai
Si no tienes una clave API de OpenAI, puedes registrarte en la Plataforma OpenAI y generar una clave API.
Leeremos esta clave desde la variable de entorno OPENAI_KEY, así que configúrala:
export OPENAI_KEY=<your OpenAI API key>
A continuación, implementemos el comando IndexCmd para crear embeddings. Abre el archivo cmd/entrag/rag.go y
añade el siguiente código:
// Run is the method called when the "index" command is executed.
func (cmd *IndexCmd) Run(cli *CLI) error {
client, err := cli.entClient()
if err != nil {
return fmt.Errorf("failed opening connection to postgres: %w", err)
}
ctx := context.Background()
chunks := client.Chunk.Query().
Where(
chunk.Not(
chunk.HasEmbedding(),
),
).
Order(ent.Asc(chunk.FieldID)).
AllX(ctx)
for _, ch := range chunks {
log.Println("Created embedding for chunk", ch.Path, ch.Nchunk)
embedding := getEmbedding(ch.Data)
_, err := client.Embedding.Create().
SetEmbedding(pgvector.NewVector(embedding)).
SetChunk(ch).
Save(ctx)
if err != nil {
return fmt.Errorf("error creating embedding: %v", err)
}
}
return nil
}
// getEmbedding invokes the OpenAI embedding API to calculate the embedding
// for the given string. It returns the embedding.
func getEmbedding(data string) []float32 {
client := openai.NewClient(os.Getenv("OPENAI_KEY"))
queryReq := openai.EmbeddingRequest{
Input: []string{data},
Model: openai.AdaEmbeddingV2,
}
queryResponse, err := client.CreateEmbeddings(context.Background(), queryReq)
if err != nil {
log.Fatalf("Error getting embedding: %v", err)
}
return queryResponse.Data[0].Embedding
}
Hemos definido el método Run para el comando IndexCmd. Este método consulta los fragmentos sin embeddings,
genera sus embeddings usando la API de OpenAI y los guarda en la base de datos.
Finalmente, ejecutemos el comando index para crear los embeddings:
go run ./cmd/entrag index
Deberías ver registros similares a:
2025/02/13 13:04:42 Created embedding for chunk /Users/home/entr/data/md/aggregate.md 0
2025/02/13 13:04:43 Created embedding for chunk /Users/home/entr/data/md/ci.mdx 0
2025/02/13 13:04:44 Created embedding for chunk /Users/home/entr/data/md/ci.mdx 1
2025/02/13 13:04:45 Created embedding for chunk /Users/home/entr/data/md/ci.mdx 2
2025/02/13 13:04:46 Created embedding for chunk /Users/home/entr/data/md/code-gen.md 0
2025/02/13 13:04:47 Created embedding for chunk /Users/home/entr/data/md/code-gen.md 1
Haciendo preguntas
Ahora que tenemos los documentos cargados y sus embeddings, podemos implementar
el comando AskCmd para hacer preguntas. Abre cmd/entrag/rag.go y añade:
// Run is the method called when the "ask" command is executed.
func (cmd *AskCmd) Run(ctx *CLI) error {
client, err := ctx.entClient()
if err != nil {
return fmt.Errorf("failed opening connection to postgres: %w", err)
}
question := cmd.Text
emb := getEmbedding(question)
embVec := pgvector.NewVector(emb)
embs := client.Embedding.
Query().
Order(func(s *sql.Selector) {
s.OrderExpr(sql.ExprP("embedding <-> $1", embVec))
}).
WithChunk().
Limit(5).
AllX(context.Background())
b := strings.Builder{}
for _, e := range embs {
chnk := e.Edges.Chunk
b.WriteString(fmt.Sprintf("From file: %v\n", chnk.Path))
b.WriteString(chnk.Data)
}
query := fmt.Sprintf(`Use the below information from the ent docs to answer the subsequent question.
Information:
%v
Question: %v`, b.String(), question)
oac := openai.NewClient(ctx.OpenAIKey)
resp, err := oac.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
Model: openai.GPT4o,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleUser,
Content: query,
},
},
},
)
if err != nil {
return fmt.Errorf("error creating chat completion: %v", err)
}
choice := resp.Choices[0]
out, err := glamour.Render(choice.Message.Content, "dark")
fmt.Print(out)
return nil
}
Aquí es donde convergen todas las partes. Tras preparar nuestra base de datos con documentos y sus embeddings,
podemos preguntar sobre ellos. Analicemos el comando AskCmd:
emb := getEmbedding(question)
embVec := pgvector.NewVector(emb)
embs := client.Embedding.
Query().
Order(func(s *sql.Selector) {
s.OrderExpr(sql.ExprP("embedding <-> $1", embVec))
}).
WithChunk().
Limit(5).
AllX(context.Background())
Comenzamos transformando la pregunta del usuario en un vector usando la API de OpenAI. Con este vector buscamos
los embeddings más similares en nuestra base de datos. Consultamos los embeddings ordenados por similitud
usando el operador <-> de pgvector y limitamos a los 5 mejores resultados.
for _, e := range embs {
chnk := e.Edges.Chunk
b.WriteString(fmt.Sprintf("From file: %v\n", chnk.Path))
b.WriteString(chnk.Data)
}
query := fmt.Sprintf(`Use the below information from the ent docs to answer the subsequent question.
Information:
%v
Question: %v`, b.String(), question)
Preparamos la información de los 5 fragmentos principales como contexto para la pregunta. Formateamos la pregunta y el contexto en una única cadena.
oac := openai.NewClient(ctx.OpenAIKey)
resp, err := oac.CreateChatCompletion(
context.Background(),
openai.ChatCompletionRequest{
Model: openai.GPT4o,
Messages: []openai.ChatCompletionMessage{
{
Role: openai.ChatMessageRoleUser,
Content: query,
},
},
},
)
if err != nil {
return fmt.Errorf("error creating chat completion: %v", err)
}
choice := resp.Choices[0]
out, err := glamour.Render(choice.Message.Content, "dark")
fmt.Print(out)
Luego usamos la API de OpenAI para generar una respuesta. Pasamos la pregunta y contexto a la API
y recibimos una respuesta. Finalmente, formateamos la respuesta con el paquete glamour para mostrarla en terminal.
Antes de ejecutar ask, instala el paquete glamour:
go get github.com/charmbracelet/glamour
Finalmente, ejecutemos ask para hacer una pregunta sobre los documentos:
go run ./cmd/entrag ask "tl;dr What is Ent?"
Y nuestro sistema RAG responde:
Ent is an open-source entity framework (ORM) for the Go programming language. It
allows developers to define data models or graph-structures in Go code. Ent
emphasizes principles such as schema as code, a statically typed and explicit
API generated through codegen, simple queries and graph traversals, statically
typed predicates, and storage agnosticism. It supports various databases,
including MySQL, MariaDB, PostgreSQL, SQLite, and Gremlin-based graph databases,
and aims to enhance productivity in Go development.
¡Increíble! Hemos construido un sistema RAG usando Ent, Atlas y pgvector. Ahora podemos hacer preguntas sobre los documentos cargados y recibir respuestas con contexto.
Aquí tienes más preguntas y respuestas con nuestro práctico sistema RAG:
go run ./cmd/entrag ask "how to define order field in entgql"
Choose a Comparable Field: Select the field in your schema that you want to be orderable. This could be a text field, a timestamp, an integer, an enum, etc.
Annotate the Field: Use the entgql.OrderField annotation on the chosen field. The order field name should be uppercase and match its corresponding enum value in the GraphQL schema.
Update your Schema: Here’s an example of how to annotate your fields in the ent/schema :
func (Todo) Fields() []ent.Field {
return []ent.Field{
field.Text("text").
NotEmpty().
Annotations(
entgql.OrderField("TEXT"),
),
field.Time("created_at").
Default(time.Now).
Immutable().
Annotations(
entgql.OrderField("CREATED_AT"),
),
field.Enum("status").
NamedValues(
"InProgress", "IN_PROGRESS",
"Completed", "COMPLETED",
).
Default("IN_PROGRESS").
Annotations(
entgql.OrderField("STATUS"),
),
field.Int("priority").
Default(0).
Annotations(
entgql.OrderField("PRIORITY"),
),
}
}Alternatively, Multiple Order Fields: If you want to enable ordering by multiple fields, you can use the entgql.MultiOrder() annotation:
func (Todo) Annotations() []schema.Annotation {
return []schema.Annotation{
entgql.MultiOrder(),
}
}Generate GraphQL Types: Ensure that the GraphQL types match this setup in your GraphQL schema, for example:
enum OrderDirection {
ASC
DESC
}
enum TodoOrderField {
CREATED_AT
PRIORITY
STATUS
TEXT
}
input TodoOrder {
direction: OrderDirection!
field: TodoOrderField
}Add orderBy Argument to Queries: Include the orderBy parameter in your GraphQL queries to allow clients to sort items:
type Query {
todos(
after: Cursor
first: Int
before: Cursor
last: Int
orderBy: TodoOrder
): TodoConnection!
}By following these steps, you ensure that your Ent-based application can support sorting in your GraphQL API based on various fields.
go run ./cmd/entrag ask "what's the difference between privacy rules and interceptors"
- Privacy Rules:
- Purpose: Privacy rules are primarily concerned with enforcing access control policies for data queries and mutations. They determine whether an operation is allowed or denied based on predefined rules.
- Implementation: They are implemented as part of a privacy policy using the ent.Policy interface which includes EvalQuery and EvalMutation methods.
- Operation: Privacy rules evaluate whether specific conditions are met before allowing or denying access to the data. They can return decisions like privacy.Allow , privacy.Deny , or privacy.Skip to control the flow of evaluation.
- Use Case: Ideal for managing access control by ensuring that users can only perform certain operations if they meet the specified criteria.
- Interceptors:
- Purpose: Interceptors act as middleware for Ent queries, allowing modification and customization of query behaviors. They can be used to augment or modify queries during different stages of their lifecycle.
- Implementation: Implemented as interfaces or using the ent.InterceptFunc adapter. They intercept and potentially modify queries by working on the read-path.
- Operation: Interceptors modify or enhance queries, typically without the access control logic inherent in privacy rules. They provide hooks to execute custom logic pre and post query execution.
- Use Case: Suitable for generic transformations or modifications to queries, such as adding default filters, query limitations, or logging operations without focusing on access control.
In summary, while privacy rules focus on access control, interceptors are about managing and modifying the query execution process.
Para finalizar
En esta entrada de blog, hemos explorado cómo construir un sistema RAG usando Ent, Atlas y pgvector. Un agradecimiento especial a Eli Bendersky por su informativa entrada de blog y por sus excelentes escritos sobre Go a lo largo de los años!
Crédito: martinfowler.com
