服务与迁移：把 Agent 变成流式 API，可切云端

// 第 4 篇演示：把 Agent 循环变成常驻 HTTP 服务（SSE 流式 + 可切云端）。
// 与第 3 篇的差异见正文改动清单表；核心是 client（OLLAMA_BASE 可切云端）、
// streamChat（按 index 累积分片 arguments）与 server（把循环输出转成 SSE 事件）。
//
// 运行：
//
//	OLLAMA_BASE="http://localhost:11434/v1" OLLAMA_MODEL=llama3.1:8b go run main.go
//	默认监听 :8899。
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

// ===========================================
// 消息模型（与 /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"`
	Type     string       `json:"type"`
	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 的含义见第 2 篇 §3（tool_calls[].function）
type FunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// ChatRequest 与第 3 篇一致（每篇都是自包含程序，所以这里重新声明一份）。
// 用结构体而不是内联 map：与第 2/3/5 篇保持同一种写法，读者不必猜"服务端是不是特殊处理"。
type ChatRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Tools       []Tool    `json:"tools,omitempty"`
	Stream      bool      `json:"stream"`
	Temperature float64   `json:"temperature"`
}

// ===========================================
// 流式 chunk 的解析结构
// ===========================================

type streamChunk struct {
	Choices []struct {
		Delta struct {
			Role      string `json:"role"`
			Content   string `json:"content"`
			ToolCalls []struct {
				Index    int          `json:"index"`
				ID       string       `json:"id"`
				Type     string       `json:"type"`
				Function FunctionCall `json:"function"`
			} `json:"tool_calls"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
}

// pendingCall 累积一个"还没发完"的工具调用（按 index 对齐）。
type pendingCall struct {
	index     int
	id        string
	name      string
	arguments strings.Builder
}

// ===========================================
// 工具注册表（与第 3 篇相同，省注释）
// ===========================================

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
		},
	}
}

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, errors.New("除数不能为 0")
			}
			return a / b, nil
		}),
}

// toolDefs 与第 3 篇相同（按名字排序，保证请求体稳定，理由见第 3 篇 §2）。
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
}

// ===========================================
// OpenAI 兼容客户端（OLLAMA_BASE 可切本地/云端）
// ===========================================

type client struct {
	baseURL string // 例如 http://localhost:11434/v1 或 https://api.openai.com/v1
	apiKey  string
	model   string
	http    *http.Client
}

func newClient() *client {
	base := os.Getenv("OLLAMA_BASE")
	if base == "" {
		base = "http://localhost:11434/v1"
	}
	key := os.Getenv("OLLAMA_API_KEY")
	if key == "" {
		key = "ollama" // Ollama 忽略 key，OpenAI 需要填真 key
	}
	model := os.Getenv("OLLAMA_MODEL")
	if model == "" {
		model = "llama3.1:8b"
	}
	return &client{baseURL: base, apiKey: key, model: model,
		http: &http.Client{Timeout: 5 * time.Minute}}
}

// streamChat 发一次流式请求，把模型输出累积成一条 assistant 消息。
// callback 用于把每个 token 实时推给下游（SSE）。
// ctx 从 handleChat 一路传下来：客户端断开或超时，上游请求随即取消
// （http.Client.Timeout 管"最长多久"，ctx 管"调用方还要不要"）。
func (c *client) streamChat(ctx context.Context, messages []Message, onDelta func(string)) (Message, string, error) {
	body, _ := json.Marshal(ChatRequest{
		Model:       c.model,
		Messages:    messages,
		Tools:       toolDefs(),
		Stream:      true,
		Temperature: 0, // 贪心解码：演示输出可复现（聊天场景可调回 0.7）
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return Message{}, "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.http.Do(req)
	if err != nil {
		return Message{}, "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return Message{}, "", fmt.Errorf("upstream %d: %s", resp.StatusCode, string(b))
	}

	var (
		acc       strings.Builder // 累积本轮 content（最终回答）
		calls     []*pendingCall  // 累积本轮 tool_calls（按 index）
		finish    string
		callByIdx = map[int]*pendingCall{}
	)

	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 1024), 1<<20)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if data == "[DONE]" {
			break
		}
		var ch streamChunk
		if err := json.Unmarshal([]byte(data), &ch); err != nil {
			continue // 忽略无法解析的行（兼容性防呆）
		}
		for _, choice := range ch.Choices {
			if choice.FinishReason != "" {
				finish = choice.FinishReason
			}
			d := choice.Delta
			if d.Content != "" {
				acc.WriteString(d.Content)
				onDelta(d.Content)
			}
			for _, tc := range d.ToolCalls {
				p := callByIdx[tc.Index]
				if p == nil {
					p = &pendingCall{index: tc.Index}
					callByIdx[tc.Index] = p
					calls = append(calls, p)
				}
				if tc.ID != "" {
					p.id = tc.ID
				}
				if tc.Function.Name != "" {
					p.name = tc.Function.Name // 分片场景下若拆名则追加，此处取简版
				}
				p.arguments.WriteString(tc.Function.Arguments) // arguments 可能被拆成多片，必须拼接
			}
		}
	}

	// 循环结束后检查 Scanner 是否有错
	if err := sc.Err(); err != nil {
		return Message{}, "", fmt.Errorf("读取流失败: %w", err)
	}

	msg := Message{Role: "assistant", Content: acc.String()}
	for _, p := range calls {
		var tc ToolCall
		tc.ID = p.id
		tc.Type = "function"
		tc.Function.Name = p.name
		tc.Function.Arguments = p.arguments.String()
		msg.ToolCalls = append(msg.ToolCalls, tc)
	}
	return msg, finish, nil
}

// ===========================================
// HTTP 服务
// ===========================================

type server struct {
	client *client
}

func main() {
	s := &server{client: newClient()}

	mux := http.NewServeMux()
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(map[string]string{"status": "ok", "model": s.client.model})
	})
	mux.HandleFunc("POST /chat", s.handleChat)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8899"
	}
	fmt.Printf("stream-server 已启动: http://localhost:%s  (base=%s model=%s)\n",
		port, s.client.baseURL, s.client.model)
	if err := http.ListenAndServe(":"+port, mux); err != nil {
		fmt.Println("启动失败:", err)
	}
}

// handleChat：多轮 Agent 循环 + SSE 实时推送。
func (s *server) handleChat(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Messages []Message `json:"messages"`
		Stream   bool      `json:"stream"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "bad json: "+err.Error(), http.StatusBadRequest)
		return
	}

	if req.Stream {
		s.runAgentSSE(r.Context(), w, req.Messages)
		return
	}
	// 非流式：跑同一套逻辑，把事件收进内存，最后返回 JSON
	events := &eventSink{}
	s.runAgent(r.Context(), events, req.Messages)
	answer, _ := events.lastAnswer()
	json.NewEncoder(w).Encode(map[string]any{"answer": answer, "events": events.list})
}

