Added message history saving
This commit is contained in:
Radu Macocian (admac)
+199
-20
@@ -1,49 +1,55 @@
|
||||
package history
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"git.estatecloud.org/radumaco/souvenir/config"
|
||||
"git.estatecloud.org/radumaco/souvenir/model"
|
||||
"github.com/gopsql/db"
|
||||
"github.com/gopsql/pgx"
|
||||
)
|
||||
|
||||
type DbClient struct {
|
||||
Cfg config.DbConfig
|
||||
Cfg config.DbConfig
|
||||
Logger slog.Logger
|
||||
}
|
||||
|
||||
type Conversation struct {
|
||||
Id string
|
||||
Messages []model.Message
|
||||
}
|
||||
|
||||
func (client DbClient) init() error {
|
||||
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()
|
||||
)`
|
||||
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,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
UNIQUE (conversation_id, seq)
|
||||
)`
|
||||
const createToolCallTable = `
|
||||
CREATE TABLE IF NOT EXISTS tool_calls (
|
||||
id UUID PRIMARY KEY,
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
type TEXT NOT NULL,
|
||||
functionName TEXT NOT NULL,
|
||||
functionArgs TEXT NOT NULL,
|
||||
message_id UUID,
|
||||
message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
)`
|
||||
|
||||
if err := client.ensureDbExists(); err != nil {
|
||||
return err
|
||||
}
|
||||
conn := pgx.MustOpen(client.connectionString())
|
||||
defer conn.Close()
|
||||
tables := []struct {
|
||||
@@ -61,29 +67,202 @@ func (client DbClient) init() error {
|
||||
(SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_name =$1)`, table.tableName).Scan(&exists); err != nil {
|
||||
client.Logger.Error("Error while checking table", "table", table.tableName)
|
||||
return err
|
||||
client.Logger.Error("Error while checking table", "table", table.tableName, "error", err)
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
client.Logger.Debug("Table found", "table", table.tableName)
|
||||
} else {
|
||||
client.Logger.Debug("Table not found. Creating", "table", table.tableName)
|
||||
if _, err := conn.Exec(table.script); err != nil {
|
||||
client.Logger.Error("Could not create table", "table", table.tableName)
|
||||
client.Logger.Error("Could not create table", "table", table.tableName, "error", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := conn.Exec(createToolCallTable); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := conn.Exec(createToolCallTable); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client DbClient) GetConversation(id string) (model.Conversation, error) {
|
||||
conn := pgx.MustOpen(client.connectionString())
|
||||
defer conn.Close()
|
||||
|
||||
client.Logger.Debug("Searching for conversation", "id", id)
|
||||
var conv model.Conversation
|
||||
if err := conn.QueryRow("SELECT id, name, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Name, &conv.Summary); err != nil {
|
||||
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)
|
||||
if err != nil {
|
||||
client.Logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
|
||||
return conv, err
|
||||
}
|
||||
defer rows.Close()
|
||||
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)
|
||||
if err != nil {
|
||||
client.Logger.Error("Could not get message", "error", err)
|
||||
}
|
||||
if len(messages) == 0 || messages[len(messages) - 1].Id != msgId {
|
||||
messages = append(messages, model.Message{
|
||||
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)
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
func (client DbClient) GetConversations() ([]model.Conversation, error) {
|
||||
client.Logger.Debug("Fetching conversations")
|
||||
conn := pgx.MustOpen(client.connectionString())
|
||||
defer conn.Close()
|
||||
rows, err := conn.Query("SELECT id, name, summary FROM conversations")
|
||||
if err != nil {
|
||||
client.Logger.Error("Could not fetch conversations", "error", err)
|
||||
return []model.Conversation{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
conversations := []model.Conversation{}
|
||||
for rows.Next() {
|
||||
var conv model.Conversation
|
||||
rows.Scan(&conv.Id, &conv.Name, &conv.Summary)
|
||||
conversations = append(conversations, conv)
|
||||
}
|
||||
client.Logger.Debug("Found conversations", "count", len(conversations))
|
||||
return conversations, nil
|
||||
}
|
||||
|
||||
func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversation, error) {
|
||||
client.Logger.Debug("Saving conversation", "conversation", conv)
|
||||
conn := pgx.MustOpen(client.connectionString())
|
||||
defer conn.Close()
|
||||
delta := 0
|
||||
if conv.Id == "" {
|
||||
if err := conn.QueryRow(`INSERT INTO conversations(name, summary) VALUES($1, $2) RETURNING id`,
|
||||
conv.Name,
|
||||
conv.Summary).Scan(&conv.Id); err != nil {
|
||||
client.Logger.Error("Could not create conversation", "name", conv.Name, "error", err)
|
||||
return conv, err
|
||||
}
|
||||
}
|
||||
if err := conn.QueryRow(`SELECT COALESCE(MAX(seq) + 1, 1) FROM messages WHERE conversation_id = $1`, conv.Id).Scan(&delta); err != nil {
|
||||
client.Logger.Warn("Could not get delta", "conversation_id", conv.Id, "error", err)
|
||||
return conv, errors.New("Could not get delta")
|
||||
}
|
||||
|
||||
client.Logger.Debug("Found delta.", "delta", delta, "Unsaved messages", len(conv.Messages[delta - 1:]))
|
||||
|
||||
if delta >= len(conv.Messages) {
|
||||
return conv, nil
|
||||
}
|
||||
for _, message := range conv.Messages[delta - 1:] {
|
||||
id, err := client.saveMessage(conv.Id, message, conn)
|
||||
if err != nil {
|
||||
client.Logger.Error("Could not save message", "message_id", message.Id, "seq", message.Seq, "error", err)
|
||||
return conv, errors.New("Could not save message")
|
||||
}
|
||||
message.Id = id
|
||||
}
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
func (client DbClient) saveMessage(conversation_id string, message model.Message, conn db.DB) (string, error) {
|
||||
var exists bool
|
||||
if message.Id != "" {
|
||||
if err := conn.QueryRow(`SELECT EXISTS
|
||||
(SELECT 1 FROM messages
|
||||
WHERE id = $1)`, message.Id).Scan(&exists); err != nil {
|
||||
client.Logger.Warn("Error while scanning for message", "id", message.Id, "error", err)
|
||||
return message.Id, err
|
||||
}
|
||||
}
|
||||
if exists {
|
||||
client.Logger.Warn("Message found. Skipping", "id", message.Id)
|
||||
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")
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
func (client DbClient) ensureDbExists() error {
|
||||
maintenanceUrl := &url.URL{
|
||||
Scheme: "postgres",
|
||||
User: url.UserPassword(client.Cfg.User, client.Cfg.Password),
|
||||
Host: net.JoinHostPort(client.Cfg.Url, client.Cfg.Port),
|
||||
Path: "postgres",
|
||||
}
|
||||
conn := pgx.MustOpen(maintenanceUrl.String())
|
||||
defer conn.Close()
|
||||
|
||||
var exists bool
|
||||
if err := conn.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)`, client.Cfg.DbName).Scan(&exists); err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
client.Logger.Debug("Database already exists", "name", client.Cfg.DbName)
|
||||
return nil
|
||||
}
|
||||
|
||||
client.Logger.Info("Creating database", "name", client.Cfg.DbName)
|
||||
_, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`)
|
||||
return err
|
||||
}
|
||||
|
||||
func (client DbClient) connectionString() string {
|
||||
u := &url.URL{
|
||||
Scheme: "postgres",
|
||||
|
||||
Reference in New Issue
Block a user