Add incremental search index appends

search_api.go

@@ -2,6 +2,7 @@ package gobed
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"math/rand"
@@ -88,6 +89,8 @@ type IndexingStats struct {
 	IndexingTime       time.Duration
 }
 
+var ErrDocumentIDExists = errors.New("document ID already exists")
+
 // DefaultSearchConfig returns optimized default configuration
 // Automatically detects and enables GPU with CAGRA when available
 func DefaultSearchConfig() SearchConfig {
@@ -239,6 +242,35 @@ func (se *SearchEngine) IndexBatchWithIDs(ids []int, texts []string) error {
 	return se.indexBatchInternal(ids, texts)
 }
 
+func (se *SearchEngine) AppendWithID(id int, text string) error {
+	return se.AppendBatchWithIDs([]int{id}, []string{text})
+}
+
+func (se *SearchEngine) AppendBatchWithIDs(ids []int, texts []string) error {
+	if len(ids) != len(texts) {
+		return fmt.Errorf("ids and texts must have the same length")
+	}
+	if len(ids) == 0 {
+		return nil
+	}
+
+	se.mu.Lock()
+	defer se.mu.Unlock()
+
+	seen := make(map[int]struct{}, len(ids))
+	for _, id := range ids {
+		if _, exists := seen[id]; exists {
+			return fmt.Errorf("%w: %d", ErrDocumentIDExists, id)
+		}
+		seen[id] = struct{}{}
+		if _, exists := se.documents[id]; exists {
+			return fmt.Errorf("%w: %d", ErrDocumentIDExists, id)
+		}
+	}
+
+	return se.indexBatchInternal(ids, texts)
+}
+
 // IndexBatchAsync asynchronously indexes multiple texts and returns a channel for the result
 func (se *SearchEngine) IndexBatchAsync(texts []string) <-chan IndexResponse {
 	ids := make([]int, len(texts))
@@ -366,6 +398,9 @@ func (se *SearchEngine) indexBatchInternal(ids []int, texts []string) error {
 	if numWorkers > 8 {
 		numWorkers = 8
 	}
+	if numWorkers > len(texts) {
+		numWorkers = len(texts)
+	}
 
 	type embeddingJob struct {
 		index int

search_api_incremental_test.go

@@ -0,0 +1,88 @@
+package gobed
+
+import (
+	"errors"
+	"sync"
+	"testing"
+)
+
+func TestSearchEngineAppendWithID(t *testing.T) {
+	model := mustLoadModelForTest(t)
+	config := DefaultSearchConfig()
+	config.AutoMode = false
+	config.EnableGPU = false
+	config.MaxExactSearchSize = 50000
+	engine := NewSearchEngineWithConfig(model, config)
+
+	if err := engine.AppendWithID(41, "calm piano ambience"); err != nil {
+		t.Fatal(err)
+	}
+	if err := engine.AppendWithID(42, "heavy metal guitar"); err != nil {
+		t.Fatal(err)
+	}
+	if err := engine.AppendWithID(41, "replacement"); !errors.Is(err, ErrDocumentIDExists) {
+		t.Fatalf("expected ErrDocumentIDExists, got %v", err)
+	}
+	if got := engine.Stats().NumDocuments; got != 2 {
+		t.Fatalf("expected 2 documents, got %d", got)
+	}
+
+	results, err := engine.Search("piano", 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	found := false
+	for _, result := range results {
+		if result.ID == 41 {
+			found = true
+		}
+	}
+	if !found {
+		t.Fatalf("appended document missing from results: %+v", results)
+	}
+}
+
+func TestSearchEngineAppendWithIDConcurrent(t *testing.T) {
+	model := mustLoadModelForTest(t)
+	config := DefaultSearchConfig()
+	config.AutoMode = false
+	config.EnableGPU = false
+	config.MaxExactSearchSize = 50000
+	engine := NewSearchEngineWithConfig(model, config)
+
+	var wg sync.WaitGroup
+	errs := make(chan error, 2)
+	for id, text := range map[int]string{10: "ocean waves", 11: "forest birds"} {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			errs <- engine.AppendWithID(id, text)
+		}()
+	}
+	wg.Wait()
+	close(errs)
+	for err := range errs {
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	if got := engine.Stats().NumDocuments; got != 2 {
+		t.Fatalf("expected 2 documents, got %d", got)
+	}
+}
+
+func TestSearchEngineAppendBatchRejectsDuplicateIDs(t *testing.T) {
+	engine := &SearchEngine{documents: make(map[int]string)}
+	if err := engine.AppendBatchWithIDs([]int{7, 7}, []string{"one", "two"}); !errors.Is(err, ErrDocumentIDExists) {
+		t.Fatalf("expected ErrDocumentIDExists, got %v", err)
+	}
+	if engine.initialized || len(engine.documents) != 0 {
+		t.Fatal("duplicate batch changed engine state")
+	}
+	if err := engine.AppendBatchWithIDs(nil, nil); err != nil {
+		t.Fatal(err)
+	}
+	if engine.initialized {
+		t.Fatal("empty batch initialized engine")
+	}
+}