Update main.go

main.go

@@ -61,14 +61,13 @@ 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"`
+	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 {
@@ -78,9 +77,9 @@ type RawChoice struct {
 }
 
 type RawDelta struct {
-	Role      string           `json:"role,omitempty"`
-	Content   *string          `json:"content,omitempty"`
-	ToolCalls []RawToolCall    `json:"tool_calls,omitempty"`
+	Role      string       `json:"role,omitempty"`
+	Content   *string      `json:"content,omitempty"`
+	ToolCalls []RawToolCall `json:"tool_calls,omitempty"`
 }
 
 type RawToolCall struct {
@@ -170,108 +169,6 @@ 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)
@@ -281,6 +178,7 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, `{"error":{"message":"Method not allowed"}}`, http.StatusMethodNotAllowed)
 		return
 	}
+
 	var req ChatRequest
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		http.Error(w, `{"error":{"message":"Invalid request body"}}`, http.StatusBadRequest)
@@ -288,28 +186,54 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 	}
 
 	modelID := resolveModel(req.Model)
-	messages := injectSystemPrompt(req.Messages, modelID)
-
-	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,
-		}
-		if modelID == "z-ai/glm4.7" {
-			u.ExtraBody = map[string]interface{}{
-				"chat_template_kwargs": map[string]interface{}{
-					"enable_thinking": false,
-				},
-			}
+	upstream := UpstreamRequest{
+		Model:       modelID,
+		Messages:    injectSystemPrompt(req.Messages, modelID),
+		Stream:      true,
+		Tools:       req.Tools,
+		ToolChoice:  req.ToolChoice,
+		Temperature: req.Temperature,
+		MaxTokens:   req.MaxTokens,
+		TopP:        req.TopP,
+		Stop:        req.Stop,
+	}
+	if modelID == "z-ai/glm4.7" {
+		upstream.ExtraBody = map[string]interface{}{
+			"chat_template_kwargs": map[string]interface{}{
+				"enable_thinking": false,
+			},
 		}
-		return u
+	}
+
+	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":"%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
 	}
 
 	w.Header().Set("Content-Type", "text/event-stream")
@@ -326,36 +250,90 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	// 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
-		}
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
 
-		if resp.StatusCode != http.StatusOK {
-			body, _ := io.ReadAll(resp.Body)
-			resp.Body.Close()
-			emit("data: " + string(body) + "\n\n")
-			return
+	// Accumulate tool_calls across delta chunks, stream content chunks immediately
+	accum := make(map[int]*AccumToolCall)
+	var lastChunk RawChunk
+
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if !strings.HasPrefix(line, "data: ") {
+			emit(line + "\n")
+			continue
 		}
 
-		chunks, accum, _ := collectStream(resp.Body)
-		resp.Body.Close()
+		data := strings.TrimPrefix(line, "data: ")
+
+		if data == "[DONE]" {
+			emit("data: [DONE]\n\n")
+			continue
+		}
 
-		if len(chunks) == 0 {
-			break
+		var chunk RawChunk
+		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+			emit(line + "\n")
+			continue
 		}
+		lastChunk = chunk
 
-		lastChunk := chunks[len(chunks)-1]
+		isToolChunk := false
+		isFinishToolCalls := false
 
-		if hasToolCallsInChunks(chunks) && len(req.Tools) > 0 {
-			// Emit tool_calls chunk to client in OpenAI format
-			assembled := assembleToolCalls(accum)
+		for _, choice := range chunk.Choices {
+			if len(choice.Delta.ToolCalls) > 0 {
+				isToolChunk = true
+				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
+				}
+			}
+			if choice.FinishReason != nil && *choice.FinishReason == "tool_calls" {
+				isFinishToolCalls = true
+			}
+		}
+
+		if isFinishToolCalls {
+			// Emit one complete tool_calls chunk with all assembled tool calls
+			indices := make([]int, 0, len(accum))
+			for idx := range accum {
+				indices = append(indices, idx)
+			}
+			sort.Ints(indices)
+
+			assembled := make([]map[string]interface{}, 0, len(indices))
+			for _, idx := range indices {
+				acc := accum[idx]
+				tcType := acc.Type
+				if tcType == "" {
+					tcType = "function"
+				}
+				assembled = append(assembled, map[string]interface{}{
+					"index": idx,
+					"id":    acc.ID,
+					"type":  tcType,
+					"function": map[string]string{
+						"name":      acc.Name,
+						"arguments": acc.Args,
+					},
+				})
+			}
 
 			fr := "tool_calls"
-			toolChunk := map[string]interface{}{
+			out, _ := json.Marshal(map[string]interface{}{
 				"id":      lastChunk.ID,
 				"object":  "chat.completion.chunk",
 				"created": lastChunk.Created,
@@ -371,37 +349,20 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 						"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
+			emit("data: " + string(out) + "\n\n")
+			accum = make(map[int]*AccumToolCall)
+			continue
 		}
 
-		// 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")
+		// Skip intermediate tool_call delta chunks (already accumulating)
+		if isToolChunk {
+			continue
 		}
-		break
-	}
 
-	emit("data: [DONE]\n\n")
+		// Regular content chunk — stream immediately as-is
+		emit("data: " + data + "\n\n")
+	}
 }
 
 func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {