kb/client/cmd/add.go

package cmd

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/kb-search/kb/internal/api"
	"github.com/kb-search/kb/internal/output"
	"github.com/spf13/cobra"
)

var supportedExts = map[string]bool{
	".pdf":  true,
	".docx": true,
	".html": true,
	".md":   true,
	".txt":  true,
	".py":   true,
	".sh":   true,
	".go":   true,
}

var addCmd = &cobra.Command{
	Use:   "add <path>",
	Short: "Add a document or directory to the knowledge base",
	Args:  cobra.MaximumNArgs(1),
	RunE:  runAdd,
}

func init() {
	addCmd.Flags().String("tags", "", "tags (comma-separated)")
	addCmd.Flags().String("type", "", "document type")
	addCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents")
	addCmd.Flags().String("note", "", "add a text note instead of a file")
	addCmd.Flags().String("title", "", "title for the note")
	rootCmd.AddCommand(addCmd)
}

func runAdd(cmd *cobra.Command, args []string) error {
	tags, _ := cmd.Flags().GetString("tags")
	docType, _ := cmd.Flags().GetString("type")
	recursive, _ := cmd.Flags().GetBool("recursive")
	note, _ := cmd.Flags().GetString("note")
	title, _ := cmd.Flags().GetString("title")

	client := api.NewClient()

	// Note mode
	if note != "" {
		fields := map[string]string{
			"note": note,
		}
		if title != "" {
			fields["title"] = title
		}
		if tags != "" {
			fields["tags"] = tags
		}
		if docType != "" {
			fields["type"] = docType
		}

		resp, err := client.PostMultipart("/api/v1/jobs", fields, nil)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		if err := api.CheckError(resp); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var result interface{}
		if err := api.DecodeJSON(resp, &result); err != nil {
			return fmt.Errorf("failed to decode response: %w", err)
		}

		if output.IsJSON() {
			output.PrintJSON(result)
		} else {
			fmt.Println("Queued: note")
		}
		return nil
	}

	if len(args) == 0 {
		return fmt.Errorf("path argument is required (or use --note)")
	}

	path := args[0]
	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("cannot access %s: %w", path, err)
	}

	if !info.IsDir() {
		// Single file upload
		result, err := uploadFile(client, path, tags, docType)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		if output.IsJSON() {
			output.PrintJSON([]interface{}{result})
		} else {
			fmt.Printf("Queued: %s\n", filepath.Base(path))
		}
		return nil
	}

	// Directory upload
	var files []string
	walkFn := func(p string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() {
			if !recursive && p != path {
				return filepath.SkipDir
			}
			return nil
		}
		ext := strings.ToLower(filepath.Ext(p))
		if supportedExts[ext] {
			files = append(files, p)
		}
		return nil
	}

	if err := filepath.WalkDir(path, walkFn); err != nil {
		return fmt.Errorf("failed to walk directory: %w", err)
	}

	var results []interface{}
	for _, f := range files {
		result, err := uploadFile(client, f, tags, docType)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err)
			continue
		}
		results = append(results, result)
		if !output.IsJSON() {
			fmt.Printf("Queued: %s\n", filepath.Base(f))
		}
	}

	if output.IsJSON() {
		output.PrintJSON(results)
	} else {
		fmt.Printf("Queued: %d files\n", len(results))
	}
	return nil
}

func uploadFile(client *api.Client, path, tags, docType string) (interface{}, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("cannot open %s: %w", path, err)
	}
	defer f.Close()

	fields := make(map[string]string)
	if tags != "" {
		fields["tags"] = tags
	}
	if docType != "" {
		fields["type"] = docType
	}

	upload := &api.FileUpload{
		FieldName: "file",
		FileName:  filepath.Base(path),
		Reader:    f,
	}

	resp, err := client.PostMultipart("/api/v1/jobs", fields, upload)
	if err != nil {
		return nil, err
	}
	if err := api.CheckError(resp); err != nil {
		return nil, err
	}

	var result interface{}
	if err := api.DecodeJSON(resp, &result); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}
	return result, nil
}