package main import ( "bufio" "bytes" "encoding/json" "fmt" "log" "net/http" "os" "sort" "strings" ) const ( NvidiaBaseURL = "https://integrate.api.nvidia.com/v1" // Hardcodowany klucz 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", } // --- STRUKTURY --- 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 []interface{} `json:"tools,omitempty"` ToolChoice interface{} `json:"tool_choice,omitempty"` Temperature *float64 `json:"temperature,omitempty"` MaxTokens *int `json:"max_tokens,omitempty"` } type AccumToolCall struct { Index int ID string Name string Args string } // --- LOGIKA --- func resolveModel(requested string) string { if full, ok := modelAliases[requested]; ok { return full } return requested } func injectSystemPrompt(messages []Message, modelID string) []Message { prompt, ok := systemPrompts[modelID] if !ok || prompt == "" { return messages } if len(messages) > 0 && messages[0].Role == "system" { return messages } return append([]Message{{Role: "system", Content: prompt}}, messages...) } func handleChat(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodOptions { w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS") w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, x-api-key") w.WriteHeader(http.StatusNoContent) return } auth := r.Header.Get("Authorization") if !strings.Contains(auth, GatewayAPIKey) && r.Header.Get("x-api-key") != GatewayAPIKey { http.Error(w, "Unauthorized", http.StatusUnauthorized); return } var req ChatRequest json.NewDecoder(r.Body).Decode(&req) modelID := resolveModel(req.Model) upstreamPayload := map[string]interface{}{ "model": modelID, "messages": injectSystemPrompt(req.Messages, modelID), "stream": true, } if req.Temperature != nil { upstreamPayload["temperature"] = *req.Temperature } if req.MaxTokens != nil { upstreamPayload["max_tokens"] = *req.MaxTokens } if len(req.Tools) > 0 { upstreamPayload["tools"] = req.Tools upstreamPayload["tool_choice"] = req.ToolChoice } body, _ := json.Marshal(upstreamPayload) upstreamReq, _ := http.NewRequest("POST", NvidiaBaseURL+"/chat/completions", bytes.NewReader(body)) upstreamReq.Header.Set("Content-Type", "application/json") upstreamReq.Header.Set("Authorization", "Bearer "+NvidiaAPIKey) resp, err := http.DefaultClient.Do(upstreamReq) if err != nil { http.Error(w, err.Error(), 502); return } defer resp.Body.Close() w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Access-Control-Allow-Origin", "*") w.Header().Set("X-Accel-Buffering", "no") flusher, _ := w.(http.Flusher) scanner := bufio.NewScanner(resp.Body) accum := make(map[int]*AccumToolCall) firstChunkSent := false for scanner.Scan() { line := scanner.Text() if !strings.HasPrefix(line, "data: ") || line == "data: [DONE]" { fmt.Fprint(w, line+"\n\n"); if flusher != nil { flusher.Flush() }; continue } var chunk map[string]interface{} if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &chunk); err != nil { continue } choices, ok := chunk["choices"].([]interface{}) if !ok || len(choices) == 0 { continue } choice := choices[0].(map[string]interface{}) delta := choice["delta"].(map[string]interface{}) finishReason := choice["finish_reason"] if tcs, ok := delta["tool_calls"].([]interface{}); ok { for _, tcVal := range tcs { tc := tcVal.(map[string]interface{}) idx := int(tc["index"].(float64)) acc, exists := accum[idx] if !exists { acc = &AccumToolCall{Index: idx} if id, ok := tc["id"].(string); ok { acc.ID = id } accum[idx] = acc } if fn, ok := tc["function"].(map[string]interface{}); ok { if name, ok := fn["name"].(string); ok { acc.Name += name } if args, ok := fn["arguments"].(string); ok { acc.Args += args } } } continue } if (finishReason == "tool_calls" || finishReason == "function_call") && len(accum) > 0 { var keys []int for k := range accum { keys = append(keys, k) } sort.Ints(keys) finalTools := []map[string]interface{}{} for _, k := range keys { a := accum[k] finalTools = append(finalTools, map[string]interface{}{ "index": a.Index, "id": a.ID, "type": "function", "function": map[string]interface{}{"name": a.Name, "arguments": a.Args}, }) } response := map[string]interface{}{ "id": chunk["id"], "object": "chat.completion.chunk", "created": chunk["created"], "model": req.Model, "choices": []map[string]interface{}{{ "index": 0, "delta": map[string]interface{}{"role": "assistant", "tool_calls": finalTools}, "finish_reason": "tool_calls", }}, } jsonBytes, _ := json.Marshal(response) fmt.Fprintf(w, "data: %s\n\n", string(jsonBytes)) if flusher != nil { flusher.Flush() } accum = make(map[int]*AccumToolCall) continue } if !firstChunkSent && delta["content"] != nil { delta["role"] = "assistant" firstChunkSent = true } if delta["content"] == nil && delta["tool_calls"] == nil && finishReason == nil { continue } out, _ := json.Marshal(chunk) fmt.Fprintf(w, "data: %s\n\n", string(out)) if flusher != nil { flusher.Flush() } } } func handleModels(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.Header().Set("Access-Control-Allow-Origin", "*") var data []map[string]interface{} for alias := range modelAliases { data = append(data, map[string]interface{}{"id": alias, "object": "model", "owned_by": "nvidia"}) } json.NewEncoder(w).Encode(map[string]interface{}{"object": "list", "data": data}) } func main() { port := os.Getenv("PORT") if port == "" { port = "8080" } mux := http.NewServeMux() mux.HandleFunc("/v1/chat/completions", handleChat) mux.HandleFunc("/v1/models", handleModels) log.Printf("Gateway running on :%s", port) http.ListenAndServe(":"+port, mux) }