mirror of https://gitee.com/godoos/godoos.git
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
64 lines
1.7 KiB
package model
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type VecList struct {
|
|
BaseModel
|
|
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
|
|
}
|
|
|