type event struct {
	Type string `json:"type"`
	Data string `json:"data"`
}

// eventSink 收集事件（非流式路径用）。
type eventSink struct{ list []event }

func (e *eventSink) emit(t, d string) { e.list = append(e.list, event{t, d}) }
func (e *eventSink) lastAnswer() (string, bool) {
	for i := len(e.list) - 1; i >= 0; i-- {
		if e.list[i].Type == "answer" {
			return e.list[i].Data, true
		}
	}
	return "", false
}

// sseWriter 直接把事件写给浏览器（流式路径用）。
type sseWriter struct {
	w http.ResponseWriter
	f http.Flusher
}

// emit 写一个 SSE 事件。data 可能含换行（模型的段落/空行 token），必须按行
// 拆成多条 data: 行：SSE 规范中连续 data 行会以 \n 重新拼接，这样既保持帧
// 合法，又能逐字保留换行（若把 \n 直接写进单行 data，空行会提前终止事件、
// 后续裸行会被客户端丢弃）。
func (s *sseWriter) emit(t, d string) {
	fmt.Fprintf(s.w, "event: %s\n", t)
	for _, line := range strings.Split(d, "\n") {
		fmt.Fprintf(s.w, "data: %s\n", line)
	}
	fmt.Fprint(s.w, "\n")
	s.f.Flush()
}

// emitter 是两种输出目标的共同接口。
type emitter interface{ emit(typ, data string) }

// runAgent 是核心多轮循环：模型输出逐 token 推送；工具调用执行后推送结果。
func (s *server) runAgent(ctx context.Context, em emitter, history []Message) {
	const maxRounds = 8
	for round := 1; round <= maxRounds; round++ {
		msg, finish, err := s.client.streamChat(ctx, history, func(token string) {
			em.emit("delta", token) // 每个 token 实时推送
		})
		if err != nil {
			em.emit("error", err.Error())
			return
		}
		history = append(history, msg)

		// 铁律：以 tool_calls 非空为准（第 2/3 篇一致）。finish_reason 只作校验：
		// 部分上游调用工具时仍返回 "stop" 或缺失，拿它当唯一判据会静默丢掉工具调用。
		if len(msg.ToolCalls) == 0 {
			em.emit("answer", msg.Content)
			return
		}
		roundMsg := fmt.Sprintf("第 %d 轮：模型请求 %d 个工具", round, len(msg.ToolCalls))
		if finish != "tool_calls" {
			roundMsg += fmt.Sprintf("（finish_reason=%q，仍按 tool_calls 执行）", finish)
		}
		em.emit("round", roundMsg)
		for _, tc := range msg.ToolCalls {
			result, err := runTool(tc)
			if err != nil {
				result = fmt.Sprintf("工具执行出错：%v。请修正参数后重试，或放弃这一步。", err)
			}
			em.emit("tool", fmt.Sprintf("%s(%s) -> %s", tc.Function.Name, tc.Function.Arguments, result))
			history = append(history, Message{Role: "tool", ToolCallID: tc.ID, Content: result})
		}
	}
	em.emit("error", "达到最大轮数")
}

func (s *server) runAgentSSE(ctx context.Context, w http.ResponseWriter, history []Message) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	f, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming unsupported", http.StatusInternalServerError)
		return
	}
	sw := &sseWriter{w: w, f: f}
	s.runAgent(ctx, sw, history)
	sw.emit("done", "[DONE]")
}

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))
}