tqdb

type CompressModelProgress func(tensorName string, tensorIdx, totalTensors int, rows, totalRows int)

type CompressedProdVector struct {
	CompressedVector
	Signs        []int8  // QJL sign bits, length = QJL projection dimension
	ResidualNorm float32 // ‖residual‖₂
}

type CompressedVector struct {
	Dim     int     // original dimension
	Bits    int     // quantization bits per coordinate
	Norm    float32 // original L2 norm
	Indices []uint8 // quantization indices (one per coordinate, values 0..2^bits-1)
}

type Config struct {
	Dim         int          // embedding dimension (required)
	Bits        int          // bits per coordinate: 1-8 (default: 4)
	Seed        uint64       // rotation matrix seed (default: 42)
	Rotation    RotationType // rotation algorithm (default: RotationQR)
	UseExactPDF bool         // use exact Beta PDF for codebook (default: false, uses Gaussian approx)
}

type Document struct {
	ID        string         // unique identifier
	Content   string         // original text for RAG (optional)
	Data      map[string]any // typed fields (VS2: data)
	Embedding []float64      // raw embedding (nil = auto-embed via EmbeddingFunc)
}

type EmbeddingFunc func(ctx context.Context, text string) ([]float32, error)

type Filter interface {
	Match(data map[string]any) bool
}

type IndexConfig struct {
	FilterFields   []string  // fields to build inverted indexes on
	Type           IndexType // index algorithm: IndexIVF (default) or IndexHNSW
	NumPartitions  int       // IVF: partitions (0 = auto √N)
	NProbe         int       // IVF: partitions to search per query (0 = auto √NumPartitions)
	SkipIVF        bool      // IVF: if true, only build filter indexes
	M              int       // HNSW: max edges per layer (default 16)
	EfConstruction int       // HNSW: build-time beam width (default 200)
}

type IndexType int

type KVCacheConfig struct {
	Layers      int          // number of transformer layers
	Heads       int          // number of attention heads
	HeadDim     int          // dimension per head (typically 64 or 128)
	Bits        int          // quantization bits per coordinate (default: 4)
	PackIndices bool         // bit-pack indices in storage (default: false)
	Rotation    RotationType // rotation algorithm (default: RotationHadamard)
	Seed        uint64       // rotation seed (default: 42)
}

type ModelConfig struct {
	Bits     int          // bits per coordinate (default: 4)
	Rotation RotationType // rotation algorithm (default: RotationHadamard)
	Seed     uint64       // rotation seed (default: 42)
	Workers  int          // concurrent workers per tensor (default: GOMAXPROCS)
}

type QueryOptions struct {
	PageSize int    // max results (default: 100)
	Filter   Filter // required
}

type Result struct {
	ID      string
	Score   float64 // inner product similarity (≈ cosine sim for unit-normalized vectors)
	Content string
	Data    map[string]any
}

type RotationType int

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

type SearchOptions struct {
	TopK     int     // max results (default: 10)
	MinScore float64 // minimum similarity threshold (0 = no filter)
	Offset   int     // skip first N results (pagination)
	Filter   Filter  // data field filter
	Rescore  int     // rescore top-N with exact dequantized distance (0 = disabled, recommended: 3×TopK)
	Ef       int     // HNSW: search beam width (0 = auto, higher = better recall, slower)
}

type StoreConfig struct {
	Dim         int          // embedding dimension (required)
	Bits        int          // bits per coordinate: 1-8 (default: 4)
	Rotation    RotationType // rotation algorithm (default: RotationQR)
	Seed        uint64       // rotation matrix seed (default: 42)
	UseExactPDF bool         // use exact Beta PDF for codebook
}

type StoreInfo struct {
	Path        string
	Dim         int
	WorkDim     int
	Bits        int
	Rotation    RotationType
	Seed        uint64
	NumVecs     int
	FileSize    int64
	IndexBytes  int64
	Compression float64 // ratio vs float32
}

type TQMTensorInfo struct {
	Shape     []int64 `json:"shape"`
	OrigDType string  `json:"orig_dtype"`
	Rows      int     `json:"rows"`
	RowDim    int     `json:"row_dim"`
	WorkDim   int     `json:"work_dim"`
	Offset    int64   `json:"offset"`
	Size      int64   `json:"size"`
	AvgCosSim float64 `json:"avg_cos_sim"`
}

type TQModelHeader struct {
	Version  int                      `json:"version"`
	Source   string                   `json:"source,omitempty"`
	Bits     int                      `json:"bits"`
	Packed   bool                     `json:"packed"`
	Rotation string                   `json:"rotation"`
	Seed     uint64                   `json:"seed"`
	Tensors  map[string]TQMTensorInfo `json:"tensors"`
}

type TensorInfo struct {
	Name        string
	DType       string   `json:"dtype"`
	Shape       []int64  `json:"shape"`
	DataOffsets [2]int64 `json:"data_offsets"`
}