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

cleanup db and added embedding coroutine

Browse Source
This commit is contained in:
Radu Macocian (admac)
2026-06-26 12:04:25 +02:00
parent 9537237a85
commit b318f4602b
6 changed files with 382 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
db/embed/embed.go
+61 -19
View File
@@ -7,6 +7,7 @@ import (
"log/slog"
"git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/llm"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
@@ -32,6 +33,35 @@ func (e EmbedStore) TableName() string {
return e.tableName
}
func (e EmbedStore) GetConversationsToEmbed(ctx context.Context) ([]string, error) {
rows, err := e.pool.Query(ctx, fmt.Sprintf(
`
SELECT DISTINC c.id
FROM conversations c
JOIN message_chunks mc
ON mc.conversation_id = c.id
LEFT JOIN %s v
ON v.chunk_id = mc.id AND v.content_hash = mc.content_hash
WHERE c.updated_at < now() - interval '5 hours'
AND v.chunk_id is NULL;
`,
e.TableName()))
if err != nil {
e.logger.Error("Error while getting conversations to embed", "error", err)
}
defer rows.Close()
var ids []string
for rows.Next() {
var id string
err = rows.Scan(&id)
if err != nil {
e.logger.Error("Could not get id", "error", err)
}
ids = append(ids, id)
}
return ids, nil
}
func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedder) (*EmbedStore, error) {
logger := *slog.Default().With("Component", "EmbedStore")
@@ -47,17 +77,16 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
pool: pool,
}
const createVectorTableTemplate = `
CREATE TABLE IF NOT EXISTS %s (
chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE,
embedding vector(%d),
content_hash text not null
)`
var exists bool
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 {
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 {
logger.Error("Error while checking table", "table", store.TableName(), "error", err)
return nil, err
}
@@ -65,7 +94,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
logger.Debug("Table found", "table", store.TableName())
} else {
logger.Debug("Table not found. Creating", "table", store.TableName())
script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim())
script := fmt.Sprintf(db.CreateVectorTableTemplate, store.TableName(), store.embedder.Dim())
if _, err := pool.Exec(ctx, script); err != nil {
logger.Error("Could not create table", "table", store.TableName(), "error", err)
return nil, err
@@ -75,11 +104,17 @@ 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, fmt.Sprintf(`SELECT c.id, c.content, c.content_hash
FROM message_chunks c
rows, err := e.pool.Query(ctx, fmt.Sprintf(
`
SELECT c.id, c.content, c.content_hash
FROM message_chunks c
LEFT JOIN %s v
ON v.chunk_id = c.id AND v.content_hash = c.content_hash
WHERE c.conversation_id = $1 and v.chunk_id IS NULL`, e.TableName()),
ON v.chunk_id = c.id
AND v.content_hash = c.content_hash
WHERE c.conversation_id = $1
AND v.chunk_id IS NULL
`,
e.TableName()),
conversationId)
if err != nil {
@@ -131,12 +166,15 @@ func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string
}
func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error {
query := fmt.Sprintf(`
query := fmt.Sprintf(
`
INSERT INTO %s (chunk_id, embedding, content_hash)
VALUES ($1, $2, $3)
VALUES ($1, $2, $3)
ON CONFLICT (chunk_id) DO UPDATE
SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash`, e.TableName())
SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash
`,
e.TableName())
batch := &pgx.Batch{}
for i, chunk := range chunks {
@@ -152,3 +190,7 @@ func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32
}
return nil
}
func (e EmbedStore) Close() {
e.pool.Close()
}
db/history/history.go
+81 -54
View File
@@ -11,6 +11,7 @@ import (
"strings"
"git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/model"
"github.com/gopsql/pgx"
"github.com/jackc/pgx/v5/pgxpool"
@@ -24,36 +25,6 @@ type DbClient struct {
func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
logger := *slog.Default().With("Component", "History")
const createConversationTable = `
CREATE TABLE IF NOT EXISTS conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
title TEXT,
title_source TEXT,
summary TEXT,
summary_through_seq int default 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
const createMessageTable = `
CREATE TABLE IF NOT EXISTS messages (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
role TEXT NOT NULL,
seq INTEGER NOT NULL,
content TEXT NOT NULL,
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
tsv tsvector generated always as (to_tsvector('english', content)) stored,
UNIQUE (conversation_id, seq)
);
CREATE INDEX on messages using gin (tsv);
CREATE INDEX on messages (conversation_id, seq);`
const createMessageChunkTable = `
CREATE TABLE IF NOT EXISTS message_chunks (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
content TEXT NOT NULL,
content_hash TEXT NOT NULL
)`
if err := ensureDbExists(cfg, logger); err != nil {
return nil, err
@@ -67,17 +38,23 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
tableName string
script string
}{
{"conversations", createConversationTable},
{"messages", createMessageTable},
{"message_chunks", createMessageChunkTable},
{"conversations", db.ConversationTableScript},
{"messages", db.MessageTableScript},
{"message_chunks", db.MessageChunkTableScript},
}
for _, table := range tables {
var exists bool
if err := pool.QueryRow(ctx, `SELECT EXISTS
(SELECT 1 FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name =$1)`, table.tableName).Scan(&exists); err != nil {
if err := pool.QueryRow(ctx,
`
SELECT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name =$1
)
`,
table.tableName).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", table.tableName, "error", err)
return nil, err
}
@@ -102,13 +79,22 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) {
cl.logger.Debug("Searching for conversation", "id", id)
var conv model.Conversation
if err := cl.pool.QueryRow(ctx, "SELECT id, title, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil {
if err := cl.pool.QueryRow(ctx,
`
SELECT id, title, summary
FROM conversations
WHERE id = $1
`, id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil {
cl.logger.Error("Could not fetch conversation", "id", id, "error", err)
return conv, err
}
rows, err := cl.pool.Query(ctx, `SELECT id, role, content, seq
FROM messages
WHERE conversation_id = $1 ORDER BY seq`, id)
rows, err := cl.pool.Query(ctx,
`
SELECT id, role, content, seq
FROM messages
WHERE conversation_id = $1
ORDER BY seq
`, id)
if err != nil {
cl.logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
return conv, err
@@ -137,7 +123,11 @@ func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conver
func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) {
cl.logger.Debug("Fetching conversations")
rows, err := cl.pool.Query(ctx, "SELECT id, title, summary FROM conversations")
rows, err := cl.pool.Query(ctx,
`
SELECT id, title, summary
FROM conversations
`)
if err != nil {
cl.logger.Error("Could not fetch conversations", "error", err)
return []model.Conversation{}, err
@@ -157,19 +147,33 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation
cl.logger.Debug("Saving conversation", "conversation", conv)
delta := 0
if conv.Id == "" {
if err := cl.pool.QueryRow(ctx, `INSERT INTO conversations(title, summary) VALUES($1, $2) RETURNING id`,
conv.Title,
conv.Summary).Scan(&conv.Id); err != nil {
if err := cl.pool.QueryRow(ctx,
`
INSERT INTO conversations(title, summary)
VALUES ($1, $2)
RETURNING id
`,
conv.Title, conv.Summary).Scan(&conv.Id); err != nil {
cl.logger.Error("Could not create conversation", "title", conv.Title, "error", err)
return conv, err
}
} else {
if _, err := cl.pool.Exec(ctx, `UPDATE conversations SET title = $1, summary = $2 WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil {
if _, err := cl.pool.Exec(ctx,
`
UPDATE conversations
SET title = $1, summary = $2
WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil {
cl.logger.Error("Could not update conversation", "id", conv.Id, "error", err)
return conv, err
}
}
if err := cl.pool.QueryRow(ctx, `SELECT COALESCE(MAX(seq) + 1, 1) FROM messages WHERE conversation_id = $1`, conv.Id).Scan(&delta); err != nil {
if err := cl.pool.QueryRow(ctx,
`
SELECT COALESCE(MAX(seq) + 1, 1)
FROM messages
WHERE conversation_id = $1
`,
conv.Id).Scan(&delta); err != nil {
cl.logger.Warn("Could not get delta", "conversation_id", conv.Id, "error", err)
return conv, errors.New("Could not get delta")
}
@@ -193,16 +197,27 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation
func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) {
var exists bool
if message.Id != "" {
if err := cl.pool.QueryRow(ctx, `SELECT EXISTS
(SELECT 1 FROM messages
WHERE id = $1)`, message.Id).Scan(&exists); err != nil {
if err := cl.pool.QueryRow(ctx,
`
SELECT EXISTS (
SELECT 1
FROM messages
WHERE id = $1
)
`,
message.Id).Scan(&exists); err != nil {
cl.logger.Warn("Error while scanning for message", "id", message.Id, "error", err)
return message.Id, err
}
}
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`,
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)
return message.Id, errors.New("Could not save message")
@@ -213,7 +228,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
for _, chunk := range chunks {
hash := sha256.Sum256([]byte(chunk))
if err := cl.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM message_chunks WHERE message_id = $1 and content_hash = $2)`,
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")
@@ -223,8 +246,12 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
continue
}
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 {
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)
}
}
db/schema.go
+44
View File
@@ -0,0 +1,44 @@
package db
const MessageTableScript = `
CREATE TABLE IF NOT EXISTS messages (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
role TEXT NOT NULL,
seq INTEGER NOT NULL,
content TEXT NOT NULL,
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
tsv tsvector generated always as (to_tsvector('english', content)) stored,
UNIQUE (conversation_id, seq)
);
CREATE INDEX on messages using gin (tsv);
CREATE INDEX on messages (conversation_id, seq);
`
const ConversationTableScript = `
CREATE TABLE IF NOT EXISTS conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
title TEXT,
title_source TEXT,
summary TEXT,
summary_through_seq int default 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
const MessageChunkTableScript = `
CREATE TABLE IF NOT EXISTS message_chunks (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
content TEXT NOT NULL,
content_hash TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
const CreateVectorTableTemplate = `
CREATE TABLE IF NOT EXISTS %s (
chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE,
embedding vector(%d),
content_hash text not null
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`