vectorstore

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

type CollectionConfig struct {
	Description string

	// 是否强制创建新的知识库，如果已经存在，会返回错误
	ForceNew bool

	// embedding 配置
	ModelName       string
	Dimension       int
	EmbeddingClient aispec.EmbeddingCaller

	// hnsw 配置
	DistanceFuncType      string
	MaxNeighbors          int
	LayerGenerationFactor float64
	EfSearch              int
	EfConstruct           int

	EnablePQ                   bool
	EnableAutoUpdateGraphInfos bool
	DisableEmbedCollectionInfo bool
	LazyLoadEmbeddingClient    bool

	DB *gorm.DB

	MaxChunkSize int
	Overlap      int
	BigTextPlan  string

	CacheSize    int
	PreCacheSize int

	KeyAsUID bool

	TryRebuildHNSWIndex bool
	// contains filtered or unexported fields
}

type CollectionConfigFunc func(config *CollectionConfig)

type CollectionInfo struct {
	Name        string
	Description string
	ModelName   string
	Dimension   int

	M                int
	Ml               float64
	EfSearch         int
	EfConstruct      int
	DistanceFuncType string
}

type CollectionQueryConfig struct {
	Ctx                  context.Context
	Limit                int // 单次子查询的结果限制。
	CollectionNumLimit   int
	QueryCollectionNames []string
	CollectionScoreLimit float64
	EnhancePlan          []string // 默认开启 HyDE 、 泛化查询 、拆分查询
	Filter               func(key string, getDoc func() *Document) bool
	Concurrent           int
	MsgCallBack          func(*RAGSearchResult)
	OnSubQueryStart      func(method string, query string)
	OnQueryFinish        func([]*ScoredResult)
	OnStatus             func(label string, value string)
	OnlyResults          bool // 仅返回最终结果，忽略中间结果和消息

	// On Stream Reader
	OnLogReader func(reader io.Reader)

	RAGSimilarityThreshold   float64 // RAG相似度限制
	EveryQueryResultCallback func(result *ScoredResult)
	RAGQueryType             []string

	EnhanceSearchHandler enhancesearch.SearchHandler

	LoadConfig []CollectionConfigFunc
}

type CollectionQueryOption func(*CollectionQueryConfig)

type Document struct {
	ID              string                 `json:"id"`   // 文档唯一标识符
	Type            schema.RAGDocumentType `json:"type"` // 文档类型
	EntityUUID      string                 `json:"entityUUID"`
	RelatedEntities []string               `json:"relatedEntities"`
	Content         string                 `json:"content"`  // 文档内容
	Metadata        schema.MetadataMap     `json:"metadata"` // 文档元数据
	Embedding       []float32              `json:"-"`        // 文档的嵌入向量，不参与 JSON 序列化
	RuntimeID       string                 `json:"runtimeID"`
}

type DocumentOption func(document *Document)

type EmbeddingClient interface {
	Embedding(text string) ([]float32, error)
	// EmbeddingRaw 返回原始的 embedding 结果，可能包含多个向量
	EmbeddingRaw(text string) ([][]float32, error)
}

type EmptyEmbedding struct{}

type ExportVectorStoreDocument struct {
	DocumentID      string                 `json:"document_id"`
	Metadata        map[string]interface{} `json:"metadata"`
	Embedding       []float32              `json:"embedding"`
	PQCode          []byte                 `json:"pq_code"`
	Content         string                 `json:"content"`
	DocumentType    string                 `json:"document_type"`
	EntityID        string                 `json:"entity_id"`
	RelatedEntities string                 `json:"related_entities"`
}

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

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

type MockEmbedder struct {
	MockEmbedderFunc func(text string) ([]float32, error)
}

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

type NodeOffsetToVectorFunc func(offset uint32) []float32

type RAGBinaryData struct {
	Collection *schema.VectorStoreCollection
	Documents  []*ExportVectorStoreDocument
	Version    uint32
}

type RAGExportConfig struct {
	Ctx               context.Context
	DB                *gorm.DB // 数据库（导入时使用）
	NoHNSWIndex       bool     // 是否不包含HNSW索引（导出时使用）
	OnlyPQCode        bool     // 是否只导出PQ编码（导出时使用）
	NoMetadata        bool     // 是否不导出元数据（导出时使用）
	OverwriteExisting bool     // 是否覆盖现有数据（导入时使用）
	NoOriginInput     bool     // 是否不导出原始输入数据（导出时使用）
	RebuildHNSWIndex  bool     // 是否重新构建HNSW索引（导入时使用）

	CollectionName    string // 指定集合名称（导入时使用，可选）
	DocumentHandler   func(doc schema.VectorStoreDocument) (schema.VectorStoreDocument, error)
	OnProgressHandler func(percent float64, message string, messageType string) // 进度回调

	SerialVersionUID string // 序列化版本号（导入时使用）
	RAGID            string // RAGID（导入时使用）
}

type RAGExportOptionFunc func(*RAGExportConfig)

type RAGSearchResult struct {
	Message     string      `json:"message"`
	Data        interface{} `json:"data"`
	Type        string      `json:"type"`      // message, mid_result, result
	Score       float64     `json:"score"`     // 相似度分数
	Source      string      `json:"source"`    // 结果来源（集合名称）
	Timestamp   int64       `json:"timestamp"` // 时间戳
	QueryMethod string      `json:"query_method"`
	QueryOrigin string      `json:"query_origin"`
	Index       int64       `json:"index"`
}

type SQLiteVectorStoreHNSW struct {
	EnableAutoUpdateGraphInfos bool

	UIDType string
	// contains filtered or unexported fields
}

type ScoredResult struct {
	Index       int64
	QueryMethod string
	QueryOrigin string
	Document    *Document
	Score       float64
	Source      string
}

type SearchResult struct {
	Document *Document `json:"document"` // 检索到的文档
	Score    float64   `json:"score"`    // 相似度得分 (-1 到 1 之间)
}

type SimpleERMAnalysisResult struct {
	Entities      []*schema.ERModelEntity `json:"entities"`
	Relationships []*SimpleRelationship   `json:"relationships"`
	OriginalData  []byte                  `json:"original_data"`
}

type SimpleRelationship struct {
	SourceTemporaryName     string `json:"source_temporary_name"`
	TargetTemporaryName     string `json:"target_temporary_name"`
	RelationshipType        string `json:"relationship_type"`
	RelationshipTypeVerbose string `json:"relationship_type_verbose"`
	DecorationAttributes    string `json:"decoration_attributes"`
}

type VectorStore interface {
	// Add 添加文档到向量存储
	Add(docs ...*Document) error

	// Search 根据查询文本检索相关文档
	Search(query string, page, limit int) ([]SearchResult, error)

	SearchWithFilter(query string, page, limit int, filter func(key string, getDoc func() *Document) bool) ([]SearchResult, error)

	// 非语义 模糊搜1
	FuzzSearch(ctx context.Context, query string, limit int) (<-chan SearchResult, error)

	// Delete 根据 ID 删除文档
	Delete(ids ...string) error

	// Get 根据 ID 获取文档
	Get(id string) (*Document, bool, error)

	// List 列出所有文档
	List() ([]*Document, error)

	// Count 返回文档总数
	Count() (int, error)
}