embeddings

package
v1.20.0 Latest Latest
Warning

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

 Go to latest Published: May 30, 2026 License: Apache-2.0 Imports: 30 Imported by: 0

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func AutoIndex 

func AutoIndex(pipeline *Pipeline, dirs []string, keyFn RepoKeyFunc, opts AutoIndexOpts)

AutoIndex scans directories for Git repositories and indexes them with a bounded worker pool plus per-repo retry-with-backoff on transient errors.

Each immediate subdirectory containing a .git folder is treated as a repo. keyFn should be codegraph.GraphNameFor (passed from caller to avoid an import cycle).

Rollback to byte-identical legacy behavior: opts.Concurrency=1, RetryMax=0.

func ExtractQueryKeywords 

func ExtractQueryKeywords(query string) []string

extractQueryKeywords splits a natural language query into meaningful search terms by removing stopwords and short tokens. Returns lowercase terms >= 3 chars.

func NewPipelineCache 

func NewPipelineCache() *kitcache.Cache

NewPipelineCache constructs a *kitcache.Cache pre-tuned for the embed pipeline (15-minute L1 TTL, modest item ceiling). Lives here rather than in register.go so wiring code stays free of cache-policy details.

Prometheus metrics are published via go-kit cache.WithMetrics using prometheus.DefaultRegisterer and the "embed_pipeline" cache name.

func PublishRRFWeights 

func PublishRRFWeights(w RRFWeights)

PublishRRFWeights records the deployed RRF weights to Prometheus. Idempotent (set, not increment); safe to call multiple times if config reloads land.

Types

type AutoIndexOpts 

type AutoIndexOpts struct {
	// Concurrency is the maximum number of repos indexed in parallel.
	// Must be >=1; values <1 are normalized to 1 (serial fallback).
	Concurrency int
	// RetryMax is the maximum number of retries per repo on transient
	// failures. RetryMax=0 disables retry (single attempt only).
	RetryMax int
	// RetryBase is the initial backoff duration; doubles on each retry.
	RetryBase time.Duration
}

AutoIndexOpts controls the bounded worker pool and retry policy used by AutoIndex. A zero/invalid value falls back to DefaultAutoIndexOpts.

func DefaultAutoIndexOpts 

func DefaultAutoIndexOpts() AutoIndexOpts

DefaultAutoIndexOpts returns sane defaults: concurrency=2, retry_max=3, retry_base=5s (exponential backoff: 5s, 10s, 20s).

type EmbeddingRecord 

type EmbeddingRecord struct {
	RepoKey    string
	FilePath   string
	SymbolName string
	SymbolKind string // function, method, class, etc.
	Language   string
	StartLine  int
	BodyHash   uint64 // for change detection
	Embedding  []float32
}

EmbeddingRecord holds a single symbol embedding for storage.

type Expander 

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

Expander enriches semantic search results with 1-hop CALLS neighbors from Apache AGE.

func NewExpander 

func NewExpander(pool *pgxpool.Pool) *Expander

NewExpander creates an Expander backed by the given connection pool.

func (*Expander) Expand 

func (e *Expander) Expand(ctx context.Context, graphName string, results []SearchResult, maxExtra int) []SearchResult

Expand queries the AGE graph for 1-hop CALLS neighbors of the symbols in results. Returns at most maxExtra additional SearchResult entries (source="graph") not already present in the input. If the graph does not exist or any query fails, returns nil gracefully.

type FileIndexResult added in v1.20.0 

type FileIndexResult struct {
	Embedded int   // newly-embedded symbol count (new or body-changed)
	Skipped  int   // hash-matched symbols — no embed call issued
	Deleted  int64 // symbols removed from DB (file shrank or was deleted)
}

FileIndexResult summarizes the outcome of a single-file incremental index.

type FileLineHit 

type FileLineHit struct {
	FilePath string
	Line     int
}

FileLineHit is a file path + line number from keyword search.

type HybridResult 

type HybridResult struct {
	SearchResult
	Source   string  // "semantic", "keyword", or "hybrid"
	RRFScore float64 // combined reciprocal rank fusion score
}

HybridResult extends SearchResult with source attribution.

func MergeRRF 

func MergeRRF(semantic []SearchResult, keyword []KeywordHit, topK int, weights RRFWeights) []HybridResult

MergeRRF combines semantic and keyword results using Reciprocal Rank Fusion. Backed by go-kit/rerank.WeightedRRF (Cormack-Clarke 2009 with per-list weights, k=60). Key = FilePath + ":" + SymbolName for dedup; results in both lists get boosted ("hybrid" source). Returns at most topK results.

