embedding

package
v0.17.21 Latest Latest
Warning

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

 Go to latest Published: Aug 29, 2026 License: MIT Imports: 45 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 (
	// DefaultDuplicateThreshold gates code-vs-code duplicate detection.
	// Measured through CheckFileForDuplicates with documentPrefix on both sides.
	// 0.65 sits between near-duplicates (0.767) and related code (0.492).
	DefaultDuplicateThreshold = 0.65

	// DefaultRelatedCodeThreshold gates "related code" injection into read_file results.
	// Related code barely separates from unrelated (0.421 vs 0.368), so 0.55
	// keeps injected context closer to genuine near-duplicates.
	DefaultRelatedCodeThreshold = 0.55

	// DefaultSemanticSearchThreshold gates natural-language search over code.
	// Asymmetric query/doc embeddings: correct hits 0.499-0.613, wrong answers ~0.32.
	DefaultSemanticSearchThreshold = 0.40

	// DefaultConversationSearchThreshold gates search over the CONVERSATION store
	// (turns, rollups, memories). No task prefix; scores not comparable to code index.
	// 0.45 matches what semantic recall uses against the same store.
	DefaultConversationSearchThreshold = 0.45

	// DefaultCodeModelDuplicateThreshold gates duplicate detection.
	DefaultCodeModelDuplicateThreshold = 0.65

	// DefaultCodeModelRelatedThreshold gates "related code" injection.
	DefaultCodeModelRelatedThreshold = 0.35

	// DefaultCodeModelSemanticSearchThreshold gates NL-code search via Jina.
	DefaultCodeModelSemanticSearchThreshold = 0.30
)

Similarity thresholds derived from retrieval_quality_test.go and prefix_symmetry_test.go measurements against the active model.

Prior values (0.90 duplicate, 0.85 file-check, 0.75 search) were unreachable: near-duplicates top out ~0.84 symmetric, correct hits near 0.5.

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 BuildLockTimeout = 2 * time.Second

BuildLockTimeout is the max time to poll (50ms intervals) for the cross-process build lock after the initial non-blocking attempt fails.

View Source 
const BuildTimeout = 45 * time.Minute

BuildTimeout is the max duration for the full index build lifecycle. Sized for this repo: ~12k units at ~5.3 units/s ≈ 40 min cold build.

View Source 
const DefaultRemoteSocketTimeout = 60 * time.Second

DefaultRemoteSocketTimeout bounds a single remote operation. Inference can be slow on a loaded daemon, but a stuck daemon must not hang the CLI.

View Source 
const EmbedBatchSize = 32

EmbedBatchSize is the number of code units per ONNX inference call.

View Source 
const MaxDepth = 15

MaxDepth limits WalkCodeFiles directory nesting (avoids pathological trees).

View Source 
const MaxFileCount = 10000

MaxFileCount caps files WalkCodeFiles will collect.

View Source 
const MaxIndexableFileBytes int64 = 1 << 20 // 1 MiB

MaxIndexableFileBytes caps the size of ANY file the index will read — code files (via ExtractFromFile) and file-level non-code files alike. Files larger than this are skipped entirely: index bodies truncate to 8 KB anyway, so reading a multi-GB file (generated code bundles, data corpora) would only waste memory and walk budget. The read paths also use this as their LimitReader bound, so even a race between stat and open can't blow memory.

View Source 
const ProgressInterval = 500

ProgressInterval is the file count between progress events during walk and embedding.

View Source 
const WalkTimeout = 30 * time.Second

WalkTimeout is the maximum time for WalkCodeFiles to enumerate files.

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)

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 DefaultIndexDir added in v0.17.17 

func DefaultIndexDir(workspaceRoot string) string

func DefaultModelDir 

func DefaultModelDir() string

DefaultModelDir returns the directory holding the embedding model weights and the ONNX Runtime shared library.

