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

// 第 4 篇演示：把 Agent 循环变成常驻 HTTP 服务（SSE 流式 + 可切云端）
//
// 设计要点：
//  1. 只实现一个 OpenAI 兼容客户端（/v1/chat/completions），通过
//     BASE_URL 指向本地 Ollama 或任意 OpenAI 兼容云服务——同一套代码可切换；
//  2. 用流式（stream:true）做多轮工具循环：每一轮的模型输出实时
//     以 SSE 转发给浏览器/客户端，工具调用与结果也以事件形式发出；
//  3. delta 解析做到"防呆"：content 可能为 "" 或 null、tool_calls 的
//     arguments 可能被拆成多个分片（OpenAI 会拆，Ollama 一次性给全），
//     按 index 累积拼接，两种服务端都能正确处理。
//
// 运行：
//
//	OLLAMA_BASE="http://localhost:11434/v1" OLLAMA_MODEL=llama3.1:8b go run ./demo/stream-server
//	默认监听 :8899。
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"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 struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `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"`
}

// ===========================================
// 流式 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 struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `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 tool struct {
	description string
	parameters  map[string]any
	run         func(args json.RawMessage) (string, error)
}

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": {
		description: "计算两个整数 a 与 b 的和",
		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: intTool(func(a, b int) (string, error) { return fmt.Sprintf("%d", a+b), nil }),
	},
	"multiply": {
		description: "计算两个整数 a 与 b 的乘积",
		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: intTool(func(a, b int) (string, error) { return fmt.Sprintf("%d", a*b), nil }),
	},
	"divide": {
		description: "计算两个整数 a 除以 b 的商（整除）",
		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: intTool(func(a, b int) (string, error) {
			if b == 0 {
				return "", errors.New("除数不能为 0")
			}
			return fmt.Sprintf("%d", a/b), nil
		}),
	},
}

func intTool(fn func(a, b int) (string, error)) func(json.RawMessage) (string, error) {
	return func(args json.RawMessage) (string, error) {
		var p struct {
			A int `json:"a"`
			B int `json:"b"`
		}
		if err := json.Unmarshal(args, &p); err != nil {
			return "", fmt.Errorf("参数解析失败（应为 {\"a\":整数,\"b\":整数}）：%v", err)
		}
		return fn(p.A, p.B)
	}
}

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
}

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

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）。
func (c *client) streamChat(messages []Message, onDelta func(string)) (Message, string, error) {
	body, _ := json.Marshal(map[string]any{
		"model":    c.model,
		"messages": messages,
		"tools":    toolDefs(),
		"stream":   true,
	})
	req, err := http.NewRequest(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 可能被拆成多片，必须拼接
			}
		}
	}

	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(w, req.Messages)
		return
	}
	// 非流式：跑同一套逻辑，把事件收进内存，最后返回 JSON
	events := &eventSink{}
	s.runAgent(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(em emitter, history []Message) {
	const maxRounds = 8
	for round := 1; round <= maxRounds; round++ {
		msg, finish, err := s.client.streamChat(history, func(token string) {
			em.emit("delta", token) // 每个 token 实时推送
		})
		if err != nil {
			em.emit("error", err.Error())
			return
		}
		history = append(history, msg)

		if len(msg.ToolCalls) == 0 || finish != "tool_calls" {
			em.emit("answer", msg.Content)
			return
		}
		em.emit("round", fmt.Sprintf("第 %d 轮：模型请求 %d 个工具", round, len(msg.ToolCalls)))
		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(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(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))
}