radumaco/souvenir
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added chunking and saving embeddings

Browse Source
This commit is contained in:
Radu Macocian (admac)
2026-06-25 09:50:26 +02:00
parent 5d5351a459
commit 491510ca78
6 changed files with 116 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config/config.go
+16 -14
View File
@@ -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,
},
}
}
db/embed/chunker.go
-10
View File
@@ -1,10 +0,0 @@
package embed
type Chuner struct {
MaxChunks int
Overlap int
}
func Chunk(text string) []uint32 {
splitSentences
}
db/embed/embed.go
+62 -16
View File
@@ -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
}
db/history/history.go
+35 -14
View File
@@ -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",
go.mod
+1
View File
@@ -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
go.sum
+2
View File
@@ -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=