From 491510ca789f928d6150acaf42da2d86fd14ea2d Mon Sep 17 00:00:00 2001 From: "Radu Macocian (admac)" Date: Thu, 25 Jun 2026 09:50:26 +0200 Subject: [PATCH] Added chunking and saving embeddings --- config/config.go | 30 +++++++++-------- db/embed/chunker.go | 10 ------ db/embed/embed.go | 78 ++++++++++++++++++++++++++++++++++--------- db/history/history.go | 49 +++++++++++++++++++-------- go.mod | 1 + go.sum | 2 ++ 6 files changed, 116 insertions(+), 54 deletions(-) diff --git a/config/config.go b/config/config.go index 447dcec..53c7460 100644 --- a/config/config.go +++ b/config/config.go @@ -43,10 +43,11 @@ func (config DbConfig) ConnectionString() string { } type EmbeddingConfig struct { - Url string - Key string - Model string - Dim int + Url string + Key string + Model string + Dim int + BatchSize int } type MemoryConfig struct { @@ -130,19 +131,20 @@ func defaultConfig() *Config { LLMConfig{Model: ""}, []LogConfig{{Level: slog.LevelInfo.Level(), Format: "text"}}, DbConfig{ - Url: "localhost", - User: "psql", - Port: "54321", - Password: "", - DbName: "souvenir", - History: HistoryConfig{Enabled: true}, - ChunkSize: 400, + Url: "localhost", + User: "psql", + Port: "54321", + Password: "", + DbName: "souvenir", + History: HistoryConfig{Enabled: true}, + ChunkSize: 400, ChunkOverlap: 40, }, EmbeddingConfig{ - Url: "", - Key: "", - Model: "", + Url: "", + Key: "", + Model: "", + BatchSize: 64, }, } } diff --git a/db/embed/chunker.go b/db/embed/chunker.go index bdccff8..e69de29 100644 --- a/db/embed/chunker.go +++ b/db/embed/chunker.go @@ -1,10 +0,0 @@ -package embed - -type Chuner struct { - MaxChunks int - Overlap int -} - -func Chunk(text string) []uint32 { - splitSentences -} diff --git a/db/embed/embed.go b/db/embed/embed.go index 155ced1..83c2e44 100644 --- a/db/embed/embed.go +++ b/db/embed/embed.go @@ -2,13 +2,15 @@ package embed import ( "context" + "errors" "fmt" "log/slog" "git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/llm" - "github.com/gopsql/pgx" + "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" + "github.com/pgvector/pgvector-go" ) type EmbedStore struct { @@ -16,7 +18,11 @@ type EmbedStore struct { logger slog.Logger embedder llm.Embedder tableName string - pool *pgxpool.Pool + pool *pgxpool.Pool +} + +type Chunk struct { + id, content, contentHash string } func (e EmbedStore) TableName() string { @@ -38,7 +44,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd cfg: cfg, logger: logger, embedder: embedder, - pool: pool, + pool: pool, } const createVectorTableTemplate = ` @@ -47,9 +53,8 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd embedding vector(%d), content_hash text not null )` - conn := pgx.MustOpen(cfg.ConnectionString()) var exists bool - if err := conn.QueryRow(`SELECT EXISTS + if err := pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name =$1)`, store.TableName()).Scan(&exists); err != nil { @@ -61,7 +66,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd } else { logger.Debug("Table not found. Creating", "table", store.TableName()) script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim()) - if _, err := conn.Exec(script); err != nil { + if _, err := pool.Exec(ctx, script); err != nil { logger.Error("Could not create table", "table", store.TableName(), "error", err) return nil, err } @@ -70,23 +75,22 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd } func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string) error { - rows, err := e.pool.Query(ctx, `SELECT c.id, c.content, c.content_hash + rows, err := e.pool.Query(ctx, fmt.Sprintf(`SELECT c.id, c.content, c.content_hash FROM message_chunks c - LEFT JOIN $1 v + LEFT JOIN %s v ON v.chunk_id = c.id AND v.content_hash = c.content_hash - WHERE c.conversation_id = $2 and v.chunk_id IS NULL`, - e.TableName(), conversationId) + WHERE c.conversation_id = $1 and v.chunk_id IS NULL`, e.TableName()), + conversationId) if err != nil { e.logger.Error("Could not get message chunks", "conversation_id", conversationId, "error", err) return err } - var pending []struct{ - id, content, contentHash string - } + var pending []Chunk + for rows.Next() { - var p struct{id, content, contentHash string} + var p Chunk if err := rows.Scan(&p.id, &p.content, &p.contentHash); err != nil { rows.Close() e.logger.Error("Error reading chunk", "error", err) @@ -101,8 +105,50 @@ func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string return nil } - for start := 0; start < len(pending); start += e.cfg.ChunkSize { - e.logger.Debug("", "a", start) + for start := 0; start < len(pending); start += e.embedder.Cfg.BatchSize { + batch := pending[start:min(start+e.embedder.Cfg.BatchSize, len(pending))] + + texts := make([]string, len(batch)) + for i, chunk := range batch { + texts[i] = chunk.content + } + vecs, err := e.embedder.EmbedBatch(texts) + if err != nil { + e.logger.Error("There was an error embedding batch", "conversation_id", conversationId, "error", err) + return err + } + if len(vecs) != len(batch) { + err := fmt.Sprintf("Embedder returned %d vecs for %d inputs", len(vecs), len(batch)) + e.logger.Error(err) + return errors.New(err) + } + if err := e.upsert(ctx, batch, vecs); err != nil { + e.logger.Error("Error inserting vectors in db", "error", err) + return err + } + } + return nil +} + +func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error { + query := fmt.Sprintf(` + INSERT INTO %s (chunk_id, embedding, content_hash) + VALUES ($1, $2, $3) + ON CONFLICT (chunk_id) DO UPDATE + SET embedding = EXCLUDED.embedding, + content_hash = EXCLUDED.content_hash`, e.TableName()) + + batch := &pgx.Batch{} + for i, chunk := range chunks { + batch.Queue(query, chunk.id, pgvector.NewVector(vecs[i]), chunk.contentHash) + } + result := e.pool.SendBatch(ctx, batch) + defer result.Close() + for range chunks { + if _, err := result.Exec(); err != nil { + e.logger.Error("Error inserting chunk", "error", err) + return err + } } return nil } diff --git a/db/history/history.go b/db/history/history.go index f7941c5..27215d7 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -2,10 +2,13 @@ package history import ( "context" + "crypto/sha256" + "encoding/hex" "errors" "log/slog" "net" "net/url" + "strings" "git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/model" @@ -192,8 +195,8 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess return message.Id, err } } - if exists { - cl.logger.Warn("Insertting message ", "seq", message.Seq) + if !exists { + cl.logger.Debug("Insertting message ", "seq", message.Seq) if err := cl.pool.QueryRow(ctx, `INSERT INTO messages (seq, role, content, conversation_id) VALUES ($1, $2, $3, $4) RETURNING id`, message.Seq, message.Role, message.Content, conversation_id).Scan(&message.Id); err != nil { cl.logger.Error("Could not save message", "seq", message.Seq, "error", err) @@ -201,19 +204,24 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess } } - if err := cl.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM message_chunks WHERE id = $1)`, - message.Id).Scan(&exists); err != nil { - cl.logger.Warn("Error while scanning for message_chunk", "message id", message.Id, "error", err) - return message.Id, errors.New("Could not save message chunk") - } - if exists { - cl.logger.Debug("Message chunk already exists. Skipping") - return message.Id, nil - } + chunks := Chunk(message.Content, cl.cfg.ChunkSize, cl.cfg.ChunkOverlap) + for _, chunk := range chunks { + hash := sha256.Sum256([]byte(chunk)) - if _, err := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) - VALUES ($1, $2, $3, $4)`); err != nil { - cl.logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err) + if err := cl.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM message_chunks WHERE message_id = $1 and content_hash = $2)`, + message.Id, hex.EncodeToString(hash[:])).Scan(&exists); err != nil { + cl.logger.Warn("Error while scanning for message_chunk", "message id", message.Id, "error", err) + return message.Id, errors.New("Could not save message chunk") + } + if exists { + cl.logger.Debug("Message chunk already exists. Skipping") + return message.Id, nil + } + + if _, err := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) + VALUES ($1, $2, $3, $4)`, conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil { + cl.logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err) + } } return message.Id, nil @@ -223,6 +231,19 @@ func (cl DbClient) Close() { cl.pool.Close() } +func Chunk(text string, size int, overlap int) []string { + words := strings.Fields(text) + var chunks []string + for start := 0; start < len(words); start += size - overlap { + end := min(start + size, len(words)) + chunks = append(chunks, strings.Join(words[start:end], " ")) + if end == len(words) { + return chunks + } + } + return chunks +} + func ensureDbExists(cfg config.DbConfig, logger slog.Logger) error { maintenanceUrl := &url.URL{ Scheme: "postgres", diff --git a/go.mod b/go.mod index 90b63c3..28c71fa 100644 --- a/go.mod +++ b/go.mod @@ -14,6 +14,7 @@ require ( github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect github.com/jackc/pgx/v5 v5.4.3 // indirect github.com/jackc/puddle/v2 v2.2.1 // indirect + github.com/pgvector/pgvector-go v0.4.0 // indirect github.com/sahilm/fuzzy v0.1.1 // indirect golang.org/x/crypto v0.9.0 // indirect golang.org/x/text v0.9.0 // indirect diff --git a/go.sum b/go.sum index 22d699f..3d8f4d4 100644 --- a/go.sum +++ b/go.sum @@ -55,6 +55,8 @@ github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3Ry github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/pgvector/pgvector-go v0.4.0 h1:879hQCnuix1bkfa5TQISnnK9ik4Fo+cHj2vuZSgW5v4= +github.com/pgvector/pgvector-go v0.4.0/go.mod h1:4fSXyjl1TYAIdByAql6JazKWRr2s7J0g4hcRY5cBFCk= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=