vectorstore

type BedrockRuntimeAPI interface {
	InvokeModel(context.Context, *bedrockruntime.InvokeModelInput, ...func(*bedrockruntime.Options)) (*bedrockruntime.InvokeModelOutput, error)
}

type Call struct {
	Operation      string
	Keys           []string
	Records        []VectorRecord
	Vector         []float32
	TopK           int
	Filter         map[string]any
	ReturnMetadata bool
}

type DeleteInput struct {
	Keys []string
}

type Embedder interface {
	Embed(context.Context, string) ([]float32, error)
	EmbedBatch(context.Context, []string) ([][]float32, error)
}

type Error struct {
	Code    string
	Message string
	Cause   error
}

type FakeEmbedder struct {
	Embeddings map[string][]float32
	Default    []float32
	Calls      []string
}

type FakeStore struct {
	Dimension            int
	RequiredMetadataKeys []string
	// contains filtered or unexported fields
}

type GetInput struct {
	Keys           []string
	ReturnMetadata bool
}

type PutInput struct {
	Records []VectorRecord
}

type QueryHit struct {
	Key      string         `json:"key"`
	Distance float32        `json:"distance"`
	Metadata map[string]any `json:"metadata,omitempty"`
}

type QueryInput struct {
	Vector         []float32
	TopK           int
	Filter         map[string]any
	ReturnMetadata bool
}

type S3VectorStore struct {
	Client           S3VectorsAPI
	VectorBucketName string
	IndexName        string
	Dimension        int
	MaxBatchSize     int
}

type S3VectorsAPI interface {
	PutVectors(context.Context, *s3vectors.PutVectorsInput, ...func(*s3vectors.Options)) (*s3vectors.PutVectorsOutput, error)
	QueryVectors(context.Context, *s3vectors.QueryVectorsInput, ...func(*s3vectors.Options)) (*s3vectors.QueryVectorsOutput, error)
	GetVectors(context.Context, *s3vectors.GetVectorsInput, ...func(*s3vectors.Options)) (*s3vectors.GetVectorsOutput, error)
	DeleteVectors(context.Context, *s3vectors.DeleteVectorsInput, ...func(*s3vectors.Options)) (*s3vectors.DeleteVectorsOutput, error)
}

type SemanticIndex struct {
	Store                Store
	Embedder             Embedder
	Dimension            int
	RequiredMetadataKeys []string
}

type SemanticRecord struct {
	Key      string
	Text     string
	Metadata map[string]any
}

type Store interface {
	PutVectors(context.Context, PutInput) error
	GetVectors(context.Context, GetInput) ([]VectorRecord, error)
	DeleteVectors(context.Context, DeleteInput) error
	QueryVectors(context.Context, QueryInput) ([]QueryHit, error)
}

type TitanEmbedder struct {
	Runtime          BedrockRuntimeAPI
	ModelID          string
	Dimensions       int
	Normalize        bool
	BatchConcurrency int
}

type VectorRecord struct {
	Key      string         `json:"key"`
	Data     []float32      `json:"data"`
	Metadata map[string]any `json:"metadata,omitempty"`
}