feat(sdk): support thinking controll for model invoking

Author: ZhaoHeh

packages/ui-tars/sdk/src/GUIAgent.ts

@@ -108,6 +108,8 @@ export class GUIAgent<T extends Operator> extends BaseGUIAgent<
 
     let loopCnt = 0;
     let snapshotErrCnt = 0;
+    let totalTokens = 0;
+    let totalTime = 0;
 
     // start running agent
     data.status = StatusEnum.RUNNING;
@@ -240,44 +242,51 @@ export class GUIAgent<T extends Operator> extends BaseGUIAgent<
           scaleFactor: snapshot.scaleFactor,
           uiTarsVersion: this.uiTarsVersion,
         };
-        const { prediction, parsedPredictions } = await asyncRetry(
-          async (bail) => {
-            try {
-              const result = await model.invoke(vlmParams);
-              return result;
-            } catch (error: unknown) {
-              if (
-                error instanceof Error &&
-                (error?.name === 'APIUserAbortError' ||
-                  error?.message?.includes('aborted'))
-              ) {
-                bail(error as unknown as Error);
+        const { prediction, parsedPredictions, costTime, costTokens } =
+          await asyncRetry(
+            async (bail) => {
+              try {
+                const result = await model.invoke(vlmParams);
+                return result;
+              } catch (error: unknown) {
+                if (
+                  error instanceof Error &&
+                  (error?.name === 'APIUserAbortError' ||
+                    error?.message?.includes('aborted'))
+                ) {
+                  bail(error as unknown as Error);
+                  return {
+                    prediction: '',
+                    parsedPredictions: [],
+                  };
+                }
+
+                Object.assign(data, {
+                  status: StatusEnum.ERROR,
+                  error: this.guiAgentErrorParser(
+                    ErrorStatusEnum.INVOKE_RETRY_ERROR,
+                    error as Error,
+                  ),
+                });
+
                 return {
                   prediction: '',
                   parsedPredictions: [],
                 };
               }
+            },
+            {
+              retries: retry?.model?.maxRetries ?? 0,
+              onRetry: retry?.model?.onRetry,
+            },
+          );
 
-              Object.assign(data, {
-                status: StatusEnum.ERROR,
-                error: this.guiAgentErrorParser(
-                  ErrorStatusEnum.INVOKE_RETRY_ERROR,
-                  error as Error,
-                ),
-              });
+        totalTokens += costTokens || 0;
+        totalTime += costTime || 0;
 
-              return {
-                prediction: '',
-                parsedPredictions: [],
-              };
-            }
-          },
-          {
-            retries: retry?.model?.maxRetries ?? 0,
-            onRetry: retry?.model?.onRetry,
-          },
+        logger.info(
+          `[GUIAgent] consumes: >>> costTime: ${costTime}, costTokens: ${costTokens} <<<`,
         );
-
         logger.info('[GUIAgent] Response:', prediction);
         logger.info(
           '[GUIAgent] Parsed Predictions:',
@@ -450,6 +459,10 @@ export class GUIAgent<T extends Operator> extends BaseGUIAgent<
             ),
         });
       }
+
+      logger.info(
+        `[GUIAgent] >>> totalTokens: ${totalTokens}, totalTime: ${totalTime}, loopCnt: ${loopCnt} <<<`,
+      );
     }
   }
 

packages/ui-tars/sdk/src/Model.ts

@@ -4,6 +4,7 @@
  */
 import OpenAI, { type ClientOptions } from 'openai';
 import {
+  type ChatCompletionCreateParamsNonStreaming,
   type ChatCompletionCreateParamsBase,
   type ChatCompletionMessageParam,
 } from 'openai/resources/chat/completions';
@@ -29,6 +30,13 @@ type OpenAIChatCompletionCreateParams = Omit<ClientOptions, 'maxRetries'> &
 
 export interface UITarsModelConfig extends OpenAIChatCompletionCreateParams {}
 
+export interface ThinkingVisionProModelConfig
+  extends ChatCompletionCreateParamsNonStreaming {
+  thinking?: {
+    type: 'enabled' | 'disabled';
+  };
+}
+
 export class UITarsModel extends Model {
   constructor(protected readonly modelConfig: UITarsModelConfig) {
     super();
@@ -60,6 +68,8 @@ export class UITarsModel extends Model {
     },
   ): Promise<{
     prediction: string;
+    costTime?: number;
+    costTokens?: number;
   }> {
     const { messages } = params;
     const {
@@ -79,25 +89,38 @@ export class UITarsModel extends Model {
       apiKey,
     });
 
-    const result = await openai.chat.completions.create(
-      {
-        model,
-        messages,
-        stream: false,
-        seed: null,
-        stop: null,
-        frequency_penalty: null,
-        presence_penalty: null,
-        // custom options
-        max_tokens,
-        temperature,
-        top_p,
+    const createCompletionPrams: ChatCompletionCreateParamsNonStreaming = {
+      model,
+      messages,
+      stream: false,
+      seed: null,
+      stop: null,
+      frequency_penalty: null,
+      presence_penalty: null,
+      // custom options
+      max_tokens,
+      temperature,
+      top_p,
+    };
+
+    const createCompletionPramsThinkingVp: ThinkingVisionProModelConfig = {
+      ...createCompletionPrams,
+      thinking: {
+        type: 'disabled',
       },
+    };
+
+    const startTime = Date.now();
+    const result = await openai.chat.completions.create(
+      createCompletionPramsThinkingVp,
       options,
     );
+    const costTime = Date.now() - startTime;
 
     return {
       prediction: result.choices?.[0]?.message?.content ?? '',
+      costTime: costTime,
+      costTokens: result.usage?.total_tokens ?? 0,
     };
   }
 
@@ -152,7 +175,7 @@ export class UITarsModel extends Model {
       throw err;
     }
 
-    const { prediction } = result;
+    const { prediction, costTime, costTokens } = result;
 
     try {
       const { parsed: parsedPredictions } = await actionParser({
@@ -165,6 +188,8 @@ export class UITarsModel extends Model {
       return {
         prediction,
         parsedPredictions,
+        costTime,
+        costTokens,
       };
     } catch (error) {
       logger?.error('[UITarsModel] error', error);

packages/ui-tars/sdk/src/types.ts

@@ -52,6 +52,8 @@ export interface InvokeParams {
 export interface InvokeOutput {
   prediction: string;
   parsedPredictions: PredictionParsed[];
+  costTime?: number;
+  costTokens?: number;
   // TODO: status: StatusEnum, status should be provided by model
 }
 export abstract class Operator extends BaseOperator {