YunShu/llm.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"
)

var (
	llmHost  = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
	llmModel = "doubao-seed-2-0-pro-260215"
	llmKey   = ""
)

func init() {
	// 1. 从配置文件加载
	cfg, err := LoadConfig()
	if err == nil {
		if cfg.LLM.Host != "" {
			llmHost = cfg.LLM.Host
		}
		if cfg.LLM.Model != "" {
			llmModel = cfg.LLM.Model
		}
		if cfg.LLM.Key != "" {
			llmKey = cfg.LLM.Key
		}
	}

	// 2. 环境变量覆盖配置文件（优先级最高）
	if v := os.Getenv("LLM_ENDPOINT"); v != "" {
		llmHost = v
	}
	if v := os.Getenv("LLM_MODEL"); v != "" {
		llmModel = v
	}
	if v := os.Getenv("LLM_API_KEY"); v != "" {
		llmKey = v
	}
	// 兼容旧环境变量名
	if v := os.Getenv("OPENAI_API_KEY"); v != "" && llmKey == "" {
		llmKey = v
	}
}

// GetLLMKey 获取 API Key，优先使用已加载的密钥
func GetLLMKey() (string, error) {
	if llmKey == "" {
		return "", fmt.Errorf("未配置 API Key。请运行 'weather-cia onboard' 初始化，或设置 LLM_API_KEY 环境变量")
	}
	return llmKey, nil
}

// CallLLM 调用大模型 API（兼容 OpenAI Chat Completion 格式）
func CallLLM(messages []Message, toolDefs []ToolDef) (*OpenAIResponse, error) {
	apiKey, err := GetLLMKey()
	if err != nil {
		return nil, err
	}

	reqBody := map[string]interface{}{
		"model":    llmModel,
		"messages": messages,
	}

	// 注册工具定义
	if len(toolDefs) > 0 {
		tools := make([]OpenAITool, 0, len(toolDefs))
		for _, td := range toolDefs {
			tools = append(tools, OpenAITool{
				Type: "function",
				Function: OpenAIToolFunc{
					Name:        td.Name,
					Description: td.Description,
					Parameters:  td.Parameters,
				},
			})
		}
		reqBody["tools"] = tools
		reqBody["tool_choice"] = "auto"
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return nil, fmt.Errorf("序列化请求失败: %w", err)
	}

	req, err := http.NewRequest("POST", llmHost, bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("创建请求失败: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+apiKey)

	client := &http.Client{Timeout: 120 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("请求 LLM 失败: %w", err)
	}
	defer resp.Body.Close()

	respData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %w", err)
	}

	if resp.StatusCode != 200 {
		var errResp OpenAIErrorResponse
		if json.Unmarshal(respData, &errResp) == nil && errResp.Error.Message != "" {
			return nil, fmt.Errorf("LLM API 错误 [%s]: %s", errResp.Error.Type, errResp.Error.Message)
		}
		return nil, fmt.Errorf("LLM API 返回 HTTP %d: %s", resp.StatusCode, string(respData))
	}

	var result OpenAIResponse
	if err := json.Unmarshal(respData, &result); err != nil {
		return nil, fmt.Errorf("解析响应失败: %w", err)
	}

	if len(result.Choices) == 0 {
		return nil, fmt.Errorf("LLM 返回空结果")
	}

	return &result, nil
}