This resolver is deliberately NOT per-workspace and deliberately has no build-tag variants. Both properties were violated, and each violation cost real disk:

  1. It used to resolve under the *config* root. `--isolated-config` (which SP-116 turns on automatically for any git repo) points SPROUT_CONFIG at <workspace>/.sprout, so every repo the user opened downloaded its own copy of the 196MB weights blob plus the 35MB runtime dylib — ~222MB per workspace, found duplicated across six of them. SP-133 moved the cgo build to the data root but left the wasm and non-cgo builds resolving off the config root and off ~/.cache respectively, so the bug stayed live and the three builds disagreed about where the model even was.

  2. Three copies of this function existed behind //go:build cgo, !cgo, and wasm, with three different env var names (SPROUT_MODELS_DIR vs SPROUT_MODEL_DIR) and three different fallbacks. Divergence was the root cause; one definition is the fix. Keep it that way — if a platform needs different behavior, branch inside this function, not with a build tag.

The contents are large, immutable, and content-addressed (the downloader verifies sha256), so one shared copy per user is always correct.

Resolution: $SPROUT_MODELS_DIR → $SPROUT_MODEL_DIR (legacy alias) → $SPROUT_DATA_DIR/models/embedding → $XDG_DATA_HOME/sprout/models/embedding → $HOME/.local/share/sprout/models/embedding.

Nested under models/embedding (not directly under models/) as of the sibling pkg/localmodel/DefaultModelsDir landing: LLM chat models moved in as models/llm alongside this, one shared "models" root with a subdirectory per kind, rather than two unrelated top-level directories under DataDir with no indication either exists. Falls back to the flat models/ directory (this function's own pre-split default) when that already has content and models/embedding doesn't, so installations that downloaded the embedding model before this split keep working without re-downloading.

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 ReleaseManager added in v0.17.17 

func ReleaseManager(m *EmbeddingManager)

ReleaseManager drops one reference taken by AcquireManager, closing the manager once the last holder releases it. Managers not obtained from AcquireManager are closed directly, so callers can release unconditionally.

Closing on the last release rather than never is deliberate: unlike the ONNX weights, a manager pins a full copy of the workspace's vectors, and a daemon that opens many workspaces over a long session should not accumulate them.

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 SetProviderFactory added in v0.17.17 

func SetProviderFactory(f func(ctx context.Context) (EmbeddingProvider, error))

SetProviderFactory installs a process-wide factory consulted by EmbeddingManager.initLocked BEFORE falling back to in-process ONNX.

SP-136 P3: a CLI that has a live daemon socket sets this to construct RemoteEmbeddingProvider instances, so the daemon owns the model copy and the CLI never loads its own 155MB model. If the factory returns an error (socket unavailable), initLocked silently falls back to in-process ONNX — sprout always works.

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
}

type ByteLevelTokenizer added in v0.17.17 

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

ByteLevelTokenizer is a GPT-2/Jina/RoBERTa BPE tokenizer with byte-level pre-tokenization.

func NewByteLevelTokenizer added in v0.17.17 

func NewByteLevelTokenizer(path string) (*ByteLevelTokenizer, error)

NewByteLevelTokenizer parses a HuggingFace tokenizer.json file.

func (*ByteLevelTokenizer) BOSID added in v0.17.17 

func (t *ByteLevelTokenizer) BOSID() int32

BOSID returns the beginning-of-sequence token ID, or -1 if unset.

func (*ByteLevelTokenizer) EOSID added in v0.17.17 

func (t *ByteLevelTokenizer) EOSID() int32

EOSID returns the end-of-sequence token ID, or -1 if unset.

func (*ByteLevelTokenizer) Encode added in v0.17.17 

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

Encode converts text into a sequence of BPE token IDs using byte-level pre-tokenization. No special tokens are added.

func (*ByteLevelTokenizer) EncodeWithBOSAndEOS added in v0.17.17 

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

EncodeWithBOSAndEOS wraps Encode with the BOS/EOS special tokens.

func (*ByteLevelTokenizer) VocabSize added in v0.17.17 

func (t *ByteLevelTokenizer) VocabSize() int

VocabSize returns the number of entries in the vocabulary.

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 extracts code units from content, embeds each, and queries the index for similar records. Filters self-matches via workspace-relative path comparison.

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 or oversized files.

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
}

