eta_api/services/llm/chat.go

package llm

import (
	"bytes"
	"encoding/json"
	"eta/eta_api/utils"
	"eta/eta_api/utils/llm/eta_llm"
	"eta/eta_api/utils/llm/eta_llm/eta_llm_http"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
)

var (
	llmService = eta_llm.GetInstance()
)

type UploadTempDocsResp struct {
	Code int                   `json:"code"`
	Msg  string                `json:"msg"`
	Data UploadTempDocDataResp `json:"data"`
}

type UploadTempDocDataResp struct {
	Id          string        `json:"id"`
	FailedFiles []interface{} `json:"failed_files"`
}

// UploadTempDocs
// @Description: 上传到临时知识库
// @author: Roc
// @datetime 2025-03-10 14:42:18
// @param filePath string
// @return resp UploadDocsResp
// @return err error
func UploadTempDocs(filePath string) (resp UploadTempDocsResp, err error) {
	postUrl := utils.LLM_SERVER + "/knowledge_base/upload_temp_docs"

	params := make(map[string]string)
	//params[`prev_id`] = ``
	params[`chunk_size`] = `750`
	params[`chunk_overlap`] = `150`
	params[`zh_title_enhance`] = `true`

	files := make(map[string]string)
	files[`files`] = filePath

	result, err := PostFormData(postUrl, params, files)
	if err != nil {
		return
	}

	str := string(result)
	fmt.Println(str)

	err = json.Unmarshal(result, &resp)
	if err != nil {
		return
	}

	return
}

type UploadDocsResp struct {
	Code int               `json:"code"`
	Msg  string            `json:"msg"`
	Data UploadDocDataResp `json:"data"`
}

type UploadDocDataResp struct {
	Id          string            `json:"id"`
	FailedFiles map[string]string `json:"failed_files"`
}

// UploadDocsToKnowledge
// @Description: 上传文章到知识库
// @author: Roc
// @datetime 2025-03-10 14:40:44
// @param filePath string
// @param knowledgeName string
// @return resp UploadTempDocsResp
// @return err error
func UploadDocsToKnowledge(filePath, knowledgeName string) (updateResp UploadDocDataResp, err error) {
	postUrl := utils.LLM_SERVER + "/knowledge_base/upload_docs"

	params := make(map[string]string)
	params[`knowledge_base_name`] = knowledgeName
	params[`override`] = `true`              // 覆盖已有文件
	params[`to_vector_store`] = `true`       // 上传文件后是否进行向量化
	params[`chunk_size`] = `750`             // 知识库中单段文本最大长度
	params[`chunk_overlap`] = `150`          // 知识库中相邻文本重合长度
	params[`zh_title_enhance`] = `true`      // 是否开启中文标题加强
	params[`docs`] = ``                      // 自定义的docs，需要转为json字符串
	params[`not_refresh_vs_cache`] = `false` // 暂不保存向量库（用于FAISS）

	files := make(map[string]string)
	files[`files`] = filePath

	result, err := PostFormData(postUrl, params, files)
	if err != nil {
		return
	}

	str := string(result)
	fmt.Println(str)

	var resp UploadDocsResp
	err = json.Unmarshal(result, &resp)
	if err != nil {
		return
	}

	if resp.Code != 200 {
		err = fmt.Errorf(`上传文件失败: %s`, resp.Msg)
		return
	}
	updateResp = resp.Data

	return
}

// ChatResp 问答响应
type ChatResp struct {
	Answer string   `json:"answer"`
	Docs   []string `json:"docs"`
}

type HistoryContent struct {
	Content string `json:"content"`
	Role    string `json:"role"`
}

func ChatByFile(knowledgeId, question string, historyList []eta_llm_http.HistoryContent) (answerStr string, answer ChatResp, err error) {
	// 没有问题那就直接返回
	if question == `` {
		return
	}

	history := make([]json.RawMessage, 0)
	for _, v := range historyList {
		tmpHistory, tmpErr := json.Marshal(v)
		if tmpErr != nil {
			return
		}
		history = append(history, json.RawMessage(string(tmpHistory)))
	}

	resp, err := llmService.DocumentChat(question, knowledgeId, history, false)
	if err != nil {
		return
	}
	defer func() {
		if resp != nil && resp.Body != nil {
			_ = resp.Body.Close()
		}
	}()
	if resp == nil {
		err = fmt.Errorf(`知识库问答失败: 无应答`)
		return
	}
	result, err := io.ReadAll(resp.Body)
	if err != nil {
		err = fmt.Errorf(`知识库问答数据解析失败: %s`, err.Error())
		return
	}

	answerStr = string(result)
	// 找到"data:"关键字的位置
	dataIndex := bytes.Index([]byte(answerStr), []byte("data: "))
	if dataIndex == -1 {
		err = fmt.Errorf(`未找到"data:"关键字`)
		return
	}

	// 提取"data:"关键字之后的部分
	answerStr = answerStr[dataIndex+len("data: "):]

	// 解析JSON数据
	err = json.Unmarshal([]byte(answerStr), &answer)
	if err != nil {
		err = fmt.Errorf(`解析JSON数据失败: %s`, err.Error())
		return
	}

	return
}

// PostFormData sends a POST request with form-data
func PostFormData(url string, params map[string]string, files map[string]string) ([]byte, error) {

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	for key, val := range params {
		if err := writer.WriteField(key, val); err != nil {
			return nil, err
		}
	}

	for fieldName, filePath := range files {
		file, err := os.Open(filePath)
		if err != nil {
			return nil, err
		}
		defer file.Close()

		part, err := writer.CreateFormFile(fieldName, filePath)
		if err != nil {
			return nil, err
		}
		_, err = io.Copy(part, file)
		if err != nil {
			return nil, err
		}
	}

	err := writer.Close()
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest("POST", url, body)
	if err != nil {
		return nil, err
	}
	//req.Header.Set("accept", `application/json`)
	req.Header.Set("Content-Type", writer.FormDataContentType())

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	result, err := io.ReadAll(resp.Body)

	return result, nil
}