Update main.go

main.go

@@ -9,6 +9,8 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"sort"
+	"strings"
 	"time"
 )
 
@@ -59,6 +61,48 @@ type UpstreamRequest struct {
 	ExtraBody   map[string]interface{} `json:"extra_body,omitempty"`
 }
 
+// SSE chunk types for parsing upstream stream
+type RawChunk struct {
+	ID      string       `json:"id"`
+	Object  string       `json:"object"`
+	Created int64        `json:"created"`
+	Model   string       `json:"model"`
+	Choices []RawChoice  `json:"choices"`
+	Usage   interface{}  `json:"usage,omitempty"`
+}
+
+type RawChoice struct {
+	Index        int      `json:"index"`
+	Delta        RawDelta `json:"delta"`
+	FinishReason *string  `json:"finish_reason"`
+}
+
+type RawDelta struct {
+	Role      string           `json:"role,omitempty"`
+	Content   *string          `json:"content,omitempty"`
+	ToolCalls []RawToolCall    `json:"tool_calls,omitempty"`
+}
+
+type RawToolCall struct {
+	Index    int         `json:"index"`
+	ID       string      `json:"id,omitempty"`
+	Type     string      `json:"type,omitempty"`
+	Function RawFunction `json:"function"`
+}
+
+type RawFunction struct {
+	Name      string `json:"name,omitempty"`
+	Arguments string `json:"arguments,omitempty"`
+}
+
+type AccumToolCall struct {
+	Index int
+	ID    string
+	Type  string
+	Name  string
+	Args  string
+}
+
 func resolveModel(requested string) string {
 	if full, ok := modelAliases[requested]; ok {
 		return full
@@ -126,6 +170,108 @@ func handleBaseURL(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, `{"url":"https://%s/v1"}`, host)
 }
 
+func doUpstream(upstream UpstreamRequest) (*http.Response, error) {
+	body, err := json.Marshal(upstream)
+	if err != nil {
+		return nil, err
+	}
+	req, err := http.NewRequest(http.MethodPost, NvidiaBaseURL+"/chat/completions", bytes.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+NvidiaAPIKey)
+	req.Header.Set("Accept", "text/event-stream")
+	client := &http.Client{Timeout: 300 * time.Second}
+	return client.Do(req)
+}
+
+// collectStream reads SSE lines, accumulates tool_calls, returns all chunks + assembled tool calls
+func collectStream(body io.Reader) ([]RawChunk, map[int]*AccumToolCall, error) {
+	scanner := bufio.NewScanner(body)
+	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
+
+	var chunks []RawChunk
+	accum := make(map[int]*AccumToolCall)
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		if !strings.HasPrefix(line, "data: ") {
+			continue
+		}
+		data := strings.TrimPrefix(line, "data: ")
+		if data == "[DONE]" {
+			break
+		}
+		var chunk RawChunk
+		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+			continue
+		}
+		chunks = append(chunks, chunk)
+
+		for _, choice := range chunk.Choices {
+			for _, tc := range choice.Delta.ToolCalls {
+				acc, ok := accum[tc.Index]
+				if !ok {
+					acc = &AccumToolCall{Index: tc.Index}
+					accum[tc.Index] = acc
+				}
+				if tc.ID != "" {
+					acc.ID = tc.ID
+				}
+				if tc.Type != "" {
+					acc.Type = tc.Type
+				}
+				acc.Name += tc.Function.Name
+				acc.Args += tc.Function.Arguments
+			}
+		}
+	}
+	return chunks, accum, scanner.Err()
+}
+
+// hasToolCalls returns true if any chunk contains tool_calls
+func hasToolCallsInChunks(chunks []RawChunk) bool {
+	for _, chunk := range chunks {
+		for _, choice := range chunk.Choices {
+			if len(choice.Delta.ToolCalls) > 0 {
+				return true
+			}
+			if choice.FinishReason != nil && *choice.FinishReason == "tool_calls" {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func assembleToolCalls(accum map[int]*AccumToolCall) []map[string]interface{} {
+	indices := make([]int, 0, len(accum))
+	for idx := range accum {
+		indices = append(indices, idx)
+	}
+	sort.Ints(indices)
+
+	result := make([]map[string]interface{}, 0, len(indices))
+	for _, idx := range indices {
+		acc := accum[idx]
+		tcType := acc.Type
+		if tcType == "" {
+			tcType = "function"
+		}
+		result = append(result, map[string]interface{}{
+			"index": idx,
+			"id":    acc.ID,
+			"type":  tcType,
+			"function": map[string]string{
+				"name":      acc.Name,
+				"arguments": acc.Args,
+			},
+		})
+	}
+	return result
+}
+
 func handleChat(w http.ResponseWriter, r *http.Request) {
 	if !authenticate(r) {
 		http.Error(w, `{"error":{"message":"Unauthorized"}}`, http.StatusUnauthorized)
@@ -142,58 +288,28 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 	}
 
 	modelID := resolveModel(req.Model)
-	req.Messages = injectSystemPrompt(req.Messages, modelID)
-
-	upstream := UpstreamRequest{
-		Model:       modelID,
-		Messages:    req.Messages,
-		Stream:      true,
-		Tools:       req.Tools,
-		ToolChoice:  req.ToolChoice,
-		Temperature: req.Temperature,
-		MaxTokens:   req.MaxTokens,
-		TopP:        req.TopP,
-		Stop:        req.Stop,
-	}
+	messages := injectSystemPrompt(req.Messages, modelID)
 
-	// GLM-4.7 requires thinking disabled via extra_body
-	if modelID == "z-ai/glm4.7" {
-		upstream.ExtraBody = map[string]interface{}{
-			"chat_template_kwargs": map[string]interface{}{
-				"enable_thinking": false,
-			},
+	buildUpstream := func(msgs []Message) UpstreamRequest {
+		u := UpstreamRequest{
+			Model:       modelID,
+			Messages:    msgs,
+			Stream:      true,
+			Tools:       req.Tools,
+			ToolChoice:  req.ToolChoice,
+			Temperature: req.Temperature,
+			MaxTokens:   req.MaxTokens,
+			TopP:        req.TopP,
+			Stop:        req.Stop,
 		}
-	}
-
-	body, err := json.Marshal(upstream)
-	if err != nil {
-		http.Error(w, `{"error":{"message":"Failed to marshal request"}}`, http.StatusInternalServerError)
-		return
-	}
-
-	upstreamReq, err := http.NewRequest(http.MethodPost, NvidiaBaseURL+"/chat/completions", bytes.NewReader(body))
-	if err != nil {
-		http.Error(w, `{"error":{"message":"Failed to create upstream request"}}`, http.StatusInternalServerError)
-		return
-	}
-	upstreamReq.Header.Set("Content-Type", "application/json")
-	upstreamReq.Header.Set("Authorization", "Bearer "+NvidiaAPIKey)
-	upstreamReq.Header.Set("Accept", "text/event-stream")
-
-	client := &http.Client{Timeout: 300 * time.Second}
-	resp, err := client.Do(upstreamReq)
-	if err != nil {
-		http.Error(w, fmt.Sprintf(`{"error":{"message":"Upstream error: %s"}}`, err.Error()), http.StatusBadGateway)
-		return
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		upstreamBody, _ := io.ReadAll(resp.Body)
-		w.Header().Set("Content-Type", "application/json")
-		w.WriteHeader(resp.StatusCode)
-		w.Write(upstreamBody)
-		return
+		if modelID == "z-ai/glm4.7" {
+			u.ExtraBody = map[string]interface{}{
+				"chat_template_kwargs": map[string]interface{}{
+					"enable_thinking": false,
+				},
+			}
+		}
+		return u
 	}
 
 	w.Header().Set("Content-Type", "text/event-stream")
@@ -203,20 +319,89 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 
 	flusher, canFlush := w.(http.Flusher)
-	scanner := bufio.NewScanner(resp.Body)
-	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
-
-	for scanner.Scan() {
-		line := scanner.Text()
-		if line != "" {
-			fmt.Fprintf(w, "%s\n", line)
-		} else {
-			fmt.Fprintf(w, "\n")
-		}
+	emit := func(s string) {
+		fmt.Fprint(w, s)
 		if canFlush {
 			flusher.Flush()
 		}
 	}
+
+	// Agentic loop — handle tool_calls rounds
+	for {
+		resp, err := doUpstream(buildUpstream(messages))
+		if err != nil {
+			emit(fmt.Sprintf("data: {\"error\":{\"message\":\"%s\"}}\n\n", err.Error()))
+			return
+		}
+
+		if resp.StatusCode != http.StatusOK {
+			body, _ := io.ReadAll(resp.Body)
+			resp.Body.Close()
+			emit("data: " + string(body) + "\n\n")
+			return
+		}
+
+		chunks, accum, _ := collectStream(resp.Body)
+		resp.Body.Close()
+
+		if len(chunks) == 0 {
+			break
+		}
+
+		lastChunk := chunks[len(chunks)-1]
+
+		if hasToolCallsInChunks(chunks) && len(req.Tools) > 0 {
+			// Emit tool_calls chunk to client in OpenAI format
+			assembled := assembleToolCalls(accum)
+
+			fr := "tool_calls"
+			toolChunk := map[string]interface{}{
+				"id":      lastChunk.ID,
+				"object":  "chat.completion.chunk",
+				"created": lastChunk.Created,
+				"model":   req.Model,
+				"choices": []map[string]interface{}{
+					{
+						"index": 0,
+						"delta": map[string]interface{}{
+							"role":       "assistant",
+							"content":    nil,
+							"tool_calls": assembled,
+						},
+						"finish_reason": fr,
+					},
+				},
+			}
+			out, _ := json.Marshal(toolChunk)
+			emit("data: " + string(out) + "\n\n")
+
+			// Add assistant tool_calls message to history
+			messages = append(messages, Message{
+				Role:      "assistant",
+				Content:   nil,
+				ToolCalls: assembled,
+			})
+
+			// Add placeholder tool results — client must re-call with results
+			// For now signal finish so client can handle tool execution
+			emit("data: [DONE]\n\n")
+			return
+		}
+
+		// No tool_calls — stream content chunks directly to client
+		for _, chunk := range chunks {
+			// Remap model alias
+			chunk.Model = req.Model
+			out, err := json.Marshal(chunk)
+			if err != nil {
+				continue
+			}
+			emit("data: " + string(out) + "\n\n")
+		}
+		break
+	}
+
+	emit("data: [DONE]\n\n")
 }
 
 func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {