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

embed work

Browse Source
This commit is contained in:
Radu Macocian (admac)
2026-06-23 14:36:34 +02:00
parent 8599de09de
commit 789f3b5dd3
5 changed files with 133 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
db/embed/embed.go
+51
View File
@@ -0,0 +1,51 @@
package embed
import (
"fmt"
"log/slog"
"git.estatecloud.org/radumaco/souvenir/config"
"github.com/gopsql/pgx"
)
type Embed struct {
Cfg config.DbConfig
Logger slog.Logger
}
func (e Embed) tableName() string {
return fmt.Sprintf(`message_chunks_vec_%s`, e.Cfg.Embedding.Model)
}
func (e Embed) Init() error {
const createVectorTableTemplate = `
CREATE TABLE IF NOT EXISTS %s (
chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE,
embedding vector(384),
content_hash text not null
)`
conn := pgx.MustOpen(e.Cfg.ConnectionString())
var exists bool
if err := conn.QueryRow(`SELECT EXISTS
(SELECT 1 FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name =$1)`, e.tableName()).Scan(&exists); err != nil {
e.Logger.Error("Error while checking table", "table", e.tableName(), "error", err)
return err
}
if exists {
e.Logger.Debug("Table found", "table", e.tableName())
} else {
e.Logger.Debug("Table not found. Creating", "table", e.tableName())
script := fmt.Sprintf(createVectorTableTemplate, e.tableName())
if _, err := conn.Exec(script); err != nil {
e.Logger.Error("Could not create table", "table", e.tableName(), "error", err)
return err
}
}
return nil
}
func (e Embed) EmbedConversation(id string) error {
return nil
}
db/embed/embed_onnx.go
+2
View File
@@ -0,0 +1,2 @@
package embed
db/embed/embed_openai.go
+2
View File
@@ -0,0 +1,2 @@
package embed
db/history/history.go
+45 -76
View File
@@ -22,35 +22,39 @@ type DbClient struct {
func (client DbClient) Init() error {
const createConversationTable = `
CREATE TABLE IF NOT EXISTS conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
name TEXT,
summary TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
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 (
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)
)`
const createToolCallTable = `
CREATE TABLE IF NOT EXISTS tool_calls (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
type TEXT NOT NULL,
functionName TEXT NOT NULL,
functionArgs TEXT NOT NULL,
message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
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 := client.ensureDbExists(); err != nil {
return err
}
conn := pgx.MustOpen(client.connectionString())
conn := pgx.MustOpen(client.Cfg.ConnectionString())
defer conn.Close()
tables := []struct {
tableName string
@@ -58,7 +62,7 @@ func (client DbClient) Init() error {
}{
{"conversations", createConversationTable},
{"messages", createMessageTable},
{"tool_calls", createToolCallTable},
{"message_chunks", createMessageChunkTable},
}
for _, table := range tables {
@@ -85,7 +89,7 @@ func (client DbClient) Init() error {
}
func (client DbClient) GetConversation(id string) (model.Conversation, error) {
conn := pgx.MustOpen(client.connectionString())
conn := pgx.MustOpen(client.Cfg.ConnectionString())
defer conn.Close()
client.Logger.Debug("Searching for conversation", "id", id)
@@ -94,11 +98,9 @@ func (client DbClient) GetConversation(id string) (model.Conversation, error) {
client.Logger.Error("Could not fetch conversation", "id", id, "error", err)
return conv, err
}
rows, err := conn.Query(`SELECT t.id, t.type, t.functionName, t.functionArgs, m.id, m.role, m.content, m.seq
FROM messages m
LEFT JOIN tool_calls t
ON m.id = t.message_id
WHERE m.conversation_id = $1 ORDER BY m.seq`, id)
rows, err := conn.Query(`SELECT id, role, content, seq
FROM messages
WHERE conversation_id = $1 ORDER BY seq`, id)
if err != nil {
client.Logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
return conv, err
@@ -107,11 +109,10 @@ func (client DbClient) GetConversation(id string) (model.Conversation, error) {
var messages []model.Message
for rows.Next() {
var (
tcId, tcType, tcName, tcArgs *string
msgId, role, content string
seq int
)
err = rows.Scan(&tcId, &tcType, &tcName, &tcArgs, &msgId, &role, &content, &seq)
err = rows.Scan(&msgId, &role, &content, &seq)
if err != nil {
client.Logger.Error("Could not get message", "error", err)
}
@@ -120,20 +121,6 @@ func (client DbClient) GetConversation(id string) (model.Conversation, error) {
Id: msgId, Role: role, Content: content, Seq: seq,
})
}
if tcId != nil {
deref := func(s *string) string {
if s == nil {
return ""
}
return *s
}
messages[len(messages) - 1].ToolCalls = append(messages[len(messages)-1].ToolCalls, model.ToolCall{
Id: *tcId,
Type: deref(tcType),
FunctionName: deref(tcName),
FunctionArguments: deref(tcArgs),
})
}
}
conv.Messages = messages
client.Logger.Debug("Found conversation", "conversation", conv)
@@ -142,7 +129,7 @@ func (client DbClient) GetConversation(id string) (model.Conversation, error) {
func (client DbClient) GetConversations() ([]model.Conversation, error) {
client.Logger.Debug("Fetching conversations")
conn := pgx.MustOpen(client.connectionString())
conn := pgx.MustOpen(client.Cfg.ConnectionString())
defer conn.Close()
rows, err := conn.Query("SELECT id, name, summary FROM conversations")
if err != nil {
@@ -162,7 +149,7 @@ func (client DbClient) GetConversations() ([]model.Conversation, error) {
func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversation, error) {
client.Logger.Debug("Saving conversation", "conversation", conv)
conn := pgx.MustOpen(client.connectionString())
conn := pgx.MustOpen(client.Cfg.ConnectionString())
defer conn.Close()
delta := 0
if conv.Id == "" {
@@ -205,37 +192,29 @@ func (client DbClient) saveMessage(conversation_id string, message model.Message
}
}
if exists {
client.Logger.Warn("Message found. Skipping", "id", message.Id)
client.Logger.Warn("Insertting message ", "seq", message.Seq)
if err := conn.QueryRow(`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 {
client.Logger.Error("Could not save message", "seq", message.Seq, "error", err)
return message.Id, errors.New("Could not save message")
}
}
if err := conn.QueryRow(`SELECT EXISTS (SELECT 1 FROM message_chunks WHERE id = $1)`,
message.Id).Scan(&exists); err != nil {
client.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 {
client.Logger.Debug("Message chunk already exists. Skipping")
return message.Id, nil
}
if err := conn.QueryRow(`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 {
client.Logger.Error("Could not save message", "seq", message.Seq, "error", err)
return message.Id, errors.New("Could not save message")
if _, err := conn.Exec(`INSERT INTO message_chunks (conversation_id, message_id, content, content_hash)
VALUES ($1, $2, $3, $4)`); err != nil {
client.Logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err)
}
placeholder := 0
var placeholders strings.Builder
args := []any{}
const cols = 4
for _, call := range message.ToolCalls {
placeholders.WriteString("($")
placeholders.WriteString(strconv.Itoa(placeholder + 1))
placeholders.WriteString(",$")
placeholders.WriteString(strconv.Itoa(placeholder + 2))
placeholders.WriteString(",$")
placeholders.WriteString(strconv.Itoa(placeholder + 3))
placeholders.WriteString(",$")
placeholders.WriteString(strconv.Itoa(placeholder + 4))
placeholders.WriteString(")")
args = append(args, call.Type, call.FunctionName, call.FunctionArguments, message.Id)
placeholder += cols
}
if _, err := conn.Exec(`INSERT INTO tool_calls(type, functionName, functionArgs, message_id) VALUES `+placeholders.String(),
args...); err != nil && placeholder > 0 {
client.Logger.Warn("Could not insert tool_calls", "message_id", message.Id, "error", err)
}
return message.Id, nil
}
@@ -262,13 +241,3 @@ func (client DbClient) ensureDbExists() error {
_, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`)
return err
}
func (client DbClient) connectionString() string {
u := &url.URL{
Scheme: "postgres",
User: url.UserPassword(client.Cfg.User, client.Cfg.Password),
Host: net.JoinHostPort(client.Cfg.Url, client.Cfg.Port),
Path: client.Cfg.DbName,
}
return u.String()
}