gpt proxy

proxy.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"
)

var serverPort = flag.String("port", "8088", "HTTP服务器端口")
var authorizationKey = flag.String("apikey", "", "用于授权的API密钥")
var targetAPI = flag.String("target", "", "目标服务器URL")
var skipWord = flag.String("skip", "q", "跳过的词")

func init() {
	flag.Parse()
}

type RequestBody struct {
	Messages         []Messages `json:"messages"`
	Stream           bool       `json:"stream"`
	Model            string     `json:"model"`
	Temperature      float64    `json:"temperature"`
	PresencePenalty  int        `json:"presence_penalty"`
	FrequencyPenalty int        `json:"frequency_penalty"`
	TopP             int        `json:"top_p"`
}

type Messages struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

func createProxyServer(targetURL *url.URL, apiKey string) *httputil.ReverseProxy {
	proxy := httputil.NewSingleHostReverseProxy(targetURL)
	proxy.Director = func(req *http.Request) {
		req.URL.Scheme = targetURL.Scheme
		req.URL.Host = targetURL.Host
		req.Host = targetURL.Host
		req.URL.Path = targetURL.Path
		if req.Method != http.MethodOptions {
			req.Header.Set("Authorization", "Bearer "+apiKey)
		}

		// 读取原始body
		if req.Body != nil {
			originalBody, err := io.ReadAll(req.Body)
			if err != nil {
				log.Printf("Error reading body: %v", err)
				return
			}

			var bodyData RequestBody
			if err := json.Unmarshal(originalBody, &bodyData); err != nil {
				log.Printf("Error unmarshalling JSON: %v", err)
				return
			}
			bodyData.Model = strings.Replace(bodyData.Model, fmt.Sprintf("(%s)", *skipWord), "", -1)
			modifiedBody, err := json.Marshal(bodyData)
			if err != nil {
				log.Printf("Error marshalling JSON: %v", err)
				return
			}

			req.Body = io.NopCloser(bytes.NewReader(modifiedBody))
			req.ContentLength = int64(len(modifiedBody))
		}
	}
	return proxy
}

func setupCORS(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
}

func handleProxyRequest(proxy *httputil.ReverseProxy) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodOptions {
			setupCORS(w)
			w.WriteHeader(http.StatusNoContent)
			return
		}
		proxy.ServeHTTP(w, r)
	}
}

func main() {
	if *authorizationKey == "" {
		log.Fatal("API密钥未设置，请通过-apikey参数提供。")
	}

	targetURL, err := url.Parse(*targetAPI)
	if err != nil {
		log.Fatalf("解析目标API URL时出错: %v", err)
	}

	proxy := createProxyServer(targetURL, *authorizationKey)

	http.HandleFunc("/v1/chat/completions", handleProxyRequest(proxy))

	log.Printf("服务器正在端口 %s 上运行", *serverPort)
	log.Fatal(http.ListenAndServe(":"+*serverPort, nil))
}