embedding

package
v0.17.8 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 23, 2026 License: MIT Imports: 37 Imported by: 0

Documentation

Overview

Package embedding provides ONNX-based embedding infrastructure for sprout.

The ONNX runtime is a shared base that any ONNX-based model can use: embedding models (EmbeddingGemma) and future generation models (Gemma 3 2B).

Package embedding provides an embedding provider interface, vector store, and similarity utilities for semantic search and duplicate detection.

Package embedding — HNSW-backed vector store for fast approximate nearest neighbor search. Uses github.com/coder/hnsw for the graph index.

Index

Constants

View Source 
const (
	// QueryPrefix is used for search/query embeddings.
	QueryPrefix = "task: search result | query: "
	// DocumentPrefix is used for document/content embeddings.
	DocumentPrefix = "title: none | text: "
	// CodeQueryPrefix is used for code-specific search queries.
	CodeQueryPrefix = "task: code retrieval | query: "
)

EmbeddingGemma prompt prefixes for task-specific embedding. These are prepended to text before tokenization to guide the model toward the right embedding space for the task.

View Source 
const MaxDepth = 15

MaxDepth is the maximum directory nesting depth WalkCodeFiles will descend into. Deeper directories are pruned to avoid pathological directory trees (e.g., deeply nested generated code).

View Source 
const MaxFileCount = 10000

MaxFileCount is the maximum number of source files WalkCodeFiles will collect before stopping. Once this limit is reached, the walk exits early and returns the files collected so far.

View Source 
const ProgressInterval = 500

ProgressInterval controls how many files must be processed before a progress event is emitted (both during walk and batch embedding).

View Source 
const WalkTimeout = 30 * time.Second

WalkTimeout is the absolute maximum time allowed for WalkCodeFiles to enumerate files across the workspace. After this duration the walk is cancelled and a partial result is returned.

Variables

View Source 
var ErrStoreClosed = errors.New("embedding: hnsw store is closed")

ErrStoreClosed is returned by mutating operations (Store, ReplaceAll, DeleteByFile, DeleteByIDs, Save) invoked after Close. It exists to turn a pre-existing nil-map panic into a recoverable error when a background goroutine races with a Disable/Close call from the foreground.

Functions

func CheckFileChanged 

func CheckFileChanged(path string, lastMtimeNano int64) (bool, error)

CheckFileChanged stats the file at path and returns true if its mtime (UnixNano) differs from lastMtimeNano. Returns (false, nil) if the file does not exist (it was deleted).

func ClearEmbeddingFiles 

func ClearEmbeddingFiles(indexDir string, fileType string) (int, error)

ClearEmbeddingFiles removes embedding index files from the given directory. fileType should be one of: "code", "conversation_turn", "memory", "all". For "memory", it clears the same files as "conversation_turn" since memories are stored in the conversation_turns index alongside conversation turns. Returns the number of files actually deleted.

func CosineSimilarity 

func CosineSimilarity(a, b []float32) float32

CosineSimilarity computes the cosine similarity between two vectors. It returns a value in [-1, 1], where 1 means identical direction. Vectors of different lengths return 0.

func DefaultModelDir 

func DefaultModelDir() string

DefaultModelDir returns the default model directory path. Priority: SPROUT_MODELS_DIR env > SPROUT_CONFIG/SPROUT_CONFIG env > ~/.config/sprout

func DownloadModel 

func DownloadModel(ctx context.Context, modelDir string, cfg ModelConfig) error

DownloadModel ensures the model and tokenizer files exist in modelDir for the given ModelConfig. If files already exist with matching checksums, download is skipped. Progress is reported via the callback (0.0 to 1.0).

func FormatDuplicateWarning 

func FormatDuplicateWarning(matches []QueryResult) string

FormatDuplicateWarning formats duplicate matches as an agent-internal note (not a user-facing warning). The format is designed for the agent to silently evaluate whether its code overlaps with existing functionality.

func HashContent 

func HashContent(content []byte) string

HashContent computes a SHA-256 hex digest of the given content.

func IsBinaryFile 

func IsBinaryFile(path string) bool

IsBinaryFile reports whether the file at path appears to be binary. It reads up to the first 8 KB and checks for a null byte (0x00).

func IsSupportedIndexableFile 

func IsSupportedIndexableFile(path string) bool

IsSupportedIndexableFile checks if a file path should be indexed at the file level. Used by the indexing pipeline to decide which files to process with FileExtractor.

func Normalize 

func Normalize(v []float32) []float32

Normalize returns a new vector that has been normalized to unit length. If the input vector is zero-length or empty, it returns an empty slice.

func NormalizePathToWorkspace 

func NormalizePathToWorkspace(workspaceRoot, path string) string

NormalizePathToWorkspace normalizes a file path to a relative path from the workspace root for consistent comparison. If workspaceRoot is empty or the path cannot be resolved relative to it, returns the cleaned path as-is.

func SaveManifest 

func SaveManifest(path string, m *BuildManifest) error

SaveManifest writes the manifest to path atomically (temp file + rename).

func ScoreWithDecay 

func ScoreWithDecay(similarity float64, timestamp time.Time, now time.Time) float64

ScoreWithDecay computes a time-decayed similarity score by combining a raw similarity value with exponential decay based on the age of the record. The decay uses a 30-day half-life, meaning the score is halved every 30 days.

Parameters:

  • similarity: The raw cosine similarity value (typically in [-1, 1])
  • timestamp: The timestamp when the record was created/indexed
  • now: The current reference time (usually time.Now())

Returns: The decayed score as a float64. Recent records (same day) have decay ≈ 1.0. Old records are deprioritized but never eliminated completely.

func SetPackageDebugLogging 

func SetPackageDebugLogging(enabled bool)

SetPackageDebugLogging toggles the debug gate at runtime. Useful for tests and for the agent's --debug flag wiring.

func ShouldIgnorePath 

func ShouldIgnorePath(path string, repoRoot string) bool

ShouldIgnorePath reports whether the given path should be excluded from indexing. It applies two layers of filtering:

