You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

64 lines
1.7 KiB

package model
import (
"fmt"
"gorm.io/gorm"
)
type VecList struct {
gorm.Model
FilePath string `json:"file_path" gorm:"not null"`
Engine string `json:"engine" gorm:"not null"`
EmbedSize int `json:"embed_size"`
EmbeddingModel string `json:"model" gorm:"not null"`
}
func (*VecList) TableName() string {
return "vec_list"
}
// BeforeCreate 在插入数据之前检查是否存在相同路径的数据
func (v *VecList) BeforeCreate(tx *gorm.DB) error {
var count int64
if err := tx.Model(&VecList{}).Where("file_path = ?", v.FilePath).Count(&count).Error; err != nil {
return err
}
if count > 0 {
return fmt.Errorf("file path already exists: %s", v.FilePath)
}
return nil
}
// AfterCreate 在插入数据之后创建虚拟表
func (v *VecList) AfterCreate(tx *gorm.DB) error {
return CreateVirtualTable(tx, v.ID, v.EmbedSize)
}
// AfterDelete 在删除数据之后删除虚拟表
func (v *VecList) AfterDelete(tx *gorm.DB) error {
// 删除 VecDoc 表中 ListID 对应的所有数据
if err := tx.Where("list_id = ?", v.ID).Delete(&VecDoc{}).Error; err != nil {
return err
}
return DropVirtualTable(tx, v.ID)
}
// CreateVirtualTable 创建虚拟表
func CreateVirtualTable(db *gorm.DB, vectorID uint, embeddingSize int) error {
sql := fmt.Sprintf(`
CREATE VIRTUAL TABLE IF NOT EXISTS [%d_vec] USING
vec0(
document_id TEXT PRIMARY KEY,
embedding float[%d] distance_metric=cosine
)
`, vectorID, embeddingSize)
//log.Printf("sql: %s", sql)
return db.Exec(sql).Error
}
// DropVirtualTable 删除虚拟表
func DropVirtualTable(db *gorm.DB, vectorID uint) error {
sql := fmt.Sprintf(`DROP TABLE IF EXISTS [%d_vec]`, vectorID)
return db.Exec(sql).Error
}