Long markdown files (>~8KB) silently failed to embed because nomic-embed-text on iguana has a 2048-token context. embed sync logged errors=1 every cycle with no useful body until #37 added per-item logging — three files exceed the ceiling: finbert source (8 KB), koala-machine-state (7.1 KB), litellm-absorption (8.8 KB). Curated knowledge entries should never be vector-blind. Approach: chunk-before-embed, no schema change. vectorstore/chunk.go (new) - ChunkMarkdown splits at H1/H2 boundaries; sections over maxBytes are further split at paragraph boundaries, packing greedily under budget. - NumberChunks assigns "<parent>#NNNN" storage paths (1-based, zero-padded to 4 digits — handles files with up to ~10k sections in stable sort order). - ParentPath strips the chunk suffix for retrieval-side dedup. vectorstore/sync.go - After ChunkMarkdown produces N pieces, each is embedded + upserted as a separate brain_embeddings row at "<parent>#NNNN". maxChunkBytes = 4000 (≈1000 nomic tokens, well under the 2048 ceiling with headroom for unicode/code blocks). - "Already embedded?" check now reduces known paths to parent set via ParentPath, so the first chunk hit short-circuits the file. - Delete walk also reduces via ParentPath; when a parent file disappears, every chunk row (and any pre-existing bare-path row, for backward compatibility with rows written before this change) gets dropped. search/search.go - hybridMerge collapses chunk-path vector hits to parent via ParentPath before scope check, RRF accumulation, and hydration. A file with three chunk hits returns one result row, not three. Backward compatibility: pre-existing bare-path rows in brain_embeddings keep working — ParentPath returns them unchanged, knownParents handles them as if they were "wiki/foo.md#NNNN" hits, sync skips re-embed, and search dedup is a no-op for them. No migration required to ship. Tests: - chunk_test.go covers short / heading split / oversized section / content preservation / chunk numbering / parent-path stripping. - sync_test.go adds long-file chunking, single-chunk-row short file, skip-if-any-chunk-known, delete-all-chunks-of-disappeared-file. Existing tests updated for #NNNN paths. - search_test.go adds chunk-paths-dedupe-to-parent. Closes gitea/mathias/infra#38.
171 lines
5.2 KiB
Go
171 lines
5.2 KiB
Go
package vectorstore
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
// Embedder produces dense vectors. The embed package's Client satisfies
|
|
// this; it's declared locally so vectorstore doesn't depend on embed.
|
|
type Embedder interface {
|
|
Embed(ctx context.Context, text string) ([]float32, error)
|
|
}
|
|
|
|
// Store is the subset of PGStore that Sync needs. Lets tests stub it.
|
|
type Store interface {
|
|
KnownPaths(ctx context.Context) (map[string]struct{}, error)
|
|
Upsert(ctx context.Context, path string, embedding []float32) error
|
|
Delete(ctx context.Context, path string) error
|
|
}
|
|
|
|
// SyncResult tallies what Sync did. Returned for logs / metrics; callers
|
|
// generally don't act on the fields directly.
|
|
type SyncResult struct {
|
|
Added int
|
|
Updated int
|
|
Deleted int
|
|
Errors []error
|
|
}
|
|
|
|
// scanDirs is the set of brainDir subdirectories whose .md files are
|
|
// embedded for vector retrieval. wiki/ holds LLM-extracted entity and
|
|
// source pages; knowledge/ holds curated hand-written entries.
|
|
var scanDirs = []string{"wiki", "knowledge"}
|
|
|
|
// maxChunkBytes is the per-chunk byte budget passed to ChunkMarkdown.
|
|
// Sized to fit comfortably under nomic-embed-text's 2048-token default
|
|
// context (~4 chars/token for English markdown → ~8 KB ceiling; we sit
|
|
// at 4 KB to leave headroom for unicode, code blocks, and tokenizer
|
|
// variance).
|
|
const maxChunkBytes = 4000
|
|
|
|
// Sync brings the embedding store in line with brain/{wiki,knowledge}/
|
|
// on disk:
|
|
// - new files (in the tree, not in the store) get embedded + upserted
|
|
// - files whose mtime exceeds the store's updated_at get re-embedded
|
|
// - files no longer on disk get deleted from the store
|
|
//
|
|
// Designed to be called on a ticker. Best-effort: per-file errors are
|
|
// collected into SyncResult.Errors and do not abort the run.
|
|
func Sync(ctx context.Context, brainDir string, store Store, embedder Embedder) (SyncResult, error) {
|
|
var res SyncResult
|
|
if store == nil || embedder == nil {
|
|
return res, nil
|
|
}
|
|
|
|
known, err := store.KnownPaths(ctx)
|
|
if err != nil {
|
|
return res, fmt.Errorf("known paths: %w", err)
|
|
}
|
|
// Build a parent → "any chunk known?" set so we can skip files that
|
|
// already have at least one chunk row in the store.
|
|
knownParents := make(map[string]struct{}, len(known))
|
|
for p := range known {
|
|
knownParents[ParentPath(p)] = struct{}{}
|
|
}
|
|
seenParents := make(map[string]struct{})
|
|
|
|
for _, sub := range scanDirs {
|
|
root := filepath.Join(brainDir, sub)
|
|
if _, err := os.Stat(root); os.IsNotExist(err) {
|
|
continue
|
|
}
|
|
|
|
err = filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if d.IsDir() || !strings.HasSuffix(path, ".md") || d.Name() == "_index.md" {
|
|
return nil
|
|
}
|
|
rel, err := filepath.Rel(brainDir, path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
relSlash := filepath.ToSlash(rel)
|
|
seenParents[relSlash] = struct{}{}
|
|
|
|
if _, ok := knownParents[relSlash]; ok {
|
|
// File has at least one chunk in the store already.
|
|
// TODO: compare mtime once Store exposes updated_at so we
|
|
// re-embed on edit. For now, skip.
|
|
return nil
|
|
}
|
|
|
|
content, readErr := os.ReadFile(path)
|
|
if readErr != nil {
|
|
res.Errors = append(res.Errors, fmt.Errorf("read %s: %w", relSlash, readErr))
|
|
return nil
|
|
}
|
|
chunks := NumberChunks(relSlash, ChunkMarkdown(string(content), maxChunkBytes))
|
|
for _, ch := range chunks {
|
|
vec, embErr := embedder.Embed(ctx, ch.Content)
|
|
if embErr != nil {
|
|
res.Errors = append(res.Errors, fmt.Errorf("embed %s: %w", ch.Path, embErr))
|
|
continue
|
|
}
|
|
if upErr := store.Upsert(ctx, ch.Path, vec); upErr != nil {
|
|
res.Errors = append(res.Errors, fmt.Errorf("upsert %s: %w", ch.Path, upErr))
|
|
continue
|
|
}
|
|
res.Added++
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return res, fmt.Errorf("walk %s: %w", sub, err)
|
|
}
|
|
}
|
|
|
|
// Drop chunk rows whose parent file is gone.
|
|
for path := range known {
|
|
if _, ok := seenParents[ParentPath(path)]; ok {
|
|
continue
|
|
}
|
|
if err := store.Delete(ctx, path); err != nil {
|
|
res.Errors = append(res.Errors, fmt.Errorf("delete %s: %w", path, err))
|
|
continue
|
|
}
|
|
res.Deleted++
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
// StartSync launches Sync on a ticker in a background goroutine. The
|
|
// goroutine exits when ctx is cancelled. Failures are logged via slog.
|
|
func StartSync(ctx context.Context, brainDir string, store Store, embedder Embedder, interval time.Duration) {
|
|
if interval <= 0 {
|
|
interval = 5 * time.Minute
|
|
}
|
|
go func() {
|
|
t := time.NewTicker(interval)
|
|
defer t.Stop()
|
|
// Run once immediately so first-boot doesn't wait a full tick.
|
|
if r, err := Sync(ctx, brainDir, store, embedder); err != nil {
|
|
slog.Error("embed sync failed", "err", err)
|
|
} else if r.Added+r.Deleted > 0 || len(r.Errors) > 0 {
|
|
slog.Info("embed sync", "added", r.Added, "deleted", r.Deleted, "errors", len(r.Errors))
|
|
for _, e := range r.Errors {
|
|
slog.Warn("embed sync item failed", "err", e)
|
|
}
|
|
}
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-t.C:
|
|
if r, err := Sync(ctx, brainDir, store, embedder); err != nil {
|
|
slog.Error("embed sync failed", "err", err)
|
|
} else if r.Added+r.Deleted > 0 || len(r.Errors) > 0 {
|
|
slog.Info("embed sync", "added", r.Added, "deleted", r.Deleted, "errors", len(r.Errors))
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|