1473 lines
44 KiB
Go
1473 lines
44 KiB
Go
// k3r053n3: Standalone OpenAI-compatible gateway for Kimi K3 (cw-105-kimi-k3-gguf-demo)
|
|
// Created by Luxferre in 2026, released into the public domain
|
|
|
|
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"os"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
DefaultUserAgent = "Mozilla/5.0 (X11; Linux x86_64; rv:153.0) Gecko/20100101 Firefox/153.0"
|
|
ConfiguredUserAgent string
|
|
fnCallSyntaxRegex = regexp.MustCompile(`^([a-zA-Z0-9_\-\.]+)\s*\((.*)\)$`)
|
|
kvPairRegex = regexp.MustCompile(`([a-zA-Z0-9_]+)\s*=\s*(?:"([^"]*)"|'([^']*)'|([^,\s]+))`)
|
|
toolTagRegexes = []*regexp.Regexp{
|
|
regexp.MustCompile(`(?s)<tool_call>\s*(.*?)\s*</tool_call>`),
|
|
regexp.MustCompile(`(?s)<tool-call>\s*(.*?)\s*</tool-call>`),
|
|
regexp.MustCompile(`(?s)<function_call>\s*(.*?)\s*</function_call>`),
|
|
regexp.MustCompile(`(?s)<tool>\s*(.*?)\s*</tool>`),
|
|
regexp.MustCompile(`(?s)<action>\s*(.*?)\s*</action>`),
|
|
regexp.MustCompile("(?s)```tool_call\\s*(.*?)\\s*```"),
|
|
regexp.MustCompile("(?s)```tool-call\\s*(.*?)\\s*```"),
|
|
regexp.MustCompile("(?s)```json\\s*(.*?)\\s*```"),
|
|
}
|
|
)
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 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 {
|
|
switch v := m.Content.(type) {
|
|
case string:
|
|
return v
|
|
case []interface{}:
|
|
var sb strings.Builder
|
|
for _, p := range v {
|
|
if s, ok := p.(string); ok {
|
|
sb.WriteString(s)
|
|
} else if tm, ok := p.(map[string]interface{}); ok {
|
|
if t, ok := tm["text"].(string); ok {
|
|
sb.WriteString(t)
|
|
}
|
|
}
|
|
}
|
|
return sb.String()
|
|
default:
|
|
if m.Content == nil {
|
|
return ""
|
|
}
|
|
b, _ := json.Marshal(m.Content)
|
|
return string(b)
|
|
}
|
|
}
|
|
|
|
type ChatCompletionRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []ChatMessage `json:"messages"`
|
|
Tools []Tool `json:"tools,omitempty"`
|
|
Functions []interface{} `json:"functions,omitempty"`
|
|
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
|
FunctionCall interface{} `json:"function_call,omitempty"`
|
|
Stream bool `json:"stream"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
MaxCompletionTokens int `json:"max_completion_tokens"`
|
|
Temperature *float64 `json:"temperature,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"`
|
|
}
|
|
|
|
type GradioJoinResponse struct {
|
|
EventID string `json:"event_id"`
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Zero-Dependency Pure Go SOCKS5 Proxy Client (RFC 1928 / RFC 1929)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type SOCKS5Config struct {
|
|
Address, Username, Password string
|
|
}
|
|
|
|
func ParseSOCKS5URL(proxyURL string) (*SOCKS5Config, error) {
|
|
u := strings.TrimSpace(proxyURL)
|
|
if u == "" {
|
|
return nil, nil
|
|
}
|
|
u = strings.TrimPrefix(strings.TrimPrefix(u, "socks5h://"), "socks5://")
|
|
cfg := &SOCKS5Config{}
|
|
if at := strings.LastIndex(u, "@"); at != -1 {
|
|
userPass := u[:at]
|
|
cfg.Address = u[at+1:]
|
|
if col := strings.Index(userPass, ":"); col != -1 {
|
|
cfg.Username, cfg.Password = userPass[:col], userPass[col+1:]
|
|
} else {
|
|
cfg.Username = userPass
|
|
}
|
|
} else {
|
|
cfg.Address = u
|
|
}
|
|
if !strings.Contains(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 || cfg == nil {
|
|
return (&net.Dialer{}).DialContext(ctx, "tcp", targetAddr)
|
|
}
|
|
conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", cfg.Address)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("socks5 dial failed: %w", err)
|
|
}
|
|
if d, ok := ctx.Deadline(); ok {
|
|
conn.SetDeadline(d)
|
|
} else {
|
|
conn.SetDeadline(time.Now().Add(30 * time.Second))
|
|
}
|
|
defer conn.SetDeadline(time.Time{})
|
|
|
|
greeting := []byte{0x05, 0x01, 0x00}
|
|
if cfg.Username != "" {
|
|
greeting = []byte{0x05, 0x02, 0x00, 0x02}
|
|
}
|
|
if _, err := conn.Write(greeting); err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
resp := make([]byte, 2)
|
|
if _, err := io.ReadFull(conn, resp); err != nil || resp[0] != 0x05 {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("socks5 greeting failed")
|
|
}
|
|
|
|
if resp[1] == 0x02 {
|
|
req := append(append([]byte{0x01, byte(len(cfg.Username))}, cfg.Username...), byte(len(cfg.Password)))
|
|
req = append(req, cfg.Password...)
|
|
if _, err := conn.Write(req); err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
if _, err := io.ReadFull(conn, resp); err != nil || resp[1] != 0x00 {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("socks5 auth failed")
|
|
}
|
|
} else if resp[1] != 0x00 {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("socks5 auth rejected: 0x%02x", resp[1])
|
|
}
|
|
|
|
host, portStr, err := net.SplitHostPort(targetAddr)
|
|
if err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
port, _ := strconv.Atoi(portStr)
|
|
|
|
req := []byte{0x05, 0x01, 0x00}
|
|
ip := net.ParseIP(host)
|
|
if ip4 := ip.To4(); ip4 != nil {
|
|
req = append(append(req, 0x01), ip4...)
|
|
} else if ip6 := ip.To16(); ip6 != nil {
|
|
req = append(append(req, 0x04), ip6...)
|
|
} else {
|
|
req = append(append(req, 0x03, byte(len(host))), host...)
|
|
}
|
|
req = append(req, byte(port>>8), byte(port&0xFF))
|
|
|
|
if _, err := conn.Write(req); err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
|
|
respHdr := make([]byte, 4)
|
|
if _, err := io.ReadFull(conn, respHdr); err != nil || respHdr[1] != 0x00 {
|
|
conn.Close()
|
|
return nil, fmt.Errorf("socks5 connect failed: 0x%02x", respHdr[1])
|
|
}
|
|
|
|
switch respHdr[3] {
|
|
case 0x01:
|
|
io.ReadFull(conn, make([]byte, 6))
|
|
case 0x03:
|
|
lb := make([]byte, 1)
|
|
io.ReadFull(conn, lb)
|
|
io.ReadFull(conn, make([]byte, int(lb[0])+2))
|
|
case 0x04:
|
|
io.ReadFull(conn, make([]byte, 18))
|
|
}
|
|
return conn, nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers & Tool Calling
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func GenerateUUID() string {
|
|
var b [16]byte
|
|
rand.Read(b[:])
|
|
b[6], b[8] = (b[6]&0x0f)|0x40, (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 DoWithFibonacciRetry(client *http.Client, makeReq func() (*http.Request, error), maxRetries int) (*http.Response, error) {
|
|
var lastErr error
|
|
a, b := 1, 1
|
|
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 {
|
|
body, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
lastErr = fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
|
|
} else {
|
|
lastErr = err
|
|
}
|
|
if attempt < maxRetries {
|
|
time.Sleep(time.Duration(a) * time.Second)
|
|
a, b = b, a+b
|
|
}
|
|
}
|
|
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")
|
|
}
|
|
|
|
func ResolveMaxTokens(req ChatCompletionRequest, def int) int {
|
|
mt := req.MaxTokens
|
|
if mt == 0 && req.MaxCompletionTokens > 0 {
|
|
mt = req.MaxCompletionTokens
|
|
}
|
|
if mt <= 0 {
|
|
mt = def
|
|
}
|
|
if mt < 256 {
|
|
return 256
|
|
}
|
|
if mt > 8192 {
|
|
return 8192
|
|
}
|
|
return mt
|
|
}
|
|
|
|
func ResolveTemperature(req ChatCompletionRequest, def float64) float64 {
|
|
if req.Temperature != nil {
|
|
t := *req.Temperature
|
|
if t < 0.0 {
|
|
return 0.0
|
|
}
|
|
if t > 1.5 {
|
|
return 1.5
|
|
}
|
|
return t
|
|
}
|
|
return def
|
|
}
|
|
|
|
func ResolveReasoningEffort(req ChatCompletionRequest, def string) string {
|
|
e := strings.ToLower(strings.TrimSpace(req.ReasoningEffort))
|
|
if e == "" {
|
|
e = strings.ToLower(strings.TrimSpace(def))
|
|
}
|
|
switch e {
|
|
case "max", "high", "low":
|
|
return e
|
|
case "default", "medium", "none", "off":
|
|
return "default"
|
|
default:
|
|
return "max"
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
func formatArgumentsString(args interface{}) string {
|
|
if args == nil {
|
|
return "{}"
|
|
}
|
|
if str, ok := args.(string); ok {
|
|
return str
|
|
}
|
|
b, err := json.Marshal(args)
|
|
if err != nil {
|
|
return "{}"
|
|
}
|
|
return string(b)
|
|
}
|
|
|
|
func parseFnArgsString(raw string) string {
|
|
raw = strings.TrimSpace(raw)
|
|
if raw == "" {
|
|
return "{}"
|
|
}
|
|
var m map[string]interface{}
|
|
if err := json.Unmarshal([]byte(raw), &m); err == nil {
|
|
b, _ := json.Marshal(m)
|
|
return string(b)
|
|
}
|
|
kvMap := make(map[string]interface{})
|
|
for _, match := range kvPairRegex.FindAllStringSubmatch(raw, -1) {
|
|
k := match[1]
|
|
v := match[2]
|
|
if v == "" {
|
|
v = match[3]
|
|
}
|
|
if v == "" {
|
|
v = match[4]
|
|
}
|
|
kvMap[k] = v
|
|
}
|
|
if len(kvMap) > 0 {
|
|
b, _ := json.Marshal(kvMap)
|
|
return string(b)
|
|
}
|
|
b, _ := json.Marshal(map[string]string{"input": raw})
|
|
return string(b)
|
|
}
|
|
|
|
func BuildToolsSystemPrompt(tools []Tool, toolChoice interface{}) string {
|
|
if len(tools) == 0 {
|
|
return ""
|
|
}
|
|
b, _ := json.MarshalIndent(tools, "", " ")
|
|
|
|
choiceInstruction := ""
|
|
if tcStr, ok := toolChoice.(string); ok {
|
|
if tcStr == "required" {
|
|
choiceInstruction = "\nIMPORTANT: You MUST invoke at least one tool to satisfy the request."
|
|
}
|
|
} else if tcMap, ok := toolChoice.(map[string]interface{}); ok {
|
|
if fn, ok := tcMap["function"].(map[string]interface{}); ok {
|
|
if fnName, ok := fn["name"].(string); ok && fnName != "" {
|
|
choiceInstruction = fmt.Sprintf("\nIMPORTANT: You MUST call the %q tool.", fnName)
|
|
}
|
|
}
|
|
}
|
|
|
|
return fmt.Sprintf("# Tools Available\nYou have access to the following tools:\n%s\n\n# Tool Calling Instructions\n1. When a task requires gathering information, inspecting files, running commands, or executing actions, you MUST emit the <tool_call>...</tool_call> block in your response.\n2. DO NOT output conversational promises or filler statements (e.g. \"I will analyze the project\", \"Let me check the files\") without emitting the tool call.\n3. Wrap tool calls inside <tool_call>...</tool_call> XML tags:\n<tool_call>\n{\"name\": \"tool_name\", \"arguments\": {\"param\": \"value\"}}\n</tool_call>%s", string(b), choiceInstruction)
|
|
}
|
|
|
|
func extractToolCallFromMap(m map[string]interface{}, idx int) *ToolCall {
|
|
var name string
|
|
var args interface{}
|
|
|
|
for _, k := range []string{"name", "tool", "tool_name", "function_name", "action"} {
|
|
if n, ok := m[k].(string); ok && n != "" {
|
|
name = n
|
|
break
|
|
}
|
|
}
|
|
if name == "" {
|
|
if fn, ok := m["function"].(map[string]interface{}); ok {
|
|
for _, k := range []string{"name", "tool", "tool_name", "function_name"} {
|
|
if n, ok := fn[k].(string); ok && n != "" {
|
|
name = n
|
|
args = fn["arguments"]
|
|
if args == nil {
|
|
args = fn["parameters"]
|
|
}
|
|
if args == nil {
|
|
args = fn["args"]
|
|
}
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if name == "" {
|
|
return nil
|
|
}
|
|
|
|
if args == nil {
|
|
for _, k := range []string{"arguments", "parameters", "args", "params", "input", "action_input"} {
|
|
if v, exists := m[k]; exists && v != nil {
|
|
args = v
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
if args == nil {
|
|
rem := make(map[string]interface{})
|
|
for k, v := range m {
|
|
if k != "name" && k != "function" && k != "tool" && k != "tool_name" && k != "function_name" && k != "action" && k != "type" && k != "id" && k != "index" {
|
|
rem[k] = v
|
|
}
|
|
}
|
|
if len(rem) > 0 {
|
|
args = rem
|
|
} else {
|
|
args = map[string]interface{}{}
|
|
}
|
|
}
|
|
|
|
id := fmt.Sprintf("call_%s_%d", GenerateUUID()[:8], idx)
|
|
if existingID, ok := m["id"].(string); ok && existingID != "" {
|
|
id = existingID
|
|
}
|
|
|
|
return &ToolCall{
|
|
Index: &idx,
|
|
ID: id,
|
|
Type: "function",
|
|
Function: ToolCallFunction{
|
|
Name: name,
|
|
Arguments: formatArgumentsString(args),
|
|
},
|
|
}
|
|
}
|
|
|
|
func parseRawToolPayload(raw string) []ToolCall {
|
|
raw = strings.TrimSpace(raw)
|
|
if raw == "" {
|
|
return nil
|
|
}
|
|
|
|
var arr []map[string]interface{}
|
|
if err := json.Unmarshal([]byte(raw), &arr); err == nil {
|
|
var res []ToolCall
|
|
for _, item := range arr {
|
|
if tc := extractToolCallFromMap(item, len(res)); tc != nil {
|
|
res = append(res, *tc)
|
|
}
|
|
}
|
|
if len(res) > 0 {
|
|
return res
|
|
}
|
|
}
|
|
|
|
var single map[string]interface{}
|
|
if err := json.Unmarshal([]byte(raw), &single); err == nil {
|
|
if tcList, ok := single["tool_calls"].([]interface{}); ok {
|
|
var res []ToolCall
|
|
for _, item := range tcList {
|
|
if m, ok := item.(map[string]interface{}); ok {
|
|
if tc := extractToolCallFromMap(m, len(res)); tc != nil {
|
|
res = append(res, *tc)
|
|
}
|
|
}
|
|
}
|
|
if len(res) > 0 {
|
|
return res
|
|
}
|
|
}
|
|
if tc := extractToolCallFromMap(single, 0); tc != nil {
|
|
return []ToolCall{*tc}
|
|
}
|
|
}
|
|
|
|
if fnMatch := fnCallSyntaxRegex.FindStringSubmatch(raw); len(fnMatch) >= 3 {
|
|
fnName := fnMatch[1]
|
|
argsRaw := strings.TrimSpace(fnMatch[2])
|
|
idx := 0
|
|
return []ToolCall{{
|
|
Index: &idx,
|
|
ID: fmt.Sprintf("call_%s_0", GenerateUUID()[:8]),
|
|
Type: "function",
|
|
Function: ToolCallFunction{
|
|
Name: fnName,
|
|
Arguments: parseFnArgsString(argsRaw),
|
|
},
|
|
}}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func DetectToolCalls(text string) ([]ToolCall, string, bool) {
|
|
var toolCalls []ToolCall
|
|
cleaned := text
|
|
|
|
for _, re := range toolTagRegexes {
|
|
matches := re.FindAllStringSubmatchIndex(cleaned, -1)
|
|
if len(matches) == 0 {
|
|
continue
|
|
}
|
|
var sb strings.Builder
|
|
lastIdx := 0
|
|
foundInRe := false
|
|
for _, m := range matches {
|
|
raw := strings.TrimSpace(cleaned[m[2]:m[3]])
|
|
parsed := parseRawToolPayload(raw)
|
|
if len(parsed) > 0 {
|
|
foundInRe = true
|
|
sb.WriteString(cleaned[lastIdx:m[0]])
|
|
lastIdx = m[1]
|
|
for _, tc := range parsed {
|
|
idx := len(toolCalls)
|
|
tc.Index = &idx
|
|
toolCalls = append(toolCalls, tc)
|
|
}
|
|
}
|
|
}
|
|
if foundInRe {
|
|
sb.WriteString(cleaned[lastIdx:])
|
|
cleaned = sb.String()
|
|
}
|
|
}
|
|
|
|
if len(toolCalls) == 0 {
|
|
for _, tag := range []string{"<tool_call>", "<tool-call>", "<function_call>", "<tool>", "<action>"} {
|
|
if idx := strings.Index(cleaned, tag); idx != -1 {
|
|
raw := strings.TrimSpace(cleaned[idx+len(tag):])
|
|
for _, tc := range parseRawToolPayload(raw) {
|
|
idxTC := len(toolCalls)
|
|
tc.Index = &idxTC
|
|
toolCalls = append(toolCalls, tc)
|
|
}
|
|
if len(toolCalls) > 0 {
|
|
cleaned = strings.TrimSpace(cleaned[:idx])
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return toolCalls, strings.TrimSpace(cleaned), len(toolCalls) > 0
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Streaming & Snapshot Parser
|
|
// ---------------------------------------------------------------------------
|
|
|
|
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) SendChunk(delta StreamDelta, finish ...string) {
|
|
var finishReason *string
|
|
if len(finish) > 0 && finish[0] != "" {
|
|
finishReason = &finish[0]
|
|
}
|
|
chunk := StreamResponse{
|
|
ID: s.id, Object: "chat.completion.chunk", Created: s.created, Model: s.model,
|
|
Choices: []StreamChoice{{Index: 0, Delta: delta, FinishReason: finishReason}},
|
|
}
|
|
b, _ := json.Marshal(chunk)
|
|
fmt.Fprintf(s.w, "data: %s\n\n", b)
|
|
if s.flusher != nil {
|
|
s.flusher.Flush()
|
|
}
|
|
}
|
|
|
|
func (s *Streamer) Done() {
|
|
fmt.Fprintf(s.w, "data: [DONE]\n\n")
|
|
if s.flusher != nil {
|
|
s.flusher.Flush()
|
|
}
|
|
}
|
|
|
|
func cleanReasoningArtifacts(s string) string {
|
|
s = strings.TrimPrefix(s, "<>")
|
|
s = strings.TrimPrefix(s, ">")
|
|
s = strings.TrimSuffix(s, "</>")
|
|
s = strings.TrimSuffix(s, "</")
|
|
s = strings.TrimSuffix(s, ">")
|
|
return s
|
|
}
|
|
|
|
func cleanContentArtifacts(s string) string {
|
|
if s == "…" || s == "<" || s == ">" || s == "<>" || s == "</" || s == "</>" {
|
|
return ""
|
|
}
|
|
s = strings.TrimPrefix(s, ">\n\n")
|
|
s = strings.TrimPrefix(s, ">\n")
|
|
s = strings.TrimPrefix(s, ">")
|
|
s = strings.TrimPrefix(s, "<>")
|
|
s = strings.TrimSuffix(s, "</>")
|
|
s = strings.TrimSuffix(s, "</")
|
|
return s
|
|
}
|
|
|
|
func parseGradioSnapshot(dataJSON string) (string, string, bool) {
|
|
var raw []interface{}
|
|
if err := json.Unmarshal([]byte(dataJSON), &raw); err != nil || len(raw) == 0 {
|
|
return "", "", false
|
|
}
|
|
msgs, ok := raw[0].([]interface{})
|
|
if !ok || len(msgs) == 0 {
|
|
return "", "", false
|
|
}
|
|
|
|
lastUserIdx := -1
|
|
for i := len(msgs) - 1; i >= 0; i-- {
|
|
if m, ok := msgs[i].(map[string]interface{}); ok && m["role"] == "user" {
|
|
lastUserIdx = i
|
|
break
|
|
}
|
|
}
|
|
|
|
var assistantMsgs []interface{}
|
|
if lastUserIdx != -1 && lastUserIdx+1 < len(msgs) {
|
|
assistantMsgs = msgs[lastUserIdx+1:]
|
|
} else {
|
|
assistantMsgs = msgs
|
|
}
|
|
|
|
var reasoningParts []string
|
|
var contentParts []string
|
|
var rawAllText strings.Builder
|
|
|
|
for _, item := range assistantMsgs {
|
|
m, ok := item.(map[string]interface{})
|
|
if !ok || m["role"] != "assistant" {
|
|
continue
|
|
}
|
|
|
|
var textParts strings.Builder
|
|
if cArr, ok := m["content"].([]interface{}); ok {
|
|
for _, part := range cArr {
|
|
if cMap, ok := part.(map[string]interface{}); ok {
|
|
if rawText, ok := cMap["text"].(string); ok {
|
|
textParts.WriteString(rawText)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
msgText := textParts.String()
|
|
rawAllText.WriteString(msgText)
|
|
|
|
if msgText == "…" || msgText == "<" || msgText == ">" || msgText == "<>" || msgText == "</" || msgText == "</>" {
|
|
continue
|
|
}
|
|
|
|
isReasoning := false
|
|
if meta, ok := m["metadata"].(map[string]interface{}); ok && meta != nil {
|
|
if t, ok := meta["title"].(string); ok && strings.EqualFold(t, "Reasoning") {
|
|
isReasoning = true
|
|
} else if st, ok := meta["status"].(string); ok && (st == "pending" || st == "done") {
|
|
isReasoning = true
|
|
}
|
|
}
|
|
|
|
if !isReasoning {
|
|
if strings.HasPrefix(msgText, "<>") || strings.HasSuffix(msgText, "</") || strings.HasSuffix(msgText, "</>") {
|
|
isReasoning = true
|
|
} else if strings.HasPrefix(msgText, ">") && !strings.HasPrefix(msgText, ">\n") {
|
|
isReasoning = true
|
|
}
|
|
}
|
|
|
|
if isReasoning {
|
|
cleaned := cleanReasoningArtifacts(msgText)
|
|
if cleaned != "" {
|
|
reasoningParts = append(reasoningParts, cleaned)
|
|
}
|
|
} else {
|
|
cleaned := cleanContentArtifacts(msgText)
|
|
if cleaned != "" {
|
|
contentParts = append(contentParts, cleaned)
|
|
}
|
|
}
|
|
}
|
|
|
|
reasoning := strings.Join(reasoningParts, "")
|
|
content := strings.Join(contentParts, "")
|
|
|
|
if reasoning != "" || content != "" {
|
|
return content, reasoning, true
|
|
}
|
|
|
|
combined := rawAllText.String()
|
|
if combined == "…" || combined == "<" || combined == ">" || combined == "<>" {
|
|
return "", "", true
|
|
}
|
|
|
|
if strings.Contains(combined, "<think>") {
|
|
tIdx := strings.Index(combined, "<think>")
|
|
endIdx := strings.Index(combined, "</think>")
|
|
if endIdx != -1 {
|
|
return strings.TrimLeft(combined[:tIdx]+combined[endIdx+8:], "\n"), combined[tIdx+7 : endIdx], true
|
|
}
|
|
return strings.TrimSpace(combined[:tIdx]), combined[tIdx+7:], true
|
|
}
|
|
|
|
if strings.HasPrefix(combined, "<>") {
|
|
inner := combined[2:]
|
|
if idx := strings.Index(inner, "</>"); idx != -1 {
|
|
return strings.TrimLeft(inner[idx+3:], "\n"), inner[:idx], true
|
|
}
|
|
if idx := strings.LastIndex(inner, "</"); idx != -1 {
|
|
return strings.TrimLeft(strings.TrimPrefix(inner[idx+2:], ">"), "\n"), inner[:idx], true
|
|
}
|
|
return "", inner, true
|
|
}
|
|
|
|
return cleanContentArtifacts(combined), "", true
|
|
}
|
|
|
|
type StreamToolInterceptor struct {
|
|
streamer *Streamer
|
|
buffer string
|
|
inToolTag bool
|
|
matchedOpen string
|
|
closingTag string
|
|
emittedContent string
|
|
toolCallCount int
|
|
}
|
|
|
|
func NewStreamToolInterceptor(streamer *Streamer) *StreamToolInterceptor {
|
|
return &StreamToolInterceptor{streamer: streamer}
|
|
}
|
|
|
|
func (si *StreamToolInterceptor) HasToolCalls() bool {
|
|
return si.toolCallCount > 0
|
|
}
|
|
|
|
func (si *StreamToolInterceptor) ProcessContentDelta(delta string) {
|
|
if delta == "" {
|
|
return
|
|
}
|
|
si.buffer += delta
|
|
|
|
tags := []struct{ open, close string }{
|
|
{"<tool_call>", "</tool_call>"},
|
|
{"<tool-call>", "</tool-call>"},
|
|
{"<function_call>", "</function_call>"},
|
|
{"<tool>", "</tool>"},
|
|
{"<action>", "</action>"},
|
|
{"```tool_call", "```"},
|
|
{"```tool-call", "```"},
|
|
{"```json", "```"},
|
|
}
|
|
|
|
for len(si.buffer) > 0 {
|
|
if !si.inToolTag {
|
|
openIdx := -1
|
|
var foundOpen, foundClose string
|
|
for _, t := range tags {
|
|
if idx := strings.Index(si.buffer, t.open); idx != -1 {
|
|
if openIdx == -1 || idx < openIdx {
|
|
openIdx = idx
|
|
foundOpen = t.open
|
|
foundClose = t.close
|
|
}
|
|
}
|
|
}
|
|
|
|
if openIdx == -1 {
|
|
maxOverlap := 0
|
|
for _, t := range tags {
|
|
for o := len(t.open) - 1; o > 0; o-- {
|
|
if strings.HasSuffix(si.buffer, t.open[:o]) && o > maxOverlap {
|
|
maxOverlap = o
|
|
}
|
|
}
|
|
}
|
|
if maxOverlap > 0 {
|
|
toEmit := si.buffer[:len(si.buffer)-maxOverlap]
|
|
if toEmit != "" {
|
|
si.streamer.SendChunk(StreamDelta{Content: toEmit})
|
|
si.emittedContent += toEmit
|
|
si.buffer = si.buffer[len(toEmit):]
|
|
}
|
|
return
|
|
}
|
|
si.streamer.SendChunk(StreamDelta{Content: si.buffer})
|
|
si.emittedContent += si.buffer
|
|
si.buffer = ""
|
|
return
|
|
}
|
|
|
|
if openIdx > 0 {
|
|
pre := si.buffer[:openIdx]
|
|
si.streamer.SendChunk(StreamDelta{Content: pre})
|
|
si.emittedContent += pre
|
|
}
|
|
si.inToolTag = true
|
|
si.matchedOpen = foundOpen
|
|
si.closingTag = foundClose
|
|
si.buffer = si.buffer[openIdx+len(foundOpen):]
|
|
}
|
|
|
|
if si.inToolTag {
|
|
closeIdx := strings.Index(si.buffer, si.closingTag)
|
|
if closeIdx == -1 {
|
|
return
|
|
}
|
|
toolJSON := strings.TrimSpace(si.buffer[:closeIdx])
|
|
si.buffer = si.buffer[closeIdx+len(si.closingTag):]
|
|
si.inToolTag = false
|
|
parsed := parseRawToolPayload(toolJSON)
|
|
if len(parsed) > 0 {
|
|
for _, tc := range parsed {
|
|
idx := si.toolCallCount
|
|
si.toolCallCount++
|
|
tc.Index = &idx
|
|
si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
|
|
}
|
|
} else {
|
|
rawReconstruct := si.matchedOpen + toolJSON + si.closingTag
|
|
si.streamer.SendChunk(StreamDelta{Content: rawReconstruct})
|
|
si.emittedContent += rawReconstruct
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (si *StreamToolInterceptor) FlushRemaining() {
|
|
if si.buffer == "" {
|
|
return
|
|
}
|
|
if si.inToolTag {
|
|
raw := strings.TrimSpace(strings.TrimSuffix(si.buffer, si.closingTag))
|
|
tcs := parseRawToolPayload(raw)
|
|
if len(tcs) > 0 {
|
|
for _, tc := range tcs {
|
|
idx := si.toolCallCount
|
|
si.toolCallCount++
|
|
tc.Index = &idx
|
|
si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
|
|
}
|
|
si.buffer = ""
|
|
return
|
|
}
|
|
}
|
|
tcs, cleaned, ok := DetectToolCalls(si.buffer)
|
|
if ok {
|
|
if cleaned != "" {
|
|
si.streamer.SendChunk(StreamDelta{Content: cleaned})
|
|
si.emittedContent += cleaned
|
|
}
|
|
for _, tc := range tcs {
|
|
idx := si.toolCallCount
|
|
si.toolCallCount++
|
|
tc.Index = &idx
|
|
si.streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
|
|
}
|
|
} else {
|
|
si.streamer.SendChunk(StreamDelta{Content: si.buffer})
|
|
si.emittedContent += si.buffer
|
|
}
|
|
si.buffer = ""
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Kimi K3 Service & Upstream Logic
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type KimiService struct {
|
|
endpoint, defaultModel, defaultBackend, defaultReason, socks5Proxy string
|
|
defaultTokens int
|
|
defaultTemp float64
|
|
client *http.Client
|
|
}
|
|
|
|
func NewKimiService(endpoint, model, backend, reasoning, socks5Proxy string, tokens int, temp float64) *KimiService {
|
|
if model == "" {
|
|
model = "kimi-k3"
|
|
}
|
|
if backend == "" {
|
|
backend = "direct:together"
|
|
}
|
|
if reasoning == "" {
|
|
reasoning = "max"
|
|
}
|
|
if tokens <= 0 {
|
|
tokens = 8192
|
|
}
|
|
if temp <= 0 {
|
|
temp = 0.7
|
|
}
|
|
cleanProxy := strings.TrimSpace(socks5Proxy)
|
|
transport := &http.Transport{
|
|
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
|
|
if cleanProxy != "" {
|
|
return DialSOCKS5(ctx, cleanProxy, addr)
|
|
}
|
|
return (&net.Dialer{}).DialContext(ctx, network, addr)
|
|
},
|
|
ForceAttemptHTTP2: true,
|
|
MaxIdleConns: 100,
|
|
IdleConnTimeout: 90 * time.Second,
|
|
TLSHandshakeTimeout: 15 * time.Second,
|
|
}
|
|
return &KimiService{
|
|
endpoint: strings.TrimRight(endpoint, "/"),
|
|
defaultModel: model,
|
|
defaultBackend: backend,
|
|
defaultReason: reasoning,
|
|
socks5Proxy: cleanProxy,
|
|
defaultTokens: tokens,
|
|
defaultTemp: temp,
|
|
client: &http.Client{Transport: transport, Timeout: 300 * time.Second},
|
|
}
|
|
}
|
|
|
|
func (s *KimiService) ListModels() []ModelItem {
|
|
now := time.Now().Unix()
|
|
return []ModelItem{
|
|
{ID: s.defaultModel, Object: "model", Created: now, OwnedBy: "moonshotai"},
|
|
{ID: "moonshotai/Kimi-K3", Object: "model", Created: now, OwnedBy: "moonshotai"},
|
|
{ID: "kimi-k3:together", Object: "model", Created: now, OwnedBy: "together"},
|
|
{ID: "kimi-k3:fireworks", Object: "model", Created: now, OwnedBy: "fireworks"},
|
|
{ID: "kimi-k3:hf-together", Object: "model", Created: now, OwnedBy: "huggingface"},
|
|
{ID: "kimi-k3:hf-fireworks", Object: "model", Created: now, OwnedBy: "huggingface"},
|
|
{ID: "kimi-k3:hf-featherless", Object: "model", Created: now, OwnedBy: "huggingface"},
|
|
{ID: "kimi-k3:hf-baseten", Object: "model", Created: now, OwnedBy: "huggingface"},
|
|
}
|
|
}
|
|
|
|
func (s *KimiService) ResolveBackend(modelName string) string {
|
|
if strings.Contains(modelName, ":") {
|
|
switch strings.SplitN(modelName, ":", 2)[1] {
|
|
case "together", "direct-together", "direct:together":
|
|
return "direct:together"
|
|
case "fireworks", "direct-fireworks", "direct:fireworks":
|
|
return "direct:fireworks"
|
|
case "hf-together", "hf:together":
|
|
return "hf:together"
|
|
case "hf-fireworks", "hf:fireworks", "hf:fireworks-ai":
|
|
return "hf:fireworks-ai"
|
|
case "hf-featherless", "hf:featherless", "hf:featherless-ai":
|
|
return "hf:featherless-ai"
|
|
case "hf-baseten", "hf:baseten":
|
|
return "hf:baseten"
|
|
}
|
|
}
|
|
return s.defaultBackend
|
|
}
|
|
|
|
func formatToolCallBlock(tc ToolCall) string {
|
|
var argsObj interface{}
|
|
argsRaw := strings.TrimSpace(tc.Function.Arguments)
|
|
if argsRaw == "" {
|
|
argsObj = map[string]interface{}{}
|
|
} else if err := json.Unmarshal([]byte(argsRaw), &argsObj); err != nil {
|
|
argsObj = map[string]string{"input": argsRaw}
|
|
}
|
|
b, _ := json.Marshal(map[string]interface{}{
|
|
"name": tc.Function.Name,
|
|
"arguments": argsObj,
|
|
})
|
|
return fmt.Sprintf("<tool_call>\n%s\n</tool_call>", string(b))
|
|
}
|
|
|
|
func formatMessageContent(m ChatMessage) (string, string) {
|
|
role := m.Role
|
|
cStr := m.GetContentString()
|
|
if role == "tool" || role == "function" {
|
|
tName := m.Name
|
|
if tName == "" {
|
|
tName = m.ToolCallID
|
|
}
|
|
if tName != "" {
|
|
cStr = fmt.Sprintf("[Tool Result for %s]: %s", tName, cStr)
|
|
} else {
|
|
cStr = fmt.Sprintf("[Tool Result]: %s", cStr)
|
|
}
|
|
role = "user"
|
|
} else if role == "assistant" && len(m.ToolCalls) > 0 {
|
|
var tcParts []string
|
|
for _, tc := range m.ToolCalls {
|
|
tcParts = append(tcParts, formatToolCallBlock(tc))
|
|
}
|
|
if cStr != "" {
|
|
cStr += "\n\n" + strings.Join(tcParts, "\n")
|
|
} else {
|
|
cStr = strings.Join(tcParts, "\n")
|
|
}
|
|
}
|
|
return role, cStr
|
|
}
|
|
|
|
func (s *KimiService) Chat(w http.ResponseWriter, r *http.Request, req ChatCompletionRequest) error {
|
|
modelName := req.Model
|
|
if modelName == "" {
|
|
modelName = s.defaultModel
|
|
}
|
|
|
|
if len(req.Tools) == 0 && len(req.Functions) > 0 {
|
|
for _, fn := range req.Functions {
|
|
req.Tools = append(req.Tools, Tool{Type: "function", Function: fn})
|
|
}
|
|
if req.ToolChoice == nil && req.FunctionCall != nil {
|
|
req.ToolChoice = req.FunctionCall
|
|
}
|
|
}
|
|
|
|
var nonSys []ChatMessage
|
|
var sysParts []string
|
|
for _, m := range req.Messages {
|
|
if m.Role == "system" {
|
|
if c := m.GetContentString(); c != "" {
|
|
sysParts = append(sysParts, c)
|
|
}
|
|
} else {
|
|
nonSys = append(nonSys, m)
|
|
}
|
|
}
|
|
userSysPrompt := strings.Join(sysParts, "\n\n")
|
|
|
|
toolsPrompt := ""
|
|
if len(req.Tools) > 0 {
|
|
toolsPrompt = BuildToolsSystemPrompt(req.Tools, req.ToolChoice)
|
|
}
|
|
|
|
var history []map[string]interface{}
|
|
var promptText string
|
|
|
|
if len(nonSys) == 0 {
|
|
promptText = "Hello"
|
|
if toolsPrompt != "" {
|
|
promptText = toolsPrompt + "\n\n" + promptText
|
|
}
|
|
if userSysPrompt != "" {
|
|
promptText = userSysPrompt + "\n\n" + promptText
|
|
}
|
|
} else {
|
|
firstTrailingTool := -1
|
|
for i := len(nonSys) - 1; i >= 0; i-- {
|
|
if nonSys[i].Role == "tool" || nonSys[i].Role == "function" {
|
|
firstTrailingTool = i
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
var historyMsgs []ChatMessage
|
|
var currentTurnMsgs []ChatMessage
|
|
|
|
if firstTrailingTool != -1 {
|
|
historyMsgs = nonSys[:firstTrailingTool]
|
|
currentTurnMsgs = nonSys[firstTrailingTool:]
|
|
} else {
|
|
historyMsgs = nonSys[:len(nonSys)-1]
|
|
currentTurnMsgs = nonSys[len(nonSys)-1:]
|
|
}
|
|
|
|
for i, m := range historyMsgs {
|
|
role, cStr := formatMessageContent(m)
|
|
if i == 0 && role == "user" {
|
|
if userSysPrompt != "" {
|
|
cStr = userSysPrompt + "\n\n" + cStr
|
|
}
|
|
if toolsPrompt != "" {
|
|
cStr = toolsPrompt + "\n\n" + cStr
|
|
}
|
|
}
|
|
history = append(history, map[string]interface{}{
|
|
"role": role, "metadata": nil,
|
|
"content": []map[string]interface{}{{"text": cStr, "type": "text"}},
|
|
"options": nil,
|
|
})
|
|
}
|
|
|
|
if firstTrailingTool != -1 {
|
|
var toolResParts []string
|
|
for _, m := range currentTurnMsgs {
|
|
_, c := formatMessageContent(m)
|
|
toolResParts = append(toolResParts, c)
|
|
}
|
|
promptText = strings.Join(toolResParts, "\n\n") + "\n\nBased on the tool results above, continue your task. If you need to invoke another tool, respond with <tool_call>{\"name\": \"tool_name\", \"arguments\": {...}}</tool_call>. Otherwise, provide your final response."
|
|
} else {
|
|
_, promptText = formatMessageContent(currentTurnMsgs[0])
|
|
}
|
|
|
|
if toolsPrompt != "" {
|
|
promptText = toolsPrompt + "\n\n" + promptText
|
|
}
|
|
if len(history) == 0 && userSysPrompt != "" {
|
|
promptText = userSysPrompt + "\n\n" + promptText
|
|
}
|
|
}
|
|
|
|
gradioPayload, _ := json.Marshal(map[string]interface{}{
|
|
"data": []interface{}{
|
|
map[string]interface{}{"text": promptText, "files": []interface{}{}},
|
|
history,
|
|
s.ResolveBackend(modelName),
|
|
ResolveMaxTokens(req, s.defaultTokens),
|
|
ResolveTemperature(req, s.defaultTemp),
|
|
ResolveReasoningEffort(req, s.defaultReason),
|
|
},
|
|
})
|
|
|
|
effUA := EffectiveUserAgent(r)
|
|
resp, err := DoWithFibonacciRetry(s.client, func() (*http.Request, error) {
|
|
req, err := http.NewRequest("POST", s.endpoint+"/gradio_api/call/on_submit", bytes.NewBuffer(gradioPayload))
|
|
if err == nil {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("User-Agent", effUA)
|
|
}
|
|
return req, err
|
|
}, 5)
|
|
if err != nil {
|
|
return fmt.Errorf("upstream 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")
|
|
}
|
|
|
|
streamResp, err := DoWithFibonacciRetry(s.client, func() (*http.Request, error) {
|
|
req, err := http.NewRequest("GET", fmt.Sprintf("%s/gradio_api/call/on_submit/%s", s.endpoint, joinRes.EventID), nil)
|
|
if err == nil {
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
req.Header.Set("User-Agent", effUA)
|
|
}
|
|
return req, err
|
|
}, 5)
|
|
if err != nil {
|
|
return fmt.Errorf("upstream stream error: %w", err)
|
|
}
|
|
defer streamResp.Body.Close()
|
|
|
|
compID := "chatcmpl-" + GenerateUUID()
|
|
created := time.Now().Unix()
|
|
|
|
if !req.Stream {
|
|
reader := bufio.NewReader(streamResp.Body)
|
|
var finalContent, finalReasoning string
|
|
for {
|
|
line, err := reader.ReadString('\n')
|
|
if err != nil {
|
|
break
|
|
}
|
|
line = strings.TrimSpace(line)
|
|
if strings.HasPrefix(line, "data: ") {
|
|
if c, r, ok := parseGradioSnapshot(strings.TrimPrefix(line, "data: ")); ok {
|
|
if c != "" {
|
|
finalContent = c
|
|
}
|
|
if r != "" {
|
|
finalReasoning = r
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
toolCalls, cleanedContent, hasTools := DetectToolCalls(finalContent)
|
|
if !hasTools && len(req.Tools) > 0 {
|
|
if rCalls, _, rHas := DetectToolCalls(finalReasoning); rHas {
|
|
toolCalls = rCalls
|
|
hasTools = true
|
|
}
|
|
}
|
|
if !hasTools && len(req.Tools) > 0 {
|
|
if directTCs := parseRawToolPayload(finalContent); len(directTCs) > 0 {
|
|
toolCalls = directTCs
|
|
hasTools = true
|
|
cleanedContent = ""
|
|
}
|
|
}
|
|
|
|
finishReason := "stop"
|
|
var msgContent interface{} = finalContent
|
|
if hasTools {
|
|
finishReason = "tool_calls"
|
|
if cleanedContent == "" {
|
|
msgContent = nil
|
|
} else {
|
|
msgContent = cleanedContent
|
|
}
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
return json.NewEncoder(w).Encode(ChatCompletionResponse{
|
|
ID: compID, Object: "chat.completion", Created: created, Model: modelName,
|
|
Choices: []ChatCompletionResponseChoice{{
|
|
Index: 0,
|
|
Message: ChatMessage{
|
|
Role: "assistant",
|
|
Content: msgContent,
|
|
ReasoningContent: finalReasoning,
|
|
ToolCalls: toolCalls,
|
|
},
|
|
FinishReason: finishReason,
|
|
}},
|
|
})
|
|
}
|
|
|
|
flusher, _ := w.(http.Flusher)
|
|
streamer := NewStreamer(w, flusher, compID, created, modelName)
|
|
streamer.SendChunk(StreamDelta{Role: "assistant"})
|
|
|
|
interceptor := NewStreamToolInterceptor(streamer)
|
|
reader := bufio.NewReader(streamResp.Body)
|
|
var emittedContent, emittedReasoning string
|
|
|
|
for {
|
|
line, err := reader.ReadString('\n')
|
|
if err != nil {
|
|
break
|
|
}
|
|
line = strings.TrimSpace(line)
|
|
if strings.HasPrefix(line, "data: ") {
|
|
if c, r, ok := parseGradioSnapshot(strings.TrimPrefix(line, "data: ")); ok {
|
|
if len(r) > len(emittedReasoning) {
|
|
streamer.SendChunk(StreamDelta{ReasoningContent: r[len(emittedReasoning):]})
|
|
emittedReasoning = r
|
|
}
|
|
if len(c) > len(emittedContent) {
|
|
interceptor.ProcessContentDelta(c[len(emittedContent):])
|
|
emittedContent = c
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
interceptor.FlushRemaining()
|
|
if !interceptor.HasToolCalls() && len(req.Tools) > 0 {
|
|
if rCalls, _, rHas := DetectToolCalls(emittedReasoning); rHas {
|
|
for _, tc := range rCalls {
|
|
idx := interceptor.toolCallCount
|
|
interceptor.toolCallCount++
|
|
tc.Index = &idx
|
|
streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
|
|
}
|
|
}
|
|
}
|
|
if !interceptor.HasToolCalls() && len(req.Tools) > 0 {
|
|
if directTCs := parseRawToolPayload(emittedContent); len(directTCs) > 0 {
|
|
for _, tc := range directTCs {
|
|
idx := interceptor.toolCallCount
|
|
interceptor.toolCallCount++
|
|
tc.Index = &idx
|
|
streamer.SendChunk(StreamDelta{ToolCalls: []ToolCall{tc}})
|
|
}
|
|
}
|
|
}
|
|
|
|
if interceptor.HasToolCalls() {
|
|
streamer.SendChunk(StreamDelta{}, "tool_calls")
|
|
} else {
|
|
streamer.SendChunk(StreamDelta{}, "stop")
|
|
}
|
|
streamer.Done()
|
|
return nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// HTTP Handlers & Main
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func NewMux(service *KimiService) *http.ServeMux {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
|
EnableCORS(w)
|
|
if r.Method == http.MethodOptions {
|
|
w.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]interface{}{
|
|
"gateway": "k3r053n3",
|
|
"status": "online",
|
|
"model": service.defaultModel,
|
|
"backend": service.defaultBackend,
|
|
"reasoning_effort": service.defaultReason,
|
|
"socks5_proxy": service.socks5Proxy,
|
|
"endpoint": service.endpoint,
|
|
"routes": []string{"GET /models", "GET /v1/models", "POST /chat/completions", "POST /v1/chat/completions"},
|
|
})
|
|
})
|
|
modelsH := func(w http.ResponseWriter, r *http.Request) {
|
|
EnableCORS(w)
|
|
if r.Method == http.MethodOptions {
|
|
w.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(ModelsResponse{Object: "list", Data: service.ListModels()})
|
|
}
|
|
mux.HandleFunc("/models", modelsH)
|
|
mux.HandleFunc("/v1/models", modelsH)
|
|
|
|
chatH := func(w http.ResponseWriter, r *http.Request) {
|
|
EnableCORS(w)
|
|
if r.Method == http.MethodOptions {
|
|
w.WriteHeader(http.StatusOK)
|
|
return
|
|
}
|
|
if r.Method != http.MethodPost {
|
|
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, "Invalid JSON payload: "+err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if err := service.Chat(w, r, req); err != nil {
|
|
if w.Header().Get("Content-Type") != "text/event-stream" {
|
|
http.Error(w, "Upstream error: "+err.Error(), http.StatusBadGateway)
|
|
}
|
|
}
|
|
}
|
|
mux.HandleFunc("/chat/completions", chatH)
|
|
mux.HandleFunc("/v1/chat/completions", chatH)
|
|
return mux
|
|
}
|
|
|
|
func resolveProxyEnv() string {
|
|
for _, k := range []string{"ALL_PROXY", "all_proxy", "SOCKS5_PROXY", "socks5_proxy", "SOCKS_PROXY", "socks_proxy"} {
|
|
if v := strings.TrimSpace(os.Getenv(k)); v != "" {
|
|
return v
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func main() {
|
|
port := flag.String("port", "8080", "Port to listen on")
|
|
endpoint := flag.String("endpoint", "https://cw-105-kimi-k3-gguf-demo.hf.space", "Root URL of the Kimi K3 Gradio space")
|
|
model := flag.String("model", "kimi-k3", "Exposed default model name")
|
|
backend := flag.String("backend", "direct:together", "Default backend (direct:together, direct:fireworks, hf:fireworks-ai, hf:together, hf:featherless-ai, hf:baseten)")
|
|
reasoning := flag.String("reasoning", "max", "Default reasoning effort (max, high, low, default)")
|
|
tokens := flag.Int("max-tokens", 8192, "Default max completion tokens (256-8192)")
|
|
temp := flag.Float64("temperature", 0.7, "Default temperature (0.0-1.5)")
|
|
socks := flag.String("socks", "", "SOCKS5 proxy URL (e.g. socks5://127.0.0.1:1080 or socks5://user:pass@host:port)")
|
|
proxy := flag.String("proxy", "", "Alias for -socks")
|
|
socks5 := flag.String("socks5", "", "Alias for -socks")
|
|
ua := flag.String("user-agent", DefaultUserAgent, "Custom User-Agent header")
|
|
uaShort := flag.String("ua", "", "Alias for -user-agent")
|
|
flag.Parse()
|
|
|
|
ConfiguredUserAgent = *ua
|
|
if *uaShort != "" {
|
|
ConfiguredUserAgent = *uaShort
|
|
}
|
|
socksProxy := *socks
|
|
if socksProxy == "" {
|
|
socksProxy = *proxy
|
|
}
|
|
if socksProxy == "" {
|
|
socksProxy = *socks5
|
|
}
|
|
if socksProxy == "" {
|
|
socksProxy = resolveProxyEnv()
|
|
}
|
|
|
|
svc := NewKimiService(*endpoint, *model, *backend, *reasoning, socksProxy, *tokens, *temp)
|
|
fmt.Printf("k3r053n3 starting on port %s...\nTarget Endpoint: %s\nModel Name: %s\nBackend: %s\nReasoning Effort: %s\nMax Tokens: %d\nTemperature: %.2f\n", *port, svc.endpoint, svc.defaultModel, svc.defaultBackend, svc.defaultReason, svc.defaultTokens, svc.defaultTemp)
|
|
if svc.socks5Proxy != "" {
|
|
fmt.Printf("SOCKS5 Proxy: %s\n", svc.socks5Proxy)
|
|
}
|
|
fmt.Printf("Endpoints:\n GET http://localhost:%s/v1/models\n POST http://localhost:%s/v1/chat/completions\n", *port, *port)
|
|
|
|
if err := http.ListenAndServe(":"+*port, NewMux(svc)); err != nil {
|
|
fmt.Fprintf(os.Stderr, "server failed: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|