|
@@ -0,0 +1,91 @@
|
|
|
+package agent
|
|
|
+
|
|
|
+import (
|
|
|
+ "context"
|
|
|
+ "encoding/csv"
|
|
|
+ "fmt"
|
|
|
+ "github.com/suyuan32/simple-admin-common/msg/errormsg"
|
|
|
+ "mime/multipart"
|
|
|
+ agentModel "wechat-api/ent/agent"
|
|
|
+ "wechat-api/hook/fastgpt"
|
|
|
+ "wechat-api/internal/utils/dberrorhandler"
|
|
|
+
|
|
|
+ "wechat-api/internal/svc"
|
|
|
+ "wechat-api/internal/types"
|
|
|
+
|
|
|
+ "github.com/zeromicro/go-zero/core/logx"
|
|
|
+)
|
|
|
+
|
|
|
+type UploadAgentDataLogic struct {
|
|
|
+ logx.Logger
|
|
|
+ ctx context.Context
|
|
|
+ svcCtx *svc.ServiceContext
|
|
|
+}
|
|
|
+
|
|
|
+func NewUploadAgentDataLogic(ctx context.Context, svcCtx *svc.ServiceContext) *UploadAgentDataLogic {
|
|
|
+ return &UploadAgentDataLogic{
|
|
|
+ Logger: logx.WithContext(ctx),
|
|
|
+ ctx: ctx,
|
|
|
+ svcCtx: svcCtx}
|
|
|
+}
|
|
|
+
|
|
|
+func (l *UploadAgentDataLogic) UploadAgentData(req *types.UploadDataReq, file multipart.File, agentId uint64) (*types.BaseDataInfo, error) {
|
|
|
+ var count uint64
|
|
|
+
|
|
|
+ reader := csv.NewReader(file)
|
|
|
+ records, err := reader.ReadAll()
|
|
|
+ if err != nil {
|
|
|
+ return nil, err
|
|
|
+ }
|
|
|
+
|
|
|
+ fmt.Printf("req=%+v", *req)
|
|
|
+
|
|
|
+ agent, err := l.svcCtx.DB.Agent.Query().Where(agentModel.ID(agentId)).Only(l.ctx)
|
|
|
+ if err != nil {
|
|
|
+ return nil, dberrorhandler.DefaultEntError(l.Logger, err, req)
|
|
|
+ }
|
|
|
+
|
|
|
+ //TODO
|
|
|
+ token := agent.CollectionID
|
|
|
+
|
|
|
+ var params fastgpt.CreateBulkDataReq
|
|
|
+ params.CollectionID = agent.CollectionID
|
|
|
+ params.TrainingMode = "chunk"
|
|
|
+
|
|
|
+ qas := make([]fastgpt.DataQuestion, 0, 100)
|
|
|
+ for idx, record := range records {
|
|
|
+ //l.Logger.Infof("idx=>%d 0=>%s 1=>%s\n", idx, record[0], record[1])
|
|
|
+ qas = append(qas, fastgpt.DataQuestion{
|
|
|
+ Q: record[0],
|
|
|
+ A: record[1],
|
|
|
+ })
|
|
|
+ if idx%100 == 0 && len(qas) > 0 {
|
|
|
+ params.Data = qas
|
|
|
+ //response, err := fastgpt.DiyCreateBulkData(token, ¶ms)
|
|
|
+ response, err := fastgpt.CreateBulkData(¶ms)
|
|
|
+ if err != nil {
|
|
|
+ l.Logger.Errorf("batch insert data to fastgpt failed. collection=%s error=%s", agent.CollectionID, err.Error())
|
|
|
+ return nil, err
|
|
|
+ }
|
|
|
+ count += response.Data.InsertLen
|
|
|
+ qas = make([]fastgpt.DataQuestion, 0, 100)
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ if len(qas) > 0 {
|
|
|
+ params.Data = qas
|
|
|
+ response, err := fastgpt.DiyCreateBulkData(token, ¶ms)
|
|
|
+ //response, err := fastgpt.CreateBulkData(¶ms)
|
|
|
+ if err != nil {
|
|
|
+ l.Logger.Errorf("batch insert data to fastgpt failed. collection=%s error=%s", agent.CollectionID, err.Error())
|
|
|
+ return nil, err
|
|
|
+ }
|
|
|
+ count += response.Data.InsertLen
|
|
|
+ }
|
|
|
+
|
|
|
+ resp := &types.BaseDataInfo{}
|
|
|
+ resp.Code = 0
|
|
|
+ resp.Msg = errormsg.Success
|
|
|
+ resp.Data = fmt.Sprintf("upload %d rows", count)
|
|
|
+ return resp, nil
|
|
|
+}
|