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
Timeout int
BatchSize int
Interval int
}
type MemoryConfig struct {
@@ -152,6 +153,7 @@ func defaultConfig() *Config {
Model: "",
BatchSize: 64,
Timeout: 60000,
Interval: 60,
},
}
}
@@ -185,6 +187,7 @@ func overrideFromEnv(cfg *Config) {
}{
{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)
+61 -19
View File
@@ -7,6 +7,7 @@ import (
"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"
@@ -32,6 +33,35 @@ func (e EmbedStore) TableName() string {
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")
@@ -47,17 +77,16 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
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
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 {
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
}
@@ -65,7 +94,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd
logger.Debug("Table found", "table", store.TableName())
} else {
logger.Debug("Table not found. Creating", "table", store.TableName())
script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim())
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
@@ -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 {
rows, err := e.pool.Query(ctx, fmt.Sprintf(`SELECT c.id, c.content, c.content_hash
FROM message_chunks c
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()),
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 {
@@ -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 {
query := fmt.Sprintf(`
query := fmt.Sprintf(
`
INSERT INTO %s (chunk_id, embedding, content_hash)
VALUES ($1, $2, $3)
VALUES ($1, $2, $3)
ON CONFLICT (chunk_id) DO UPDATE
SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash`, e.TableName())
SET embedding = EXCLUDED.embedding,
content_hash = EXCLUDED.content_hash
`,
e.TableName())
batch := &pgx.Batch{}
for i, chunk := range chunks {
@@ -152,3 +190,7 @@ func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32
}
return nil
}
func (e EmbedStore) Close() {
e.pool.Close()
}
+84 -54
View File
@@ -11,6 +11,7 @@ import (
"strings"
"git.estatecloud.org/radumaco/souvenir/config"
"git.estatecloud.org/radumaco/souvenir/db"
"git.estatecloud.org/radumaco/souvenir/model"
"github.com/gopsql/pgx"
"github.com/jackc/pgx/v5/pgxpool"
@@ -24,36 +25,6 @@ type DbClient struct {
func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
logger := *slog.Default().With("Component", "History")
const createConversationTable = `
CREATE TABLE IF NOT EXISTS conversations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
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 {
return nil, err
@@ -67,17 +38,23 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) {
tableName string
script string
}{
{"conversations", createConversationTable},
{"messages", createMessageTable},
{"message_chunks", createMessageChunkTable},
{"conversations", db.ConversationTableScript},
{"messages", db.MessageTableScript},
{"message_chunks", db.MessageChunkTableScript},
}
for _, table := range tables {
var exists bool
if err := pool.QueryRow(ctx, `SELECT EXISTS
(SELECT 1 FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name =$1)`, table.tableName).Scan(&exists); err != nil {
if err := pool.QueryRow(ctx,
`
SELECT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name =$1
)
`,
table.tableName).Scan(&exists); err != nil {
logger.Error("Error while checking table", "table", table.tableName, "error", 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) {
cl.logger.Debug("Searching for conversation", "id", id)
var conv model.Conversation
if err := cl.pool.QueryRow(ctx, "SELECT id, title, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil {
if err := cl.pool.QueryRow(ctx,
`
SELECT id, title, summary
FROM conversations
WHERE id = $1
`, id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil {
cl.logger.Error("Could not fetch conversation", "id", id, "error", err)
return conv, err
}
rows, err := cl.pool.Query(ctx, `SELECT id, role, content, seq
FROM messages
WHERE conversation_id = $1 ORDER BY seq`, id)
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
@@ -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) {
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 {
cl.logger.Error("Could not fetch conversations", "error", 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)
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 {
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 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)
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)
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) {
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 {
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`,
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")
@@ -213,7 +228,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
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)`,
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")
@@ -223,8 +246,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess
continue
}
if _, err := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash)
VALUES ($1, $2, $3, $4)`, conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil {
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)
}
}
+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 (
"context"
"errors"
"log/slog"
"os"
"time"
tea "charm.land/bubbletea/v2"
"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"
)
func main() {
ctx := context.Background()
config, err := config.Load(".config.json")
@@ -28,6 +31,10 @@ func main() {
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 {
@@ -35,3 +42,46 @@ func main() {
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
}
+75 -38
View File
@@ -8,6 +8,7 @@ import (
"log/slog"
"net/http"
"strconv"
"strings"
"time"
config "git.estatecloud.org/radumaco/souvenir/config"
@@ -29,6 +30,16 @@ 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 //
@@ -129,7 +140,12 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) {
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 {
@@ -137,63 +153,45 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) {
return RenameResponse{}, err
}
var meta RenameResponse
if err := json.Unmarshal([]byte(out[0].Content), &meta); err != nil {
cl.logger.Error("Could not parse the rename response", "content", out[0].Content)
return meta, err
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
}
}
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{
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) {
serialized, err := json.Marshal(requestBody)
if err != nil {
cl.logger.Error("There was an error marshelling the request body",
"request", requestBody,
"error", err)
return []model.Message{}, err
}
resp, err := cl.CallApi(requestBody)
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 []model.Message{}, 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
cl.logger.Error("There was an error during the API call", "error", err)
return nil, err
}
defer resp.Body.Close()
cl.logger.Debug("LLM Call returned", "response", resp)
data, err := io.ReadAll(resp.Body)
if err != nil {
cl.logger.Error("There was an error reading the response body", "error", err.Error())
return []model.Message{}, err
}
if resp.StatusCode != 200 && resp.StatusCode != 202 {
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)
return nil, err
}
cl.logger.Debug("Received data", "data", data)
@@ -215,3 +213,42 @@ func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) {
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}}
}
+166 -60
View File
@@ -7,6 +7,7 @@ import (
"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"
@@ -26,30 +27,40 @@ const (
)
type uiModel struct {
logger slog.Logger
viewport viewport.Model
conversation model.Conversation
textarea textarea.Model
focus focusState
picker picker
senderStyle lipgloss.Style
agentStyle lipgloss.Style
errorStyle lipgloss.Style
statusStyle lipgloss.Style
client llm.ChatClient
statusMessage string
commands commandList
history history.DbClient
ctx context.Context
width int
height int
waiting bool
err error
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 agentResponseMessage struct {
response []model.Message
err error
type streamEventMessage struct {
event llm.StreamEvent
}
type streamClosedMessage struct {
}
type streamStartedMessage struct {
ch <-chan llm.StreamEvent
}
type conversationSavedMessage struct {
@@ -69,7 +80,7 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl
ta.SetVirtualCursor(false)
ta.Focus()
ta.Prompt = " "
ta.Prompt = " "
ta.CharLimit = 280
ta.SetWidth(30)
@@ -88,19 +99,25 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl
ta.KeyMap.InsertNewline.SetEnabled(false)
sp := spinner.New()
sp.Spinner = spinner.Dot
ui := uiModel{
textarea: ta,
conversation: model.Conversation{},
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.NewLLMClient(config),
focus: focusChat,
history: client,
ctx: ctx,
logger: *slog.Default().With("Component", "TUI"),
err: nil,
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)
@@ -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
for _, c := range m.commands.available {
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.textarea.Reset()
m.viewport.GotoBottom()
m.waiting = true
return m, tea.Batch(m.callAgent(input, m.conversation.Messages), m.saveConversation())
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)
@@ -242,7 +259,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case conversationRenamedMessage:
m.logger.Debug("Received messageRenamedMessage")
m.waiting = false
m.stopWait()
if msg.err != nil {
m.logger.Error("There was an error", "error", msg.err)
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
return m, nil
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)
m.setErrorMessage(msg.err.Error())
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)
}
for _, r := range resp {
r.Seq = len(m.conversation.Messages) + 1
m.logger.Debug("Appending response message", "message", r)
m.conversation.Messages = append(m.conversation.Messages, r)
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, 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:
switch m.focus {
case focusModels:
m.logger.Debug("Received pickerChosenMsg", "model", msg.id)
m.client.Cfg.Llm.Model = msg.id
m.waiting = false
m.focus = focusChat
case focusHistory:
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.conversation = conv
m.waiting = false
m.focus = focusChat
var messages = m.renderMessages()
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:
m.logger.Debug("Received pickerDismissMsg")
m.waiting = false
m.stopWait()
m.focus = focusChat
case conversationSavedMessage:
m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation)
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())
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)
@@ -313,7 +358,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, cmd
}
return m, nil
return m, spinnerCmd
}
func (m uiModel) View() tea.View {
@@ -327,6 +372,9 @@ func (m uiModel) View() tea.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())
}
@@ -344,6 +392,9 @@ func (m uiModel) View() tea.View {
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)
}
@@ -353,11 +404,25 @@ func (m uiModel) View() tea.View {
return view
}
func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd {
m.logger.Debug("Calling agent", "query", input, "messages", messages)
func (m uiModel) callAgent(messages []model.Message) tea.Cmd {
m.logger.Debug("Calling agent", "messages", messages)
return func() tea.Msg {
resp, err := m.client.Query(input, messages)
return agentResponseMessage{response: resp, err: err}
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}
}
}
@@ -373,6 +438,7 @@ 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}
}
@@ -453,9 +519,49 @@ func (m *uiModel) renderMessages() string {
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)
}