weights are env-driven (RRF_WEIGHT_SEMANTIC / RRF_WEIGHT_KEYWORD) and surfaced via Prometheus gauge gocode_rrf_weights{retriever}. Pass DefaultRRFWeights() for byte-identical legacy behavior.

type IncrementalSyncMode added in v1.20.0 

type IncrementalSyncMode string

IncrementalSyncMode is the typed label for which code path IncrementalSync took. Values are stable wire strings used in Prometheus labels — do not rename. 

const (
	// IncrementalSyncIncremental is the normal git-diff path.
	IncrementalSyncIncremental IncrementalSyncMode = "incremental"
	// IncrementalSyncSkipSHAMatch is the fast-path when SHA is unchanged.
	IncrementalSyncSkipSHAMatch IncrementalSyncMode = "skip-sha-match"
	// IncrementalSyncFullFallbackBootstrap is triggered when no prior SHA exists in DB.
	IncrementalSyncFullFallbackBootstrap IncrementalSyncMode = "full-fallback-bootstrap"
	// IncrementalSyncFullFallbackNoGit is triggered when the path is not a git repo.
	IncrementalSyncFullFallbackNoGit IncrementalSyncMode = "full-fallback-no-git"
	// IncrementalSyncFullFallbackDiffError is triggered when git diff exec fails.
	IncrementalSyncFullFallbackDiffError IncrementalSyncMode = "full-fallback-diff-error"
)

type IncrementalSyncResult added in v1.20.0 

type IncrementalSyncResult struct {
	// Mode describes which code path was taken.
	Mode IncrementalSyncMode

	PrevSHA    string
	CurrentSHA string

	// FilesChanged is the count of files returned by git diff (0 on full-fallback paths).
	FilesChanged int
	// FilesEmbedded is the sum of FileIndexResult.Embedded across processed files
	// (or IndexResult.Indexed on full-fallback paths).
	FilesEmbedded int
	// FilesSkipped is the sum of FileIndexResult.Skipped (hash-matched symbols).
	FilesSkipped int
	// FilesDeleted is the sum of FileIndexResult.Deleted (tombstoned symbols).
	FilesDeleted int64

	// Errors collects per-file errors. When non-empty, SHA is NOT advanced so the
	// next boot retries the affected files (hash-skip from IndexFile makes retries cheap).
	Errors []error
}

IncrementalSyncResult is the return value of Pipeline.IncrementalSync.

type IndexResult 

type IndexResult struct {
	Indexed int
	Skipped int
	Total   int
}

IndexResult summarizes the outcome of an indexing run.

type KeywordHit 

type KeywordHit struct {
	FilePath   string
	SymbolName string
	Line       int
}

KeywordHit represents a grep/code_search match mapped to a symbol.

type Pipeline 

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

Pipeline orchestrates embedding indexing for repository symbols.

func NewPipeline 

func NewPipeline(client *embed.Client, store *Store, opts ...PipelineOpt) *Pipeline

NewPipeline creates a Pipeline backed by the given client and store.

Pass a non-nil fileCache via WithFileCache to enable per-file symbol-entry caching keyed on (repoKey, file.RelPath) and validated by file modTime+size. When fileCache is nil, behavior is byte-identical to the v0.32.0 baseline.

func (*Pipeline) IncrementalSync added in v1.20.0 

func (p *Pipeline) IncrementalSync(ctx context.Context, repoKey, root string) (*IncrementalSyncResult, error)

IncrementalSync reconciles a repo's embedding index via git-diff against the last-indexed SHA, calling Pipeline.IndexFile per changed file. Bumps SetRepoState only on full success (Errors empty).

Falls back to full IndexRepo when:

  • currentSHA == "" (non-git path)
  • prevSHA == "" (never indexed — bootstrap)
  • git diff exec failure (corrupted state — safer to full re-scan)

Same-SHA path (prevSHA == currentSHA) bumps indexed_at timestamp via SetRepoState but does no file work. Returns mode="skip-sha-match".

Returns non-nil result describing the path taken. Returns error only on catastrophic failures (DB connection lost). Per-file errors are collected in Result.Errors, not returned — caller decides retry policy.

func (*Pipeline) IndexFile added in v1.20.0 

func (p *Pipeline) IndexFile(ctx context.Context, repoKey, root, relPath string) (*FileIndexResult, error)

