Feat: add memory [wip]

2024-11-20 20:47:49 +03:00
parent fc517c2c69
commit 5ccad20bd6
8 changed files with 167 additions and 19 deletions
--- a/README.md
+++ b/README.md
@@ -7,7 +7,6 @@
 - delete last message; +
 - edit message? (including from bot); +
 - ability to copy message; +
- aility to copy selected text; (I can do it though vim mode of the terminal, so +)
 - menu with old chats (chat files); +
 - fullscreen textarea option (for long prompt);
 - tab to switch selection between textview and textarea (input and chat); +
@@ -20,3 +19,7 @@
 - bot responding (or haninging) blocks everything; +
 - programm requires history folder, but it is .gitignore; +
 - at first run chat table does not exist; run migrations sql on startup; +
+- Tab is needed to copy paste text into textarea box, use shift+tab to switch focus; (changed tp pgup) +
+- delete last msg: can have unexpected behavior (deletes what appears to be two messages);
+- EOF from llama, possibly broken json in request;
+- chat upsert does not work;
--- a/main.go
+++ b/main.go
@@ -16,7 +16,7 @@ var (
 	editMode      = false
 	botMsg        = "no"
 	selectedIndex = int(-1)
-	indexLine     = "Esc: send msg; Tab: switch focus; F1: manage chats; F2: regen last; F3:delete last msg; F4: edit msg; F5: toggle system; F6: interrupt bot resp; Row: [yellow]%d[white], Column: [yellow]%d; bot resp mode: %v"
+	indexLine     = "Esc: send msg; PgUp/Down: switch focus; F1: manage chats; F2: regen last; F3:delete last msg; F4: edit msg; F5: toggle system; F6: interrupt bot resp; Row: [yellow]%d[white], Column: [yellow]%d; bot resp mode: %v"
 	focusSwitcher = map[tview.Primitive]tview.Primitive{}
 )

@@ -56,7 +56,7 @@ func main() {
 		if fromRow == toRow && fromColumn == toColumn {
 			position.SetText(fmt.Sprintf(indexLine, fromRow, fromColumn, botRespMode))
 		} else {
-			position.SetText(fmt.Sprintf("Esc: send msg; Tab: switch focus; F1: manage chats; F2: regen last; F3:delete last msg; F4: edit msg; F5: toggle system; F6: interrupt bot resp; Row: [yellow]%d[white], Column: [yellow]%d[white] - [red]To[white] Row: [yellow]%d[white], To Column: [yellow]%d; bot resp mode: %v", fromRow, fromColumn, toRow, toColumn, botRespMode))
+			position.SetText(fmt.Sprintf("Esc: send msg; PgUp/Down: switch focus; F1: manage chats; F2: regen last; F3:delete last msg; F4: edit msg; F5: toggle system; F6: interrupt bot resp; Row: [yellow]%d[white], Column: [yellow]%d[white] - [red]To[white] Row: [yellow]%d[white], To Column: [yellow]%d; bot resp mode: %v", fromRow, fromColumn, toRow, toColumn, botRespMode))
 		}
 	}
 	chatOpts := []string{"cancel", "new"}
