From 789f3b5dd3e157bc764e384719078bc3e579a39a Mon Sep 17 00:00:00 2001 From: "Radu Macocian (admac)" Date: Tue, 23 Jun 2026 14:36:34 +0200 Subject: [PATCH] embed work --- config/config.go | 41 ++++++++++--- db/embed/embed.go | 51 +++++++++++++++++ db/embed/embed_onnx.go | 2 + db/embed/embed_openai.go | 2 + db/history/history.go | 121 +++++++++++++++------------------------ 5 files changed, 133 insertions(+), 84 deletions(-) create mode 100644 db/embed/embed.go create mode 100644 db/embed/embed_onnx.go create mode 100644 db/embed/embed_openai.go diff --git a/config/config.go b/config/config.go index 783ceb6..b4e4a1d 100644 --- a/config/config.go +++ b/config/config.go @@ -5,6 +5,8 @@ import ( "errors" "io" "log/slog" + "net" + "net/url" "os" ) @@ -18,13 +20,30 @@ type Config struct { } type DbConfig struct { - Url string - Port string - DbName string - User string - Password string - History HistoryConfig - Memory MemoryConfig + Url string + Port string + DbName string + User string + Password string + History HistoryConfig + Memory MemoryConfig + Embedding EmbeddingConfig +} + +func (config DbConfig) ConnectionString() string { + u := &url.URL{ + Scheme: "postgres", + User: url.UserPassword(config.User, config.Password), + Host: net.JoinHostPort(config.Url, config.Port), + Path: config.DbName, + } + return u.String() +} + +type EmbeddingConfig struct { + Url string + key string + Model string } type MemoryConfig struct { @@ -126,13 +145,16 @@ func overrideFromEnv(cfg *Config) { {conf: &cfg.Api.Url, key: ENV_PREFIX + "api__url"}, {conf: &cfg.Api.Key, key: ENV_PREFIX + "api__key"}, {conf: &cfg.Llm.Model, key: ENV_PREFIX + "llm__model"}, + {conf: &cfg.Db.Embedding.Url, key: ENV_PREFIX + "db__embedding__url"}, + {conf: &cfg.Db.Embedding.key, key: ENV_PREFIX + "db__embedding__key"}, + {conf: &cfg.Db.Embedding.Model, key: ENV_PREFIX + "db__embedding__model"}, {conf: &cfg.Db.DbName, key: ENV_PREFIX + "db__dbName"}, {conf: &cfg.Db.Url, key: ENV_PREFIX + "db__url"}, {conf: &cfg.Db.User, key: ENV_PREFIX + "db__user"}, {conf: &cfg.Db.Password, key: ENV_PREFIX + "db__password"}, {conf: &cfg.Db.Port, key: ENV_PREFIX + "db__port"}, } - for _, override := range(overrides) { + for _, override := range overrides { val := os.Getenv(override.key) if val != "" { *override.conf = val @@ -144,5 +166,8 @@ func (cfg Config) validate() error { if cfg.Api.Url == "" { return errors.New("No API url configured") } + if cfg.Db.Embedding.Model == "" { + return errors.New("No embedding model configured") + } return nil } diff --git a/db/embed/embed.go b/db/embed/embed.go new file mode 100644 index 0000000..313404c --- /dev/null +++ b/db/embed/embed.go @@ -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 +} diff --git a/db/embed/embed_onnx.go b/db/embed/embed_onnx.go new file mode 100644 index 0000000..e0566e3 --- /dev/null +++ b/db/embed/embed_onnx.go @@ -0,0 +1,2 @@ +package embed + diff --git a/db/embed/embed_openai.go b/db/embed/embed_openai.go new file mode 100644 index 0000000..e0566e3 --- /dev/null +++ b/db/embed/embed_openai.go @@ -0,0 +1,2 @@ +package embed + diff --git a/db/history/history.go b/db/history/history.go index 624c44b..2c40f80 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -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() -}