grokkit/cmd/chat.go

package cmd

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"gmgauthier.com/grokkit/config"
	"gmgauthier.com/grokkit/internal/grok"
)

type ChatHistory struct {
	Messages []map[string]string `json:"messages"`
}

func loadChatHistory() []map[string]string {
	histFile := getChatHistoryFile()
	data, err := os.ReadFile(histFile)
	if err != nil {
		return nil
	}

	var hist ChatHistory
	if err := json.Unmarshal(data, &hist); err != nil {
		return nil
	}
	return hist.Messages
}

func saveChatHistory(messages []map[string]string) error {
	histFile := getChatHistoryFile()
	hist := ChatHistory{Messages: messages}
	data, err := json.MarshalIndent(hist, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(histFile, data, 0644)
}

func getChatHistoryFile() string {
	configFile := viper.GetString("chat.history_file")
	if configFile != "" {
		return configFile
	}

	home, _ := os.UserHomeDir()
	if home == "" {
		home = "."
	}
	histDir := filepath.Join(home, ".config", "grokkit")
	_ = os.MkdirAll(histDir, 0755) // Ignore error, WriteFile will catch it
	return filepath.Join(histDir, "chat_history.json")
}

var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Simple interactive CLI chat with Grok (full history + streaming)",
	Run: func(cmd *cobra.Command, args []string) {
		modelFlag, _ := cmd.Flags().GetString("model")
		model := config.GetModel(modelFlag)

		client := grok.NewClient()

		// Strong system prompt to lock in correct model identity
		systemPrompt := map[string]string{
			"role":    "system",
			"content": fmt.Sprintf("You are Grok 4, the latest and most powerful model from xAI (2026). You are currently running as `%s`. Be helpful, truthful, and a little irreverent. Never claim to be an older model.", model),
		}

		// Load history or start fresh
		history := loadChatHistory()
		if history == nil {
			history = []map[string]string{systemPrompt}
		} else {
			// Update system prompt in loaded history
			if len(history) > 0 && history[0]["role"] == "system" {
				history[0] = systemPrompt
			} else {
				history = append([]map[string]string{systemPrompt}, history...)
			}
		}

		color.Cyan("┌──────────────────────────────────────────────────────────────┐")
		color.Cyan("│ Grokkit Chat  —  Model: %s                                 │", model)
		color.Cyan("│ Type /quit or Ctrl+C to exit                                 │")
		color.Cyan("└──────────────────────────────────────────────────────────────┘\n")

		scanner := bufio.NewScanner(os.Stdin)

		for {
			color.Yellow("You > ")
			if !scanner.Scan() {
				break
			}

			input := strings.TrimSpace(scanner.Text())
			if input == "" {
				continue
			}
			if input == "/quit" || input == "/q" || input == "exit" {
				color.Cyan("\nGoodbye 👋\n")
				break
			}

			history = append(history, map[string]string{"role": "user", "content": input})

			color.Green("Grok > ")
			reply := client.Stream(history, model)

			history = append(history, map[string]string{"role": "assistant", "content": reply})

			// Save history after each exchange
			_ = saveChatHistory(history)

			fmt.Println()
		}
	},
}