func AcquireManager added in v0.17.17 

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

AcquireManager returns the process-wide manager for the index that (cfg, workspaceRoot) resolves to, creating it on first use. Every successful call must be paired with exactly one ReleaseManager.

The key includes both the index directory and the workspace root. The index directory alone is not enough: two workspaces that explicitly configure the same embedding_index.index_dir must not silently share a build whose stale-record sweep is scoped to one of them. The default index directory is already workspace-derived, so ordinary callers key on it uniquely.

func NewEmbeddingManager 

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

func (*EmbeddingManager) AutoBuildWhenReady 

func (m *EmbeddingManager) AutoBuildWhenReady()

AutoBuildWhenReady runs a background index build after a short startup delay.

func (*EmbeddingManager) BuildIndex 

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

func (*EmbeddingManager) BuildIndexBackground 

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

func (*EmbeddingManager) CheckDuplicates 

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

func (*EmbeddingManager) Close 

func (m *EmbeddingManager) Close() error

func (*EmbeddingManager) CloseNotify added in v0.16.18 

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

func (*EmbeddingManager) Dimensions added in v0.17.17 

func (m *EmbeddingManager) Dimensions() int

Dimensions returns the embedding dimensionality (0 before Init).

func (*EmbeddingManager) Embed added in v0.17.17 

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

Embed embeds a single text with the manager's provider, initializing it first. Used by the daemon socket service (SP-136 P3).

func (*EmbeddingManager) EmbedBatch added in v0.17.17 

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

EmbedBatch embeds multiple texts with the manager's provider, initializing it first. The returned slice matches the input order.

func (*EmbeddingManager) GetConversationStore 

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

func (*EmbeddingManager) IndexSize 

func (m *EmbeddingManager) IndexSize() int

func (*EmbeddingManager) Init 

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

Init initializes the ONNX embedding provider and opens the vector store. Idempotent.

func (*EmbeddingManager) InitError 

func (m *EmbeddingManager) InitError() error

func (*EmbeddingManager) IsBuilding 

func (m *EmbeddingManager) IsBuilding() bool

func (*EmbeddingManager) IsInitialized 

func (m *EmbeddingManager) IsInitialized() bool

func (*EmbeddingManager) ModelHash 

func (m *EmbeddingManager) ModelHash() string

func (*EmbeddingManager) Name added in v0.17.17 

func (m *EmbeddingManager) Name() string

Name returns the provider's human-readable name ("" before Init).

func (*EmbeddingManager) QuerySimilar 

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

func (*EmbeddingManager) QuerySimilarCode added in v0.17.17 

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

QuerySimilarCode searches using source code as the input rather than NL.

func (*EmbeddingManager) Readiness added in v0.17.17 

func (m *EmbeddingManager) Readiness() IndexReadiness

func (*EmbeddingManager) RelatedCodeThreshold added in v0.17.17 

func (m *EmbeddingManager) RelatedCodeThreshold() float32

func (*EmbeddingManager) SemanticSearchThreshold added in v0.17.17 

func (m *EmbeddingManager) SemanticSearchThreshold() float32

func (*EmbeddingManager) SetForTesting 

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

func (*EmbeddingManager) UpdateFile 

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

func (*EmbeddingManager) UpdateFromGitDiff 

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

func (*EmbeddingManager) UpdateFromGitDiffBackground added in v0.17.10 

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

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.

Extract is memory-bounded: it never materializes the full input as a string or rune slice. Newline counting and truncation operate on the raw []byte, so callers may pass arbitrarily large content without spiking memory.

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 with mtime-based manifest to skip unchanged files.

func (*IndexManager) CheckDuplicates 

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

CheckDuplicates finds indexed code similar to codeText.

It embeds with documentPrefix, NOT the query prefix: this compares code against code, and the index stores code as documents. Routing this through QuerySimilar (which is for natural-language questions) put the two sides in different subspaces and cost roughly 0.10 of similarity — enough that, on top of an already unreachable 0.90 gate, duplicate detection could not fire at all. See TestPrefixSymmetryAffectsDuplicateThresholds.

