rag

type AutoSplitter struct {
	// Size, Overlap and Length are as in RecursiveSplitter.
	Size, Overlap int
	Length        LengthFunc
}

type CSVParser struct {
	// Comma is the field delimiter. Defaults to ','.
	Comma rune
	// ContentColumns form the content. Defaults to every column.
	ContentColumns []string
	// MetadataColumns are copied to the document metadata as strings.
	MetadataColumns []string
	// IDColumn names a column that identifies rows.
	IDColumn string
}

type Cache interface {
	Get(key string) (string, bool)
	Put(key, value string)
}

type ChangeLoader interface {
	LoadChanges(ctx context.Context, cursor string) (*Changes, error)
}

type Changes struct {
	Documents      []*Document `json:"documents"`
	DeletedSources []string    `json:"deleted_sources"`
	Cursor         string      `json:"cursor"`
}

type Chunk struct {
	// ID is stable: it is derived from the source, the document ID and the
	// offset, so re-splitting unchanged content yields the same IDs.
	ID         string `json:"id"`
	DocumentID string `json:"document_id"`
	Source     string `json:"source,omitempty"`
	// Index is the position of the chunk within its document.
	Index int `json:"index"`
	// Offset is the byte offset of Content within the document content.
	Offset  int    `json:"offset"`
	Content string `json:"content"`
	// Context situates the chunk within its document, such as a Markdown
	// header path or text written by ContextualEnricher. It is prepended to
	// Content for embedding and keyword search; see Text.
	Context  string         `json:"context,omitempty"`
	Metadata map[string]any `json:"metadata,omitempty"`
	// Embedding is set by the Index. Stores share it between the copies
	// they return, so it must not be modified.
	Embedding []float32 `json:"embedding,omitempty"`
}

type Citation struct {
	// N is the number of the hit in the output, starting at 1.
	N      int     `json:"n"`
	ID     string  `json:"id"`
	Source string  `json:"source"`
	Score  float64 `json:"score"`
}

type CodeSplitter struct {
	// Language selects the separators. When empty, the document's
	// MetaLanguage metadata is used; unknown languages use the separators
	// of RecursiveSplitter.
	Language string
	// Size, Overlap and Length are as in RecursiveSplitter.
	Size, Overlap int
	Length        LengthFunc
}

type Config struct {
	// Store stores the chunks. Defaults to a new MemoryStore.
	Store VectorStore
	// Writer optionally selects a non-atomic ingestion backend. Store must be
	// nil when Writer is supplied. Set Retriever to query that backend; without
	// it the index supports ingestion only and retrieval returns an error.
	Writer SourceWriter
	// Embedder embeds chunks. Without it chunks are stored without
	// embeddings and only keyword retrieval works.
	Embedder memory.Embedder
	// QueryEmbedder embeds queries. Defaults to Embedder. Use it with
	// embedders that distinguish queries from documents; both must produce
	// vectors in the same space.
	QueryEmbedder memory.Embedder
	// Splitter splits documents. Defaults to AutoSplitter{}.
	Splitter Splitter
	// Enricher, if set, runs on the chunks of each document before they
	// are embedded, such as a ContextualEnricher.
	Enricher Enricher
	// BatchSize is the number of chunks embedded per request. Defaults to
	// DefaultBatchSize.
	BatchSize int
	// Version is included in the content hash of every source. Change it
	// after changing the splitter, enricher or embedder so that the next
	// Add or Sync re-indexes sources whose content did not change.
	Version string
	// Retriever replaces the default retriever: Keyword, fused with Vector
	// by Hybrid when an embedder is set.
	Retriever Retriever
	// Reranker, if set, reranks the candidates of the retriever; see
	// Reranked.
	Reranker Reranker
	// Candidates is the number of hits to rerank; see Reranked.
	Candidates int
	// OnError receives retrieval errors that do not fail a query, such as
	// a failing retriever of Hybrid or a failing Reranker. If nil, they are
	// logged with slog.Default.
	OnError func(err error)
}

type ContextualEnricher struct {
	Model model.Model
	// Instruction follows the chunk in the user message. Defaults to
	// DefaultContextInstruction.
	Instruction string
	// MaxDocumentBytes truncates long documents. Defaults to 100000, about
	// 25000 tokens of English text.
	MaxDocumentBytes int
	// MaxTokens bounds each context. Defaults to 200.
	MaxTokens int
	// Concurrency is the number of chunks enriched in parallel. Defaults to
	// 4.
	Concurrency int
	// Cache stores contexts. Defaults to a MapCache owned by the enricher;
	// supply a persistent cache to reuse contexts across processes.
	Cache Cache
	// contains filtered or unexported fields
}

