From 5d5351a4598bff79065b86d38a9dcf8ddcba0aa4 Mon Sep 17 00:00:00 2001 From: "Radu Macocian (admac)" Date: Tue, 23 Jun 2026 14:36:34 +0200 Subject: [PATCH] Added call to embedding openai --- config/config.go | 62 ++++++-- db/embed/chunker.go | 10 ++ db/embed/embed.go | 108 ++++++++++++++ db/history/history.go | 240 +++++++++++++----------------- hello.go | 8 +- llm/{client.go => chat_client.go} | 64 ++++---- llm/embedder.go | 137 +++++++++++++++++ ui/chat.go | 15 +- 8 files changed, 463 insertions(+), 181 deletions(-) create mode 100644 db/embed/chunker.go create mode 100644 db/embed/embed.go rename llm/{client.go => chat_client.go} (61%) create mode 100644 llm/embedder.go diff --git a/config/config.go b/config/config.go index 783ceb6..447dcec 100644 --- a/config/config.go +++ b/config/config.go @@ -5,26 +5,48 @@ import ( "errors" "io" "log/slog" + "net" + "net/url" "os" ) var ENV_PREFIX = "SOUV__" type Config struct { - Api ApiConfig - Llm LLMConfig - Logging []LogConfig - Db DbConfig + Api ApiConfig + Llm LLMConfig + Logging []LogConfig + Db DbConfig + Embedding EmbeddingConfig } 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 + ChunkSize int + ChunkOverlap int +} + +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 + Dim int } type MemoryConfig struct { @@ -114,6 +136,13 @@ func defaultConfig() *Config { Password: "", DbName: "souvenir", History: HistoryConfig{Enabled: true}, + ChunkSize: 400, + ChunkOverlap: 40, + }, + EmbeddingConfig{ + Url: "", + Key: "", + Model: "", }, } } @@ -126,13 +155,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.Embedding.Url, key: ENV_PREFIX + "embedding__url"}, + {conf: &cfg.Embedding.Key, key: ENV_PREFIX + "embedding__key"}, + {conf: &cfg.Embedding.Model, key: ENV_PREFIX + "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 +176,11 @@ func (cfg Config) validate() error { if cfg.Api.Url == "" { return errors.New("No API url configured") } + if cfg.Embedding.Model == "" { + return errors.New("No embedding model configured") + } + if cfg.Embedding.Url == "" { + return errors.New("No embedding url configured") + } return nil } diff --git a/db/embed/chunker.go b/db/embed/chunker.go new file mode 100644 index 0000000..bdccff8 --- /dev/null +++ b/db/embed/chunker.go @@ -0,0 +1,10 @@ +package embed + +type Chuner struct { + MaxChunks int + Overlap int +} + +func Chunk(text string) []uint32 { + splitSentences +} diff --git a/db/embed/embed.go b/db/embed/embed.go new file mode 100644 index 0000000..155ced1 --- /dev/null +++ b/db/embed/embed.go @@ -0,0 +1,108 @@ +package embed + +import ( + "context" + "fmt" + "log/slog" + + "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/llm" + "github.com/gopsql/pgx" + "github.com/jackc/pgx/v5/pgxpool" +) + +type EmbedStore struct { + cfg config.DbConfig + logger slog.Logger + embedder llm.Embedder + tableName string + pool *pgxpool.Pool +} + +func (e EmbedStore) TableName() string { + if e.tableName == "" { + e.tableName = fmt.Sprintf(`message_chunks_vec_%s`, e.embedder.ID()) + } + return e.tableName +} + +func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedder) (*EmbedStore, error) { + logger := *slog.Default().With("Component", "EmbedStore") + + pool, err := pgxpool.New(ctx, cfg.ConnectionString()) + if err != nil { + logger.ErrorContext(ctx, "Could not create pgxpool") + return nil, err + } + store := EmbedStore{ + cfg: cfg, + logger: logger, + embedder: embedder, + pool: pool, + } + + const createVectorTableTemplate = ` + CREATE TABLE IF NOT EXISTS %s ( + chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE, + embedding vector(%d), + content_hash text not null + )` + conn := pgx.MustOpen(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)`, store.TableName()).Scan(&exists); err != nil { + logger.Error("Error while checking table", "table", store.TableName(), "error", err) + return nil, err + } + if exists { + logger.Debug("Table found", "table", store.TableName()) + } else { + logger.Debug("Table not found. Creating", "table", store.TableName()) + script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim()) + if _, err := conn.Exec(script); err != nil { + logger.Error("Could not create table", "table", store.TableName(), "error", err) + return nil, err + } + } + return &store, nil +} + +func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string) error { + rows, err := e.pool.Query(ctx, `SELECT c.id, c.content, c.content_hash + FROM message_chunks c + LEFT JOIN $1 v + ON v.chunk_id = c.id AND v.content_hash = c.content_hash + WHERE c.conversation_id = $2 and v.chunk_id IS NULL`, + e.TableName(), conversationId) + + if err != nil { + e.logger.Error("Could not get message chunks", "conversation_id", conversationId, "error", err) + return err + } + + var pending []struct{ + id, content, contentHash string + } + for rows.Next() { + var p struct{id, content, contentHash string} + if err := rows.Scan(&p.id, &p.content, &p.contentHash); err != nil { + rows.Close() + e.logger.Error("Error reading chunk", "error", err) + return err + } + pending = append(pending, p) + } + rows.Close() + + if len(pending) == 0 { + e.logger.Debug("Nothing to embed", "conversation_id", conversationId) + return nil + } + + for start := 0; start < len(pending); start += e.cfg.ChunkSize { + e.logger.Debug("", "a", start) + } + return nil +} diff --git a/db/history/history.go b/db/history/history.go index 624c44b..f7941c5 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -1,152 +1,142 @@ package history import ( + "context" "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" + "github.com/jackc/pgx/v5/pgxpool" ) type DbClient struct { - Cfg config.DbConfig - Logger slog.Logger + cfg config.DbConfig + logger slog.Logger + pool *pgxpool.Pool } -func (client DbClient) Init() error { +func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) { + logger := *slog.Default().With("Component", "History") 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 + if err := ensureDbExists(cfg, logger); err != nil { + return nil, err + } + pool, err := pgxpool.New(ctx, cfg.ConnectionString()) + if err != nil { + logger.ErrorContext(ctx, "Could not create pgxpool") + return nil, err } - conn := pgx.MustOpen(client.connectionString()) - defer conn.Close() tables := []struct { tableName string script string }{ {"conversations", createConversationTable}, {"messages", createMessageTable}, - {"tool_calls", createToolCallTable}, + {"message_chunks", createMessageChunkTable}, } for _, table := range tables { var exists bool - if err := conn.QueryRow(`SELECT EXISTS + if err := pool.QueryRow(ctx, `SELECT EXISTS (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, "error", err) - return err + logger.Error("Error while checking table", "table", table.tableName, "error", err) + return nil, err } if exists { - client.Logger.Debug("Table found", "table", table.tableName) + 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, "error", err) - return err + logger.Debug("Table not found. Creating", "table", table.tableName) + if _, err := pool.Exec(ctx, table.script); err != nil { + logger.Error("Could not create table", "table", table.tableName, "error", err) + return nil, err } } } - return nil + return &DbClient{ + logger: logger, + cfg: cfg, + pool: pool, + }, 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) +func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) { + cl.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) + if err := cl.pool.QueryRow(ctx, "SELECT id, name, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Name, &conv.Summary); err != nil { + cl.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 := cl.pool.Query(ctx, `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) + cl.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 + 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) + cl.logger.Error("Could not get message", "error", err) } - if len(messages) == 0 || messages[len(messages) - 1].Id != msgId { + 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) + cl.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") +func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) { + cl.logger.Debug("Fetching conversations") + rows, err := cl.pool.Query(ctx, "SELECT id, name, summary FROM conversations") if err != nil { - client.Logger.Error("Could not fetch conversations", "error", err) + cl.logger.Error("Could not fetch conversations", "error", err) return []model.Conversation{}, err } defer rows.Close() @@ -156,37 +146,35 @@ func (client DbClient) GetConversations() ([]model.Conversation, error) { rows.Scan(&conv.Id, &conv.Name, &conv.Summary) conversations = append(conversations, conv) } - client.Logger.Debug("Found conversations", "count", len(conversations)) + cl.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() +func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation) (model.Conversation, error) { + cl.logger.Debug("Saving conversation", "conversation", conv) delta := 0 if conv.Id == "" { - if err := conn.QueryRow(`INSERT INTO conversations(name, summary) VALUES($1, $2) RETURNING id`, + if err := cl.pool.QueryRow(ctx, `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) + cl.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) + if err := cl.pool.QueryRow(ctx, `SELECT COALESCE(MAX(seq) + 1, 1) FROM messages WHERE conversation_id = $1`, conv.Id).Scan(&delta); err != nil { + cl.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:])) + cl.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) + for _, message := range conv.Messages[delta-1:] { + id, err := cl.saveMessage(ctx, conv.Id, message) if err != nil { - client.Logger.Error("Could not save message", "message_id", message.Id, "seq", message.Seq, "error", err) + cl.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 @@ -194,81 +182,67 @@ func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversa return conv, nil } -func (client DbClient) saveMessage(conversation_id string, message model.Message, conn db.DB) (string, error) { +func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) { var exists bool if message.Id != "" { - if err := conn.QueryRow(`SELECT EXISTS + if err := cl.pool.QueryRow(ctx, `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) + cl.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) + cl.logger.Warn("Insertting message ", "seq", message.Seq) + if err := cl.pool.QueryRow(ctx, `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 { + cl.logger.Error("Could not save message", "seq", message.Seq, "error", err) + return message.Id, errors.New("Could not save message") + } + } + + if err := cl.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM message_chunks WHERE id = $1)`, + message.Id).Scan(&exists); err != nil { + cl.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 { + cl.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 := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) + VALUES ($1, $2, $3, $4)`); err != nil { + cl.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 } -func (client DbClient) ensureDbExists() error { +func (cl DbClient) Close() { + cl.pool.Close() +} + +func ensureDbExists(cfg config.DbConfig, logger slog.Logger) error { maintenanceUrl := &url.URL{ Scheme: "postgres", - User: url.UserPassword(client.Cfg.User, client.Cfg.Password), - Host: net.JoinHostPort(client.Cfg.Url, client.Cfg.Port), + User: url.UserPassword(cfg.User, cfg.Password), + Host: net.JoinHostPort(cfg.Url, 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 { + if err := conn.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)`, cfg.DbName).Scan(&exists); err != nil { return err } if exists { - client.Logger.Debug("Database already exists", "name", client.Cfg.DbName) + logger.Debug("Database already exists", "name", cfg.DbName) return nil } - client.Logger.Info("Creating database", "name", client.Cfg.DbName) - _, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`) + logger.Info("Creating database", "name", cfg.DbName) + _, err := conn.Exec(`CREATE DATABASE "` + 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() -} diff --git a/hello.go b/hello.go index 8002993..fab7c39 100644 --- a/hello.go +++ b/hello.go @@ -1,6 +1,7 @@ package main import ( + "context" "log/slog" "os" @@ -12,6 +13,7 @@ import ( func main() { + ctx := context.Background() config, err := config.Load(".config.json") if err != nil { slog.Error("Config error:", "message", err.Error()) @@ -20,14 +22,14 @@ func main() { var logger = slog.Default().With("Compoent", "Main") logger.Debug("Config initialized") - db := history.DbClient{Cfg: config.Db, Logger: *slog.Default().With("Component", "History")}; - err = db.Init() + db, err := history.Init(ctx, config.Db) if err != nil { logger.Error("Error while initializing database", "message", err) os.Exit(1) } + defer db.Close() - p := tea.NewProgram(ui.InitialModel(*config, db)) + p := tea.NewProgram(ui.InitialModel(ctx, *config, *db)) if _, err := p.Run(); err != nil { logger.Error("Alas, there's been an error:", "message", err) os.Exit(1) diff --git a/llm/client.go b/llm/chat_client.go similarity index 61% rename from llm/client.go rename to llm/chat_client.go index 51c5513..8a03cfd 100644 --- a/llm/client.go +++ b/llm/chat_client.go @@ -16,9 +16,9 @@ import ( const COMPLETIONS_API = "/v1/chat/completions" const MODELS_API = "/v1/models" -type LLMClient struct { +type ChatClient struct { Cfg config.Config - Logger slog.Logger + logger slog.Logger } // Request Types @@ -52,22 +52,29 @@ type Usage struct { TotalTokens int `json:"total_tokens"` } -func (client LLMClient) Models() ([]string, error) { - resp, err := http.Get(client.Cfg.Api.Url + MODELS_API) +func NewLLMClient(cfg config.Config) *ChatClient { + return &ChatClient{ + Cfg: cfg, + logger: *slog.Default().With("Component", "ChatClient"), + } +} + +func (cl ChatClient) Models() ([]string, error) { + resp, err := http.Get(cl.Cfg.Api.Url + MODELS_API) if err != nil { - client.Logger.Error("There was an error during the network request", "error", err.Error()) + cl.logger.Error("There was an error during the network request", "error", err.Error()) return nil, err } defer resp.Body.Close() data, err := io.ReadAll(resp.Body) if err != nil { - client.Logger.Error("There was an error reading the response body", "error", err.Error()) + cl.logger.Error("There was an error reading the response body", "error", err.Error()) return nil, err } if resp.StatusCode != 200 && resp.StatusCode != 202 { - client.Logger.Error("Call returned non 200 status", + cl.logger.Error("Call returned non 200 status", "statusCode", resp.StatusCode, "status", resp.Status) return nil, @@ -75,18 +82,18 @@ func (client LLMClient) Models() ([]string, error) { strconv.Itoa(resp.StatusCode) + resp.Status) } - client.Logger.Debug("Received data", "data", data) + cl.logger.Debug("Received data", "data", data) - var response struct{ - Data []struct{ + var response struct { + Data []struct { Id string } } if err := json.Unmarshal(data, &response); err != nil { - client.Logger.Error("There was an error parsing the response", "error", err.Error()) + cl.logger.Error("There was an error parsing the response", "error", err.Error()) return nil, err } - client.Logger.Debug("Unmarshal response", "data", response) + cl.logger.Debug("Unmarshal response", "data", response) var ret []string for _, m := range response.Data { ret = append(ret, m.Id) @@ -94,43 +101,46 @@ func (client LLMClient) Models() ([]string, error) { return ret, nil } -func (client LLMClient) Call(query string, messages []model.Message) ([]model.Message, error) { +func (cl ChatClient) Call(query string, messages []model.Message) ([]model.Message, error) { requestBody := ChatRequest{ - Model: client.Cfg.Llm.Model, + Model: cl.Cfg.Llm.Model, Messages: messages, Tools: config.AvailableTools(), } serialized, err := json.Marshal(requestBody) if err != nil { + cl.logger.Error("There was an error marshelling the request body", + "request", requestBody, + "error", err) return []model.Message{}, err } body := bytes.NewBuffer(serialized) - req, err := http.NewRequest("POST", client.Cfg.Api.Url + COMPLETIONS_API, body) + req, err := http.NewRequest("POST", cl.Cfg.Api.Url+COMPLETIONS_API, body) if err != nil { - client.Logger.Error("There was an error creating the request", "error", err.Error()) + cl.logger.Error("There was an error creating the request", "error", err.Error()) return []model.Message{}, err } req.Header.Set("Content-Type", "application/json") - if client.Cfg.Api.Key != "" { - req.Header.Set("Authorization", "Bearer "+client.Cfg.Api.Key) + if cl.Cfg.Api.Key != "" { + req.Header.Set("Authorization", "Bearer "+cl.Cfg.Api.Key) } - client.Logger.Debug("Executing llm call", "req", req) + cl.logger.Debug("Executing llm call", "req", req) resp, err := http.DefaultClient.Do(req) if err != nil { - client.Logger.Error("There was an error during http call", "error", err.Error()) + cl.logger.Error("There was an error during http call", "error", err.Error()) return []model.Message{}, err } - client.Logger.Debug("LLM Call returned", "response", resp) defer resp.Body.Close() + cl.logger.Debug("LLM Call returned", "response", resp) data, err := io.ReadAll(resp.Body) if err != nil { - client.Logger.Error("There was an error reading the response body", "error", err.Error()) + cl.logger.Error("There was an error reading the response body", "error", err.Error()) return []model.Message{}, err } if resp.StatusCode != 200 && resp.StatusCode != 202 { - client.Logger.Error("Call returned non 200 status", + cl.logger.Error("Call returned non 200 status", "statusCode", resp.StatusCode, "status", resp.Status) return []model.Message{}, @@ -139,20 +149,20 @@ func (client LLMClient) Call(query string, messages []model.Message) ([]model.Me resp.Status) } - client.Logger.Debug("Received data", "data", data) + cl.logger.Debug("Received data", "data", data) var response ChatResponse if err := json.Unmarshal(data, &response); err != nil { - client.Logger.Error("There was an error parsing the response", "error", err.Error()) + cl.logger.Error("There was an error parsing the response", "error", err.Error()) return []model.Message{}, err } var result []model.Message if len(response.Choices) == 0 { - client.Logger.Error("The call returned an empty response") + cl.logger.Error("The call returned an empty response") return []model.Message{}, errors.New("The call returned an empty response") } for _, choice := range response.Choices { - client.Logger.Debug("Appending response message", "message", choice.Message) + cl.logger.Debug("Appending response message", "message", choice.Message) result = append(result, choice.Message) } diff --git a/llm/embedder.go b/llm/embedder.go new file mode 100644 index 0000000..a923427 --- /dev/null +++ b/llm/embedder.go @@ -0,0 +1,137 @@ +package llm + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "strings" + + config "git.estatecloud.org/radumaco/souvenir/config" +) + +const EMBEDDINGS_API = "/v1/embeddings" + +type Embedder struct { + Cfg config.EmbeddingConfig + logger slog.Logger + id string +} + +type EmbedRequest struct { + Input []string `json:"input"` + Model string `json:"model"` + EncodingFormat string `json:"encoding_format"` +} + +type EmbedResponse struct { + Data []Embedding `json:"data"` + Model string `json:"model"` + Object string `json:"object"` + Usage Usage `json:"usage"` +} + +type Embedding struct { + Embedding []float32 `json:"embeddgin"` + Index int `json:"index"` + Object string `json:"object"` +} + +func NewEmbedder(cfg config.EmbeddingConfig) *Embedder { + return &Embedder{ + Cfg: cfg, + logger: *slog.Default().With("Component", "Embedder"), + } +} + +func (e Embedder) Dim() int { + return e.Cfg.Dim +} + +func (e Embedder) ID() string { + if e.id == "" { + e.id = strings.Map(func(r rune) rune { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r <= '0' && r <= '9': + return r + default: + return '_' + } + }, e.Cfg.Model) + } + return e.id +} + +func (e Embedder) EmbedBatch(text []string) ([][]float32, error) { + requestBody := EmbedRequest{ + Model: e.Cfg.Model, + Input: text, + EncodingFormat: "float", + } + serialized, err := json.Marshal(requestBody) + if err != nil { + e.logger.Error("There was an error marshelling the request body", + "request", requestBody, + "error", err) + return [][]float32{}, err + } + body := bytes.NewBuffer(serialized) + req, err := http.NewRequest("POST", e.Cfg.Url+EMBEDDINGS_API, body) + if err != nil { + e.logger.Error("There was an error creating the request", "error", err.Error()) + return [][]float32{}, err + } + req.Header.Set("Content-Type", "application/json") + if e.Cfg.Key != "" { + req.Header.Set("Authorization", "Bearer "+e.Cfg.Key) + } + e.logger.Debug("Executing embedding call", "req", req) + resp, err := http.DefaultClient.Do(req) + if err != nil { + e.logger.Error("There was an error during http call", "error", err.Error()) + return [][]float32{}, err + } + defer resp.Body.Close() + + e.logger.Debug("LLM Call returned", "response", resp) + data, err := io.ReadAll(resp.Body) + if err != nil { + e.logger.Error("There was an error reading the response body", "error", err.Error()) + return [][]float32{}, err + } + if resp.StatusCode != 200 && resp.StatusCode != 202 { + e.logger.Error("Call returned non 200 status", + "statusCode", resp.StatusCode, + "status", resp.Status) + return [][]float32{}, + errors.New("Call returned invalid status " + + resp.Status) + } + + e.logger.Debug("Received data", "data", data) + var response EmbedResponse + if err := json.Unmarshal(data, &response); err != nil { + e.logger.Error("There was an error parsing the response", "error", err.Error()) + return [][]float32{}, err + } + + if len(response.Data) == 0 { + e.logger.Error("The call returned empty data") + return [][]float32{}, errors.New("The call returned empty data") + } + + result := make([][]float32, len(text)) + for _, embed := range response.Data { + if embed.Index < 0 || embed.Index >= len(result) { + return nil, fmt.Errorf("Embedding index out of range. Index:%d Length:%d", + embed.Index, len(result)) + } + e.logger.Debug("Appening embbeding", "embed", embed.Embedding) + result[embed.Index] = embed.Embedding + } + + return result, nil +} diff --git a/ui/chat.go b/ui/chat.go index accf27f..63cb2c7 100644 --- a/ui/chat.go +++ b/ui/chat.go @@ -1,6 +1,7 @@ package ui import ( + "context" "fmt" "log/slog" "strings" @@ -34,9 +35,10 @@ type uiModel struct { senderStyle lipgloss.Style agentStyle lipgloss.Style errorStyle lipgloss.Style - client llm.LLMClient + client llm.ChatClient errorMessage string history history.DbClient + ctx context.Context width int height int waiting bool @@ -53,7 +55,7 @@ type conversationSavedMessage struct { err error } -func InitialModel(config config.Config, client history.DbClient) uiModel { +func InitialModel(ctx context.Context, config config.Config, client history.DbClient) uiModel { ta := textarea.New() ta.Placeholder = "Send a message..." ta.SetVirtualCursor(false) @@ -85,9 +87,10 @@ func InitialModel(config config.Config, client history.DbClient) uiModel { senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")), errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")), agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")), - client: llm.LLMClient{Cfg: config, Logger: *slog.Default().With("Component", "LLM")}, + client: *llm.NewLLMClient(config), focus: focusChat, history: client, + ctx: ctx, logger: *slog.Default().With("Component", "TUI"), err: nil, } @@ -212,7 +215,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.focus = focusChat case focusHistory: m.logger.Debug("Received pickerChosenMsg", "history", msg.id) - conv, err := m.history.GetConversation(msg.id) + conv, err := m.history.GetConversation(m.ctx, msg.id) if err != nil { m.logger.Error("Could not retrieve conversation", "id", msg.id) } @@ -291,7 +294,7 @@ func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd { func (m uiModel) saveConversation() tea.Cmd { m.logger.Debug("Saving conversation", "conversation", m.conversation) return func() tea.Msg { - conv, err := m.history.SaveConversation(m.conversation) + conv, err := m.history.SaveConversation(m.ctx, m.conversation) return conversationSavedMessage{conversation: conv, err: err} } } @@ -315,7 +318,7 @@ func (m uiModel) getModels() tea.Cmd { func (m uiModel) getHistory() tea.Cmd { m.logger.Debug("Querying history") return func() tea.Msg { - resp, err := m.history.GetConversations() + resp, err := m.history.GetConversations(m.ctx) if err != nil { m.logger.Error("Could not fetch history", "error", err) }