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
+91 -25
View File
@@ -5,26 +5,52 @@ 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 {
Url string Url string
Port string Port string
DbName string DbName string
User string User string
Password string Password string
History HistoryConfig History HistoryConfig
Memory MemoryConfig Memory MemoryConfig
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 {
@@ -36,12 +62,14 @@ 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)
@@ -103,17 +131,29 @@ func getLogTarget(target string) (io.Writer, error) {
func defaultConfig() *Config { 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{
Url: "localhost", Url: "localhost",
User: "psql", User: "psql",
Port: "54321", Port: "54321",
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()
}
+272 -60
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
}
func (client DbClient) init() error { if err := ensureDbExists(cfg, logger); err != nil {
const createConversationTable = ` return nil, err
CREATE TABLE IF NOT EXISTS conversations ( }
id UUID PRIMARY KEY DEFAULT gen_random_uuid(), pool, err := pgxpool.New(ctx, cfg.ConnectionString())
created_at TIMESTAMPTZ NOT NULL DEFAULT now() if err != nil {
)` logger.ErrorContext(ctx, "Could not create pgxpool")
const createMessageTable = ` return nil, err
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 `
WHERE table_schema = 'public' SELECT EXISTS (
AND table_name =$1)`, table.tableName).Scan(&exists); err != nil { SELECT 1
client.Logger.Error("Error while checking table", "table", table.tableName) FROM information_schema.tables
return err WHERE table_schema = 'public'
AND table_name =$1
)
`,
table.tableName).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", table.tableName, "error", err)
return nil, err
} }
if exists { if exists {
client.Logger.Debug("Table found", "table", table.tableName) logger.Debug("Table found", "table", table.tableName)
} else { } else {
client.Logger.Debug("Table not found. Creating", "table", table.tableName) logger.Debug("Table not found. Creating", "table", table.tableName)
if _, err := conn.Exec(table.script); err != nil { if _, err := pool.Exec(ctx, table.script); err != nil {
client.Logger.Error("Could not create table", "table", table.tableName) logger.Error("Could not create table", "table", table.tableName, "error", err)
return err return nil, err
} }
} }
} }
if _, err := conn.Exec(createToolCallTable); err != nil { return &DbClient{
return err logger: logger,
} cfg: cfg,
if _, err := conn.Exec(createToolCallTable); err != nil { pool: pool,
return err }, nil
}
return nil
} }
func (client DbClient) connectionString() string { func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) {
u := &url.URL{ cl.logger.Debug("Searching for conversation", "id", id)
Scheme: "postgres", var conv model.Conversation
User: url.UserPassword(client.Cfg.User, client.Cfg.Password), if err := cl.pool.QueryRow(ctx,
Host: net.JoinHostPort(client.Cfg.Url, client.Cfg.Port), `
Path: client.Cfg.DbName, 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
} }
return u.String() 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
}
if exists {
logger.Debug("Database already exists", "name", cfg.DbName)
return nil
}
logger.Info("Creating database", "name", cfg.DbName)
_, err := conn.Exec(`CREATE DATABASE "` + cfg.DbName + `"`)
return err
} }
+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)
}
}