循环：让 Agent 自己决定调几次工具

// 第 3 篇演示：让 Agent 自己决定调几次工具（多轮循环版）
//
// 与第 2 篇最小案例的差异：
//  1. 多个工具（时间 / 加法 / 乘法 / 除法），按 JSON Schema 声明参数；
//  2. while 循环：只要模型还在请求工具就继续，直到它输出普通回答；
//  3. 一次响应可能带多个 tool_calls（并行调用），全部执行后一起回传；
//  4. 工具执行失败时把错误信息作为 role="tool" 内容回喂，让模型自行修正或放弃；
//  5. 简单的历史裁剪（防止上下文无限增长）。
package main

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

// ===========================================
// 与 Ollama /v1/chat/completions 对应的结构
// ===========================================

type Message struct {
	Role       string     `json:"role"`
	Content    string     `json:"content,omitempty"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
}

type ToolCall struct {
	ID       string `json:"id"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"` // 模型给的是 JSON 字符串，执行前要再解析
	} `json:"function"`
}

type Tool struct {
	Type     string `json:"type"`
	Function struct {
		Name        string         `json:"name"`
		Description string         `json:"description"`
		Parameters  map[string]any `json:"parameters"`
	} `json:"function"`
}

type ChatRequest struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
	Tools    []Tool    `json:"tools,omitempty"` // 空则不传，遵守第 2 篇的兼容铁律
	Stream   bool      `json:"stream"`
}

type Choice struct {
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

type ChatResponse struct {
	Choices []Choice `json:"choices"`
	Usage   struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
	} `json:"usage"`
}

// ===========================================
// 工具注册表：名字 -> 描述 + 参数 Schema + 实现
// ===========================================

// tool 描述一个可被模型调用的工具。
// Run 收到的 args 是模型输出的 arguments JSON 字符串，需要自己解析。
type tool struct {
	description string
	parameters  map[string]any
	run         func(args json.RawMessage) (string, error)
}

// binaryOp 提取了三个算术工具的公共模式：解析 {a,b} 整数参数 → 执行运算 → 返回结果。
// 每个工具只需传入一行运算函数（如 func(a,b int)(int,error){ return a+b, nil }）。
func binaryOp(name, desc string, fn func(int, int) (int, error)) *tool {
	return &tool{
		description: desc,
		parameters: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"a": map[string]any{"type": "integer"},
				"b": map[string]any{"type": "integer"},
			},
			"required": []string{"a", "b"},
		},
		run: func(args json.RawMessage) (string, error) {
			var p struct{ A, B int }
			if err := json.Unmarshal(args, &p); err != nil {
				return "", fmt.Errorf("参数解析失败（应为 {\"a\":整数,\"b\":整数}）：%v", err)
			}
			v, err := fn(p.A, p.B)
			if err != nil {
				return "", err
			}
			return fmt.Sprintf("%d", v), nil
		},
	}
}

// tools 是全局注册表，新增工具只需在这里加一项。
var tools = map[string]*tool{
	"get_current_time": {
		description: "获取当前的日期和时间",
		parameters:  map[string]any{"type": "object", "properties": map[string]any{}},
		run: func(args json.RawMessage) (string, error) {
			return time.Now().Format("2006-01-02 15:04:05"), nil
		},
	},
	"add": binaryOp("add", "计算两个整数 a 与 b 的和",
		func(a, b int) (int, error) { return a + b, nil }),
	"multiply": binaryOp("multiply", "计算两个整数 a 与 b 的乘积",
		func(a, b int) (int, error) { return a * b, nil }),
	"divide": binaryOp("divide", "计算两个整数 a 除以 b 的商（整除）",
		func(a, b int) (int, error) {
			if b == 0 {
				return 0, fmt.Errorf("除数不能为 0") // 刻意制造一次"工具执行失败"，演示错误回喂
			}
			return a / b, nil
		}),
}

// toolDefs 把注册表转成请求里的 tools 数组。
func toolDefs() []Tool {
	var out []Tool
	for name, t := range tools {
		var td Tool
		td.Type = "function"
		td.Function.Name = name
		td.Function.Description = t.description
		td.Function.Parameters = t.parameters
		out = append(out, td)
	}
	return out
}

// ===========================================
// Agent 循环
// ===========================================

const maxRounds = 8 // 保险丝：防止模型陷入"调用-报错-再调用"的死循环

func main() {
	if len(os.Args) > 1 {
		runAgent(strings.Join(os.Args[1:], " "))
		return
	}
	// 默认问题：多轮 + 刻意触发一次除零错误 + 顺带取时间
	runAgent("请先计算 10 除以 0，再计算 10 除以 2，然后把两个结果与当前时间一起告诉我。")
}