IndexFile incrementally indexes one file: parses it, extracts symbols, computes the diff against the currently-indexed set (Store.GetSymbolsForFile), embeds only NEW or CHANGED symbols (body_hash differs), and DELETEs symbols that disappeared from this file in the new parse.

root is the repo's absolute path on the host (for symbol path resolution); relPath is the file path relative to root (used as code_embeddings.file_path). repoKey is the canonical repo identifier (see codegraph.GraphNameFor).

Returns counts of newly-embedded, unchanged (skipped), and deleted symbols.

Side effects: Postgres writes; HTTP call to embed-server ONLY for new/changed symbols (zero HTTP calls if file unchanged from last index).

IndexFile does NOT touch repo_state / repoMainBranchSHA — that is repo-level fingerprinting owned by IndexRepo. IndexFile is file-level only.

func (*Pipeline) IndexProgress 

func (p *Pipeline) IndexProgress(repoKey string) (done, total int, running bool)

IndexProgress returns (done, total, running) for the given repo.

func (*Pipeline) IndexRepo 

func (p *Pipeline) IndexRepo(ctx context.Context, repoKey, root string) (*IndexResult, error)

IndexRepo indexes all functions and methods in a repository for semantic search.

func (*Pipeline) IndexRepoAsync 

func (p *Pipeline) IndexRepoAsync(repoKey, root string) bool

IndexRepoAsync starts background indexing if not already running. Returns true if indexing was started, false if already in progress.

func (*Pipeline) IsIndexing 

func (p *Pipeline) IsIndexing(repoKey string) bool

IsIndexing returns true if background indexing is running for the given repo.

type PipelineOpt 

type PipelineOpt func(*Pipeline)

PipelineOpt configures a Pipeline at construction time.

func WithFileCache 

func WithFileCache(c *kitcache.Cache) PipelineOpt

WithFileCache wires a *kitcache.Cache to memoize per-file symbol entries. Validator (modTime+size) ensures stale entries are evicted on the next indexRepo pass after a file is touched. Pass nil to disable explicitly.

type RRFWeights 

type RRFWeights struct {
	Semantic float64
	Keyword  float64
}

RRFWeights are per-retriever weights for the WeightedRRF fusion in MergeRRF. Both fields must be ≥ 0 (negative values panic in go-kit/rerank.WeightedRRF — programmer error, not runtime). Weights == (1.0, 1.0) is mathematically identical to plain RRF (Cormack-Clarke 2009); use that for byte-identical rollback.

func DefaultRRFWeights 

func DefaultRRFWeights() RRFWeights

DefaultRRFWeights returns the (1.0, 1.0) baseline that reproduces the pre-Stream-1 unweighted rerank.RRF behavior. Tests and callers that don't thread per-deployment config should use this.

type RepoKeyFunc 

type RepoKeyFunc func(root string) string

RepoKeyFunc generates a graph key from a repo root path.

type SearchOpts 

type SearchOpts struct {
	RepoKey     string  // optional filter
	Language    string  // optional filter
	TopK        int     // default 20, max 100
	MaxDistance float32 // 0 = no filter; cosine distance threshold (0.0-1.0)
}

SearchOpts controls search filtering and result count.

type SearchResult 

type SearchResult struct {
	RepoKey    string
	FilePath   string
	SymbolName string
	SymbolKind string
	Language   string
	StartLine  int
	Distance   float32 // cosine distance (lower = more similar)
	Source     string  // "semantic", "keyword", "hybrid", "graph" — set by caller
	PageRank   float32 // structural importance from graph analysis (0 if not available)
}

SearchResult is a single semantic search hit.

type SimilarPair 

type SimilarPair struct {
	SymbolA    string
	FileA      string
	LineA      int
	SymbolB    string
	FileB      string
	LineB      int
	Similarity float32
}

SimilarPair represents two semantically similar symbols within the same repo.

type SimilarPairOpts 

type SimilarPairOpts struct {
	RepoKey   string
	Threshold float32 // minimum cosine similarity (default 0.92)
	Limit     int     // max pairs returned (default 50)
}

SimilarPairOpts controls the similarity search parameters.

type Store 

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

Store manages vector embeddings in PostgreSQL with pgvector.

func NewStore 

func NewStore(pool *pgxpool.Pool) *Store

NewStore creates a Store backed by the given connection pool.

func (*Store) DeleteRepo 

func (s *Store) DeleteRepo(ctx context.Context, repoKey string) error

DeleteRepo removes all embeddings for a given repo_key.

func (*Store) DeleteSymbolsForFile added in v1.20.0 

