Skip to content

Commit 459be86

Browse files
sjy3sjy3
committed
feat: add ark api key fetch by IAM role
1 parent b4a9778 commit 459be86

File tree

4 files changed

+174
-1
lines changed

4 files changed

+174
-1
lines changed

‎agent/llmagent/agent.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ package llmagent
1717
import (
1818
"context"
1919

20+
"github.com/volcengine/veadk-go/auth/veauth"
2021
"github.com/volcengine/veadk-go/common"
2122
"github.com/volcengine/veadk-go/configs"
2223
"github.com/volcengine/veadk-go/knowledgebase"
@@ -55,7 +56,7 @@ func New(cfg *Config) (agent.Agent, error) {
5556
cfg.ModelName = utils.GetEnvWithDefault(common.MODEL_AGENT_NAME, configs.GetGlobalConfig().Model.Agent.Name, common.DEFAULT_MODEL_AGENT_NAME)
5657
}
5758
if cfg.ModelAPIKey == "" {
58-
cfg.ModelAPIKey = utils.GetEnvWithDefault(common.MODEL_AGENT_API_KEY, configs.GetGlobalConfig().Model.Agent.ApiKey)
59+
cfg.ModelAPIKey = utils.GetEnvWithDefault(common.MODEL_AGENT_API_KEY, configs.GetGlobalConfig().Model.Agent.ApiKey, utils.Must(veauth.GetArkToken(common.DEFAULT_MODEL_REGION)))
5960
}
6061
if cfg.ModelAPIBase == "" {
6162
cfg.ModelAPIBase = utils.GetEnvWithDefault(common.MODEL_AGENT_API_BASE, configs.GetGlobalConfig().Model.Agent.ApiBase, common.DEFAULT_MODEL_AGENT_API_BASE)

‎auth/veauth/ark_veauth.go‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package veauth
16+
17+
import (
18+
"encoding/json"
19+
"fmt"
20+
"log"
21+
"net/http"
22+
23+
"github.com/volcengine/veadk-go/common"
24+
"github.com/volcengine/veadk-go/integrations/ve_sign"
25+
"github.com/volcengine/veadk-go/utils"
26+
)
27+
28+
type GetRawApiKeyResponse struct {
29+
Result struct {
30+
ApiKey string `json:"ApiKey"`
31+
} `json:"Result"`
32+
}
33+
34+
type ListApiKeysResponse struct {
35+
Result struct {
36+
Items []struct {
37+
ID int `json:"Id"`
38+
Name string `json:"Name"`
39+
} `json:"Items"`
40+
} `json:"Result"`
41+
}
42+
43+
func GetArkToken(region string) (string, error) {
44+
if region == "" {
45+
region = common.DEFAULT_MODEL_REGION
46+
}
47+
log.Println("Fetching ARK token...")
48+
49+
accessKey := utils.GetEnvWithDefault(common.VOLCENGINE_ACCESS_KEY)
50+
secretKey := utils.GetEnvWithDefault(common.VOLCENGINE_SECRET_KEY)
51+
sessionToken := ""
52+
53+
if accessKey == "" || secretKey == "" {
54+
// try to get from vefaas iam
55+
cred, err := GetCredentialFromVeFaaSIAM()
56+
if err != nil {
57+
// If we can't get credentials from env or IAM, we can't proceed.
58+
// Here we return error immediately if IAM fails and no env vars.
59+
return "", fmt.Errorf("failed to get credential from vefaas iam: %w", err)
60+
}
61+
accessKey = cred.AccessKeyID
62+
secretKey = cred.SecretAccessKey
63+
sessionToken = cred.SessionToken
64+
}
65+
66+
header := make(map[string]string)
67+
if sessionToken != "" {
68+
header["X-Security-Token"] = sessionToken
69+
}
70+
71+
// ListApiKeys
72+
req1 := ve_sign.VeRequest{
73+
AK: accessKey,
74+
SK: secretKey,
75+
Method: "POST",
76+
Scheme: "https",
77+
Host: "open.volcengineapi.com",
78+
Path: "/",
79+
Service: "ark",
80+
Region: region,
81+
Action: "ListApiKeys",
82+
Version: "2024-01-01",
83+
Header: header,
84+
Body: map[string]interface{}{
85+
"ProjectName": "default",
86+
"Filter": map[string]interface{}{"AllowAll": true},
87+
},
88+
}
89+
90+
resp1Body, err := req1.DoRequest()
91+
if err != nil {
92+
return "", fmt.Errorf("failed to list api keys: %w", err)
93+
}
94+
var listResp ListApiKeysResponse
95+
if err := json.Unmarshal(resp1Body, &listResp); err != nil {
96+
return "", fmt.Errorf("failed to unmarshal list api keys response: %w", err)
97+
}
98+
99+
if len(listResp.Result.Items) == 0 {
100+
return "", fmt.Errorf("failed to get ARK api key list: empty items")
101+
}
102+
103+
firstApiKeyId := listResp.Result.Items[0].ID
104+
log.Println("By default, VeADK fetches the first API Key in the list.")
105+
log.Printf("Try to fetch ARK API Key with id=%d, name=%s\n", firstApiKeyId, listResp.Result.Items[0].Name)
106+
107+
// GetRawApiKey
108+
req2 := ve_sign.VeRequest{
109+
AK: accessKey,
110+
SK: secretKey,
111+
Method: http.MethodPost,
112+
Scheme: "https",
113+
Host: "open.volcengineapi.com",
114+
Path: "/",
115+
Service: "ark",
116+
Region: region,
117+
Action: "GetRawApiKey",
118+
Version: "2024-01-01",
119+
Header: header,
120+
Body: map[string]interface{}{
121+
"Id": firstApiKeyId,
122+
},
123+
}
124+
125+
resp2Body, err := req2.DoRequest()
126+
if err != nil {
127+
return "", fmt.Errorf("failed to get raw api key: %w", err)
128+
}
129+
130+
var getResp GetRawApiKeyResponse
131+
if err := json.Unmarshal(resp2Body, &getResp); err != nil {
132+
return "", fmt.Errorf("failed to unmarshal get raw api key response: %w", err)
133+
}
134+
135+
apiKey := getResp.Result.ApiKey
136+
if apiKey == "" {
137+
return "", fmt.Errorf("failed to get ARK api key: key not found in response")
138+
}
139+
140+
if len(apiKey) >= 8 {
141+
log.Printf("Successfully fetched ARK API Key (starts with %s).\n", apiKey[:8])
142+
} else {
143+
log.Println("Successfully fetched ARK API Key.")
144+
}
145+
146+
return apiKey, nil
147+
}

‎common/defaults.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ package common
1616

1717
const (
1818
// Agent
19+
DEFAULT_MODEL_REGION = "cn-beijing"
1920
DEFAULT_MODEL_AGENT_NAME = "doubao-seed-1-6-250615"
2021
DEFAULT_MODEL_AGENT_PROVIDER = "openai"
2122
DEFAULT_MODEL_AGENT_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"

‎utils/must.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package utils
16+
17+
import "log"
18+
19+
func Must[T any](obj T, err error) T {
20+
if err != nil {
21+
log.Printf("Must internal err:%v", err)
22+
}
23+
return obj
24+
}

0 commit comments

Comments
 (0)