Service & Migration: Turning the Agent into a Streaming API You Can Move to the Cloud (Part 4)

// Part 4 demo: turning the agent loop into a long-running HTTP service (SSE streaming + cloud-switchable)
//
// Design points:
//  1. only one OpenAI-compatible client (/v1/chat/completions) is implemented, and
//     BASE_URL points at local Ollama or any OpenAI-compatible cloud service — the
//     same code can switch between them;
//  2. streaming (stream:true) drives the multi-round tool loop: each round's model
//     output is forwarded to the browser/client over SSE in real time, and tool
//     calls and their results are emitted as events too;
//  3. delta parsing is defensive: content may be "" or null, and a tool call's
//     arguments may be split into several fragments (OpenAI splits them, Ollama
//     sends them whole), so they are accumulated by index — both kinds of server
//     are handled correctly.
//
// Run:
//
//	OLLAMA_BASE="http://localhost:11434/v1" OLLAMA_MODEL=llama3.1:8b go run ./demo/stream-server
//	Listens on :8899 by default.
package main

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

// ===========================================
// Message model (matching /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"`
}

// ===========================================
// Structures for parsing streaming chunks
// ===========================================

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 accumulates one tool call that has not finished arriving (aligned by index).
type pendingCall struct {
	index     int
	id        string
	name      string
	arguments strings.Builder
}

// ===========================================
// Tool registry (same as Part 3, comments omitted)
// ===========================================

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-compatible client (BASE_URL switches local/cloud)
// ===========================================

type client struct {
	baseURL string // e.g. http://localhost:11434/v1 or 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 ignores the key; OpenAI needs a real one
	}
	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 sends one streaming request and accumulates the model output into
// a single assistant message.
// callback pushes every token downstream (SSE) in real time.
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 // accumulates this round's content (the final answer)
		calls     []*pendingCall  // accumulates this round's tool_calls (by 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 // ignore lines that cannot be parsed (compatibility guard)
		}
		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 // with fragmented names this should append; the simple version takes it as-is
				}
				p.arguments.WriteString(tc.Function.Arguments) // arguments may be split across fragments and must be concatenated
			}
		}
	}

	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 service
// ===========================================

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: the multi-round agent loop plus real-time SSE pushes.
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
	}
	// Non-streaming: run the same logic, collect the events in memory, return JSON at the end
	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 collects events (used by the non-streaming path).
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 writes events straight to the browser (used by the streaming path).
type sseWriter struct {
	w http.ResponseWriter
	f http.Flusher
}

// emit writes one SSE event. data may contain newlines (the model's paragraph
// and blank-line tokens), so it must be split into several data: lines: in the
// SSE spec consecutive data lines are rejoined with \n, which keeps the frame
// legal while preserving newlines exactly (writing \n directly into a single
// data line would let a blank line terminate the event early and the client
// would drop the bare lines that follow).
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 is the shared interface of both output targets.
type emitter interface{ emit(typ, data string) }

// runAgent is the core multi-round loop: model output is pushed token by token,
// and tool results are pushed once the tool has run.
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) // push every token in real time
		})
		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))
}