Update main.go

main.go

@@ -9,7 +9,6 @@ import (
 	"log"
 	"net/http"
 	"os"
-	"strings"
 	"time"
 )
 
@@ -48,15 +47,22 @@ type ChatRequest struct {
 }
 
 type UpstreamRequest struct {
-	Model       string        `json:"model"`
-	Messages    []Message     `json:"messages"`
-	Stream      bool          `json:"stream"`
-	Tools       []interface{} `json:"tools,omitempty"`
-	ToolChoice  interface{}   `json:"tool_choice,omitempty"`
-	Temperature *float64      `json:"temperature,omitempty"`
-	MaxTokens   *int          `json:"max_tokens,omitempty"`
-	TopP        *float64      `json:"top_p,omitempty"`
-	Stop        interface{}   `json:"stop,omitempty"`
+	Model       string                 `json:"model"`
+	Messages    []Message              `json:"messages"`
+	Stream      bool                   `json:"stream"`
+	Tools       []interface{}          `json:"tools,omitempty"`
+	ToolChoice  interface{}            `json:"tool_choice,omitempty"`
+	Temperature *float64               `json:"temperature,omitempty"`
+	MaxTokens   *int                   `json:"max_tokens,omitempty"`
+	TopP        *float64               `json:"top_p,omitempty"`
+	Stop        interface{}            `json:"stop,omitempty"`
+	ExtraBody   map[string]interface{} `json:"extra_body,omitempty"`
+}
+
+type StreamChoice struct {
+	Index        int         `json:"index"`
+	Delta        StreamDelta `json:"delta"`
+	FinishReason *string     `json:"finish_reason"`
 }
 
 func resolveModel(requested string) string {
@@ -87,20 +93,15 @@ func injectSystemPrompt(messages []Message, modelID string) []Message {
 
 func authenticate(r *http.Request) bool {
 	auth := r.Header.Get("Authorization")
-	if len(auth) > 7 && auth[:7] == "Bearer " {
-		if auth[7:] == GatewayAPIKey {
-			return true
-		}
-	}
-	if r.Header.Get("x-api-key") == GatewayAPIKey {
+	if len(auth) > 7 && auth[:7] == "Bearer " && auth[7:] == GatewayAPIKey {
 		return true
 	}
-	return false
+	return r.Header.Get("x-api-key") == GatewayAPIKey
 }
 
 func handleModels(w http.ResponseWriter, r *http.Request) {
 	if !authenticate(r) {
-		http.Error(w, `{"error":{"message":"Unauthorized","type":"auth_error"}}`, http.StatusUnauthorized)
+		http.Error(w, `{"error":{"message":"Unauthorized"}}`, http.StatusUnauthorized)
 		return
 	}
 	type ModelObj struct {
@@ -116,61 +117,24 @@ func handleModels(w http.ResponseWriter, r *http.Request) {
 	models := ModelsResponse{Object: "list"}
 	now := time.Now().Unix()
 	for alias := range modelAliases {
-		models.Data = append(models.Data, ModelObj{
-			ID: alias, Object: "model", Created: now, OwnedBy: "nvidia",
-		})
+		models.Data = append(models.Data, ModelObj{ID: alias, Object: "model", Created: now, OwnedBy: "nvidia"})
 	}
 	w.Header().Set("Content-Type", "application/json")
 	json.NewEncoder(w).Encode(models)
 }
 
-// StreamChoice represents a single choice in a streaming chunk
-type StreamChoice struct {
-	Index        int             `json:"index"`
-	Delta        StreamDelta     `json:"delta"`
-	FinishReason *string         `json:"finish_reason"`
-}
-
-// StreamDelta is the delta object inside a streaming chunk
-type StreamDelta struct {
-	Role      string          `json:"role,omitempty"`
-	Content   *string         `json:"content,omitempty"`
-	ToolCalls []ToolCallChunk `json:"tool_calls,omitempty"`
-}
-
-// ToolCallChunk is a partial tool call in a streaming delta
-type ToolCallChunk struct {
-	Index    int              `json:"index"`
-	ID       string           `json:"id,omitempty"`
-	Type     string           `json:"type,omitempty"`
-	Function ToolCallFunction `json:"function,omitempty"`
-}
-
-type ToolCallFunction struct {
-	Name      string `json:"name,omitempty"`
-	Arguments string `json:"arguments,omitempty"`
-}
-
-// StreamChunk is a full SSE data chunk from upstream
-type StreamChunk struct {
-	ID      string         `json:"id"`
-	Object  string         `json:"object"`
-	Created int64          `json:"created"`
-	Model   string         `json:"model"`
-	Choices []StreamChoice `json:"choices"`
-}
-
-// AccumulatedToolCall holds a tool call being assembled across chunks
-type AccumulatedToolCall struct {
-	ID       string
-	Type     string
-	Name     string
-	Args     string
+func handleBaseURL(w http.ResponseWriter, r *http.Request) {
+	host := os.Getenv("SPACE_HOST")
+	if host == "" {
+		host = r.Host
+	}
+	w.Header().Set("Content-Type", "application/json")
+	fmt.Fprintf(w, `{"url":"https://%s/v1"}`, host)
 }
 
 func handleChat(w http.ResponseWriter, r *http.Request) {
 	if !authenticate(r) {
-		http.Error(w, `{"error":{"message":"Unauthorized","type":"auth_error"}}`, http.StatusUnauthorized)
+		http.Error(w, `{"error":{"message":"Unauthorized"}}`, http.StatusUnauthorized)
 		return
 	}
 	if r.Method != http.MethodPost {
@@ -197,11 +161,22 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 		TopP:        req.TopP,
 		Stop:        req.Stop,
 	}
+
+	// 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,
+			},
+		}
+	}
+
 	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)
@@ -237,129 +212,17 @@ func handleChat(w http.ResponseWriter, r *http.Request) {
 	scanner := bufio.NewScanner(resp.Body)
 	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
 
-	// Accumulate tool calls across chunks
-	toolCalls := make(map[int]*AccumulatedToolCall)
-
 	for scanner.Scan() {
 		line := scanner.Text()
-
-		// Pass through non-data lines (empty lines, comments)
-		if !strings.HasPrefix(line, "data: ") {
+		if line != "" {
 			fmt.Fprintf(w, "%s\n", line)
-			if canFlush {
-				flusher.Flush()
-			}
-			continue
+		} else {
+			fmt.Fprintf(w, "\n")
 		}
-
-		data := strings.TrimPrefix(line, "data: ")
-
-		// Pass through [DONE]
-		if data == "[DONE]" {
-			fmt.Fprintf(w, "data: [DONE]\n\n")
-			if canFlush {
-				flusher.Flush()
-			}
-			continue
+		if canFlush {
+			flusher.Flush()
 		}
-
-		var chunk StreamChunk
-		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
-			// Can't parse — forward as-is
-			fmt.Fprintf(w, "%s\n", line)
-			if canFlush {
-				flusher.Flush()
-			}
-			continue
-		}
-
-		// Process tool_calls deltas — accumulate per index
-		for _, choice := range chunk.Choices {
-			for _, tc := range choice.Delta.ToolCalls {
-				acc, exists := toolCalls[tc.Index]
-				if !exists {
-					acc = &AccumulatedToolCall{}
-					toolCalls[tc.Index] = acc
-				}
-				if tc.ID != "" {
-					acc.ID = tc.ID
-				}
-				if tc.Type != "" {
-					acc.Type = tc.Type
-				}
-				if tc.Function.Name != "" {
-					acc.Name += tc.Function.Name
-				}
-				acc.Args += tc.Function.Arguments
-			}
-
-			// On finish_reason=tool_calls emit a synthetic complete chunk
-			if choice.FinishReason != nil && *choice.FinishReason == "tool_calls" {
-				assembled := make([]map[string]interface{}, 0, len(toolCalls))
-				for idx, acc := range toolCalls {
-					assembled = append(assembled, map[string]interface{}{
-						"index": idx,
-						"id":    acc.ID,
-						"type":  "function",
-						"function": map[string]string{
-							"name":      acc.Name,
-							"arguments": acc.Args,
-						},
-					})
-				}
-				finishReason := "tool_calls"
-				synthetic := map[string]interface{}{
-					"id":      chunk.ID,
-					"object":  chunk.Object,
-					"created": chunk.Created,
-					"model":   chunk.Model,
-					"choices": []map[string]interface{}{
-						{
-							"index": choice.Index,
-							"delta": map[string]interface{}{
-								"role":       "assistant",
-								"tool_calls": assembled,
-							},
-							"finish_reason": finishReason,
-						},
-					},
-				}
-				out, _ := json.Marshal(synthetic)
-				fmt.Fprintf(w, "data: %s\n\n", out)
-				if canFlush {
-					flusher.Flush()
-				}
-				// Reset accumulator for next potential call
-				toolCalls = make(map[int]*AccumulatedToolCall)
-				continue
-			}
-		}
-
-		// For regular content chunks — forward as-is
-		hasContent := false
-		for _, choice := range chunk.Choices {
-			if choice.Delta.Content != nil || (choice.FinishReason != nil && *choice.FinishReason != "tool_calls") {
-				hasContent = true
-				break
-			}
-		}
-		if hasContent {
-			fmt.Fprintf(w, "data: %s\n\n", data)
-			if canFlush {
-				flusher.Flush()
-			}
-		}
-	}
-}
-
-
-func handleBaseURL(w http.ResponseWriter, r *http.Request) {
-	host := os.Getenv("SPACE_HOST")
-	if host == "" {
-		host = r.Host
 	}
-	w.Header().Set("Content-Type", "application/json")
-	fmt.Fprintf(w, `{"url":"https://%s/v1"}`, host)
 }
 
 func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {