From d1d742209935e8857e977c97dfa9d15774b7fb93 Mon Sep 17 00:00:00 2001 From: Luxferre Date: Mon, 7 Sep 2026 07:45:50 +0300 Subject: [PATCH] feat: implement universal Gradio to OpenAI proxy gateway --- .gitignore | 2 + Makefile | 13 + README.md | 251 ++++++ go.mod | 3 + gr2gw.go | 2199 +++++++++++++++++++++++++++++++++++++++++++++++++ gr2gw_test.go | 249 ++++++ 6 files changed, 2717 insertions(+) create mode 100644 .gitignore create mode 100644 Makefile create mode 100644 README.md create mode 100644 go.mod create mode 100644 gr2gw.go create mode 100644 gr2gw_test.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..50e1cda --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +bin/ +gr2gw diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..db37e3c --- /dev/null +++ b/Makefile @@ -0,0 +1,13 @@ +.PHONY: all build clean test + +all: build + +build: + mkdir -p bin + go build -trimpath -ldflags="-s -w" -o bin/gr2gw gr2gw.go + +clean: + rm -rf bin + +test: + go test -v ./... diff --git a/README.md b/README.md new file mode 100644 index 0000000..6c379ee --- /dev/null +++ b/README.md @@ -0,0 +1,251 @@ +# gr2gw: Universal Gradio to OpenAI LLM Gateway + +A zero-dependency, high-performance Go proxy server that introspects any Gradio chat space (such as Hugging Face Spaces or custom deployments) and exposes a standards-compliant OpenAI `/v1/chat/completions` and `/v1/models` HTTP API. + +Default demo space: `https://ghost2513-openai-gpt-oss-120b.hf.space` + +--- + +## Features + +- **Zero External Dependencies**: Pure Go standard library (`net/http`, `encoding/json`, `bufio`, etc.). +- **Automatic Space Introspection**: Dynamically queries `/gradio_api/info`, `/config`, and Hugging Face space metadata to discover models, endpoints, and input parameter mappings. +- **Universal Multi-turn Handling**: + - Automatically formats conversation history into structured inputs when the space supports them. + - Transparently composes multi-turn dialogue (`System`, `User`, `Assistant`) into single prompt inputs when the space only accepts a single message textbox. + - Automatically pads hidden/State inputs (e.g. Gradio State components) to prevent backend argument count mismatches. +- **Real-Time Streaming & Accumulation Filter**: + - Automatically computes token deltas from cumulative or incremental Gradio SSE output streams. + - Emits standards-compliant `chat.completion.chunk` SSE events in real time. +- **Thinking & Reasoning Token Separation**: + - Detects `...` tags in real time. + - Separates reasoning into `delta.reasoning_content` (streaming) and `message.reasoning_content` (non-streaming). + - Keeps `content` clean without tag leakage. +- **Full Tool Calling & Function Interception**: + - Formats schemas into system prompts with strict function calling instructions. + - **`StreamToolCallFilter`**: Stateful sliding-window filter that prevents `` tags from leaking into `delta.content`. Emits structured OpenAI `delta.tool_calls` chunks and sets `finish_reason: "tool_calls"`. + - Seamlessly maintains multi-turn context when tool results are submitted back via `role: "tool"`. +- **Built-in SOCKS5 Proxy Client**: + - Full RFC 1928 / RFC 1929 implementation with domain resolution (`socks5h://`), IPv4, IPv6, and username/password auth. +- **Dynamic Space Override**: + - Switch the target Gradio space on-the-fly per request using the `X-Gradio-Space` or `X-Space-URL` HTTP headers. +- **Fibonacci Retry Engine**: + - Resilient backoff retry mechanism (1s, 1s, 2s, 3s, 5s) for transient network hiccups. + +--- + +## Build + +```bash +make build +``` + +Binary will be compiled to `bin/gr2gw`. + +To run tests: +```bash +make test +``` + +--- + +## Usage + +### Quick Start + +Run with the default space (`https://ghost2513-openai-gpt-oss-120b.hf.space`): +```bash +./bin/gr2gw -port 8080 +``` + +Target any other Gradio space: +```bash +./bin/gr2gw -space https://ericsqin-hy3.hf.space -port 8080 +``` + +With SOCKS5 proxy: +```bash +./bin/gr2gw -space https://ghost2513-openai-gpt-oss-120b.hf.space -socks socks5://127.0.0.1:1080 +``` + +### CLI Flags + +| Flag | Default | Description | +|------|---------|-------------| +| `-space`, `-url` | `https://ghost2513-openai-gpt-oss-120b.hf.space` | Target Gradio space URL | +| `-port` | `8080` | Port to listen on | +| `-host` | `0.0.0.0` | Host interface to bind to | +| `-socks`, `-proxy`, `-socks5` | `""` | SOCKS5 proxy URL (`socks5://user:pass@host:port`) | +| `-user-agent`, `-ua` | Firefox string | Custom User-Agent header | +| `-timeout` | `300` | Upstream request timeout in seconds | + +### Environment Variables + +- `GRADIO_SPACE_URL`: Default Gradio space URL fallback. +- `ALL_PROXY`, `SOCKS5_PROXY`, `SOCKS_PROXY`: Default SOCKS5 proxy URL fallback. + +--- + +## API Examples + +### List Models + +```bash +curl http://localhost:8080/v1/models +``` + +Response: +```json +{ + "object": "list", + "data": [ + { + "id": "openai/gpt-oss-120b", + "object": "model", + "created": 1788756307, + "owned_by": "gradio" + }, + { + "id": "gpt-oss-120b", + "object": "model", + "created": 1788756307, + "owned_by": "gradio" + } + ] +} +``` + +### Chat Completions (Non-Streaming) + +```bash +curl http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openai/gpt-oss-120b", + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ] + }' +``` + +Response: +```json +{ + "id": "chatcmpl-16425f9d-c350-47a1-9a6d-e9ce10871545", + "object": "chat.completion", + "created": 1788756310, + "model": "openai/gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Paris is the capital of France." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } +} +``` + +### Chat Completions (Streaming) + +```bash +curl -N http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openai/gpt-oss-120b", + "messages": [ + {"role": "user", "content": "Count from 1 to 5."} + ], + "stream": true + }' +``` + +### Tool Calling + +```bash +curl http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "openai/gpt-oss-120b", + "messages": [ + {"role": "user", "content": "What is the weather in Tokyo?"} + ], + "tools": [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather in a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name"} + }, + "required": ["location"] + } + } + }] + }' +``` + +Response: +```json +{ + "id": "chatcmpl-32727c62-ef2a-4866-855b-f1c7ec2b8023", + "object": "chat.completion", + "created": 1788756322, + "model": "openai/gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": null, + "tool_calls": [ + { + "id": "call_3d4c016a", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"location\":\"Tokyo\"}" + } + } + ] + }, + "finish_reason": "tool_calls" + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0 + } +} +``` + +### Dynamic Target Space Override + +Override the target space per request without restarting the server: + +```bash +curl http://localhost:8080/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "X-Gradio-Space: https://ericsqin-hy3.hf.space" \ + -d '{ + "messages": [ + {"role": "user", "content": "Hello!"} + ] + }' +``` + +--- + +## License + +Released into the public domain under Creative Commons Zero (CC0) or Unlicense. diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..66d5521 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module gr2gw + +go 1.26 diff --git a/gr2gw.go b/gr2gw.go new file mode 100644 index 0000000..cbe6c5d --- /dev/null +++ b/gr2gw.go @@ -0,0 +1,2199 @@ +// gr2gw: Universal Gradio to OpenAI LLM proxy gateway in Go +// Created by Luxferre in 2026, released into the public domain + +package main + +import ( + "bufio" + "bytes" + "context" + "crypto/rand" + "encoding/json" + "flag" + "fmt" + "io" + "log" + "net" + "net/http" + "net/url" + "os" + "regexp" + "strconv" + "strings" + "sync" + "time" +) + +var ( + DefaultSpaceURL = "https://ghost2513-openai-gpt-oss-120b.hf.space" + DefaultUserAgent = "Mozilla/5.0 (X11; Linux x86_64; rv:153.0) Gecko/20100101 Firefox/153.0" + ConfiguredUserAgent string +) + +// --------------------------------------------------------------------------- +// OpenAI API Data Structures +// --------------------------------------------------------------------------- + +type ModelItem 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 []ModelItem `json:"data"` +} + +type ToolCallFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type ToolCall struct { + Index *int `json:"index,omitempty"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function ToolCallFunction `json:"function"` +} + +type Tool struct { + Type string `json:"type"` + Function interface{} `json:"function"` +} + +type ChatMessage struct { + Role string `json:"role"` + Content interface{} `json:"content"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + Name string `json:"name,omitempty"` +} + +func (m *ChatMessage) GetContentString() string { + if m.Content == nil { + return "" + } + if str, ok := m.Content.(string); ok { + return str + } + if parts, ok := m.Content.([]interface{}); ok { + var sb strings.Builder + for _, p := range parts { + if str, ok := p.(string); ok { + sb.WriteString(str) + } else if itemMap, ok := p.(map[string]interface{}); ok { + if textVal, ok := itemMap["text"].(string); ok { + sb.WriteString(textVal) + } + } + } + return sb.String() + } + b, err := json.Marshal(m.Content) + if err == nil { + return string(b) + } + return fmt.Sprintf("%v", m.Content) +} + +type ChatCompletionRequest struct { + Model string `json:"model"` + Messages []ChatMessage `json:"messages"` + Tools []Tool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Stream bool `json:"stream"` + MaxTokens int `json:"max_tokens"` + MaxCompletionTokens int `json:"max_completion_tokens"` + Temperature *float64 `json:"temperature,omitempty"` + TopP *float64 `json:"top_p,omitempty"` + ReasoningEffort string `json:"reasoning_effort,omitempty"` +} + +type ChatCompletionResponseChoice struct { + Index int `json:"index"` + Message ChatMessage `json:"message"` + FinishReason string `json:"finish_reason"` +} + +type Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +type ChatCompletionResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []ChatCompletionResponseChoice `json:"choices"` + Usage Usage `json:"usage"` +} + +type StreamDelta struct { + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` +} + +type StreamChoice struct { + Index int `json:"index"` + Delta StreamDelta `json:"delta"` + FinishReason *string `json:"finish_reason,omitempty"` +} + +type StreamResponse struct { + ID string `json:"id"` + Object string `json:"object"` + Created int64 `json:"created"` + Model string `json:"model"` + Choices []StreamChoice `json:"choices"` +} + +// --------------------------------------------------------------------------- +// SOCKS5 Proxy Client (RFC 1928 / RFC 1929) +// --------------------------------------------------------------------------- + +type SOCKS5Config struct { + Address string + Username string + Password string +} + +func ParseSOCKS5URL(proxyURL string) (*SOCKS5Config, error) { + cleanURL := strings.TrimSpace(proxyURL) + if cleanURL == "" { + return nil, nil + } + if strings.HasPrefix(cleanURL, "socks5://") { + cleanURL = strings.TrimPrefix(cleanURL, "socks5://") + } else if strings.HasPrefix(cleanURL, "socks5h://") { + cleanURL = strings.TrimPrefix(cleanURL, "socks5h://") + } + + cfg := &SOCKS5Config{} + if atIdx := strings.LastIndex(cleanURL, "@"); atIdx != -1 { + userPass := cleanURL[:atIdx] + cfg.Address = cleanURL[atIdx+1:] + if colonIdx := strings.Index(userPass, ":"); colonIdx != -1 { + cfg.Username = userPass[:colonIdx] + cfg.Password = userPass[colonIdx+1:] + } else { + cfg.Username = userPass + } + } else { + cfg.Address = cleanURL + } + + if !strings.Contains(cfg.Address, ":") { + cfg.Address = cfg.Address + ":1080" + } + return cfg, nil +} + +func DialSOCKS5(ctx context.Context, proxyURL, targetAddr string) (net.Conn, error) { + cfg, err := ParseSOCKS5URL(proxyURL) + if err != nil { + return nil, fmt.Errorf("invalid socks5 proxy configuration: %w", err) + } + if cfg == nil { + var d net.Dialer + return d.DialContext(ctx, "tcp", targetAddr) + } + + var d net.Dialer + conn, err := d.DialContext(ctx, "tcp", cfg.Address) + if err != nil { + return nil, fmt.Errorf("failed to connect to socks5 proxy at %s: %w", cfg.Address, err) + } + + deadline, ok := ctx.Deadline() + if !ok { + deadline = time.Now().Add(30 * time.Second) + } + conn.SetDeadline(deadline) + defer conn.SetDeadline(time.Time{}) + + // 1. Negotiation Greeting (RFC 1928) + var greeting []byte + if cfg.Username != "" { + greeting = []byte{0x05, 0x02, 0x00, 0x02} + } else { + greeting = []byte{0x05, 0x01, 0x00} + } + + if _, err := conn.Write(greeting); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to write socks5 greeting: %w", err) + } + + resp := make([]byte, 2) + if _, err := io.ReadFull(conn, resp); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to read socks5 greeting response: %w", err) + } + + if resp[0] != 0x05 { + conn.Close() + return nil, fmt.Errorf("unsupported socks version: 0x%02x", resp[0]) + } + + // 2. Authentication if required (RFC 1929) + if resp[1] == 0x02 { + if cfg.Username == "" { + conn.Close() + return nil, fmt.Errorf("socks5 proxy requires authentication, but no credentials provided") + } + uLen := byte(len(cfg.Username)) + pLen := byte(len(cfg.Password)) + authReq := []byte{0x01, uLen} + authReq = append(authReq, []byte(cfg.Username)...) + authReq = append(authReq, pLen) + authReq = append(authReq, []byte(cfg.Password)...) + + if _, err := conn.Write(authReq); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to send socks5 authentication: %w", err) + } + + authResp := make([]byte, 2) + if _, err := io.ReadFull(conn, authResp); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to read socks5 auth response: %w", err) + } + if authResp[1] != 0x00 { + conn.Close() + return nil, fmt.Errorf("socks5 authentication failed with status 0x%02x", authResp[1]) + } + } else if resp[1] != 0x00 { + conn.Close() + return nil, fmt.Errorf("socks5 proxy rejected authentication methods: 0x%02x", resp[1]) + } + + // 3. Connection Request (CONNECT command) + host, portStr, err := net.SplitHostPort(targetAddr) + if err != nil { + conn.Close() + return nil, fmt.Errorf("invalid target address %s: %w", targetAddr, err) + } + + port, err := strconv.Atoi(portStr) + if err != nil || port < 1 || port > 65535 { + conn.Close() + return nil, fmt.Errorf("invalid port in target address %s", targetAddr) + } + + reqBuf := []byte{0x05, 0x01, 0x00} + ip := net.ParseIP(host) + if ip4 := ip.To4(); ip4 != nil { + reqBuf = append(reqBuf, 0x01) + reqBuf = append(reqBuf, ip4...) + } else if ip6 := ip.To16(); ip6 != nil { + reqBuf = append(reqBuf, 0x04) + reqBuf = append(reqBuf, ip6...) + } else { + reqBuf = append(reqBuf, 0x03, byte(len(host))) + reqBuf = append(reqBuf, []byte(host)...) + } + reqBuf = append(reqBuf, byte(port>>8), byte(port&0xFF)) + + if _, err := conn.Write(reqBuf); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to send socks5 connect request: %w", err) + } + + // 4. Connection Response + respHdr := make([]byte, 4) + if _, err := io.ReadFull(conn, respHdr); err != nil { + conn.Close() + return nil, fmt.Errorf("failed to read socks5 connect response: %w", err) + } + + if respHdr[1] != 0x00 { + conn.Close() + return nil, fmt.Errorf("socks5 connect failed with reply code 0x%02x", respHdr[1]) + } + + switch respHdr[3] { + case 0x01: + bnd := make([]byte, 6) + io.ReadFull(conn, bnd) + case 0x03: + lenBuf := make([]byte, 1) + io.ReadFull(conn, lenBuf) + bnd := make([]byte, int(lenBuf[0])+2) + io.ReadFull(conn, bnd) + case 0x04: + bnd := make([]byte, 18) + io.ReadFull(conn, bnd) + } + + return conn, nil +} + +// --------------------------------------------------------------------------- +// Helper Utilities +// --------------------------------------------------------------------------- + +func GenerateUUID() string { + var b [16]byte + _, err := rand.Read(b[:]) + if err != nil { + return fmt.Sprintf("%d", time.Now().UnixNano()) + } + b[6] = (b[6] & 0x0f) | 0x40 + b[8] = (b[8] & 0x3f) | 0x80 + return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) +} + +func GenerateSessionHash() string { + const chars = "abcdefghijklmnopqrstuvwxyz0123456789" + var b [12]byte + rand.Read(b[:]) + var sb strings.Builder + for _, v := range b { + sb.WriteByte(chars[int(v)%len(chars)]) + } + return sb.String() +} + +func FibonacciDelay(attempt int) time.Duration { + if attempt <= 0 { + return 1 * time.Second + } + a, b := 1, 1 + for i := 1; i < attempt; i++ { + a, b = b, a+b + } + return time.Duration(a) * time.Second +} + +func DoWithFibonacciRetry(client *http.Client, makeReq func() (*http.Request, error), maxRetries int) (*http.Response, error) { + var lastErr error + for attempt := 1; attempt <= maxRetries; attempt++ { + req, err := makeReq() + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err == nil && resp.StatusCode == http.StatusOK { + return resp, nil + } + + if resp != nil { + respBody, _ := io.ReadAll(resp.Body) + resp.Body.Close() + lastErr = fmt.Errorf("HTTP status %d: %s", resp.StatusCode, string(respBody)) + } else { + lastErr = err + } + + if attempt < maxRetries { + delay := FibonacciDelay(attempt) + time.Sleep(delay) + } + } + return nil, fmt.Errorf("request failed after %d retries: %v", maxRetries, lastErr) +} + +func EnableCORS(w http.ResponseWriter) { + 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, api-key, X-User-Agent, X-Gradio-Space, X-Space-URL") +} + +func ResolveMaxTokens(req ChatCompletionRequest) int { + mt := req.MaxTokens + if mt == 0 && req.MaxCompletionTokens > 0 { + mt = req.MaxCompletionTokens + } + if mt <= 0 { + mt = 131072 + } + return mt +} + +func EffectiveUserAgent(r *http.Request) string { + if r != nil { + if c := r.Header.Get("X-User-Agent"); c != "" { + return c + } + } + if ConfiguredUserAgent != "" { + return ConfiguredUserAgent + } + return DefaultUserAgent +} + +// --------------------------------------------------------------------------- +// Tool and Message Processing +// --------------------------------------------------------------------------- + +func BuildToolInstruction(tools []Tool) string { + if len(tools) == 0 { + return "" + } + toolsBytes, _ := json.MarshalIndent(tools, "", " ") + return fmt.Sprintf("\n\n# Tool Calling Instructions\n\nYou have access to the following functions:\n\n%s\n\n\nWhen you need to call a function, respond ONLY with a block formatted exactly as follows:\n\n{\"name\": \"\", \"arguments\": {}}\n\n\nDo not include conversational filler before or after the tool call.", string(toolsBytes)) +} + +func TransformMessages(req ChatCompletionRequest) (processed []ChatMessage, toolInstruction string, hasSystem bool) { + toolInstruction = BuildToolInstruction(req.Tools) + for _, msg := range req.Messages { + contentStr := msg.GetContentString() + m := ChatMessage{Role: msg.Role, Content: contentStr} + switch msg.Role { + case "system": + hasSystem = true + m.Content = contentStr + case "assistant": + var sb strings.Builder + if contentStr != "" { + sb.WriteString(contentStr) + } + for _, tc := range msg.ToolCalls { + if sb.Len() > 0 { + sb.WriteString("\n") + } + args := tc.Function.Arguments + if strings.TrimSpace(args) == "" { + args = "{}" + } + sb.WriteString(fmt.Sprintf("\n{\"name\": %q, \"arguments\": %s}\n", tc.Function.Name, args)) + } + m.Content = sb.String() + case "tool", "function": + m.Role = "user" + toolName := msg.Name + if toolName == "" { + toolName = msg.ToolCallID + } + var contentJSON []byte + if json.Valid([]byte(contentStr)) { + contentJSON = []byte(contentStr) + } else { + contentJSON, _ = json.Marshal(contentStr) + } + m.Content = fmt.Sprintf("\n{\"name\": %q, \"content\": %s}\n", toolName, string(contentJSON)) + } + processed = append(processed, m) + } + + if toolInstruction != "" { + if hasSystem { + for i, m := range processed { + if m.Role == "system" { + processed[i].Content = m.GetContentString() + "\n" + strings.TrimSpace(toolInstruction) + break + } + } + } else { + processed = append([]ChatMessage{ + {Role: "user", Content: strings.TrimSpace(toolInstruction)}, + }, processed...) + } + } + + return processed, toolInstruction, hasSystem +} + +func cleanJSONBlock(input string) string { + s := strings.TrimSpace(input) + if strings.HasPrefix(s, "```") { + lines := strings.Split(s, "\n") + if len(lines) >= 2 { + if strings.HasPrefix(lines[len(lines)-1], "```") { + lines = lines[1 : len(lines)-1] + } else { + lines = lines[1:] + } + s = strings.TrimSpace(strings.Join(lines, "\n")) + } + } + return s +} + +func sanitizeJSONValue(v interface{}) interface{} { + switch val := v.(type) { + case string: + return strings.TrimSpace(val) + case map[string]interface{}: + res := make(map[string]interface{}) + for k, item := range val { + res[k] = sanitizeJSONValue(item) + } + return res + case []interface{}: + res := make([]interface{}, len(val)) + for i, item := range val { + res[i] = sanitizeJSONValue(item) + } + return res + default: + return v + } +} + +func repairToolCallJSON(input string) (ToolCall, bool) { + s := strings.TrimSpace(input) + reName := regexp.MustCompile(`"(?:name|function|action|call)"\s*:\s*"([^"]+)"`) + matches := reName.FindStringSubmatch(s) + if len(matches) < 2 { + return ToolCall{}, false + } + fnName := matches[1] + + reArgs := regexp.MustCompile(`"(?:arguments|parameters|args|input)"\s*:\s*(\{[\s\S]*\})`) + argMatches := reArgs.FindStringSubmatch(s) + argsStr := "{}" + if len(argMatches) >= 2 { + candidate := argMatches[1] + var dummy map[string]interface{} + if json.Unmarshal([]byte(candidate), &dummy) == nil { + argsStr = candidate + } + } + + return ToolCall{ + ID: "call_" + GenerateUUID()[:8], + Type: "function", + Function: ToolCallFunction{ + Name: fnName, + Arguments: argsStr, + }, + }, true +} + +func parseSingleToolCall(jsonStr string) (ToolCall, bool) { + cleaned := cleanJSONBlock(jsonStr) + var raw map[string]interface{} + if err := json.Unmarshal([]byte(cleaned), &raw); err == nil { + sanitizedRaw, ok := sanitizeJSONValue(raw).(map[string]interface{}) + if !ok { + sanitizedRaw = raw + } + + for _, wrapperKey := range []string{"function", "function_call", "tool_call"} { + if fnObj, ok := sanitizedRaw[wrapperKey].(map[string]interface{}); ok { + if nameVal, ok := fnObj["name"].(string); ok && nameVal != "" { + argsStr := "{}" + var argsVal interface{} + if a, hasA := fnObj["arguments"]; hasA { + argsVal = a + } else if p, hasP := fnObj["parameters"]; hasP { + argsVal = p + } else if args, hasArgs := fnObj["args"]; hasArgs { + argsVal = args + } + if argsVal != nil { + if s, isStr := argsVal.(string); isStr { + var innerObj interface{} + if json.Unmarshal([]byte(s), &innerObj) == nil { + b, _ := json.Marshal(sanitizeJSONValue(innerObj)) + argsStr = string(b) + } else { + argsStr = strings.TrimSpace(s) + } + } else { + b, _ := json.Marshal(sanitizeJSONValue(argsVal)) + argsStr = string(b) + } + } + return ToolCall{ + ID: "call_" + GenerateUUID()[:8], + Type: "function", + Function: ToolCallFunction{ + Name: nameVal, + Arguments: argsStr, + }, + }, true + } + } + } + + nameVal := "" + for _, key := range []string{"name", "function", "action", "call"} { + if n, ok := sanitizedRaw[key].(string); ok && n != "" { + nameVal = n + break + } + } + + if nameVal != "" { + argsStr := "{}" + var argsVal interface{} + for _, key := range []string{"arguments", "parameters", "args", "input"} { + if a, ok := sanitizedRaw[key]; ok { + argsVal = a + break + } + } + if argsVal != nil { + if s, isStr := argsVal.(string); isStr { + var innerObj interface{} + if json.Unmarshal([]byte(s), &innerObj) == nil { + b, _ := json.Marshal(sanitizeJSONValue(innerObj)) + argsStr = string(b) + } else { + argsStr = strings.TrimSpace(s) + } + } else { + b, _ := json.Marshal(sanitizeJSONValue(argsVal)) + argsStr = string(b) + } + } else { + argsMap := make(map[string]interface{}) + for k, v := range sanitizedRaw { + if k != "name" && k != "function" && k != "type" && k != "action" && k != "call" { + argsMap[k] = v + } + } + if len(argsMap) > 0 { + b, _ := json.Marshal(sanitizeJSONValue(argsMap)) + argsStr = string(b) + } + } + return ToolCall{ + ID: "call_" + GenerateUUID()[:8], + Type: "function", + Function: ToolCallFunction{ + Name: nameVal, + Arguments: argsStr, + }, + }, true + } + } + + return repairToolCallJSON(cleaned) +} + +func parseXMLToolCall(block string) (ToolCall, bool) { + inner := strings.TrimSpace(block) + if strings.HasPrefix(inner, "") { + inner = strings.TrimPrefix(inner, "") + } + if strings.HasSuffix(inner, "") { + inner = strings.TrimSuffix(inner, "") + } + inner = cleanJSONBlock(inner) + + if tc, ok := parseSingleToolCall(inner); ok { + return tc, true + } + + var fnName string + if strings.Contains(inner, "") && strings.Contains(inner, "") { + nStart := strings.Index(inner, "") + len("") + nEnd := strings.Index(inner, "") + if nStart < nEnd { + fnName = strings.TrimSpace(inner[nStart:nEnd]) + } + } + + var argsStr string + if strings.Contains(inner, "") && strings.Contains(inner, "") { + aStart := strings.Index(inner, "") + len("") + aEnd := strings.Index(inner, "") + if aStart < aEnd { + argsStr = strings.TrimSpace(inner[aStart:aEnd]) + } + } + + if fnName != "" { + if argsStr == "" { + argsStr = "{}" + } + return ToolCall{ + ID: "call_" + GenerateUUID()[:8], + Type: "function", + Function: ToolCallFunction{ + Name: fnName, + Arguments: argsStr, + }, + }, true + } + + return ToolCall{}, false +} + +func ExtractToolCallBlocks(content string) (blocks []string, remaining string) { + s := content + remaining = content + + for strings.Contains(s, "") { + sIdx := strings.Index(s, "") + rest := s[sIdx+len(""):] + + relNextSIdx := strings.Index(rest, "") + var nextSIdx int + if relNextSIdx != -1 { + nextSIdx = sIdx + len("") + relNextSIdx + } else { + nextSIdx = -1 + } + + relEIdx := strings.Index(rest, "") + var eIdx int + if relEIdx != -1 { + eIdx = sIdx + len("") + relEIdx + } else { + eIdx = -1 + } + + var blockText string + var blockEndPos int + + if eIdx != -1 && (nextSIdx == -1 || eIdx < nextSIdx) { + blockEndPos = eIdx + len("") + blockText = s[sIdx:blockEndPos] + s = s[blockEndPos:] + } else if nextSIdx != -1 { + blockEndPos = nextSIdx + blockText = s[sIdx:blockEndPos] + s = s[blockEndPos:] + } else { + blockText = s[sIdx:] + s = "" + } + + blocks = append(blocks, blockText) + } + + for strings.Contains(remaining, "") { + st := strings.Index(remaining, "") + rest := remaining[st+len(""):] + + relNext := strings.Index(rest, "") + var nextSt int + if relNext != -1 { + nextSt = st + len("") + relNext + } else { + nextSt = -1 + } + + relEn := strings.Index(rest, "") + var en int + if relEn != -1 { + en = st + len("") + relEn + } else { + en = -1 + } + + if en != -1 && (nextSt == -1 || en < nextSt) { + remaining = strings.TrimSpace(remaining[:st] + remaining[en+len(""):]) + } else if nextSt != -1 { + remaining = strings.TrimSpace(remaining[:st] + remaining[nextSt:]) + } else { + remaining = strings.TrimSpace(remaining[:st]) + } + } + + return blocks, remaining +} + +func DetectToolCalls(content string) ([]ToolCall, string, bool) { + blocks, remaining := ExtractToolCallBlocks(content) + var calls []ToolCall + + for _, block := range blocks { + if toolCall, ok := parseXMLToolCall(block); ok { + calls = append(calls, toolCall) + } + } + + if len(calls) > 0 { + return calls, remaining, true + } + + if tc, ok := parseSingleToolCall(strings.TrimSpace(content)); ok { + return []ToolCall{tc}, "", true + } + + return nil, content, false +} + +func ExtractThinking(content string) (string, string) { + if strings.Contains(content, "") && strings.Contains(content, "") { + start := strings.Index(content, "") + end := strings.Index(content, "") + if start < end { + reasoning := content[start+len("") : end] + rem := content[:start] + content[end+len(""):] + rem = strings.TrimPrefix(rem, "\n\n") + rem = strings.TrimPrefix(rem, "\n") + return rem, reasoning + } + } + return content, "" +} + +// --------------------------------------------------------------------------- +// Response Framing & Streamer +// --------------------------------------------------------------------------- + +type FinalOutput struct { + Content interface{} + ReasoningContent string + ToolCalls []ToolCall + FinishReason string +} + +func WriteCompletionResponse(w http.ResponseWriter, completionID string, created int64, model string, out FinalOutput) { + finish := out.FinishReason + if finish == "" { + finish = "stop" + } + resp := ChatCompletionResponse{ + ID: completionID, + Object: "chat.completion", + Created: created, + Model: model, + Choices: []ChatCompletionResponseChoice{ + { + Index: 0, + Message: ChatMessage{ + Role: "assistant", + Content: out.Content, + ReasoningContent: out.ReasoningContent, + ToolCalls: out.ToolCalls, + }, + FinishReason: finish, + }, + }, + Usage: Usage{ + PromptTokens: 0, + CompletionTokens: 0, + TotalTokens: 0, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) +} + +type Streamer struct { + w http.ResponseWriter + flusher http.Flusher + id string + created int64 + model string +} + +func NewStreamer(w http.ResponseWriter, flusher http.Flusher, id string, created int64, model string) *Streamer { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + return &Streamer{w: w, flusher: flusher, id: id, created: created, model: model} +} + +func (s *Streamer) Role() { + sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{Role: "assistant"}) +} + +func (s *Streamer) Reasoning(text string) { + sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{ReasoningContent: text}) +} + +func (s *Streamer) Content(text string) { + sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{Content: text}) +} + +func (s *Streamer) ToolCallDelta(tc ToolCall) { + sendStreamDelta(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{ToolCalls: []ToolCall{tc}}) +} + +func (s *Streamer) Finish(reason string) { + sendStreamChunk(s.w, s.flusher, s.id, s.created, s.model, StreamDelta{}, &reason) +} + +func (s *Streamer) Done() { + fmt.Fprintf(s.w, "data: [DONE]\n\n") + if s.flusher != nil { + s.flusher.Flush() + } +} + +func sendStreamDelta(w http.ResponseWriter, flusher http.Flusher, completionID string, createdTime int64, modelName string, delta StreamDelta) { + sendStreamChunk(w, flusher, completionID, createdTime, modelName, delta, nil) +} + +func sendStreamChunk(w http.ResponseWriter, flusher http.Flusher, completionID string, createdTime int64, modelName string, delta StreamDelta, finishReason *string) { + chunk := StreamResponse{ + ID: completionID, + Object: "chat.completion.chunk", + Created: createdTime, + Model: modelName, + Choices: []StreamChoice{ + { + Index: 0, + Delta: delta, + FinishReason: finishReason, + }, + }, + } + b, _ := json.Marshal(chunk) + fmt.Fprintf(w, "data: %s\n\n", b) + if flusher != nil { + flusher.Flush() + } +} + +// --------------------------------------------------------------------------- +// Stateful Thinking Tag Filter for Streaming +// --------------------------------------------------------------------------- + +type StreamThinkingFilter struct { + inThinking bool + buf string +} + +func NewStreamThinkingFilter() *StreamThinkingFilter { + return &StreamThinkingFilter{} +} + +func hasPrefixOf(target string, prefixes []string) int { + for _, p := range prefixes { + if strings.HasSuffix(target, p) { + return len(p) + } + } + return 0 +} + +func (f *StreamThinkingFilter) Feed(chunk string, onContent func(string), onReasoning func(string)) { + f.buf += chunk + thinkStartTag := "" + thinkEndTag := "" + + thinkStartPrefixes := []string{"<", " 0 { + if !f.inThinking { + if idx := strings.Index(f.buf, thinkStartTag); idx != -1 { + before := f.buf[:idx] + if before != "" { + onContent(before) + } + f.inThinking = true + f.buf = f.buf[idx+len(thinkStartTag):] + } else if matchLen := hasPrefixOf(f.buf, thinkStartPrefixes); matchLen > 0 { + safe := f.buf[:len(f.buf)-matchLen] + if safe != "" { + onContent(safe) + } + f.buf = f.buf[len(f.buf)-matchLen:] + break + } else { + onContent(f.buf) + f.buf = "" + break + } + } else { + if idx := strings.Index(f.buf, thinkEndTag); idx != -1 { + before := f.buf[:idx] + if before != "" { + onReasoning(before) + } + f.inThinking = false + f.buf = f.buf[idx+len(thinkEndTag):] + f.buf = strings.TrimPrefix(f.buf, "\n\n") + f.buf = strings.TrimPrefix(f.buf, "\n") + } else if matchLen := hasPrefixOf(f.buf, thinkEndPrefixes); matchLen > 0 { + safe := f.buf[:len(f.buf)-matchLen] + if safe != "" { + onReasoning(safe) + } + f.buf = f.buf[len(f.buf)-matchLen:] + break + } else { + onReasoning(f.buf) + f.buf = "" + break + } + } + } +} + +func (f *StreamThinkingFilter) Flush(onContent func(string), onReasoning func(string)) { + if len(f.buf) > 0 { + if f.inThinking { + onReasoning(f.buf) + } else { + onContent(f.buf) + } + f.buf = "" + } +} + +// --------------------------------------------------------------------------- +// Stateful Tool Call Tag Filter for Streaming +// --------------------------------------------------------------------------- + +type StreamToolCallFilter struct { + inToolCall bool + buf string + toolCallBuf string + toolIndex int + emittedCall bool +} + +func NewStreamToolCallFilter() *StreamToolCallFilter { + return &StreamToolCallFilter{} +} + +func (f *StreamToolCallFilter) Feed(chunk string, onContent func(string), onToolCall func(ToolCall)) { + f.buf += chunk + toolStartTag := "" + toolEndTag := "" + + startPrefixes := []string{"<", " 0 { + if !f.inToolCall { + if idx := strings.Index(f.buf, toolStartTag); idx != -1 { + before := f.buf[:idx] + if before != "" { + onContent(before) + } + f.inToolCall = true + f.buf = f.buf[idx+len(toolStartTag):] + } else if matchLen := hasPrefixOf(f.buf, startPrefixes); matchLen > 0 { + safe := f.buf[:len(f.buf)-matchLen] + if safe != "" { + onContent(safe) + } + f.buf = f.buf[len(f.buf)-matchLen:] + break + } else { + onContent(f.buf) + f.buf = "" + break + } + } else { + if idx := strings.Index(f.buf, toolEndTag); idx != -1 { + f.toolCallBuf += f.buf[:idx] + f.buf = f.buf[idx+len(toolEndTag):] + f.inToolCall = false + + if tc, ok := parseSingleToolCall(f.toolCallBuf); ok { + idxCopy := f.toolIndex + tc.Index = &idxCopy + f.toolIndex++ + f.emittedCall = true + onToolCall(tc) + } else if tc2, ok2 := parseXMLToolCall("" + f.toolCallBuf + ""); ok2 { + idxCopy := f.toolIndex + tc2.Index = &idxCopy + f.toolIndex++ + f.emittedCall = true + onToolCall(tc2) + } else { + onContent("" + f.toolCallBuf + "") + } + f.toolCallBuf = "" + } else if matchLen := hasPrefixOf(f.buf, endPrefixes); matchLen > 0 { + safe := f.buf[:len(f.buf)-matchLen] + f.toolCallBuf += safe + f.buf = f.buf[len(f.buf)-matchLen:] + break + } else { + f.toolCallBuf += f.buf + f.buf = "" + break + } + } + } +} + +func (f *StreamToolCallFilter) Flush(onContent func(string), onToolCall func(ToolCall)) { + if f.inToolCall && len(f.toolCallBuf) > 0 { + if tc, ok := parseSingleToolCall(f.toolCallBuf); ok { + idxCopy := f.toolIndex + tc.Index = &idxCopy + f.emittedCall = true + onToolCall(tc) + } else if tc2, ok2 := parseXMLToolCall("" + f.toolCallBuf + ""); ok2 { + idxCopy := f.toolIndex + tc2.Index = &idxCopy + f.emittedCall = true + onToolCall(tc2) + } else { + onContent("" + f.toolCallBuf) + } + f.toolCallBuf = "" + } + if len(f.buf) > 0 { + onContent(f.buf) + f.buf = "" + } +} + +// --------------------------------------------------------------------------- +// Universal Gradio Space Inspector & Metadata Discovery +// --------------------------------------------------------------------------- + +type GradioParamInfo struct { + Label string `json:"label"` + ParameterName string `json:"parameter_name"` + Component string `json:"component"` +} + +type GradioEndpointInfo struct { + Parameters []GradioParamInfo `json:"parameters"` + Returns []GradioParamInfo `json:"returns"` + APIVisibility string `json:"api_visibility"` + Description string `json:"description"` +} + +type GradioAPIInfoResponse struct { + NamedEndpoints map[string]GradioEndpointInfo `json:"named_endpoints"` + UnnamedEndpoints map[string]GradioEndpointInfo `json:"unnamed_endpoints"` +} + +type GradioComponent struct { + ID int `json:"id"` + Type string `json:"type"` + Props map[string]interface{} `json:"props"` + SkipAPI bool `json:"skip_api"` +} + +type GradioDependencyTypes struct { + Generator bool `json:"generator"` + Cancel bool `json:"cancel"` +} + +type GradioDependency struct { + ID int `json:"id"` + APIName interface{} `json:"api_name"` + Inputs []int `json:"inputs"` + Outputs []int `json:"outputs"` + Queue interface{} `json:"queue"` + Types GradioDependencyTypes `json:"types"` + APIVisibility string `json:"api_visibility"` +} + +type GradioConfigResponse struct { + Version string `json:"version"` + APIPrefix string `json:"api_prefix"` + Mode string `json:"mode"` + Title string `json:"title"` + Components []GradioComponent `json:"components"` + Dependencies []GradioDependency `json:"dependencies"` +} + +type HFSpaceCardData struct { + Title string `json:"title"` + ShortDescription string `json:"short_description"` +} + +type HFSpaceInfoResponse struct { + ID string `json:"id"` + Models []string `json:"models"` + CardData HFSpaceCardData `json:"cardData"` +} + +type SpaceParamMapping struct { + InputIndex int + ComponentID int + ParamType string // "message", "history", "system_prompt", "temperature", "max_tokens", "top_p", "state", "other" + DefaultValue interface{} +} + +type SpaceDiscovery struct { + SpaceURL string + Title string + Models []string + PrimaryModel string + APIPrefix string // e.g. "/gradio_api" or "" + Endpoint string // e.g. "/chat_fn" or "/chat" + CleanEndpoint string // e.g. "chat_fn" or "chat" + Protocol string // "call", "queue", "predict" + TotalInputs int + ParamMappings []SpaceParamMapping + HistoryIndex int // -1 if none + MessageIndex int // index for user message text + SystemIndex int // -1 if none + TempIndex int // -1 if none + MaxTokensIndex int // -1 if none + TopPIndex int // -1 if none + HistoryFormat string // "messages", "pairs", "none" + LastDiscovered time.Time +} + +func (d *SpaceDiscovery) GetModelList() []ModelItem { + now := time.Now().Unix() + var items []ModelItem + seen := make(map[string]bool) + + for _, m := range d.Models { + if m != "" && !seen[m] { + seen[m] = true + items = append(items, ModelItem{ + ID: m, + Object: "model", + Created: now, + OwnedBy: "gradio", + }) + } + } + + if d.PrimaryModel != "" && !seen[d.PrimaryModel] { + seen[d.PrimaryModel] = true + items = append(items, ModelItem{ + ID: d.PrimaryModel, + Object: "model", + Created: now, + OwnedBy: "gradio", + }) + } + + if len(items) == 0 { + items = append(items, ModelItem{ + ID: "default", + Object: "model", + Created: now, + OwnedBy: "gradio", + }) + } + + return items +} + +// InspectSpace queries Gradio's /gradio_api/info, /config, and HuggingFace Space APIs +// to build an adaptive schema mapping for any Gradio space. +func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscovery, error) { + cleanURL := strings.TrimRight(rawURL, "/") + if !strings.HasPrefix(cleanURL, "http://") && !strings.HasPrefix(cleanURL, "https://") { + cleanURL = "https://" + cleanURL + } + + discovery := &SpaceDiscovery{ + SpaceURL: cleanURL, + APIPrefix: "/gradio_api", + Endpoint: "/chat_fn", + CleanEndpoint: "chat_fn", + Protocol: "call", + TotalInputs: 1, + HistoryIndex: -1, + MessageIndex: 0, + SystemIndex: -1, + TempIndex: -1, + MaxTokensIndex: -1, + TopPIndex: -1, + HistoryFormat: "messages", + LastDiscovered: time.Now(), + } + + // 1. Try fetching /gradio_api/info or /info + var infoResp GradioAPIInfoResponse + infoFetched := false + + for _, path := range []string{"/gradio_api/info", "/info"} { + infoURL := cleanURL + path + req, err := http.NewRequest("GET", infoURL, nil) + if err == nil { + req.Header.Set("User-Agent", userAgent) + resp, err := client.Do(req) + if err == nil && resp.StatusCode == http.StatusOK { + if json.NewDecoder(resp.Body).Decode(&infoResp) == nil { + infoFetched = true + if path == "/gradio_api/info" { + discovery.APIPrefix = "/gradio_api" + } else { + discovery.APIPrefix = "" + } + } + resp.Body.Close() + if infoFetched { + break + } + } else if resp != nil { + resp.Body.Close() + } + } + } + + // 2. Try fetching /config + var configResp GradioConfigResponse + configFetched := false + + for _, path := range []string{"/config", "/gradio_api/config"} { + cfgURL := cleanURL + path + req, err := http.NewRequest("GET", cfgURL, nil) + if err == nil { + req.Header.Set("User-Agent", userAgent) + resp, err := client.Do(req) + if err == nil && resp.StatusCode == http.StatusOK { + if json.NewDecoder(resp.Body).Decode(&configResp) == nil { + configFetched = true + if configResp.APIPrefix != "" { + discovery.APIPrefix = configResp.APIPrefix + } + if configResp.Title != "" { + discovery.Title = configResp.Title + } + } + resp.Body.Close() + if configFetched { + break + } + } else if resp != nil { + resp.Body.Close() + } + } + } + + // 3. Inspect Hugging Face Space Metadata if hosted on HF + parsedURL, _ := url.Parse(cleanURL) + if parsedURL != nil && (strings.HasSuffix(parsedURL.Host, ".hf.space") || strings.Contains(parsedURL.Host, "huggingface.co")) { + subdomain := strings.TrimSuffix(parsedURL.Host, ".hf.space") + var owner, name string + dashIdx := strings.Index(subdomain, "-") + if dashIdx != -1 { + owner = subdomain[:dashIdx] + name = subdomain[dashIdx+1:] + } + + if owner != "" && name != "" { + hfAPIURL := fmt.Sprintf("https://huggingface.co/api/spaces/%s/%s", owner, name) + req, err := http.NewRequest("GET", hfAPIURL, nil) + if err == nil { + req.Header.Set("User-Agent", userAgent) + resp, err := client.Do(req) + if err == nil && resp.StatusCode == http.StatusOK { + var hfResp HFSpaceInfoResponse + if json.NewDecoder(resp.Body).Decode(&hfResp) == nil { + for _, m := range hfResp.Models { + discovery.Models = append(discovery.Models, m) + cleanM := strings.TrimPrefix(m, "openai/") + cleanM = strings.TrimPrefix(cleanM, "models/") + if cleanM != m { + discovery.Models = append(discovery.Models, cleanM) + } + } + if hfResp.CardData.Title != "" && discovery.Title == "" { + discovery.Title = hfResp.CardData.Title + } + if len(discovery.Models) > 0 { + discovery.PrimaryModel = discovery.Models[0] + } + } + resp.Body.Close() + } else if resp != nil { + resp.Body.Close() + } + } + } + } + + // Fallback model names if not discovered + if len(discovery.Models) == 0 { + if parsedURL != nil && strings.HasSuffix(parsedURL.Host, ".hf.space") { + sub := strings.TrimSuffix(parsedURL.Host, ".hf.space") + parts := strings.Split(sub, "-") + if len(parts) > 1 { + cleanModel := strings.Join(parts[1:], "-") + discovery.Models = append(discovery.Models, cleanModel) + discovery.PrimaryModel = cleanModel + } + } + } + if discovery.PrimaryModel == "" { + if len(discovery.Models) > 0 { + discovery.PrimaryModel = discovery.Models[0] + } else { + discovery.PrimaryModel = "gradio-chat" + discovery.Models = append(discovery.Models, "gradio-chat") + } + } + + // 4. Score and select the best chat endpoint + bestEndpoint := "" + bestScore := -1000 + var bestEndpointInfo *GradioEndpointInfo + + if infoFetched && len(infoResp.NamedEndpoints) > 0 { + for epName, epInfo := range infoResp.NamedEndpoints { + score := 0 + lowerName := strings.ToLower(epName) + + if strings.Contains(lowerName, "chat") { + score += 100 + } + if strings.Contains(lowerName, "predict") || strings.Contains(lowerName, "respond") || strings.Contains(lowerName, "generate") { + score += 50 + } + + for _, p := range epInfo.Parameters { + pLower := strings.ToLower(p.ParameterName) + if strings.Contains(pLower, "message") || strings.Contains(pLower, "text") || strings.Contains(pLower, "prompt") { + score += 40 + } + if strings.Contains(pLower, "history") || strings.Contains(pLower, "chat") { + score += 20 + } + } + + if score > bestScore { + bestScore = score + bestEndpoint = epName + epCopy := epInfo + bestEndpointInfo = &epCopy + } + } + } + + if bestEndpoint != "" { + discovery.Endpoint = bestEndpoint + discovery.CleanEndpoint = strings.TrimPrefix(bestEndpoint, "/") + } + + // 5. Correlate with config.dependencies to determine exact input count & state padding + compMap := make(map[int]GradioComponent) + if configFetched { + for _, comp := range configResp.Components { + compMap[comp.ID] = comp + } + + var matchingDep *GradioDependency + cleanTarget := strings.TrimPrefix(discovery.Endpoint, "/") + + for _, dep := range configResp.Dependencies { + depAPIName := "" + if s, ok := dep.APIName.(string); ok { + depAPIName = strings.TrimPrefix(s, "/") + } + if depAPIName == cleanTarget { + depCopy := dep + matchingDep = &depCopy + break + } + } + + if matchingDep != nil { + discovery.TotalInputs = len(matchingDep.Inputs) + for idx, compID := range matchingDep.Inputs { + mapping := SpaceParamMapping{ + InputIndex: idx, + ComponentID: compID, + ParamType: "other", + } + if comp, exists := compMap[compID]; exists { + cType := strings.ToLower(comp.Type) + switch cType { + case "textbox", "multimodaltextbox": + if discovery.MessageIndex == 0 && idx == 0 { + mapping.ParamType = "message" + } else if discovery.SystemIndex == -1 { + mapping.ParamType = "system_prompt" + discovery.SystemIndex = idx + } + case "state": + mapping.ParamType = "state" + if idx == 1 && len(matchingDep.Inputs) == 2 { + // Standard Gradio ChatInterface: [textbox, state] + // Component 13 is state + } + case "slider", "number": + label := "" + if comp.Props != nil { + if l, ok := comp.Props["label"].(string); ok { + label = strings.ToLower(l) + } + } + if strings.Contains(label, "temp") { + mapping.ParamType = "temperature" + discovery.TempIndex = idx + } else if strings.Contains(label, "max") || strings.Contains(label, "token") { + mapping.ParamType = "max_tokens" + discovery.MaxTokensIndex = idx + } else if strings.Contains(label, "top_p") { + mapping.ParamType = "top_p" + discovery.TopPIndex = idx + } + } + } + discovery.ParamMappings = append(discovery.ParamMappings, mapping) + } + } else if bestEndpointInfo != nil { + discovery.TotalInputs = len(bestEndpointInfo.Parameters) + } + } + + // Check parameters in bestEndpointInfo for history support + if bestEndpointInfo != nil { + for idx, p := range bestEndpointInfo.Parameters { + pName := strings.ToLower(p.ParameterName) + if strings.Contains(pName, "message") && discovery.MessageIndex == 0 { + discovery.MessageIndex = idx + } else if strings.Contains(pName, "history") { + discovery.HistoryIndex = idx + } else if strings.Contains(pName, "system") { + discovery.SystemIndex = idx + } else if strings.Contains(pName, "temp") { + discovery.TempIndex = idx + } else if strings.Contains(pName, "token") { + discovery.MaxTokensIndex = idx + } else if strings.Contains(pName, "top_p") { + discovery.TopPIndex = idx + } + } + } + + // Ensure total inputs is at least 1 + if discovery.TotalInputs < 1 { + discovery.TotalInputs = 1 + } + + return discovery, nil +} + +// --------------------------------------------------------------------------- +// Universal Gradio Gateway Engine +// --------------------------------------------------------------------------- + +type GradioJoinResponse struct { + EventID string `json:"event_id"` +} + +type GradioGateway struct { + mu sync.RWMutex + defaultURL string + client *http.Client + discoveries map[string]*SpaceDiscovery + proxyURL string +} + +func NewGradioGateway(defaultSpaceURL, proxyURL string, timeout time.Duration) *GradioGateway { + cleanDefault := strings.TrimRight(defaultSpaceURL, "/") + if !strings.HasPrefix(cleanDefault, "http://") && !strings.HasPrefix(cleanDefault, "https://") { + cleanDefault = "https://" + cleanDefault + } + + transport := &http.Transport{ + DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { + if proxyURL != "" { + return DialSOCKS5(ctx, proxyURL, addr) + } + var d net.Dialer + return d.DialContext(ctx, network, addr) + }, + MaxIdleConns: 100, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 15 * time.Second, + } + + gw := &GradioGateway{ + defaultURL: cleanDefault, + client: &http.Client{Transport: transport, Timeout: timeout}, + discoveries: make(map[string]*SpaceDiscovery), + proxyURL: proxyURL, + } + + // Pre-discover the default space + disc, err := InspectSpace(gw.client, cleanDefault, DefaultUserAgent) + if err == nil && disc != nil { + gw.discoveries[cleanDefault] = disc + } + + return gw +} + +func (g *GradioGateway) GetDiscovery(spaceURL, userAgent string) *SpaceDiscovery { + target := spaceURL + if target == "" { + target = g.defaultURL + } + cleanTarget := strings.TrimRight(target, "/") + + g.mu.RLock() + disc, exists := g.discoveries[cleanTarget] + g.mu.RUnlock() + + if exists && disc != nil && time.Since(disc.LastDiscovered) < 30*time.Minute { + return disc + } + + g.mu.Lock() + defer g.mu.Unlock() + + // Double-check under lock + disc, exists = g.discoveries[cleanTarget] + if exists && disc != nil && time.Since(disc.LastDiscovered) < 30*time.Minute { + return disc + } + + newDisc, err := InspectSpace(g.client, cleanTarget, userAgent) + if err == nil && newDisc != nil { + g.discoveries[cleanTarget] = newDisc + return newDisc + } + + if disc != nil { + return disc + } + + // Fallback discovery + fallback := &SpaceDiscovery{ + SpaceURL: cleanTarget, + APIPrefix: "/gradio_api", + Endpoint: "/chat_fn", + CleanEndpoint: "chat_fn", + Protocol: "call", + TotalInputs: 2, + HistoryIndex: -1, + MessageIndex: 0, + SystemIndex: -1, + PrimaryModel: "gradio-chat", + Models: []string{"gradio-chat"}, + LastDiscovered: time.Now(), + } + g.discoveries[cleanTarget] = fallback + return fallback +} + +// BuildGradioPayload packages OpenAI messages and parameters into the target Gradio input array. +func (g *GradioGateway) BuildGradioPayload(disc *SpaceDiscovery, req ChatCompletionRequest) ([]interface{}, error) { + transformed, _, _ := TransformMessages(req) + + var systemPromptStr string + var historyArray []map[string]interface{} + var lastUserMessage string + + var nonSystem []ChatMessage + for _, m := range transformed { + cStr := m.GetContentString() + if m.Role == "system" && systemPromptStr == "" { + systemPromptStr = cStr + } else { + nonSystem = append(nonSystem, m) + } + } + + if len(nonSystem) > 0 { + for i := 0; i < len(nonSystem)-1; i++ { + m := nonSystem[i] + historyArray = append(historyArray, map[string]interface{}{ + "role": m.Role, + "content": m.GetContentString(), + }) + } + lastUserMessage = nonSystem[len(nonSystem)-1].GetContentString() + } else if systemPromptStr != "" { + lastUserMessage = systemPromptStr + } + + var promptMessageText string + if disc.HistoryIndex != -1 { + promptMessageText = lastUserMessage + } else { + // Single message space: compose multi-turn history into the prompt + if len(transformed) <= 1 && systemPromptStr == "" { + promptMessageText = lastUserMessage + } else { + var sb strings.Builder + if systemPromptStr != "" { + sb.WriteString("System: " + systemPromptStr + "\n\n") + } + for i := 0; i < len(nonSystem)-1; i++ { + m := nonSystem[i] + roleLabel := "User" + if m.Role == "assistant" { + roleLabel = "Assistant" + } + sb.WriteString(fmt.Sprintf("%s: %s\n\n", roleLabel, m.GetContentString())) + } + sb.WriteString(lastUserMessage) + promptMessageText = sb.String() + } + } + + // Allocate input array matching TotalInputs + totalInputs := disc.TotalInputs + if totalInputs < 1 { + totalInputs = 1 + } + data := make([]interface{}, totalInputs) + + // Populate mapped fields + msgIdx := disc.MessageIndex + if msgIdx >= 0 && msgIdx < len(data) { + data[msgIdx] = promptMessageText + } + + if disc.HistoryIndex >= 0 && disc.HistoryIndex < len(data) { + if disc.HistoryFormat == "pairs" { + var pairs [][]string + for i := 0; i < len(historyArray); i += 2 { + u := "" + a := "" + if i < len(historyArray) { + u, _ = historyArray[i]["content"].(string) + } + if i+1 < len(historyArray) { + a, _ = historyArray[i+1]["content"].(string) + } + pairs = append(pairs, []string{u, a}) + } + data[disc.HistoryIndex] = pairs + } else { + data[disc.HistoryIndex] = historyArray + } + } + + if disc.SystemIndex >= 0 && disc.SystemIndex < len(data) { + data[disc.SystemIndex] = systemPromptStr + } + + if disc.TempIndex >= 0 && disc.TempIndex < len(data) { + if req.Temperature != nil { + data[disc.TempIndex] = *req.Temperature + } else { + data[disc.TempIndex] = 0.7 + } + } + + if disc.MaxTokensIndex >= 0 && disc.MaxTokensIndex < len(data) { + data[disc.MaxTokensIndex] = ResolveMaxTokens(req) + } + + if disc.TopPIndex >= 0 && disc.TopPIndex < len(data) { + if req.TopP != nil { + data[disc.TopPIndex] = *req.TopP + } else { + data[disc.TopPIndex] = 1.0 + } + } + + return data, nil +} + +// ExtractTextFromGradioOutput extracts the assistant text string from Gradio output chunks +func ExtractTextFromGradioOutput(rawJSON string) (string, bool) { + var val interface{} + if err := json.Unmarshal([]byte(rawJSON), &val); err != nil { + return "", false + } + + switch v := val.(type) { + case string: + return v, true + case []interface{}: + if len(v) == 0 { + return "", false + } + // Check if first element is string + if s, ok := v[0].(string); ok { + return s, true + } + // Check if it's a list of messages: [{"role":..., "content":...}] + if len(v) > 0 { + lastItem := v[len(v)-1] + if m, ok := lastItem.(map[string]interface{}); ok { + if c, ok := m["content"].(string); ok { + return c, true + } + if parts, ok := m["content"].([]interface{}); ok && len(parts) > 0 { + for _, p := range parts { + if pm, ok := p.(map[string]interface{}); ok { + if t, ok := pm["text"].(string); ok { + return t, true + } + } + } + } + } + // Check if it's pairs: [[u1, a1], [u2, a2]] + if pair, ok := lastItem.([]interface{}); ok && len(pair) >= 2 { + if aStr, ok := pair[1].(string); ok { + return aStr, true + } + } + } + case map[string]interface{}: + for _, key := range []string{"text", "content", "response", "data", "value"} { + if s, ok := v[key].(string); ok { + return s, true + } + } + } + + return "", false +} + +// ExecuteChatCompletion handles both streaming and non-streaming requests. +func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Request, req ChatCompletionRequest) error { + effUA := EffectiveUserAgent(r) + + // Target space selection: check request headers or fallback to default space + spaceURL := g.defaultURL + if hdr := r.Header.Get("X-Gradio-Space"); hdr != "" { + spaceURL = hdr + } else if hdr := r.Header.Get("X-Space-URL"); hdr != "" { + spaceURL = hdr + } + + disc := g.GetDiscovery(spaceURL, effUA) + + modelName := req.Model + if modelName == "" { + modelName = disc.PrimaryModel + } + + gradioData, err := g.BuildGradioPayload(disc, req) + if err != nil { + return fmt.Errorf("failed to build Gradio payload: %w", err) + } + + payloadMap := map[string]interface{}{"data": gradioData} + jsonPayload, err := json.Marshal(payloadMap) + if err != nil { + return fmt.Errorf("failed to encode request: %w", err) + } + + // 1. Submit to /call/{endpoint} + callURL := fmt.Sprintf("%s%s/call/%s", disc.SpaceURL, disc.APIPrefix, disc.CleanEndpoint) + makeCallReq := func() (*http.Request, error) { + r, err := http.NewRequest("POST", callURL, bytes.NewBuffer(jsonPayload)) + if err != nil { + return nil, err + } + r.Header.Set("Content-Type", "application/json") + r.Header.Set("User-Agent", effUA) + return r, nil + } + + resp, err := DoWithFibonacciRetry(g.client, makeCallReq, 5) + if err != nil { + // If call failed, try without APIPrefix or try /call/v2 + altCallURL := fmt.Sprintf("%s/call/%s", disc.SpaceURL, disc.CleanEndpoint) + makeAltReq := func() (*http.Request, error) { + r, err := http.NewRequest("POST", altCallURL, bytes.NewBuffer(jsonPayload)) + if err != nil { + return nil, err + } + r.Header.Set("Content-Type", "application/json") + r.Header.Set("User-Agent", effUA) + return r, nil + } + resp, err = DoWithFibonacciRetry(g.client, makeAltReq, 3) + if err != nil { + return fmt.Errorf("upstream Gradio call error: %w", err) + } + } + defer resp.Body.Close() + + var joinRes GradioJoinResponse + if err := json.NewDecoder(resp.Body).Decode(&joinRes); err != nil || joinRes.EventID == "" { + return fmt.Errorf("failed to parse Gradio event ID from response") + } + + // 2. Connect to Gradio SSE EventStream + streamURL := fmt.Sprintf("%s%s/call/%s/%s", disc.SpaceURL, disc.APIPrefix, disc.CleanEndpoint, joinRes.EventID) + makeStreamReq := func() (*http.Request, error) { + r, err := http.NewRequest("GET", streamURL, nil) + if err != nil { + return nil, err + } + r.Header.Set("Accept", "text/event-stream") + r.Header.Set("User-Agent", effUA) + return r, nil + } + + streamResp, err := DoWithFibonacciRetry(g.client, makeStreamReq, 5) + if err != nil { + return fmt.Errorf("upstream Gradio stream error: %w", err) + } + defer streamResp.Body.Close() + + completionID := "chatcmpl-" + GenerateUUID() + createdTime := time.Now().Unix() + + // 3. Handle Non-Streaming vs Streaming + if !req.Stream { + reader := bufio.NewReader(streamResp.Body) + var latestFullText string + currentEvent := "" + + for { + line, err := reader.ReadString('\n') + if err != nil { + break + } + line = strings.TrimRight(line, "\r\n") + + if strings.HasPrefix(line, "event: ") { + currentEvent = strings.TrimPrefix(line, "event: ") + continue + } + + if strings.HasPrefix(line, "data: ") { + dataStr := strings.TrimPrefix(line, "data: ") + if currentEvent == "error" { + return fmt.Errorf("gradio stream error: %s", dataStr) + } + if txt, ok := ExtractTextFromGradioOutput(dataStr); ok { + latestFullText = txt + } + if currentEvent == "complete" { + break + } + } + } + + cleanText, reasoning := ExtractThinking(latestFullText) + toolCalls, remainingText, hasTools := DetectToolCalls(cleanText) + + finishReason := "stop" + var finalContent interface{} = remainingText + if hasTools && len(toolCalls) > 0 { + finishReason = "tool_calls" + if strings.TrimSpace(remainingText) == "" { + finalContent = nil + } + } + + WriteCompletionResponse(w, completionID, createdTime, modelName, FinalOutput{ + Content: finalContent, + ReasoningContent: reasoning, + ToolCalls: toolCalls, + FinishReason: finishReason, + }) + return nil + } + + // 4. Streaming Mode + flusher, _ := w.(http.Flusher) + streamer := NewStreamer(w, flusher, completionID, createdTime, modelName) + streamer.Role() + + thinkFilter := NewStreamThinkingFilter() + toolFilter := NewStreamToolCallFilter() + + reader := bufio.NewReader(streamResp.Body) + var prevText string + currentEvent := "" + + for { + line, err := reader.ReadString('\n') + if err != nil { + break + } + line = strings.TrimRight(line, "\r\n") + + if strings.HasPrefix(line, "event: ") { + currentEvent = strings.TrimPrefix(line, "event: ") + continue + } + + if strings.HasPrefix(line, "data: ") { + dataStr := strings.TrimPrefix(line, "data: ") + if currentEvent == "error" { + break + } + + if currentText, ok := ExtractTextFromGradioOutput(dataStr); ok { + var delta string + if strings.HasPrefix(currentText, prevText) { + delta = currentText[len(prevText):] + } else if prevText == "" { + delta = currentText + } else { + delta = currentText + } + prevText = currentText + + if delta != "" { + thinkFilter.Feed(delta, func(contentChunk string) { + toolFilter.Feed(contentChunk, func(cleanChunk string) { + if cleanChunk != "" { + streamer.Content(cleanChunk) + } + }, func(tc ToolCall) { + streamer.ToolCallDelta(tc) + }) + }, func(reasoningChunk string) { + if reasoningChunk != "" { + streamer.Reasoning(reasoningChunk) + } + }) + } + } + + if currentEvent == "complete" { + break + } + } + } + + // Flush remaining tokens in filters + thinkFilter.Flush(func(contentChunk string) { + toolFilter.Feed(contentChunk, func(cleanChunk string) { + if cleanChunk != "" { + streamer.Content(cleanChunk) + } + }, func(tc ToolCall) { + streamer.ToolCallDelta(tc) + }) + }, func(reasoningChunk string) { + if reasoningChunk != "" { + streamer.Reasoning(reasoningChunk) + } + }) + + toolFilter.Flush(func(cleanChunk string) { + if cleanChunk != "" { + streamer.Content(cleanChunk) + } + }, func(tc ToolCall) { + streamer.ToolCallDelta(tc) + }) + + if toolFilter.emittedCall { + streamer.Finish("tool_calls") + } else { + streamer.Finish("stop") + } + streamer.Done() + + return nil +} + +// --------------------------------------------------------------------------- +// HTTP Routes & Main Server +// --------------------------------------------------------------------------- + +func main() { + spaceFlag := flag.String("space", DefaultSpaceURL, "Target Gradio Space URL") + flag.StringVar(spaceFlag, "url", DefaultSpaceURL, "Alias for -space") + portFlag := flag.Int("port", 8080, "Gateway HTTP server port") + hostFlag := flag.String("host", "0.0.0.0", "Gateway HTTP server host") + socksFlag := flag.String("socks", "", "Optional SOCKS5 proxy URL (e.g. socks5://127.0.0.1:1080)") + flag.StringVar(socksFlag, "proxy", "", "Alias for -socks") + flag.StringVar(socksFlag, "socks5", "", "Alias for -socks") + uaFlag := flag.String("user-agent", "", "Custom User-Agent header") + flag.StringVar(uaFlag, "ua", "", "Alias for -user-agent") + timeoutFlag := flag.Int("timeout", 300, "Upstream timeout in seconds") + flag.Parse() + + // Environment variable fallbacks + if envSpace := os.Getenv("GRADIO_SPACE_URL"); envSpace != "" && *spaceFlag == DefaultSpaceURL { + *spaceFlag = envSpace + } + if *socksFlag == "" { + for _, envName := range []string{"ALL_PROXY", "all_proxy", "SOCKS5_PROXY", "socks5_proxy", "SOCKS_PROXY", "socks_proxy"} { + if p := os.Getenv(envName); p != "" { + *socksFlag = p + break + } + } + } + if *uaFlag != "" { + ConfiguredUserAgent = *uaFlag + } + + gateway := NewGradioGateway(*spaceFlag, *socksFlag, time.Duration(*timeoutFlag)*time.Second) + + mux := http.NewServeMux() + + // Health and Info + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + EnableCORS(w) + if r.Method == "OPTIONS" { + w.WriteHeader(http.StatusOK) + return + } + if r.URL.Path != "/" { + http.NotFound(w, r) + return + } + disc := gateway.GetDiscovery("", EffectiveUserAgent(r)) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]interface{}{ + "status": "running", + "service": "gr2gw", + "space_url": disc.SpaceURL, + "title": disc.Title, + "endpoint": disc.Endpoint, + "primary_model": disc.PrimaryModel, + "models": disc.Models, + "total_inputs": disc.TotalInputs, + "history_format": disc.HistoryFormat, + }) + }) + + // Models list + handleModels := func(w http.ResponseWriter, r *http.Request) { + EnableCORS(w) + if r.Method == "OPTIONS" { + w.WriteHeader(http.StatusOK) + return + } + spaceURL := gateway.defaultURL + if hdr := r.Header.Get("X-Gradio-Space"); hdr != "" { + spaceURL = hdr + } else if hdr := r.Header.Get("X-Space-URL"); hdr != "" { + spaceURL = hdr + } + disc := gateway.GetDiscovery(spaceURL, EffectiveUserAgent(r)) + resp := ModelsResponse{ + Object: "list", + Data: disc.GetModelList(), + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + } + + mux.HandleFunc("/models", handleModels) + mux.HandleFunc("/v1/models", handleModels) + + // Chat completions + handleCompletions := func(w http.ResponseWriter, r *http.Request) { + EnableCORS(w) + if r.Method == "OPTIONS" { + w.WriteHeader(http.StatusOK) + return + } + if r.Method != "POST" { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + var req ChatCompletionRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, fmt.Sprintf("Invalid JSON request: %v", err), http.StatusBadRequest) + return + } + + if err := gateway.ExecuteChatCompletion(w, r, req); err != nil { + log.Printf("Chat completion error: %v", err) + http.Error(w, fmt.Sprintf("Gateway error: %v", err), http.StatusBadGateway) + return + } + } + + mux.HandleFunc("/chat/completions", handleCompletions) + mux.HandleFunc("/v1/chat/completions", handleCompletions) + + addr := fmt.Sprintf("%s:%d", *hostFlag, *portFlag) + log.Printf("gr2gw listening on %s (target space: %s)", addr, *spaceFlag) + if *socksFlag != "" { + log.Printf("Using SOCKS5 proxy: %s", *socksFlag) + } + + server := &http.Server{ + Addr: addr, + Handler: mux, + ReadTimeout: time.Duration(*timeoutFlag+30) * time.Second, + WriteTimeout: time.Duration(*timeoutFlag+30) * time.Second, + } + + if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed { + log.Fatalf("Server failed: %v", err) + } +} diff --git a/gr2gw_test.go b/gr2gw_test.go new file mode 100644 index 0000000..c01399d --- /dev/null +++ b/gr2gw_test.go @@ -0,0 +1,249 @@ +package main + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func TestParseSOCKS5URL(t *testing.T) { + tests := []struct { + input string + expected *SOCKS5Config + }{ + {"", nil}, + {"127.0.0.1:1080", &SOCKS5Config{Address: "127.0.0.1:1080"}}, + {"socks5://127.0.0.1:9050", &SOCKS5Config{Address: "127.0.0.1:9050"}}, + {"socks5h://user:pass@10.0.0.1:1080", &SOCKS5Config{Address: "10.0.0.1:1080", Username: "user", Password: "pass"}}, + } + + for _, tc := range tests { + cfg, err := ParseSOCKS5URL(tc.input) + if err != nil { + t.Fatalf("unexpected error for %q: %v", tc.input, err) + } + if tc.expected == nil { + if cfg != nil { + t.Errorf("expected nil config, got %+v", cfg) + } + continue + } + if cfg.Address != tc.expected.Address || cfg.Username != tc.expected.Username || cfg.Password != tc.expected.Password { + t.Errorf("for %q, expected %+v, got %+v", tc.input, tc.expected, cfg) + } + } +} + +func TestChatMessageGetContentString(t *testing.T) { + m1 := ChatMessage{Role: "user", Content: "hello world"} + if m1.GetContentString() != "hello world" { + t.Errorf("expected 'hello world', got %q", m1.GetContentString()) + } + + m2 := ChatMessage{ + Role: "user", + Content: []interface{}{ + map[string]interface{}{"type": "text", "text": "part 1 "}, + map[string]interface{}{"type": "text", "text": "part 2"}, + }, + } + if m2.GetContentString() != "part 1 part 2" { + t.Errorf("expected 'part 1 part 2', got %q", m2.GetContentString()) + } +} + +func TestExtractThinking(t *testing.T) { + content := "Let me calculate 2+2.The answer is 4." + clean, reasoning := ExtractThinking(content) + if reasoning != "Let me calculate 2+2." { + t.Errorf("expected reasoning 'Let me calculate 2+2.', got %q", reasoning) + } + if clean != "The answer is 4." { + t.Errorf("expected clean 'The answer is 4.', got %q", clean) + } +} + +func TestDetectToolCalls(t *testing.T) { + xmlContent := ` +{"name": "get_weather", "arguments": {"city": "Paris"}} +` + calls, rem, ok := DetectToolCalls(xmlContent) + if !ok || len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d (ok: %v)", len(calls), ok) + } + if calls[0].Function.Name != "get_weather" { + t.Errorf("expected function name get_weather, got %q", calls[0].Function.Name) + } + if rem != "" { + t.Errorf("expected empty remaining content, got %q", rem) + } + + jsonContent := `{"name": "calculator", "arguments": {"expr": "1+1"}}` + calls2, rem2, ok2 := DetectToolCalls(jsonContent) + if !ok2 || len(calls2) != 1 { + t.Fatalf("expected 1 tool call from JSON, got %d", len(calls2)) + } + if calls2[0].Function.Name != "calculator" { + t.Errorf("expected function calculator, got %q", calls2[0].Function.Name) + } + if rem2 != "" { + t.Errorf("expected empty remaining, got %q", rem2) + } +} + +func TestStreamThinkingFilter(t *testing.T) { + filter := NewStreamThinkingFilter() + var contentParts []string + var reasoningParts []string + + onContent := func(s string) { contentParts = append(contentParts, s) } + onReasoning := func(s string) { reasoningParts = append(reasoningParts, s) } + + chunks := []string{"Thinking de", "eplyHere is your answer."} + for _, c := range chunks { + filter.Feed(c, onContent, onReasoning) + } + filter.Flush(onContent, onReasoning) + + fullReasoning := strings.Join(reasoningParts, "") + fullContent := strings.Join(contentParts, "") + + if fullReasoning != "Thinking deeply" { + t.Errorf("expected reasoning 'Thinking deeply', got %q", fullReasoning) + } + if fullContent != "Here is your answer." { + t.Errorf("expected content 'Here is your answer.', got %q", fullContent) + } +} + +func TestStreamToolCallFilter(t *testing.T) { + filter := NewStreamToolCallFilter() + var contentParts []string + var toolCalls []ToolCall + + onContent := func(s string) { contentParts = append(contentParts, s) } + onToolCall := func(tc ToolCall) { toolCalls = append(toolCalls, tc) } + + chunks := []string{ + "Searching now: ", + "\n{\"name\": \"search_web\", \"arguments\": {\"query\": \"golang\"}}\n", + " Done.", + } + + for _, c := range chunks { + filter.Feed(c, onContent, onToolCall) + } + filter.Flush(onContent, onToolCall) + + if len(toolCalls) != 1 { + t.Fatalf("expected 1 emitted tool call, got %d", len(toolCalls)) + } + if toolCalls[0].Function.Name != "search_web" { + t.Errorf("expected tool name 'search_web', got %q", toolCalls[0].Function.Name) + } + fullContent := strings.Join(contentParts, "") + if fullContent != "Searching now: Done." { + t.Errorf("expected 'Searching now: Done.', got %q", fullContent) + } +} + +func TestMockGradioServerCompletion(t *testing.T) { + // Setup a mock Gradio server + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/gradio_api/info" { + resp := GradioAPIInfoResponse{ + NamedEndpoints: map[string]GradioEndpointInfo{ + "/chat_fn": { + Parameters: []GradioParamInfo{ + {ParameterName: "message", Component: "Textbox"}, + }, + Returns: []GradioParamInfo{ + {ParameterName: "response", Component: "Json"}, + }, + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + return + } + + if r.URL.Path == "/gradio_api/call/chat_fn" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_123"}) + return + } + + if r.URL.Path == "/gradio_api/call/chat_fn/evt_123" { + w.Header().Set("Content-Type", "text/event-stream") + flusher, ok := w.(http.Flusher) + if !ok { + t.Fatal("expected flusher") + } + fmt.Fprintf(w, "event: generating\ndata: [\"Hello \", null]\n\n") + flusher.Flush() + fmt.Fprintf(w, "event: generating\ndata: [\"Hello world!\", null]\n\n") + flusher.Flush() + fmt.Fprintf(w, "event: complete\ndata: [\"Hello world!\", null]\n\n") + flusher.Flush() + return + } + + http.NotFound(w, r) + })) + defer ts.Close() + + gw := NewGradioGateway(ts.URL, "", 10*time.Second) + + // 1. Test Non-streaming request + reqBody := ChatCompletionRequest{ + Model: "test-model", + Messages: []ChatMessage{ + {Role: "user", Content: "Hi"}, + }, + Stream: false, + } + b, _ := json.Marshal(reqBody) + httpReq := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(b)) + httpReq.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + err := gw.ExecuteChatCompletion(rec, httpReq, reqBody) + if err != nil { + t.Fatalf("unexpected completion error: %v", err) + } + + var resp ChatCompletionResponse + if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil { + t.Fatalf("failed to decode completion response: %v", err) + } + if len(resp.Choices) != 1 { + t.Fatalf("expected 1 choice, got %d", len(resp.Choices)) + } + if resp.Choices[0].Message.GetContentString() != "Hello world!" { + t.Errorf("expected 'Hello world!', got %q", resp.Choices[0].Message.GetContentString()) + } + + // 2. Test Streaming request + reqBodyStream := reqBody + reqBodyStream.Stream = true + recStream := httptest.NewRecorder() + err = gw.ExecuteChatCompletion(recStream, httpReq, reqBodyStream) + if err != nil { + t.Fatalf("unexpected streaming error: %v", err) + } + streamOutput := recStream.Body.String() + if !strings.Contains(streamOutput, "data: [DONE]") { + t.Errorf("expected stream to contain [DONE], got:\n%s", streamOutput) + } + if !strings.Contains(streamOutput, "Hello world!") && !strings.Contains(streamOutput, "world!") { + t.Errorf("expected stream output to contain delta tokens, got:\n%s", streamOutput) + } +}