Layer 1 — Hard-coded directory and filename patterns. Layer 2 — Binary file detection (null byte in first 8 KB).

The repoRoot parameter is reserved for future gitignore-based filtering.

func StaticEmbed added in v0.16.19 

func StaticEmbed(_ context.Context, _ string) ([]float32, error)

StaticEmbed is a no-op on native builds. The static provider is only used in WASM builds as a fallback when the ONNX bridge isn't available.

func StaticProviderName added in v0.16.19 

func StaticProviderName() string

StaticProviderName returns the name of the static embedding provider. On native builds, the static provider is not used (ONNX is preferred), but the function exists for API compatibility with WASM builds.

func WalkAllIndexableFiles 

func WalkAllIndexableFiles(ctx context.Context, root string) ([]string, error)

WalkAllIndexableFiles walks the directory tree rooted at root and returns all file paths that should be indexed — both code files (for symbol extraction) and non-code files (for file-level embedding). It includes all extensions from supportedCodeExtensions plus supported non-code extensions (.md, .yaml, .json, .sh, etc.) and special filenames (Makefile, Dockerfile, .gitignore, etc.).

It accepts a context for cancellation and applies three protections:

  • A 30-second absolute timeout (WalkTimeout).
  • A maximum directory depth of 15 (MaxDepth).
  • A cap of 10,000 collected files (MaxFileCount).

Progress is logged every ProgressInterval files. If the context is cancelled or any limit is hit, the files collected so far are returned with no error (partial result).

func WalkCodeFiles 

func WalkCodeFiles(ctx context.Context, root string) ([]string, error)

WalkCodeFiles walks the directory tree rooted at root and returns all file paths that should be indexed (i.e., those that pass ShouldIgnorePath). Only files with recognized extensions (.go, .ts, .tsx, .js, .tsx, .mjs, .py) are included. Directories matching Layer 1 skip patterns are pruned (no recursion).

It accepts a context for cancellation and applies three protections:

  • A 30-second absolute timeout (WalkTimeout).
  • A maximum directory depth of 15 (MaxDepth).
  • A cap of 10,000 collected files (MaxFileCount).

Progress is logged every ProgressInterval files. If the context is cancelled or any limit is hit, the files collected so far are returned with no error (partial result).

Types

type BuildManifest 

type BuildManifest struct {
	// Files maps file path → mtime UnixNano at last successful build.
	Files map[string]int64 `json:"files"`

	// ModelHash is the model hash used when this manifest was created.
	// If the model changes, the manifest is invalidated.
	ModelHash string `json:"modelHash"`
}

BuildManifest tracks file modification times from the last successful BuildIndex call. It allows subsequent builds to skip parsing unchanged files, turning a multi-minute full parse into a ~2-second stat sweep.

func BuildManifestFromFiles 

func BuildManifestFromFiles(files []string, modelHash string) *BuildManifest

BuildManifestFromFiles creates a new manifest by statting all given files.

func LoadManifest 

func LoadManifest(path string) (*BuildManifest, error)

LoadManifest loads a manifest from a JSON file. Returns (nil, nil) if the file does not exist. The caller should check both return values.

type BuildResult 

type BuildResult struct {
	Stats *IndexStats
	Err   error
}

BuildResult carries the result of a background index build.

type CheckDuplicatesResult 

type CheckDuplicatesResult struct {
	// Duplicates is a list of potential duplicate matches, sorted by similarity.
	Duplicates []QueryResult
	// WarningText is a formatted warning message for the agent, or empty if no duplicates.
	WarningText string
}

CheckDuplicatesResult holds the result of a duplicate check.

func CheckFileForDuplicates 

func CheckFileForDuplicates(ctx context.Context, mgr *IndexManager, filePath string, content string, workspaceRoot string, threshold float32, topK int) (*CheckDuplicatesResult, error)

CheckFileForDuplicates checks if any functions in the given file content have semantically similar existing code in the index.

It works by extracting code units from the content, embedding each one, and querying the existing index for similar records. Self-matches (same ID or same file path) are filtered out using workspace-relative path comparison.

The top-K overall matches above the threshold are returned, sorted by similarity descending.

type CodeUnit 

type CodeUnit struct {
	// ID is a unique identifier produced by makeUnitID:
	// "<file>:<name>#L<startLine>" for code-unit extractors (Go/Python/TS).
	// File-level extractors (extractor_file.go) use the bare file path
	// instead since one record represents the whole file.
	ID string `json:"id"`

	// File is the file path the code unit comes from.
	File string `json:"file"`

	// Name is the symbol name (e.g., "funcName" or "(*Receiver).Method").
	Name string `json:"name"`

	// Signature is the full function signature text.
	Signature string `json:"signature"`

	// Body is the source text of the function body.
	Body string `json:"body"`

	// StartLine is the 1-based starting line number.
	StartLine int `json:"startLine"`

	// EndLine is the 1-based ending line number.
	EndLine int `json:"endLine"`

	// Language is the programming language identifier (e.g., "go").
	Language string `json:"language"`

	// Hash is a SHA-256 hex digest of Signature+Body for deduplication.
	Hash string `json:"hash"`
}

CodeUnit represents a single unit of code (e.g., a function) extracted from a source file.

func ExtractFromFile 

func ExtractFromFile(path string, opts ...ExtractOption) ([]CodeUnit, error)

ExtractFromFile extracts code units from the given file path using the language-specific extractor determined by file extension. Returns an empty slice (no error) for unsupported file types.

func ExtractGoFile 

func ExtractGoFile(path string, opts ...ExtractOption) ([]CodeUnit, error)

ExtractGoFile parses a Go source file and extracts all top-level function declarations as CodeUnit values. Test functions (Test*, Benchmark*, Fuzz*) are excluded by default; use WithIncludeTests to change this.

func ExtractPyFile 

func ExtractPyFile(path string, opts ...ExtractOption) ([]CodeUnit, error)

ExtractPyFile parses a Python source file and extracts code units (functions, classes, methods) as CodeUnit values. Test functions (prefixed with test_) are excluded by default; use WithIncludeTests to change this.

