1

2026-08-05 10:28:44 +08:00
commit 0bacf2e0a4
103 changed files with 9415 additions and 0 deletions
@@ -0,0 +1,32 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+)
+
+type chunk struct{}
+
+var Chunk = &chunk{}
+
+func (c *chunk) List(ctx context.Context, req *dto.ListChunkReq) (*dto.ListChunkRes, error) {
+	list, total, err := service.ChunkService.List(ctx, req.DocumentId, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListChunkRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
+
+func (c *chunk) Update(ctx context.Context, req *dto.UpdateChunkReq) (*dto.UpdateChunkRes, error) {
+	if err := service.ChunkService.Update(ctx, req.Id, req.Content); err != nil {
+		return nil, err
+	}
+	return &dto.UpdateChunkRes{}, nil
+}
@@ -0,0 +1,40 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/model/entity"
+	"rag-local/kb/service"
+)
+
+type conversation struct{}
+
+var Conversation = &conversation{}
+
+func (c *conversation) List(ctx context.Context, req *dto.ListConversationReq) (*dto.ListConversationRes, error) {
+	list, err := service.ConversationService.List(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListConversationRes{List: list}, nil
+}
+
+func (c *conversation) Save(ctx context.Context, req *dto.SaveConversationReq) (*dto.SaveConversationRes, error) {
+	id, err := service.ConversationService.Save(ctx, &entity.Conversation{
+		Id:        req.Id,
+		DatasetId: req.DatasetId,
+		Title:     req.Title,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &dto.SaveConversationRes{Id: id}, nil
+}
+
+func (c *conversation) Delete(ctx context.Context, req *dto.DeleteConversationReq) (*dto.DeleteConversationRes, error) {
+	if err := service.ConversationService.Delete(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.DeleteConversationRes{}, nil
+}
@@ -0,0 +1,42 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/model/entity"
+	"rag-local/kb/service"
+)
+
+type dataset struct{}
+
+var Dataset = &dataset{}
+
+func (c *dataset) List(ctx context.Context, _ *dto.ListDatasetReq) (*dto.ListDatasetRes, error) {
+	list, err := service.DatasetService.List(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListDatasetRes{List: list}, nil
+}
+
+func (c *dataset) Save(ctx context.Context, req *dto.SaveDatasetReq) (*dto.SaveDatasetRes, error) {
+	id, err := service.DatasetService.Save(ctx, &entity.Dataset{
+		Id:             req.Id,
+		Name:           req.Name,
+		Description:    req.Description,
+		EmbeddingCfgId: req.EmbeddingCfgId,
+		Status:         1,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &dto.SaveDatasetRes{Id: id}, nil
+}
+
+func (c *dataset) Delete(ctx context.Context, req *dto.DeleteDatasetReq) (*dto.DeleteDatasetRes, error) {
+	if err := service.DatasetService.Delete(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.DeleteDatasetRes{}, nil
+}
@@ -0,0 +1,62 @@
+package controller
+
+import (
+	"context"
+	"io"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+
+	"github.com/gogf/gf/v2/errors/gerror"
+)
+
+type document struct{}
+
+var Document = &document{}
+
+func (c *document) Upload(ctx context.Context, req *dto.UploadDocumentReq) (*dto.UploadDocumentRes, error) {
+	if req.File == nil {
+		return nil, gerror.New("请选择文件")
+	}
+	f, err := req.File.Open()
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = f.Close() }()
+	data, err := io.ReadAll(f)
+	if err != nil {
+		return nil, err
+	}
+	doc, err := service.DocumentService.Upload(ctx, req.DatasetId, req.File.Filename, data)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.UploadDocumentRes{Id: doc.Id}, nil
+}
+
+func (c *document) List(ctx context.Context, req *dto.ListDocumentReq) (*dto.ListDocumentRes, error) {
+	list, total, err := service.DocumentService.List(ctx, req.DatasetId, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListDocumentRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
+
+func (c *document) Delete(ctx context.Context, req *dto.DeleteDocumentReq) (*dto.DeleteDocumentRes, error) {
+	if err := service.DocumentService.Delete(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.DeleteDocumentRes{}, nil
+}
+
+func (c *document) Reembed(ctx context.Context, req *dto.ReembedDocumentReq) (*dto.ReembedDocumentRes, error) {
+	if err := service.DocumentService.Reembed(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.ReembedDocumentRes{}, nil
+}
@@ -0,0 +1,25 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+)
+
+type kgEntity struct{}
+
+var KgEntity = &kgEntity{}
+
+func (c *kgEntity) List(ctx context.Context, req *dto.ListKgEntityReq) (*dto.ListKgEntityRes, error) {
+	list, total, err := service.KgEntityService.List(ctx, req.DatasetId, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListKgEntityRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
@@ -0,0 +1,25 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+)
+
+type kgRelation struct{}
+
+var KgRelation = &kgRelation{}
+
+func (c *kgRelation) List(ctx context.Context, req *dto.ListKgRelationReq) (*dto.ListKgRelationRes, error) {
+	list, total, err := service.KgRelationService.List(ctx, req.DatasetId, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListKgRelationRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
@@ -0,0 +1,78 @@
+package controller
+
+import (
+	"context"
+	"encoding/json"
+	"time"
+
+	"rag-local/kb/model/domain"
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+
+	"github.com/gogf/gf/v2/frame/g"
+)
+
+type message struct{}
+
+var Message = &message{}
+
+func (c *message) List(ctx context.Context, req *dto.ListMessageReq) (*dto.ListMessageRes, error) {
+	list, err := service.MessageService.List(ctx, req.ConversationId)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListMessageRes{List: list}, nil
+}
+
+// Chat RAG 问答 SSE 流式：citations → delta* → done / error
+func (c *message) Chat(ctx context.Context, req *dto.ChatReq) (*dto.ChatRes, error) {
+	r := g.RequestFromCtx(ctx)
+	r.Response.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
+	r.Response.Header().Set("Cache-Control", "no-cache")
+	r.Response.Header().Set("X-Accel-Buffering", "no")
+	r.Response.Header().Set("Connection", "keep-alive")
+
+	send := func(event string, data any) bool {
+		buf, err := json.Marshal(data)
+		if err != nil {
+			return false
+		}
+		if event != "" {
+			r.Response.Write("event: " + event + "\n")
+		}
+		r.Response.Write("data: " + string(buf) + "\n\n")
+		r.Response.Flush()
+		return true
+	}
+
+	// 心跳：长回答期间保持连接，避免网关断流
+	stopHeartbeat := make(chan struct{})
+	defer close(stopHeartbeat)
+	go func() {
+		ticker := time.NewTicker(15 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-stopHeartbeat:
+				return
+			case <-ticker.C:
+				r.Response.Write(": ping\n\n")
+				r.Response.Flush()
+			}
+		}
+	}()
+
+	_, _, _, err := service.MessageService.Chat(ctx, req.ConversationId, req.DatasetId, req.Question,
+		func(citations []domain.Citation, conversationId int64) {
+			send("citations", map[string]any{"conversation_id": conversationId, "citations": citations})
+		},
+		func(delta string) {
+			send("delta", map[string]string{"content": delta})
+		})
+	if err != nil {
+		send("error", map[string]string{"message": err.Error()})
+		return nil, nil
+	}
+	send("done", map[string]string{"status": "ok"})
+	return nil, nil
+}
@@ -0,0 +1,49 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/model/entity"
+	"rag-local/kb/service"
+)
+
+type modelConfig struct{}
+
+var ModelConfig = new(modelConfig)
+
+func (c *modelConfig) List(ctx context.Context, req *dto.ListModelConfigReq) (res *dto.ListModelConfigRes, err error) {
+	list, err := service.ModelConfigService.List(ctx, req.ModelType)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListModelConfigRes{List: list}, nil
+}
+
+func (c *modelConfig) Save(ctx context.Context, req *dto.SaveModelConfigReq) (res *dto.SaveModelConfigRes, err error) {
+	id, err := service.ModelConfigService.Save(ctx, &entity.ModelConfig{
+		Id:          req.Id,
+		Name:        req.Name,
+		ModelType:   req.ModelType,
+		ModelName:   req.ModelName,
+		EndpointUrl: req.EndpointUrl,
+		ApiKey:      req.ApiKey,
+		Dimension:   req.Dimension,
+		Extra:       req.Extra,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &dto.SaveModelConfigRes{Id: id}, nil
+}
+
+func (c *modelConfig) Delete(ctx context.Context, req *dto.DeleteModelConfigReq) (res *dto.DeleteModelConfigRes, err error) {
+	return nil, service.ModelConfigService.Delete(ctx, req.Id)
+}
+
+func (c *modelConfig) Test(ctx context.Context, req *dto.TestModelConfigReq) (res *dto.TestModelConfigRes, err error) {
+	if err := service.ModelConfigService.Test(ctx, req.Id); err != nil {
+		return &dto.TestModelConfigRes{Ok: false, Msg: err.Error()}, nil
+	}
+	return &dto.TestModelConfigRes{Ok: true, Msg: "连接正常"}, nil
+}
@@ -0,0 +1,32 @@
+package controller
+
+import (
+	"context"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+)
+
+type parseTask struct{}
+
+var ParseTask = &parseTask{}
+
+func (c *parseTask) List(ctx context.Context, req *dto.ListParseTaskReq) (*dto.ListParseTaskRes, error) {
+	list, total, err := service.ParseTaskService.List(ctx, req.Page, req.PageSize)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.ListParseTaskRes{
+		List:     list,
+		Total:    total,
+		Page:     req.Page,
+		PageSize: req.PageSize,
+	}, nil
+}
+
+func (c *parseTask) Retry(ctx context.Context, req *dto.RetryParseTaskReq) (*dto.RetryParseTaskRes, error) {
+	if err := service.ParseTaskService.Retry(ctx, req.Id); err != nil {
+		return nil, err
+	}
+	return &dto.RetryParseTaskRes{}, nil
+}
@@ -0,0 +1,53 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+
+	"rag-local/kb/model/dto"
+	"rag-local/kb/service"
+)
+
+type systemConfig struct{}
+
+var SystemConfig = new(systemConfig)
+
+func (c *systemConfig) Login(ctx context.Context, req *dto.LoginReq) (res *dto.LoginRes, err error) {
+	token, err := service.SystemConfigService.Login(ctx, req.Token)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.LoginRes{Token: token}, nil
+}
+
+func (c *systemConfig) Get(ctx context.Context, req *dto.GetSystemConfigReq) (res *dto.GetSystemConfigRes, err error) {
+	chatModel, dataset, err := service.SystemConfigService.GetSettings(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.GetSystemConfigRes{DefaultChatModel: chatModel, DefaultDataset: dataset}, nil
+}
+
+func (c *systemConfig) Update(ctx context.Context, req *dto.UpdateSystemConfigReq) (res *dto.UpdateSystemConfigRes, err error) {
+	return nil, service.SystemConfigService.UpdateSettings(ctx, req.DefaultChatModel, req.DefaultDataset)
+}
+
+func (c *systemConfig) GetToken(ctx context.Context, req *dto.GetTokenReq) (*dto.GetTokenRes, error) {
+	token, err := service.SystemConfigService.GetToken(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &dto.GetTokenRes{Token: token}, nil
+}
+
+func (c *systemConfig) RegenerateToken(ctx context.Context, req *dto.RegenerateTokenReq) (res *dto.RegenerateTokenRes, err error) {
+	token, err := service.SystemConfigService.RegenerateToken(ctx)
+	if err != nil {
+		return nil, err
+	}
+	fmt.Printf("\n============================================\n")
+	fmt.Printf("访问令牌(登录用)已更新: %s\n", token)
+	fmt.Printf("旧令牌已失效，请在登录页重新输入\n")
+	fmt.Printf("============================================\n\n")
+	return &dto.RegenerateTokenRes{Token: token}, nil
+}