wechat-api/internal/utils/compapi/result.go

package compapi

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"wechat-api/internal/types"
)

// 包装了原始的 API 响应，并提供了解析助手方法
type ChatResult struct {
	*types.CompOpenApiResp
	err error
}

// NewChatResult 是 ChatResult 的构造函数
func NewChatResult(resp any) *ChatResult {
	nresp, err := AnyToCompOpenApiResp(resp)
	res := ChatResult{nresp, err}
	if nresp == nil {
		res.err = errors.New("CompOpenApiRes is nil")
	}
	return &res
}

// 以下是ChatResul助手方法

// GetContentString 返回第一个 Choice 的 Message Content (标准方式)
func (r *ChatResult) GetContentString() (string, error) {
	var (
		content string = ""
		err     error
	)
	if r.err == nil && len(r.Choices) > 0 {
		content = r.Choices[0].Message.Content
	} else if r.err == nil && len(r.Choices) == 0 {
		err = errors.New("choices empty")
	}
	return content, err
}

func (r *ChatResult) GetContentJsonStr() (string, error) {
	var (
		content string = ""
		err     error
	)
	if r.err == nil && len(r.Choices) > 0 {
		content = r.Choices[0].Message.Content
	} else if r.err == nil && len(r.Choices) == 0 {
		err = errors.New("choices empty")
	}
	if !IsOpenaiModel(r.Model) { //不支持Response Schema的要特殊处理一下
		var isJsonContent bool
		content, isJsonContent = ExtractJSONContent(content)
		if !isJsonContent {
			return "", errors.New("invalid json content")
		}
	}

	return content, err
}

// ParseContentAs 解析 Message Content 中的 JSON 到指定的 Go 结构体
// target 必须是一个指向目标结构体实例的指针 (e.g., &MyStruct{})
func (r *ChatResult) ParseContentAs(target any) error {
	content, err := r.GetContentString()
	if err != nil {
		return fmt.Errorf("parseContent err: %s", err)
	} else if content == "" {
		return errors.New("parseContent err: content is empty or unavailable")
	}

	if !IsOpenaiModel(r.Model) { //不支持Response Schema的要特殊处理一下
		var isJsonContent bool
		content, isJsonContent = ExtractJSONContent(content)
		if !isJsonContent {
			return errors.New("invalid json content")
		}
	}
	return ParseContentAs(content, target, false)
}

func AnyToBytes(in any) ([]byte, error) {
	switch v := in.(type) {
	case string:
		return []byte(v), nil
	case []byte:
		return v, nil
	default:
		return json.Marshal(v)
	}
}

func AnyToCompOpenApiResp(in any) (*types.CompOpenApiResp, error) {
	if resp, ok := in.(*types.CompOpenApiResp); ok {
		return resp, nil
	}
	if resp, ok := in.(types.CompOpenApiResp); ok {
		return &resp, nil
	}
	bs, err := AnyToBytes(in)
	if err != nil {
		return nil, err
	}
	nresp := &types.CompOpenApiResp{}
	err = json.Unmarshal(bs, nresp)
	if err != nil {
		return nil, err
	}
	return nresp, nil
}

func AnyToCompApiReq(in any) (*types.CompApiReq, error) {
	if req, ok := in.(*types.CompApiReq); ok {
		return req, nil
	}
	if req, ok := in.(types.CompApiReq); ok {
		return &req, nil
	}
	bs, err := AnyToBytes(in)
	if err != nil {
		return nil, err
	}
	nreq := &types.CompApiReq{}
	err = json.Unmarshal(bs, nreq)
	if err != nil {
		return nil, err
	}
	return nreq, nil
}

func CheckJSON(input any, checkEmpty bool) (bool, error) {
	inputBytes, err := AnyToBytes(input)
	if err != nil {
		return false, err
	}
	var raw json.RawMessage
	err = json.Unmarshal(inputBytes, &raw)
	if err != nil {
		return false, fmt.Errorf("input is not valid JSON: %w", err)
	}
	if checkEmpty {
		trimmed := bytes.TrimSpace(inputBytes)
		if len(trimmed) == 0 {
			return false, fmt.Errorf("input is empty")
		}
	}
	return true, nil
}

func WrapJSON(input any, warpKey string, checkValid bool) ([]byte, error) {
	var (
		inputBytes  []byte
		outputBytes []byte
		err         error
	)
	if inputBytes, err = AnyToBytes(input); err != nil {
		return nil, err
	}

	if checkValid {
		if _, err = CheckJSON(inputBytes, true); err != nil {
			return nil, err
		}
	}
	if len(warpKey) == 0 {
		return inputBytes, nil
	}
	wrapper := map[string]json.RawMessage{
		warpKey: json.RawMessage(inputBytes),
	}
	if outputBytes, err = json.Marshal(wrapper); err != nil {
		return nil, fmt.Errorf("failed to marshal wrapper structure: %w", err)
	}
	return outputBytes, nil
}

func ParseContentAs(content string, target any, removeJsonBlock bool) error {

	if removeJsonBlock &&
		strings.HasPrefix(content, "```json") && strings.HasSuffix(content, "```") {
		content = strings.TrimSuffix(strings.TrimPrefix(content, "```json"), "```")
		content = strings.TrimSpace(content)
	}

	if err := json.Unmarshal([]byte(content), target); err != nil {
		return fmt.Errorf("parseContent err:failed to unmarshal"+
			" content JSON into target type '%w'", err)
	}
	return nil
}

func ExtractJSONContent(s string) (string, bool) {
	startMarker := "```json"
	endMarker := "```"

	// 寻找起始标记
	startIdx := strings.Index(s, startMarker)
	if startIdx == -1 {
		return "", false // 没有起始标记
	}

	// 寻找结束标记（需在起始标记之后查找）
	endIdx := strings.LastIndex(s, endMarker)
	if endIdx == -1 || endIdx <= startIdx {
		return "", false // 没有结束标记或标记顺序错误
	}

	// 计算内容范围
	contentStart := startIdx + len(startMarker)
	contentEnd := endIdx

	// 提取内容并去除前后空白
	content := strings.TrimSpace(s[contentStart:contentEnd])

	// 若内容为空视为无效
	if content == "" {
		return "", false
	}

	return content, true
}