@@ -172,7 +172,7 @@ func main() {
 			return nil
 		}
 		if event.Key() == tcell.KeyF3 {
-			// modal window with input field
+			// delete last msg
 			chatBody.Messages = chatBody.Messages[:len(chatBody.Messages)-1]
 			textView.SetText(chatToText(showSystemMsgs))
 			botRespMode = false // hmmm; is that correct?
@@ -195,6 +195,15 @@ func main() {
 			return nil
 		}
 		if event.Key() == tcell.KeyF7 {
+			// copy msg to clipboard
+			editMode = false
+			m := chatBody.Messages[len(chatBody.Messages)-1]
+			copyToClipboard(m.Content)
+			notification := fmt.Sprintf("msg '%s' was copied to the clipboard", m.Content[:30])
+			notifyUser("copied", notification)
+			return nil
+		}
+		if event.Key() == tcell.KeyF8 {
 			// copy msg to clipboard
 			editMode = false
 			pages.AddPage("getIndex", indexPickWindow, true, true)
@@ -215,9 +224,10 @@ func main() {
 			go chatRound(msgText, userRole, textView)
 			return nil
 		}
-		if event.Key() == tcell.KeyTab {
+		if event.Key() == tcell.KeyPgUp || event.Key() == tcell.KeyPgDn {
 			currentF := app.GetFocus()
 			app.SetFocus(focusSwitcher[currentF])
+			return nil
 		}
 		if isASCII(string(event.Rune())) && !botRespMode {
 			// botRespMode = false
--- a/models/db.go
+++ b/models/db.go
@@ -21,7 +21,16 @@ func (c Chat) ToHistory() ([]MessagesStory, error) {
 	return resp, nil
 }

+/*
+memories should have two key system
+to be able to store different perspectives
+agent -> topic -> data
+agent is somewhat similar to a char
+*/
 type Memory struct {
+	Agent     string    `db:"agent" json:"agent"`
 	Topic     string    `db:"topic" json:"topic"`
-	Data  string `db:"data" json:"data"`
+	Mind      string    `db:"mind" json:"mind"`
+	CreatedAt time.Time `db:"created_at" json:"created_at"`
+	UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
 }
--- a/session.go
+++ b/session.go
@@ -69,14 +69,14 @@ func loadHistoryChat(chatName string) ([]models.MessagesStory, error) {
 }

 func loadOldChatOrGetNew() []models.MessagesStory {
-	// find last chat
-	chat, err := store.GetLastChat()
 	newChat := &models.Chat{
 		ID:        0,
 		CreatedAt: time.Now(),
 		UpdatedAt: time.Now(),
 	}
-	newChat.Name = fmt.Sprintf("%d_%v", chat.ID, chat.CreatedAt.Unix())
+	newChat.Name = fmt.Sprintf("%d_%v", newChat.ID, newChat.CreatedAt.Unix())
+	// find last chat
+	chat, err := store.GetLastChat()
 	if err != nil {
 		logger.Warn("failed to load history chat", "error", err)
 		activeChatName = newChat.Name
--- a/storage/memory.go
+++ b/storage/memory.go
@@ -0,0 +1,44 @@
+package storage
+
+import "elefant/models"
+
+type Memories interface {
+	Memorise(m *models.Memory) (*models.Memory, error)
+	Recall(agent, topic string) (string, error)
+	RecallTopics(agent string) ([]string, error)
+}
+
+func (p ProviderSQL) Memorise(m *models.Memory) (*models.Memory, error) {
+	query := "INSERT INTO memories (agent, topic, mind) VALUES (:agent, :topic, :mind) RETURNING *;"
+	stmt, err := p.db.PrepareNamed(query)
+	if err != nil {
+		return nil, err
+	}
+	defer stmt.Close()
+	var memory models.Memory
+	err = stmt.Get(&memory, m)
+	if err != nil {
+		return nil, err
+	}
+	return &memory, nil
+}
+
+func (p ProviderSQL) Recall(agent, topic string) (string, error) {
+	query := "SELECT mind FROM memories WHERE agent = $1 AND topic = $2"
+	var mind string
+	err := p.db.Get(&mind, query, agent, topic)
+	if err != nil {
+		return "", err
+	}
+	return mind, nil
+}
+
+func (p ProviderSQL) RecallTopics(agent string) ([]string, error) {
+	query := "SELECT DISTINCT topic FROM memories WHERE agent = $1"
+	var topics []string
+	err := p.db.Select(&topics, query, agent)
+	if err != nil {
+		return nil, err
+	}
+	return topics, nil
+}
--- a/storage/migrations/001_init.up.sql
+++ b/storage/migrations/001_init.up.sql
@@ -1,7 +1,16 @@
-CREATE TABLE IF NOT EXISTS chat (
+CREATE TABLE IF NOT EXISTS chats (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
-    msgs TEXT NOT NULL,  -- Store messages as a comma-separated string
+    msgs TEXT NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
 );
+
+CREATE TABLE IF NOT EXISTS memories (
+    agent TEXT NOT NULL,
+    topic TEXT NOT NULL,
+    mind TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (agent, topic)
+);
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -23,26 +23,26 @@ type ProviderSQL struct {

 func (p ProviderSQL) ListChats() ([]models.Chat, error) {
 	resp := []models.Chat{}
-	err := p.db.Select(&resp, "SELECT * FROM chat;")
+	err := p.db.Select(&resp, "SELECT * FROM chats;")
 	return resp, err
 }

 func (p ProviderSQL) GetChatByID(id uint32) (*models.Chat, error) {
 	resp := models.Chat{}
-	err := p.db.Get(&resp, "SELECT * FROM chat WHERE id=$1;", id)
+	err := p.db.Get(&resp, "SELECT * FROM chats WHERE id=$1;", id)
 	return &resp, err
 }

 func (p ProviderSQL) GetLastChat() (*models.Chat, error) {
 	resp := models.Chat{}
-	err := p.db.Get(&resp, "SELECT * FROM chat ORDER BY updated_at DESC LIMIT 1")
+	err := p.db.Get(&resp, "SELECT * FROM chats ORDER BY updated_at DESC LIMIT 1")
 	return &resp, err
 }

 func (p ProviderSQL) UpsertChat(chat *models.Chat) (*models.Chat, error) {
 	// Prepare the SQL statement
 	query := `
-        INSERT OR REPLACE INTO chat (id, name, msgs, created_at, updated_at)
+        INSERT OR REPLACE INTO chats (id, name, msgs, created_at, updated_at)
        VALUES (:id, :name, :msgs, :created_at, :updated_at)
        RETURNING *;`
 	stmt, err := p.db.PrepareNamed(query)
@@ -56,7 +56,7 @@ func (p ProviderSQL) UpsertChat(chat *models.Chat) (*models.Chat, error) {
 }

 func (p ProviderSQL) RemoveChat(id uint32) error {
-	query := "DELETE FROM chat WHERE ID = $1;"
+	query := "DELETE FROM chats WHERE ID = $1;"
 	_, err := p.db.Exec(query, id)
 	return err
 }
--- a/storage/storage_test.go
+++ b/storage/storage_test.go
@@ -2,6 +2,9 @@ package storage

 import (
 	"elefant/models"
+	"fmt"
+	"log/slog"
+	"os"
 	"testing"
 	"time"

@@ -9,6 +12,76 @@ import (
 	"github.com/jmoiron/sqlx"
 )

+func TestMemories(t *testing.T) {
+	db, err := sqlx.Open("sqlite", ":memory:")
+	if err != nil {
+		t.Fatalf("Failed to open SQLite in-memory database: %v", err)
+	}
+	defer db.Close()
+	_, err = db.Exec(`
+CREATE TABLE IF NOT EXISTS memories (
+    agent TEXT NOT NULL,
+    topic TEXT NOT NULL,
+    mind TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (agent, topic)
+);`)
+	if err != nil {
+		t.Fatalf("Failed to create chat table: %v", err)
+	}
+	provider := ProviderSQL{
+		db:     db,
+		logger: slog.New(slog.NewJSONHandler(os.Stdout, nil)),
+	}
+	// Create a sample memory for testing
+	sampleMemory := &models.Memory{
+		Agent:     "testAgent",
+		Topic:     "testTopic",
+		Mind:      "testMind",
+		CreatedAt: time.Now(),
+		UpdatedAt: time.Now(),
+	}
+	cases := []struct {
+		memory *models.Memory
+	}{
+		{memory: sampleMemory},
+	}
+	for i, tc := range cases {
+		t.Run(fmt.Sprintf("run_%d", i), func(t *testing.T) {
+			// Recall topics: get no rows
+			topics, err := provider.RecallTopics(tc.memory.Agent)
+			if err != nil {
+				t.Fatalf("Failed to recall topics: %v", err)
+			}
+			if len(topics) != 0 {
+				t.Fatalf("Expected no topics, got: %v", topics)
+			}
+			// Memorise
+			_, err = provider.Memorise(tc.memory)
+			if err != nil {
+				t.Fatalf("Failed to memorise: %v", err)
+			}
+			// Recall topics: has topics
+			topics, err = provider.RecallTopics(tc.memory.Agent)
+			if err != nil {
+				t.Fatalf("Failed to recall topics: %v", err)
+			}
+			if len(topics) == 0 {
+				t.Fatalf("Expected topics, got none")
+			}
+			// Recall
+			content, err := provider.Recall(tc.memory.Agent, tc.memory.Topic)
+			if err != nil {
+				t.Fatalf("Failed to recall: %v", err)
+			}
+			if content != tc.memory.Mind {
+				t.Fatalf("Expected content: %v, got: %v", tc.memory.Mind, content)
+			}
+		})
+	}
+}
+
 func TestChatHistory(t *testing.T) {
 	// Create an in-memory SQLite database
 	db, err := sqlx.Open("sqlite", ":memory:")
@@ -18,10 +91,10 @@ func TestChatHistory(t *testing.T) {
 	defer db.Close()
 	// Create the chat table
 	_, err = db.Exec(`
-		CREATE TABLE chat (
+		CREATE TABLE chats (
 	    id INTEGER PRIMARY KEY AUTOINCREMENT,
 	    name TEXT NOT NULL,
-	    msgs TEXT NOT NULL,  -- Store messages as a comma-separated string
+	    msgs TEXT NOT NULL, 
 	    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
 	    updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
 		);`)