embed work

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
+	Url       string
-	Port     string
+	Port      string
-	DbName   string
+	DbName    string
-	User     string
+	User      string
-	Password string
+	Password  string
-	History  HistoryConfig
+	History   HistoryConfig
-	Memory   MemoryConfig
+	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
 }

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
+}

db/embed/embed_onnx.go

@@ -0,0 +1,2 @@
+package embed
+

db/embed/embed_openai.go

@@ -0,0 +1,2 @@
+package embed
+

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(),
+		id           		UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-		name        TEXT,
+		title 		 		TEXT,
-		summary     TEXT,
+		title_source 		TEXT,
-		created_at  TIMESTAMPTZ NOT NULL DEFAULT now()
+		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 INDEX on messages using gin (tsv);
-	CREATE TABLE IF	 NOT EXISTS tool_calls (
+	CREATE INDEX on messages (conversation_id, seq);`
-		id           UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+	const createMessageChunkTable = `
-		type 		 TEXT NOT NULL,
+	CREATE TABLE IF NOT EXISTS message_chunks (
-		functionName TEXT NOT NULL,
+		id				UUID PRIMARY KEY DEFAULT gen_random_uuid(),
-		functionArgs TEXT NOT NULL,
+		conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
-		message_id 	 UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
+		message_id		UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
-		created_at   TIMESTAMPTZ NOT NULL DEFAULT now()
+		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
+	rows, err := conn.Query(`SELECT id, role, content, seq
-		FROM messages m 
+		FROM messages
-		LEFT JOIN tool_calls t
+		WHERE conversation_id = $1 ORDER BY seq`, id)
-		ON m.id = t.message_id
-		WHERE m.conversation_id = $1 ORDER BY m.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`,
+	if _, err := conn.Exec(`INSERT INTO message_chunks (conversation_id, message_id, content, content_hash)
-		message.Seq, message.Role, message.Content, conversation_id).Scan(&message.Id); err != nil {
+		VALUES ($1, $2, $3, $4)`); err != nil {
-		client.Logger.Error("Could not save message", "seq", message.Seq, "error", err)
+		client.Logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err)
-		return message.Id, errors.New("Could not save message")
 	}
 
-	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()
-}