From 95f59b72f9afd2479a294b1c54e1a182cb3deb7f Mon Sep 17 00:00:00 2001 From: Lee Penkman Date: Tue, 11 Aug 2026 14:52:21 +1200 Subject: [PATCH] Add incremental search index appends --- search_api.go | 35 ++++++++++++++ search_api_incremental_test.go | 88 ++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) create mode 100644 search_api_incremental_test.go diff --git a/search_api.go b/search_api.go index 1353cf3..10bcd04 100644 --- a/search_api.go +++ b/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 diff --git a/search_api_incremental_test.go b/search_api_incremental_test.go new file mode 100644 index 0000000..4b22c0e --- /dev/null +++ b/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") + } +}