feat: let max_tokens configurable (#212)

* feat: let max_tokens configurable * fix: update ci test case
web-infra-dev · Dec 26, 2024 · 4a82e9b · 4a82e9b
1 parent 7e5b4c9
commit 4a82e9b
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 3 deletions.
diff --git a/apps/site/docs/en/model-provider.md b/apps/site/docs/en/model-provider.md
@@ -28,6 +28,9 @@ export MIDSCENE_OPENAI_INIT_CONFIG_JSON='{"baseURL":"....","defaultHeaders":{"ke
 
 # if you want to use proxy. Midscene uses `socks-proxy-agent` under the hood.
 export MIDSCENE_OPENAI_SOCKS_PROXY="socks5://127.0.0.1:1080"
+
+# if you want to specify the max tokens for the model
+export OPENAI_MAX_TOKENS=2048
 ```
 
 ## Using Azure OpenAI Service

diff --git a/apps/site/docs/zh/model-provider.md b/apps/site/docs/zh/model-provider.md
@@ -25,6 +25,9 @@ export MIDSCENE_OPENAI_INIT_CONFIG_JSON='{"baseURL":"....","defaultHeaders":{"ke
 
 # 可选, 如果你想使用代理。Midscene 使用 `socks-proxy-agent` 作为底层库。
 export MIDSCENE_OPENAI_SOCKS_PROXY="socks5://127.0.0.1:1080"
+
+# 可选, 如果你想指定模型 max_tokens
+export OPENAI_MAX_TOKENS=2048
 ```
 
 ## 使用 Azure OpenAI 服务时的配置

diff --git a/packages/midscene/src/ai-model/openai/index.ts b/packages/midscene/src/ai-model/openai/index.ts
@@ -24,6 +24,7 @@ import {
   MIDSCENE_USE_AZURE_OPENAI,
   OPENAI_API_KEY,
   OPENAI_BASE_URL,
+  OPENAI_MAX_TOKENS,
   OPENAI_USE_AZURE,
   allAIConfig,
   getAIConfig,
@@ -149,14 +150,19 @@ export async function call(
   const shouldPrintTiming =
     typeof getAIConfig(MIDSCENE_DEBUG_AI_PROFILE) === 'string';
 
+  const maxTokens = getAIConfig(OPENAI_MAX_TOKENS);
+
   const startTime = Date.now();
   const model = getModelName();
   let content: string | undefined;
   let usage: OpenAI.CompletionUsage | undefined;
   const commonConfig = {
     temperature: 0.1,
     stream: false,
-    max_tokens: 3000,
+    max_tokens:
+      typeof maxTokens === 'number'
+        ? maxTokens
+        : Number.parseInt(maxTokens || '2048', 10),
   };
   if (style === 'openai') {
     const result = await completion.create({

diff --git a/packages/midscene/src/env.ts b/packages/midscene/src/env.ts
@@ -10,6 +10,7 @@ export const MIDSCENE_DEBUG_MODE = 'MIDSCENE_DEBUG_MODE';
 export const MIDSCENE_OPENAI_SOCKS_PROXY = 'MIDSCENE_OPENAI_SOCKS_PROXY';
 export const OPENAI_API_KEY = 'OPENAI_API_KEY';
 export const OPENAI_BASE_URL = 'OPENAI_BASE_URL';
+export const OPENAI_MAX_TOKENS = 'OPENAI_MAX_TOKENS';
 export const MIDSCENE_MODEL_TEXT_ONLY = 'MIDSCENE_MODEL_TEXT_ONLY';
 
 export const MIDSCENE_CACHE = 'MIDSCENE_CACHE';
@@ -43,6 +44,7 @@ const allConfigFromEnv = () => {
     [OPENAI_BASE_URL]: process.env[OPENAI_BASE_URL] || undefined,
     [MIDSCENE_MODEL_TEXT_ONLY]:
       process.env[MIDSCENE_MODEL_TEXT_ONLY] || undefined,
+    [OPENAI_MAX_TOKENS]: process.env[OPENAI_MAX_TOKENS] || undefined,
     [OPENAI_USE_AZURE]: process.env[OPENAI_USE_AZURE] || undefined,
     [MIDSCENE_CACHE]: process.env[MIDSCENE_CACHE] || undefined,
     [MATCH_BY_POSITION]: process.env[MATCH_BY_POSITION] || undefined,

diff --git a/packages/web-integration/tests/ai/web/puppeteer/showcase.test.ts b/packages/web-integration/tests/ai/web/puppeteer/showcase.test.ts
@@ -86,7 +86,7 @@ describe(
     });
 
     it('search engine', async () => {
-      const { originPage, reset } = await launchPage('https://www.baidu.com/');
+      const { originPage, reset } = await launchPage('https://www.bing.com/');
       const mid = new PuppeteerAgent(originPage);
       await mid.aiAction(
         'type "AI 101" in search box, hit Enter, wait 2s, click the second result, wait 4s',
@@ -100,7 +100,6 @@ describe(
     it('scroll', async () => {
       const htmlPath = path.join(__dirname, 'scroll.html');
       const { originPage, reset } = await launchPage(`file://${htmlPath}`);
-      // const { originPage, reset } = await launchPage('https://news.baidu.com/');
       const mid = new PuppeteerAgent(originPage);
       await mid.aiAction(
         'find the "Vertical 2" element, scroll down 200px, find the "Horizontal 2" element, scroll right 100px',