feat: refactor chat history storage and enhance conversation handling

- Added a new dependency `github.com/charmbracelet/x/exp/ordered v0.1.0` to `go.mod`.
- Removed unused imports (`context` and `convo`) from `cli.go`.
- Simplified the `postRunHook` function by removing unnecessary conversation storage logic.
- Renamed methods in `chat_history_store.go` from `Persistent` and `Invalidate` to `PersistentMessages` and `InvalidateMessages` respectively.
- Added directory creation logic before saving messages in `chat_history_store.go`.
- Updated test cases in `chat_history_store_test.go` to use the renamed methods.
- Added a new test case for reading messages in `chat_history_store_test.go`.
- Updated the `ChatMessageHistory` interface in `conversation.go` to reflect method name changes.
- Exported the `Sha1reg` variable and added a `MatchSha1` function in `sha.go`.
- Updated the SQLite store path to include a `conversations` subdirectory in `convo_store.go`.
- Added a `GetConvoStore` method to the `Engine` struct in `engine.go`.
- Modified `CreateStreamCompletion` to accept a `context.Context` parameter and added chat context setup logic.
- Added new flags (`show`, `show-last`, `continue`, `continue-last`, `title`) in `basic_flags.go`.
- Added new configuration fields (`ContinueLast`, `Continue`, `Title`, `Show`, `ShowLast`, `CacheReadFromID`, `CacheWriteToID`, `CacheWriteToTitle`) in `config.go`.
- Added logic to handle conversation saving and retrieval in `chat.go`.
- Added a `saveConversation` function to persist conversations in `chat.go`.
- Added utility functions (`lastPrompt`, `firstLine`) to extract the first line of the last human message in `chat.go`.

Signed-off-by: codiing-hui <wecoding@yeah.net>

Author: coding-hui

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/charmbracelet/bubbletea v1.2.5-0.20241205214244-9306010a31ee
 	github.com/charmbracelet/glamour v0.8.0
 	github.com/charmbracelet/lipgloss v1.0.0
+	github.com/charmbracelet/x/exp/ordered v0.1.0
 	github.com/coding-hui/common v0.8.7
 	github.com/coding-hui/go-prompt v0.2.8
 	github.com/coding-hui/wecoding-sdk-go v0.8.7

go.sum

@@ -58,6 +58,8 @@ github.com/charmbracelet/x/ansi v0.4.5 h1:LqK4vwBNaXw2AyGIICa5/29Sbdq58GbGdFngSe
 github.com/charmbracelet/x/ansi v0.4.5/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b h1:MnAMdlwSltxJyULnrYbkZpp4k58Co7Tah3ciKhSNo0Q=
 github.com/charmbracelet/x/exp/golden v0.0.0-20240815200342-61de596daa2b/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
+github.com/charmbracelet/x/exp/ordered v0.1.0 h1:55/qLwjIh0gL0Vni+QAWk7T/qRVP6sBf+2agPBgnOFE=
+github.com/charmbracelet/x/exp/ordered v0.1.0/go.mod h1:5UHwmG+is5THxMyCJHNPCn2/ecI07aKNrW+LcResjJ8=
 github.com/charmbracelet/x/term v0.2.1 h1:AQeHeLZ1OqSXhrAWpYUtZyX1T3zVxfpZuEQMIQaGIAQ=
 github.com/charmbracelet/x/term v0.2.1/go.mod h1:oQ4enTYFV7QN4m0i9mzHrViD7TQKvNEEkHUMCmsxdUg=
 github.com/coding-hui/common v0.8.7 h1:f9iHZcdQLgRFW/nIJHVfBINuzzCKsUcwVIhfboxXe2s=

internal/cli/cli.go

