From f8b364a052f042f2bbf87abeed4a223f32d33ae7 Mon Sep 17 00:00:00 2001
From: Sakurasan <26715255+Sakurasan@users.noreply.github.com>
Date: Wed, 18 Dec 2024 01:00:45 +0800
Subject: [PATCH] fix: custom_endpoint

---
 pkg/openai/chat.go | 20 ++++++++++----------
 router/router.go   | 11 -----------
 2 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/pkg/openai/chat.go b/pkg/openai/chat.go
index 78f0f10..cd87ebb 100644
--- a/pkg/openai/chat.go
+++ b/pkg/openai/chat.go
@@ -25,13 +25,13 @@ const (
 )
 
 var (
-	BaseURL            string // "https://api.openai.com"
-	AIGateWay_Endpoint = "https://gateway.ai.cloudflare.com/v1/431ba10f11200d544922fbca177aaa7f/openai/openai/chat/completions"
+	Custom_Endpoint    string
+	AIGateWay_Endpoint string // "https://gateway.ai.cloudflare.com/v1/431ba10f11200d544922fbca177aaa7f/openai/openai/chat/completions"
 )
 
 func init() {
 	if os.Getenv("OpenAI_Endpoint") != "" {
-		BaseURL = os.Getenv("OpenAI_Endpoint")
+		Custom_Endpoint = os.Getenv("OpenAI_Endpoint")
 	}
 	if os.Getenv("AIGateWay_Endpoint") != "" {
 		AIGateWay_Endpoint = os.Getenv("AIGateWay_Endpoint")
@@ -247,16 +247,16 @@ func ChatProxy(c *gin.Context, chatReq *ChatCompletionRequest) {
 		req.Header = c.Request.Header
 		req.Header.Set("api-key", onekey.Key)
 	default:
-		req, err = http.NewRequest(c.Request.Method, OpenAI_Endpoint, bytes.NewReader(chatReq.ToByteJson()))
-		if onekey.EndPoint != "" { // 优先key的endpoint
-			req, err = http.NewRequest(c.Request.Method, onekey.EndPoint+c.Request.RequestURI, bytes.NewReader(chatReq.ToByteJson()))
-		}
+		req, err = http.NewRequest(c.Request.Method, OpenAI_Endpoint, bytes.NewReader(chatReq.ToByteJson())) // default endpoint
+
 		if AIGateWay_Endpoint != "" { // cloudflare gateway的endpoint
 			req, err = http.NewRequest(c.Request.Method, AIGateWay_Endpoint, bytes.NewReader(chatReq.ToByteJson()))
 		}
-		customEndpoint := os.Getenv("CUSTOM_ENDPOINT") // 最后是用户自定义的endpoint CUSTOM_ENDPOINT=true OpenAI_Endpoint
-		if customEndpoint == "true" && OpenAI_Endpoint != "" {
-			req, err = http.NewRequest(c.Request.Method, BaseURL, bytes.NewReader(chatReq.ToByteJson()))
+		if Custom_Endpoint != "" { // 自定义endpoint
+			req, err = http.NewRequest(c.Request.Method, Custom_Endpoint, bytes.NewReader(chatReq.ToByteJson()))
+		}
+		if onekey.EndPoint != "" { // 优先key的endpoint
+			req, err = http.NewRequest(c.Request.Method, onekey.EndPoint+c.Request.RequestURI, bytes.NewReader(chatReq.ToByteJson()))
 		}
 
 		req.Header = c.Request.Header
diff --git a/router/router.go b/router/router.go
index d0ba2ba..9930b78 100644
--- a/router/router.go
+++ b/router/router.go
@@ -2,22 +2,18 @@ package router
 
 import (
 	"crypto/tls"
-	"fmt"
-	"log"
 	"net"
 	"net/http"
 	"net/http/httputil"
 	"opencatd-open/pkg/claude"
 	oai "opencatd-open/pkg/openai"
 	"opencatd-open/store"
-	"os"
 	"time"
 
 	"github.com/gin-gonic/gin"
 )
 
 var (
-	baseUrl       = "https://api.openai.com"
 	GPT3Dot5Turbo = "gpt-3.5-turbo"
 	GPT4          = "gpt-4"
 )
@@ -62,13 +58,6 @@ var (
 // 	} `json:"usage"`
 // }
 
-func init() {
-	if openai_endpoint := os.Getenv("openai_endpoint"); openai_endpoint != "" {
-		log.Println(fmt.Sprintf("replace %s to %s", baseUrl, openai_endpoint))
-		baseUrl = openai_endpoint
-	}
-}
-
 func HandleProxy(c *gin.Context) {
 	var (
 		localuser bool