vectorstores

type ContextNetwork struct {
	// Dependencies are documents that the current code imports/depends on
	Dependencies []schema.Document
	// Dependents are documents that import/depend on the current code (impact analysis)
	Dependents []schema.Document
}

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

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

type DocumentWithScore struct {
	Document schema.Document
	Score    float32
}

type HyDEOption func(*HyDERetriever)

type HyDERetriever struct {
	// BaseRetriever performs the actual similarity search
	BaseRetriever schema.Retriever
	// Generator produces a hypothetical document from a query
	Generator func(ctx context.Context, query string) (string, error)
	// NumGenerations controls how many hypothetical docs to generate (default 1).
	// When > 1, results from all generated docs are deduplicated.
	NumGenerations int
}

type MultiQueryRetriever struct {
	Store        VectorStore
	LLM          llms.Model
	NumDocuments int
	Count        int
	// Max results to return after deduplication across all query variations.
	// Defaults to NumDocuments when zero.
	MaxResults int
	// Hook to generate sparse vectors for the newly generated queries
	SparseGenFunc func(ctx context.Context, queries []string) ([]*schema.SparseVector, error)
}

type Option func(*Options)

type Options struct {
	Embedder       embeddings.Embedder
	NameSpace      string
	CollectionName string
	ScoreThreshold float32
	Filters        map[string]any
	SparseQuery    *schema.SparseVector
	SparseQueries  []*schema.SparseVector
}

type RerankingRetriever struct {
	Retriever schema.Retriever
	Reranker  schema.Reranker
	TopK      int // Final number of documents to return after reranking

	// Pre-filter candidates before sending to the reranker. Useful for cheap
	// filtering (e.g. BM25 scoring) to reduce the number of documents the
	// expensive LLM-based reranker has to process.
	CandidateFilter func(query string, docs []schema.Document) []schema.Document

	// MinScore filters out documents that have a reranked score below this threshold.
	// If zero, no threshold is applied.
	MinScore float32
}

type Retriever interface {
	GetRelevantDocuments(ctx context.Context, query string) ([]schema.Document, error)
}

type ScoredRetriever interface {
	GetRelevantScoredDocuments(ctx context.Context, query string) ([]schema.ScoredDocument, error)
}

type VectorStore interface {
	AddDocuments(ctx context.Context, docs []schema.Document, options ...Option) ([]string, error)
	SimilaritySearch(ctx context.Context, query string, numDocuments int, options ...Option) ([]schema.Document, error)
	SimilaritySearchBatch(ctx context.Context, queries []string, numDocuments int, options ...Option) ([][]schema.Document, error)
	SimilaritySearchWithScores(ctx context.Context, query string, numDocuments int, options ...Option) ([]DocumentWithScore, error)
	ListCollections(ctx context.Context) ([]string, error)
	DeleteCollection(ctx context.Context, collectionName string) error
	DeleteDocumentsByFilter(ctx context.Context, filters map[string]any, options ...Option) error
}