From 3d518efd6f52c5e4bde77f2342af5af969263083 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 5 Nov 2023 07:41:43 +0000
Subject: [PATCH] Implement persistence for `lmcli new`

---
 pkg/cli/cli.go |  4 ++++
 pkg/cli/cmd.go | 50 ++++++++++++++++++++++++++++++++++++--------------
 2 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go
index 9e7fa26..59fcc56 100644
--- a/pkg/cli/cli.go
+++ b/pkg/cli/cli.go
@@ -12,3 +12,7 @@ func Fatal(format string, args ...any) {
 	fmt.Fprintf(os.Stderr, format, args...)
 	os.Exit(1)
 }
+
+func Warn(format string, args ...any) {
+	fmt.Fprintf(os.Stderr, format, args...)
+}
diff --git a/pkg/cli/cmd.go b/pkg/cli/cmd.go
index 9d42f94..63c69a9 100644
--- a/pkg/cli/cmd.go
+++ b/pkg/cli/cmd.go
@@ -71,29 +71,53 @@ var newCmd = &cobra.Command{
 		messageContents, err := InputFromEditor("# What would you like to say?\n", "message.*.md")
 		if err != nil {
 			Fatal("Failed to get input: %v\n", err)
-			return
 		}
 
 		if messageContents == "" {
 			Fatal("No message was provided.\n")
-			return
 		}
 
-		fmt.Printf("> %s\n", messageContents)
-
-		messages := []Message{
-			{
-				Role:            "user",
-				OriginalContent: messageContents,
-			},
+		// TODO: set title if --title provided, otherwise defer for later(?)
+		conversation := Conversation{}
+		err = store.SaveConversation(&conversation)
+		if err != nil {
+			Fatal("Could not save new conversation: %v\n", err)
 		}
 
+		message := Message{
+			ConversationID:  conversation.ID,
+			Role:            "user",
+			OriginalContent: messageContents,
+		}
+		err = store.SaveMessage(&message)
+		if err != nil {
+			Warn("Could not save message: %v\n", err)
+		}
+
+		const system = "You are a helpful assistant."
+		fmt.Printf("<System>\n\n%s\n\n", system)
+		fmt.Printf("<You>\n\n%s\n\n", messageContents)
+		fmt.Print("<Assistant>\n\n")
+
 		receiver := make(chan string)
-		go HandleDelayedResponse(receiver)
-		err = CreateChatCompletionStream("You are a helpful assistant.", messages, receiver)
+		response := make(chan string)
+		go func() {
+			response <- HandleDelayedResponse(receiver)
+		}()
+
+		err = CreateChatCompletionStream(system, []Message{message}, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
-			return
+		}
+
+		reply := Message{
+			ConversationID:  conversation.ID,
+			Role:            "assistant",
+			OriginalContent: <-response,
+		}
+		err = store.SaveMessage(&reply)
+		if err != nil {
+			Fatal("Could not save reply: %v\n", err)
 		}
 
 		fmt.Println()
@@ -108,7 +132,6 @@ var promptCmd = &cobra.Command{
 		message := strings.Join(args, " ")
 		if len(strings.Trim(message, " \t\n")) == 0 {
 			Fatal("No message was provided.\n")
-			return
 		}
 
 		messages := []Message{
@@ -123,7 +146,6 @@ var promptCmd = &cobra.Command{
 		err := CreateChatCompletionStream("You are a helpful assistant.", messages, receiver)
 		if err != nil {
 			Fatal("%v\n", err)
-			return
 		}
 
 		fmt.Println()