@@ -6,7 +6,6 @@
 package cli
 
 import (
-	"context"
 	"flag"
 	"io"
 	"os"
@@ -28,7 +27,6 @@ import (
 	"github.com/coding-hui/ai-terminal/internal/cli/manpage"
 	"github.com/coding-hui/ai-terminal/internal/cli/review"
 	"github.com/coding-hui/ai-terminal/internal/cli/version"
-	"github.com/coding-hui/ai-terminal/internal/convo"
 	"github.com/coding-hui/ai-terminal/internal/errbook"
 	"github.com/coding-hui/ai-terminal/internal/options"
 	"github.com/coding-hui/ai-terminal/internal/util/debug"
@@ -142,14 +140,5 @@ func postRunHook(cfg *options.Config) error {
 	if err := flushProfiling(); err != nil {
 		return err
 	}
-
-	convoStore, _ := convo.GetConversationStore(cfg)
-	if !cfg.NoCache && convoStore != nil {
-		err := convoStore.SaveConversation(context.Background(), cfg.ConversationID, "convo", cfg.Model)
-		if err != nil {
-			return err
-		}
-	}
-
 	return nil
 }

internal/convo/chat_history_store.go

@@ -51,15 +51,15 @@ func (h *SimpleChatHistoryStore) AddMessage(_ context.Context, convoID string, m
 
 func (h *SimpleChatHistoryStore) SetMessages(ctx context.Context, convoID string, messages []llms.ChatMessage) error {
 	h.messages[convoID] = messages
-	if err := h.Invalidate(ctx, convoID); err != nil && !errors.Is(err, os.ErrNotExist) {
+	if err := h.InvalidateMessages(ctx, convoID); err != nil && !errors.Is(err, os.ErrNotExist) {
 		return err
 	}
-	return h.Persistent(ctx, convoID, h.messages[convoID])
+	return h.PersistentMessages(ctx, convoID)
 }
 
 func (h *SimpleChatHistoryStore) Messages(_ context.Context, convoID string) ([]llms.ChatMessage, error) {
 	if !h.loaded[convoID] {
-		if err := h.load(convoID); err != nil {
+		if err := h.load(convoID); err != nil && !errors.Is(err, os.ErrNotExist) {
 			return nil, err
 		}
 		h.loaded[convoID] = true
@@ -91,19 +91,24 @@ func (h *SimpleChatHistoryStore) load(convoID string) error {
 	return nil
 }
 
-func (h *SimpleChatHistoryStore) Persistent(_ context.Context, convoID string, messages []llms.ChatMessage) error {
+func (h *SimpleChatHistoryStore) PersistentMessages(_ context.Context, convoID string) error {
 	if convoID == "" {
 		return fmt.Errorf("write: %w", errInvalidID)
 	}
 
+	// Ensure directory exists
+	if err := os.MkdirAll(h.dir, 0755); err != nil {
+		return fmt.Errorf("create directory: %w", err)
+	}
+
 	file, err := os.Create(filepath.Join(h.dir, convoID+cacheExt))
 	if err != nil {
 		return fmt.Errorf("write: %w", err)
 	}
 	defer file.Close() //nolint:errcheck
 
 	var rawMessages []llms.ChatMessageModel
-	for _, v := range messages {
+	for _, v := range h.messages[convoID] {
 		if v != nil {
 			rawMessages = append(rawMessages, llms.ConvertChatMessageToModel(v))
 		}
@@ -115,7 +120,7 @@ func (h *SimpleChatHistoryStore) Persistent(_ context.Context, convoID string, m
 	return nil
 }
 
-func (h *SimpleChatHistoryStore) Invalidate(_ context.Context, convoID string) error {
+func (h *SimpleChatHistoryStore) InvalidateMessages(_ context.Context, convoID string) error {
 	if convoID == "" {
 		return fmt.Errorf("delete: %w", errInvalidID)
 	}

internal/convo/chat_history_store_test.go

@@ -13,19 +13,28 @@ import (
 func TestStore(t *testing.T) {
 	convoID := NewConversationID()
 
+	t.Run("read", func(t *testing.T) {
+		store := NewSimpleChatHistoryStore("/Users/bytedance/Codes/ai-terminal/bin/conversations/")
+		err := store.load("3b7e9fb1ae86660c32cef53346a3230f4bfc8797")
+		require.NoError(t, err)
+		messages, err := store.Messages(context.Background(), "3b7e9fb1ae86660c32cef53346a3230f4bfc8797")
+		require.NoError(t, err)
+		require.Len(t, messages, 1)
+	})
+
 	t.Run("read non-existent", func(t *testing.T) {
 		store := NewSimpleChatHistoryStore(t.TempDir())
 		err := store.load("super-fake")
 		require.ErrorIs(t, err, os.ErrNotExist)
 		defer func() {
-			_ = store.Invalidate(context.Background(), convoID)
+			_ = store.InvalidateMessages(context.Background(), convoID)
 		}()
 	})
 
 	t.Run("set messages", func(t *testing.T) {
 		ctx := context.Background()
 		store := NewSimpleChatHistoryStore(t.TempDir())
-		_ = store.Invalidate(context.Background(), convoID)
+		_ = store.InvalidateMessages(context.Background(), convoID)
 		require.NoError(t, store.AddUserMessage(ctx, convoID, "hello"))
 		require.NoError(t, store.AddAIMessage(ctx, convoID, "hi"))
 		require.NoError(t, store.AddAIMessage(ctx, convoID, "bye"))
@@ -36,13 +45,13 @@ func TestStore(t *testing.T) {
 		require.Equal(t, 3, len(messages))
 
 		// After persist, messages should be saved
-		require.NoError(t, store.Persistent(ctx, convoID, messages))
+		require.NoError(t, store.PersistentMessages(ctx, convoID))
 		persistedMessages, err := store.Messages(ctx, convoID)
 		require.NoError(t, err)
 		require.Equal(t, messages, persistedMessages)
 
 		defer func() {
-			_ = store.Invalidate(context.Background(), convoID)
+			_ = store.InvalidateMessages(context.Background(), convoID)
 		}()
 	})
 
@@ -53,21 +62,21 @@ func TestStore(t *testing.T) {
 			llms.HumanChatMessage{Content: "bar"},
 			llms.AIChatMessage{Content: "zoo"},
 		}
-		require.NoError(t, store.Persistent(context.Background(), convoID, messages))
+		require.NoError(t, store.SetMessages(context.Background(), convoID, messages))
 		require.NoError(t, store.load(convoID))
 		require.ElementsMatch(t, messages, store.messages[convoID])
 		defer func() {
-			_ = store.Invalidate(context.Background(), convoID)
+			_ = store.InvalidateMessages(context.Background(), convoID)
 		}()
 	})
 
 	t.Run("delete", func(t *testing.T) {
 		store := NewSimpleChatHistoryStore(t.TempDir())
-		require.NoError(t, store.Persistent(context.Background(), convoID, []llms.ChatMessage{}))
-		require.NoError(t, store.Invalidate(context.Background(), convoID))
+		require.NoError(t, store.PersistentMessages(context.Background(), convoID))
+		require.NoError(t, store.InvalidateMessages(context.Background(), convoID))
 		require.ErrorIs(t, store.load(convoID), os.ErrNotExist)
 		defer func() {
-			_ = store.Invalidate(context.Background(), convoID)
+			_ = store.InvalidateMessages(context.Background(), convoID)
 		}()
 	})
 

internal/convo/conversation.go

@@ -80,10 +80,10 @@ type ChatMessageHistory interface {
 	SetMessages(ctx context.Context, convoID string, messages []llms.ChatMessage) error
 	// Messages retrieves all messages from the store
 	Messages(ctx context.Context, convoID string) ([]llms.ChatMessage, error)
-	// Persistent saves messages to persistent storage
-	Persistent(ctx context.Context, convoID string, messages []llms.ChatMessage) error
-	// Invalidate removes messages from persistent storage
-	Invalidate(ctx context.Context, convoID string) error
+	// PersistentMessages saves messages to persistent storage
+	PersistentMessages(ctx context.Context, convoID string) error
+	// InvalidateMessages removes messages from persistent storage
+	InvalidateMessages(ctx context.Context, convoID string) error
 }
 
 // Store is the interface for chat history convo store.

internal/convo/sha.go

@@ -13,10 +13,14 @@ const (
 	Sha1ReadBlockSize = 4096
 )
 
-var sha1reg = regexp.MustCompile(`\b[0-9a-f]{40}\b`)
+var Sha1reg = regexp.MustCompile(`\b[0-9a-f]{40}\b`)
 
 func NewConversationID() string {
 	b := make([]byte, Sha1ReadBlockSize)
 	_, _ = rand.Read(b)
 	return fmt.Sprintf("%x", sha1.Sum(b)) //nolint: gosec
 }
+
+func MatchSha1(s string) bool {
+	return Sha1reg.MatchString(s)
+}

internal/convo/sqlite3/convo_store.go

@@ -33,7 +33,7 @@ func (s *sqliteStoreFactor) Type() string {
 
 func (s *sqliteStoreFactor) Create(options *options.Config) (convo.Store, error) {
 	return NewSqliteStore(
-		WithDataPath(options.DataStore.CachePath),
+		WithDataPath(filepath.Join(options.DataStore.CachePath, "conversations")),
 		WithConversation(options.ConversationID),
 		WithDBAddress(filepath.Join(options.DataStore.CachePath, "convo.db")),
 	), nil

internal/llm/engine.go

@@ -132,6 +132,10 @@ func (e *Engine) GetChannel() chan StreamCompletionOutput {
 	return e.channel
 }
 
+func (e *Engine) GetConvoStore() convo.Store {
+	return e.convoStore
+}
+
 func (e *Engine) Interrupt() {
 	e.channel <- StreamCompletionOutput{
 		content:    "[Interrupt]",
@@ -169,9 +173,7 @@ func (e *Engine) CreateCompletion(messages []llms.ChatMessage) (*EngineExecOutpu
 	return &output, nil
 }
 
-func (e *Engine) CreateStreamCompletion(messages []llms.ChatMessage) tea.Msg {
-	ctx := context.Background()
-
+func (e *Engine) CreateStreamCompletion(ctx context.Context, messages []llms.ChatMessage) tea.Msg {
 	e.running = true
 
 	streamingFunc := func(ctx context.Context, chunk []byte) error {
@@ -184,6 +186,17 @@ func (e *Engine) CreateStreamCompletion(messages []llms.ChatMessage) tea.Msg {
 		return nil
 	}
 
+	if err := e.setupChatContext(ctx, &messages); err != nil {
+		return err
+	}
+
+	for _, v := range messages {
+		err := e.convoStore.AddMessage(ctx, e.config.CacheWriteToID, v)
+		if err != nil {
+			errbook.HandleError(errbook.Wrap("Failed to add user chat input message to history", err))
+		}
+	}
+
 	messageParts := slices.Map(messages, convert)
 	rsp, err := e.Model.GenerateContent(ctx, messageParts, e.callOptions(streamingFunc)...)
 	if err != nil {
@@ -235,16 +248,25 @@ func (e *Engine) callOptions(streamingFunc ...func(ctx context.Context, chunk []
 	return opts
 }
 
-func (e *Engine) appendUserMessage(content string) {
-	if len(strings.TrimSpace(content)) == 0 {
-		errbook.HandleError(errbook.New("empty input is not allowed."))
-		return
+func (e *Engine) setupChatContext(ctx context.Context, messages *[]llms.ChatMessage) error {
+	store := e.convoStore
+	if store == nil {
+		return errbook.New("no chat history store found")
 	}
-	if e.convoStore != nil {
-		if err := e.convoStore.AddUserMessage(context.Background(), e.config.ConversationID, content); err != nil {
-			errbook.HandleError(errbook.Wrap("failed to add user chat input message to history", err))
+
+	if !e.config.NoCache && e.config.CacheReadFromID != "" {
+		history, err := store.Messages(ctx, e.config.CacheReadFromID)
+		if err != nil {
+			return errbook.Wrap(fmt.Sprintf(
+				"There was a problem reading the cache. Use %s / %s to disable it.",
+				console.StderrStyles().InlineCode.Render("--no-cache"),
+				console.StderrStyles().InlineCode.Render("NO_CACHE"),
+			), err)
 		}
+		*messages = append(*messages, history...)
 	}
+
+	return nil
 }
 
 func (e *Engine) appendAssistantMessage(content string) {

internal/options/basic_flags.go

@@ -30,7 +30,11 @@ func AddBasicFlags(flags *pflag.FlagSet, cfg *Config) {
 	flags.UintVar(&cfg.Fanciness, "fanciness", cfg.Fanciness, console.StdoutStyles().FlagDesc.Render(help["fanciness"]))
 	flags.StringVar(&cfg.LoadingText, "loading-text", cfg.LoadingText, console.StdoutStyles().FlagDesc.Render(help["status-text"]))
 	flags.BoolVar(&cfg.NoCache, "no-cache", cfg.NoCache, console.StdoutStyles().FlagDesc.Render(help["no-cache"]))
-	//flags.BoolVar(&cfg.Dirs, "dirs", false, console.StdoutStyles().FlagDesc.Render(help["dirs"]))
+	flags.StringVarP(&cfg.Show, "show", "s", cfg.Show, console.StdoutStyles().FlagDesc.Render(help["show"]))
+	flags.BoolVarP(&cfg.ShowLast, "show-last", "S", false, console.StdoutStyles().FlagDesc.Render(help["show-last"]))
+	flags.StringVarP(&cfg.Continue, "continue", "c", "", console.StdoutStyles().FlagDesc.Render(help["continue"]))
+	flags.BoolVarP(&cfg.ContinueLast, "continue-last", "C", false, console.StdoutStyles().FlagDesc.Render(help["continue-last"]))
+	flags.StringVarP(&cfg.Title, "title", "T", cfg.Title, console.StdoutStyles().FlagDesc.Render(help["title"]))
 	//flags.StringVarP(&cfg.Role, "role", "R", cfg.Role, console.StdoutStyles().FlagDesc.Render(help["role"]))
 	//flags.BoolVar(&cfg.ListRoles, "list-roles", cfg.ListRoles, console.StdoutStyles().FlagDesc.Render(help["list-roles"]))
 	//flags.StringVar(&cfg.Theme, "theme", "charm", console.StdoutStyles().FlagDesc.Render(help["theme"]))