func (*IndexManager) QuerySimilar 

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

QuerySimilar embeds a natural-language query and returns the top-K most similar records above threshold. Use CheckDuplicates when the input is code rather than a question.

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 indexes non-code files (markdown, configs, etc.) at file level.
	IndexFileLevel bool
	// ManifestPath is the path to the build manifest file for incremental rebuilds.
	ManifestPath string
	// IndexDir is the directory containing the HNSW index files. Used to place
	// the cross-process .build.lock file. Empty disables locking.
	IndexDir string
}

IndexOptions configures IndexManager behavior.

type IndexReadiness added in v0.17.17 

type IndexReadiness struct {
	Initialized bool
	Building    bool
	Records     int
}

func (IndexReadiness) CanAnswerQueries added in v0.17.17 

func (r IndexReadiness) CanAnswerQueries() bool

type IndexStats 

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

IndexStats reports the results of an indexing operation.

type JinaONNXEmbeddingProvider added in v0.17.17 

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

JinaONNXEmbeddingProvider implements EmbeddingProvider using Jina Code v2 via ONNX Runtime. Outputs last_hidden_state [batch, seq_len, 768] which is mean-pooled and L2-normalized in Go. Symmetric model: query and document embeddings use identical encoding (no task prefixes).

func NewJinaONNXEmbeddingProvider added in v0.17.17 

func NewJinaONNXEmbeddingProvider(ctx context.Context, runtime *ONNXRuntime, modelPath, tokenizerPath string) (*JinaONNXEmbeddingProvider, error)

NewJinaONNXEmbeddingProvider creates a Jina Code v2 embedding provider.

func (*JinaONNXEmbeddingProvider) Close added in v0.17.17 

func (p *JinaONNXEmbeddingProvider) Close() error

Close releases the ONNX session.

func (*JinaONNXEmbeddingProvider) Dimensions added in v0.17.17 

func (p *JinaONNXEmbeddingProvider) Dimensions() int

Dimensions returns 768.

func (*JinaONNXEmbeddingProvider) Embed added in v0.17.17 

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

Embed returns an L2-normalized, mean-pooled 768-dim embedding for text.

func (*JinaONNXEmbeddingProvider) EmbedBatch added in v0.17.17 

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

EmbedBatch returns L2-normalized embeddings for multiple texts. Sequences are right-padded to the longest in the batch; attention mask ensures padded positions don't affect mean pooling.

func (*JinaONNXEmbeddingProvider) EmbedBatchWithPrefix added in v0.17.17 

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

EmbedBatchWithPrefix ignores the prefix (Jina is symmetric) and delegates to EmbedBatch.

func (*JinaONNXEmbeddingProvider) EmbedWithPrefix added in v0.17.17 

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

EmbedWithPrefix ignores the prefix (Jina is symmetric) and delegates to Embed.

func (*JinaONNXEmbeddingProvider) ModelHash added in v0.17.17 

func (p *JinaONNXEmbeddingProvider) ModelHash() string

ModelHash returns the SHA-256 of the model file.

func (*JinaONNXEmbeddingProvider) Name added in v0.17.17 

func (p *JinaONNXEmbeddingProvider) Name() string

Name returns the model identifier.

type MLXEmbeddingProvider added in v0.17.17 

type MLXEmbeddingProvider struct{}

MLXEmbeddingProvider is a non-functional stub on platforms without Apple Silicon GPU. The constructor returns an error so callers fall back to the ONNX provider.

func NewMLXEmbeddingProvider added in v0.17.17 

func NewMLXEmbeddingProvider(ctx context.Context, modelPath, tokenizerPath string) (*MLXEmbeddingProvider, error)

func (*MLXEmbeddingProvider) Close added in v0.17.17 

func (p *MLXEmbeddingProvider) Close() error

func (*MLXEmbeddingProvider) Dimensions added in v0.17.17 

func (p *MLXEmbeddingProvider) Dimensions() int

func (*MLXEmbeddingProvider) Embed added in v0.17.17 

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

func (*MLXEmbeddingProvider) EmbedBatch added in v0.17.17 

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

