477 lines
14 KiB
Go
477 lines
14 KiB
Go
package service
|
||
|
||
import (
|
||
"bufio"
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"sort"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"rag-local/common"
|
||
"rag-local/kb/consts"
|
||
"rag-local/kb/dao"
|
||
"rag-local/kb/model/domain"
|
||
"rag-local/kb/model/entity"
|
||
|
||
eembedding "github.com/cloudwego/eino/components/embedding"
|
||
emodel "github.com/cloudwego/eino/components/model"
|
||
eretriever "github.com/cloudwego/eino/components/retriever"
|
||
"github.com/cloudwego/eino/schema"
|
||
"github.com/gogf/gf/v2/errors/gerror"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
)
|
||
|
||
var httpClient = &http.Client{Timeout: 2 * time.Minute}
|
||
|
||
// ---------- OpenAI 兼容 HTTP 组件 ----------
|
||
|
||
type openAIMessage struct {
|
||
Role string `json:"role"`
|
||
Content string `json:"content"`
|
||
}
|
||
|
||
type openAIChatResponse struct {
|
||
Choices []struct {
|
||
Message openAIMessage `json:"message"`
|
||
FinishReason string `json:"finish_reason"`
|
||
} `json:"choices"`
|
||
Usage *schema.TokenUsage `json:"usage"`
|
||
}
|
||
|
||
type openAIStreamChunk struct {
|
||
Choices []struct {
|
||
Delta openAIMessage `json:"delta"`
|
||
FinishReason string `json:"finish_reason"`
|
||
} `json:"choices"`
|
||
}
|
||
|
||
// OpenAIChatModel 基于 OpenAI 兼容 /chat/completions 接口的对话模型，实现 eino model.ChatModel
|
||
type OpenAIChatModel struct {
|
||
cfg *entity.ModelConfig
|
||
}
|
||
|
||
func NewOpenAIChatModel(cfg *entity.ModelConfig) *OpenAIChatModel {
|
||
return &OpenAIChatModel{cfg: cfg}
|
||
}
|
||
|
||
func (m *OpenAIChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...emodel.Option) (*schema.Message, error) {
|
||
payload := map[string]any{
|
||
"model": m.cfg.ModelName,
|
||
"messages": buildOpenAIMessages(input),
|
||
"stream": false,
|
||
}
|
||
body, err := postOpenAI(ctx, m.cfg, m.endpoint("/chat/completions"), payload)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
var resp openAIChatResponse
|
||
if err := json.Unmarshal(body, &resp); err != nil {
|
||
return nil, gerror.Wrap(err, "解析模型响应失败")
|
||
}
|
||
if len(resp.Choices) == 0 {
|
||
return nil, gerror.New("模型返回空响应")
|
||
}
|
||
choice := resp.Choices[0]
|
||
msg := &schema.Message{Role: schema.Assistant, Content: choice.Message.Content}
|
||
if choice.FinishReason != "" {
|
||
msg.ResponseMeta = &schema.ResponseMeta{FinishReason: choice.FinishReason, Usage: resp.Usage}
|
||
}
|
||
return msg, nil
|
||
}
|
||
|
||
func (m *OpenAIChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...emodel.Option) (*schema.StreamReader[*schema.Message], error) {
|
||
payload := map[string]any{
|
||
"model": m.cfg.ModelName,
|
||
"messages": buildOpenAIMessages(input),
|
||
"stream": true,
|
||
}
|
||
reader, writer := schema.Pipe[*schema.Message](16)
|
||
go func() {
|
||
defer writer.Close()
|
||
body, err := postOpenAIStream(ctx, m.cfg, m.endpoint("/chat/completions"), payload)
|
||
if err != nil {
|
||
writer.Send(nil, err)
|
||
return
|
||
}
|
||
defer body.Close()
|
||
br := bufio.NewReader(body)
|
||
for {
|
||
line, err := br.ReadBytes('\n')
|
||
text := strings.TrimSpace(string(line))
|
||
if strings.HasPrefix(text, "data:") {
|
||
data := strings.TrimSpace(strings.TrimPrefix(text, "data:"))
|
||
if data == "[DONE]" {
|
||
break
|
||
}
|
||
var chunk openAIStreamChunk
|
||
if json.Unmarshal([]byte(data), &chunk) == nil && len(chunk.Choices) > 0 {
|
||
if delta := chunk.Choices[0].Delta.Content; delta != "" {
|
||
if closed := writer.Send(&schema.Message{Role: schema.Assistant, Content: delta}, nil); closed {
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if err != nil {
|
||
break
|
||
}
|
||
}
|
||
}()
|
||
return reader, nil
|
||
}
|
||
|
||
func (m *OpenAIChatModel) BindTools(tools []*schema.ToolInfo) error {
|
||
return nil
|
||
}
|
||
|
||
func (m *OpenAIChatModel) endpoint(path string) string {
|
||
return strings.TrimRight(m.cfg.EndpointUrl, "/") + path
|
||
}
|
||
|
||
// OpenAIEmbedder 基于 OpenAI 兼容 /embeddings 接口的向量模型，实现 eino embedding.Embedder
|
||
type OpenAIEmbedder struct {
|
||
cfg *entity.ModelConfig
|
||
}
|
||
|
||
func NewOpenAIEmbedder(cfg *entity.ModelConfig) *OpenAIEmbedder {
|
||
return &OpenAIEmbedder{cfg: cfg}
|
||
}
|
||
|
||
// Dim 配置的向量维度（vec0 表建表维度需一致）
|
||
func (e *OpenAIEmbedder) Dim() int {
|
||
if e.cfg.Dimension > 0 {
|
||
return e.cfg.Dimension
|
||
}
|
||
return consts.DefaultEmbeddingDim
|
||
}
|
||
|
||
func (e *OpenAIEmbedder) EmbedStrings(ctx context.Context, texts []string, opts ...eembedding.Option) ([][]float64, error) {
|
||
payload := map[string]any{"model": e.cfg.ModelName, "input": texts}
|
||
body, err := postOpenAI(ctx, e.cfg, strings.TrimRight(e.cfg.EndpointUrl, "/")+"/embeddings", payload)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
var resp struct {
|
||
Data []struct {
|
||
Embedding []float64 `json:"embedding"`
|
||
Index int `json:"index"`
|
||
} `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(body, &resp); err != nil {
|
||
return nil, gerror.Wrap(err, "解析 embedding 响应失败")
|
||
}
|
||
if len(resp.Data) == 0 {
|
||
return nil, gerror.New("embedding 接口返回空数据")
|
||
}
|
||
out := make([][]float64, len(texts))
|
||
for _, d := range resp.Data {
|
||
if d.Index >= 0 && d.Index < len(out) {
|
||
out[d.Index] = d.Embedding
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// BuildChatModel 按配置 id 构建对话模型
|
||
func BuildChatModel(ctx context.Context, cfgId int64) (*OpenAIChatModel, error) {
|
||
cfg, err := dao.ModelConfig.GetOne(ctx, cfgId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if cfg == nil {
|
||
return nil, gerror.New("模型配置不存在")
|
||
}
|
||
if cfg.ModelType != consts.ModelTypeChat {
|
||
return nil, gerror.New("该模型配置不是对话模型（model_type=chat）")
|
||
}
|
||
if cfg.EndpointUrl == "" || cfg.ModelName == "" {
|
||
return nil, gerror.New("模型配置缺少 endpoint_url 或 model_name")
|
||
}
|
||
return NewOpenAIChatModel(cfg), nil
|
||
}
|
||
|
||
// BuildEmbedder 按配置 id 构建向量模型
|
||
func BuildEmbedder(ctx context.Context, cfgId int64) (*OpenAIEmbedder, error) {
|
||
cfg, err := dao.ModelConfig.GetOne(ctx, cfgId)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if cfg == nil {
|
||
return nil, gerror.New("模型配置不存在")
|
||
}
|
||
if cfg.ModelType != consts.ModelTypeEmbedding {
|
||
return nil, gerror.New("该模型配置不是向量模型（model_type=embedding）")
|
||
}
|
||
if cfg.EndpointUrl == "" || cfg.ModelName == "" {
|
||
return nil, gerror.New("模型配置缺少 endpoint_url 或 model_name")
|
||
}
|
||
return NewOpenAIEmbedder(cfg), nil
|
||
}
|
||
|
||
// HybridRetriever 混合检索器：向量 KNN + FTS5 BM25，RRF 融合，实现 eino retriever.Retriever
|
||
type HybridRetriever struct {
|
||
embedder eembedding.Embedder
|
||
datasetId int64
|
||
}
|
||
|
||
func NewHybridRetriever(embedder eembedding.Embedder, datasetId int64) *HybridRetriever {
|
||
return &HybridRetriever{embedder: embedder, datasetId: datasetId}
|
||
}
|
||
|
||
func (r *HybridRetriever) Retrieve(ctx context.Context, query string, opts ...eretriever.Option) ([]*schema.Document, error) {
|
||
o := eretriever.GetCommonOptions(nil, opts...)
|
||
topK := consts.HybridTopK
|
||
if o.TopK != nil && *o.TopK > 0 {
|
||
topK = *o.TopK
|
||
}
|
||
|
||
scores := make(map[int64]float64)
|
||
srcs := make(map[int64][]string)
|
||
|
||
if r.embedder != nil {
|
||
vecs, err := r.embedder.EmbedStrings(ctx, []string{query})
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "query embed failed: %v", err)
|
||
} else if len(vecs) > 0 {
|
||
hits, err := dao.Chunk.VecSearch(ctx, r.datasetId, domain.VecJsonF64(vecs[0]), consts.VectorTopK)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "vec search failed: %v", err)
|
||
} else {
|
||
for i, h := range hits {
|
||
scores[h.ChunkId] += 1 / (float64(consts.RrfK) + float64(i) + 1)
|
||
srcs[h.ChunkId] = append(srcs[h.ChunkId], "vector")
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
ftsHits, err := dao.Chunk.FtsSearch(ctx, r.datasetId, common.TokenizeQuery(query), consts.FtsTopK)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "fts search failed: %v", err)
|
||
} else {
|
||
for i, h := range ftsHits {
|
||
scores[h.ChunkId] += 1 / (float64(consts.RrfK) + float64(i) + 1)
|
||
srcs[h.ChunkId] = append(srcs[h.ChunkId], "fts")
|
||
}
|
||
}
|
||
|
||
type scoredChunk struct {
|
||
id int64
|
||
score float64
|
||
sources []string
|
||
}
|
||
items := make([]scoredChunk, 0, len(scores))
|
||
for id, s := range scores {
|
||
items = append(items, scoredChunk{id: id, score: s, sources: srcs[id]})
|
||
}
|
||
sort.Slice(items, func(i, j int) bool { return items[i].score > items[j].score })
|
||
if len(items) > topK {
|
||
items = items[:topK]
|
||
}
|
||
|
||
docs := make([]*schema.Document, 0, len(items))
|
||
for _, it := range items {
|
||
chunk, err := dao.Chunk.GetOne(ctx, it.id)
|
||
if err != nil || chunk == nil {
|
||
continue
|
||
}
|
||
docs = append(docs, &schema.Document{
|
||
ID: strconv.FormatInt(chunk.Id, 10),
|
||
Content: chunk.Content,
|
||
MetaData: map[string]any{
|
||
"chunk_id": chunk.Id,
|
||
"document_id": chunk.DocumentId,
|
||
"seq": chunk.Seq,
|
||
"score": it.score,
|
||
"sources": it.sources,
|
||
},
|
||
})
|
||
}
|
||
return docs, nil
|
||
}
|
||
|
||
// ---------- RAG 问答工作流 ----------
|
||
|
||
var ChatService = &chatService{}
|
||
|
||
type chatService struct{}
|
||
|
||
// MaxHistoryRounds 携带进模型的历史对话轮数（每条消息算一条，含用户与助手）
|
||
const MaxHistoryRounds = 10
|
||
|
||
// Ask RAG 问答工作流：混合检索 → 组装提示（含引用编号）→ 对话模型流式生成。
|
||
// history 需已包含最新一条用户问题；onCitations 在检索完成后先于流式输出回调；onDelta 接收增量文本，均可为 nil。
|
||
func (s *chatService) Ask(ctx context.Context, datasetId int64, question string, history []*schema.Message, onCitations func([]domain.Citation), onDelta func(string)) (string, []domain.Citation, error) {
|
||
docs, err := s.retrieve(ctx, datasetId, question)
|
||
if err != nil {
|
||
return "", nil, err
|
||
}
|
||
citations := buildCitations(docs)
|
||
if onCitations != nil {
|
||
onCitations(citations)
|
||
}
|
||
|
||
triples, err := KgRelationService.GraphEnhance(ctx, datasetId, question)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "graph enhance failed: %v", err)
|
||
}
|
||
|
||
defaultChatModel, _, err := SystemConfigService.GetSettings(ctx)
|
||
if err != nil {
|
||
return "", nil, err
|
||
}
|
||
if defaultChatModel <= 0 {
|
||
return "", nil, gerror.New("请先在设置中选择默认对话模型")
|
||
}
|
||
model, err := BuildChatModel(ctx, defaultChatModel)
|
||
if err != nil {
|
||
return "", nil, err
|
||
}
|
||
|
||
msgs := make([]*schema.Message, 0, len(history)+1)
|
||
msgs = append(msgs, &schema.Message{Role: schema.System, Content: buildSystemPrompt(citations, triples)})
|
||
if start := len(history) - MaxHistoryRounds*2; start > 0 {
|
||
history = history[start:]
|
||
}
|
||
msgs = append(msgs, history...)
|
||
|
||
sr, err := model.Stream(ctx, msgs)
|
||
if err != nil {
|
||
return "", nil, gerror.Wrap(err, "调用对话模型失败")
|
||
}
|
||
defer sr.Close()
|
||
var full strings.Builder
|
||
for {
|
||
m, err := sr.Recv()
|
||
if errors.Is(err, io.EOF) {
|
||
break
|
||
}
|
||
if err != nil {
|
||
return "", nil, gerror.Wrap(err, "流式输出中断")
|
||
}
|
||
full.WriteString(m.Content)
|
||
if onDelta != nil {
|
||
onDelta(m.Content)
|
||
}
|
||
}
|
||
return full.String(), citations, nil
|
||
}
|
||
|
||
// retrieve 构建数据集绑定的混合检索器并执行检索（无 embedding 配置时仅全文）
|
||
func (s *chatService) retrieve(ctx context.Context, datasetId int64, question string) ([]*schema.Document, error) {
|
||
var emb eembedding.Embedder
|
||
if cfgId, err := dao.Dataset.GetEmbeddingCfgId(ctx, datasetId); err == nil && cfgId > 0 {
|
||
if em, err := BuildEmbedder(ctx, cfgId); err == nil {
|
||
emb = em
|
||
} else {
|
||
g.Log().Warningf(ctx, "build embedder failed, retrieve fts only: %v", err)
|
||
}
|
||
}
|
||
return NewHybridRetriever(emb, datasetId).Retrieve(ctx, question)
|
||
}
|
||
|
||
// buildCitations 从检索结果生成引用列表（编号从 1 开始，与提示词 [编号] 对应）
|
||
func buildCitations(docs []*schema.Document) []domain.Citation {
|
||
cits := make([]domain.Citation, 0, len(docs))
|
||
for i, d := range docs {
|
||
c := domain.Citation{Index: i + 1, Content: d.Content}
|
||
if id, ok := d.MetaData["chunk_id"].(int64); ok {
|
||
c.ChunkId = id
|
||
}
|
||
if id, ok := d.MetaData["document_id"].(int64); ok {
|
||
c.DocumentId = id
|
||
}
|
||
if s, ok := d.MetaData["score"].(float64); ok {
|
||
c.Score = s
|
||
}
|
||
if srcs, ok := d.MetaData["sources"].([]string); ok {
|
||
c.Sources = srcs
|
||
}
|
||
cits = append(cits, c)
|
||
}
|
||
return cits
|
||
}
|
||
|
||
// buildSystemPrompt 系统提示词：引用资料编号 + 检索片段 + 知识图谱三元组（M5 图增强）
|
||
func buildSystemPrompt(citations []domain.Citation, triples []string) string {
|
||
var sb strings.Builder
|
||
sb.WriteString("你是一个本地知识库助手。请仅根据以下资料回答用户问题；若资料不足以回答，请明确说明。")
|
||
sb.WriteString("回答引用资料时，在对应位置标注 [编号]。\n\n【资料】\n")
|
||
for _, c := range citations {
|
||
sb.WriteString(fmt.Sprintf("[%d] %s\n", c.Index, c.Content))
|
||
}
|
||
if len(triples) > 0 {
|
||
sb.WriteString("\n【知识图谱】以下为与问题相关的实体关系，可辅助回答关系类问题：\n")
|
||
for _, t := range triples {
|
||
sb.WriteString(t + "\n")
|
||
}
|
||
}
|
||
return sb.String()
|
||
}
|
||
|
||
// ---------- HTTP 辅助 ----------
|
||
|
||
func buildOpenAIMessages(input []*schema.Message) []openAIMessage {
|
||
out := make([]openAIMessage, 0, len(input))
|
||
for _, m := range input {
|
||
if m == nil {
|
||
continue
|
||
}
|
||
role := string(m.Role)
|
||
if role == "" {
|
||
role = string(schema.User)
|
||
}
|
||
out = append(out, openAIMessage{Role: role, Content: m.Content})
|
||
}
|
||
return out
|
||
}
|
||
|
||
func postOpenAI(ctx context.Context, cfg *entity.ModelConfig, url string, payload any) ([]byte, error) {
|
||
body, err := doOpenAIRequest(ctx, cfg, url, payload)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer body.Close()
|
||
resp, err := io.ReadAll(body)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return resp, nil
|
||
}
|
||
|
||
func postOpenAIStream(ctx context.Context, cfg *entity.ModelConfig, url string, payload any) (io.ReadCloser, error) {
|
||
return doOpenAIRequest(ctx, cfg, url, payload)
|
||
}
|
||
|
||
func doOpenAIRequest(ctx context.Context, cfg *entity.ModelConfig, url string, payload any) (io.ReadCloser, error) {
|
||
buf, err := json.Marshal(payload)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf))
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
req.Header.Set("Content-Type", "application/json")
|
||
if cfg.ApiKey != "" {
|
||
req.Header.Set("Authorization", "Bearer "+cfg.ApiKey)
|
||
}
|
||
resp, err := httpClient.Do(req)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if resp.StatusCode >= 400 {
|
||
msg, _ := io.ReadAll(resp.Body)
|
||
resp.Body.Close()
|
||
return nil, gerror.Newf("模型接口 %s 返回 %d: %s", url, resp.StatusCode, strings.TrimSpace(string(msg)))
|
||
}
|
||
return resp.Body, nil
|
||
}
|