Limour

package main

import (
    "io"
    "io/ioutil"
    "log"
    "net/http"
    "net/url"
    "strings"
    "bytes"
    "encoding/json"
    "time"
    "sync"
)

var (
    target    = "https://api.openai.com" // Target domain
    mu sync.Mutex
    count int
)

func main() {
    http.HandleFunc("/", handleRequest)
    http.ListenAndServe(":9000", nil)
}

func get1Key(key string) string {
    mu.Lock()
    defer mu.Unlock()

    arr := strings.Split(key, "|")
    randomIndex := count % len(arr)
    count++
    if count > 999999 {
        count = 0
    }
    randomSubstr := arr[randomIndex]
    log.Println("Authorization", randomSubstr)
    return randomSubstr
}

// Get a json decoder for a given requests body
func requestBodyDecoder(request *http.Request) *json.Decoder {
    // Read body to buffer
    body, err := ioutil.ReadAll(request.Body)
    if err != nil {
        log.Printf("Error reading body: %v", err)
        panic(err)
    }

    // Because go lang is a pain in the ass if you read the body then any subsequent calls
    // are unable to read the body again....
    // request.Body = ioutil.NopCloser(bytes.NewBuffer(body))

    return json.NewDecoder(ioutil.NopCloser(bytes.NewBuffer(body)))
}

// Call setResponseHeader function where response headers need to be set
func setResponseHeader(w http.ResponseWriter) {
    w.Header().Set("Access-Control-Allow-Methods", "GET,HEAD,POST,OPTIONS")
    w.Header().Set("Access-Control-Allow-Origin", "*")
    w.Header().Set("Access-Control-Max-Age", "86400")
    w.Header().Set("Access-Control-Allow-Headers", "authorization,content-type")
}

func handleBody(r *http.Request) io.ReadCloser {
    // Read request body
    decoder := requestBodyDecoder(r)
    // Parse JSON data
    var requestData map[string]interface{}
    err := decoder.Decode(&requestData)
    if err != nil {
        // Handle JSON data parsing exception
        log.Printf("Error decoding body: %v", err)
        panic(err)
    }
    // Get the "messages" list
    messages, ok := requestData["messages"].([]interface{})
    if !ok {
        // "messages" field does not exist or type is incorrect
        log.Printf("Error reading messages: %v", ok)
        panic(ok)
    }
    // Modify the "messages" list as needed
    // log.Println("debug 5", len(messages), len(messages) > 4)
    if len(messages) > 4 {
        firstMessage, ok := messages[0].(map[string]interface{})
        if !ok {
            // First message type is incorrect
            log.Printf("Error reading firstMessage: %v", ok)
            panic(ok)
        }

        role, roleOk := firstMessage["role"].(string)
        // log.Println("debug 6", roleOk, role, role == "system", strings.EqualFold(role, "system"))
        if roleOk && strings.EqualFold(role, "system") {
            // Move the first message to the third-to-last position
            thirdToLastIndex := len(messages) - 3
            messages_copy := make([]interface{}, 0)
            messages_copy = append(messages_copy, messages[0])
            messages_copy = append(messages_copy, messages[thirdToLastIndex:]...)
            messages = append(messages[1:thirdToLastIndex], messages_copy...)
            
            // Update the "messages" list
            requestData["messages"] = messages
            log.Println("move system role to ", thirdToLastIndex)
        }
    }
    // Encode the updated data as JSON
    var buf bytes.Buffer
    encoder := json.NewEncoder(&buf)
    err = encoder.Encode(requestData)
    if err != nil {
        // Handle JSON data encoding exception
        log.Printf("Error encoding body: %v", err)
        panic(err)
    }

    return ioutil.NopCloser(&buf)
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
    // Filter invalid URLs
    _, err := url.Parse(r.URL.String())
    if err != nil {
        log.Println("Error parsing URL: ", err.Error())
        http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
        return
    }

    // Remove environment prefix (for Tencent Cloud, if included, currently I only use test and release)
    newPath := strings.Replace(r.URL.Path, "/release", "", 1)
    newPath = strings.Replace(newPath, "/test", "", 1)

    // Concatenate target URL
    targetURL := target + newPath

    // Create proxy HTTP request
    var proxyReq *http.Request
    // log.Println("debug 1", targetURL, r.Method)
    // log.Println("debug 2", strings.Contains(targetURL, "chat/completions"))
    // log.Println("debug 3", strings.EqualFold(r.Method, "POST"))
    if strings.Contains(targetURL, "chat/completions") && strings.EqualFold(r.Method, "POST")  {
        // log.Println("debug 4-0")
        proxyReq, err = http.NewRequest(r.Method, targetURL, handleBody(r))
    } else {
        // log.Println("debug 4-1")
        proxyReq, err = http.NewRequest(r.Method, targetURL, r.Body)
    }

    if err != nil {
        log.Println("Error creating proxy request: ", err.Error())
        http.Error(w, "Error creating proxy request", http.StatusInternalServerError)
        return
    }

    // Copy original request headers to new request
    keys := strings.Split(r.Header.Get("Authorization"), " ")
    if len(keys) == 2 {
           r.Header.Set("Authorization", "Bearer " + get1Key(keys[1]))
    }
    if _, ok := r.Header["User-Agent"]; !ok {
        // explicitly disable User-Agent so it's not set to default value
        r.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 6.1; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/88.0.4324.96 Safari/537.36")
    }
    proxyReq.Header = http.Header{
        "Content-Type":  []string{"application/json"},
        "Authorization": []string{r.Header.Get("Authorization")},
        "User-Agent":    []string{r.Header.Get("User-Agent")},
    }

    
    // Default timeout set to 60s
    client := &http.Client{
        Timeout: 60 * time.Second,
    }

    // Send proxy request to OpenAI
    resp, err := client.Do(proxyReq)
    if err != nil {
        log.Println("Error sending proxy request: ", err.Error())
        http.Error(w, err.Error(), http.StatusInternalServerError)
        return
    }
    defer resp.Body.Close()

    // Set response headers
    setResponseHeader(w)
    
    // Set response status code to original response status code
    w.WriteHeader(resp.StatusCode)

    // Write response body to response stream (supports streaming response)
    buf := make([]byte, 1024)
    for {
        if n, err := resp.Body.Read(buf); err == io.EOF || n == 0 {
            return
        } else if err != nil {
            log.Println("error while reading respbody: ", err.Error())
            http.Error(w, err.Error(), http.StatusInternalServerError)
            return
        } else {
            if _, err = w.Write(buf[:n]); err != nil {
                log.Println("error while writing resp: ", err.Error())
                http.Error(w, err.Error(), http.StatusInternalServerError)
                return
            }
            w.(http.Flusher).Flush()
        }
    }
}

CC=musl-gcc /home/jovyan/go/bin/go1.20.1 build -tags musl -o openai -trimpath -ldflags '-linkmode "external" -extldflags "-static" -s -w -buildid=' ./openai-proxy.go

mkdir -p ~/app/apio && cd ~/app/apio && nano docker-compose.yml
chmod 777 openai
sudo docker-compose up -d && sudo docker-compose logs

version: '3.3'
services:
    apio:
        restart: always
        image: alpine:latest
        volumes:
          - ./openai:/bin/openai
        entrypoint: ["/bin/openai"]
 
networks:
  default:
    external: true
    name: ngpm