red-future/rag
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 支持多数据库配置与PGVector检索

Browse Source
This commit is contained in:
曲皓东
2026-04-03 17:59:05 +08:00
parent 86c2b7d66e
commit 026beea4d9
12 changed files with 304 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
dao/document_chunk.go
View File

@@ -20,7 +20,7 @@ func (d *documentChunkDao) BatchInsert(ctx context.Context, req []*dto.VectorDoc
if err = gconv.Structs(req, &res); err != nil {
return
}
r, err := gfdb.DB(ctx).Model(ctx, public.TableNameDocumentChunk).Data(&res).Insert()
r, err := gfdb.DB(ctx, public.DbNameVector).Model(ctx, public.TableNameDocumentChunk).Data(&res).Insert()
if err != nil {
return
}
@@ -29,7 +29,7 @@ func (d *documentChunkDao) BatchInsert(ctx context.Context, req []*dto.VectorDoc
// Update 更新文件块
func (d *documentChunkDao) Update(ctx context.Context, req *dto.UpdateDocumentChunkReq) (rows int64, err error) {
model := gfdb.DB(ctx).Model(ctx, public.TableNameDocumentChunk)
model := gfdb.DB(ctx, public.DbNameVector).Model(ctx, public.TableNameDocumentChunk)
r, err := model.Data(&req).Where(entity.DocumentChunkCol.Id, req.Id).Update()
if err != nil {
return
@@ -39,7 +39,7 @@ func (d *documentChunkDao) Update(ctx context.Context, req *dto.UpdateDocumentCh
// List 文件块列表
func (d *documentChunkDao) List(ctx context.Context, req *dto.ListDocumentChunkReq, fields ...string) (res []*entity.DocumentChunk, total int, err error) {
model := gfdb.DB(ctx).Model(ctx, public.TableNameDocumentChunk).Fields(fields).OmitEmpty().
model := gfdb.DB(ctx, public.DbNameVector).Model(ctx, public.TableNameDocumentChunk).Fields(fields).OmitEmpty().
Where(entity.DocumentChunkCol.DatasetId, req.DatasetId).
Where(entity.DocumentChunkCol.DocumentId, req.DocumentId).
Where(entity.DocumentChunkCol.Status, req.Status).
@@ -55,50 +55,3 @@ func (d *documentChunkDao) List(ctx context.Context, req *dto.ListDocumentChunkR
err = r.Structs(&res)
return
}
//// Insert 插入向量文档
//func (d *vectorDocumentDao) Insert(ctx context.Context, docs []*entity.DocumentChunk) (ids []interface{}, err error) {
// if len(docs) == 0 {
// return
// }
// interfaces := make([]interface{}, len(docs))
// for i := range docs {
// interfaces[i] = docs[i]
// }
// return mongoDB.Insert(ctx, interfaces, CollectionVectorDoc)
//}
//
//// DeleteByIDs 根据ID删除向量文档
//func (d *vectorDocumentDao) DeleteByIDs(ctx context.Context, ids []string) (err error) {
// if len(ids) == 0 {
// return
// }
// objectIDs := make([]bson.ObjectID, len(ids))
// for i, id := range ids {
// objectIDs[i], err = bson.ObjectIDFromHex(id)
// if err != nil {
// return err
// }
// }
// filter := bson.M{"_id": bson.M{"$in": objectIDs}}
// _, err = mongoDB.Delete(ctx, filter, CollectionVectorDoc)
// return
//}
//
//// GetByIndexID 根据索引ID获取向量文档
//func (d *vectorDocumentDao) GetByIndexID(ctx context.Context, indexID string, limit int) (result []*entity.DocumentChunk, err error) {
// filter := bson.M{"indexId": indexID}
// page := &beans.Page{PageNum: 1, PageSize: int64(limit)}
// _, err = mongoDB.Find(ctx, filter, &result, CollectionVectorDoc, page, nil)
// return
//}
//
//// GetByVectorIDs 根据向量ID获取向量文档
//func (d *vectorDocumentDao) GetByVectorIDs(ctx context.Context, vectorIDs []string) (result []*entity.DocumentChunk, err error) {
// if len(vectorIDs) == 0 {
// return
// }
// filter := bson.M{"vectorId": bson.M{"$in": vectorIDs}}
// _, err = mongoDB.Find(ctx, filter, &result, CollectionVectorDoc, &beans.Page{PageSize: -1}, nil)
// return
//}