func (*MLXEmbeddingProvider) EmbedBatchWithPrefix added in v0.17.17 

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

func (*MLXEmbeddingProvider) EmbedWithPrefix added in v0.17.17 

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

func (*MLXEmbeddingProvider) ModelHash added in v0.17.17 

func (p *MLXEmbeddingProvider) ModelHash() string

func (*MLXEmbeddingProvider) Name added in v0.17.17 

func (p *MLXEmbeddingProvider) Name() string

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 JinaCodeV2Config added in v0.17.17 

func JinaCodeV2Config() ModelConfig

JinaCodeV2Config returns the config for Jina AI's jina-embeddings-v2-base-code — a 137M parameter BERT encoder purpose-built for code retrieval (trained on The Stack v2 code corpus, Apache-2.0). Used as the code-specific provider in the dual-model architecture (SP-135): code retrieval + duplicate detection routes here; Gemma handles conversation/NL semantics.

Ships the quantized (int8) ONNX export (~162 MB). Jina's ONNX graph uses standard ops (MatMulInteger, LayerNormalization, GELU — no com.microsoft custom ops), making it CoreML-EP-friendly (unlike Gemma's export — see SP-134). The model outputs last_hidden_state [batch, seq, 768]; mean pooling is done in Go (JinaProvider.runInference), not by the graph.

Source: https://huggingface.co/jinaai/jina-embeddings-v2-base-code

func JinaCodeV2SafetensorsConfig added in v0.17.17 

func JinaCodeV2SafetensorsConfig() ModelConfig

JinaCodeV2SafetensorsConfig returns the config for the fp16 safetensors weights of Jina Code v2, used by the MLX Metal provider (SP-134). The safetensors format is loaded directly by the Go code (no ONNX), keeping the model in fp16 for GPU inference.

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 ModelInfo added in v0.17.17 

type ModelInfo struct {
	Name         string `json:"name"`
	Quantization string `json:"quantization"`
	Dims         int    `json:"dims"`      // emitted vector width (after MRL truncation)
	FullDims     int    `json:"full_dims"` // model's native output width
	Truncated    bool   `json:"truncated"` // Dims < FullDims (Matryoshka truncation active)
}

ModelInfo describes the embedding model this build actually loads. It is derived from the same ModelConfig the provider is constructed from, so a UI rendering it can never drift from what is running.

The webui settings panel previously hardcoded these strings and reported a model that had not been in use for some time (wrong name, wrong quantization, and 256 dims against an actual 768) — which made the panel actively misleading when reasoning about index size and memory.

func ActiveModelInfo added in v0.17.17 

func ActiveModelInfo() ModelInfo

ActiveModelInfo returns the ModelInfo for the model this build uses.

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 RemoteClient added in v0.17.17 

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

RemoteClient provides the non-provider operations of the daemon socket protocol: Query, BuildIndex, and CheckDuplicates. It shares the provider's connection so a CLI keeps a single socket.

func NewRemoteClient added in v0.17.17 

func NewRemoteClient(p *RemoteEmbeddingProvider) *RemoteClient

NewRemoteClient creates a client over the given provider's connection.

func (*RemoteClient) BuildIndex added in v0.17.17 

func (c *RemoteClient) BuildIndex(ctx context.Context, workspaceRoot string) (*IndexStats, error)

BuildIndex asks the daemon to (re)build the workspace index and returns the build stats.

func (*RemoteClient) CheckDuplicates added in v0.17.17 

func (c *RemoteClient) CheckDuplicates(ctx context.Context, workspaceRoot, filePath, content string) (*CheckDuplicatesResult, error)

CheckDuplicates asks the daemon to check file content against its index.

func (*RemoteClient) QuerySimilar added in v0.17.17 

func (c *RemoteClient) QuerySimilar(ctx context.Context, workspaceRoot, text string, topK int, threshold float32) ([]QueryResult, error)

QuerySimilar runs a semantic query on the daemon-owned index for the workspace root.

type RemoteEmbeddingProvider added in v0.17.17 

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

