88 lines
2.3 KiB
Go
88 lines
2.3 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"rag/dao"
|
||
|
|
"rag/model/dto"
|
||
|
|
|
||
|
|
"github.com/gogf/gf/v2/util/gconv"
|
||
|
|
)
|
||
|
|
|
||
|
|
var Dataset = new(datasetService)
|
||
|
|
|
||
|
|
type datasetService struct{}
|
||
|
|
|
||
|
|
// Create 创建数据集
|
||
|
|
func (s *datasetService) Create(ctx context.Context, req *dto.CreateDatasetReq) (res *dto.CreateDatasetRes, err error) {
|
||
|
|
id, err := dao.Dataset.Insert(ctx, req)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
return &dto.CreateDatasetRes{Id: id}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Update 更新数据集
|
||
|
|
func (s *datasetService) Update(ctx context.Context, req *dto.UpdateDatasetReq) (err error) {
|
||
|
|
_, err = dao.Dataset.Update(ctx, req)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// Delete 删除数据集
|
||
|
|
func (s *datasetService) Delete(ctx context.Context, req *dto.DeleteDatasetReq) (err error) {
|
||
|
|
_, err = dao.Dataset.Delete(ctx, req)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// List 数据集列表
|
||
|
|
func (s *datasetService) List(ctx context.Context, req *dto.ListDatasetReq) (res *dto.ListDatasetRes, err error) {
|
||
|
|
list, total, err := dao.Dataset.List(ctx, req)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
res = &dto.ListDatasetRes{
|
||
|
|
Total: total,
|
||
|
|
}
|
||
|
|
err = gconv.Struct(list, &res.List)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
//// Search 搜索(示例,实际需要调用向量库)
|
||
|
|
//func (s *datasetService) Search(ctx context.Context, req *dto.SearchReq) (res *dto.SearchRes, err error) {
|
||
|
|
// // 1. 获取数据集信息
|
||
|
|
// kb, err := dao.Dataset.GetByID(ctx, req)
|
||
|
|
// if err != nil {
|
||
|
|
// return nil, err
|
||
|
|
// }
|
||
|
|
//
|
||
|
|
// // 2. 获取文件块
|
||
|
|
// chunks, err := dao.Chunk.FindChunksByKBIDWithLimit(ctx, req.KBID, 0, req.TopK)
|
||
|
|
// if err != nil {
|
||
|
|
// return nil, err
|
||
|
|
// }
|
||
|
|
//
|
||
|
|
// // 3. TODO: 使用向量检索(需要集成向量库)
|
||
|
|
// // 暂时使用简单的关键词匹配
|
||
|
|
// results := make([]dto.SearchResult, 0)
|
||
|
|
// for _, chunk := range chunks {
|
||
|
|
// results = append(results, dto.SearchResult{
|
||
|
|
// Content: chunk.Content,
|
||
|
|
// Score: 0.8, // TODO: 计算实际向量相似度
|
||
|
|
// DocumentID: chunk.DocumentID,
|
||
|
|
// ChunkIndex: chunk.Index,
|
||
|
|
// })
|
||
|
|
// }
|
||
|
|
//
|
||
|
|
// g.Log().Infof(ctx, "数据集[%s]搜索完成,查询:%s,结果数:%d", kb.Name, req.Query, len(results))
|
||
|
|
//
|
||
|
|
// return &dto.SearchRes{Results: results}, nil
|
||
|
|
//}
|
||
|
|
//
|
||
|
|
//// formatChunks 格式化文件块为上下文
|
||
|
|
//func (s *datasetService) formatChunks(chunks []*entity.DocumentChunk) string {
|
||
|
|
// var sb strings.Builder
|
||
|
|
// for i, chunk := range chunks {
|
||
|
|
// sb.WriteString(fmt.Sprintf("[%d] %s\n\n", i+1, chunk.Content))
|
||
|
|
// }
|
||
|
|
// return sb.String()
|
||
|
|
//}
|