Update Dockerfile

Dockerfile

@@ -1,13 +1,407 @@
-FROM golang:1.21-alpine AS builder
-WORKDIR /app
-COPY go.mod ./
-COPY *.go ./
-RUN go build -o gateway .
-
-FROM alpine:latest
-RUN apk --no-cache add ca-certificates
-WORKDIR /app
-COPY --from=builder /app/gateway .
-
-EXPOSE 7860
-CMD ["./gateway"]
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"sort"
+	"strings"
+	"time"
+)
+
+const (
+	NvidiaBaseURL = "https://integrate.api.nvidia.com/v1"
+	NvidiaAPIKey  = "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"
+	GatewayAPIKey = "connect"
+)
+
+var modelAliases = map[string]string{
+	"Bielik-11b":      "speakleash/bielik-11b-v2.6-instruct",
+	"GLM-4.7":         "z-ai/glm4.7",
+	"Mistral-Small-4": "mistralai/mistral-small-4-119b-2603",
+	"DeepSeek-V3.1":   "deepseek-ai/deepseek-v3.1",
+	"Kimi-K2":         "moonshotai/kimi-k2-instruct",
+}
+
+type Message struct {
+	Role       string      `json:"role"`
+	Content    interface{} `json:"content"`
+	ToolCallID string      `json:"tool_call_id,omitempty"`
+	ToolCalls  interface{} `json:"tool_calls,omitempty"`
+	Name       string      `json:"name,omitempty"`
+}
+
+type ChatRequest struct {
+	Model       string        `json:"model"`
+	Messages    []Message     `json:"messages"`
+	Stream      *bool         `json:"stream,omitempty"`
+	Tools       json.RawMessage `json:"tools,omitempty"`
+	ToolChoice  json.RawMessage `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"`
+}
+
+type UpstreamRequest struct {
+	Model       string                 `json:"model"`
+	Messages    []Message              `json:"messages"`
+	Stream      bool                   `json:"stream"`
+	Tools       json.RawMessage        `json:"tools,omitempty"`
+	ToolChoice  json.RawMessage        `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 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
+	}
+	for _, full := range modelAliases {
+		if full == requested {
+			return requested
+		}
+	}
+	return requested
+}
+
+func injectSystemPrompt(messages []Message, modelID string) []Message {
+	filtered := make([]Message, 0, len(messages))
+	for _, m := range messages {
+		if m.Role != "system" {
+			filtered = append(filtered, m)
+		}
+	}
+	prompt, ok := systemPrompts[modelID]
+	if !ok || prompt == "" {
+		return filtered
+	}
+	return append([]Message{{Role: "system", Content: prompt}}, filtered...)
+}
+
+func authenticate(r *http.Request) bool {
+	auth := r.Header.Get("Authorization")
+	if len(auth) > 7 && auth[:7] == "Bearer " && auth[7:] == GatewayAPIKey {
+		return true
+	}
+	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"}}`, http.StatusUnauthorized)
+		return
+	}
+	type ModelObj struct {
+		ID      string `json:"id"`
+		Object  string `json:"object"`
+		Created int64  `json:"created"`
+		OwnedBy string `json:"owned_by"`
+	}
+	type ModelsResponse struct {
+		Object string     `json:"object"`
+		Data   []ModelObj `json:"data"`
+	}
+	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"})
+	}
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(models)
+}
+
+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"}}`, http.StatusUnauthorized)
+		return
+	}
+	if r.Method != http.MethodPost {
+		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)
+		return
+	}
+
+	modelID := resolveModel(req.Model)
+	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,
+			},
+		}
+	}
+
+	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")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+	w.Header().Set("X-Accel-Buffering", "no")
+	w.WriteHeader(http.StatusOK)
+
+	flusher, canFlush := w.(http.Flusher)
+	emit := func(s string) {
+		fmt.Fprint(w, s)
+		if canFlush {
+			flusher.Flush()
+		}
+	}
+
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)
+
+	// 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
+		}
+
+		data := strings.TrimPrefix(line, "data: ")
+
+		if data == "[DONE]" {
+			emit("data: [DONE]\n\n")
+			continue
+		}
+
+		var chunk RawChunk
+		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
+			emit(line + "\n")
+			continue
+		}
+		lastChunk = chunk
+
+		isToolChunk := false
+		isFinishToolCalls := false
+
+		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"
+			out, _ := json.Marshal(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,
+					},
+				},
+			})
+			emit("data: " + string(out) + "\n\n")
+			accum = make(map[int]*AccumToolCall)
+			continue
+		}
+
+		// Skip intermediate tool_call delta chunks (already accumulating)
+		if isToolChunk {
+			continue
+		}
+
+		// Regular content chunk — stream immediately as-is
+		emit("data: " + data + "\n\n")
+	}
+}
+
+func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		start := time.Now()
+		log.Printf("[%s] %s %s", r.Method, r.URL.Path, r.RemoteAddr)
+		next(w, r)
+		log.Printf("[%s] %s done in %s", r.Method, r.URL.Path, time.Since(start))
+	}
+}
+
+func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Access-Control-Allow-Origin", "*")
+		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, x-api-key")
+		if r.Method == http.MethodOptions {
+			w.WriteHeader(http.StatusNoContent)
+			return
+		}
+		next(w, r)
+	}
+}
+
+func main() {
+	port := os.Getenv("PORT")
+	if port == "" {
+		port = "7860"
+	}
+	mux := http.NewServeMux()
+	mux.HandleFunc("/v1/chat/completions", corsMiddleware(loggingMiddleware(handleChat)))
+	mux.HandleFunc("/v1/models", corsMiddleware(loggingMiddleware(handleModels)))
+	mux.HandleFunc("/v1/base-url", corsMiddleware(handleBaseURL))
+	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte(`{"status":"ok"}`))
+	})
+	log.Printf("Gateway starting on :%s", port)
+	if err := http.ListenAndServe(":"+port, mux); err != nil {
+		log.Fatal(err)
+	}
+}