Merge remote-tracking branch 'origin/dev' into dev
# Conflicts: # service/document.go
This commit is contained in:
@@ -14,91 +14,122 @@ import (
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
)
|
||||
|
||||
var globalChatModel *qwen.ChatModel
|
||||
const (
|
||||
MaxHistoryTurns = 5 // 最大历史轮数
|
||||
)
|
||||
|
||||
var (
|
||||
globalChatModel *qwen.ChatModel
|
||||
ragPromptTemplate prompt.ChatTemplate // EINO 官方模板
|
||||
)
|
||||
|
||||
func init() {
|
||||
ctx := context.Background()
|
||||
// 初始化大模型
|
||||
if err := initChatModel(ctx); err != nil {
|
||||
glog.Errorf(ctx, "初始化大模型失败: %v", err)
|
||||
}
|
||||
// 初始化 EINO 提示词模板
|
||||
initRAGPromptTemplate()
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// 初始化通义千问
|
||||
func initChatModel(ctx context.Context) error {
|
||||
if globalChatModel != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
apiKey := g.Cfg().MustGet(ctx, "eino.chatmodel.apiKey").String()
|
||||
model := g.Cfg().MustGet(ctx, "eino.chatmodel.model").String()
|
||||
|
||||
var err error
|
||||
globalChatModel, err = qwen.NewChatModel(ctx, &qwen.ChatModelConfig{
|
||||
cm, err := qwen.NewChatModel(ctx, &qwen.ChatModelConfig{
|
||||
APIKey: apiKey,
|
||||
Model: model,
|
||||
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
Timeout: 60 * 1e9,
|
||||
Temperature: gconv.PtrFloat32(0.7), // 客服最佳
|
||||
MaxTokens: gconv.PtrInt(1024), // 最长回答
|
||||
TopP: gconv.PtrFloat32(1.0),
|
||||
})
|
||||
if err != nil {
|
||||
glog.Errorf(ctx, "初始化大模型失败: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return
|
||||
globalChatModel = cm
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewChatModel 只处理逻辑,不复用创建模型
|
||||
func NewChatModel(ctx context.Context, content string, docs []*schema.Document) (replyMsg *schema.Message, sources []string, err error) {
|
||||
// 1. 构建参考知识
|
||||
knowledge, sources := buildKnowledgeAndSources(docs)
|
||||
|
||||
// 2. 构建提示词
|
||||
msgs, err := buildPromptMessages(ctx, knowledge, content)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 🔥 直接使用全局单例,不重复创建
|
||||
replyMsg, err = streamGenerateAnswer(ctx, globalChatModel, msgs)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// buildKnowledgeAndSources 拼接参考知识 + 提取文档来源
|
||||
func buildKnowledgeAndSources(docs []*schema.Document) (string, []string) {
|
||||
var knowledge string
|
||||
var sources []string
|
||||
|
||||
for i, doc := range docs {
|
||||
knowledge += fmt.Sprintf("[参考%d] %s\n", i+1, doc.Content)
|
||||
|
||||
// 提取 document_id
|
||||
if docID, ok := doc.MetaData["document_id"].(int64); ok && docID > 0 {
|
||||
sources = append(sources, gconv.String(docID))
|
||||
}
|
||||
}
|
||||
return knowledge, sources
|
||||
}
|
||||
|
||||
// buildPromptMessages 构建提示词模板
|
||||
func buildPromptMessages(ctx context.Context, knowledge string, question string) (msgs []*schema.Message, err error) {
|
||||
promptTpl := prompt.FromMessages(
|
||||
// 初始化 EINO 官方提示词模板(最关键!)
|
||||
func initRAGPromptTemplate() {
|
||||
ragPromptTemplate = prompt.FromMessages(
|
||||
schema.FString,
|
||||
// 系统提示(带参考知识)
|
||||
&schema.Message{
|
||||
Role: schema.System,
|
||||
// Content: `你是专业的客服助手,语气友好。
|
||||
//如果参考知识中有相关信息,请优先依据参考知识回答。
|
||||
//如果没有相关信息,就正常回答,不要说无法回答。
|
||||
//
|
||||
//参考知识:
|
||||
//{knowledge}`,
|
||||
Content: `你是专业的客服助手,语气友好。
|
||||
请根据参考知识回答用户问题,无法回答则说:抱歉,我暂时无法回答这个问题。
|
||||
|
||||
参考知识:
|
||||
{knowledge}`,
|
||||
Content: `你是专业客服,语气友好简洁。
|
||||
请严格依据参考知识回答,不知道就说:抱歉,我暂时无法回答这个问题。
|
||||
|
||||
参考知识:
|
||||
{knowledge}`,
|
||||
},
|
||||
// 用户问题
|
||||
&schema.Message{
|
||||
Role: schema.User,
|
||||
Content: "{question}",
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
return promptTpl.Format(ctx, map[string]any{
|
||||
// NewChatModel 只处理逻辑,不复用创建模型
|
||||
func NewChatModel(ctx context.Context, question string, docs []*schema.Document, history []*schema.Message) (replyMsg *schema.Message, err error) {
|
||||
// 1. 构建参考知识
|
||||
knowledge := buildKnowledgeAndSources(docs)
|
||||
// 2. 历史精简
|
||||
history = limitHistory(history)
|
||||
// 3. ✅ EINO 官方模板格式化(超级干净)
|
||||
msgs, err := ragPromptTemplate.Format(ctx, map[string]any{
|
||||
"knowledge": knowledge,
|
||||
"question": question,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 4. 历史插入到模板消息中间(标准EINO用法)
|
||||
if len(history) > 0 {
|
||||
msgs = append(msgs[:1], append(history, msgs[1:]...)...)
|
||||
}
|
||||
// 5. 🔥 直接使用全局单例,不重复创建
|
||||
replyMsg, err = streamGenerateAnswer(ctx, globalChatModel, msgs)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func limitHistory(history []*schema.Message) []*schema.Message {
|
||||
valid := make([]*schema.Message, 0, len(history))
|
||||
for _, m := range history {
|
||||
if m.Role == schema.User || m.Role == schema.Assistant {
|
||||
valid = append(valid, m)
|
||||
}
|
||||
}
|
||||
|
||||
keep := 2 * MaxHistoryTurns
|
||||
if len(valid) > keep {
|
||||
valid = valid[len(valid)-keep:]
|
||||
}
|
||||
return valid
|
||||
}
|
||||
|
||||
// buildKnowledgeAndSources 拼接参考知识
|
||||
func buildKnowledgeAndSources(docs []*schema.Document) string {
|
||||
var knowledge string
|
||||
|
||||
for i, doc := range docs {
|
||||
knowledge += fmt.Sprintf("[参考%d] %s\n", i+1, doc.Content)
|
||||
|
||||
}
|
||||
return knowledge
|
||||
}
|
||||
|
||||
// streamGenerateAnswer 流式生成
|
||||
|
||||
@@ -5,11 +5,14 @@ import (
|
||||
"errors"
|
||||
"rag/dao"
|
||||
"sort"
|
||||
"time"
|
||||
|
||||
"github.com/cloudwego/eino/callbacks"
|
||||
"github.com/cloudwego/eino/components/embedding"
|
||||
"github.com/cloudwego/eino/components/retriever"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
"github.com/gogf/gf/v2/os/grpool"
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
"github.com/pgvector/pgvector-go"
|
||||
)
|
||||
@@ -53,43 +56,139 @@ func (r *PGVectorRetriever) Retrieve(ctx context.Context, query string, opts ...
|
||||
}
|
||||
options = retriever.GetCommonOptions(options, opts...)
|
||||
|
||||
// 安全保护:防止 nil 指针 panic
|
||||
topK := 10
|
||||
if options.TopK != nil {
|
||||
topK = *options.TopK
|
||||
}
|
||||
|
||||
ctx = callbacks.OnStart(ctx, &retriever.CallbackInput{
|
||||
Query: query,
|
||||
TopK: *options.TopK,
|
||||
})
|
||||
|
||||
// ==========================================
|
||||
// 🔥 双路检索:向量 + 全文
|
||||
// 🔥 优化版:grpool 并行双路检索(安全、健壮、无泄漏)
|
||||
// ==========================================
|
||||
docsVector, err := r.doRetrieveVector(ctx, query, options)
|
||||
if err != nil {
|
||||
callbacks.OnError(ctx, err)
|
||||
var (
|
||||
docsVector []*schema.Document
|
||||
docsFulltext []*schema.Document
|
||||
errVector error
|
||||
errFulltext error
|
||||
|
||||
// 缓冲通道=2,确保无死锁等待
|
||||
done = make(chan struct{}, 2)
|
||||
)
|
||||
|
||||
// 上下文:超时 + 可取消双保障(建议5s超时,根据业务调整)
|
||||
taskCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// 封装并行任务函数,消除重复代码
|
||||
runTask := func(task func() error, errTarget *error) {
|
||||
defer func() {
|
||||
// 任务结束必发信号,确保通道不阻塞
|
||||
done <- struct{}{}
|
||||
}()
|
||||
|
||||
// 捕获 panic + 执行业务逻辑
|
||||
g.TryCatch(taskCtx, func(ctx context.Context) {
|
||||
*errTarget = task()
|
||||
}, func(ctx context.Context, panicErr error) {
|
||||
*errTarget = panicErr
|
||||
})
|
||||
|
||||
// 任务失败:立即取消另一个任务(快速失败)
|
||||
if *errTarget != nil {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------
|
||||
// 并行提交两个检索任务
|
||||
// ----------------------
|
||||
// 任务1:向量检索
|
||||
grpool.Add(taskCtx, func(ctx context.Context) {
|
||||
runTask(func() error {
|
||||
docsVector, errVector = r.doRetrieveVector(ctx, query, options)
|
||||
return errVector
|
||||
}, &errVector)
|
||||
})
|
||||
|
||||
// 任务2:全文检索
|
||||
grpool.Add(taskCtx, func(ctx context.Context) {
|
||||
runTask(func() error {
|
||||
docsFulltext, errFulltext = r.doRetrieveMeilisearch(ctx, query, options)
|
||||
return errFulltext
|
||||
}, &errFulltext)
|
||||
})
|
||||
|
||||
// ----------------------
|
||||
// 安全等待所有任务完成
|
||||
// ----------------------
|
||||
<-done
|
||||
<-done
|
||||
|
||||
// ----------------------
|
||||
// 统一错误处理
|
||||
// ----------------------
|
||||
// 用 errors.Join 合并所有错误,不丢失信息
|
||||
if err := errors.Join(errVector, errFulltext); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
docsFulltext, err := r.doRetrieveMeilisearch(ctx, query, options)
|
||||
if err != nil {
|
||||
callbacks.OnError(ctx, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 合并 + 去重
|
||||
// 合并 + 智能去重(保留最优分数)
|
||||
docs := mergeAndDeduplicate(docsVector, docsFulltext)
|
||||
|
||||
// 排序(distance 越小越靠前)
|
||||
// 排序:向量优先,同类型按距离升序
|
||||
sort.Slice(docs, func(i, j int) bool {
|
||||
//byI, okI := docs[i].MetaData["retrieve_by"].(string)
|
||||
//byJ, okJ := docs[j].MetaData["retrieve_by"].(string)
|
||||
//
|
||||
//// 有类型标记的优先
|
||||
//if okI && !okJ {
|
||||
// return true
|
||||
//}
|
||||
//if !okI && okJ {
|
||||
// return false
|
||||
//}
|
||||
//
|
||||
//// 向量永远排前面
|
||||
//if byI == "vector" && byJ == "fulltext" {
|
||||
// return true
|
||||
//}
|
||||
//if byI == "fulltext" && byJ == "vector" {
|
||||
// return false
|
||||
//}
|
||||
|
||||
// 同类型按 distance 升序(越小越相似)
|
||||
d1 := gconv.Float64(docs[i].MetaData["distance"])
|
||||
d2 := gconv.Float64(docs[j].MetaData["distance"])
|
||||
return d1 < d2
|
||||
})
|
||||
|
||||
// 最多保留 topK
|
||||
if len(docs) > *options.TopK {
|
||||
docs = docs[:*options.TopK]
|
||||
// 在Retrieve方法末尾,增加相关性校验
|
||||
validDocs := make([]*schema.Document, 0)
|
||||
for i, d := range docs {
|
||||
// 过滤distance过大的垃圾结果(比如distance>0.8的直接丢弃)
|
||||
if gconv.Float64(docs[i].MetaData["distance"]) < 0.8 {
|
||||
validDocs = append(validDocs, d)
|
||||
}
|
||||
}
|
||||
|
||||
callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: docs})
|
||||
return docs, nil
|
||||
// 如果没有有效结果,返回空,让LLM回答「暂无相关信息」
|
||||
if len(validDocs) == 0 {
|
||||
callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: validDocs})
|
||||
return validDocs, nil
|
||||
}
|
||||
|
||||
// 最多保留 topK
|
||||
if len(validDocs) > topK {
|
||||
validDocs = validDocs[:topK]
|
||||
}
|
||||
|
||||
callbacks.OnEnd(ctx, &retriever.CallbackOutput{Docs: validDocs})
|
||||
return validDocs, nil
|
||||
}
|
||||
|
||||
// ==========================================
|
||||
@@ -105,7 +204,10 @@ func (r *PGVectorRetriever) doRetrieveVector(ctx context.Context, query string,
|
||||
}
|
||||
|
||||
queryVec := pgvector.NewVector(gconv.Float32s(vectors[0]))
|
||||
topK := *opts.TopK
|
||||
topK := 10
|
||||
if opts.TopK != nil {
|
||||
topK = *opts.TopK
|
||||
}
|
||||
datasetIds := gconv.Int64s(opts.DSLInfo["dataset_ids"])
|
||||
|
||||
rows, err := dao.DocumentChunk.GetAllByVector(ctx, datasetIds, queryVec, topK)
|
||||
@@ -144,13 +246,16 @@ func (r *PGVectorRetriever) doRetrieveMeilisearch(ctx context.Context, query str
|
||||
|
||||
docs := make([]*schema.Document, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
score := gconv.Float64(row["_rankingScore"])
|
||||
distance := score
|
||||
|
||||
docs = append(docs, &schema.Document{
|
||||
ID: gconv.String(row["id"]),
|
||||
Content: gconv.String(row["content"]),
|
||||
MetaData: map[string]any{
|
||||
"dataset_id": gconv.Int64(row["dataset_id"]),
|
||||
"document_id": gconv.Int64(row["document_id"]),
|
||||
"distance": 0.1, // 全文结果给高分
|
||||
"distance": distance,
|
||||
"retrieve_by": "fulltext",
|
||||
},
|
||||
})
|
||||
@@ -159,18 +264,26 @@ func (r *PGVectorRetriever) doRetrieveMeilisearch(ctx context.Context, query str
|
||||
}
|
||||
|
||||
// ==========================================
|
||||
// 合并去重
|
||||
// 合并去重(智能版:两路都命中时,保留向量结果 + 全文标记)
|
||||
// ==========================================
|
||||
func mergeAndDeduplicate(vecDocs, fullDocs []*schema.Document) []*schema.Document {
|
||||
idMap := make(map[string]*schema.Document)
|
||||
|
||||
// 先存入向量结果
|
||||
for _, d := range vecDocs {
|
||||
idMap[d.ID] = d
|
||||
}
|
||||
|
||||
// 再处理全文:不存在则添加;存在则标记“双路命中”,不覆盖向量分数
|
||||
for _, d := range fullDocs {
|
||||
if _, exists := idMap[d.ID]; !exists {
|
||||
if existDoc, ok := idMap[d.ID]; ok {
|
||||
// 标记同时被向量和全文检索到
|
||||
existDoc.MetaData["retrieve_by"] = "both"
|
||||
} else {
|
||||
idMap[d.ID] = d
|
||||
}
|
||||
}
|
||||
|
||||
merged := make([]*schema.Document, 0, len(idMap))
|
||||
for _, d := range idMap {
|
||||
merged = append(merged, d)
|
||||
|
||||
@@ -64,7 +64,7 @@ func (d *documentChunkDao) List(ctx context.Context, req *dto.ListDocumentChunkR
|
||||
func (d *documentChunkDao) GetAllByVector(ctx context.Context, datasetId []int64, queryVec pgvector.Vector, topK int) (list gdb.List, err error) {
|
||||
sql := `
|
||||
SELECT id, content, dataset_id, document_id,
|
||||
vector <-> ? AS distance
|
||||
vector <=> ? AS distance
|
||||
FROM rag_vector_document_chunk
|
||||
WHERE dataset_id IN (?)
|
||||
AND vector IS NOT NULL
|
||||
@@ -84,8 +84,9 @@ func (d *documentChunkDao) GetAllByVector(ctx context.Context, datasetId []int64
|
||||
func (d *documentChunkDao) SearchByKeywords(ctx context.Context, query string, datasetIds []int64, topK int) (list gdb.List, err error) {
|
||||
// 构建 meilisearch 查询参数
|
||||
searchParams := &meilisearch.SearchParams{
|
||||
Query: query,
|
||||
Limit: int64(topK),
|
||||
Query: query,
|
||||
Limit: int64(topK),
|
||||
ShowRankingScore: true,
|
||||
}
|
||||
|
||||
// 构建 datasetIds 过滤条件
|
||||
|
||||
@@ -8,14 +8,18 @@ import (
|
||||
type RAGQueryReq struct {
|
||||
g.Meta `path:"/ragQuery" method:"post" tags:"RAG查询" summary:"执行RAG查询" dc:"执行RAG查询"`
|
||||
|
||||
Content string `json:"content" v:"required#查询内容不能为空" dc:"用户问题"`
|
||||
DatasetIds []int64 `json:"datasetIds" dc:"数据集ID"`
|
||||
TopK int `json:"topK" d:"5" dc:"检索topK,默认5"`
|
||||
Content string `json:"content" v:"required#查询内容不能为空" dc:"用户问题"`
|
||||
DatasetIds []int64 `json:"datasetIds" dc:"数据集ID"`
|
||||
History []*Message `json:"history" dc:"历史对话"`
|
||||
TopK int `json:"topK" d:"5" dc:"检索topK,默认5"`
|
||||
}
|
||||
|
||||
type Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// RAGQueryRes RAG查询响应
|
||||
type RAGQueryRes struct {
|
||||
Answer string `json:"answer" dc:"生成的答案"`
|
||||
DatasetId string `json:"datasetId" dc:"使用的数据集ID"`
|
||||
Sources []string `json:"sources" dc:"参考来源"`
|
||||
Answer string `json:"answer" dc:"生成的答案"`
|
||||
}
|
||||
|
||||
@@ -12,10 +12,10 @@ import (
|
||||
"rag/model/dto"
|
||||
"rag/model/entity"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.com/red-future/common/db/gfdb"
|
||||
"gitea.com/red-future/common/full-text-search/meilisearch"
|
||||
"gitea.com/red-future/common/http"
|
||||
"gitea.com/red-future/common/utils"
|
||||
gmq "github.com/bjang03/gmq/core/gmq"
|
||||
"github.com/bjang03/gmq/mq"
|
||||
@@ -159,12 +159,16 @@ func (s *documentService) Process(ctx context.Context, req *dto.ProcessDocumentR
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
user, err := utils.GetUserInfo(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// ======================
|
||||
// 核心:grpool + g.Try 最佳实践
|
||||
// ======================
|
||||
taskCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
// 使用带超时的background context,避免HTTP请求完成后context被取消
|
||||
taskCtx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
|
||||
taskCtx = context.WithValue(taskCtx, "user", user)
|
||||
// 任务1: SQL 切分文档
|
||||
grpool.Add(taskCtx, func(ctx context.Context) {
|
||||
g.TryCatch(ctx, func(ctx context.Context) {
|
||||
@@ -597,7 +601,7 @@ func (s *documentService) loadDocument(ctx context.Context, doc *entity.Document
|
||||
func (s *documentService) getHistoryData(ctx context.Context, doc *entity.Document, lockKey, contentKey string) (err error) {
|
||||
docsLockKey := fmt.Sprintf(lockKey, doc.DatasetId)
|
||||
success, err := utils.Lock(ctx, docsLockKey, int64(60), func(ctx context.Context) error {
|
||||
// 1. 扫描 Redis 中所有 前缀为 rag_binary:knowledge:xxx:contentHash 的 key
|
||||
// 1. 扫描 Redis 中所有 前缀为 rag:knowledge:xxx:contentHash 的 key
|
||||
pattern := fmt.Sprintf(contentKey, "*")
|
||||
keys, err := g.Redis().Keys(ctx, pattern)
|
||||
if err != nil {
|
||||
@@ -655,23 +659,12 @@ func (s *documentService) getHistoryData(ctx context.Context, doc *entity.Docume
|
||||
|
||||
// getHistoryDataFromHttp 通过 HTTP 接口查询历史数据
|
||||
func (s *documentService) getHistoryDataFromHttp(ctx context.Context, doc *entity.Document) (dictData []*dto.DocumentChunkRPC, err error) {
|
||||
headers := make(map[string]string)
|
||||
if r := g.RequestFromCtx(ctx); r != nil {
|
||||
for k, v := range r.Request.Header {
|
||||
if len(v) > 0 {
|
||||
headers[k] = v[0]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 调用接口获取数据
|
||||
d := &dto.ListDocumentChunkRPC{}
|
||||
if err = http.Get(ctx, "rag_binary-vector/document/chunk/listDocumentChunk", headers, &d,
|
||||
"datasetId", gconv.String(doc.DatasetId),
|
||||
"status", 1); err != nil {
|
||||
return
|
||||
}
|
||||
dictData = d.List
|
||||
res, _, err := dao.DocumentChunk.List(ctx, &dto.ListDocumentChunkReq{
|
||||
DatasetId: doc.DatasetId,
|
||||
Status: gconv.PtrInt8(1),
|
||||
})
|
||||
err = gconv.Struct(res, &dictData)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"rag/model/dto"
|
||||
"rag/model/entity"
|
||||
|
||||
"gitea.com/red-future/common/beans"
|
||||
"github.com/cloudwego/eino/components/indexer"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
@@ -48,7 +49,10 @@ func (s *documentChunkService) DocsChunkMsg(ctx context.Context, msg any) (err e
|
||||
g.Log().Error(ctx, "DocsChunkMsg err:", "msg is empty")
|
||||
return
|
||||
}
|
||||
|
||||
ctx = context.WithValue(ctx, "user", &beans.User{
|
||||
TenantId: gconv.Uint64(docs[0].MetaData[entity.DocumentChunkCol.TenantId]),
|
||||
UserName: gconv.String(docs[0].MetaData[entity.DocumentChunkCol.Creator]),
|
||||
})
|
||||
idx := eino.NewPGVectorIndexer(&eino.PGVectorIndexerOptions{
|
||||
BatchSize: 10,
|
||||
})
|
||||
|
||||
@@ -7,7 +7,9 @@ import (
|
||||
"rag/model/dto"
|
||||
|
||||
"github.com/cloudwego/eino/components/retriever"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"github.com/gogf/gf/v2/os/glog"
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
)
|
||||
|
||||
var RAGQuery = new(ragQueryService)
|
||||
@@ -39,14 +41,20 @@ func (s *ragQueryService) Query(ctx context.Context, req *dto.RAGQueryReq) (*dto
|
||||
return nil, fmt.Errorf("向量检索失败: %w", err)
|
||||
}
|
||||
|
||||
replyMsg, sources, err := eino.NewChatModel(ctx, req.Content, docs)
|
||||
messages := make([]*schema.Message, 0)
|
||||
err = gconv.Struct(req.History, &messages)
|
||||
if err != nil {
|
||||
glog.Errorf(ctx, "转换历史消息失败: %v", err)
|
||||
return nil, fmt.Errorf("转换历史消息失败: %w", err)
|
||||
}
|
||||
|
||||
replyMsg, err := eino.NewChatModel(ctx, req.Content, docs, messages)
|
||||
if err != nil {
|
||||
glog.Errorf(ctx, "向量检索失败: %v", err)
|
||||
return nil, fmt.Errorf("向量检索失败: %w", err)
|
||||
}
|
||||
|
||||
return &dto.RAGQueryRes{
|
||||
Answer: replyMsg.Content,
|
||||
Sources: sources,
|
||||
Answer: replyMsg.Content,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -7,6 +7,8 @@ import (
|
||||
|
||||
"rag/common/task"
|
||||
|
||||
"gitea.com/red-future/common/db/gfdb"
|
||||
"github.com/gogf/gf/v2/database/gdb"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
"github.com/gogf/gf/v2/util/gconv"
|
||||
)
|
||||
@@ -24,17 +26,20 @@ func (s *taskService) WriteTaskProgress(ctx context.Context, req *dto.WriteTaskP
|
||||
g.Log().Errorf(ctx, "查询任务失败: %v", err)
|
||||
return err
|
||||
}
|
||||
taskVO := make([]dto.TaskVO, 0, total)
|
||||
err = gconv.Struct(t, taskVO)
|
||||
if err != nil {
|
||||
g.Log().Errorf(ctx, "转换任务失败: %v", err)
|
||||
return err
|
||||
completed := false
|
||||
if total != 0 {
|
||||
taskVO := make([]*dto.TaskVO, 0, total)
|
||||
err = gconv.Struct(t, &taskVO)
|
||||
if err != nil {
|
||||
g.Log().Errorf(ctx, "转换任务失败: %v", err)
|
||||
return err
|
||||
}
|
||||
taskVO = append(taskVO, &dto.TaskVO{
|
||||
TaskType: req.TaskType,
|
||||
Status: req.Status,
|
||||
})
|
||||
completed = IsAllSubTasksCompleted(taskVO)
|
||||
}
|
||||
taskVO = append(taskVO, dto.TaskVO{
|
||||
TaskType: req.TaskType,
|
||||
Status: req.Status,
|
||||
})
|
||||
completed := IsAllSubTasksCompleted(taskVO)
|
||||
|
||||
// 1. 查询是否已存在该文档的该类型任务
|
||||
existTask, _, err := dao.Task.Get(ctx, &dto.GetTaskReq{
|
||||
@@ -45,44 +50,48 @@ func (s *taskService) WriteTaskProgress(ctx context.Context, req *dto.WriteTaskP
|
||||
g.Log().Errorf(ctx, "查询任务失败: %v", err)
|
||||
return err
|
||||
}
|
||||
err = gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
|
||||
// 2. 如果不存在,则创建新任务
|
||||
if g.IsEmpty(existTask) {
|
||||
createReq := &dto.CreateTaskReq{
|
||||
TaskId: req.TaskId,
|
||||
TaskType: req.TaskType,
|
||||
Status: req.Status,
|
||||
Remark: req.Remark,
|
||||
}
|
||||
_, err = dao.Task.Insert(ctx, createReq)
|
||||
} else {
|
||||
// 3. 如果已存在,则更新任务
|
||||
updateReq := &dto.UpdateTaskReq{
|
||||
Id: existTask[0].Id,
|
||||
Status: req.Status,
|
||||
Remark: req.Remark,
|
||||
}
|
||||
_, err = dao.Task.Update(ctx, updateReq)
|
||||
if err != nil {
|
||||
g.Log().Errorf(ctx, "更新任务失败: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 2. 如果不存在,则创建新任务
|
||||
if g.IsEmpty(existTask) {
|
||||
createReq := &dto.CreateTaskReq{
|
||||
TaskId: req.TaskId,
|
||||
TaskType: req.TaskType,
|
||||
Status: req.Status,
|
||||
Remark: req.Remark,
|
||||
if completed {
|
||||
// 3. 如果已存在,则更新任务
|
||||
_, err = dao.Task.Update(ctx, &dto.UpdateTaskReq{
|
||||
TaskId: req.TaskId,
|
||||
Status: task.TaskStatusCompleted,
|
||||
Remark: "文档解析完成",
|
||||
})
|
||||
}
|
||||
_, err = dao.Task.Insert(ctx, createReq)
|
||||
} else {
|
||||
// 3. 如果已存在,则更新任务
|
||||
updateReq := &dto.UpdateTaskReq{
|
||||
Id: existTask[0].Id,
|
||||
Status: req.Status,
|
||||
Remark: req.Remark,
|
||||
}
|
||||
_, err = dao.Task.Update(ctx, updateReq)
|
||||
if err != nil {
|
||||
g.Log().Errorf(ctx, "更新任务失败: %v", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if completed {
|
||||
// 3. 如果已存在,则更新任务
|
||||
_, err = dao.Task.Update(ctx, &dto.UpdateTaskReq{
|
||||
TaskId: req.TaskId,
|
||||
Status: task.TaskStatusCompleted,
|
||||
})
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// IsAllSubTasksCompleted 判断三个子任务是否全部完成
|
||||
// 参数:传入当前文档的所有子任务列表
|
||||
func IsAllSubTasksCompleted(subTasks []dto.TaskVO) bool {
|
||||
func IsAllSubTasksCompleted(subTasks []*dto.TaskVO) bool {
|
||||
// 必须包含 3 种任务类型
|
||||
hasKeywords := false
|
||||
hasVector := false
|
||||
|
||||
@@ -159,7 +159,7 @@ COMMENT ON COLUMN rag_knowledge_keyword.dataset_id IS '数据集ID';
|
||||
COMMENT ON COLUMN rag_knowledge_keyword.document_id IS '文档ID';
|
||||
COMMENT ON COLUMN rag_knowledge_keyword.word IS '关键词';
|
||||
COMMENT ON COLUMN rag_knowledge_keyword.weight IS '权重';
|
||||
|
||||
CREATE UNIQUE INDEX uk_rag_knowledge_keyword_tenant_dataset_doc_word ON rag_knowledge_keyword (tenant_id, dataset_id, document_id, word);
|
||||
--------------------pgsql创建rag_knowledge_keyword表语句---------------------------
|
||||
|
||||
--------------------pgsql创建rag_knowledge_task表语句---------------------------
|
||||
|
||||
Reference in New Issue
Block a user