131 lines
4.3 KiB
Go
131 lines
4.3 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
)
|
|
|
|
var Document = &documentDao{}
|
|
|
|
type documentDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameDocument+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
dataset_id INTEGER NOT NULL DEFAULT 0,
|
|
filename TEXT NOT NULL DEFAULT '',
|
|
file_path TEXT NOT NULL DEFAULT '',
|
|
file_size INTEGER NOT NULL DEFAULT 0,
|
|
file_type TEXT NOT NULL DEFAULT '',
|
|
status INTEGER NOT NULL DEFAULT 0,
|
|
chunk_count INTEGER NOT NULL DEFAULT 0,
|
|
error_msg TEXT NOT NULL DEFAULT '',
|
|
content TEXT NOT NULL DEFAULT '',
|
|
created_at DATETIME DEFAULT (datetime('now','localtime')),
|
|
updated_at DATETIME DEFAULT (datetime('now','localtime'))
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create kb_document table failed: %v", err)
|
|
}
|
|
// 迁移：旧库补 content 列（原始全文，抽屉展示用）
|
|
cnt, err := g.DB(consts.DbGroupDefault).Ctx(ctx).GetValue(ctx,
|
|
"SELECT COUNT(*) FROM pragma_table_info('"+consts.TableNameDocument+"') WHERE name=?", "content")
|
|
if err == nil && cnt.Int64() == 0 {
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx,
|
|
"ALTER TABLE "+consts.TableNameDocument+" ADD COLUMN content TEXT NOT NULL DEFAULT ''"); err != nil {
|
|
g.Log().Warningf(ctx, "alter kb_document add content failed: %v", err)
|
|
}
|
|
}
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, "CREATE INDEX IF NOT EXISTS idx_kb_document_dataset ON "+consts.TableNameDocument+"(dataset_id)"); err != nil {
|
|
g.Log().Warningf(ctx, "create index idx_kb_document_dataset failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *documentDao) GetOne(ctx context.Context, id int64) (*entity.Document, error) {
|
|
var m entity.Document
|
|
err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).Where("id", id).Scan(&m)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
func (d *documentDao) List(ctx context.Context, datasetId int64, page, pageSize int) ([]*entity.Document, int, error) {
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if pageSize < 1 {
|
|
pageSize = 20
|
|
}
|
|
total, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).Where("dataset_id", datasetId).Count()
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
var list []*entity.Document
|
|
err = g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).
|
|
FieldsEx("content"). // 列表不返回全文，详情接口单独取
|
|
Where("dataset_id", datasetId).Page(page, pageSize).OrderDesc("id").Scan(&list)
|
|
if list == nil {
|
|
list = make([]*entity.Document, 0)
|
|
}
|
|
return list, total, err
|
|
}
|
|
|
|
// ListByIds 批量按主键查文档（单表约束：IN ≤100 分批；调用方保证 ids 非空）
|
|
func (d *documentDao) ListByIds(ctx context.Context, ids []int64) ([]*entity.Document, error) {
|
|
var list []*entity.Document
|
|
for start := 0; start < len(ids); start += 100 {
|
|
end := min(start+100, len(ids))
|
|
var part []*entity.Document
|
|
if err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).
|
|
FieldsEx("content").WhereIn("id", ids[start:end]).Scan(&part); err != nil {
|
|
return nil, err
|
|
}
|
|
list = append(list, part...)
|
|
}
|
|
if list == nil {
|
|
list = make([]*entity.Document, 0)
|
|
}
|
|
return list, nil
|
|
}
|
|
|
|
func (d *documentDao) Insert(ctx context.Context, data *entity.Document) (int64, error) {
|
|
now := gtime.Now().Format("Y-m-d H:i:s")
|
|
r, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).Data(g.Map{
|
|
"dataset_id": data.DatasetId,
|
|
"filename": data.Filename,
|
|
"file_path": data.FilePath,
|
|
"file_size": data.FileSize,
|
|
"file_type": data.FileType,
|
|
"status": data.Status,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}).Insert()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return r.LastInsertId()
|
|
}
|
|
|
|
func (d *documentDao) UpdateFields(ctx context.Context, id int64, data g.Map) error {
|
|
data["updated_at"] = gtime.Now().Format("Y-m-d H:i:s")
|
|
_, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).Data(data).Where("id", id).Update()
|
|
return err
|
|
}
|
|
|
|
func (d *documentDao) Delete(ctx context.Context, id int64) error {
|
|
_, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameDocument).Ctx(ctx).Where("id", id).Delete()
|
|
return err
|
|
}
|