feat: add util function for getting bear token (#776)

* update readme, add some cases * update readme, and add some cases * feat: add util function for getting bear token * feat: add util function for getting bear token * feat: add util function for getting bear token * update defaultConfig --------- Co-authored-by: zhuping04 <[email protected]>
baidubce · Sep 5, 2024 · 0038584 · 0038584
1 parent d4ab537
commit 0038584
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 2 deletions.
diff --git a/javascript/src/constant.ts b/javascript/src/constant.ts
@@ -26,6 +26,7 @@ export const DEFAULT_CONFIG: DefaultConfig = {
     QIANFAN_SECRET_KEY: '',
     QIANFAN_BASE_URL: 'https://aip.baidubce.com',
     QIANFAN_CONSOLE_API_BASE_URL: 'https://qianfan.baidubce.com',
+    QIANFAN_BEAR_TOKEN_URL: 'http://iam.bj.baidubce.com/v1/BCE-BEARER/token',
     QIANFAN_LLM_API_RETRY_TIMEOUT: '600000',
     QIANFAN_LLM_API_RETRY_BACKOFF_FACTOR: '0',
     QIANFAN_LLM_RETRY_MAX_WAIT_INTERVAL: '120000',
@@ -46,4 +47,4 @@ export const RETRY_CODE = [
 
 export const SERVER_LIST_API = '/wenxinworkshop/service/list';
 
-export const DYNAMIC_INVALID = ['reranker'];
+export const DYNAMIC_INVALID = ['reranker'];
diff --git a/javascript/src/index.ts b/javascript/src/index.ts
@@ -18,7 +18,7 @@ import Embedding from './Embedding';
 import Plugin from './Plugin';
 import {Text2Image, Image2Text} from './Images';
 import Reranker from './Reranker';
-import {setEnvVariable, setBrowserVariable, consoleAction} from './utils';
+import {setEnvVariable, setBrowserVariable, consoleAction, getBearToken} from './utils';
 
 export {
     ChatCompletion,
@@ -31,4 +31,5 @@ export {
     setEnvVariable,
     setBrowserVariable,
     consoleAction,
+    getBearToken
 };
diff --git a/javascript/src/interface.ts b/javascript/src/interface.ts
@@ -19,6 +19,7 @@ export interface DefaultConfig {
     QIANFAN_SECRET_KEY: string;
     QIANFAN_BASE_URL: string;
     QIANFAN_CONSOLE_API_BASE_URL: string;
+    QIANFAN_BEAR_TOKEN_URL: string;
     QIANFAN_LLM_API_RETRY_TIMEOUT: string;
     QIANFAN_LLM_API_RETRY_BACKOFF_FACTOR: string;
     QIANFAN_LLM_RETRY_MAX_WAIT_INTERVAL: string;

diff --git a/javascript/src/utils.ts b/javascript/src/utils.ts
@@ -14,6 +14,7 @@
 
 import HttpClient from './HttpClient';
 import Fetch from './Fetch';
+import {Headers} from './Fetch/nodeFetch';
 import {BASE_PATH, DEFAULT_CONFIG, DEFAULT_HEADERS} from './constant';
 import {IAMConfig, QfLLMInfoMap, ReqBody, DefaultConfig} from './interface';
 import * as packageJson from '../package.json';
@@ -356,3 +357,67 @@ export async function consoleAction({base_api_route, data, action}: ConsoleActio
         throw error;
     }
 }
+
+interface GetTokenProps {
+    expireInSeconds?: number;
+}
+
+interface TokenResp {
+    headers?: Headers,
+    userId?: string,
+    status?: string,
+    createTime?: string,
+    token: string,
+    expireTime: string
+}
+
+async function fetchBearToken(props?: GetTokenProps): Promise<TokenResp> {
+    const {expireInSeconds: expireInSecondsInProps} = props || {};
+    const config = getDefaultConfig();
+    const {QIANFAN_BEAR_TOKEN_URL} = config;
+    try {
+        // 鉴权
+        const httpClientConfig = getIAMConfig(
+            config.QIANFAN_ACCESS_KEY,
+            config.QIANFAN_SECRET_KEY,
+            QIANFAN_BEAR_TOKEN_URL
+        );
+        const client = new HttpClient(httpClientConfig);
+        const expireInSeconds = typeof expireInSecondsInProps === 'number' ? expireInSecondsInProps : 100000
+        const fetchOptions = await client.getSignature({
+            httpMethod: 'GET',
+            path: QIANFAN_BEAR_TOKEN_URL,
+            params: {expireInSeconds},
+            headers: {
+                ...DEFAULT_HEADERS,
+            }
+        });
+        const fetchInstance = new Fetch();
+        const {url, ...rest} = fetchOptions;
+        const resp = await fetchInstance.makeRequest(`${QIANFAN_BEAR_TOKEN_URL}?expireInSeconds=${expireInSeconds}`, rest);
+        return resp;
+    }
+    catch (error) {
+        const error_msg = `Failed to get access token: ${error && error.message}`;
+        throw new Error(error_msg);
+    }
+}
+
+function _getBearToken() {
+    let expire_time: string | number = 0, data;
+    return async function getToken(props?: GetTokenProps): Promise<TokenResp> {
+        try{
+            if(!expire_time || new Date(expire_time) <= new Date()){
+                const resp = await fetchBearToken(props);
+                const {expireTime} = resp || {};
+                expire_time = expireTime;
+                data = resp;
+            }
+            return data;
+        } catch(error) {
+            throw new Error(error?.message);
+        }
+    }
+}
+
+export const getBearToken: (props?: GetTokenProps) => Promise<TokenResp> = _getBearToken();