type Document struct {
	// ID identifies the document. It defaults to Source.
	ID      string `json:"id"`
	Content string `json:"content"`
	// Metadata is copied to every chunk of the document.
	Metadata map[string]any `json:"metadata,omitempty"`
	// Source names where the document came from, such as a file path or a
	// URL. Documents with the same source are indexed and replaced
	// together; a source defaults to the document ID.
	Source string `json:"source,omitempty"`
}

type Enricher interface {
	Enrich(ctx context.Context, doc *Document, chunks []*Chunk) error
}

type EnricherFunc func(ctx context.Context, doc *Document, chunks []*Chunk) error

type FSLoader struct {
	FS fs.FS
	// Root is the directory to load, relative to FS. Defaults to ".".
	Root string
	// Include, when non-empty, loads only files matching one of these
	// globs. Patterns without a slash match file names; others match paths
	// relative to Root. "**" matches any number of directories and "{a,b}"
	// alternatives.
	Include []string
	// Exclude skips files and directories matching one of these globs,
	// which work like Include.
	Exclude []string
	// GitIgnore skips paths excluded by .gitignore files within Root.
	// Ignore files above Root and git's global excludes do not apply.
	GitIgnore bool
	// Parsers maps lower-case file extensions, such as ".md", to parsers.
	// Defaults to DefaultParsers. Files without a parser are skipped.
	Parsers map[string]Parser
	// MaxFileSize skips larger files. Defaults to DefaultMaxFileSize.
	MaxFileSize int64
	// SourcePrefix is prepended to document sources, for example to tell
	// apart several directories in one index.
	SourcePrefix string
}

type Filter struct {
	// Sources matches chunks whose source is one of these when non-empty.
	Sources []string
	// SourcePrefix matches chunks whose source starts with it.
	SourcePrefix string
	// Metadata matches chunks whose metadata has every key with an equal
	// value. Numbers are compared by value whatever their type, and slices
	// element by element, so values decoded from JSON match the originals.
	Metadata map[string]any
}

type FormatOptions struct {
	// MaxChunkBytes bounds each hit's content, including its truncation note.
	// Defaults to 2000.
	MaxChunkBytes int
	// MaxBytes bounds the whole text: hits that would exceed it are
	// replaced by an omission note when space permits, except the first,
	// which is truncated instead. All notes count toward the byte limit.
	// Defaults to 10000.
	MaxBytes int
}

type HTMLParser struct{}

type Hit struct {
	Chunk *Chunk
	Score float64
}

type Hybrid struct {
	Retrievers []Retriever
	// Weights scale the contribution of each retriever. Missing weights
	// are 1.
	Weights []float64
	// K is the RRF constant. Defaults to 60.
	K float64
	// OnError receives the error of each retriever that fails while
	// another succeeds.
	OnError func(err error)
}

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

type JSONParser struct {
	// Lines parses JSON Lines (NDJSON).
	Lines bool
	// ContentFields form the content. One field yields its value, several
	// yield "field: value" lines. Strings are used as-is and other values
	// as JSON. When empty, the content is the whole record as JSON.
	ContentFields []string
	// MetadataFields are copied to the document metadata under their
	// paths.
	MetadataFields []string
	// IDField names a field that identifies records.
	IDField string
}

type Keyword struct {
	Store VectorStore
	// K1 and B are the BM25 parameters. Zero selects the defaults, 1.2 and
	// 0.75; a negative value selects 0.
	K1, B float64
	// contains filtered or unexported fields
}

type LLMReranker struct {
	Model model.Model
	// Prompt is the system prompt. Defaults to DefaultRerankPrompt.
	Prompt string
	// MaxPassageBytes truncates each passage. Defaults to 2000.
	MaxPassageBytes int
	// MaxTokens bounds the reply. Defaults to 256.
	MaxTokens int
}

type LengthFunc func(string) int

type Loader interface {
	Load(ctx context.Context) ([]*Document, error)
}

type LoaderFunc func(ctx context.Context) ([]*Document, error)

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

type MarkdownSplitter struct {
	// Size, Overlap and Length are as in RecursiveSplitter.
	Size, Overlap int
	Length        LengthFunc
	// MaxLevel is the deepest header level that starts a section. Defaults
	// to 6.
	MaxLevel int
	// HeaderContext sets Chunk.Context to the header path, joined with
	// " > ", so it is embedded and searched with each chunk.
	HeaderContext bool
}

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

type Parser interface {
	Parse(ctx context.Context, source string, data []byte) ([]*Document, error)
}

