Compare commits

..

2 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
8 changed files with 630 additions and 172 deletions
+3
View File
@@ -50,6 +50,7 @@ type EmbeddingConfig struct {
Dim int Dim int
Timeout int Timeout int
BatchSize int BatchSize int
Interval int
} }
type MemoryConfig struct { type MemoryConfig struct {
@@ -152,6 +153,7 @@ func defaultConfig() *Config {
Model: "", Model: "",
BatchSize: 64, BatchSize: 64,
Timeout: 60000, Timeout: 60000,
Interval: 60,
}, },
} }
} }
@@ -185,6 +187,7 @@ func overrideFromEnv(cfg *Config) {
}{ }{
{conf: &cfg.Api.Timeout, key: ENV_PREFIX + "api__timeout"}, {conf: &cfg.Api.Timeout, key: ENV_PREFIX + "api__timeout"},
{conf: &cfg.Embedding.Timeout, key: ENV_PREFIX + "embedding__timeout"}, {conf: &cfg.Embedding.Timeout, key: ENV_PREFIX + "embedding__timeout"},
{conf: &cfg.Embedding.Interval, key: ENV_PREFIX + "embedding__interval"},
} }
for _, override := range durationOverrides { for _, override := range durationOverrides {
val := os.Getenv(override.key) val := os.Getenv(override.key)
+57 -15
View File
@@ -7,6 +7,7 @@ import (
"log/slog" "log/slog"
"git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/llm" "git.estatecloud.org/radumaco/souvenir/llm"
"github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/pgxpool"
@@ -32,6 +33,35 @@ func (e EmbedStore) TableName() string {
return e.tableName 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) { func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedder) (*EmbedStore, error) {
logger := *slog.Default().With("Component", "EmbedStore") logger := *slog.Default().With("Component", "EmbedStore")
@@ -47,17 +77,16 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
pool: pool, pool: pool,
} }
const createVectorTableTemplate = `
CREATE TABLE IF NOT EXISTS %s (
chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE,
embedding vector(%d),
content_hash text not null
)`
var exists bool var exists bool
if err := pool.QueryRow(ctx, `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)`, store.TableName()).Scan(&exists); err != nil { AND table_name =$1)
`,
store.TableName()).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", store.TableName(), "error", err) logger.Error("Error while checking table", "table", store.TableName(), "error", err)
return nil, err return nil, err
} }
@@ -65,7 +94,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
logger.Debug("Table found", "table", store.TableName()) logger.Debug("Table found", "table", store.TableName())
} else { } else {
logger.Debug("Table not found. Creating", "table", store.TableName()) logger.Debug("Table not found. Creating", "table", store.TableName())
script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim()) script := fmt.Sprintf(db.CreateVectorTableTemplate, store.TableName(), store.embedder.Dim())
if _, err := pool.Exec(ctx, script); err != nil { if _, err := pool.Exec(ctx, script); err != nil {
logger.Error("Could not create table", "table", store.TableName(), "error", err) logger.Error("Could not create table", "table", store.TableName(), "error", err)
return nil, err return nil, err
@@ -75,11 +104,17 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
} }
func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string) error { 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 rows, err := e.pool.Query(ctx, fmt.Sprintf(
`
SELECT c.id, c.content, c.content_hash
FROM message_chunks c FROM message_chunks c
LEFT JOIN %s v LEFT JOIN %s v
ON v.chunk_id = c.id AND v.content_hash = c.content_hash ON v.chunk_id = c.id
WHERE c.conversation_id = $1 and v.chunk_id IS NULL`, e.TableName()), AND v.content_hash = c.content_hash
WHERE c.conversation_id = $1
AND v.chunk_id IS NULL
`,
e.TableName()),
conversationId) conversationId)
if err != nil { if err != nil {
@@ -131,12 +166,15 @@ func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string
} }
func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error { func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error {
query := fmt.Sprintf(` query := fmt.Sprintf(
`
INSERT INTO %s (chunk_id, embedding, content_hash) INSERT INTO %s (chunk_id, embedding, content_hash)
VALUES ($1, $2, $3) VALUES ($1, $2, $3)
ON CONFLICT (chunk_id) DO UPDATE ON CONFLICT (chunk_id) DO UPDATE
SET embedding = EXCLUDED.embedding, SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash`, e.TableName()) content_hash = EXCLUDED.content_hash
`,
e.TableName())
batch := &pgx.Batch{} batch := &pgx.Batch{}
for i, chunk := range chunks { for i, chunk := range chunks {
@@ -152,3 +190,7 @@ func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32
} }
return nil return nil
} }
func (e EmbedStore) Close() {
e.pool.Close()
}
+82 -52
View File
@@ -11,6 +11,7 @@ import (
"strings" "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" "github.com/jackc/pgx/v5/pgxpool"
@@ -24,36 +25,6 @@ type DbClient struct {
func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) { func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
logger := *slog.Default().With("Component", "History") logger := *slog.Default().With("Component", "History")
const createConversationTable = `
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()
)`
const createMessageTable = `
CREATE TABLE IF NOT EXISTS messages (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
role TEXT NOT NULL,
seq INTEGER NOT NULL,
content TEXT NOT NULL,
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
tsv tsvector generated always as (to_tsvector('english', content)) stored,
UNIQUE (conversation_id, seq)
);
CREATE INDEX on messages using gin (tsv);
CREATE INDEX on messages (conversation_id, seq);`
const createMessageChunkTable = `
CREATE TABLE IF NOT EXISTS message_chunks (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE,
content TEXT NOT NULL,
content_hash TEXT NOT NULL
)`
if err := ensureDbExists(cfg, logger); err != nil { if err := ensureDbExists(cfg, logger); err != nil {
return nil, err return nil, err
@@ -67,17 +38,23 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
tableName string tableName string
script string script string
}{ }{
{"conversations", createConversationTable}, {"conversations", db.ConversationTableScript},
{"messages", createMessageTable}, {"messages", db.MessageTableScript},
{"message_chunks", createMessageChunkTable}, {"message_chunks", db.MessageChunkTableScript},
} }
for _, table := range tables { for _, table := range tables {
var exists bool var exists bool
if err := pool.QueryRow(ctx, `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
)
`,
table.tableName).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", table.tableName, "error", err) logger.Error("Error while checking table", "table", table.tableName, "error", err)
return nil, err return nil, err
} }
@@ -102,13 +79,22 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) { func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) {
cl.logger.Debug("Searching for conversation", "id", id) cl.logger.Debug("Searching for conversation", "id", id)
var conv model.Conversation 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 { 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) cl.logger.Error("Could not fetch conversation", "id", id, "error", err)
return conv, err return conv, err
} }
rows, err := cl.pool.Query(ctx, `SELECT id, role, content, seq rows, err := cl.pool.Query(ctx,
`
SELECT id, role, content, seq
FROM messages FROM messages
WHERE conversation_id = $1 ORDER BY seq`, id) WHERE conversation_id = $1
ORDER BY seq
`, id)
if err != nil { if err != nil {
cl.logger.Error("Could not fetch messages", "conversation_id", id, "error", err) cl.logger.Error("Could not fetch messages", "conversation_id", id, "error", err)
return conv, err return conv, err
@@ -137,7 +123,11 @@ func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conver
func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) { func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) {
cl.logger.Debug("Fetching conversations") cl.logger.Debug("Fetching conversations")
rows, err := cl.pool.Query(ctx, "SELECT id, title, summary FROM conversations") rows, err := cl.pool.Query(ctx,
`
SELECT id, title, summary
FROM conversations
`)
if err != nil { if err != nil {
cl.logger.Error("Could not fetch conversations", "error", err) cl.logger.Error("Could not fetch conversations", "error", err)
return []model.Conversation{}, err return []model.Conversation{}, err
@@ -157,19 +147,33 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation
cl.logger.Debug("Saving conversation", "conversation", conv) cl.logger.Debug("Saving conversation", "conversation", conv)
delta := 0 delta := 0
if conv.Id == "" { if conv.Id == "" {
if err := cl.pool.QueryRow(ctx, `INSERT INTO conversations(title, summary) VALUES($1, $2) RETURNING id`, if err := cl.pool.QueryRow(ctx,
conv.Title, `
conv.Summary).Scan(&conv.Id); err != nil { 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) cl.logger.Error("Could not create conversation", "title", conv.Title, "error", err)
return conv, err return conv, err
} }
} else { } else {
if _, err := cl.pool.Exec(ctx, `UPDATE conversations SET title = $1, summary = $2 WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil { 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) cl.logger.Error("Could not update conversation", "id", conv.Id, "error", err)
return conv, 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 { 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) cl.logger.Warn("Could not get delta", "conversation_id", conv.Id, "error", err)
return conv, errors.New("Could not get delta") return conv, errors.New("Could not get delta")
} }
@@ -193,16 +197,27 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation
func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) { func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) {
var exists bool var exists bool
if message.Id != "" { if message.Id != "" {
if err := cl.pool.QueryRow(ctx, `SELECT EXISTS if err := cl.pool.QueryRow(ctx,
(SELECT 1 FROM messages `
WHERE id = $1)`, message.Id).Scan(&exists); err != nil { 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) cl.logger.Warn("Error while scanning for message", "id", message.Id, "error", err)
return message.Id, err return message.Id, err
} }
} }
if !exists { if !exists {
cl.logger.Debug("Insertting message ", "seq", message.Seq) 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`, 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 { 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) cl.logger.Error("Could not save message", "seq", message.Seq, "error", err)
return message.Id, errors.New("Could not save message") return message.Id, errors.New("Could not save message")
@@ -213,7 +228,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
for _, chunk := range chunks { for _, chunk := range chunks {
hash := sha256.Sum256([]byte(chunk)) 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)`, 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 { message.Id, hex.EncodeToString(hash[:])).Scan(&exists); err != nil {
cl.logger.Warn("Error while scanning for message_chunk", "message id", message.Id, "error", err) 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") return message.Id, errors.New("Could not save message chunk")
@@ -223,8 +246,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
continue continue
} }
if _, err := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) if _, err := cl.pool.Exec(ctx,
VALUES ($1, $2, $3, $4)`, conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil { `
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) cl.logger.Warn("Error while creating message chunk", "message id", message.Id, "error", 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()
)`
+51 -1
View File
@@ -2,16 +2,19 @@ package main
import ( import (
"context" "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/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() ctx := context.Background()
config, err := config.Load(".config.json") config, err := config.Load(".config.json")
@@ -28,6 +31,10 @@ func main() {
os.Exit(1) os.Exit(1)
} }
defer db.Close() defer db.Close()
es, _ := setupEmbedding(config, logger)
if es != nil {
defer es.Close()
}
p := tea.NewProgram(ui.InitialModel(ctx, *config, *db)) p := tea.NewProgram(ui.InitialModel(ctx, *config, *db))
if _, err := p.Run(); err != nil { if _, err := p.Run(); err != nil {
@@ -35,3 +42,46 @@ func main() {
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
}
+72 -35
View File
@@ -8,6 +8,7 @@ import (
"log/slog" "log/slog"
"net/http" "net/http"
"strconv" "strconv"
"strings"
"time" "time"
config "git.estatecloud.org/radumaco/souvenir/config" config "git.estatecloud.org/radumaco/souvenir/config"
@@ -29,6 +30,16 @@ type ChatRequest struct {
Model string `json:"model"` Model string `json:"model"`
Messages []model.Message `json:"messages"` Messages []model.Message `json:"messages"`
Tools []config.Tool `json:"tools,omitempty"` 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 // // Response types //
@@ -129,7 +140,12 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) {
Role: "system", Role: "system",
Content: "Reply with ONLY JSON with the form: {\"title\":<title>, \"summary\":<summary>}. Title <= 10 words. Summary 2-3 sentances.", Content: "Reply with ONLY JSON with the form: {\"title\":<title>, \"summary\":<summary>}. Title <= 10 words. Summary 2-3 sentances.",
}}, messages...), }}, 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) cl.logger.Debug("Calling rename", "request", request)
out, err := cl.Call(request) out, err := cl.Call(request)
if err != nil { if err != nil {
@@ -137,63 +153,45 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) {
return RenameResponse{}, err return RenameResponse{}, err
} }
var meta RenameResponse var meta RenameResponse
if err := json.Unmarshal([]byte(out[0].Content), &meta); err != nil { 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) cl.logger.Error("Could not parse the rename response", "content", out[0].Content)
return meta, err return meta, err
} }
return meta, nil 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(query string, messages []model.Message) ([]model.Message, error) { func (cl ChatClient) Query(messages []model.Message) ([]model.Message, error) {
requestBody := ChatRequest{ requestBody := ChatRequest{
Model: cl.Cfg.Llm.Model, Model: cl.Cfg.Llm.Model,
Messages: messages, Messages: messages,
Tools: config.AvailableTools(), Tools: config.AvailableTools(),
Stream: true,
Kwargs: Kwargs{Thinking: false},
} }
return cl.Call(requestBody) return cl.Call(requestBody)
} }
func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) { func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) {
serialized, err := json.Marshal(requestBody) resp, err := cl.CallApi(requestBody)
if err != nil {
cl.logger.Error("There was an error marshelling the request body",
"request", requestBody,
"error", err)
return []model.Message{}, err
}
body := bytes.NewBuffer(serialized)
req, err := http.NewRequest("POST", cl.Cfg.Api.Url+COMPLETIONS_API, body)
if err != nil { if err != nil {
cl.logger.Error("There was an error creating the request", "error", err.Error()) cl.logger.Error("There was an error during the API call", "error", err)
return []model.Message{}, err 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 []model.Message{}, err
} }
defer resp.Body.Close() defer resp.Body.Close()
cl.logger.Debug("LLM Call returned", "response", resp)
data, err := io.ReadAll(resp.Body) data, err := io.ReadAll(resp.Body)
if err != nil { if err != nil {
cl.logger.Error("There was an error reading the response body", "error", err.Error()) cl.logger.Error("There was an error reading the response body", "error", err.Error())
return []model.Message{}, err 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 []model.Message{},
errors.New("Call returned invalid status " +
strconv.Itoa(resp.StatusCode) +
resp.Status)
} }
cl.logger.Debug("Received data", "data", data) cl.logger.Debug("Received data", "data", data)
@@ -215,3 +213,42 @@ func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) {
return result, nil 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
}
+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}}
}
+132 -26
View File
@@ -7,6 +7,7 @@ import (
"strings" "strings"
"charm.land/bubbles/v2/cursor" "charm.land/bubbles/v2/cursor"
"charm.land/bubbles/v2/spinner"
"charm.land/bubbles/v2/textarea" "charm.land/bubbles/v2/textarea"
"charm.land/bubbles/v2/viewport" "charm.land/bubbles/v2/viewport"
tea "charm.land/bubbletea/v2" tea "charm.land/bubbletea/v2"
@@ -28,10 +29,13 @@ const (
type uiModel struct { type uiModel struct {
logger slog.Logger logger slog.Logger
viewport viewport.Model viewport viewport.Model
spinner spinner.Model
conversation model.Conversation conversation model.Conversation
textarea textarea.Model textarea textarea.Model
focus focusState focus focusState
picker picker picker picker
thinkingBuffer *strings.Builder
answerBuffer *strings.Builder
senderStyle lipgloss.Style senderStyle lipgloss.Style
agentStyle lipgloss.Style agentStyle lipgloss.Style
errorStyle lipgloss.Style errorStyle lipgloss.Style
@@ -40,6 +44,7 @@ type uiModel struct {
statusMessage string statusMessage string
commands commandList commands commandList
history history.DbClient history history.DbClient
streamCh <-chan llm.StreamEvent
ctx context.Context ctx context.Context
width int width int
height int height int
@@ -47,9 +52,15 @@ type uiModel struct {
err error err error
} }
type agentResponseMessage struct { type streamEventMessage struct {
response []model.Message event llm.StreamEvent
err error }
type streamClosedMessage struct {
}
type streamStartedMessage struct {
ch <-chan llm.StreamEvent
} }
type conversationSavedMessage struct { type conversationSavedMessage struct {
@@ -69,7 +80,7 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl
ta.SetVirtualCursor(false) ta.SetVirtualCursor(false)
ta.Focus() ta.Focus()
ta.Prompt = " " ta.Prompt = " "
ta.CharLimit = 280 ta.CharLimit = 280
ta.SetWidth(30) ta.SetWidth(30)
@@ -88,13 +99,19 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl
ta.KeyMap.InsertNewline.SetEnabled(false) ta.KeyMap.InsertNewline.SetEnabled(false)
sp := spinner.New()
sp.Spinner = spinner.Dot
ui := uiModel{ ui := uiModel{
textarea: ta, textarea: ta,
conversation: model.Conversation{}, conversation: model.Conversation{},
viewport: vp, viewport: vp,
spinner: sp,
senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")), senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")),
errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")), errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")),
agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")), agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")),
thinkingBuffer: &strings.Builder{},
answerBuffer: &strings.Builder{},
client: *llm.NewLLMClient(config), client: *llm.NewLLMClient(config),
focus: focusChat, focus: focusChat,
history: client, history: client,
@@ -195,7 +212,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
} }
} }
if after, ok :=strings.CutPrefix(input, "/"); ok { if after, ok := strings.CutPrefix(input, "/"); ok {
name := after name := after
for _, c := range m.commands.available { for _, c := range m.commands.available {
if c.name == name { if c.name == name {
@@ -220,8 +237,8 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.textarea.Reset() m.textarea.Reset()
m.viewport.GotoBottom() m.viewport.GotoBottom()
m.waiting = true m.startWait()
return m, tea.Batch(m.callAgent(input, m.conversation.Messages), m.saveConversation()) return m, tea.Batch(m.callAgent(m.conversation.Messages), m.saveConversation(), m.spinner.Tick)
default: default:
var cmd tea.Cmd var cmd tea.Cmd
m.textarea, cmd = m.textarea.Update(msg) m.textarea, cmd = m.textarea.Update(msg)
@@ -242,7 +259,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case conversationRenamedMessage: case conversationRenamedMessage:
m.logger.Debug("Received messageRenamedMessage") m.logger.Debug("Received messageRenamedMessage")
m.waiting = false m.stopWait()
if msg.err != nil { if msg.err != nil {
m.logger.Error("There was an error", "error", msg.err) m.logger.Error("There was an error", "error", msg.err)
m.setErrorMessage(msg.err.Error()) m.setErrorMessage(msg.err.Error())
@@ -252,31 +269,56 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.conversation.Summary = msg.summary m.conversation.Summary = msg.summary
return m, nil return m, nil
case agentResponseMessage: case streamStartedMessage:
m.logger.Debug("Received agentResponseMessage") m.logger.Debug("Received streamStartedMessage")
m.waiting = false m.streamCh = msg.ch
resp := msg.response return m, waitForEvent(m.streamCh)
if msg.err != nil {
m.logger.Error("Message is error", "error", msg.err) case streamEventMessage:
m.setErrorMessage(msg.err.Error()) m.logger.Debug("Received streamEventMessage")
if msg.event.Err != nil {
m.logger.Error("There was an error mid stream", "error", msg.event.Err)
} }
for _, r := range resp { 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 r.Seq = len(m.conversation.Messages) + 1
m.logger.Debug("Appending response message", "message", r) m.logger.Debug("Appending response message", "message", r)
m.conversation.Messages = append(m.conversation.Messages, r) m.conversation.Messages = append(m.conversation.Messages, r)
} }
return m, func() tea.Msg { return streamClosedMessage{} }
}
var messages = m.renderMessages() var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
m.viewport.GotoBottom() m.viewport.GotoBottom()
return m, m.saveConversation() 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: case pickerChosenMsg:
switch m.focus { switch m.focus {
case focusModels: case focusModels:
m.logger.Debug("Received pickerChosenMsg", "model", msg.id) m.logger.Debug("Received pickerChosenMsg", "model", msg.id)
m.client.Cfg.Llm.Model = msg.id m.client.Cfg.Llm.Model = msg.id
m.waiting = false
m.focus = focusChat m.focus = focusChat
case focusHistory: case focusHistory:
m.logger.Debug("Received pickerChosenMsg", "history", msg.id) m.logger.Debug("Received pickerChosenMsg", "history", msg.id)
@@ -285,7 +327,6 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.logger.Error("Could not retrieve conversation", "id", msg.id) m.logger.Error("Could not retrieve conversation", "id", msg.id)
} }
m.conversation = conv m.conversation = conv
m.waiting = false
m.focus = focusChat m.focus = focusChat
var messages = m.renderMessages() var messages = m.renderMessages()
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
@@ -294,18 +335,22 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case pickerDismissMsg: case pickerDismissMsg:
m.logger.Debug("Received pickerDismissMsg") m.logger.Debug("Received pickerDismissMsg")
m.waiting = false m.stopWait()
m.focus = focusChat m.focus = focusChat
case conversationSavedMessage: case conversationSavedMessage:
m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation) m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation)
if msg.err != nil { if msg.err != nil {
m.logger.Error("Message is error", "error", msg.err) m.logger.Error("Could not save message", "error", msg.err)
m.setErrorMessage(msg.err.Error()) m.setErrorMessage(msg.err.Error())
return m, nil return m, nil
} }
m.conversation = msg.conversation m.conversation = msg.conversation
} }
var spinnerCmd tea.Cmd
if m.waiting {
m.spinner, spinnerCmd = m.spinner.Update(msg)
}
if m.focus != focusChat { if m.focus != focusChat {
picker, cmd := m.picker.Update(msg) picker, cmd := m.picker.Update(msg)
@@ -313,7 +358,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, cmd return m, cmd
} }
return m, nil return m, spinnerCmd
} }
func (m uiModel) View() tea.View { func (m uiModel) View() tea.View {
@@ -327,6 +372,9 @@ func (m uiModel) View() tea.View {
if m.statusMessage != "" { if m.statusMessage != "" {
parts = append(parts, m.statusMessage) parts = append(parts, m.statusMessage)
} }
if m.waiting {
parts = append(parts, m.spinner.View())
}
if m.commands.open { if m.commands.open {
parts = append(parts, m.commands.view()) parts = append(parts, m.commands.view())
} }
@@ -344,6 +392,9 @@ func (m uiModel) View() tea.View {
if m.commands.open { if m.commands.open {
c.Y += lipgloss.Height(m.commands.view()) c.Y += lipgloss.Height(m.commands.view())
} }
if m.waiting {
c.Y++
}
gap := m.width - lipgloss.Width(ui) gap := m.width - lipgloss.Width(ui)
c.X += max(gap/2, 0) c.X += max(gap/2, 0)
} }
@@ -353,11 +404,25 @@ func (m uiModel) View() tea.View {
return view return view
} }
func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd { func (m uiModel) callAgent(messages []model.Message) tea.Cmd {
m.logger.Debug("Calling agent", "query", input, "messages", messages) m.logger.Debug("Calling agent", "messages", messages)
return func() tea.Msg { return func() tea.Msg {
resp, err := m.client.Query(input, messages) resp, err := m.client.QueryStream(messages)
return agentResponseMessage{response: resp, err: err} 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}
} }
} }
@@ -373,6 +438,7 @@ func (m uiModel) renameConversation() tea.Cmd {
m.logger.Debug("Renaming conversation", "conversation", m.conversation) m.logger.Debug("Renaming conversation", "conversation", m.conversation)
return func() tea.Msg { return func() tea.Msg {
meta, err := m.client.Rename(m.conversation.Messages) meta, err := m.client.Rename(m.conversation.Messages)
m.startWait()
if err != nil { if err != nil {
return conversationRenamedMessage{title: meta.Title, summary: meta.Summary, err: err} return conversationRenamedMessage{title: meta.Title, summary: meta.Summary, err: err}
} }
@@ -453,9 +519,49 @@ func (m *uiModel) renderMessages() string {
result.WriteString("\n") 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() 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) { func (m *uiModel) setErrorMessage(message string) {
m.statusMessage = m.errorStyle.Render(message) m.statusMessage = m.errorStyle.Render(message)
} }