func runAgent(userPrompt string) {
	messages := []Message{{Role: "user", Content: userPrompt}}
	fmt.Println("🧑 用户:", userPrompt)
	fmt.Println(strings.Repeat("-", 56))

	for round := 1; round <= maxRounds; round++ {
		resp, err := chat(messages)
		if err != nil {
			fmt.Println("❌ 请求失败:", err)
			return
		}
		if len(resp.Choices) == 0 {
			fmt.Println("❌ 无响应")
			return
		}

		msg := resp.Choices[0].Message
		finish := resp.Choices[0].FinishReason
		messages = append(messages, msg)

		// 铁律：以 tool_calls 非空 / finish_reason=="tool_calls" 为准
		if len(msg.ToolCalls) == 0 || finish != "tool_calls" {
			fmt.Printf("🤖 第 %d 轮 最终回答: %s\n", round, msg.Content)
			fmt.Printf("📊 tokens: prompt=%d completion=%d total=%d\n",
				resp.Usage.PromptTokens, resp.Usage.CompletionTokens,
				resp.Usage.PromptTokens+resp.Usage.CompletionTokens)
			return
		}

		fmt.Printf("🔧 第 %d 轮 模型请求 %d 个工具:\n", round, len(msg.ToolCalls))
		for _, tc := range msg.ToolCalls {
			result, err := runTool(tc)
			if err != nil {
				result = fmt.Sprintf("工具执行出错：%v。请修正参数后重试，或放弃这一步。", err)
			}
			fmt.Printf("   - %s(%s) -> %s\n", tc.Function.Name, tc.Function.Arguments, result)
			// 结果必须以 role="tool" + 对应 ID 回传（第 2 篇的机制）
			messages = append(messages, Message{
				Role:       "tool",
				ToolCallID: tc.ID,
				Content:    result,
			})
		}
		messages = trimHistory(messages, resp.Usage.PromptTokens)
	}
	fmt.Println("⚠️ 达到最大轮数，可能陷入循环。")
}

func runTool(tc ToolCall) (string, error) {
	t, ok := tools[tc.Function.Name]
	if !ok {
		return "", fmt.Errorf("未知工具 %q（模型幻觉了工具名）", tc.Function.Name)
	}
	return t.run(json.RawMessage(tc.Function.Arguments))
}

// trimHistory：极简历史裁剪——当某轮 prompt 已经很长时，丢掉最早的
// 一轮完整对话（该轮 user 及其后的 assistant 与 role=tool 消息），
// 保留最近的上下文。生产中一般按 token 阈值触发并配合摘要压缩，
// 这里只演示思路。
func trimHistory(messages []Message, lastPromptTokens int) []Message {
	if lastPromptTokens < 4000 || len(messages) <= 4 {
		return messages
	}
	// 定位最早一轮的边界：从第一条 user 开始，到下一个 user 之前结束
	// （期间是这条 user 引发的 assistant 与 role=tool 消息，必须整组删除，
	// 否则会留下无 assistant 对应的孤儿 tool 消息，上游会校验失败）。
	start := 0
	for start < len(messages) && messages[start].Role != "user" {
		start++
	}
	end := len(messages)
	for i := start + 1; i < len(messages); i++ {
		if messages[i].Role == "user" {
			end = i
			break
		}
	}
	if start >= len(messages) || end == len(messages) {
		// 找不到第二条 user：当前只有一轮对话在进行中（如本 demo 的
		// 单个问题多轮调工具），此时不裁剪，避免删掉正在使用的上下文。
		return messages
	}
	fmt.Printf("   ✂️ 历史已超过 %d tokens，丢弃最早一轮对话（%d 条消息）\n", lastPromptTokens, end-start)
	return append([]Message{}, messages[end:]...)
}

// ===========================================
// HTTP（与 Ollama 通信）
// ===========================================

var httpClient = &http.Client{Timeout: 5 * time.Minute} // 请求超时：上游卡住时不至于无限挂起

func chat(messages []Message) (ChatResponse, error) {
	reqBody := ChatRequest{
		Model:    "llama3.1:8b",
		Messages: messages,
		Tools:    toolDefs(),
		Stream:   false,
	}
	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return ChatResponse{}, err
	}
	resp, err := httpClient.Post("http://localhost:11434/v1/chat/completions", "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return ChatResponse{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return ChatResponse{}, fmt.Errorf("upstream %d: %s", resp.StatusCode, string(b))
	}
	body, _ := io.ReadAll(resp.Body)
	var result ChatResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return ChatResponse{}, fmt.Errorf("解析失败，原始响应: %s", string(body))
	}
	return result, nil
}