会话

controllers/rag/kb_controller.go

@@ -38,13 +38,13 @@ func (kbctrl *KbController) SearchDocs() {
 		br.ErrMsg = "参数解析失败，Err:" + err.Error()
 		return
 	}
-	docs, err := facade.LLMKnowledgeBaseSearchDocs(req)
+	searchResp, err := facade.LLMKnowledgeBaseSearchDocs(req)
 	if err != nil {
 		br.Msg = "搜索知识库失败"
 		br.ErrMsg = "搜索知识库失败:" + err.Error()
 		return
 	}
-	br.Data = docs
+	br.Data = searchResp
 	br.Ret = 200
 	br.Success = true
 	br.Msg = "获取成功"

services/llm/facade/bus_response/bus_response.go

@@ -0,0 +1,14 @@
+package bus_response
+
+type SearchDocsResponse struct {
+	PageContent string   `json:"page_content"`
+	Metadata    Metadata `json:"metadata"`
+	Type        string   `json:"type"`
+	Id          string   `json:"id"`
+	Score       float32  `json:"score"`
+}
+
+type Metadata struct {
+	Source string `json:"source"`
+	Id     string `json:"id"`
+}

services/llm/facade/bus_response/eta_response.go

@@ -0,0 +1,6 @@
+package bus_response
+
+type SearchDocsEtaResponse struct {
+	Content string
+	Docs    []SearchDocsResponse
+}

services/llm/facade/llm_service.go

@@ -1,6 +1,7 @@
 package facade
 
 import (
+	"eta/eta_api/services/llm/facade/bus_response"
 	"eta/eta_api/utils/llm"
 	"eta/eta_api/utils/ws"
 	"fmt"
@@ -22,8 +23,16 @@ func AddSession(userId int, conn *websocket.Conn) {
 	ws.Manager().AddSession(session)
 }
 
-func LLMKnowledgeBaseSearchDocs(search LLMKnowledgeSearch) (resp string, err error) {
-	return llmService.SearchKbDocs(search.Query, search.KnowledgeBaseName)
+func LLMKnowledgeBaseSearchDocs(search LLMKnowledgeSearch) (resp bus_response.SearchDocsEtaResponse, err error) {
+	docs, err := llmService.SearchKbDocs(search.Query, search.KnowledgeBaseName)
+	if err != nil {
+		return
+	}
+	for _, doc := range docs.([]bus_response.SearchDocsResponse) {
+		resp.Content = resp.Content + doc.PageContent
+	}
+	resp.Docs = docs.([]bus_response.SearchDocsResponse)
+	return
 }
 
 type LLMKnowledgeSearch struct {

utils/llm/eta_llm/eta_llm_client.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"eta/eta_api/services/llm/facade/bus_response"
 	"eta/eta_api/utils"
 	"eta/eta_api/utils/llm"
 	"eta/eta_api/utils/llm/eta_llm/eta_llm_http"
@@ -47,15 +48,17 @@ func (ds *ETALLMClient) KnowledgeBaseChat() string {
 	return ""
 }
 
-func (ds *ETALLMClient) SearchKbDocs(query string, KnowledgeBaseName string) (content string, err error) {
+func (ds *ETALLMClient) SearchKbDocs(query string, KnowledgeBaseName string) (content interface{}, err error) {
 	// 类型断言
+
 	kbReq := eta_llm_http.KbSearchDocsRequest{
 		Query:             query,
 		KnowledgeBaseName: KnowledgeBaseName,
-		Model:             ds.LlmModel,
-		TopK:              3,
-		ScoreThreshold:    2,
+		TopK:              10,
+		ScoreThreshold:    0.5,
+		Metadata:          struct{}{},
 	}
+
 	body, err := json.Marshal(kbReq)
 	if err != nil {
 		return
@@ -65,7 +68,18 @@ func (ds *ETALLMClient) SearchKbDocs(query string, KnowledgeBaseName string) (co
 		err = errors.New(resp.Msg)
 		return
 	}
-	return "", nil
+	if resp.Data != nil {
+		var kbSearchRes []bus_response.SearchDocsResponse
+		err = json.Unmarshal(resp.Data, &kbSearchRes)
+		if err != nil {
+			err = errors.New("搜索知识库失败")
+			return
+		}
+		content = kbSearchRes
+		return
+	}
+	err = errors.New("搜索知识库失败")
+	return
 }
 func init() {
 	err := llm.Register(llm.ETA_LLM_CLIENT, GetInstance())
@@ -97,6 +111,7 @@ func parseResponse(response *http.Response) (baseResp eta_llm_http.BaseResponse,
 		err = fmt.Errorf("读取响应体失败: %w", err)
 		return
 	}
+	baseResp.Success = true
 	baseResp.Data = bodyBytes
 	return
 }

utils/llm/eta_llm/eta_llm_http/request.go

@@ -1,7 +1,5 @@
 package eta_llm_http
 
-import "encoding/json"
-
 type KbChatRequest struct {
 	Query          string           `json:"query"`
 	Mode           string           `json:"mode"`
@@ -23,11 +21,10 @@ type HistoryContent struct {
 }
 
 type KbSearchDocsRequest struct {
-	Query             string          `json:"query"`
-	KnowledgeBaseName string          `json:"Knowledge_base_name"`
-	TopK              int             `json:"top_k"`
-	ScoreThreshold    int             `json:"score_threshold"`
-	FileName          string            `json:"file_name"`
-	Model             string          `json:"model"`
-	Metadata          json.RawMessage `json:"metadata"`
+	Query             string      `json:"query"`
+	KnowledgeBaseName string      `json:"knowledge_base_name"`
+	TopK              int         `json:"top_k"`
+	ScoreThreshold    float32         `json:"score_threshold"`
+	FileName          string      `json:"file_name"`
+	Metadata          interface{} `json:"metadata"`
 }

utils/llm/eta_llm/eta_llm_http/response.go

@@ -8,3 +8,5 @@ type BaseResponse struct {
 	Success bool            `json:"success"`
 	Data    json.RawMessage `json:"data"`
 }
+
+

utils/llm/llm_client.go

@@ -21,5 +21,5 @@ func NewLLMClient(baseURL string, timeout time.Duration) *LLMClient {
 
 type LLMService interface {
 	KnowledgeBaseChat() string
-	SearchKbDocs(query string, KnowledgeBaseName string) (string, error)
+	SearchKbDocs(query string, KnowledgeBaseName string) (data interface{}, err error)
 }