diff --git a/app/http/controller/web/rag_controller.go b/app/http/controller/web/rag_controller.go index 879c2e5..7b51562 100644 --- a/app/http/controller/web/rag_controller.go +++ b/app/http/controller/web/rag_controller.go @@ -6,6 +6,7 @@ import ( "catface/app/model" "catface/app/model_es" "catface/app/service/nlp" + "catface/app/service/rag/curd" "catface/app/utils/llm_factory" "catface/app/utils/micro_service" "catface/app/utils/response" @@ -138,15 +139,6 @@ func (r *Rag) ChatWebSocket(context *gin.Context) { response.Fail(context, errcode.ErrWebsocketUpgradeFail, errcode.ErrMsg[errcode.ErrWebsocketUpgradeFail], "") return } - defer func() { // UPDATE 临时"持久化"方案,之后考虑结合 jwt 维护的 token 处理。 - tokenMsg := model.CreateNlpWebSocketResult("token", token) - - err := ws.WriteMessage(websocket.TextMessage, tokenMsg.JsonMarshal()) - if err != nil { - variable.ZapLog.Error("Failed to send token message via WebSocket", zap.Error(err)) - } - ws.Close() - }() // 0-2. 测试 Python 微服务是否启动 if !micro_service.TestLinkPythonService() { @@ -181,8 +173,8 @@ func (r *Rag) ChatWebSocket(context *gin.Context) { return } - // 2. ES TopK - docs, err := model_es.CreateDocESFactory().TopK(embedding, 1) + // 2. ES TopK // TODO 这里需要特化选取不同知识库的文档;目前是依靠显式的路由。 + docs, err := curd.CreateDocCurdFactory().TopK(embedding, 1) if err != nil || len(docs) == 0 { variable.ZapLog.Error("ES TopK error", zap.Error(err)) @@ -194,6 +186,24 @@ func (r *Rag) ChatWebSocket(context *gin.Context) { return } + // STAGE websocket 的 defer 关闭函数,但是需要 ES 拿到的 doc—id + defer func() { // UPDATE 临时"持久化"方案,之后考虑结合 jwt 维护的 token 处理。 + // 0. 传递参考资料的信息 + docMsg := model.CreateNlpWebSocketResult(docs[0].Type, docs) + err := ws.WriteMessage(websocket.TextMessage, docMsg.JsonMarshal()) + if err != nil { + variable.ZapLog.Error("Failed to send doc message via WebSocket", zap.Error(err)) + } + + // 1. 传递 token 信息; // UPDATE 临时方案 + tokenMsg := model.CreateNlpWebSocketResult("token", token) + err = ws.WriteMessage(websocket.TextMessage, tokenMsg.JsonMarshal()) + if err != nil { + variable.ZapLog.Error("Failed to send token message via WebSocket", zap.Error(err)) + } + ws.Close() + }() + // 3. closeEventFromVue := context.Request.Context().Done() // 接收前端传来的中断信号。 ch := make(chan string) // TIP 建立通道。 diff --git a/app/model/doc.go b/app/model/doc.go index 7e6d679..8061ca2 100644 --- a/app/model/doc.go +++ b/app/model/doc.go @@ -35,3 +35,30 @@ func (d *Doc) InsertDocumentData(c *gin.Context) (int64, bool) { } return 0, false } + +func (d *Doc) ShowById(id int64, attrs ...string) *Doc { + var temp Doc + db := d.DB.Table(d.TableName()).Where("id = ?", id) + + if len(attrs) > 0 { + db.Select(attrs) + } + + err := db.First(&temp) + if err.Error != nil { + variable.ZapLog.Error("Doc ShowById Error", zap.Error(err.Error)) + } + return &temp +} + +func (d *Doc) ShowByIds(ids []int64, attrs ...string) (temp []Doc) { + db := d.DB.Table(d.TableName()).Where("id in (?)", ids) + if len(attrs) > 0 { + db.Select(attrs) + } + err := db.Find(&temp) + if err.Error != nil { + variable.ZapLog.Error("Doc ShowByIds Error", zap.Error(err.Error)) + } + return +} diff --git a/app/model_res/doc.go b/app/model_res/doc.go new file mode 100644 index 0000000..565b0a1 --- /dev/null +++ b/app/model_res/doc.go @@ -0,0 +1,23 @@ +package model_res + +import ( + "catface/app/model" + "catface/app/model_es" +) + +// BUG 存在 依賴循環 +func NewDocResult(doc *model.Doc, doc_es *model_es.Doc) *DocResult { + return &DocResult{ + Type: "doc", + Id: doc.Id, + Name: doc.Name, + Content: doc_es.Content, + } +} + +type DocResult struct { + Type string `json:"type"` + Id int64 `json:"id"` + Name string `json:"name"` + Content string `json:"content"` +} diff --git a/app/service/rag/curd/doc_curd.go b/app/service/rag/curd/doc_curd.go new file mode 100644 index 0000000..1c3be36 --- /dev/null +++ b/app/service/rag/curd/doc_curd.go @@ -0,0 +1,44 @@ +package curd + +import ( + "catface/app/model" + "catface/app/model_es" + "catface/app/model_res" +) + +func CreateDocCurdFactory() *DocCurd { + return &DocCurd{ + doc: model.CreateDocFactory(""), + doc_es: model_es.CreateDocESFactory()} +} + +type DocCurd struct { // 组合方法的使用 + doc *model.Doc + doc_es *model_es.Doc +} + +func (d *DocCurd) TopK(embedding []float64, k int) (temp []model_res.DocResult, err error) { + // ES:TopK + docs_es, err := d.doc_es.TopK(embedding, k) + if err != nil { + return + } + + // MySQL:补充基本信息 + var ids []int64 + for _, doc := range docs_es { + ids = append(ids, doc.Id) + } + docs := d.doc.ShowByIds(ids, "id", "name") + + // 装载 + for _, doc := range docs { + for _, doc_es := range docs_es { + if doc.Id == doc_es.Id { + temp = append(temp, *model_res.NewDocResult(&doc, &doc_es)) + } + } + } + + return +} diff --git a/config/prompts.yml b/config/prompts.yml index 536f0c8..05ce48f 100644 --- a/config/prompts.yml +++ b/config/prompts.yml @@ -12,10 +12,10 @@ Prompt: Title: "请根据以下长文本生成一个合适的标题,不需要书名号,长度10字内:" - KnoledgeRAG: "使用以上下文来回答用户的问题,如果无法回答,请回答知识库中未找到符合的资料,我不知道。 + KnoledgeRAG: "使用以知识库来回答用户的问题,如果无法回答,请回答知识库中未找到符合的资料,我不知道。 问题: {question} - 可参考的上下文: + 可参考的知识库: ··· {context} ··· - 如果给定的上下文无法让你做出回答,请回答知识库中未找到符合的资料,我不知道。" \ No newline at end of file + 如果给定的知识库无法让你做出回答,请回答知识库中未找到符合的资料,我不知道。" \ No newline at end of file