embed work
This commit is contained in:
+26
-1
@@ -5,6 +5,8 @@ import (
|
|||||||
"errors"
|
"errors"
|
||||||
"io"
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
|
"net"
|
||||||
|
"net/url"
|
||||||
"os"
|
"os"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -25,6 +27,23 @@ type DbConfig struct {
|
|||||||
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 {
|
type MemoryConfig struct {
|
||||||
@@ -126,13 +145,16 @@ func overrideFromEnv(cfg *Config) {
|
|||||||
{conf: &cfg.Api.Url, key: ENV_PREFIX + "api__url"},
|
{conf: &cfg.Api.Url, key: ENV_PREFIX + "api__url"},
|
||||||
{conf: &cfg.Api.Key, key: ENV_PREFIX + "api__key"},
|
{conf: &cfg.Api.Key, key: ENV_PREFIX + "api__key"},
|
||||||
{conf: &cfg.Llm.Model, key: ENV_PREFIX + "llm__model"},
|
{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.DbName, key: ENV_PREFIX + "db__dbName"},
|
||||||
{conf: &cfg.Db.Url, key: ENV_PREFIX + "db__url"},
|
{conf: &cfg.Db.Url, key: ENV_PREFIX + "db__url"},
|
||||||
{conf: &cfg.Db.User, key: ENV_PREFIX + "db__user"},
|
{conf: &cfg.Db.User, key: ENV_PREFIX + "db__user"},
|
||||||
{conf: &cfg.Db.Password, key: ENV_PREFIX + "db__password"},
|
{conf: &cfg.Db.Password, key: ENV_PREFIX + "db__password"},
|
||||||
{conf: &cfg.Db.Port, key: ENV_PREFIX + "db__port"},
|
{conf: &cfg.Db.Port, key: ENV_PREFIX + "db__port"},
|
||||||
}
|
}
|
||||||
for _, override := range(overrides) {
|
for _, override := range overrides {
|
||||||
val := os.Getenv(override.key)
|
val := os.Getenv(override.key)
|
||||||
if val != "" {
|
if val != "" {
|
||||||
*override.conf = val
|
*override.conf = val
|
||||||
@@ -144,5 +166,8 @@ func (cfg Config) validate() error {
|
|||||||
if cfg.Api.Url == "" {
|
if cfg.Api.Url == "" {
|
||||||
return errors.New("No API url configured")
|
return errors.New("No API url configured")
|
||||||
}
|
}
|
||||||
|
if cfg.Db.Embedding.Model == "" {
|
||||||
|
return errors.New("No embedding model configured")
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
package embed
|
||||||
|
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
package embed
|
||||||
|
|
||||||
+36
-67
@@ -23,8 +23,10 @@ func (client DbClient) Init() error {
|
|||||||
const createConversationTable = `
|
const createConversationTable = `
|
||||||
CREATE TABLE IF NOT EXISTS conversations (
|
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,
|
||||||
|
title_source TEXT,
|
||||||
summary TEXT,
|
summary TEXT,
|
||||||
|
summary_through_seq int default 0
|
||||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||||
)`
|
)`
|
||||||
const createMessageTable = `
|
const createMessageTable = `
|
||||||
@@ -35,22 +37,24 @@ func (client DbClient) Init() error {
|
|||||||
content TEXT NOT NULL,
|
content TEXT NOT NULL,
|
||||||
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||||
|
tsv tsvector generated always as (to_tsvector('english', content)) stored,
|
||||||
UNIQUE (conversation_id, seq)
|
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);`
|
||||||
|
const createMessageChunkTable = `
|
||||||
|
CREATE TABLE IF NOT EXISTS message_chunks (
|
||||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||||
type TEXT NOT NULL,
|
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||||
functionName TEXT NOT NULL,
|
|
||||||
functionArgs TEXT NOT NULL,
|
|
||||||
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 {
|
if err := client.ensureDbExists(); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
conn := pgx.MustOpen(client.connectionString())
|
conn := pgx.MustOpen(client.Cfg.ConnectionString())
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
tables := []struct {
|
tables := []struct {
|
||||||
tableName string
|
tableName string
|
||||||
@@ -58,7 +62,7 @@ func (client DbClient) Init() error {
|
|||||||
}{
|
}{
|
||||||
{"conversations", createConversationTable},
|
{"conversations", createConversationTable},
|
||||||
{"messages", createMessageTable},
|
{"messages", createMessageTable},
|
||||||
{"tool_calls", createToolCallTable},
|
{"message_chunks", createMessageChunkTable},
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, table := range tables {
|
for _, table := range tables {
|
||||||
@@ -85,7 +89,7 @@ func (client DbClient) Init() error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (client DbClient) GetConversation(id string) (model.Conversation, error) {
|
func (client DbClient) GetConversation(id string) (model.Conversation, error) {
|
||||||
conn := pgx.MustOpen(client.connectionString())
|
conn := pgx.MustOpen(client.Cfg.ConnectionString())
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
|
|
||||||
client.Logger.Debug("Searching for conversation", "id", id)
|
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)
|
client.Logger.Error("Could not fetch conversation", "id", id, "error", err)
|
||||||
return conv, 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 {
|
if err != nil {
|
||||||
client.Logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
|
client.Logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
|
||||||
return conv, err
|
return conv, err
|
||||||
@@ -107,11 +109,10 @@ func (client DbClient) GetConversation(id string) (model.Conversation, error) {
|
|||||||
var messages []model.Message
|
var messages []model.Message
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var (
|
var (
|
||||||
tcId, tcType, tcName, tcArgs *string
|
|
||||||
msgId, role, content string
|
msgId, role, content string
|
||||||
seq int
|
seq int
|
||||||
)
|
)
|
||||||
err = rows.Scan(&tcId, &tcType, &tcName, &tcArgs, &msgId, &role, &content, &seq)
|
err = rows.Scan(&msgId, &role, &content, &seq)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
client.Logger.Error("Could not get message", "error", err)
|
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,
|
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
|
conv.Messages = messages
|
||||||
client.Logger.Debug("Found conversation", "conversation", conv)
|
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) {
|
func (client DbClient) GetConversations() ([]model.Conversation, error) {
|
||||||
client.Logger.Debug("Fetching conversations")
|
client.Logger.Debug("Fetching conversations")
|
||||||
conn := pgx.MustOpen(client.connectionString())
|
conn := pgx.MustOpen(client.Cfg.ConnectionString())
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
rows, err := conn.Query("SELECT id, name, summary FROM conversations")
|
rows, err := conn.Query("SELECT id, name, summary FROM conversations")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -162,7 +149,7 @@ func (client DbClient) GetConversations() ([]model.Conversation, error) {
|
|||||||
|
|
||||||
func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversation, error) {
|
func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversation, error) {
|
||||||
client.Logger.Debug("Saving conversation", "conversation", conv)
|
client.Logger.Debug("Saving conversation", "conversation", conv)
|
||||||
conn := pgx.MustOpen(client.connectionString())
|
conn := pgx.MustOpen(client.Cfg.ConnectionString())
|
||||||
defer conn.Close()
|
defer conn.Close()
|
||||||
delta := 0
|
delta := 0
|
||||||
if conv.Id == "" {
|
if conv.Id == "" {
|
||||||
@@ -205,37 +192,29 @@ func (client DbClient) saveMessage(conversation_id string, message model.Message
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if exists {
|
if exists {
|
||||||
client.Logger.Warn("Message found. Skipping", "id", message.Id)
|
client.Logger.Warn("Insertting message ", "seq", message.Seq)
|
||||||
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.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 {
|
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)
|
client.Logger.Error("Could not save message", "seq", message.Seq, "error", err)
|
||||||
return message.Id, errors.New("Could not save message")
|
return message.Id, errors.New("Could not save message")
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
|
||||||
placeholder := 0
|
if err := conn.QueryRow(`SELECT EXISTS (SELECT 1 FROM message_chunks WHERE id = $1)`,
|
||||||
var placeholders strings.Builder
|
message.Id).Scan(&exists); err != nil {
|
||||||
args := []any{}
|
client.Logger.Warn("Error while scanning for message_chunk", "message id", message.Id, "error", err)
|
||||||
const cols = 4
|
return message.Id, errors.New("Could not save message chunk")
|
||||||
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(),
|
if exists {
|
||||||
args...); err != nil && placeholder > 0 {
|
client.Logger.Debug("Message chunk already exists. Skipping")
|
||||||
client.Logger.Warn("Could not insert tool_calls", "message_id", message.Id, "error", err)
|
return message.Id, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
return message.Id, nil
|
return message.Id, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -262,13 +241,3 @@ func (client DbClient) ensureDbExists() error {
|
|||||||
_, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`)
|
_, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`)
|
||||||
return err
|
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()
|
|
||||||
}
|
|
||||||
|
|||||||
Reference in New Issue
Block a user