diff --git a/README.md b/README.md index 6e67a3a..3f361f3 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,10 @@ - [Watch Mode](#watch-mode) - [Activating Watch Mode](#activating-watch-mode) - [Example Use Cases](#example-use-cases) +- [Knowledge Base](#knowledge-base) + - [Creating Knowledge Bases](#creating-knowledge-bases) + - [Using Knowledge Bases](#using-knowledge-bases) + - [Auto-Loading Knowledge Bases](#auto-loading-knowledge-bases) - [Squashing](#squashing) - [What is Squashing?](#what-is-squashing) - [Manual Squashing](#manual-squashing) @@ -314,6 +318,97 @@ If you'd like to manage your context before reaching the automatic threshold, yo TmuxAI » /squash ``` +## Knowledge Base + +The Knowledge Base feature allows you to create pre-defined context files in markdown format that can be loaded into TmuxAI's conversation context. This is useful for sharing common patterns, workflows, or project-specific information with the AI across sessions. + +### Creating Knowledge Bases + +Knowledge bases are markdown files stored in `~/.config/tmuxai/kb/`. To create one: + +1. Create the knowledge base directory if it doesn't exist: + ```bash + mkdir -p ~/.config/tmuxai/kb + ``` + +2. Create a markdown file with your knowledge base content: + ```bash + cat > ~/.config/tmuxai/kb/docker-workflows.md << 'EOF' + # Docker Workflows + + ## Common Commands + - Always use `docker compose` (not `docker-compose`) + - Prefer named volumes over bind mounts for databases + - Use `.env` files for environment-specific configuration + + ## Project Structure + - Development: `docker compose -f docker-compose.dev.yml up` + - Production: `docker compose -f docker-compose.prod.yml up -d` + EOF + ``` + +### Using Knowledge Bases + +Once created, you can load knowledge bases into your TmuxAI session: + +```bash +# List available knowledge bases +TmuxAI » /kb +Available knowledge bases: + [ ] docker-workflows + [ ] git-conventions + [ ] testing-procedures + +# Load a knowledge base +TmuxAI » /kb load docker-workflows +✓ Loaded knowledge base: docker-workflows (850 tokens) + +# List again to see loaded status +TmuxAI » /kb +Available knowledge bases: + [✓] docker-workflows (850 tokens) + [ ] git-conventions + [ ] testing-procedures + +Loaded: 1 KB(s), 850 tokens + +# Unload a knowledge base +TmuxAI » /kb unload docker-workflows +✓ Unloaded knowledge base: docker-workflows + +# Unload all knowledge bases +TmuxAI » /kb unload --all +✓ Unloaded all knowledge bases (2 KB(s)) +``` + +You can also load knowledge bases directly from the command line when starting TmuxAI: + +```bash +# Load single knowledge base +tmuxai --kb docker-workflows + +# Load multiple knowledge bases (comma-separated) +tmuxai --kb docker-workflows,git-conventions +``` + +### Auto-Loading Knowledge Bases + +You can configure knowledge bases to load automatically on startup by adding them to your `~/.config/tmuxai/config.yaml`: + +```yaml +knowledge_base: + auto_load: + - docker-workflows + - git-conventions + # path: /custom/path # Optional: use custom KB directory +``` + +**Important Notes:** +- Loaded knowledge bases consume tokens from your context budget +- Use `/info` to see how many tokens your loaded KBs are using +- Knowledge bases are injected after the system prompt but before conversation history +- Unloading a KB removes it from future messages immediately + ## Core Commands | Command | Description | @@ -324,8 +419,12 @@ TmuxAI » /squash | `/config` | View current configuration settings | | `/config set ` | Override configuration for current session | | `/squash` | Manually trigger context summarization | -| `/prepare [shell]` | Initialize Prepared Mode for the Exec Pane (e.g., bash, zsh) | +| `/prepare [shell]` | Initialize Prepared Mode for the Exec Pane (e.g., bash, zsh) | | `/watch ` | Enable Watch Mode with specified goal | +| `/kb` | List available knowledge bases with loaded status | +| `/kb load ` | Load a knowledge base into conversation context | +| `/kb unload ` | Unload a specific knowledge base | +| `/kb unload --all` | Unload all knowledge bases | | `/exit` | Exit TmuxAI | ## Command-Line Usage @@ -343,6 +442,15 @@ You can start `tmuxai` with an initial message or task file from the command lin tmuxai -f path/to/your_task.txt ``` +- **Load Knowledge Bases:** + ```sh + # Single knowledge base + tmuxai --kb docker-workflows + + # Multiple knowledge bases + tmuxai --kb docker-workflows,git-conventions + ``` + ## Configuration The configuration can be managed through a YAML file, environment variables, or via runtime commands. diff --git a/cli/cli.go b/cli/cli.go index 6caba20..aa049b4 100644 --- a/cli/cli.go +++ b/cli/cli.go @@ -16,6 +16,7 @@ import ( var ( initMessage string taskFileFlag string + kbFlag string ) var rootCmd = &cobra.Command{ @@ -56,6 +57,13 @@ var rootCmd = &cobra.Command{ logger.Error("manager.NewManager failed: %v", err) os.Exit(1) } + + // Load knowledge bases from CLI flag + if kbFlag != "" { + kbNames := strings.Split(kbFlag, ",") + mgr.LoadKBsFromCLI(kbNames) + } + if initMessage != "" { logger.Info("Starting with initial subcommand: %s", initMessage) } @@ -69,6 +77,7 @@ var rootCmd = &cobra.Command{ func init() { rootCmd.Flags().StringVarP(&taskFileFlag, "file", "f", "", "Read request from specified file") + rootCmd.Flags().StringVar(&kbFlag, "kb", "", "Comma-separated list of knowledge bases to load (e.g., --kb docker,git)") rootCmd.Flags().BoolP("version", "v", false, "Print version information") } diff --git a/config/config.go b/config/config.go index 1076887..48ee54c 100644 --- a/config/config.go +++ b/config/config.go @@ -12,19 +12,20 @@ import ( // Config holds the application configuration type Config struct { - Debug bool `mapstructure:"debug"` - MaxCaptureLines int `mapstructure:"max_capture_lines"` - MaxContextSize int `mapstructure:"max_context_size"` - WaitInterval int `mapstructure:"wait_interval"` - SendKeysConfirm bool `mapstructure:"send_keys_confirm"` - PasteMultilineConfirm bool `mapstructure:"paste_multiline_confirm"` - ExecConfirm bool `mapstructure:"exec_confirm"` - WhitelistPatterns []string `mapstructure:"whitelist_patterns"` - BlacklistPatterns []string `mapstructure:"blacklist_patterns"` - OpenRouter OpenRouterConfig `mapstructure:"openrouter"` - OpenAI OpenAIConfig `mapstructure:"openai"` - AzureOpenAI AzureOpenAIConfig `mapstructure:"azure_openai"` - Prompts PromptsConfig `mapstructure:"prompts"` + Debug bool `mapstructure:"debug"` + MaxCaptureLines int `mapstructure:"max_capture_lines"` + MaxContextSize int `mapstructure:"max_context_size"` + WaitInterval int `mapstructure:"wait_interval"` + SendKeysConfirm bool `mapstructure:"send_keys_confirm"` + PasteMultilineConfirm bool `mapstructure:"paste_multiline_confirm"` + ExecConfirm bool `mapstructure:"exec_confirm"` + WhitelistPatterns []string `mapstructure:"whitelist_patterns"` + BlacklistPatterns []string `mapstructure:"blacklist_patterns"` + OpenRouter OpenRouterConfig `mapstructure:"openrouter"` + OpenAI OpenAIConfig `mapstructure:"openai"` + AzureOpenAI AzureOpenAIConfig `mapstructure:"azure_openai"` + Prompts PromptsConfig `mapstructure:"prompts"` + KnowledgeBase KnowledgeBaseConfig `mapstructure:"knowledge_base"` } // OpenRouterConfig holds OpenRouter API configuration @@ -57,6 +58,12 @@ type PromptsConfig struct { Watch string `mapstructure:"watch"` } +// KnowledgeBaseConfig holds knowledge base configuration +type KnowledgeBaseConfig struct { + AutoLoad []string `mapstructure:"auto_load"` + Path string `mapstructure:"path"` +} + // DefaultConfig returns a configuration with default values func DefaultConfig() *Config { return &Config{ @@ -81,6 +88,10 @@ func DefaultConfig() *Config { BaseSystem: ``, ChatAssistant: ``, }, + KnowledgeBase: KnowledgeBaseConfig{ + AutoLoad: []string{}, + Path: "", + }, } } @@ -174,6 +185,25 @@ func GetConfigFilePath(filename string) string { return filepath.Join(configDir, filename) } +// GetKBDir returns the path to the knowledge base directory +func GetKBDir() string { + // Try to load config to check for custom path + cfg, err := Load() + if err == nil && cfg.KnowledgeBase.Path != "" { + // Use custom path if specified + return cfg.KnowledgeBase.Path + } + + // Default to ~/.config/tmuxai/kb/ + configDir, _ := GetConfigDir() + kbDir := filepath.Join(configDir, "kb") + + // Create KB directory if it doesn't exist + _ = os.MkdirAll(kbDir, 0o755) + + return kbDir +} + func TryInferType(key, value string) any { var typedValue any = value // Only basic type inference for bool/int/string diff --git a/internal/chat.go b/internal/chat.go index 2bbaf38..3d66b81 100644 --- a/internal/chat.go +++ b/internal/chat.go @@ -184,6 +184,32 @@ func (c *CLIInterface) newCompleter() *completion.CmdCompletionOrList2 { return []string{"bash", "zsh", "fish"}, []string{"bash", "zsh", "fish"} } } + + // Handle /kb subcommands + if len(field) > 0 && field[0] == "/kb" { + if len(field) == 1 || (len(field) == 2 && !strings.HasSuffix(field[1], " ")) { + return []string{"list", "load", "unload"}, []string{"list", "load", "unload"} + } else if (len(field) == 2 && field[1] == "load") || (len(field) >= 3 && field[1] == "load") { + // Get available knowledge bases for completion + kbs, err := c.manager.listKBs() + if err != nil { + return nil, nil + } + // Disable autocompletion when there's only one KB, bug with readline + if len(kbs) == 1 { + return nil, nil + } + return kbs, kbs + } else if (len(field) == 2 && field[1] == "unload") || (len(field) >= 3 && field[1] == "unload") { + // For unload, show loaded knowledge bases and --all option + var kbNames []string + for name := range c.manager.LoadedKBs { + kbNames = append(kbNames, name) + } + kbNames = append(kbNames, "--all") + return kbNames, kbNames + } + } return nil, nil }, } diff --git a/internal/chat_command.go b/internal/chat_command.go index 28d6df2..8f45214 100644 --- a/internal/chat_command.go +++ b/internal/chat_command.go @@ -18,6 +18,10 @@ const helpMessage = `Available commands: - /prepare: Prepare the pane for TmuxAI automation - /watch : Start watch mode - /squash: Summarize the chat history +- /kb: List available knowledge bases +- /kb load : Load a knowledge base +- /kb unload : Unload a knowledge base +- /kb unload --all: Unload all knowledge bases - /exit: Exit the application` var commands = []string{ @@ -30,6 +34,7 @@ var commands = []string{ "/prepare", "/config", "/squash", + "/kb", } // checks if the given content is a command @@ -191,6 +196,99 @@ Watch for: ` + watchDesc return } + case prefixMatch(commandPrefix, "/kb"): + // Handle KB commands: /kb, /kb list, /kb load , /kb unload + if len(parts) == 1 || (len(parts) == 2 && parts[1] == "list") { + // List all available knowledge bases + kbs, err := m.listKBs() + if err != nil { + m.Println(fmt.Sprintf("Error listing knowledge bases: %v", err)) + return + } + + if len(kbs) == 0 { + m.Println("No knowledge bases found in " + config.GetKBDir()) + return + } + + m.Println("Available knowledge bases:") + totalTokens := 0 + loadedCount := 0 + + for _, name := range kbs { + _, loaded := m.LoadedKBs[name] + status := "[ ]" + tokens := "" + if loaded { + status = "[✓]" + tokenCount := system.EstimateTokenCount(m.LoadedKBs[name]) + tokens = fmt.Sprintf(" (%d tokens)", tokenCount) + totalTokens += tokenCount + loadedCount++ + } + m.Println(fmt.Sprintf(" %s %s%s", status, name, tokens)) + } + + if loadedCount > 0 { + m.Println("") + m.Println(fmt.Sprintf("Loaded: %d KB(s), %d tokens", loadedCount, totalTokens)) + } + return + + } else if len(parts) >= 2 && parts[1] == "load" { + if len(parts) < 3 { + m.Println("Usage: /kb load ") + return + } + + name := parts[2] + if _, loaded := m.LoadedKBs[name]; loaded { + m.Println(fmt.Sprintf("Knowledge base '%s' is already loaded", name)) + return + } + + if err := m.loadKB(name); err != nil { + m.Println(fmt.Sprintf("Error loading KB '%s': %v", name, err)) + return + } + + tokenCount := system.EstimateTokenCount(m.LoadedKBs[name]) + m.Println(fmt.Sprintf("✓ Loaded knowledge base: %s (%d tokens)", name, tokenCount)) + return + + } else if len(parts) >= 2 && parts[1] == "unload" { + if len(parts) >= 3 && parts[2] == "--all" { + // Unload all KBs + if len(m.LoadedKBs) == 0 { + m.Println("No knowledge bases are currently loaded") + return + } + + count := len(m.LoadedKBs) + m.LoadedKBs = make(map[string]string) + m.Println(fmt.Sprintf("✓ Unloaded all knowledge bases (%d KB(s))", count)) + return + } + + if len(parts) < 3 { + m.Println("Usage: /kb unload or /kb unload --all") + return + } + + name := parts[2] + if err := m.unloadKB(name); err != nil { + m.Println(fmt.Sprintf("Error: %v", err)) + return + } + + m.Println(fmt.Sprintf("✓ Unloaded knowledge base: %s", name)) + return + + } else { + m.Println("Usage: /kb [list|load |unload |unload --all]") + return + } + default: m.Println(fmt.Sprintf("Unknown command: %s. Type '/help' to see available commands.", command)) return @@ -235,6 +333,12 @@ func (m *Manager) formatInfo() { fmt.Printf("%-*s %s\n", labelWidth, "", formatter.FormatProgressBar(usagePercent, 10)) formatLine("Max Size", fmt.Sprintf("%d tokens", m.GetMaxContextSize())) + // Display knowledge base information + if len(m.LoadedKBs) > 0 { + kbTokens := m.getTotalLoadedKBTokens() + formatLine("Loaded KBs", fmt.Sprintf("%d (%d tokens)", len(m.LoadedKBs), kbTokens)) + } + // Display tmux panes section fmt.Println() fmt.Println(formatter.FormatSection("Tmux Window Panes")) diff --git a/internal/knowledge_base.go b/internal/knowledge_base.go new file mode 100644 index 0000000..cd14307 --- /dev/null +++ b/internal/knowledge_base.go @@ -0,0 +1,108 @@ +package internal + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/alvinunreal/tmuxai/config" + "github.com/alvinunreal/tmuxai/logger" + "github.com/alvinunreal/tmuxai/system" +) + +// loadKB loads a knowledge base file by name +func (m *Manager) loadKB(name string) error { + kbDir := config.GetKBDir() + kbPath := filepath.Join(kbDir, name+".md") + + content, err := os.ReadFile(kbPath) + if err != nil { + return fmt.Errorf("failed to read KB file '%s': %w", name, err) + } + + m.LoadedKBs[name] = string(content) + logger.Info("Loaded knowledge base: %s", name) + return nil +} + +// unloadKB removes a knowledge base from memory +func (m *Manager) unloadKB(name string) error { + if _, exists := m.LoadedKBs[name]; !exists { + return fmt.Errorf("knowledge base '%s' is not loaded", name) + } + + delete(m.LoadedKBs, name) + logger.Info("Unloaded knowledge base: %s", name) + return nil +} + +// listKBs returns a list of all available knowledge bases with their loaded status +func (m *Manager) listKBs() ([]string, error) { + kbDir := config.GetKBDir() + + entries, err := os.ReadDir(kbDir) + if err != nil { + if os.IsNotExist(err) { + // KB directory doesn't exist yet + return []string{}, nil + } + return nil, fmt.Errorf("failed to read KB directory: %w", err) + } + + var kbs []string + for _, entry := range entries { + if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".md") { + name := strings.TrimSuffix(entry.Name(), ".md") + kbs = append(kbs, name) + } + } + + return kbs, nil +} + +// autoLoadKBs loads knowledge bases specified in the config +func (m *Manager) autoLoadKBs() { + if len(m.Config.KnowledgeBase.AutoLoad) == 0 { + return + } + + logger.Info("Auto-loading knowledge bases: %v", m.Config.KnowledgeBase.AutoLoad) + + for _, name := range m.Config.KnowledgeBase.AutoLoad { + if err := m.loadKB(name); err != nil { + logger.Error("Failed to auto-load KB '%s': %v", name, err) + m.Println(fmt.Sprintf("Warning: Failed to auto-load KB '%s': %v", name, err)) + } + } +} + +// LoadKBsFromCLI loads knowledge bases specified via CLI flag +func (m *Manager) LoadKBsFromCLI(kbNames []string) { + if len(kbNames) == 0 { + return + } + + logger.Info("Loading knowledge bases from CLI: %v", kbNames) + + for _, name := range kbNames { + name = strings.TrimSpace(name) + if name == "" { + continue + } + + if err := m.loadKB(name); err != nil { + logger.Error("Failed to load KB '%s' from CLI: %v", name, err) + m.Println(fmt.Sprintf("Warning: Failed to load KB '%s': %v", name, err)) + } + } +} + +// getTotalLoadedKBTokens calculates the total token count of all loaded KBs +func (m *Manager) getTotalLoadedKBTokens() int { + total := 0 + for _, content := range m.LoadedKBs { + total += system.EstimateTokenCount(content) + } + return total +} diff --git a/internal/knowledge_base_test.go b/internal/knowledge_base_test.go new file mode 100644 index 0000000..bec7753 --- /dev/null +++ b/internal/knowledge_base_test.go @@ -0,0 +1,97 @@ +package internal + +import ( + "os" + "path/filepath" + "testing" + + "github.com/alvinunreal/tmuxai/config" +) + + +// TestLoadKBNonExistent tests loading a non-existent KB +func TestLoadKBNonExistent(t *testing.T) { + tmpDir := t.TempDir() + kbDir := filepath.Join(tmpDir, "kb") + if err := os.MkdirAll(kbDir, 0755); err != nil { + t.Fatalf("Failed to create KB directory: %v", err) + } + + cfg := config.DefaultConfig() + cfg.KnowledgeBase.Path = kbDir + cfg.OpenRouter.APIKey = "test-key" + + mgr := &Manager{ + Config: cfg, + LoadedKBs: make(map[string]string), + } + + // Try to load non-existent KB + err := mgr.loadKB("nonexistent") + if err == nil { + t.Fatal("Expected error when loading non-existent KB, got nil") + } +} + +// TestUnloadKB tests unloading a knowledge base +func TestUnloadKB(t *testing.T) { + cfg := config.DefaultConfig() + cfg.OpenRouter.APIKey = "test-key" + + mgr := &Manager{ + Config: cfg, + LoadedKBs: map[string]string{ + "test": "test content", + }, + } + + // Test unloading KB + err := mgr.unloadKB("test") + if err != nil { + t.Fatalf("unloadKB() failed: %v", err) + } + + // Verify KB was unloaded + if _, exists := mgr.LoadedKBs["test"]; exists { + t.Fatal("KB still exists in LoadedKBs after unloading") + } +} + +// TestUnloadKBNonLoaded tests unloading a KB that isn't loaded +func TestUnloadKBNonLoaded(t *testing.T) { + cfg := config.DefaultConfig() + cfg.OpenRouter.APIKey = "test-key" + + mgr := &Manager{ + Config: cfg, + LoadedKBs: make(map[string]string), + } + + // Try to unload non-loaded KB + err := mgr.unloadKB("test") + if err == nil { + t.Fatal("Expected error when unloading non-loaded KB, got nil") + } +} + + +// TestGetTotalLoadedKBTokens tests token counting for loaded KBs +func TestGetTotalLoadedKBTokens(t *testing.T) { + cfg := config.DefaultConfig() + cfg.OpenRouter.APIKey = "test-key" + + mgr := &Manager{ + Config: cfg, + LoadedKBs: map[string]string{ + "kb1": "Short content", + "kb2": "Another piece of content with more words", + }, + } + + tokens := mgr.getTotalLoadedKBTokens() + + // We can't test exact token count, but it should be > 0 + if tokens <= 0 { + t.Errorf("Expected positive token count, got %d", tokens) + } +} diff --git a/internal/manager.go b/internal/manager.go index 9041d36..4bb1ec0 100644 --- a/internal/manager.go +++ b/internal/manager.go @@ -42,6 +42,7 @@ type Manager struct { WatchMode bool OS string SessionOverrides map[string]interface{} // session-only config overrides + LoadedKBs map[string]string // Loaded knowledge bases (name -> content) // Functions for mocking confirmedToExec func(command string, prompt string, edit bool) (bool, string) @@ -86,12 +87,17 @@ func NewManager(cfg *config.Config) (*Manager, error) { ExecPane: &system.TmuxPaneDetails{}, OS: os, SessionOverrides: make(map[string]interface{}), + LoadedKBs: make(map[string]string), } manager.confirmedToExec = manager.confirmedToExecFn manager.getTmuxPanesInXml = manager.getTmuxPanesInXmlFn manager.InitExecPane() + + // Auto-load knowledge bases from config + manager.autoLoadKBs() + return manager, nil } diff --git a/internal/process_message.go b/internal/process_message.go index cbc1cb2..7587cb0 100644 --- a/internal/process_message.go +++ b/internal/process_message.go @@ -50,6 +50,15 @@ func (m *Manager) ProcessUserMessage(ctx context.Context, message string) bool { history = []ChatMessage{m.chatAssistantPrompt(false)} } + // Inject loaded knowledge bases after system prompt + for kbName, kbContent := range m.LoadedKBs { + history = append(history, ChatMessage{ + Content: fmt.Sprintf("=== Knowledge Base: %s ===\n%s", kbName, kbContent), + FromUser: false, + Timestamp: time.Now(), + }) + } + history = append(history, m.Messages...) sending := append(history, currentMessage)