Complete outstanding ingestion and document UX work

client/cmd/add.go

@@ -8,6 +8,7 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
+	"time"
 
 	"github.com/kb-search/kb/internal/api"
 	"github.com/kb-search/kb/internal/output"
@@ -54,12 +55,16 @@ var addfileCmd = &cobra.Command{
 func init() {
 	addfileCmd.Flags().String("tags", "", "tags (comma-separated)")
 	addfileCmd.Flags().BoolP("recursive", "r", false, "recursively add directory contents")
+	addfileCmd.Flags().Bool("wait", false, "wait for ingestion to finish")
+	addfileCmd.Flags().Duration("wait-timeout", 10*time.Minute, "maximum time to wait per ingestion job")
 	rootCmd.AddCommand(addfileCmd)
 }
 
 func runAddfile(cmd *cobra.Command, args []string) error {
 	tags, _ := cmd.Flags().GetString("tags")
 	recursive, _ := cmd.Flags().GetBool("recursive")
+	wait, _ := cmd.Flags().GetBool("wait")
+	timeout, _ := cmd.Flags().GetDuration("wait-timeout")
 
 	client := api.NewClient()
 
@@ -89,12 +94,25 @@ func runAddfile(cmd *cobra.Command, args []string) error {
 		}
 
 		if output.IsJSON() {
-			output.PrintJSON([]interface{}{result.Raw})
+			if !wait || result.Duplicate {
+				output.PrintJSON([]interface{}{result.Raw})
+			}
 		} else if result.Duplicate {
 			fmt.Println(result.duplicateMsg())
 		} else {
 			fmt.Printf("Queued: %s\n", filepath.Base(path))
 		}
+		if wait && !result.Duplicate {
+			job, err := waitForJob(client, int(result.JobID), timeout)
+			if err != nil {
+				return err
+			}
+			if output.IsJSON() {
+				output.PrintJSON([]interface{}{job})
+			} else {
+				fmt.Printf("Ingested: %s (doc ID: %d, chunks: %d)\n", filepath.Base(path), job.DocumentID, job.ChunkCount)
+			}
+		}
 		return nil
 	}
 
@@ -122,6 +140,7 @@ func runAddfile(cmd *cobra.Command, args []string) error {
 	}
 
 	var results []interface{}
+	var pending []*uploadResult
 	queued := 0
 	duplicates := 0
 	for _, f := range files {
@@ -130,7 +149,9 @@ func runAddfile(cmd *cobra.Command, args []string) error {
 			fmt.Fprintf(os.Stderr, "Error uploading %s: %v\n", f, err)
 			continue
 		}
-		results = append(results, result.Raw)
+		if !wait || result.Duplicate {
+			results = append(results, result.Raw)
+		}
 		if result.Duplicate {
 			duplicates++
 			if !output.IsJSON() {
@@ -138,11 +159,25 @@ func runAddfile(cmd *cobra.Command, args []string) error {
 			}
 		} else {
 			queued++
+			pending = append(pending, result)
 			if !output.IsJSON() {
 				fmt.Printf("Queued: %s\n", filepath.Base(f))
 			}
 		}
 	}
+	if wait {
+		for _, result := range pending {
+			job, err := waitForJob(client, int(result.JobID), timeout)
+			if err != nil {
+				return err
+			}
+			if output.IsJSON() {
+				results = append(results, job)
+			} else {
+				fmt.Printf("Ingested job %d (doc ID: %d, chunks: %d)\n", int(result.JobID), job.DocumentID, job.ChunkCount)
+			}
+		}
+	}
 
 	if output.IsJSON() {
 		output.PrintJSON(results)
@@ -203,10 +238,19 @@ func uploadFile(client *api.Client, path, tags string) (*uploadResult, error) {
 		return nil, err
 	}
 
-	var result interface{}
-	if err := api.DecodeJSON(resp, &result); err != nil {
+	var raw json.RawMessage
+	if err := api.DecodeJSON(resp, &raw); err != nil {
 		return nil, fmt.Errorf("failed to decode response: %w", err)
 	}
-	return &uploadResult{Raw: result}, nil
+	var queued struct {
+		JobID float64 `json:"job_id"`
+	}
+	if err := json.Unmarshal(raw, &queued); err != nil {
+		return nil, fmt.Errorf("failed to decode queued job: %w", err)
+	}
+	var result interface{}
+	if err := json.Unmarshal(raw, &result); err != nil {
+		return nil, fmt.Errorf("failed to decode queued response: %w", err)
+	}
+	return &uploadResult{Raw: result, JobID: queued.JobID}, nil
 }
-