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

// 第 3 篇演示：让 Agent 自己决定调几次工具（多轮循环版）。
// 与第 2 篇的差异见正文开头的「改动清单表」；本文件与第 2 篇同构，
// 新增部分集中在：工具注册表 toolEntry/binaryOp、runAgent 循环、错误回喂、trimHistory。
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"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 FunctionCall `json:"function"`
}

type Tool struct {
	Type     string      `json:"type"`
	Function FunctionDef `json:"function"`
}

type FunctionDef struct {
	Name        string         `json:"name"`
	Description string         `json:"description"`
	Parameters  map[string]any `json:"parameters"`
}

// FunctionCall 是一次工具调用里函数的形状（对应 tool_calls[].function）
type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"` // 模型给的是 JSON 字符串，执行前要再解析
}

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

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 + 实现
// ===========================================

// toolEntry 描述一个可被模型调用的工具。
// Run 收到的 args 是模型输出的 arguments JSON 字符串，需要自己解析。
type toolEntry 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(desc string, fn func(int, int) (int, error)) *toolEntry {
	return &toolEntry{
		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]*toolEntry{
	"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("计算两个整数 a 与 b 的和",
		func(a, b int) (int, error) { return a + b, nil }),
	"multiply": binaryOp("计算两个整数 a 与 b 的乘积",
		func(a, b int) (int, error) { return a * b, nil }),
	"divide": binaryOp("计算两个整数 a 除以 b 的商（整除）",
		func(a, b int) (int, error) {
			if b == 0 {
				return 0, fmt.Errorf("除数不能为 0") // 刻意制造一次"工具执行失败"，演示错误回喂
			}
			return a / b, nil
		}),
}

// toolDefs 把注册表转成请求里的 tools 数组。
// 按名字排序：Go 的 map 遍历顺序是随机的，排序后请求体稳定，
// 便于对比抓包结果，也更容易命中上游的 prompt 缓存。
func toolDefs() []Tool {
	names := make([]string, 0, len(tools))
	for name := range tools {
		names = append(names, name)
	}
	sort.Strings(names)

	out := make([]Tool, 0, len(names))
	for _, name := range names {
		t := tools[name]
		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 只作校验——不同上游在
		// 调用工具时可能仍返回 "stop" 或干脆不带该字段，拿它当唯一判据会
		// 静默丢掉工具调用、只打印一个空回答。
		if len(msg.ToolCalls) == 0 {
			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
		}
		if finish != "tool_calls" {
			fmt.Printf("   ⚠️ finish_reason=%q 但本轮有 tool_calls，按 tool_calls 继续执行\n", finish)
		}
		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)
			}
			if result == "" {
				result = "(工具无输出)" // role="tool" 消息必须带 content，否则上游 400（见第 2 篇 §3）
			}
			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))
}

// trimThreshold 是裁剪阈值（prompt tokens）。取 4000 是有来历的：Ollama 默认
// 上下文窗口就是 4k（`ollama ps` 的 CONTEXT 列实测 4096；可用 OLLAMA_CONTEXT_LENGTH
// 调大）。留出的 96 不是随手写的余量——它要覆盖本轮输出与模板/系统提示的开销，
// 阈值贴着"当前窗口"取才有意义：窗口调大后这个数要同步放大。
const trimThreshold = 4000

// trimHistory：极简历史裁剪——当某轮 prompt 已经很长时，丢掉最早的
// 一轮完整对话（该轮 user 及其后的 assistant 与 role=tool 消息），
// 保留最近的上下文。生产中一般按 token 阈值触发并配合摘要压缩，
// 这里只演示思路。
func trimHistory(messages []Message, lastPromptTokens int) []Message {
	if lastPromptTokens < trimThreshold || 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,
		Temperature: 0, // 贪心解码：演示输出可复现；聊天场景可调回 0.7
	}
	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
}