Browse Source

新增ai接口

xyxie 1 week ago
parent
commit
dffae3836f

+ 39 - 0
utils/llm/eta_llm/eta_llm_client.go

@@ -34,6 +34,7 @@ const (
 	CONTENT_TYPE_JSON              = "application/json"
 	KNOWLEDGE_BASE_CHAT_API        = "/chat/kb_chat"
 	DOCUMENT_CHAT_API              = "/chat/file_chat"
+	COMPLETION_CHAT_API            = "/chat/completion_chat"
 	KNOWLEDGE_BASE_SEARCH_DOCS_API = "/knowledge_base/search_docs"
 	UPLOAD_TEMP_DOCS_API           = "/knowledge_base/upload_temp_docs"
 )
@@ -109,6 +110,44 @@ func (ds *ETALLMClient) DocumentChat(query string, KnowledgeId string, history [
 	return ds.DoStreamPost(DOCUMENT_CHAT_API, body)
 }
 
+func (ds *ETALLMClient) CompletionChat(query string, messages []json.RawMessage) (llmRes *http.Response, err error) {
+    ChatHistory := make([]eta_llm_http.HistoryContent, 0)
+    for _, historyItemStr := range messages {
+        var historyItem eta_llm_http.HistoryContentWeb
+        parseErr := json.Unmarshal(historyItemStr, &historyItem)
+        if parseErr != nil {
+            continue
+        }
+        ChatHistory = append(ChatHistory, eta_llm_http.HistoryContent{
+            Content: historyItem.Content,
+            Role:    historyItem.Role,
+        })
+    }
+    ChatHistory = append(ChatHistory, eta_llm_http.HistoryContent{
+        Content: query,
+        Role:    "user",
+    })
+    kbReq := eta_llm_http.CompletionChatRequest{
+        Mode:           KNOWLEDEG_CHAT_MODE,
+        Messages:       ChatHistory,
+        TopK:           3,
+        ScoreThreshold: 0.5,
+        Stream:         true,
+        Model:           ds.LlmModel,
+        Temperature:    0.7,
+        MaxTokens:      0,
+        PromptName:     DEFALUT_PROMPT_NAME,
+        ReturnDirect:   false,
+    }
+    fmt.Printf("%v", kbReq.Messages)
+    utils.FileLog.Info("kbReq.Messages: %v", kbReq.Messages)
+    body, err := json.Marshal(kbReq)
+    if err != nil {
+        return
+    }
+    return ds.DoStreamPost(COMPLETION_CHAT_API, body)
+}
+
 func (ds *ETALLMClient) KnowledgeBaseChat(query string, KnowledgeBaseName string, history []json.RawMessage) (llmRes *http.Response, err error) {
 	ChatHistory := make([]eta_llm_http.HistoryContent, 0)
 	for _, historyItemStr := range history {

+ 14 - 0
utils/llm/eta_llm/eta_llm_http/request.go

@@ -49,3 +49,17 @@ type UploadTempDocsRequest struct {
 	ChunkOverlap   string `json:"chunk_overlap"`
 	ZhTitleEnhance string `json:"zh_title_enhance"`
 }
+
+type CompletionChatRequest struct {
+    Mode           string           `json:"mode"`
+    KbName         string           `json:"kb_name"`
+    TopK           int              `json:"top_k"`
+    ScoreThreshold float32          `json:"score_threshold"`
+    Messages        []HistoryContent `json:"messages"`
+    Stream         bool             `json:"stream"`
+    Model          string           `json:"model"`
+    Temperature    float32          `json:"temperature"`
+    MaxTokens      int              `json:"max_tokens"`
+    PromptName     string           `json:"prompt_name"`
+    ReturnDirect   bool             `json:"return_direct"`
+}

+ 1 - 0
utils/llm/llm_client.go

@@ -28,4 +28,5 @@ type LLMService interface {
 	SearchKbDocs(query string, KnowledgeBaseName string) (data interface{}, err error)
 	UploadFileToTemplate(files []*os.File, param map[string]interface{}) (data interface{}, err error)
 	FileChat(query string, KnowledgeId string, llmModel string, history []json.RawMessage) (resp eta_llm_http.BaseResponse, err error)
+	CompletionChat(query string, messages []json.RawMessage) (llmRes *http.Response, err error)
 }