使用Ent、Atlas和pgvector在Go中构建RAG系统

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

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

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(),
	}
}

go get github.com/pgvector/pgvector-go

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"),
			),
	}
}

go mod tidy
go generate ./...

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

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

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)
	}
}

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.'"`
	}
)

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)
}

// 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
}

// 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
}

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())

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)