Support change api endpoint (#17)

* Support change api endpoint * chatgpt: Fix NewWithBaseURL
sunshineplan · Mar 14, 2024 · 3d54273 · 3d54273
1 parent cb0d560
commit 3d54273
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 6 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -23,8 +23,27 @@ jobs:
           go-version: stable
 
       - name: Test Code
-        env: 
-          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
         run: |
           go build -v ./...
           go test -v -race ./...
+
+  ai:
+    if: ${{ github.actor != 'dependabot[bot]' }}
+    runs-on: ubuntu-latest
+    needs: test
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: stable
+
+      - name: Test Code
+        env: 
+          CHATGPT_API_KEY: ${{ secrets.CHATGPT_API_KEY }}
+          CHATGPT_BASE_URL: ${{ secrets.CHATGPT_BASE_URL }}
+          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
+          GEMINI_ENDPOINT: ${{ secrets.GEMINI_ENDPOINT }}
+        run: go test -v
diff --git a/ai_test.go b/ai_test.go
@@ -97,7 +97,7 @@ func TestGemini(t *testing.T) {
 	if apiKey == "" {
 		return
 	}
-	gemini, err := gemini.New(apiKey)
+	gemini, err := gemini.NewWithEndpoint(apiKey, os.Getenv("GEMINI_ENDPOINT"))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -118,7 +118,7 @@ func TestChatGPT(t *testing.T) {
 	if apiKey == "" {
 		return
 	}
-	chatgpt := chatgpt.New(apiKey)
+	chatgpt := chatgpt.NewWithBaseURL(apiKey, os.Getenv("CHATGPT_BASE_URL"))
 	defer chatgpt.Close()
 	if err := testChat(chatgpt, "Who are you?"); err != nil {
 		t.Error(err)

diff --git a/chatgpt/chatgpt.go b/chatgpt/chatgpt.go
@@ -26,7 +26,15 @@ type ChatGPT struct {
 }
 
 func New(authToken string) ai.AI {
-	return NewWithClient(openai.NewClient(authToken))
+	return NewWithBaseURL(authToken, "")
+}
+
+func NewWithBaseURL(authToken, baseURL string) ai.AI {
+	cfg := openai.DefaultConfig(authToken)
+	if baseURL != "" {
+		cfg.BaseURL = baseURL
+	}
+	return NewWithClient(openai.NewClientWithConfig(cfg))
 }
 
 func NewWithClient(client *openai.Client) ai.AI {

diff --git a/gemini/gemini.go b/gemini/gemini.go
@@ -27,7 +27,15 @@ type Gemini struct {
 }
 
 func New(apiKey string) (ai.AI, error) {
-	client, err := genai.NewClient(context.Background(), option.WithAPIKey(apiKey))
+	return NewWithEndpoint(apiKey, "")
+}
+
+func NewWithEndpoint(apiKey, endpoint string) (ai.AI, error) {
+	opts := []option.ClientOption{option.WithAPIKey(apiKey)}
+	if endpoint != "" {
+		opts = append(opts, option.WithEndpoint(endpoint))
+	}
+	client, err := genai.NewClient(context.Background(), opts...)
 	if err != nil {
 		return nil, err
 	}