func ExtractTSFile 

func ExtractTSFile(path string, opts ...ExtractOption) ([]CodeUnit, error)

ExtractTSFile parses a TypeScript or JavaScript source file and extracts code units (functions, arrow functions, methods, classes) as CodeUnit values. Test files (.test.ts, .spec.js, etc.) are excluded by default; use WithIncludeTests to change this.

func (*CodeUnit) ComputeHash 

func (c *CodeUnit) ComputeHash()

ComputeHash calculates a SHA-256 hex digest of the signature and body.

type ConversationStore 

type ConversationStore struct {
	// contains filtered or unexported fields
}

ConversationStore wraps a VectorStore for storing and querying conversation turn embeddings. It provides a user-scoped persistent store that survives across workspace changes.

The store uses the same static embedding provider as the code index, and maintains its own in-memory cache of records for fast queries.

func NewConversationStore 

func NewConversationStore(provider EmbeddingProvider, filePath string, modelHash string) (*ConversationStore, error)

NewConversationStore creates a new conversation store at the given path. The parent directory is created if it does not exist. If a file already exists at path, its records are loaded into memory.

The provider is used for embedding operations and is not closed when the store is closed (the provider's lifecycle is managed externally). The modelHash is used to detect model changes and invalidate stale records.

func (*ConversationStore) Close 

func (s *ConversationStore) Close() error

Close releases any resources held by the store. The embedding provider is not closed (its lifecycle is managed externally).

func (*ConversationStore) DeleteMemoryByName 

func (s *ConversationStore) DeleteMemoryByName(name string) error

DeleteMemoryByName removes all memory records with the given name. This is useful when a memory file is deleted or updated and its old embedding should be removed from the store.

func (*ConversationStore) LoadAll 

func (s *ConversationStore) LoadAll() ([]VectorRecord, error)

LoadAll returns a copy of all records currently in the store.

func (*ConversationStore) Provider 

func (s *ConversationStore) Provider() EmbeddingProvider

Provider returns the embedding provider used by this store. The provider remains usable after Close() since its lifecycle is managed externally by the EmbeddingManager.

func (*ConversationStore) Query 

func (s *ConversationStore) Query(vec []float32, topK int, threshold float32) ([]QueryResult, error)

Query returns the top-K records most similar to vec, with similarity >= threshold.

func (*ConversationStore) QueryMemories 

func (s *ConversationStore) QueryMemories(ctx context.Context, query string, topK int, threshold float32) ([]QueryResult, error)

QueryMemories searches memory records by embedding the query and returning top-K results. Results are filtered to only include records with Type "memory".

func (*ConversationStore) Size 

func (s *ConversationStore) Size() int

Size returns the number of records currently in the store.

func (*ConversationStore) Store 

func (s *ConversationStore) Store(records []VectorRecord) error

Store adds records to the conversation store. If a record with the same ID already exists, it is replaced. Records are kept sorted by ID for deterministic output.

func (*ConversationStore) StoreMemory 

func (s *ConversationStore) StoreMemory(ctx context.Context, name string, content string) error

StoreMemory embeds memory content and stores it as a VectorRecord with Type "memory". The record ID is "memory:" + name to ensure unique naming and easy lookup.

type EmbedOptions 

type EmbedOptions struct {
	// Prefix is prepended to each text before tokenization.
	// This is used for task-specific prompting (e.g., query vs document prefixes).
	Prefix string
}

EmbedOptions holds optional parameters for embedding operations.

type EmbeddingManager 

type EmbeddingManager struct {
	// contains filtered or unexported fields
}

EmbeddingManager manages the embedding index lifecycle. It lazily initializes the ONNX embedding provider and IndexManager on first use, and caches them for subsequent calls.

func NewEmbeddingManager 

func NewEmbeddingManager(cfg *configuration.EmbeddingIndexConfig, workspaceRoot string) *EmbeddingManager

NewEmbeddingManager creates a new manager with the given config. The manager is NOT initialized until Init() or a query method is called.

func (*EmbeddingManager) AutoBuildWhenReady 

func (m *EmbeddingManager) AutoBuildWhenReady()

AutoBuildWhenReady runs a background index build after a short delay. This is called at agent startup so the index is ready for duplicate detection and context enrichment without waiting for an explicit query. A 2-minute timeout prevents the build from hanging indefinitely.

Two teardown paths are honored so a DisableEmbeddingIndex call arriving during the startup sleep (or during Init/Build) does not race into a closed store and panic:

  1. The 3-second startup sleep selects on m.closeCh() so Close() can wake it early.
  2. After the sleep returns, m.closeCh() is re-checked *before* the BuildIndex call. This catches the case where Close() ran while the sleep was in flight (sleep saw the wake-up but the goroutine still proceeded because the select picked the timer branch first).

As a last line of defense, HNSWStore.Store/ReplaceAll/DeleteByFile/ DeleteByIDs/Save return ErrStoreClosed instead of panicking on a nil records map if the goroutine still loses the race.

func (*EmbeddingManager) BuildIndex 

func (m *EmbeddingManager) BuildIndex(ctx context.Context) (*IndexStats, error)

BuildIndex runs a full index build for the workspace. It acquires the building lock, validates workspace size, and delegates to buildIndexLocked for the actual work.

func (*EmbeddingManager) BuildIndexBackground 

func (m *EmbeddingManager) BuildIndexBackground(ctx context.Context) <-chan *BuildResult

BuildIndexBackground starts an index build in a background goroutine and returns a channel on which the result (or error) will be delivered. This must be used when called from HTTP handlers or other code paths where blocking would cause a timeout.

The returned channel is non-buffered and the caller should read from it once to retrieve the result. The context passed to the caller is used for cancellation; if the context is cancelled, the build is interrupted gracefully (partial results may be stored).

func (*EmbeddingManager) CheckDuplicates 

func (m *EmbeddingManager) CheckDuplicates(ctx context.Context, filePath string, content string) (*CheckDuplicatesResult, error)

CheckDuplicates checks if file content duplicates existing code.

func (*EmbeddingManager) Close 

func (m *EmbeddingManager) Close() error

Close releases all resources.

func (*EmbeddingManager) CloseNotify added in v0.16.18 

func (m *EmbeddingManager) CloseNotify() <-chan struct{}

CloseNotify returns a channel that is closed when the manager is closed. Long-running goroutines owned by other packages (e.g. agent.MigrateMemories) select on this channel so they can abort when DisableEmbeddingIndex tears the manager down. The returned channel is the same one internal goroutines see, so a single Close() wakes every waiter.

func (*EmbeddingManager) GetConversationStore 

func (m *EmbeddingManager) GetConversationStore(ctx context.Context) (*ConversationStore, error)

GetConversationStore returns the conversation store, creating it lazily on first use. The store is user-scoped and lives at {indexDir}/conversation_turns.hnsw. Multiple calls return the same instance.

func (*EmbeddingManager) IndexSize 

func (m *EmbeddingManager) IndexSize() int

IndexSize returns the number of records in the vector store. Returns 0 and a nil error if the manager is not yet initialized.

func (*EmbeddingManager) Init 

func (m *EmbeddingManager) Init(ctx context.Context) error

Init initializes the ONNX embedding provider and opens the vector store. This is idempotent — calling it multiple times is safe. If a previous Init() failed, the cached error is returned immediately.

func (*EmbeddingManager) InitError 

func (m *EmbeddingManager) InitError() error

InitError returns the error from a previous failed Init() call, or nil if initialization succeeded or has never been attempted.

func (*EmbeddingManager) IsBuilding 

func (m *EmbeddingManager) IsBuilding() bool

IsBuilding returns true if an index build is currently in progress.

func (*EmbeddingManager) IsInitialized 

func (m *EmbeddingManager) IsInitialized() bool

IsInitialized returns whether the manager has been initialized. Safe to call without holding m.mu — initialized is an atomic so this never blocks, even while Init() is running and holding m.mu during ONNX loading.

func (*EmbeddingManager) ModelHash 

func (m *EmbeddingManager) ModelHash() string

ModelHash returns the active embedding provider's model hash, or "" if no provider is currently initialized. Used by tests to re-open persisted stores with the same hash so the model-change invalidation logic doesn't wipe them.

func (*EmbeddingManager) QuerySimilar 

func (m *EmbeddingManager) QuerySimilar(ctx context.Context, query string, topK int, threshold float32) ([]QueryResult, error)

QuerySimilar searches for code similar to the given query text.

func (*EmbeddingManager) SetForTesting 

func (m *EmbeddingManager) SetForTesting(provider EmbeddingProvider, store VectorStore, indexMgr *IndexManager)

SetForTesting injects mock provider, store, and indexManager for testing. This bypasses Init() so tests can run without an ONNX runtime. NOT for production use.

It also resolves indexDir (mirroring the logic in initLocked) so that GetConversationStore creates the conversation store in the expected location rather than leaking a file into the process working directory.

func (*EmbeddingManager) UpdateFile 

func (m *EmbeddingManager) UpdateFile(ctx context.Context, filePath string) error

UpdateFile incrementally updates the index for a single file.

func (*EmbeddingManager) UpdateFromGitDiff 

func (m *EmbeddingManager) UpdateFromGitDiff(ctx context.Context) (*IndexStats, error)

UpdateFromGitDiff incrementally updates the index by examining git-tracked files that have changed, been added, or been created since the last build.

type EmbeddingProvider 

type EmbeddingProvider interface {
	// Embed returns a fixed-dimension embedding vector for the given text.
	Embed(ctx context.Context, text string) ([]float32, error)

	// EmbedBatch returns embedding vectors for multiple texts.
	// The returned slice has the same length and order as input.
	EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

	// Dimensions returns the dimensionality of vectors produced by this provider.
	Dimensions() int

	// Name returns a human-readable identifier for the provider.
	Name() string

	// ModelHash returns a SHA-256 hex digest of the model data. Used to detect
	// model changes and invalidate stale store records.
	ModelHash() string

	// EmbedWithPrefix returns an embedding with a task-specific prefix prepended
	// to the text before tokenization. Implementations that don't support prefixes
	// should prepend the prefix to the text and delegate to Embed.
	EmbedWithPrefix(ctx context.Context, text string, prefix string) ([]float32, error)

	// EmbedBatchWithPrefix returns embeddings with a task-specific prefix
	// prepended to each text before tokenization.
	EmbedBatchWithPrefix(ctx context.Context, texts []string, prefix string) ([][]float32, error)

	// Close releases any resources held by the provider.
	Close() error
}

EmbeddingProvider produces vector embeddings for text input. Implementations typically wrap an external model (e.g., OpenAI, local Ollama).

type ExtractConfig 

type ExtractConfig struct {
	// IncludeTests controls whether test functions (Test*, Benchmark*, Fuzz*)
	// are included in the extraction. Default: false.
	IncludeTests bool
}

ExtractConfig holds options for code extraction.

func (*ExtractConfig) ApplyOptions 

func (c *ExtractConfig) ApplyOptions(opts ...ExtractOption)

ApplyOptions applies a list of ExtractOption functions to an ExtractConfig.

type ExtractOption 

type ExtractOption func(*ExtractConfig)

ExtractOption configures behavior for code extraction.

func WithIncludeTests 

func WithIncludeTests(include bool) ExtractOption

WithIncludeTests returns an ExtractOption that sets whether test functions are included in extraction.

type FileExtractor 

type FileExtractor struct {
	// contains filtered or unexported fields
}

FileExtractor produces file-level embeddings for non-code files.

func NewFileExtractor 

func NewFileExtractor(maxFileBytes int) *FileExtractor

NewFileExtractor creates a FileExtractor that truncates files larger than maxFileBytes. If maxFileBytes is 0, a default of 8000 bytes is used.

func (*FileExtractor) Extract 

func (e *FileExtractor) Extract(path string, content []byte) ([]CodeUnit, error)

Extract produces a single CodeUnit representing the entire file content. Returns an empty slice (no error) for unsupported file types.

type GemmaTokenizer 

type GemmaTokenizer struct {
	// contains filtered or unexported fields
}

GemmaTokenizer implements the HuggingFace BPE tokenization pipeline used by Google's EmbeddingGemma-300M model. It is intentionally narrow: it covers the exact normalizer / pre-tokenizer / model / added-tokens combination that EmbeddingGemma ships with, rather than the full HF tokenizers schema.

Pipeline (matches HuggingFace tokenizers semantics):

  1. Split input around added-token strings (e.g. "\n", "\t", "<bos>"). The matched runs are emitted directly as their IDs; everything else falls through to the BPE path.
  2. Normalize each non-added segment by replacing " " (U+0020) with the SentencePiece space marker "▁" (U+2581), per the Replace normalizer.
  3. Apply rank-ordered BPE merges to the normalized text, treated as a sequence of single-rune symbols.
  4. Map each resulting symbol to its vocab id, falling back to <unk> on miss.

Decoding is not implemented — the embedding path only needs Encode().

func NewGemmaTokenizer 

func NewGemmaTokenizer(path string) (*GemmaTokenizer, error)

NewGemmaTokenizer parses a HuggingFace tokenizer.json file produced for an EmbeddingGemma-class model and returns an encoder.

func (*GemmaTokenizer) Encode 

func (t *GemmaTokenizer) Encode(text string) []int32

Encode tokenizes text into a sequence of token ids matching what the HuggingFace `tokenizers` reference produces for the same input (modulo the BOS/EOS markers, which Encode does NOT add — use EncodeWithBOSAndEOS for that).

func (*GemmaTokenizer) EncodeBatch 

func (t *GemmaTokenizer) EncodeBatch(texts []string, padID int32) [][]int32

EncodeBatch tokenizes multiple texts and right-pads each to the longest length with padID, returning a rectangular [][]int32. Used by the ONNX provider to feed batched input tensors.

func (*GemmaTokenizer) EncodeWithBOS 

func (t *GemmaTokenizer) EncodeWithBOS(text string) []int32

EncodeWithBOS prepends BOS but does not append EOS. Kept for callers that only want the prefix marker.

func (*GemmaTokenizer) EncodeWithBOSAndEOS 

func (t *GemmaTokenizer) EncodeWithBOSAndEOS(text string) []int32

EncodeWithBOSAndEOS returns Encode(text) with the BOS id prepended and the EOS id appended (when they were resolved from added_tokens). This matches what HuggingFace's encode() returns for EmbeddingGemma with default post-processing — including the [BOS, EOS] pair for empty input.

func (*GemmaTokenizer) MaskBatch 

func (t *GemmaTokenizer) MaskBatch(encoded [][]int32, padID int32) [][]int64

MaskBatch produces attention masks (1 = real token, 0 = padding) for a batch of encoded sequences. Length matches each sequence row in encoded.

func (*GemmaTokenizer) TokenIDs 

func (t *GemmaTokenizer) TokenIDs(text string) []int32

TokenIDs is an alias for Encode, kept for callers that prefer the name.

func (*GemmaTokenizer) VocabSize 

func (t *GemmaTokenizer) VocabSize() int

VocabSize returns the model vocabulary size (not counting added tokens that share ids with vocab entries, which is rare).

type HNSWStore 

type HNSWStore struct {
	// contains filtered or unexported fields
}

HNSWStore is a thread-safe VectorStore backed by an HNSW index. The graph stores vectors keyed by record ID; a separate map holds full VectorRecord metadata. The index is persisted to disk via hnsw.SavedGraph and a sidecar .meta file tracks the model hash.

func NewHNSWStore 

func NewHNSWStore(indexPath string, modelHash string) (*HNSWStore, error)

NewHNSWStore creates or loads an HNSW-backed vector store.

indexPath is the path to the persisted HNSW index file. modelHash is the current provider's model hash; if it differs from the stored hash, the index is cleared to force a full rebuild.

func (*HNSWStore) Close 

func (s *HNSWStore) Close() error

Close saves the graph and records to disk if there are pending changes, then clears internal state. It is safe to call multiple times — subsequent calls return ErrStoreClosed-on-mutate semantics without re-persisting.

After Close, mutating operations (Store, ReplaceAll, DeleteByFile, DeleteByIDs, Save) return ErrStoreClosed rather than panicking on the nil records map. Read operations (Size, LoadAll, Query) continue to observe whatever state was on disk before Close, since Close is responsible for the final flush.

func (*HNSWStore) DeleteByFile 

func (s *HNSWStore) DeleteByFile(filePath string) error

DeleteByFile removes all records whose File path matches filePath.

func (*HNSWStore) DeleteByIDs 

func (s *HNSWStore) DeleteByIDs(ids []string) error

DeleteByIDs removes records with the given IDs in a single batched operation. IDs not present in the store are silently skipped. Saves to disk only once, regardless of how many IDs are deleted.

func (*HNSWStore) LoadAll 

func (s *HNSWStore) LoadAll() ([]VectorRecord, error)

LoadAll returns all records currently in the store.

func (*HNSWStore) Query 

func (s *HNSWStore) Query(vec []float32, topK int, threshold float32) ([]QueryResult, error)

Query returns the top-K records most similar to vec, with similarity >= threshold. Uses cosine distance from the hnsw library; similarity = 1 - distance.

func (*HNSWStore) ReplaceAll 

func (s *HNSWStore) ReplaceAll(records []VectorRecord) error

ReplaceAll discards the existing index and builds a new one from records.

Records are deduplicated by ID before hitting the HNSW graph. The upstream extractors now produce line-disambiguated IDs (see `extractor.go:makeUnitID` — every code-unit ID is `<file>:<name>#L<startLine>`), so collisions should be impossible at the source. This dedupe is kept as belt-and-suspenders: any future extractor / caller that hands us a slice with duplicate IDs (intentional or otherwise) won't trigger the coder/hnsw library's `g.Len() == preLen+1` invariant panic at graph.go:405. Replacement semantics match the on-disk record store: the LAST record with a given ID wins.

History: this branch was added after a real "node not added" panic during sprout's first-run auto-build on a workspace where the Python/TS extractors produced `path:methodname` collisions across classes in the same file.

func (*HNSWStore) Save 

func (s *HNSWStore) Save() error

Save persists the graph and metadata to disk. Records are written before the graph so a crash mid-save leaves records ahead of the graph (recoverable) rather than vice versa (panic-prone).

func (*HNSWStore) Size 

func (s *HNSWStore) Size() int

Size returns the number of records in the store.

func (*HNSWStore) Store 

func (s *HNSWStore) Store(records []VectorRecord) error

Store adds records to the store. Existing records with the same ID are replaced.

type IndexManager 

type IndexManager struct {
	// contains filtered or unexported fields
}

IndexManager orchestrates code extraction, embedding, and storage.

func NewIndexManager 

func NewIndexManager(provider EmbeddingProvider, store VectorStore, opts IndexOptions) *IndexManager

NewIndexManager creates an IndexManager with the given provider, store, and options. Default BatchSize is 32, default MaxBodyLen is 2000.

func (*IndexManager) BuildIndex 

func (m *IndexManager) BuildIndex(ctx context.Context, rootDir string) (*IndexStats, error)

BuildIndex walks rootDir, extracts code units, embeds them, and stores them. Uses incremental rebuild: loads existing records, compares content hashes, and only re-embeds changed or new files. Deleted files have their records removed from the store. When ManifestPath is set, uses an mtime-based manifest to skip parsing unchanged files entirely, turning a multi-minute full parse into a ~2-second stat sweep on warm indexes. When IndexFileLevel is enabled, also indexes non-code files at the file level.

func (*IndexManager) CheckDuplicates 

func (m *IndexManager) CheckDuplicates(ctx context.Context, codeText string, topK int, threshold float32) ([]QueryResult, error)

CheckDuplicates is like QuerySimilar but uses a default threshold of 0.90.

func (*IndexManager) QuerySimilar 

func (m *IndexManager) QuerySimilar(ctx context.Context, query string, topK int, threshold float32) ([]QueryResult, error)

QuerySimilar embeds query text and returns the top-K most similar records above threshold.

func (*IndexManager) UpdateFile 

func (m *IndexManager) UpdateFile(ctx context.Context, filePath string) error

UpdateFile re-indexes a single file: deletes old records, extracts, embeds, and stores. Handles both code files (symbol extraction) and non-code files (file-level embedding) when IndexFileLevel is enabled.

func (*IndexManager) UpdateFromGitDiff 

func (m *IndexManager) UpdateFromGitDiff(ctx context.Context, repoRoot string) (*IndexStats, error)

UpdateFromGitDiff incrementally updates the index by examining files changed since the last index build. It uses git diff to detect modified, added, and deleted files. Deleted files have their records removed from the store, while changed/new files are re-indexed.

type IndexOptions 

type IndexOptions struct {
	// IncludeTests controls whether test functions are indexed.
	IncludeTests bool
	// BatchSize controls how many code units are embedded per batch.
	BatchSize int
	// MaxBodyLen truncates CodeUnit.Body to this many bytes before embedding (0 = no limit).
	MaxBodyLen int
	// IndexFileLevel controls whether non-code files (markdown, configs, etc.)
	// are indexed at the file level. When true, files like README.md, package.json,
	// Dockerfile, etc. are indexed as single records with Type="file".
	IndexFileLevel bool
	// ManifestPath is the path to the build manifest file that tracks file
	// modification times from the last successful build. When set, BuildIndex
	// uses the manifest to skip parsing unchanged files.
	ManifestPath string
}

IndexOptions configures the behavior of IndexManager.

type IndexStats 

type IndexStats struct {
	FilesProcessed int
	UnitsExtracted int
	UnitsEmbedded  int
	Duration       time.Duration
}

IndexStats reports the results of an indexing operation.

type ManifestDiff 

type ManifestDiff struct {
	ChangedFiles   []string
	UnchangedFiles []string
	DeletedFiles   []string

	// ManifestInvalidated is true when the model hash changed and all
	// files must be re-embedded even if their content hashes match.
	ManifestInvalidated bool
}

ManifestDiff holds the result of comparing the current workspace state against a stored manifest.

func DiffManifest 

func DiffManifest(ctx context.Context, manifest *BuildManifest, currentModelHash, rootDir string, indexFileLevel bool) (*ManifestDiff, error)

DiffManifest compares the current workspace state against a stored manifest. It returns a ManifestDiff indicating which files are changed, unchanged, or deleted. If manifest.ModelHash differs from currentModelHash, all files are treated as changed (manifest invalidated).

type ModelConfig 

type ModelConfig struct {
	Name          string // e.g. "embeddinggemma-300m"
	ModelURL      string // HuggingFace download URL for the .onnx graph file
	TokenizerURL  string // HuggingFace download URL for tokenizer.json
	ModelHash     string // SHA256 hex of model file (empty = skip verification)
	TokenizerHash string // SHA256 hex of tokenizer file

	// ModelDataURL is the URL for the external weights blob (e.g.
	// model_fp16.onnx_data) that ONNX Runtime loads as a sibling of the
	// .onnx graph file. Required for models that use external data;
	// leave empty for self-contained .onnx files.
	ModelDataURL  string
	ModelDataHash string

	// ModelFilename / ModelDataFilename are the on-disk basenames the
	// downloader writes the graph + weights blob to (and the manager
	// reads them back from). Defaults to "model_q4.onnx" and
	// "model_q4.onnx_data" for back-compat with the original q4 layout;
	// callers using a different quantization (e.g. fp16) set these to
	// match the variant they ship.
	ModelFilename     string
	ModelDataFilename string

	// FullDims is the model's native output dimensionality (e.g., 768 for EmbeddingGemma).
	// Used to allocate the output tensor in runInference.
	FullDims int

	// Dims is the desired output dimensionality after optional MRL truncation.
	// Must be <= FullDims. When equal to FullDims, no truncation is applied.
	Dims int
}

ModelConfig describes an ONNX model to download.

func EmbeddingGemma300MConfig 

func EmbeddingGemma300MConfig() ModelConfig

EmbeddingGemma300MConfig returns the predefined config for Google's EmbeddingGemma-300M (308M parameter) model — the actual published model name. The .onnx graph is small (~692 KB) but references an external weights blob (~167 MB) that must be downloaded into the same directory; both URLs are set.

Source: the community ONNX export at onnx-community/embeddinggemma-300m-ONNX, since the official Google repo ships SafeTensors only.

Ships the Q4f16 variant: 4-bit weights with FP16 activations/compute (via ONNX Runtime's MatMulNBits operator). Background: the original Q4 export (model_q4.onnx) dequantizes through BFloat16, and ORT 1.25.1's SimplifiedLayerNormalization kernel on macOS arm64 doesn't have type traits for BF16 — every embed crashed with "GetElementType is not implemented" on M1+. The Q4f16 path stays in FP16 throughout and dodges the whole class of platform-fragile ops, while keeping the disk footprint near the original Q4 (168 MB vs 197 MB) AND improving latency over FP16 (single-call 11.4 ms vs 16.1 ms p50 on M1+; quality margin +0.31 vs +0.30 — statistically identical). See `embeddings-bench/results/2026-06-07-broad-sweep.md` for the wider model comparison that informed this choice.

Hashes pin the upstream files we validated end-to-end. The downloader rejects mismatched files, so a poisoned mirror or MITM swap on the HuggingFace path fails closed instead of silently feeding a tampered model into the embedding pipeline. To update for a new upstream revision, regenerate by downloading the three files and running `sha256sum` against them.

func (ModelConfig) ModelDataFilenameOrDefault added in v0.16.6 

func (c ModelConfig) ModelDataFilenameOrDefault() string

ModelDataFilenameOrDefault returns the on-disk basename for the external weights blob, mirroring ModelFilenameOrDefault.

func (ModelConfig) ModelFilenameOrDefault added in v0.16.6 

func (c ModelConfig) ModelFilenameOrDefault() string

ModelFilenameOrDefault returns the on-disk basename for the model graph, falling back to the historical "model_q4.onnx" when the config left it empty. Centralized so the back-compat default lives in one place instead of being duplicated across downloader, manager, and helpers.

type ModelDownloader 

type ModelDownloader struct {
	// contains filtered or unexported fields
}

ModelDownloader downloads ONNX models and tokenizers from HuggingFace.

func NewModelDownloader 

func NewModelDownloader() *ModelDownloader

NewModelDownloader creates a downloader that stores models in the default directory.

func NewModelDownloaderWithDir 

func NewModelDownloaderWithDir(modelDir string) *ModelDownloader

NewModelDownloaderWithDir creates a downloader with a specific model directory.

func (*ModelDownloader) Download 

func (d *ModelDownloader) Download(ctx context.Context, cfg ModelConfig, progress func(float64)) error

Download downloads the model and tokenizer files, validating checksums. If files already exist and checksums match, they are skipped. Progress is reported via the progress callback (0.0 to 1.0).

func (*ModelDownloader) GetModelPath 

func (d *ModelDownloader) GetModelPath(name string) string

GetModelPath returns the path to the model file for the given model name.

Returns the q4 path for back-compat with callers that don't know the active variant. Variant-aware code should join modelDir, name, and ModelConfig.ModelFilenameOrDefault() directly.

func (*ModelDownloader) GetTokenizerPath 

func (d *ModelDownloader) GetTokenizerPath(name string) string

GetTokenizerPath returns the path to the tokenizer file for the given model name.

func (*ModelDownloader) IsDownloaded 

func (d *ModelDownloader) IsDownloaded(name string) bool

IsDownloaded returns true if both the model and tokenizer files exist for the given name.

type ONNXEmbeddingProvider 

type ONNXEmbeddingProvider struct {
	// contains filtered or unexported fields
}

ONNXEmbeddingProvider implements EmbeddingProvider using EmbeddingGemma via ONNX Runtime.

It loads a Gemma-based embedding model, tokenizes input text with a BPE tokenizer, runs ONNX inference, and extracts mean-pooled, L2-normalized embeddings with optional MRL dimension truncation.

func NewONNXEmbeddingProvider 

func NewONNXEmbeddingProvider(ctx context.Context, runtime *ONNXRuntime, modelPath, tokenizerPath string, dims, fullDims int) (*ONNXEmbeddingProvider, error)

NewONNXEmbeddingProvider creates an embedding provider from ONNX model files.

The runtime must already be initialized. modelPath points to the .onnx model file, tokenizerPath points to tokenizer.json. dims specifies the output dimensionality (e.g., 256 for MRL truncation from 768, or 768 for full). fullDims is the model's native output dimension (used for tensor allocation).

func (*ONNXEmbeddingProvider) Close 

func (p *ONNXEmbeddingProvider) Close() error

Close releases the ONNX session and associated resources.

func (*ONNXEmbeddingProvider) Dimensions 

func (p *ONNXEmbeddingProvider) Dimensions() int

Dimensions returns the dimensionality of vectors produced by this provider.

func (*ONNXEmbeddingProvider) Embed 

func (p *ONNXEmbeddingProvider) Embed(ctx context.Context, text string) ([]float32, error)

Embed returns a L2-normalized embedding vector for the given text.

func (*ONNXEmbeddingProvider) EmbedBatch 

func (p *ONNXEmbeddingProvider) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch returns L2-normalized embeddings for multiple texts. Tokenizes all texts up-front, then runs ONNX inference in batches of defaultBatchChunkSize through a single padded [batch, seq] tensor — much faster than the previous per-text loop (verified ~6× speedup during codebase indexing on M1+, since memory bandwidth amortizes across the chunk).

func (*ONNXEmbeddingProvider) EmbedBatchWithPrefix 

func (p *ONNXEmbeddingProvider) EmbedBatchWithPrefix(ctx context.Context, texts []string, prefix string) ([][]float32, error)

EmbedBatchWithPrefix returns L2-normalized embeddings for multiple texts with the specified prefix prepended to each text before tokenization. Same batched ONNX execution as EmbedBatch — see that method's doc.

func (*ONNXEmbeddingProvider) EmbedWithPrefix 

func (p *ONNXEmbeddingProvider) EmbedWithPrefix(ctx context.Context, text, prefix string) ([]float32, error)

EmbedWithPrefix returns a L2-normalized embedding vector for the given text with the specified prefix prepended before tokenization.

func (*ONNXEmbeddingProvider) ModelHash 

func (p *ONNXEmbeddingProvider) ModelHash() string

ModelHash returns a SHA-256 hex digest of the model file.

func (*ONNXEmbeddingProvider) Name 

func (p *ONNXEmbeddingProvider) Name() string

Name returns a human-readable identifier for the provider.

type ONNXRuntime 

type ONNXRuntime struct {
	// contains filtered or unexported fields
}

ONNXRuntime provides shared ONNX Runtime infrastructure. It initializes the ONNX environment and creates dynamic inference sessions for models. Designed to be shared between embedding providers (EmbeddingGemma) and future local LLM providers (Gemma 3 2B) — they all use the same runtime environment and shared library.

func NewONNXRuntime 

func NewONNXRuntime() (*ONNXRuntime, error)

NewONNXRuntime creates a new ONNX runtime with the default model directory. The ONNX environment is initialized globally on first creation.

func NewONNXRuntimeWithDir 

func NewONNXRuntimeWithDir(modelDir string) (*ONNXRuntime, error)

NewONNXRuntimeWithDir creates a new ONNX runtime with a specific model directory. Useful for testing with isolated temp directories.

func (*ONNXRuntime) Close 

func (r *ONNXRuntime) Close() error

Close marks this runtime instance as closed. The underlying yalue ONNX environment is a PROCESS-WIDE singleton shared by every ONNXRuntime in the program, so we deliberately do NOT call DestroyEnvironment here: doing so invalidates in-flight sessions in OTHER managers and crashes their lingering init goroutines inside CGO (observed as SIGSEGV during the test suite when multiple managers were created and destroyed in quick succession).

The environment is freed at process exit; that's enough.

func (*ONNXRuntime) NewDynamicSession 

func (r *ONNXRuntime) NewDynamicSession(modelPath string, inputNames, outputNames []string, opts ...SessionOption) (*onnxruntime.DynamicAdvancedSession, error)

NewDynamicSession creates a dynamic inference session for the given ONNX model file. Dynamic sessions allow flexible input/output tensors per Run() call, which is useful when batch sizes vary.

The inputNames and outputNames can be nil to use the model's defaults. Session options (threading, GPU) can be passed via opts.

func (*ONNXRuntime) Ready 

func (r *ONNXRuntime) Ready() bool

Ready returns true if the runtime has been successfully initialized.

type QueryResult 

type QueryResult struct {
	Record     VectorRecord
	Similarity float32
}

QueryResult pairs a VectorRecord with its similarity score for ranking.

func TopK 

func TopK(query []float32, candidates []VectorRecord, k int, threshold float32) []QueryResult

TopK returns the top-K VectorRecord matches for a query vector from the given candidates, filtering out results below threshold. Results are sorted by descending similarity. If k <= 0, all results above threshold are returned.

type SessionOption 

type SessionOption struct {
	// IntraOpNumThreads sets the number of threads for intra-op parallelism.
	// 0 means use default.
	IntraOpNumThreads int
	// InterOpNumThreads sets the number of threads for inter-op parallelism.
	// 0 means use default.
	InterOpNumThreads int
}

SessionOption configures an inference session.

type VectorRecord 

type VectorRecord struct {
	// ID is a unique identifier for this record.
	ID string `json:"id"`

	// File is the file path the record comes from.
	File string `json:"file"`

	// Name is the symbol or block name (e.g., function name).
	Name string `json:"name"`

	// Signature is the function or block signature text.
	Signature string `json:"signature"`

	// StartLine is the 1-based starting line number of the record.
	StartLine int `json:"startLine"`

	// EndLine is the 1-based ending line number of the record.
	EndLine int `json:"endLine"`

	// Language is the programming language (e.g., "go", "python").
	Language string `json:"language"`

	// Embedding is the vector embedding of the record's content.
	Embedding []float32 `json:"embedding"`

	// Hash is a content hash for duplicate detection.
	Hash string `json:"hash"`

	// IndexedAt is the time when the record was indexed.
	IndexedAt time.Time `json:"indexedAt"`

	// Type is the record type: "code_unit" for extracted code symbols, or "file" for full-file embeddings.
	// Empty string for backward compatibility with legacy records (treated as "code_unit").
	Type string `json:"type"`

	// Metadata holds arbitrary key-value data for non-code-unit record types
	// (e.g., conversation turns, memories). Nil for code_unit/file records.
	Metadata map[string]interface{} `json:"metadata,omitempty"`
}

VectorRecord represents a single vectorized item with metadata. Used for duplicate detection and semantic search over code files.

type VectorStore 

type VectorStore interface {
	// Store adds records to the store. If a record with the same ID exists,
	// it is replaced.
	Store(records []VectorRecord) error

	// LoadAll returns all records currently stored.
	LoadAll() ([]VectorRecord, error)

	// Query returns the top-K records whose embeddings are most similar to
	// the query vector, with similarity >= threshold.
	Query(vec []float32, topK int, threshold float32) ([]QueryResult, error)

	// DeleteByFile removes all records whose File field matches filePath.
	DeleteByFile(filePath string) error

	// DeleteByIDs removes records with the given IDs in a single batched
	// operation. Implementations should issue at most one disk write
	// regardless of how many IDs are deleted. IDs that are not present in
	// the store are silently skipped.
	DeleteByIDs(ids []string) error

	// Size returns the total number of records in the store.
	Size() int

	// ReplaceAll replaces all records in the store with the given slice.
	// Use this when you need to perform a full replacement rather than a merge.
	ReplaceAll(records []VectorRecord) error

	// Close releases any resources held by the store.
	Close() error
}

VectorStore persists and queries vector embeddings with metadata.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.