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