RemoteEmbeddingProvider implements EmbeddingProvider by proxying embedding operations to the sprout daemon over a Unix socket. The daemon owns the model copy and inference gate; this provider is a thin client.

If the socket is unavailable at construction, NewRemoteEmbeddingProvider returns an error and callers fall back to in-process ONNX. If the socket dies later, operations fail with a descriptive error; the provider re-dials on the next operation (transient-failure reconnect).

func NewRemoteEmbeddingProvider added in v0.17.17 

func NewRemoteEmbeddingProvider(socketPath string) (*RemoteEmbeddingProvider, error)

NewRemoteEmbeddingProvider creates a provider backed by the daemon socket at socketPath. It performs a meta handshake immediately so configuration errors surface at construction time.

func (*RemoteEmbeddingProvider) Close added in v0.17.17 

func (p *RemoteEmbeddingProvider) Close() error

Close implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) Dimensions added in v0.17.17 

func (p *RemoteEmbeddingProvider) Dimensions() int

Dimensions implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) Embed added in v0.17.17 

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

Embed implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) EmbedBatch added in v0.17.17 

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

EmbedBatch implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) EmbedBatchWithPrefix added in v0.17.17 

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

EmbedBatchWithPrefix implements EmbeddingProvider by prepending the prefix to each text client-side and delegating to EmbedBatch.

func (*RemoteEmbeddingProvider) EmbedWithPrefix added in v0.17.17 

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

EmbedWithPrefix implements EmbeddingProvider by prepending the prefix client-side and delegating to Embed (the daemon tokenizes the final text).

func (*RemoteEmbeddingProvider) ModelHash added in v0.17.17 

func (p *RemoteEmbeddingProvider) ModelHash() string

ModelHash implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) Name added in v0.17.17 

func (p *RemoteEmbeddingProvider) Name() string

Name implements EmbeddingProvider.

func (*RemoteEmbeddingProvider) SocketPath added in v0.17.17 

func (p *RemoteEmbeddingProvider) SocketPath() string

SocketPath returns the daemon socket path this provider talks to.

type RemoteOp added in v0.17.17 

type RemoteOp string

RemoteOp identifies a protocol operation. 

const (
	RemoteOpMeta            RemoteOp = "meta"
	RemoteOpEmbed           RemoteOp = "embed"
	RemoteOpEmbedBatch      RemoteOp = "embed_batch"
	RemoteOpQuery           RemoteOp = "query"
	RemoteOpBuildIndex      RemoteOp = "build_index"
	RemoteOpCheckDuplicates RemoteOp = "check_duplicates"
)

type RemoteRequest added in v0.17.17 

type RemoteRequest struct {
	ID        string   `json:"id"`
	Op        RemoteOp `json:"op"`
	Text      string   `json:"text,omitempty"`
	Texts     []string `json:"texts,omitempty"`
	K         int      `json:"k,omitempty"`
	Threshold float32  `json:"threshold,omitempty"`
	Workspace string   `json:"workspace_root,omitempty"`
	FilePath  string   `json:"file_path,omitempty"`
	Content   string   `json:"content,omitempty"`
	TopK      int      `json:"top_k,omitempty"`
}

RemoteRequest is a single protocol request.

type RemoteResponse added in v0.17.17 

type RemoteResponse struct {
	ID         string                 `json:"id"`
	Error      string                 `json:"error,omitempty"`
	Name       string                 `json:"name,omitempty"`
	Dimensions int                    `json:"dimensions,omitempty"`
	ModelHash  string                 `json:"model_hash,omitempty"`
	Vector     []float32              `json:"vector,omitempty"`
	Vectors    [][]float32            `json:"vectors,omitempty"`
	Results    []QueryResult          `json:"results,omitempty"`
	Stats      *IndexStats            `json:"stats,omitempty"`
	Duplicates *CheckDuplicatesResult `json:"duplicates,omitempty"`
}

RemoteResponse is a single protocol response.

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
	// CPUMemArena and MemPattern override the ORT allocator defaults. nil
	// leaves the package default in place; see newSessionOptions.
	CPUMemArena *bool
	MemPattern  *bool
}

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.