package cmd import ( "encoding/json" "fmt" "net/http" "os" "path/filepath" "sort" "strings" "github.com/kb-search/kb/internal/api" "github.com/kb-search/kb/internal/output" "github.com/spf13/cobra" ) type uploadResult struct { Raw interface{} Duplicate bool DocID float64 JobID float64 Title string } func (r *uploadResult) duplicateMsg() string { if r.DocID > 0 { return fmt.Sprintf("Already imported: %s (doc ID: %.0f)", r.Title, r.DocID) } return fmt.Sprintf("Already queued: %s (job ID: %.0f)", r.Title, r.JobID) } var supportedExts = map[string]bool{ ".pdf": true, ".docx": true, ".html": true, ".md": true, ".txt": true, ".py": true, ".sh": true, ".go": true, } var addfileCmd = &cobra.Command{ Use: "addfile ", Short: "Upload a file or directory to the knowledge base", Args: cobra.ExactArgs(1), RunE: runAddfile, } func init() { addfileCmd.Flags().String("tags", "", "tags (comma-separated)") addfileCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents") rootCmd.AddCommand(addfileCmd) } func runAddfile(cmd *cobra.Command, args []string) error { tags, _ := cmd.Flags().GetString("tags") recursive, _ := cmd.Flags().GetBool("recursive") client := api.NewClient() path := args[0] info, err := os.Stat(path) if err != nil { return fmt.Errorf("cannot access %s: %w", path, err) } if !info.IsDir() { // Validate extension ext := strings.ToLower(filepath.Ext(path)) if !supportedExts[ext] { supported := make([]string, 0, len(supportedExts)) for e := range supportedExts { supported = append(supported, e) } sort.Strings(supported) return fmt.Errorf("unsupported file type %q — supported: %s", ext, strings.Join(supported, ", ")) } // Single file upload result, err := uploadFile(client, path, tags) if err != nil { fmt.Fprintln(os.Stderr, err) os.Exit(1) } if output.IsJSON() { output.PrintJSON([]interface{}{result.Raw}) } else if result.Duplicate { fmt.Println(result.duplicateMsg()) } 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{} queued := 0 duplicates := 0 for _, f := range files { result, err := uploadFile(client, f, tags) if err != nil { fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err) continue } results = append(results, result.Raw) if result.Duplicate { duplicates++ if !output.IsJSON() { fmt.Println(result.duplicateMsg()) } } else { queued++ if !output.IsJSON() { fmt.Printf("Queued: %s\n", filepath.Base(f)) } } } if output.IsJSON() { output.PrintJSON(results) } else if duplicates > 0 { fmt.Printf("Queued: %d files, %d duplicates skipped\n", queued, duplicates) } else { fmt.Printf("Queued: %d files\n", queued) } return nil } func uploadFile(client *api.Client, path, tags string) (*uploadResult, 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 } 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 resp.StatusCode == http.StatusConflict { var raw json.RawMessage if err := api.DecodeJSON(resp, &raw); err != nil { return nil, fmt.Errorf("failed to decode response: %w", err) } var dupResp struct { DocumentID float64 `json:"document_id"` JobID float64 `json:"job_id"` Title string `json:"title"` } json.Unmarshal(raw, &dupResp) var rawIface interface{} json.Unmarshal(raw, &rawIface) return &uploadResult{ Raw: rawIface, Duplicate: true, DocID: dupResp.DocumentID, JobID: dupResp.JobID, Title: dupResp.Title, }, nil } 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 &uploadResult{Raw: result}, nil }