embeddings

package
v0.1.0 Latest Latest
Warning

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

Go to latest
Published: May 8, 2026 License: MIT Imports: 28 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 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 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) 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) EnsureSchema

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

EnsureSchema creates the pgvector extension and embeddings table if needed.

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.

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) 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.

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.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL