embedder

package
v0.35.0 Latest Latest
Warning

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

 Go to latest Published: Mar 16, 2026 License: MIT Imports: 18 Imported by: 0

Documentation

Index

Constants

View Source 
const MaxBatchSize = 2000

MaxBatchSize is the maximum number of inputs per OpenAI embedding API call. OpenAI allows 2048, but we use 2000 as a safety margin.

View Source 
const MaxBatchTokens = 280000

MaxBatchTokens is the maximum total tokens per OpenAI embedding API batch. OpenAI has a 300,000 token limit. We use 280,000 for safety margin.

Variables

This section is empty.

Functions

func EstimateTokens added in v0.24.0 

func EstimateTokens(text string) int

EstimateTokens estimates the token count for a text string. Uses a conservative estimate of ~4 characters per token for English text. This is intentionally conservative to avoid hitting API limits.

func IsContextLengthError added in v0.25.0 

func IsContextLengthError(err error) bool

IsContextLengthError checks if an error is or wraps a ContextLengthError.

func IsRetryable added in v0.24.0 

func IsRetryable(statusCode int) bool

IsRetryable returns true if the HTTP status code indicates a retryable error. Retryable: 429 (rate limit), 5xx (server errors) Non-retryable: 4xx client errors (except 429)

func MapResultsToFiles added in v0.24.0 

func MapResultsToFiles(batches []Batch, results []BatchResult, numFiles int) [][][]float32

MapResultsToFiles maps batch results back to per-file embeddings. Returns a slice where each index corresponds to a file, containing embeddings for that file's chunks.

Types

type AdaptiveRateLimiter added in v0.24.0 

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

AdaptiveRateLimiter coordinates adaptive parallelism adjustment based on rate limit feedback. It reduces parallelism when rate limits are detected and gradually restores it after successful requests.

func NewAdaptiveRateLimiter added in v0.24.0 

func NewAdaptiveRateLimiter(maxWorkers int) *AdaptiveRateLimiter

NewAdaptiveRateLimiter creates a new AdaptiveRateLimiter with the given maximum parallelism. Default thresholds: reductionThreshold=3, restorationThreshold=10

func (*AdaptiveRateLimiter) CurrentWorkers added in v0.24.0 

func (arl *AdaptiveRateLimiter) CurrentWorkers() int

CurrentWorkers returns the current parallelism level.

func (*AdaptiveRateLimiter) MaxWorkers added in v0.24.0 

func (arl *AdaptiveRateLimiter) MaxWorkers() int

MaxWorkers returns the configured maximum parallelism.

func (*AdaptiveRateLimiter) OnRateLimitHit added in v0.24.0 

func (arl *AdaptiveRateLimiter) OnRateLimitHit() bool

OnRateLimitHit should be called when a 429 response is received. It increments the rate limit hit counter and may reduce parallelism. Returns true if parallelism was reduced.

func (*AdaptiveRateLimiter) OnSuccess added in v0.24.0 

func (arl *AdaptiveRateLimiter) OnSuccess() bool

OnSuccess should be called when a request succeeds. It increments the success streak counter and may restore parallelism. Returns true if parallelism was restored.

type Batch added in v0.24.0 

type Batch struct {
	// Entries contains chunks with source file tracking
	Entries []BatchEntry
	// Index is the batch number for progress reporting (0-indexed)
	Index int
}

Batch represents a collection of chunks to be embedded in a single API call.

func FormBatches added in v0.24.0 

func FormBatches(files []FileChunks) []Batch

FormBatches splits chunks from multiple files into batches respecting both MaxBatchSize (input count) and MaxBatchTokens (token limit). Chunks maintain their file/chunk index tracking for result mapping.

func (*Batch) Contents added in v0.24.0 

func (b *Batch) Contents() []string

Contents returns the text contents of all entries for embedding.

func (*Batch) Size added in v0.24.0 

func (b *Batch) Size() int

Size returns the number of entries in the batch.

type BatchEmbedder added in v0.24.0 

type BatchEmbedder interface {
	Embedder

	// EmbedBatches processes multiple batches of chunks concurrently.
	// It returns results mapped back to their source files, or an error if any batch fails.
	// The progress callback is called for each batch completion or retry attempt.
	EmbedBatches(ctx context.Context, batches []Batch, progress BatchProgress) ([]BatchResult, error)
}

BatchEmbedder extends Embedder with cross-file batch embedding capabilities. Providers that support advanced batching (like OpenAI) implement this interface to enable parallel processing of multiple batches.

type BatchEntry added in v0.24.0 

type BatchEntry struct {
	// FileIndex is the index of the source file in the files slice
	FileIndex int
	// ChunkIndex is the index of the chunk within the file's chunks
	ChunkIndex int
	// Content is the text content to embed
	Content string
}

BatchEntry represents a single chunk with metadata for tracking its source.

type BatchProgress added in v0.24.0 

type BatchProgress func(batchIndex, totalBatches, completedChunks, totalChunks int, retrying bool, attempt int, statusCode int)

BatchProgress is a callback for reporting batch embedding progress. It receives the batch index, total batches, chunk progress info, and optional retry information. completedChunks and totalChunks track overall progress across all batches. statusCode is the HTTP status code when retrying (429 = rate limited, 5xx = server error).

type BatchResult added in v0.24.0 

type BatchResult struct {
	// BatchIndex is the index of the batch this result belongs to
	BatchIndex int
	// Embeddings contains the embedding vectors in the same order as batch entries
	Embeddings [][]float32
}

BatchResult contains the embeddings for a batch with file/chunk index mapping.

type ContextLengthError added in v0.25.0 

type ContextLengthError struct {
	ChunkIndex      int    // Index of the chunk that exceeded the limit
	EstimatedTokens int    // Estimated number of tokens in the chunk
	MaxTokens       int    // Maximum tokens allowed by the model (if known)
	Message         string // Original error message from the provider
}

ContextLengthError indicates that the input text exceeds the model's context limit. This error is retryable by re-chunking the input into smaller pieces.

func AsContextLengthError added in v0.25.0 

func AsContextLengthError(err error) *ContextLengthError

AsContextLengthError extracts a ContextLengthError from an error chain. Returns nil if the error is not a ContextLengthError.

func NewContextLengthError added in v0.25.0 

func NewContextLengthError(chunkIndex, estimatedTokens, maxTokens int, message string) *ContextLengthError

NewContextLengthError creates a new ContextLengthError.

func (*ContextLengthError) Error added in v0.25.0 

func (e *ContextLengthError) Error() string

type Embedder 

type Embedder interface {
	// Embed converts text into a vector embedding
	Embed(ctx context.Context, text string) ([]float32, error)

	// EmbedBatch converts multiple texts into vector embeddings
	EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

	// Dimensions returns the vector dimension size for this embedder
	Dimensions() int

	// Close cleanly shuts down the embedder
	Close() error
}

Embedder defines the interface for text embedding providers

func NewFromConfig added in v0.32.0 

func NewFromConfig(cfg *config.Config) (Embedder, error)

NewFromConfig creates an Embedder based on the provided configuration. This factory function centralizes provider initialization and eliminates code duplication across CLI commands and MCP server.

func NewFromWorkspaceConfig added in v0.32.0 

func NewFromWorkspaceConfig(ws *config.Workspace) (Embedder, error)

NewFromWorkspaceConfig creates an Embedder from workspace configuration. This is a convenience wrapper for workspace-specific embedder creation.

type FileChunks added in v0.24.0 

type FileChunks struct {
	// FileIndex is the index of this file in the original files slice
	FileIndex int
	// Chunks is the list of text chunks from this file
	Chunks []string
}

FileChunks represents chunks from a single file for batch formation.

type LMStudioEmbedder added in v0.4.0 

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

func NewLMStudioEmbedder added in v0.4.0 

func NewLMStudioEmbedder(opts ...LMStudioOption) *LMStudioEmbedder

func (*LMStudioEmbedder) Close added in v0.4.0 

func (e *LMStudioEmbedder) Close() error

func (*LMStudioEmbedder) Dimensions added in v0.4.0 

func (e *LMStudioEmbedder) Dimensions() int

func (*LMStudioEmbedder) Embed added in v0.4.0 

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

func (*LMStudioEmbedder) EmbedBatch added in v0.4.0 

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

func (*LMStudioEmbedder) Ping added in v0.4.0 

func (e *LMStudioEmbedder) Ping(ctx context.Context) error

type LMStudioOption added in v0.4.0 

type LMStudioOption func(*LMStudioEmbedder)

func WithLMStudioDimensions added in v0.12.0 

func WithLMStudioDimensions(dimensions int) LMStudioOption

func WithLMStudioEndpoint added in v0.4.0 

func WithLMStudioEndpoint(endpoint string) LMStudioOption

func WithLMStudioModel added in v0.4.0 

func WithLMStudioModel(model string) LMStudioOption

type OllamaEmbedder 

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

func NewOllamaEmbedder 

func NewOllamaEmbedder(opts ...OllamaOption) *OllamaEmbedder

func (*OllamaEmbedder) Close 

func (e *OllamaEmbedder) Close() error

func (*OllamaEmbedder) Dimensions 

func (e *OllamaEmbedder) Dimensions() int

func (*OllamaEmbedder) Embed 

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

func (*OllamaEmbedder) EmbedBatch 

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

func (*OllamaEmbedder) Ping 

func (e *OllamaEmbedder) Ping(ctx context.Context) error

Ping checks if Ollama is reachable

type OllamaOption 

type OllamaOption func(*OllamaEmbedder)

func WithOllamaDimensions added in v0.12.0 

func WithOllamaDimensions(dimensions int) OllamaOption

func WithOllamaEndpoint 

func WithOllamaEndpoint(endpoint string) OllamaOption

func WithOllamaModel 

func WithOllamaModel(model string) OllamaOption

type OpenAIEmbedder 

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

func NewOpenAIEmbedder 

func NewOpenAIEmbedder(opts ...OpenAIOption) (*OpenAIEmbedder, error)

func (*OpenAIEmbedder) Close 

func (e *OpenAIEmbedder) Close() error

func (*OpenAIEmbedder) Dimensions 

func (e *OpenAIEmbedder) Dimensions() int

func (*OpenAIEmbedder) Embed 

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

func (*OpenAIEmbedder) EmbedBatch 

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

func (*OpenAIEmbedder) EmbedBatches added in v0.24.0 

func (e *OpenAIEmbedder) EmbedBatches(ctx context.Context, batches []Batch, progress BatchProgress) ([]BatchResult, error)

EmbedBatches implements the BatchEmbedder interface. It processes multiple batches concurrently using a bounded worker pool and retries failed requests with exponential backoff.

type OpenAIOption 

type OpenAIOption func(*OpenAIEmbedder)

func WithOpenAIDimensions added in v0.12.0 

func WithOpenAIDimensions(dimensions int) OpenAIOption

func WithOpenAIEndpoint 

func WithOpenAIEndpoint(endpoint string) OpenAIOption

func WithOpenAIKey 

func WithOpenAIKey(key string) OpenAIOption

func WithOpenAIModel 

func WithOpenAIModel(model string) OpenAIOption

func WithOpenAIParallelism added in v0.24.0 

func WithOpenAIParallelism(parallelism int) OpenAIOption

func WithOpenAIRetryPolicy added in v0.24.0 

func WithOpenAIRetryPolicy(policy RetryPolicy) OpenAIOption

func WithOpenAITPMLimit added in v0.24.0 

func WithOpenAITPMLimit(tpm int64) OpenAIOption

WithOpenAITPMLimit sets the tokens-per-minute limit for proactive rate limiting. When set > 0, the embedder will pace requests to stay within this limit. Default: 0 (disabled). OpenAI Tier 1 limit is 1,000,000 TPM for embeddings.

type OpenRouterEmbedder added in v0.32.0 

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

OpenRouterEmbedder implements the Embedder interface for the OpenRouter API. Note: Parallelism was intentionally removed to keep the implementation simple. OpenRouter processes batches efficiently as-is. If parallel processing is needed in the future, consider implementing the BatchEmbedder interface similar to OpenAIEmbedder with AdaptiveRateLimiter and worker pools.

func NewOpenRouterEmbedder added in v0.32.0 

func NewOpenRouterEmbedder(opts ...OpenRouterOption) (*OpenRouterEmbedder, error)

func (*OpenRouterEmbedder) Close added in v0.32.0 

func (e *OpenRouterEmbedder) Close() error

func (*OpenRouterEmbedder) Dimensions added in v0.32.0 

func (e *OpenRouterEmbedder) Dimensions() int

func (*OpenRouterEmbedder) Embed added in v0.32.0 

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

func (*OpenRouterEmbedder) EmbedBatch added in v0.32.0 

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

func (*OpenRouterEmbedder) Ping added in v0.32.0 

func (e *OpenRouterEmbedder) Ping(ctx context.Context) error

Ping checks if OpenRouter API is reachable

type OpenRouterOption added in v0.32.0 

type OpenRouterOption func(*OpenRouterEmbedder)

func WithOpenRouterDimensions added in v0.32.0 

func WithOpenRouterDimensions(dimensions int) OpenRouterOption

func WithOpenRouterEndpoint added in v0.32.0 

func WithOpenRouterEndpoint(endpoint string) OpenRouterOption

func WithOpenRouterKey added in v0.32.0 

func WithOpenRouterKey(key string) OpenRouterOption

func WithOpenRouterModel added in v0.32.0 

func WithOpenRouterModel(model string) OpenRouterOption

type RateLimitHeaders added in v0.24.0 

type RateLimitHeaders struct {
	// RetryAfter is the duration to wait before retrying (from Retry-After header).
	// Zero if header not present.
	RetryAfter time.Duration

	// RemainingTokens from x-ratelimit-remaining-tokens header.
	// -1 if header not present.
	RemainingTokens int

	// RemainingRequests from x-ratelimit-remaining-requests header.
	// -1 if header not present.
	RemainingRequests int

	// ResetTokens duration from x-ratelimit-reset-tokens header.
	// Zero if header not present.
	ResetTokens time.Duration
}

RateLimitHeaders contains parsed rate limit information from HTTP response headers.

func ParseRateLimitHeadersForTest added in v0.24.0 

func ParseRateLimitHeadersForTest(header http.Header) RateLimitHeaders

ParseRateLimitHeadersForTest exports parseRateLimitHeaders for testing.

type RetryPolicy added in v0.24.0 

type RetryPolicy struct {
	// BaseDelay is the initial delay before the first retry (default: 1s)
	BaseDelay time.Duration
	// Multiplier is the factor to multiply delay by on each retry (default: 2.0)
	Multiplier float64
	// MaxDelay is the maximum delay cap (default: 32s)
	MaxDelay time.Duration
	// MaxAttempts is the maximum number of retry attempts (default: 5)
	MaxAttempts int
}

RetryPolicy configures exponential backoff retry behavior for API calls.

func DefaultRetryPolicy added in v0.24.0 

func DefaultRetryPolicy() RetryPolicy

DefaultRetryPolicy returns a RetryPolicy with sensible defaults for OpenAI API. BaseDelay: 1s, Multiplier: 2x, MaxDelay: 32s, MaxAttempts: 5

func (RetryPolicy) Calculate added in v0.24.0 

func (p RetryPolicy) Calculate(attempt int) time.Duration

Calculate returns the delay for the given attempt number (0-indexed) with jitter. Jitter adds random variance (0-100% of calculated delay) to prevent thundering herd.

func (RetryPolicy) ShouldRetry added in v0.24.0 

func (p RetryPolicy) ShouldRetry(attempt int) bool

ShouldRetry returns true if more attempts are available. attempt is 0-indexed (0 = first attempt, 1 = first retry, etc.)

type RetryableError added in v0.24.0 

type RetryableError struct {
	StatusCode       int
	Message          string
	Retryable        bool
	RateLimitHeaders *RateLimitHeaders
}

RetryableError wraps an error with its HTTP status code for retry decisions.

func NewRetryableError added in v0.24.0 

func NewRetryableError(statusCode int, message string) *RetryableError

NewRetryableError creates a RetryableError from an HTTP status code and message.

func (*RetryableError) Error added in v0.24.0 

func (e *RetryableError) Error() string

type SyntheticEmbedder added in v0.32.0 

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

func NewSyntheticEmbedder added in v0.32.0 

func NewSyntheticEmbedder(opts ...SyntheticOption) (*SyntheticEmbedder, error)

func (*SyntheticEmbedder) Close added in v0.32.0 

func (e *SyntheticEmbedder) Close() error

func (*SyntheticEmbedder) Dimensions added in v0.32.0 

func (e *SyntheticEmbedder) Dimensions() int

func (*SyntheticEmbedder) Embed added in v0.32.0 

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

func (*SyntheticEmbedder) EmbedBatch added in v0.32.0 

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

func (*SyntheticEmbedder) Ping added in v0.32.0 

func (e *SyntheticEmbedder) Ping(ctx context.Context) error

Ping checks if Synthetic API is reachable

type SyntheticOption added in v0.32.0 

type SyntheticOption func(*SyntheticEmbedder)

func WithSyntheticDimensions added in v0.32.0 

func WithSyntheticDimensions(dimensions int) SyntheticOption

func WithSyntheticEndpoint added in v0.32.0 

func WithSyntheticEndpoint(endpoint string) SyntheticOption

func WithSyntheticKey added in v0.32.0 

func WithSyntheticKey(key string) SyntheticOption

func WithSyntheticModel added in v0.32.0 

func WithSyntheticModel(model string) SyntheticOption

type TokenBucket added in v0.24.0 

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

TokenBucket tracks token usage for proactive rate limiting using a sliding window.

func NewTokenBucket added in v0.24.0 

func NewTokenBucket(tokensPerMinute int64) *TokenBucket

NewTokenBucket creates a new TokenBucket with the given tokens-per-minute limit. Default for OpenAI Tier 1 embeddings is 1,000,000 TPM.

func (*TokenBucket) AddTokens added in v0.24.0 

func (tb *TokenBucket) AddTokens(tokens int64) bool

AddTokens records token usage. Returns true if tokens were successfully added.

func (*TokenBucket) TokensAvailable added in v0.24.0 

func (tb *TokenBucket) TokensAvailable() int64

TokensAvailable returns the number of tokens available in the current window. Resets the window if a minute has passed.

func (*TokenBucket) WaitForTokens added in v0.24.0 

func (tb *TokenBucket) WaitForTokens(tokens int64) time.Duration

WaitForTokens checks if the requested tokens are available. Returns the duration to wait if tokens are not available, or 0 if they are.

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.