ソースを参照

feat: vertexai support apikey usertype (#389)

* feat: vertexai support apikey usertype

* fix: ci lint

* feat: vertexai apikey support project id

* fix: ci lint
zijiren 2 ヶ月 前
コミット
37a48f9e09
2 ファイル変更70 行追加7 行削除
  1. 40 3
      core/relay/adaptor/vertexai/adaptor.go
  2. 30 4
      core/relay/adaptor/vertexai/key.go

+ 40 - 3
core/relay/adaptor/vertexai/adaptor.go

@@ -28,6 +28,7 @@ func (a *Adaptor) SupportMode(m mode.Mode) bool {
 
 type Config struct {
 	Region    string
+	Key       string
 	ProjectID string
 	ADCJSON   string
 }
@@ -68,7 +69,7 @@ func (a *Adaptor) Metadata() adaptor.Metadata {
 		Features: []string{
 			"Claude support native Endpoint: /v1/messages",
 		},
-		KeyHelp: "region|adcJSON",
+		KeyHelp: "region|adcJSON or region|apikey or region|project_id|apikey",
 		Models:  modelList,
 	}
 }
@@ -95,6 +96,18 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta, _ adaptor.Store) (adaptor.Reque
 	}
 
 	if meta.Channel.BaseURL != "" {
+		if config.ProjectID == "" || config.Region == "" {
+			return adaptor.RequestURL{
+				Method: http.MethodPost,
+				URL: fmt.Sprintf(
+					"%s/v1/publishers/google/models/%s:%s",
+					meta.Channel.BaseURL,
+					meta.ActualModel,
+					suffix,
+				),
+			}, nil
+		}
+
 		return adaptor.RequestURL{
 			Method: http.MethodPost,
 			URL: fmt.Sprintf(
@@ -108,11 +121,30 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta, _ adaptor.Store) (adaptor.Reque
 		}, nil
 	}
 
+	var requestDoamin string
+	if config.Region == "" {
+		requestDoamin = "aiplatform.googleapis.com"
+	} else {
+		requestDoamin = config.Region + "-aiplatform.googleapis.com"
+	}
+
+	if config.ProjectID == "" || config.Region == "" {
+		return adaptor.RequestURL{
+			Method: http.MethodPost,
+			URL: fmt.Sprintf(
+				"https://%s/v1/publishers/google/models/%s:%s",
+				requestDoamin,
+				meta.ActualModel,
+				suffix,
+			),
+		}, nil
+	}
+
 	return adaptor.RequestURL{
 		Method: http.MethodPost,
 		URL: fmt.Sprintf(
-			"https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
-			config.Region,
+			"https://%s/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
+			requestDoamin,
 			config.ProjectID,
 			config.Region,
 			meta.ActualModel,
@@ -132,6 +164,11 @@ func (a *Adaptor) SetupRequestHeader(
 		return err
 	}
 
+	if config.Key != "" {
+		req.Header.Set("X-Goog-Api-Key", config.Key)
+		return nil
+	}
+
 	token, err := getToken(context.Background(), config.ADCJSON)
 	if err != nil {
 		return err

+ 30 - 4
core/relay/adaptor/vertexai/key.go

@@ -19,14 +19,40 @@ func (a *Adaptor) ValidateKey(key string) error {
 	return nil
 }
 
-// region|adcJSON
+// region|adcJSON or region|apikey or region|project_id|apikey
 func getConfigFromKey(key string) (Config, error) {
-	region, adcJSON, ok := strings.Cut(key, "|")
+	region, gkey, ok := strings.Cut(key, "|")
 	if !ok {
 		return Config{}, errors.New("invalid key format")
 	}
 
-	node, err := sonic.GetFromString(adcJSON, "project_id")
+	if region == gkey {
+		region = ""
+	}
+
+	if !strings.HasPrefix(gkey, "{") {
+		projectid, ngkey, ok := strings.Cut(gkey, "|")
+		if ok {
+			// region|project_id|apikey
+			if projectid == ngkey {
+				projectid = ""
+			}
+
+			return Config{
+				Region:    region,
+				Key:       ngkey,
+				ProjectID: projectid,
+			}, nil
+		}
+		// region|apikey
+		return Config{
+			Region: region,
+			Key:    gkey,
+		}, nil
+	}
+
+	// region|adcJSON
+	node, err := sonic.GetFromString(gkey, "project_id")
 	if err != nil {
 		return Config{}, err
 	}
@@ -39,6 +65,6 @@ func getConfigFromKey(key string) (Config, error) {
 	return Config{
 		Region:    region,
 		ProjectID: projectID,
-		ADCJSON:   adcJSON,
+		ADCJSON:   gkey,
 	}, nil
 }