Compare commits

..

6 Commits

Author SHA1 Message Date
Radu Macocian (admac) e57b8f428e Added streaming 2026-06-27 22:06:06 +02:00
Radu Macocian (admac) b318f4602b cleanup db and added embedding coroutine 2026-06-26 16:15:13 +02:00
Radu Macocian (admac) 9537237a85 Added command list 2026-06-26 10:31:02 +02:00
Radu Macocian (admac) 491510ca78 Added chunking and saving embeddings 2026-06-25 09:50:26 +02:00
Radu Macocian (admac) 5d5351a459 Added call to embedding openai 2026-06-24 16:12:42 +02:00
Radu Macocian (admac) 8599de09de Added message history saving 2026-06-18 11:42:29 +02:00
21 changed files with 2074 additions and 613 deletions
+70 -4
View File
@@ -5,16 +5,20 @@ import (
"errors" "errors"
"io" "io"
"log/slog" "log/slog"
"net"
"net/url"
"os" "os"
"strconv"
) )
var ENV_PREFIX = "SOUV__" const ENV_PREFIX = "SOUV__"
type Config struct { type Config struct {
Api ApiConfig Api ApiConfig
Llm LLMConfig Llm LLMConfig
Logging []LogConfig Logging []LogConfig
Db DbConfig Db DbConfig
Embedding EmbeddingConfig
} }
type DbConfig struct { type DbConfig struct {
@@ -25,6 +29,28 @@ type DbConfig struct {
Password string Password string
History HistoryConfig History HistoryConfig
Memory MemoryConfig 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
Timeout int
BatchSize int
Interval int
} }
type MemoryConfig struct { type MemoryConfig struct {
@@ -38,10 +64,12 @@ type HistoryConfig struct {
type ApiConfig struct { type ApiConfig struct {
Url string Url string
Key string Key string
Timeout int
} }
type LLMConfig struct { type LLMConfig struct {
Model string Model string
TitleModel string
} }
type LogConfig struct { type LogConfig struct {
@@ -93,7 +121,7 @@ func getLogTarget(target string) (io.Writer, error) {
return os.Stdout, nil return os.Stdout, nil
case "stderr": case "stderr":
return os.Stderr, nil return os.Stderr, nil
case "discrd": case "discard":
return io.Discard, nil return io.Discard, nil
default: default:
return os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644) return os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
@@ -104,7 +132,9 @@ func defaultConfig() *Config {
return &Config{ return &Config{
ApiConfig{ ApiConfig{
Url: "", Url: "",
Key: ""}, Key: "",
Timeout: 300000,
},
LLMConfig{Model: ""}, LLMConfig{Model: ""},
[]LogConfig{{Level: slog.LevelInfo.Level(), Format: "text"}}, []LogConfig{{Level: slog.LevelInfo.Level(), Format: "text"}},
DbConfig{ DbConfig{
@@ -114,6 +144,16 @@ func defaultConfig() *Config {
Password: "", Password: "",
DbName: "souvenir", DbName: "souvenir",
History: HistoryConfig{Enabled: true}, History: HistoryConfig{Enabled: true},
ChunkSize: 400,
ChunkOverlap: 40,
},
EmbeddingConfig{
Url: "",
Key: "",
Model: "",
BatchSize: 64,
Timeout: 60000,
Interval: 60,
}, },
} }
} }
@@ -126,23 +166,49 @@ 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.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.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
} }
} }
durationOverrides := []struct {
conf *int
key string
}{
{conf: &cfg.Api.Timeout, key: ENV_PREFIX + "api__timeout"},
{conf: &cfg.Embedding.Timeout, key: ENV_PREFIX + "embedding__timeout"},
{conf: &cfg.Embedding.Interval, key: ENV_PREFIX + "embedding__interval"},
}
for _, override := range durationOverrides {
val := os.Getenv(override.key)
if val != "" {
parsed, err := strconv.Atoi(val)
if err == nil {
*override.conf = parsed
}
}
}
} }
func (cfg Config) validate() error { 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.Embedding.Model == "" {
return errors.New("No embedding model configured")
}
if cfg.Embedding.Url == "" {
return errors.New("No embedding url configured")
}
return nil return nil
} }
+196
View File
@@ -0,0 +1,196 @@
package embed
import (
"context"
"errors"
"fmt"
"log/slog"
"git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/llm"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/pgvector/pgvector-go"
)
type EmbedStore struct {
cfg config.DbConfig
logger slog.Logger
embedder llm.Embedder
tableName string
pool *pgxpool.Pool
}
type Chunk struct {
id, content, contentHash string
}
func (e EmbedStore) TableName() string {
if e.tableName == "" {
e.tableName = fmt.Sprintf(`message_chunks_vec_%s`, e.embedder.ID())
}
return e.tableName
}
func (e EmbedStore) GetConversationsToEmbed(ctx context.Context) ([]string, error) {
rows, err := e.pool.Query(ctx, fmt.Sprintf(
`
SELECT DISTINC c.id
FROM conversations c
JOIN message_chunks mc
ON mc.conversation_id = c.id
LEFT JOIN %s v
ON v.chunk_id = mc.id AND v.content_hash = mc.content_hash
WHERE c.updated_at < now() - interval '5 hours'
AND v.chunk_id is NULL;
`,
e.TableName()))
if err != nil {
e.logger.Error("Error while getting conversations to embed", "error", err)
}
defer rows.Close()
var ids []string
for rows.Next() {
var id string
err = rows.Scan(&id)
if err != nil {
e.logger.Error("Could not get id", "error", err)
}
ids = append(ids, id)
}
return ids, nil
}
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,
}
var exists bool
if err := pool.QueryRow(ctx,
`
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(db.CreateVectorTableTemplate, store.TableName(), store.embedder.Dim())
if _, err := pool.Exec(ctx, 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, fmt.Sprintf(
`
SELECT c.id, c.content, c.content_hash
FROM message_chunks c
LEFT JOIN %s v
ON v.chunk_id = c.id
AND v.content_hash = c.content_hash
WHERE c.conversation_id = $1
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 []Chunk
for rows.Next() {
var p Chunk
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.embedder.Cfg.BatchSize {
batch := pending[start:min(start+e.embedder.Cfg.BatchSize, len(pending))]
texts := make([]string, len(batch))
for i, chunk := range batch {
texts[i] = chunk.content
}
vecs, err := e.embedder.EmbedBatch(texts)
if err != nil {
e.logger.Error("There was an error embedding batch", "conversation_id", conversationId, "error", err)
return err
}
if len(vecs) != len(batch) {
err := fmt.Sprintf("Embedder returned %d vecs for %d inputs", len(vecs), len(batch))
e.logger.Error(err)
return errors.New(err)
}
if err := e.upsert(ctx, batch, vecs); err != nil {
e.logger.Error("Error inserting vectors in db", "error", err)
return err
}
}
return nil
}
func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error {
query := fmt.Sprintf(
`
INSERT INTO %s (chunk_id, embedding, content_hash)
VALUES ($1, $2, $3)
ON CONFLICT (chunk_id) DO UPDATE
SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash
`,
e.TableName())
batch := &pgx.Batch{}
for i, chunk := range chunks {
batch.Queue(query, chunk.id, pgvector.NewVector(vecs[i]), chunk.contentHash)
}
result := e.pool.SendBatch(ctx, batch)
defer result.Close()
for range chunks {
if _, err := result.Exec(); err != nil {
e.logger.Error("Error inserting chunk", "error", err)
return err
}
}
return nil
}
func (e EmbedStore) Close() {
e.pool.Close()
}
+275 -63
View File
@@ -1,95 +1,307 @@
package history package history
import ( import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"log/slog" "log/slog"
"net" "net"
"net/url" "net/url"
"strings"
"git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/model" "git.estatecloud.org/radumaco/souvenir/model"
"github.com/gopsql/pgx" "github.com/gopsql/pgx"
"github.com/jackc/pgx/v5/pgxpool"
) )
type DbClient struct { type DbClient struct {
Cfg config.DbConfig cfg config.DbConfig
Logger slog.Logger logger slog.Logger
pool *pgxpool.Pool
} }
type Conversation struct { func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
Id string logger := *slog.Default().With("Component", "History")
Messages []model.Message
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
} }
func (client DbClient) init() error {
const createConversationTable = `
CREATE TABLE IF NOT EXISTS conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
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,
content TEXT NOT NULL,
conversation_id UUID,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
const createToolCallTable = `
CREATE TABLE IF NOT EXISTS tool_calls (
id UUID PRIMARY KEY,
type TEXT NOT NULL,
functionName TEXT NOT NULL,
functionArgs TEXT NOT NULL,
message_id UUID,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
conn := pgx.MustOpen(client.connectionString())
defer conn.Close()
tables := []struct { tables := []struct {
tableName string tableName string
script string script string
}{ }{
{"conversations", createConversationTable}, {"conversations", db.ConversationTableScript},
{"messages", createMessageTable}, {"messages", db.MessageTableScript},
{"tool_calls", createToolCallTable}, {"message_chunks", db.MessageChunkTableScript},
} }
for _, table := range tables { for _, table := range tables {
var exists bool var exists bool
if err := conn.QueryRow(`SELECT EXISTS if err := pool.QueryRow(ctx,
(SELECT 1 FROM information_schema.tables `
SELECT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public' WHERE table_schema = 'public'
AND table_name =$1)`, table.tableName).Scan(&exists); err != nil { AND table_name =$1
client.Logger.Error("Error while checking table", "table", table.tableName) )
`,
table.tableName).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", table.tableName, "error", err)
return nil, err
}
if exists {
logger.Debug("Table found", "table", table.tableName)
} else {
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 &DbClient{
logger: logger,
cfg: cfg,
pool: pool,
}, nil
}
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 := cl.pool.QueryRow(ctx,
`
SELECT id, title, summary
FROM conversations
WHERE id = $1
`, id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil {
cl.logger.Error("Could not fetch conversation", "id", id, "error", err)
return conv, err
}
rows, err := cl.pool.Query(ctx,
`
SELECT id, role, content, seq
FROM messages
WHERE conversation_id = $1
ORDER BY seq
`, id)
if err != nil {
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 (
msgId, role, content string
seq int
)
err = rows.Scan(&msgId, &role, &content, &seq)
if err != nil {
cl.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,
})
}
}
conv.Messages = messages
cl.logger.Debug("Found conversation", "conversation", conv)
return conv, nil
}
func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) {
cl.logger.Debug("Fetching conversations")
rows, err := cl.pool.Query(ctx,
`
SELECT id, title, summary
FROM conversations
`)
if err != nil {
cl.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.Title, &conv.Summary)
conversations = append(conversations, conv)
}
cl.logger.Debug("Found conversations", "count", len(conversations))
return conversations, nil
}
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 := cl.pool.QueryRow(ctx,
`
INSERT INTO conversations(title, summary)
VALUES ($1, $2)
RETURNING id
`,
conv.Title, conv.Summary).Scan(&conv.Id); err != nil {
cl.logger.Error("Could not create conversation", "title", conv.Title, "error", err)
return conv, err
}
} else {
if _, err := cl.pool.Exec(ctx,
`
UPDATE conversations
SET title = $1, summary = $2, updated_at = now()
WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil {
cl.logger.Error("Could not update conversation", "id", conv.Id, "error", err)
return conv, 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")
}
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 := cl.saveMessage(ctx, conv.Id, message)
if err != nil {
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
}
return conv, nil
}
func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) {
var exists bool
if message.Id != "" {
if err := cl.pool.QueryRow(ctx,
`
SELECT EXISTS (
SELECT 1
FROM messages
WHERE id = $1
)
`,
message.Id).Scan(&exists); err != nil {
cl.logger.Warn("Error while scanning for message", "id", message.Id, "error", err)
return message.Id, err
}
}
if !exists {
cl.logger.Debug("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")
}
}
chunks := Chunk(message.Content, cl.cfg.ChunkSize, cl.cfg.ChunkOverlap)
for _, chunk := range chunks {
hash := sha256.Sum256([]byte(chunk))
if err := cl.pool.QueryRow(ctx,
`
SELECT EXISTS (
SELECT 1
FROM message_chunks
WHERE message_id = $1
AND content_hash = $2
)
`,
message.Id, hex.EncodeToString(hash[:])).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")
continue
}
if _, err := cl.pool.Exec(ctx,
`
INSERT INTO message_chunks (conversation_id, message_id, content, content_hash)
VALUES ($1, $2, $3, $4);
UPDATE messages
SET updated_at = now()
WHERE id = $5
`,
conversation_id, message.Id, chunk, hex.EncodeToString(hash[:]), message.Id); err != nil {
cl.logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err)
}
}
return message.Id, nil
}
func (cl DbClient) Close() {
cl.pool.Close()
}
func Chunk(text string, size int, overlap int) []string {
if size <= overlap {
return []string{}
}
words := strings.Fields(text)
var chunks []string
for start := 0; start < len(words); start += size - overlap {
end := min(start+size, len(words))
chunks = append(chunks, strings.Join(words[start:end], " "))
if end == len(words) {
return chunks
}
}
return chunks
}
func ensureDbExists(cfg config.DbConfig, logger slog.Logger) error {
maintenanceUrl := &url.URL{
Scheme: "postgres",
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)`, cfg.DbName).Scan(&exists); err != nil {
return err return err
} }
if exists { if exists {
client.Logger.Debug("Table found", "table", table.tableName) logger.Debug("Database already exists", "name", cfg.DbName)
} 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)
return err
}
}
}
if _, err := conn.Exec(createToolCallTable); err != nil {
return err
}
if _, err := conn.Exec(createToolCallTable); err != nil {
return err
}
return nil return nil
} }
func (client DbClient) connectionString() string { logger.Info("Creating database", "name", cfg.DbName)
u := &url.URL{ _, err := conn.Exec(`CREATE DATABASE "` + cfg.DbName + `"`)
Scheme: "postgres", return err
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()
} }
+44
View File
@@ -0,0 +1,44 @@
package db
const MessageTableScript = `
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(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
tsv tsvector generated always as (to_tsvector('english', content)) stored,
UNIQUE (conversation_id, seq)
);
CREATE INDEX on messages using gin (tsv);
CREATE INDEX on messages (conversation_id, seq);
`
const ConversationTableScript = `
CREATE TABLE IF NOT EXISTS conversations (
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(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
const MessageChunkTableScript = `
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,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
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,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
)`
+3 -2
View File
@@ -6,15 +6,16 @@ require (
charm.land/bubbles/v2 v2.1.0 charm.land/bubbles/v2 v2.1.0
charm.land/bubbletea/v2 v2.0.7 charm.land/bubbletea/v2 v2.0.7
charm.land/lipgloss/v2 v2.0.3 charm.land/lipgloss/v2 v2.0.3
github.com/jackc/pgx/v5 v5.4.3
github.com/pgvector/pgvector-go v0.4.0
github.com/sahilm/fuzzy v0.1.1
) )
require ( require (
github.com/gopsql/db v1.2.1 // indirect github.com/gopsql/db v1.2.1 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
github.com/jackc/pgx/v5 v5.4.3 // indirect
github.com/jackc/puddle/v2 v2.2.1 // indirect github.com/jackc/puddle/v2 v2.2.1 // indirect
github.com/sahilm/fuzzy v0.1.1 // indirect
golang.org/x/crypto v0.9.0 // indirect golang.org/x/crypto v0.9.0 // indirect
golang.org/x/text v0.9.0 // indirect golang.org/x/text v0.9.0 // indirect
) )
+8
View File
@@ -30,6 +30,7 @@ github.com/clipperhouse/uax29/v2 v2.7.0 h1:+gs4oBZ2gPfVrKPthwbMzWZDaAFPGYK72F0NJ
github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM= github.com/clipperhouse/uax29/v2 v2.7.0/go.mod h1:EFJ2TJMRUaplDxHKj1qAEhCtQPW2tJSwu5BF98AuoVM=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/gopsql/db v1.2.1 h1:dzi13FK9OCeV2pARpmn4kSp351Uzg/Le2mEKo43H+EQ= github.com/gopsql/db v1.2.1 h1:dzi13FK9OCeV2pARpmn4kSp351Uzg/Le2mEKo43H+EQ=
github.com/gopsql/db v1.2.1/go.mod h1:plazrQDxoOfbv749q6AqZyR0wtPak19s4uw8J8pnMYA= github.com/gopsql/db v1.2.1/go.mod h1:plazrQDxoOfbv749q6AqZyR0wtPak19s4uw8J8pnMYA=
@@ -49,12 +50,17 @@ github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NB
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4= github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4=
github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw= github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw=
github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= github.com/mattn/go-runewidth v0.0.23/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
github.com/pgvector/pgvector-go v0.4.0 h1:879hQCnuix1bkfa5TQISnnK9ik4Fo+cHj2vuZSgW5v4=
github.com/pgvector/pgvector-go v0.4.0/go.mod h1:4fSXyjl1TYAIdByAql6JazKWRr2s7J0g4hcRY5cBFCk=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
@@ -68,6 +74,7 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
@@ -119,4 +126,5 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+64 -3
View File
@@ -1,26 +1,87 @@
package main package main
import ( import (
"context"
"errors"
"log/slog" "log/slog"
"os" "os"
"time"
tea "charm.land/bubbletea/v2" tea "charm.land/bubbletea/v2"
"git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db/embed"
"git.estatecloud.org/radumaco/souvenir/db/history"
"git.estatecloud.org/radumaco/souvenir/llm"
ui "git.estatecloud.org/radumaco/souvenir/ui" ui "git.estatecloud.org/radumaco/souvenir/ui"
) )
func main() { func main() {
ctx := context.Background()
config, err := config.Load(".config.json") config, err := config.Load(".config.json")
if err != nil { if err != nil {
slog.Error("Config error:", "message", err.Error()) slog.Error("Config error:", "message", err.Error())
os.Exit(1) os.Exit(1)
} }
var logger = slog.Default().With("Compoent", "Main") var logger = slog.Default().With("Component", "Main")
logger.Debug("Config initialized") logger.Debug("Config initialized")
p := tea.NewProgram(ui.InitialModel(*config))
db, err := history.Init(ctx, config.Db)
if err != nil {
logger.Error("Error while initializing database", "message", err)
os.Exit(1)
}
defer db.Close()
es, _ := setupEmbedding(config, logger)
if es != nil {
defer es.Close()
}
p := tea.NewProgram(ui.InitialModel(ctx, *config, *db))
if _, err := p.Run(); err != nil { if _, err := p.Run(); err != nil {
logger.Error("Alas, there's been an error:", "message", err) logger.Error("Alas, there's been an error:", "message", err)
os.Exit(1) os.Exit(1)
} }
} }
func setupEmbedding(cfg *config.Config, logger *slog.Logger) (*embed.EmbedStore, error) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
embedder := llm.NewEmbedder(cfg.Embedding)
embedStore, err := embed.NewEmbedStore(ctx, cfg.Db, *embedder)
if err != nil {
logger.Error("Could not initialize embedder", "error", err)
return nil, err
}
ticker := time.NewTicker(time.Duration(cfg.Embedding.Interval) * time.Minute)
go func() {
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
runEmbedding(ctx, embedStore, logger)
}
}
}()
return embedStore, nil
}
func runEmbedding(ctx context.Context, em *embed.EmbedStore, logger *slog.Logger) error {
convs, err := em.GetConversationsToEmbed(ctx)
if err != nil {
logger.Error("Could not get conversations to embed", "error", err)
return err
}
var errs []error
for _, conv := range convs {
if err := em.EmbedConversation(ctx, conv); err != nil {
logger.Error("Could not embed conversation", "conversation_id", conv, "error", err)
errs = append(errs, err)
}
}
if errs != nil {
return errors.Join(errs...)
}
return nil
}
+254
View File
@@ -0,0 +1,254 @@
package llm
import (
"bytes"
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"strconv"
"strings"
"time"
config "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/model"
)
const COMPLETIONS_API = "/v1/chat/completions"
const MODELS_API = "/v1/models"
type ChatClient struct {
Cfg config.Config
logger slog.Logger
client *http.Client
}
// Request Types
type ChatRequest struct {
Model string `json:"model"`
Messages []model.Message `json:"messages"`
Tools []config.Tool `json:"tools,omitempty"`
Stream bool `json:"stream,omitempty"`
Kwargs Kwargs `json:"chat_template_kwargs"`
}
type Kwargs struct {
Thinking bool `json:"enable_thinking"`
}
type Reasoning struct {
Effort string `json:"effort"`
}
// Response types //
type ChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []Choice `json:"choices"`
Usage Usage `json:"usage"`
}
type RenameResponse struct {
Title string `json:"title"`
Summary string `json:"summary"`
}
type Choice struct {
Index int `json:"index"`
Message model.Message `json:"message"`
FinishReason string `json:"finish_reason"`
}
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}
func NewLLMClient(cfg config.Config) *ChatClient {
return &ChatClient{
Cfg: cfg,
logger: *slog.Default().With("Component", "ChatClient"),
client: &http.Client{Timeout: time.Duration(cfg.Api.Timeout) * time.Millisecond},
}
}
func (cl ChatClient) Models() ([]string, error) {
req, err := http.NewRequest("GET", cl.Cfg.Api.Url+MODELS_API, nil)
if err != nil {
cl.logger.Error("There was an error creating the request", "error", err.Error())
return nil, err
}
req.Header.Set("Content-Type", "application/json")
if cl.Cfg.Api.Key != "" {
req.Header.Set("Authorization", "Bearer "+cl.Cfg.Api.Key)
}
resp, err := cl.client.Do(req)
if err != nil {
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 {
cl.logger.Error("There was an error reading the response body", "error", err.Error())
return nil, err
}
if resp.StatusCode != 200 && resp.StatusCode != 202 {
cl.logger.Error("Call returned non 200 status",
"statusCode", resp.StatusCode,
"status", resp.Status)
return nil,
errors.New("Call returned invalid status " +
strconv.Itoa(resp.StatusCode) +
resp.Status)
}
cl.logger.Debug("Received data", "data", data)
var response struct {
Data []struct {
Id string
}
}
if err := json.Unmarshal(data, &response); err != nil {
cl.logger.Error("There was an error parsing the response", "error", err.Error())
return nil, err
}
cl.logger.Debug("Unmarshal response", "data", response)
var ret []string
for _, m := range response.Data {
ret = append(ret, m.Id)
}
return ret, nil
}
func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) {
llm := cl.Cfg.Llm.TitleModel
if llm == "" {
llm = cl.Cfg.Llm.Model
}
request := ChatRequest{
Model: llm,
Messages: append([]model.Message{{
Role: "system",
Content: "Reply with ONLY JSON with the form: {\"title\":<title>, \"summary\":<summary>}. Title <= 10 words. Summary 2-3 sentances.",
}}, messages...),
Kwargs: Kwargs{Thinking: false},
}
request.Messages = append(request.Messages, model.Message{
Role: "user",
Content: "Give me a title and description for this conversation. Reply with just a json and nothing else. The format is exactly {\"title\":<title>, \"summary\":<summary>}. Add nothing. Change nothing",
})
cl.logger.Debug("Calling rename", "request", request)
out, err := cl.Call(request)
if err != nil {
cl.logger.Error("Error while asking for rename", "error", err)
return RenameResponse{}, err
}
var meta RenameResponse
cl.logger.Debug("Received rename answer", "content", out[0].Content)
if idx := strings.IndexByte(out[0].Content, '{'); idx >= 0 {
if lidx := strings.LastIndexByte(out[0].Content, '}'); lidx <= len(out[0].Content) {
trimmed := out[0].Content[idx:lidx + 1]
cl.logger.Debug("Trimmed answer", "content", trimmed)
if err := json.Unmarshal([]byte(trimmed), &meta); err != nil {
cl.logger.Error("Could not parse the rename response", "content", out[0].Content)
return meta, err
}
return meta, nil
}
}
cl.logger.Error("Response is not standard json", "content", out[0].Content)
return meta, errors.New("Response is not standard json")
}
func (cl ChatClient) Query(messages []model.Message) ([]model.Message, error) {
requestBody := ChatRequest{
Model: cl.Cfg.Llm.Model,
Messages: messages,
Tools: config.AvailableTools(),
Stream: true,
Kwargs: Kwargs{Thinking: false},
}
return cl.Call(requestBody)
}
func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) {
resp, err := cl.CallApi(requestBody)
if err != nil {
cl.logger.Error("There was an error during the API call", "error", err)
return nil, err
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
cl.logger.Error("There was an error reading the response body", "error", err.Error())
return nil, err
}
cl.logger.Debug("Received data", "data", data)
var response ChatResponse
if err := json.Unmarshal(data, &response); err != nil {
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 {
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 {
cl.logger.Debug("Appending response message", "message", choice.Message)
result = append(result, choice.Message)
}
return result, nil
}
func (cl ChatClient) CallApi(requestBody ChatRequest) (*http.Response, error) {
serialized, err := json.Marshal(requestBody)
if err != nil {
cl.logger.Error("There was an error marshelling the request body",
"request", requestBody,
"error", err)
return nil, err
}
body := bytes.NewBuffer(serialized)
req, err := http.NewRequest("POST", cl.Cfg.Api.Url+COMPLETIONS_API, body)
if err != nil {
cl.logger.Error("There was an error creating the request", "error", err.Error())
return nil, err
}
req.Header.Set("Content-Type", "application/json")
if cl.Cfg.Api.Key != "" {
req.Header.Set("Authorization", "Bearer "+cl.Cfg.Api.Key)
}
cl.logger.Debug("Executing llm call")
resp, err := cl.client.Do(req)
if err != nil {
cl.logger.Error("There was an error during http call", "error", err.Error())
return nil, err
}
cl.logger.Debug("LLM Call returned", "response", resp.Status)
if resp.StatusCode != 200 && resp.StatusCode != 202 {
cl.logger.Error("Call returned non 200 status",
"statusCode", resp.StatusCode,
"status", resp.Status)
return nil,
errors.New("Call returned invalid status " +
strconv.Itoa(resp.StatusCode) +
resp.Status)
}
return resp, nil
}
-160
View File
@@ -1,160 +0,0 @@
package llm
import (
"bytes"
"encoding/json"
"errors"
"io"
"log/slog"
"net/http"
"strconv"
config "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/model"
)
const COMPLETIONS_API = "/v1/chat/completions"
const MODELS_API = "/v1/models"
type LLMClient struct {
Cfg config.Config
Logger slog.Logger
}
// Request Types
type ChatRequest struct {
Model string `json:"model"`
Messages []model.Message `json:"messages"`
Tools []config.Tool `json:"tools,omitempty"`
}
// Response types //
type ChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []Choice `json:"choices"`
Usage Usage `json:"usage"`
}
type Choice struct {
Index int `json:"index"`
Message model.Message `json:"message"`
FinishReason string `json:"finish_reason"`
}
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}
func (client LLMClient) Models() ([]string, error) {
resp, err := http.Get(client.Cfg.Api.Url + MODELS_API)
if err != nil {
client.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())
return nil, err
}
if resp.StatusCode != 200 && resp.StatusCode != 202 {
client.Logger.Error("Call returned non 200 status",
"statusCode", resp.StatusCode,
"status", resp.Status)
return nil,
errors.New("Call returned invalid status " +
strconv.Itoa(resp.StatusCode) +
resp.Status)
}
client.Logger.Debug("Received data", "data", data)
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())
return nil, err
}
client.Logger.Debug("Unmarshal response", "data", response)
var ret []string
for _, m := range response.Data {
ret = append(ret, m.Id)
}
return ret, nil
}
func (client LLMClient) Call(query string, messages []model.Message) ([]model.Message, error) {
requestBody := ChatRequest{
Model: client.Cfg.Llm.Model,
Messages: messages,
Tools: config.AvailableTools(),
}
serialized, err := json.Marshal(requestBody)
if err != nil {
return []model.Message{}, err
}
body := bytes.NewBuffer(serialized)
req, err := http.NewRequest("POST", client.Cfg.Api.Url + COMPLETIONS_API, body)
if err != nil {
client.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)
}
client.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())
return []model.Message{}, err
}
client.Logger.Debug("LLM Call returned", "response", resp)
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())
return []model.Message{}, err
}
if resp.StatusCode != 200 && resp.StatusCode != 202 {
client.Logger.Error("Call returned non 200 status",
"statusCode", resp.StatusCode,
"status", resp.Status)
return []model.Message{},
errors.New("Call returned invalid status " +
strconv.Itoa(resp.StatusCode) +
resp.Status)
}
client.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())
return []model.Message{}, err
}
var result []model.Message
if len(response.Choices) == 0 {
client.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.Error("Appending response message", "message", choice.Message)
result = append(result, choice.Message)
}
return result, nil
}
+140
View File
@@ -0,0 +1,140 @@
package llm
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strings"
"time"
config "git.estatecloud.org/radumaco/souvenir/config"
)
const EMBEDDINGS_API = "/v1/embeddings"
type Embedder struct {
Cfg config.EmbeddingConfig
logger slog.Logger
id string
client *http.Client
}
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:"embedding"`
Index int `json:"index"`
Object string `json:"object"`
}
func NewEmbedder(cfg config.EmbeddingConfig) *Embedder {
return &Embedder{
Cfg: cfg,
logger: *slog.Default().With("Component", "Embedder"),
client: &http.Client{Timeout: time.Duration(cfg.Timeout) * time.Millisecond},
}
}
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 := e.client.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
}
+146
View File
@@ -0,0 +1,146 @@
package llm
import (
"bufio"
"encoding/json"
"net/http"
"strings"
config "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/model"
)
type StreamEvent struct {
Reasoning string
Content string
Messages []model.Message
Done bool
Err error
}
type ChatStreamResponse struct {
Choices []StreamChoice `json:"choices"`
}
type StreamChoice struct {
Delta StreamDelta `json:"delta"`
FinishReason *string `json:"finish_reason"`
}
type StreamDelta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
ToolCalls []ToolCallDelta `json:"tool_calls,omitempty"`
}
type ToolCallDelta struct {
Index int `json:"index"`
Id string `json:"id,omitempty"`
Type string `json:"type,omitempty"`
Function struct {
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
} `json:"function"`
}
func (cl ChatClient) QueryStream(messages []model.Message) (<-chan StreamEvent, error) {
requestBody := ChatRequest{
Model: cl.Cfg.Llm.Model,
Messages: messages,
Tools: config.AvailableTools(),
Stream: true,
Kwargs: Kwargs{Thinking: false},
}
return cl.CallStream(requestBody)
}
func (cl ChatClient) CallStream(requestBody ChatRequest) (<-chan StreamEvent, error) {
resp, err := cl.CallApi(requestBody)
if err != nil {
cl.logger.Error("There was an error during the API call", "error", err)
return nil, err
}
events := make(chan StreamEvent)
go cl.readStream(resp, events)
return events, nil
}
func (cl ChatClient) readStream(resp *http.Response, events chan<- StreamEvent) {
cl.logger.Debug("Started readStrea")
defer resp.Body.Close()
defer close(events)
scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
var content strings.Builder
toolAcc := map[int]*model.ToolCall{}
var order []int
for scanner.Scan() {
line := scanner.Text()
cl.logger.Debug("Stream got new line", "line", line)
if !strings.HasPrefix(line, "data:") {
continue
}
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload == "[DONE]" {
cl.logger.Debug("Stream done")
break
}
var chunk ChatStreamResponse
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
cl.logger.Error("Error parsing stream chunk", "payload", payload, "error", err)
events <- StreamEvent{Err: err}
return
}
if len(chunk.Choices) == 0 {
continue
}
delta := chunk.Choices[0].Delta
if delta.ReasoningContent != "" {
cl.logger.Debug("Stream got reasoning content", "content", delta.ReasoningContent)
events <- StreamEvent{Reasoning: delta.ReasoningContent}
}
if delta.Content != "" {
cl.logger.Debug("Stream got content", "content", delta.Content)
content.WriteString(delta.Content)
events <- StreamEvent{Content: delta.Content}
}
for _, tc := range delta.ToolCalls {
acc, ok := toolAcc[tc.Index]
if !ok {
acc = &model.ToolCall{}
toolAcc[tc.Index] = acc
order = append(order, tc.Index)
}
if tc.Id != "" {
acc.Id = tc.Id
}
if tc.Function.Name != "" {
acc.FunctionName = tc.Function.Name
}
acc.FunctionArguments += tc.Function.Arguments
}
}
if err := scanner.Err(); err != nil {
cl.logger.Error("error reading stream", "error", err)
events <- StreamEvent{Err: err}
return
}
final := model.Message{Role: "assistant", Content: content.String()}
for _, idx := range order {
final.ToolCalls = append(final.ToolCalls, *toolAcc[idx])
}
events <- StreamEvent{Done: true, Messages: []model.Message{final}}
}
-14
View File
@@ -1,14 +0,0 @@
package model
type Message struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
FunctionName string `json:"function"`
FunctionArguments string `json:"arguments"`
}
+8
View File
@@ -0,0 +1,8 @@
package model
type Conversation struct {
Id string
Title string
Summary string
Messages []Message
}
+9
View File
@@ -0,0 +1,9 @@
package model
type Message struct {
Id string `json:"-"`
Seq int `json:"-"`
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
+8
View File
@@ -0,0 +1,8 @@
package model
type ToolCall struct {
Id string `json:"-"`
Type string `json:"type"`
FunctionName string `json:"function"`
FunctionArguments string `json:"arguments"`
}
-280
View File
@@ -1,280 +0,0 @@
package ui
import (
"fmt"
"log/slog"
"strings"
"charm.land/bubbles/v2/cursor"
"charm.land/bubbles/v2/textarea"
"charm.land/bubbles/v2/viewport"
tea "charm.land/bubbletea/v2"
"charm.land/lipgloss/v2"
"git.estatecloud.org/radumaco/souvenir/config"
llm "git.estatecloud.org/radumaco/souvenir/llm"
"git.estatecloud.org/radumaco/souvenir/model"
)
type focusState int
const (
focusChat focusState = iota
focusOverlay
)
type uiModel struct {
logger slog.Logger
viewport viewport.Model
messages []model.Message
textarea textarea.Model
focus focusState
picker modelPicker
senderStyle lipgloss.Style
agentStyle lipgloss.Style
errorStyle lipgloss.Style
client llm.LLMClient
width int
height int
waiting bool
err error
}
type agentResponseMessage struct {
response []model.Message
err error
}
type modelsResponseMessage struct {
models []string
err error
}
func InitialModel(config config.Config) uiModel {
ta := textarea.New()
ta.Placeholder = "Send a message..."
ta.SetVirtualCursor(false)
ta.Focus()
ta.Prompt = "⦀ "
ta.CharLimit = 280
ta.SetWidth(30)
ta.SetHeight(3)
s := ta.Styles()
s.Focused.CursorLine = lipgloss.NewStyle()
ta.SetStyles(s)
ta.ShowLineNumbers = false
vp := viewport.New(viewport.WithWidth(30), viewport.WithHeight(5))
vp.SetContent(renderLanding(vp.Width(), vp.Height()))
vp.KeyMap.Left.SetEnabled(false)
vp.KeyMap.Right.SetEnabled(false)
ta.KeyMap.InsertNewline.SetEnabled(false)
model := uiModel{
textarea: ta,
messages: []model.Message{},
viewport: vp,
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")},
focus: focusChat,
logger: *slog.Default().With("Component", "TUI"),
err: nil,
}
model.picker = newModelPicker(&model.focus)
return model
}
func (m uiModel) Init() tea.Cmd {
return textarea.Blink
}
func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.logger.Debug("Received message", "msg", msg)
switch msg := msg.(type) {
case tea.WindowSizeMsg:
m.viewport.SetWidth(msg.Width * 80 / 100)
m.textarea.SetWidth(msg.Width * 80 / 100)
m.viewport.SetHeight(msg.Height - lipgloss.Height(m.textarea.View()) - 1)
m.width = msg.Width
m.height = msg.Height
m.picker.SetSize(msg.Width, msg.Height)
var messages = m.renderMessages()
if len(m.messages) > 0 {
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
} else {
m.viewport.SetContent(renderLanding(m.viewport.Width(), m.viewport.Height()))
}
m.logger.Debug("New sizes",
"vp width", m.viewport.Width(), "vp height", m.viewport.Height(),
"ta width", m.textarea.Width(), "ta height", m.textarea.Height())
m.viewport.GotoBottom()
case tea.KeyPressMsg:
if m.focus == focusOverlay {
picker, cmd := m.picker.Update(msg)
m.picker = picker
return m, cmd
}
switch msg.String() {
case "esc":
if m.focus == focusOverlay {
m.focus = focusChat
}
return m, nil
case "ctrl+c":
fmt.Println(m.textarea.Value())
return m, tea.Quit
case "shift+enter", "ctrl+j", "alt+enter":
m.textarea.InsertRune('\n')
return m, nil
case "enter":
input := m.textarea.Value()
m.messages = append(m.messages, model.Message{
Content: input,
Role: "user",
})
m.logger.Debug("Received enter. Creating new message", "message", input)
switch input {
case "/models":
m.logger.Debug("Received models sequence")
m.focus = focusOverlay
return m, m.getModels()
case "/exit":
m.logger.Debug("Received exit sequence. Quiting")
fmt.Println(m.textarea.Value())
return m, tea.Quit
default:
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.textarea.Reset()
m.viewport.GotoBottom()
m.waiting = true
return m, m.callAgent(input, m.messages)
}
default:
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
return m, cmd
}
case cursor.BlinkMsg:
m.logger.Debug("Received BlinkMsg")
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
return m, cmd
case agentResponseMessage:
m.logger.Debug("Received agentResponseMessage")
m.waiting = false
resp := msg.response
if msg.err != nil {
m.logger.Error("Message is error", "error", msg.err)
resp = []model.Message{{
Role: "Error",
Content: msg.err.Error(),
}}
}
for _, r := range resp {
m.logger.Debug("Appending response message", "message", r)
m.messages = append(m.messages, r)
}
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.viewport.GotoBottom()
return m, nil
case modelChosenMsg:
m.logger.Debug("Received modelChosenMsg", "model", msg.id)
m.client.Cfg.Llm.Model = msg.id
m.waiting = false
m.focus = focusChat
}
if m.focus == focusOverlay {
picker, cmd := m.picker.Update(msg)
m.picker = picker
return m, cmd
}
return m, nil
}
func (m uiModel) View() tea.View {
if m.focus == focusOverlay {
return tea.NewView(lipgloss.Place(m.width, m.height,
lipgloss.Center, lipgloss.Center,
m.picker.View()))
}
ui := lipgloss.JoinVertical(
lipgloss.Left,
m.viewport.View(),
"model: " + m.client.Cfg.Llm.Model + "; url: " + m.client.Cfg.Api.Url,
m.textarea.View())
c := m.textarea.Cursor()
if c != nil {
c.Y += lipgloss.Height(m.viewport.View()) + 1
gap := m.width - lipgloss.Width(ui)
c.X += max(gap/2, 0)
}
view := tea.NewView(lipgloss.PlaceHorizontal(m.width, lipgloss.Center, ui))
view.Cursor = c
view.AltScreen = true
return view
}
func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd {
m.logger.Debug("Calling agent", "query", input, "messages", messages)
return func() tea.Msg {
resp, err := m.client.Call(input, messages)
return agentResponseMessage{response: resp, err: err}
}
}
func (m uiModel) getModels() tea.Cmd {
m.logger.Debug("Querying models")
return func() tea.Msg {
resp, _ := m.client.Models()
m.logger.Debug("Received models from llm", "models", resp)
return modelsLoadedMsg{models: resp}
}
}
func (m uiModel) renderMessages() string {
var result strings.Builder
var width = m.viewport.Width()
for _, message := range m.messages {
switch message.Role {
case "user":
var line strings.Builder
line.WriteString(m.senderStyle.Render("You: "))
line.WriteString(message.Content)
result.WriteString(lipgloss.PlaceHorizontal(width, lipgloss.Right, line.String()))
case "assistant":
result.WriteString(m.agentStyle.Render("Agent: "))
result.WriteString(message.Content)
case "Error":
result.WriteString(m.errorStyle.Render("Error:"))
result.WriteString(message.Content)
}
result.WriteString("\n")
var sepStyle lipgloss.Style = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
MarginTop(1).
MarginBottom(1)
sep := sepStyle.Render(strings.Repeat("┈", width))
result.WriteString(sep)
result.WriteString("\n")
}
return result.String()
}
+167
View File
@@ -0,0 +1,167 @@
package ui
import (
"fmt"
"io"
"strings"
"charm.land/bubbles/v2/list"
tea "charm.land/bubbletea/v2"
"charm.land/lipgloss/v2"
"github.com/sahilm/fuzzy"
)
const commandDropdownHeight = 8
var (
dropdownStyle = lipgloss.NewStyle().
Background(lipgloss.Color("236")).
Padding(0, 1)
inlineTitleStyle = lipgloss.NewStyle().Background(lipgloss.Color("236")).Foreground(lipgloss.Color("86")).Bold(true)
inlineTitleSelStyle = lipgloss.NewStyle().Background(lipgloss.Color("60")).Foreground(lipgloss.Color("86")).Bold(true)
inlineDescStyle = lipgloss.NewStyle().Background(lipgloss.Color("236")).Foreground(lipgloss.Color("245"))
inlineDescSelStyle = lipgloss.NewStyle().Background(lipgloss.Color("60")).Foreground(lipgloss.Color("255"))
inlinePadStyle = lipgloss.NewStyle().Background(lipgloss.Color("236"))
inlinePadSelStyle = lipgloss.NewStyle().Background(lipgloss.Color("60"))
)
type inlineDelegate struct{}
func (inlineDelegate) Height() int { return 1 }
func (inlineDelegate) Spacing() int { return 0 }
func (inlineDelegate) Update(_ tea.Msg, _ *list.Model) tea.Cmd { return nil }
func (inlineDelegate) Render(w io.Writer, m list.Model, index int, item list.Item) {
c, ok := item.(command)
if !ok {
return
}
if m.Width() <= 0 {
return
}
title := "/" + c.name
desc := c.description
titleStyle := inlineTitleStyle
descStyle := inlineDescStyle
padStyle := inlinePadStyle
if index == m.Index() {
titleStyle = inlineTitleSelStyle
descStyle = inlineDescSelStyle
padStyle = inlinePadSelStyle
}
renderedTitle := titleStyle.Render(title)
renderedDesc := descStyle.Render(desc)
pad := max(m.Width()-1-lipgloss.Width(renderedTitle)-lipgloss.Width(renderedDesc), 1)
fmt.Fprintf(w, "%s%s%s%s", renderedTitle, padStyle.Render(" "), renderedDesc, padStyle.Render(strings.Repeat(" ", pad)))
}
type command struct {
name string
description string
handler func(m uiModel) (uiModel, tea.Cmd)
}
func (c command) FilterValue() string { return c.name }
func (c command) Title() string { return "/" + c.name }
func (c command) Description() string { return c.description }
type commandList struct {
list list.Model
open bool
available []command
}
func newCommandList(width int) commandList {
l := list.New(nil, inlineDelegate{}, width, 1)
l.Title = "Commands"
l.SetShowFilter(false)
l.SetFilteringEnabled(false)
l.SetShowStatusBar(false)
l.SetShowPagination(false)
l.SetShowHelp(false)
l.DisableQuitKeybindings()
return commandList{list: l}
}
func (m uiModel) buildCommands() []command {
return []command{
{
name: "exit",
description: "Quit the application",
handler: func(m uiModel) (uiModel, tea.Cmd) {
m.logger.Debug("Received exit sequence. Quiting")
return m, tea.Quit
},
},
{
name: "history",
description: "Open conversation history",
handler: func(m uiModel) (uiModel, tea.Cmd) {
m.logger.Debug("Received history call")
m.focus = focusHistory
return m, m.getHistory()
},
},
{
name: "models",
description: "Choose an LLM model",
handler: func(m uiModel) (uiModel, tea.Cmd) {
m.logger.Debug("Received models sequence")
m.focus = focusModels
return m, m.getModels()
},
},
{
name: "rename",
description: "Rename the current conversation",
handler: func(m uiModel) (uiModel, tea.Cmd) {
m.logger.Debug("Received rename request")
m.waiting = true
return m, m.renameConversation()
},
},
}
}
func (cl *commandList) setAvailable(cmds []command) {
cl.available = cmds
items := make([]list.Item, len(cmds))
for i, c := range cmds {
items[i] = c
}
cl.list.SetItems(items)
cl.list.SetHeight(len(items) + 1)
}
func (cl *commandList) filter(query string) {
var matches []list.Item
if query == "" {
matches = make([]list.Item, len(cl.available))
for i, c := range cl.available {
matches[i] = c
}
} else {
names := make([]string, len(cl.available))
for i, c := range cl.available {
names[i] = c.name
}
result := fuzzy.Find(query, names)
matches = make([]list.Item, len(result))
for i, r := range result {
matches[i] = cl.available[r.Index]
}
}
cl.list.SetItems(matches)
cl.list.ResetSelected()
}
func (cl commandList) view() string {
return dropdownStyle.Render(cl.list.View())
}
func (cl commandList) selectedItem() (command, bool) {
c, ok := cl.list.SelectedItem().(command)
return c, ok
}
+1 -8
View File
@@ -38,18 +38,11 @@ const Wordmark = `
\__ \ (_) | |_| |\ V / __/ | | | | | \__ \ (_) | |_| |\ V / __/ | | | | |
|___/\___/ \__,_| \_/ \___|_| |_|_|_|` |___/\___/ \__,_| \_/ \___|_| |_|_|_|`
const Commands = `
/models
/exit
`
func renderLanding(width int, height int) string { func renderLanding(width int, height int) string {
return lipgloss.PlaceVertical(height, 0.7, return lipgloss.PlaceVertical(height, 0.7,
lipgloss.JoinVertical(lipgloss.Center, lipgloss.JoinVertical(lipgloss.Center,
lipgloss.PlaceHorizontal(width, lipgloss.PlaceHorizontal(width,
lipgloss.Center, lipgloss.Center,
lipgloss.JoinHorizontal(lipgloss.Center, Logo, Wordmark)), lipgloss.JoinHorizontal(lipgloss.Center, Logo, Wordmark)),
lipgloss.PlaceHorizontal(width, ))
lipgloss.Center,
Commands)))
} }
-61
View File
@@ -1,61 +0,0 @@
package ui
import (
"log/slog"
"charm.land/bubbles/v2/list"
tea "charm.land/bubbletea/v2"
)
type modelPicker struct {
focus *focusState
logger slog.Logger
list list.Model
}
type modelsLoadedMsg struct{ models []string }
type modelsErrMsg struct{ err error }
type modelChosenMsg struct{ id string }
type modelItem struct{ name string }
func (i modelItem) FilterValue() string { return i.name }
func (i modelItem) Title() string { return i.name }
func (i modelItem) Description() string { return i.name }
func newModelPicker(focus *focusState) modelPicker {
l := list.New(nil, list.NewDefaultDelegate(), 0, 0)
l.Title = "Choose a model"
return modelPicker{list: l, logger: *slog.Default().With("Component", "ModelPicker"), focus: focus}
}
func (p *modelPicker) SetSize(w, h int) {
p.list.SetSize(w, h)
}
func (p modelPicker) Update(msg tea.Msg) (modelPicker, tea.Cmd) {
switch msg := msg.(type) {
case modelsLoadedMsg:
p.logger.Debug("Received modelsLoadedMsg", "models", msg.models)
items := make([]list.Item, len(msg.models))
for i, m := range msg.models {
items[i] = modelItem{ name: m}
}
return p, p.list.SetItems(items)
case tea.KeyPressMsg:
if msg.String() == "enter" {
p.logger.Debug("Selecting model", "model", p.list.SelectedItem())
if it, ok := p.list.SelectedItem().(modelItem); ok {
*p.focus = focusChat
p.logger.Debug("Sending modelChosenMsg")
return p, func() tea.Msg { return modelChosenMsg{it.name} }
}
}
}
var cmd tea.Cmd
p.list, cmd = p.list.Update(msg)
return p, cmd
}
func (p modelPicker) View() string {
return p.list.View()
}
+73
View File
@@ -0,0 +1,73 @@
package ui
import (
"log/slog"
"charm.land/bubbles/v2/list"
tea "charm.land/bubbletea/v2"
)
type picker struct {
logger slog.Logger
list list.Model
title string
}
type modelsLoadedMsg struct {
models []modelItem
title string
}
type pickerChosenMsg struct{ id string }
type pickerDismissMsg struct {}
type modelItem struct {
name string
description string
}
func (i modelItem) FilterValue() string { return i.name }
func (i modelItem) Title() string { return i.name }
func (i modelItem) Description() string { return i.description }
func newPicker() picker {
l := list.New(nil, list.NewDefaultDelegate(), 0, 0)
l.DisableQuitKeybindings()
return picker{list: l, logger: *slog.Default().With("Component", "Picker")}
}
func (p *picker) SetSize(w, h int) {
p.list.SetSize(w, h)
}
func (p picker) Update(msg tea.Msg) (picker, tea.Cmd) {
switch msg := msg.(type) {
case modelsLoadedMsg:
p.logger.Debug("Received modelsLoadedMsg", "models", msg.models)
p.list.Title = msg.title
items := make([]list.Item, len(msg.models))
for i, m := range msg.models {
items[i] = modelItem{name: m.name, description: m.description}
}
return p, p.list.SetItems(items)
case tea.KeyPressMsg:
if msg.String() == "q" || msg.String() == "esc" {
if p.list.FilterState() != list.Filtering {
return p, func() tea.Msg { return pickerDismissMsg{} }
}
}
if msg.String() == "enter" {
p.logger.Debug("Selecting model", "model", p.list.SelectedItem())
if it, ok := p.list.SelectedItem().(modelItem); ok {
p.logger.Debug("Sending modelChosenMsg")
return p, func() tea.Msg { return pickerChosenMsg{it.name} }
}
}
}
var cmd tea.Cmd
p.list, cmd = p.list.Update(msg)
return p, cmd
}
func (p picker) View() string {
return p.list.View()
}
+590
View File
@@ -0,0 +1,590 @@
package ui
import (
"context"
"fmt"
"log/slog"
"strings"
"charm.land/bubbles/v2/cursor"
"charm.land/bubbles/v2/spinner"
"charm.land/bubbles/v2/textarea"
"charm.land/bubbles/v2/viewport"
tea "charm.land/bubbletea/v2"
"charm.land/lipgloss/v2"
"git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db/history"
llm "git.estatecloud.org/radumaco/souvenir/llm"
"git.estatecloud.org/radumaco/souvenir/model"
)
type focusState int
const (
focusChat focusState = iota
focusModels
focusHistory
)
type uiModel struct {
logger slog.Logger
viewport viewport.Model
spinner spinner.Model
conversation model.Conversation
textarea textarea.Model
focus focusState
picker picker
thinkingBuffer *strings.Builder
answerBuffer *strings.Builder
senderStyle lipgloss.Style
agentStyle lipgloss.Style
errorStyle lipgloss.Style
statusStyle lipgloss.Style
client llm.ChatClient
statusMessage string
commands commandList
history history.DbClient
streamCh <-chan llm.StreamEvent
ctx context.Context
width int
height int
waiting bool
err error
}
type streamEventMessage struct {
event llm.StreamEvent
}
type streamClosedMessage struct {
}
type streamStartedMessage struct {
ch <-chan llm.StreamEvent
}
type conversationSavedMessage struct {
conversation model.Conversation
err error
}
type conversationRenamedMessage struct {
title string
summary string
err error
}
func InitialModel(ctx context.Context, config config.Config, client history.DbClient) uiModel {
ta := textarea.New()
ta.Placeholder = "Send a message..."
ta.SetVirtualCursor(false)
ta.Focus()
ta.Prompt = "│ "
ta.CharLimit = 280
ta.SetWidth(30)
ta.SetHeight(3)
s := ta.Styles()
s.Focused.CursorLine = lipgloss.NewStyle()
ta.SetStyles(s)
ta.ShowLineNumbers = false
vp := viewport.New(viewport.WithWidth(30), viewport.WithHeight(5))
vp.SetContent(renderLanding(vp.Width(), vp.Height()))
vp.KeyMap.Left.SetEnabled(false)
vp.KeyMap.Right.SetEnabled(false)
ta.KeyMap.InsertNewline.SetEnabled(false)
sp := spinner.New()
sp.Spinner = spinner.Dot
ui := uiModel{
textarea: ta,
conversation: model.Conversation{},
viewport: vp,
spinner: sp,
senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")),
errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")),
agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")),
thinkingBuffer: &strings.Builder{},
answerBuffer: &strings.Builder{},
client: *llm.NewLLMClient(config),
focus: focusChat,
history: client,
ctx: ctx,
logger: *slog.Default().With("Component", "TUI"),
err: nil,
}
ui.picker = newPicker()
ui.commands = newCommandList(30)
ui.commands.setAvailable(ui.buildCommands())
return ui
}
func (m uiModel) Init() tea.Cmd {
return textarea.Blink
}
func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.logger.Debug("Received message", "msg", msg)
switch msg := msg.(type) {
case tea.WindowSizeMsg:
m.viewport.SetWidth(msg.Width * 80 / 100)
m.textarea.SetWidth(msg.Width * 80 / 100)
m.viewport.SetHeight(msg.Height - lipgloss.Height(m.textarea.View()) - 1)
m.width = msg.Width
m.height = msg.Height
m.picker.SetSize(msg.Width, msg.Height)
m.commands.list.SetSize(msg.Width*80/100, commandDropdownHeight)
if m.commands.open {
m.viewport.SetHeight(m.viewport.Height() - commandDropdownHeight)
}
var messages = m.renderMessages()
if len(m.conversation.Messages) > 0 {
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
} else {
m.viewport.SetContent(renderLanding(m.viewport.Width(), m.viewport.Height()))
}
m.logger.Debug("New sizes",
"vp width", m.viewport.Width(), "vp height", m.viewport.Height(),
"ta width", m.textarea.Width(), "ta height", m.textarea.Height())
m.viewport.GotoBottom()
case tea.KeyPressMsg:
if m.focus != focusChat {
picker, cmd := m.picker.Update(msg)
m.picker = picker
return m, cmd
}
switch msg.String() {
case "esc":
if m.commands.open {
m.closeCommands()
return m, nil
}
if m.focus != focusChat {
m.focus = focusChat
}
return m, nil
case "ctrl+c":
fmt.Println(m.textarea.Value())
return m, tea.Quit
case "shift+enter", "ctrl+j", "alt+enter":
m.textarea.InsertRune('\n')
return m, nil
case "up", "down":
if m.commands.open {
var cmd tea.Cmd
m.commands.list, cmd = m.commands.list.Update(msg)
return m, cmd
}
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
return m, cmd
case "tab":
if m.commands.open {
if sel, ok := m.commands.selectedItem(); ok {
m.textarea.SetValue("/" + sel.name + " ")
}
m.closeCommands()
return m, nil
}
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
return m, cmd
case "enter":
m.setStatusMessage("")
input := m.textarea.Value()
if m.commands.open {
if sel, ok := m.commands.selectedItem(); ok {
m.textarea.SetValue("")
m.closeCommands()
mm, cmd := sel.handler(m)
return mm, cmd
}
}
if after, ok := strings.CutPrefix(input, "/"); ok {
name := after
for _, c := range m.commands.available {
if c.name == name {
m.textarea.SetValue("")
m.closeCommands()
mm, cmd := c.handler(m)
return mm, cmd
}
}
}
m.textarea.SetValue("")
m.closeCommands()
m.logger.Debug("Received enter. Creating new message", "message", input)
m.conversation.Messages = append(m.conversation.Messages, model.Message{
Content: input,
Role: "user",
Seq: len(m.conversation.Messages) + 1,
})
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.textarea.Reset()
m.viewport.GotoBottom()
m.startWait()
return m, tea.Batch(m.callAgent(m.conversation.Messages), m.saveConversation(), m.spinner.Tick)
default:
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
v := m.textarea.Value()
if strings.HasPrefix(v, "/") {
m.openCommands()
} else {
m.closeCommands()
}
return m, cmd
}
case cursor.BlinkMsg:
m.logger.Debug("Received BlinkMsg")
var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg)
return m, cmd
case conversationRenamedMessage:
m.logger.Debug("Received messageRenamedMessage")
m.stopWait()
if msg.err != nil {
m.logger.Error("There was an error", "error", msg.err)
m.setErrorMessage(msg.err.Error())
}
m.setStatusMessage("Conversation renamed to: " + msg.title)
m.conversation.Title = msg.title
m.conversation.Summary = msg.summary
return m, nil
case streamStartedMessage:
m.logger.Debug("Received streamStartedMessage")
m.streamCh = msg.ch
return m, waitForEvent(m.streamCh)
case streamEventMessage:
m.logger.Debug("Received streamEventMessage")
if msg.event.Err != nil {
m.logger.Error("There was an error mid stream", "error", msg.event.Err)
}
if msg.event.Reasoning != "" {
m.thinkingBuffer.WriteString(msg.event.Reasoning)
}
if msg.event.Content != "" {
m.answerBuffer.WriteString(msg.event.Content)
}
if msg.event.Done {
for _, r := range msg.event.Messages {
r.Seq = len(m.conversation.Messages) + 1
m.logger.Debug("Appending response message", "message", r)
m.conversation.Messages = append(m.conversation.Messages, r)
}
return m, func() tea.Msg { return streamClosedMessage{} }
}
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.viewport.GotoBottom()
return m, waitForEvent(m.streamCh)
case streamClosedMessage:
m.logger.Debug("Received streamClosedMessage")
m.stopWait()
m.answerBuffer.Reset()
m.thinkingBuffer.Reset()
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.viewport.GotoBottom()
return m, nil
case pickerChosenMsg:
switch m.focus {
case focusModels:
m.logger.Debug("Received pickerChosenMsg", "model", msg.id)
m.client.Cfg.Llm.Model = msg.id
m.focus = focusChat
case focusHistory:
m.logger.Debug("Received pickerChosenMsg", "history", msg.id)
conv, err := m.history.GetConversation(m.ctx, msg.id)
if err != nil {
m.logger.Error("Could not retrieve conversation", "id", msg.id)
}
m.conversation = conv
m.focus = focusChat
var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.viewport.GotoBottom()
}
case pickerDismissMsg:
m.logger.Debug("Received pickerDismissMsg")
m.stopWait()
m.focus = focusChat
case conversationSavedMessage:
m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation)
if msg.err != nil {
m.logger.Error("Could not save message", "error", msg.err)
m.setErrorMessage(msg.err.Error())
return m, nil
}
m.conversation = msg.conversation
}
var spinnerCmd tea.Cmd
if m.waiting {
m.spinner, spinnerCmd = m.spinner.Update(msg)
}
if m.focus != focusChat {
picker, cmd := m.picker.Update(msg)
m.picker = picker
return m, cmd
}
return m, spinnerCmd
}
func (m uiModel) View() tea.View {
if m.focus != focusChat {
return tea.NewView(lipgloss.Place(m.width, m.height,
lipgloss.Center, lipgloss.Center,
m.picker.View()))
}
parts := []string{m.viewport.View()}
if m.statusMessage != "" {
parts = append(parts, m.statusMessage)
}
if m.waiting {
parts = append(parts, m.spinner.View())
}
if m.commands.open {
parts = append(parts, m.commands.view())
}
parts = append(parts, m.textarea.View())
ui := lipgloss.JoinVertical(
lipgloss.Left,
parts...)
c := m.textarea.Cursor()
if c != nil {
c.Y += lipgloss.Height(m.viewport.View())
if m.statusMessage != "" {
c.Y++
}
if m.commands.open {
c.Y += lipgloss.Height(m.commands.view())
}
if m.waiting {
c.Y++
}
gap := m.width - lipgloss.Width(ui)
c.X += max(gap/2, 0)
}
view := tea.NewView(lipgloss.PlaceHorizontal(m.width, lipgloss.Center, ui))
view.Cursor = c
view.AltScreen = true
return view
}
func (m uiModel) callAgent(messages []model.Message) tea.Cmd {
m.logger.Debug("Calling agent", "messages", messages)
return func() tea.Msg {
resp, err := m.client.QueryStream(messages)
if err != nil {
m.logger.Error("Error while calling stream query", "error", err)
return streamClosedMessage{}
}
return streamStartedMessage{ch: resp}
}
}
func waitForEvent(ch <-chan llm.StreamEvent) tea.Cmd {
return func() tea.Msg {
ev, ok := <-ch
if !ok {
return streamClosedMessage{}
}
return streamEventMessage{event: ev}
}
}
func (m uiModel) saveConversation() tea.Cmd {
m.logger.Debug("Saving conversation", "conversation", m.conversation)
return func() tea.Msg {
conv, err := m.history.SaveConversation(m.ctx, m.conversation)
return conversationSavedMessage{conversation: conv, err: err}
}
}
func (m uiModel) renameConversation() tea.Cmd {
m.logger.Debug("Renaming conversation", "conversation", m.conversation)
return func() tea.Msg {
meta, err := m.client.Rename(m.conversation.Messages)
m.startWait()
if err != nil {
return conversationRenamedMessage{title: meta.Title, summary: meta.Summary, err: err}
}
m.conversation.Title = meta.Title
m.conversation.Summary = meta.Summary
m.conversation, err = m.history.SaveConversation(m.ctx, m.conversation)
return conversationRenamedMessage{title: m.conversation.Title, summary: m.conversation.Summary, err: err}
}
}
func (m uiModel) getModels() tea.Cmd {
m.logger.Debug("Querying models")
return func() tea.Msg {
resp, err := m.client.Models()
if err != nil {
m.logger.Error("Could not fetch models", "error", err)
}
m.logger.Debug("Received models from llm", "models", resp)
items := []modelItem{}
for _, model := range resp {
items = append(items, modelItem{name: model, description: model})
}
return modelsLoadedMsg{models: items, title: "Choose a model"}
}
}
func (m uiModel) getHistory() tea.Cmd {
m.logger.Debug("Querying history")
return func() tea.Msg {
resp, err := m.history.GetConversations(m.ctx)
if err != nil {
m.logger.Error("Could not fetch history", "error", err)
}
m.logger.Debug("Received conversations from psql", "conversations", resp)
items := []modelItem{}
for _, conv := range resp {
name := conv.Id
if conv.Title != "" {
name = conv.Title
}
items = append(items, modelItem{name: name, description: conv.Summary})
}
return modelsLoadedMsg{models: items, title: "Select a conversation to continue from where you left off"}
}
}
func (m *uiModel) renderMessages() string {
var result strings.Builder
var width = m.viewport.Width()
for _, message := range m.conversation.Messages {
switch message.Role {
case "user":
var line strings.Builder
line.WriteString(m.senderStyle.Render("You: "))
line.WriteString(message.Content)
result.WriteString(lipgloss.PlaceHorizontal(width, lipgloss.Right, line.String()))
result.WriteString("\n")
var sepStyle lipgloss.Style = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
MarginTop(1).
MarginBottom(1)
sep := sepStyle.Render(strings.Repeat("┈", width))
result.WriteString(sep)
result.WriteString("\n")
case "assistant":
result.WriteString(m.agentStyle.Render("Agent: "))
result.WriteString(message.Content)
result.WriteString("\n")
var sepStyle lipgloss.Style = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
MarginTop(1).
MarginBottom(1)
sep := sepStyle.Render(strings.Repeat("┈", width))
result.WriteString(sep)
result.WriteString("\n")
}
}
if m.answerBuffer.String() != "" {
result.WriteString(m.agentStyle.Render("Agent: "))
result.WriteString(m.answerBuffer.String())
result.WriteString("\n")
var sepStyle lipgloss.Style = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
MarginTop(1).
MarginBottom(1)
sep := sepStyle.Render(strings.Repeat("┈", width))
result.WriteString(sep)
result.WriteString("\n")
}
if m.thinkingBuffer.String() != "" && m.answerBuffer.String() == "" {
result.WriteString(m.agentStyle.Render("Thinking: "))
result.WriteString(m.answerBuffer.String())
result.WriteString("\n")
var sepStyle lipgloss.Style = lipgloss.NewStyle().
Foreground(lipgloss.Color("240")).
MarginTop(1).
MarginBottom(1)
sep := sepStyle.Render(strings.Repeat("┈", width))
result.WriteString(sep)
result.WriteString("\n")
}
return result.String()
}
func (m *uiModel) startWait() {
if !m.waiting {
m.viewport.SetHeight(m.viewport.Height() - 1)
m.waiting = true
}
}
func (m *uiModel) stopWait() {
if m.waiting {
m.waiting = false
m.viewport.SetHeight(m.viewport.Height() + 1)
}
}
func (m *uiModel) setErrorMessage(message string) {
m.statusMessage = m.errorStyle.Render(message)
}
func (m *uiModel) setStatusMessage(message string) {
if message == "" {
m.statusMessage = m.statusStyle.Render("model: " + m.client.Cfg.Llm.Model + " url: " + m.client.Cfg.Api.Url)
return
}
m.statusMessage = m.statusStyle.Render(message)
}
func (m *uiModel) openCommands() {
if !m.commands.open {
m.commands.open = true
m.viewport.SetHeight(m.viewport.Height() - commandDropdownHeight)
}
m.commands.filter(strings.TrimPrefix(m.textarea.Value(), "/"))
}
func (m *uiModel) closeCommands() {
if m.commands.open {
m.commands.open = false
m.viewport.SetHeight(m.viewport.Height() + commandDropdownHeight)
}
}