spider_creater/backend/ai/ai.go

package ai

import (
	"errors"
	"fmt"

	be "spider_creator/backend"
	"strings"

	zhipu "github.com/itcwc/go-zhipu/model_api"
)

const (
	MODEL_NAME = "glm-4-flash" //"glm-4-air"
)

var (
	expireAtTime = int64(1719803252) // token 过期时间
	apiKey       = "5343038d9d934536456f281f8487866a.YUmO7HK9xNb990j9"
)

// UpdateResultDateStr
func UpdateResultDateStr(rs be.ResultItems) (err error) {
	tmp := make([]string, len(rs))
	for i, v := range rs {
		tmp[i] = v.ListPubTime
	}

	prompt := fmt.Sprintf(`根据我提供的内容，识别每行文本中的日期，按照YYYY-MM-DD形式输出,如：2024-01-01；找不到日期数据，输出NULL。不要联网，不要解释，不要说明，直接输出结果。
				---------------------------
				%s				
				`, strings.Join(tmp, "\n"))
	mssage := zhipu.PostParams{
		Model: MODEL_NAME,
		Messages: []zhipu.Message{
			{
				Role:    "user", // 消息的角色信息 详见文档
				Content: prompt, // 消息内容
			},
		},
	}
	postResponse, err := zhipu.BeCommonModel(expireAtTime, mssage, apiKey)
	if err != nil {
		return err
	}

	//解析数据
	choices, _ := postResponse["choices"].([]interface{})
	obj, _ := choices[0].(map[string]interface{})
	message, _ := obj["message"].(map[string]interface{})
	value, _ := message["content"].(string)
	// qu.Debug("提示语", prompt)
	// qu.Debug("AI调用结果", value)
	results := strings.Split(value, "\n")
	if len(results) < len(rs) {
		err = errors.New("调用大模型失败")
		return
	}
	//更新
	for i, v := range rs {
		if results[i] != "NULL" {
			v.ListPubTime = results[i]
		} else {
			v.ListPubTime = ""
		}
	}
	return nil
}

// Question
func Question(prompt string) (ret interface{}, err error) {
	mssage := zhipu.PostParams{
		Model: MODEL_NAME,
		Messages: []zhipu.Message{
			{
				Role:    "user", // 消息的角色信息 详见文档
				Content: prompt, // 消息内容
			},
		},
	}
	postResponse, err := zhipu.BeCommonModel(expireAtTime, mssage, apiKey)
	if err != nil {
		return nil, err
	}

	//解析数据
	choices, _ := postResponse["choices"].([]interface{})
	obj, _ := choices[0].(map[string]interface{})
	message, _ := obj["message"].(map[string]interface{})
	value := message["content"]

	return value, nil
}