package ai

import (
	"encoding/json"
	"errors"
	"kefu/lib"
	"log"
	"strings"
)

// 训练功能
func Train(openaiUrl string, openaiKey string, pointId interface{}, collectName, content, fileId, title, url string) (string, error) {

	gpt := lib.NewChatGptTool(openaiUrl, openaiKey)
	str := strings.TrimSpace(content)
	response, err := gpt.GetEmbedding(str, "text-embedding-ada-002")
	if err != nil {
		log.Println("向量接口失败：", err)
		return "", err
	}
	var embeddingResponse lib.EmbeddingResponse
	json.Unmarshal([]byte(response), &embeddingResponse)
	if len(embeddingResponse.Data) == 0 {
		log.Println("向量接口失败：", response)
		return "", errors.New(response)
	}

	points := []map[string]interface{}{
		{
			"id":      pointId,
			"payload": map[string]interface{}{"text": str, "title": title, "url": url, "fileid": fileId},
			"vector":  embeddingResponse.Data[0].Embedding,
		},
	}
	res, err := lib.PutPoints(collectName, points)
	return res, err
}

// 长文本分块
func SplitTextByLength(text string, length int) []string {
	var blocks []string
	runes := []rune(text)
	for i := 0; i < len(runes); i += length {
		j := i + length
		if j > len(runes) {
			j = len(runes)
		}
		blocks = append(blocks, string(runes[i:j]))
	}
	return blocks
}