248 lines
8.3 KiB
Go
248 lines
8.3 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"strconv"
|
|
"strings"
|
|
|
|
_ "modernc.org/sqlite/vec"
|
|
|
|
"rag-local/common"
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/model/domain"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
"github.com/gogf/gf/v2/text/gstr"
|
|
)
|
|
|
|
var Chunk = &chunkDao{}
|
|
|
|
type chunkDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameChunk+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
dataset_id INTEGER NOT NULL DEFAULT 0,
|
|
document_id INTEGER NOT NULL DEFAULT 0,
|
|
seq INTEGER NOT NULL DEFAULT 0,
|
|
content TEXT NOT NULL DEFAULT '',
|
|
meta TEXT NOT NULL DEFAULT '',
|
|
created_at DATETIME DEFAULT (datetime('now','localtime'))
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create kb_chunk table failed: %v", err)
|
|
}
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, "CREATE INDEX IF NOT EXISTS idx_kb_chunk_document ON "+consts.TableNameChunk+"(document_id)"); err != nil {
|
|
g.Log().Warningf(ctx, "create index idx_kb_chunk_document failed: %v", err)
|
|
}
|
|
// 向量虚拟表（维度取配置 vector.dim，切换维度需删表重建）
|
|
dim := g.Cfg().MustGet(ctx, "vector.dim", consts.DefaultEmbeddingDim).Int()
|
|
if dim < 1 {
|
|
dim = consts.DefaultEmbeddingDim
|
|
}
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE VIRTUAL TABLE IF NOT EXISTS `+consts.TableNameChunkVec+
|
|
` USING vec0(chunk_id INTEGER PRIMARY KEY, embedding float[`+strconv.Itoa(dim)+`])`); err != nil {
|
|
g.Log().Warningf(ctx, "create vec0 table failed: %v", err)
|
|
}
|
|
// 全文索引虚拟表（默认 unicode61 tokenizer；中文分词在应用层完成，content_tokens 存分词后空格连接文本）
|
|
if _, err := g.DB(consts.DbGroupDefault).Exec(ctx, `CREATE VIRTUAL TABLE IF NOT EXISTS `+consts.TableNameChunkFts+
|
|
` USING fts5(chunk_id UNINDEXED, dataset_id UNINDEXED, title, content_tokens)`); err != nil {
|
|
g.Log().Warningf(ctx, "create fts5 table failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *chunkDao) GetOne(ctx context.Context, id int64) (*entity.Chunk, error) {
|
|
var m entity.Chunk
|
|
err := g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).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 *chunkDao) ListByDocument(ctx context.Context, documentId int64, page, pageSize int) ([]*entity.Chunk, int, error) {
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if pageSize < 1 {
|
|
pageSize = 20
|
|
}
|
|
total, err := g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).Ctx(ctx).Where("document_id", documentId).Count()
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
var list []*entity.Chunk
|
|
err = g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).Ctx(ctx).
|
|
Where("document_id", documentId).Page(page, pageSize).OrderAsc("seq").Scan(&list)
|
|
if list == nil {
|
|
list = make([]*entity.Chunk, 0)
|
|
}
|
|
return list, total, err
|
|
}
|
|
|
|
// InsertWithVec 事务内写入 chunk + 向量 + 全文索引，返回 chunk id
|
|
func (d *chunkDao) InsertWithVec(ctx context.Context, datasetId, documentId int64, seq int, content, meta, title string, vecJson string, dim int) (int64, error) {
|
|
tx, err := g.DB(consts.DbGroupDefault).Begin(ctx)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
r, err := tx.Model(consts.TableNameChunk).Ctx(ctx).Data(g.Map{
|
|
"dataset_id": datasetId,
|
|
"document_id": documentId,
|
|
"seq": seq,
|
|
"content": content,
|
|
"meta": meta,
|
|
"created_at": gtime.Now().Format("Y-m-d H:i:s"),
|
|
}).Insert()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
chunkId, _ := r.LastInsertId()
|
|
if chunkId == 0 {
|
|
return 0, gerror.New("chunk insert failed")
|
|
}
|
|
if vecJson != "" {
|
|
if _, err := tx.Exec("INSERT INTO "+consts.TableNameChunkVec+" (chunk_id, embedding) VALUES (?, vec_f32(?))", chunkId, vecJson); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
if _, err := tx.Exec("INSERT INTO "+consts.TableNameChunkFts+" (chunk_id, dataset_id, title, content_tokens) VALUES (?, ?, ?, ?)",
|
|
chunkId, datasetId, title, common.Tokenize(content)); err != nil {
|
|
return 0, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return 0, err
|
|
}
|
|
return chunkId, nil
|
|
}
|
|
|
|
// DeleteByDocument 事务内删除文档全部分块（chunk + 向量 + 全文索引）
|
|
func (d *chunkDao) DeleteByDocument(ctx context.Context, documentId int64) error {
|
|
tx, err := g.DB(consts.DbGroupDefault).Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
ids := tx.Model(consts.TableNameChunk).Ctx(ctx).Fields("id").Where("document_id", documentId)
|
|
r, err := ids.Array()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
chunkIds := make([]int64, 0, len(r))
|
|
for _, v := range r {
|
|
chunkIds = append(chunkIds, v.Int64())
|
|
}
|
|
if len(chunkIds) > 0 {
|
|
placeholders := make([]string, 0, len(chunkIds))
|
|
args := make([]interface{}, 0, len(chunkIds))
|
|
for _, id := range chunkIds {
|
|
placeholders = append(placeholders, "?")
|
|
args = append(args, id)
|
|
}
|
|
in := gstr.Join(placeholders, ",")
|
|
if _, err := tx.Exec("DELETE FROM "+consts.TableNameChunkVec+" WHERE chunk_id IN ("+in+")", args...); err != nil {
|
|
return err
|
|
}
|
|
if _, err := tx.Exec("DELETE FROM "+consts.TableNameChunkFts+" WHERE chunk_id IN ("+in+")", args...); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if _, err := tx.Model(consts.TableNameChunk).Ctx(ctx).Where("document_id", documentId).Delete(); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit()
|
|
}
|
|
|
|
func (d *chunkDao) UpdateContent(ctx context.Context, id int64, content, vecJson, tokens string) error {
|
|
tx, err := g.DB(consts.DbGroupDefault).Begin(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
|
|
if _, err := tx.Model(consts.TableNameChunk).Ctx(ctx).Data(g.Map{"content": content}).Where("id", id).Update(); err != nil {
|
|
return err
|
|
}
|
|
if vecJson != "" {
|
|
if _, err := tx.Exec("UPDATE "+consts.TableNameChunkVec+" SET embedding = vec_f32(?) WHERE chunk_id = ?", vecJson, id); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if _, err := tx.Exec("UPDATE "+consts.TableNameChunkFts+" SET content_tokens = ? WHERE chunk_id = ?", tokens, id); err != nil {
|
|
return err
|
|
}
|
|
return tx.Commit()
|
|
}
|
|
|
|
// UpdateVec 仅更新向量（重新向量化，不动文本与 FTS）
|
|
func (d *chunkDao) UpdateVec(ctx context.Context, id int64, vecJson string) error {
|
|
if vecJson == "" {
|
|
return nil
|
|
}
|
|
_, err := g.DB(consts.DbGroupDefault).Exec(ctx,
|
|
"UPDATE "+consts.TableNameChunkVec+" SET embedding = vec_f32(?) WHERE chunk_id = ?", vecJson, id)
|
|
return err
|
|
}
|
|
|
|
func (d *chunkDao) CountByDocument(ctx context.Context, documentId int64) (int, error) {
|
|
return g.DB(consts.DbGroupDefault).Model(consts.TableNameChunk).Ctx(ctx).Where("document_id", documentId).Count()
|
|
}
|
|
|
|
// VecSearch 向量 KNN 检索：vec0 取最近 topK*4 后按数据集过滤（vec0 无 dataset 列）
|
|
func (d *chunkDao) VecSearch(ctx context.Context, datasetId int64, vecJson string, topK int) ([]domain.VecHit, error) {
|
|
r, err := g.DB(consts.DbGroupDefault).Ctx(ctx).Raw(
|
|
`SELECT v.chunk_id, v.distance FROM (
|
|
SELECT chunk_id, distance FROM `+consts.TableNameChunkVec+`
|
|
WHERE embedding MATCH ? ORDER BY distance LIMIT ?
|
|
) v INNER JOIN `+consts.TableNameChunk+` c ON c.id = v.chunk_id
|
|
WHERE c.dataset_id = ? ORDER BY v.distance LIMIT ?`,
|
|
vecJson, topK*4, datasetId, topK,
|
|
).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
hits := make([]domain.VecHit, 0, len(r))
|
|
for _, row := range r {
|
|
hits = append(hits, domain.VecHit{
|
|
ChunkId: row["chunk_id"].Int64(),
|
|
Distance: row["distance"].Float64(),
|
|
})
|
|
}
|
|
return hits, nil
|
|
}
|
|
|
|
// FtsSearch 全文检索（BM25）：中文分词已在应用层完成，query 为空格连接的引号词串
|
|
func (d *chunkDao) FtsSearch(ctx context.Context, datasetId int64, query string, topK int) ([]domain.FtsHit, error) {
|
|
if strings.TrimSpace(query) == "" {
|
|
return nil, nil
|
|
}
|
|
r, err := g.DB(consts.DbGroupDefault).Ctx(ctx).Raw(
|
|
`SELECT chunk_id, bm25(`+consts.TableNameChunkFts+`) AS score FROM `+consts.TableNameChunkFts+
|
|
` WHERE `+consts.TableNameChunkFts+` MATCH ? AND dataset_id = ? ORDER BY score LIMIT ?`,
|
|
query, datasetId, topK,
|
|
).All()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
hits := make([]domain.FtsHit, 0, len(r))
|
|
for _, row := range r {
|
|
hits = append(hits, domain.FtsHit{
|
|
ChunkId: row["chunk_id"].Int64(),
|
|
Score: row["score"].Float64(),
|
|
})
|
|
}
|
|
return hits, nil
|
|
}
|