Compare commits
2 Commits
9537237a85
..
master
| Author | SHA1 | Date | |
|---|---|---|---|
| e57b8f428e | |||
| b318f4602b |
@@ -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
@@ -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
@@ -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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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()
|
||||||
|
)`
|
||||||
@@ -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
@@ -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
|
||||||
|
}
|
||||||
|
|||||||
@@ -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}}
|
||||||
|
}
|
||||||
+131
-25
@@ -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,
|
||||||
@@ -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)
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user