diff --git a/README.md b/README.md index 88bd021..157f52c 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ Ingest can also pass the prompt directly to an LLM such as Ollama for processing - Copy output to clipboard (when available) - Export to file or print to console - Optional JSON output +- Optionally save output to a file in ~/ingest - Shell completions for Bash, Zsh, and Fish ## Installation @@ -87,6 +88,12 @@ You can also provide individual files or multiple paths: ingest /path/to/file /path/to/directory ``` +Save output to to ~/ingest/.md: + +```shell +ingest --save /path/to/project +``` + ### VRAM Estimation and Model Compatibility Ingest includes a feature to estimate VRAM requirements and check model compatibility using the [Gollama](https://github.com/sammcj/gollama)'s vramestimator package. This helps you determine if your generated content will fit within the specified model, VRAM, and quantisation constraints. diff --git a/config/config.go b/config/config.go index ad0de9e..7771acc 100644 --- a/config/config.go +++ b/config/config.go @@ -19,6 +19,7 @@ type OllamaConfig struct { type Config struct { Ollama []OllamaConfig `json:"ollama"` LLM LLMConfig `json:"llm"` + AutoSave bool `json:"auto_save"` } type LLMConfig struct { @@ -89,6 +90,7 @@ func createDefaultConfig(configPath string) (*Config, error) { Model: "llama3.1:8b-instruct-q6_K", MaxTokens: 2048, }, + AutoSave: false, } err := os.MkdirAll(filepath.Dir(configPath), 0750) diff --git a/main.go b/main.go index ee28ab7..a0ea8ab 100644 --- a/main.go +++ b/main.go @@ -84,7 +84,7 @@ func init() { rootCmd.Flags().BoolVar(&printDefaultExcludes, "print-default-excludes", false, "Print the default exclude patterns") rootCmd.Flags().BoolVar(&printDefaultTemplate, "print-default-template", false, "Print the default template") rootCmd.Flags().BoolVar(&relativePaths, "relative-paths", false, "Use relative paths instead of absolute paths, including the parent directory") - rootCmd.Flags().BoolVar(&report, "report", true, "Report the top 5 largest files included in the output") + rootCmd.Flags().BoolVar(&report, "report", true, "Report the top 10 largest files included in the output") rootCmd.Flags().BoolVar(&tokens, "tokens", true, "Display the token count of the generated prompt") rootCmd.Flags().BoolVarP(&diff, "diff", "d", false, "Include git diff") rootCmd.Flags().BoolVarP(&lineNumber, "line-number", "l", false, "Add line numbers to the source code") @@ -99,6 +99,7 @@ func init() { rootCmd.Flags().StringVarP(&output, "output", "o", "", "Optional output file path") rootCmd.Flags().StringArrayP("prompt", "p", nil, "Prompt suffix to append to the generated content") rootCmd.Flags().StringVarP(&templatePath, "template", "t", "", "Optional Path to a custom Handlebars template") + rootCmd.Flags().Bool("save", false, "Automatically save the generated markdown to ~/ingest/.md") // VRAM estimation flags rootCmd.Flags().BoolVar(&vramFlag, "vram", false, "Estimate vRAM usage") @@ -107,7 +108,6 @@ func init() { rootCmd.Flags().IntVar(&contextFlag, "context", 0, "vRAM Estimation - Context length for vRAM estimation") rootCmd.Flags().StringVar(&kvCacheFlag, "kvcache", "fp16", "vRAM Estimation - KV cache quantization: fp16, q8_0, or q4_0") rootCmd.Flags().Float64Var(&memoryFlag, "memory", 0, "vRAM Estimation - Available memory in GB for context calculation") - rootCmd.Flags().Float64Var(&memoryFlag, "fits", 0, "vRAM Estimation - Available memory in GB for context calculation") rootCmd.Flags().StringVar(&quantTypeFlag, "quanttype", "gguf", "vRAM Estimation - Quantization type: gguf or exl2") // Add completion command @@ -299,6 +299,19 @@ func run(cmd *cobra.Command, args []string) error { return fmt.Errorf("failed to render template: %w", err) } + // Check if save is set in config or flag + autoSave, _ := cmd.Flags().GetBool("save") + if cfg.AutoSave || autoSave { + currentDir, err := os.Getwd() + if err != nil { + return fmt.Errorf("failed to get current directory: %w", err) + } + + if err := autoSaveOutput(rendered, currentDir); err != nil { + utils.PrintColouredMessage("❌", fmt.Sprintf("Error saving file: %v", err), color.FgRed) + } + } + // VRAM estimation if vramFlag { fmt.Println() @@ -319,6 +332,9 @@ func run(cmd *cobra.Command, args []string) error { } } + // Print all collected messages at the end + utils.PrintMessages() + return nil } @@ -327,18 +343,42 @@ func reportLargestFiles(files []filesystem.FileInfo) { return len(files[i].Code) > len(files[j].Code) }) - fmt.Println("\nTop 5 largest files (by estimated token count):") - for i, file := range files[:min(5, len(files))] { + // fmt.Println("Top 10 largest files (by estimated token count):") + // print this in colour + utils.PrintColouredMessage("ℹ️", "Top 10 largest files (by estimated token count):", color.FgCyan) + colourRange := []*color.Color{ + color.New(color.FgRed), + color.New(color.FgRed), + color.New(color.FgRed), + color.New(color.FgRed), + color.New(color.FgRed), + color.New(color.FgYellow), + color.New(color.FgYellow), + color.New(color.FgYellow), + color.New(color.FgYellow), + color.New(color.FgYellow), + } + + // print the files + for i, file := range files { tokenCount := token.CountTokens(file.Code, encoding) - fmt.Printf("%d. %s (%s tokens)\n", i+1, file.Path, utils.FormatNumber(tokenCount)) + // get the colour + colour := colourRange[i] + fmt.Printf("- %d. %s (%s tokens)\n", i+1, file.Path, colour.Sprint(utils.FormatNumber(tokenCount))) + // break after 10 + if i == 9 { + break + } } + + fmt.Println() } func handleOutput(rendered string, countTokens bool, encoding string, noClipboard bool, output string, jsonOutput bool, report bool, files []filesystem.FileInfo) error { if countTokens { tokenCount := token.CountTokens(rendered, encoding) println() - utils.PrintColouredMessage("ℹ️", fmt.Sprintf("Tokens (Approximate): %v", utils.FormatNumber(tokenCount)), color.FgYellow) + utils.AddMessage("ℹ️", fmt.Sprintf("Tokens (Approximate): %v", utils.FormatNumber(tokenCount)), color.FgYellow, 1) } if report { @@ -360,7 +400,7 @@ func handleOutput(rendered string, countTokens bool, encoding string, noClipboar if !noClipboard { err := utils.CopyToClipboard(rendered) if err == nil { - utils.PrintColouredMessage("✅", "Copied to clipboard successfully.", color.FgGreen) + utils.AddMessage("✅", "Copied to clipboard successfully.", color.FgGreen, 5) return nil } // If clipboard copy failed, fall back to console output @@ -372,7 +412,7 @@ func handleOutput(rendered string, countTokens bool, encoding string, noClipboar if err != nil { return fmt.Errorf("failed to write to file: %w", err) } - utils.PrintColouredMessage("✅", fmt.Sprintf("Written to file: %s", output), color.FgGreen) + utils.AddMessage("✅", fmt.Sprintf("Written to file: %s", output), color.FgGreen, 20) } else { // If no output file is specified, print to console fmt.Print(rendered) @@ -435,7 +475,7 @@ func printExcludePatterns(patterns []string) { func handleLLMOutput(rendered string, llmConfig config.LLMConfig, countTokens bool, encoding string) error { if countTokens { tokenCount := token.CountTokens(rendered, encoding) - utils.PrintColouredMessage("ℹ️", fmt.Sprintf("Tokens (Approximate): %v", utils.FormatNumber(tokenCount)), color.FgYellow) + utils.AddMessage("ℹ️", fmt.Sprintf("Tokens (Approximate): %v", utils.FormatNumber(tokenCount)), color.FgYellow, 40) } if promptPrefix != "" { @@ -583,13 +623,18 @@ func performVRAMEstimation(content string) error { return fmt.Errorf("error estimating vRAM: %w", err) } - // Print the estimation results - fmt.Printf("\nVRAM Estimation Results:\n") - fmt.Printf("Model: %s\n", estimation.ModelName) - fmt.Printf("Estimated vRAM Required: %.2f GB\n", estimation.EstimatedVRAM) - fmt.Printf("Fits Available vRAM: %v\n", estimation.FitsAvailable) - fmt.Printf("Max Context Size: %d\n", estimation.MaxContextSize) - fmt.Printf("Maximum Quantisation: %s\n", estimation.MaximumQuant) + utils.AddMessage("ℹ️", fmt.Sprintf("Model: %s", estimation.ModelName), color.FgCyan, 10) + utils.AddMessage("ℹ️", fmt.Sprintf("Estimated vRAM Required: %.2f GB", estimation.EstimatedVRAM), color.FgCyan, 3) + // print the vram available + utils.AddMessage("ℹ️", fmt.Sprintf("Available vRAM: %.2f GB", memoryFlag), color.FgCyan, 10) + if estimation.FitsAvailable { + utils.AddMessage("✅", "Fits Available vRAM", color.FgGreen, 2) + } else { + utils.AddMessage("❌", "Does Not Fit Available vRAM", color.FgYellow, 2) + } + utils.AddMessage("ℹ️", fmt.Sprintf("Max Context Size: %d", estimation.MaxContextSize), color.FgCyan, 8) + // utils.AddMessage("ℹ️", fmt.Sprintf("Maximum Quantisation: %s", estimation.MaximumQuant), color.FgCyan, 10) + // TODO: - this isn't that useful, come up with something smarter // Generate and print the quant table table, err := quantest.GenerateQuantTable(estimation.ModelConfig, memoryFlag) @@ -601,15 +646,38 @@ func performVRAMEstimation(content string) error { // Check if the content fits within the specified constraints if memoryFlag > 0 { if tokenCount > estimation.MaxContextSize { - utils.PrintColouredMessage("❗️", fmt.Sprintf("Generated content (%d tokens) exceeds maximum context (%d tokens).", tokenCount, estimation.MaxContextSize), color.FgYellow) + utils.AddMessage("❗️", fmt.Sprintf("Generated content (%d tokens) exceeds maximum context (%d tokens).", tokenCount, estimation.MaxContextSize), color.FgYellow, 2) } else { - utils.PrintColouredMessage("✅", fmt.Sprintf("Generated content (%d tokens) fits within maximum context (%d tokens).", tokenCount, estimation.MaxContextSize), color.FgGreen) + utils.AddMessage("✅", fmt.Sprintf("Generated content (%d tokens) fits within maximum context (%d tokens).", tokenCount, estimation.MaxContextSize), color.FgGreen, 2) } } return nil } +func autoSaveOutput(content string, sourcePath string) error { + homeDir, err := os.UserHomeDir() + if err != nil { + return fmt.Errorf("failed to get user home directory: %w", err) + } + + ingestDir := filepath.Join(homeDir, "ingest") + if err := os.MkdirAll(ingestDir, 0700); err != nil { + return fmt.Errorf("failed to create ingest directory: %w", err) + } + + fileName := filepath.Base(sourcePath) + ".md" + filePath := filepath.Join(ingestDir, fileName) + + if err := os.WriteFile(filePath, []byte(content), 0600); err != nil { + return fmt.Errorf("failed to write file: %w", err) + } + + utils.AddMessage("✅", fmt.Sprintf("Automatically saved to %s", filePath), color.FgGreen, 10) + + return nil +} + func runCompletion(cmd *cobra.Command, args []string) { switch args[0] { case "bash": diff --git a/utils/output_manager.go b/utils/output_manager.go new file mode 100644 index 0000000..496c9d7 --- /dev/null +++ b/utils/output_manager.go @@ -0,0 +1,50 @@ +package utils + +import ( + "sort" + "sync" + + "github.com/fatih/color" +) + +type OutputMessage struct { + Symbol string + Message string + Color color.Attribute + Priority int +} + +var ( + messages []OutputMessage + mutex sync.Mutex +) + +// AddMessage adds a message to the output queue +func AddMessage(symbol string, message string, messageColor color.Attribute, priority int) { + mutex.Lock() + defer mutex.Unlock() + messages = append(messages, OutputMessage{ + Symbol: symbol, + Message: message, + Color: messageColor, + Priority: priority, + }) +} + +// PrintMessages prints all collected messages sorted by priority +func PrintMessages() { + mutex.Lock() + defer mutex.Unlock() + + // Sort messages by priority (lower priority prints later) + sort.Slice(messages, func(i, j int) bool { + return messages[i].Priority > messages[j].Priority + }) + + for _, msg := range messages { + PrintColouredMessage(msg.Symbol, msg.Message, msg.Color) + } + + // Clear messages after printing + messages = nil +}