type ParserFunc func(ctx context.Context, source string, data []byte) ([]*Document, error)

type Query struct {
	Text string
	// Limit is the maximum number of hits; non-positive means no limit.
	Limit  int
	Filter Filter
	// MinScore drops hits scoring below it. Scores are retriever specific.
	MinScore float64
}

type RecursiveSplitter struct {
	// Size is the maximum chunk length measured by Length. Defaults to
	// DefaultChunkSize. A single character longer than Size is kept whole.
	Size int
	// Overlap is the length shared by consecutive chunks. Zero selects
	// DefaultChunkOverlap (at most a quarter of Size); a negative value
	// disables overlap.
	Overlap int
	// Length measures text. Defaults to the token estimate of the context
	// package (four ASCII bytes or one other character per token).
	Length LengthFunc
	// Separators are tried in order. Defaults to DefaultSeparators.
	Separators []string
	// Regexp interprets the separators as regular expressions.
	Regexp bool
}

type Reranked struct {
	Retriever Retriever
	Reranker  Reranker
	// Candidates is the number of first-stage hits to rerank. Defaults to
	// four times the limit, and at least 20.
	Candidates int
	OnError    func(err error)
}

type Reranker interface {
	Rerank(ctx context.Context, query string, hits []Hit) ([]Hit, error)
}

type RerankerFunc func(ctx context.Context, query string, hits []Hit) ([]Hit, error)

type Result struct {
	// Added lists new sources, Updated sources whose content changed,
	// Unchanged sources that were skipped and Deleted removed sources.
	Added, Updated, Unchanged, Deleted []string
	// Chunks is the number of chunks stored.
	Chunks int
}

type Retriever interface {
	Retrieve(ctx context.Context, q Query) ([]Hit, error)
}

type RetrieverFunc func(ctx context.Context, q Query) ([]Hit, error)

type Revisioner interface {
	Revision() uint64
}

type SourceChunks struct {
	Info   SourceInfo `json:"info"`
	Chunks []*Chunk   `json:"chunks"`
}

type SourceInfo struct {
	Source string `json:"source"`
	// Hash identifies the indexed content of the source; see Index.
	Hash      string    `json:"hash"`
	Chunks    int       `json:"chunks"`
	IndexedAt time.Time `json:"indexed_at"`
}

type SourceWriter interface {
	Write(ctx context.Context, sources ...SourceChunks) error
	Delete(ctx context.Context, sources ...string) error
	Sources(ctx context.Context) ([]SourceInfo, error)
}

type Splitter interface {
	Split(ctx context.Context, doc *Document) ([]*Chunk, error)
}

type SyncOption func(*syncConfig)

type TextParser struct {
	// Format is stored as MetaFormat. Defaults to "text".
	Format string
	// Language is stored as MetaLanguage when non-empty.
	Language string
}

type ToolOptions struct {
	// Name defaults to "search_knowledge".
	Name string
	// Description defaults to a generic description of a knowledge base
	// search. Describe what the knowledge base holds so the model knows
	// when to search it.
	Description string
	// Limit is the default number of results. Defaults to 5.
	Limit int
	// MaxLimit caps the number of results the model may request. Defaults
	// to 20.
	MaxLimit int
	// Filter restricts every search.
	Filter Filter
	// Format bounds the output.
	Format FormatOptions
}

type URLLoader struct {
	URLs []string
	// Client sends the requests. Defaults to a client with a 60 second
	// timeout.
	Client *http.Client
	// Header is added to every request.
	Header http.Header
	// Parser, when set, parses every response.
	Parser Parser
	// MaxBytes fails on larger responses. Defaults to DefaultMaxFileSize.
	MaxBytes int64
}

type Vector struct {
	Store VectorStore
	// Embedder embeds the query. It must produce vectors in the same space
	// as the embedder of the index.
	Embedder memory.Embedder
}

type VectorStore interface {
	// Replace replaces all chunks of each source, recording its info. Each
	// source is replaced atomically.
	Replace(ctx context.Context, sources ...SourceChunks) error
	// Delete removes sources and their chunks. Unknown sources are
	// ignored.
	Delete(ctx context.Context, sources ...string) error
	// Sources returns the indexed sources ordered by name.
	Sources(ctx context.Context) ([]SourceInfo, error)
	// List returns the chunks matching f, ordered by source, document and
	// position.
	List(ctx context.Context, f Filter) ([]*Chunk, error)
	// Search returns up to limit chunks matching f with an embedding,
	// ranked by cosine similarity to vector. A non-positive limit returns
	// all of them.
	Search(ctx context.Context, vector []float32, limit int, f Filter) ([]Hit, error)
}