func (s *Store) DeleteSymbolsForFile(ctx context.Context, repoKey, filePath string, keepSymbolNames []string) (int64, error)

DeleteSymbolsForFile removes index rows for symbols in (repoKey, filePath) EXCEPT those whose names appear in keepSymbolNames. Used to reconcile a file after re-parsing: symbols missing from the new parse (deleted / renamed) must be evicted. Returns rows-affected count.

If keepSymbolNames is empty, all symbols for the file are deleted (use this for IndexFile when the file is removed from the source tree).

func (*Store) EnsureSchema 

func (s *Store) EnsureSchema(ctx context.Context) error

EnsureSchema creates the pgvector extension and embeddings table if needed. After creating tables it attempts to transfer ownership to CURRENT_USER (best-effort: warns instead of failing if the role is not the table owner).

func (*Store) FindSimilarPairs 

func (s *Store) FindSimilarPairs(ctx context.Context, opts SimilarPairOpts) ([]SimilarPair, error)

FindSimilarPairs finds semantically similar function pairs within a repo using pgvector cosine distance self-join. Results are ordered by similarity descending.

The self-join is wrapped in a short transaction with SET LOCAL statement_timeout to prevent it from pinning a CPU core for minutes on large repos. SQLSTATE 57014 (query_canceled) is treated as "no pairs found" and logged at Debug level.

func (*Store) GetHashes 

func (s *Store) GetHashes(ctx context.Context, repoKey string) (map[string]uint64, error)

GetHashes returns a map of symbol_name -> body_hash for all embeddings in a repo.

func (*Store) GetRepoState 

func (s *Store) GetRepoState(ctx context.Context, repoKey string) (string, error)

GetRepoState returns the last-indexed sha for a repo, or "" when no row exists.

func (*Store) GetSymbolsForFile added in v1.20.0 

func (s *Store) GetSymbolsForFile(ctx context.Context, repoKey, filePath string) ([]SymbolIdentity, error)

GetSymbolsForFile returns the symbol identity rows currently indexed for a specific file in a repo. Used by incremental update paths (Pipeline.IndexFile) to compute the diff: symbols present in DB but no longer in the file source must be DELETEd; symbols whose BodyHash changed must be re-embedded.

Returns rows sorted by symbol_name for determinism.

func (*Store) MatchKeywordHits 

func (s *Store) MatchKeywordHits(ctx context.Context, repoKey string, hits []FileLineHit) ([]KeywordHit, error)

MatchKeywordHits maps keyword search hits to the nearest indexed symbol. For each hit, finds the symbol in the same file with start_line <= hit.Line (i.e., the function containing that line). Returns deduplicated KeywordHit results.

func (*Store) Search 

func (s *Store) Search(ctx context.Context, query []float32, opts SearchOpts) ([]SearchResult, error)

Search finds the top-K most similar symbols to the query embedding using cosine distance.

func (*Store) SearchBySymbolName 

func (s *Store) SearchBySymbolName(
	ctx context.Context,
	repoKey string,
	keywords []string,
	language string,
	limit int,
) ([]SearchResult, error)

SearchBySymbolName searches the embedding index by symbol name and file path using pg_trgm trigram similarity — correctly handles code abbreviations like "init" matching "initializes", "llm" matching "llms" etc. Returns results sorted by combined trigram similarity score DESC.

Rune token secondary symbols are stored as "$state:L<line>" (e.g. "$state:L7"). The trigram similarity between "$state" and "$state:L7" is ≈0.57 (well above the 0.1 threshold), so rune token queries work without special-casing.

func (*Store) SetRepoState 

func (s *Store) SetRepoState(ctx context.Context, repoKey, sha string) error

SetRepoState records the sha that was successfully indexed.

func (*Store) Stats 

func (s *Store) Stats(ctx context.Context) (map[string]int, error)

Stats returns embedding counts per repo.

func (*Store) Upsert 

func (s *Store) Upsert(ctx context.Context, records []EmbeddingRecord) error

Upsert stores embeddings for symbols using multi-row INSERT with ON CONFLICT.

type SymbolIdentity added in v1.20.0 

type SymbolIdentity struct {
	SymbolName string
	BodyHash   uint64 // matches EmbeddingRecord.BodyHash (BIGINT in DB)
	StartLine  int
}

SymbolIdentity holds the identity fields for a single indexed symbol. Used by incremental update paths (Pipeline.IndexFile) to compute diffs: symbols present in DB but absent from new parse → DELETE; symbols whose BodyHash changed → re-embed.

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.