refactor!: optimize type CodeReviewResult

Author: mys1024

README.md

@@ -38,7 +38,7 @@ Supported AI service providers:
 ```javascript
 import { codeReview } from 'cr-asst';
 
-const { content } = await codeReview({
+const { reviewReport } = await codeReview({
   headRef: 'branch-to-review'
   baseRef: 'main',
   model: 'gpt-4',

README.zh.md

@@ -38,7 +38,7 @@ npx cr-asst --head-ref branch-to-review --base-ref main --model gpt-4 --provider
 ```javascript
 import { codeReview } from 'cr-asst';
 
-const { content } = await codeReview({
+const { reviewReport } = await codeReview({
   headRef: 'branch-to-review'
   baseRef: 'main',
   model: 'gpt-4',

src/code_review/index.ts

@@ -1,16 +1,15 @@
 import { writeFile } from 'node:fs/promises';
 import chalk from 'chalk';
-import { stepCountIs, type LanguageModel, type ModelMessage } from 'ai';
+import { stepCountIs, streamText, type LanguageModel, type ModelMessage } from 'ai';
 import type { CodeReviewOptions, CodeReviewResult } from '../types';
 import {
-  getUserPrompt,
-  getSystemPrompt,
-  getApprovalCheckPrompt,
-  getApprovalCheckStatusPrompt,
+  getReviewReportMessages,
+  getApprovalCheckCommentMessages,
+  getApprovalCheckStatusMessages,
 } from './prompts/index';
 import { runCmd } from './utils';
 import { reviewReportTools } from './tools';
-import { initModel, callModel } from './model';
+import { initModel, handleStreamTextResult } from './model';
 
 export async function codeReview(options: CodeReviewOptions): Promise<CodeReviewResult> {
   // init model
@@ -22,8 +21,8 @@ export async function codeReview(options: CodeReviewOptions): Promise<CodeReview
     model,
   });
 
-  // generate approval check
-  const approvalCheck = options.approvalCheck
+  // generate approval check comment
+  const approvalCheckComment = options.approvalCheck
     ? await generateApprovalCheck({
         ...options,
         model,
@@ -33,30 +32,22 @@ export async function codeReview(options: CodeReviewOptions): Promise<CodeReview
 
   // generate approval check status
   const approvalCheckStatus =
-    options.approvalCheck && approvalCheck
+    options.approvalCheck && approvalCheckComment
       ? await generateApprovalCheckStatus({
           ...options,
           model,
-          prevMessages: approvalCheck.messages,
+          prevMessages: approvalCheckComment.messages,
         })
       : undefined;
 
   // return
   return {
-    content: reviewReport.text,
-    reasoningContent: reviewReport.reasoning,
-    debug: {
-      diffsCmd: reviewReport.diffsCmd,
-      diffs: reviewReport.diffs,
-      stats: reviewReport.stats,
-      usage: reviewReport.usage,
-    },
+    reviewReport,
     approvalCheck:
-      approvalCheck && approvalCheckStatus
+      approvalCheckComment && approvalCheckStatus
         ? {
-            content: approvalCheck.text,
-            reasoningContent: approvalCheck.reasoning,
-            approved: approvalCheckStatus.approved,
+            approvalCheckComment,
+            approvalCheckStatus,
           }
         : undefined,
   };
@@ -100,29 +91,32 @@ async function generateReviewReport(
   }
   const diffs = await runCmd('git', diffArgs);
 
-  // generate review report
-  const result = await callModel({
+  // messages
+  const messages = await getReviewReportMessages({
+    systemPromptFile,
+    promptFile,
+    disableTools,
+    diffs,
+    baseRef,
+    headRef,
+  });
+
+  // call model
+  const result = await handleStreamTextResult({
     title: 'Review Report',
-    model,
-    tools: disableTools ? undefined : reviewReportTools,
-    stopWhen: stepCountIs(maxSteps),
-    messages: [
-      {
-        role: 'system',
-        content: await getSystemPrompt({
-          systemPromptFile,
-          disableTools,
-          diffs,
-          baseRef,
-          headRef,
-        }),
-      },
-      { role: 'user', content: await getUserPrompt(promptFile) },
-    ],
     print,
-    temperature,
-    topP,
-    topK,
+    streamTextResult: streamText({
+      model,
+      messages,
+      tools: disableTools ? undefined : reviewReportTools,
+      stopWhen: stepCountIs(maxSteps),
+      temperature,
+      topP,
+      topK,
+      onError: ({ error }) => {
+        throw new Error('failed to call the model', { cause: error });
+      },
+    }),
   });
 
   // write output file
@@ -133,6 +127,7 @@ async function generateReviewReport(
   // return
   return {
     ...result,
+    messages: [...messages, ...result.messages],
     diffsCmd,
     diffs,
   };
@@ -147,26 +142,32 @@ async function generateApprovalCheck(
   // options
   const { model, prevMessages, approvalCheck, print, temperature, topP, topK } = options;
 
-  // generate approval check
-  const result = await callModel({
-    title: 'Approval Check',
-    model,
-    messages: [
-      ...prevMessages,
-      {
-        role: 'user',
-        content: await getApprovalCheckPrompt(approvalCheck),
-      },
-    ],
+  // messages
+  const messages = await getApprovalCheckCommentMessages({
+    prevMessages,
+    approvalCheck,
+  });
+
+  // call model
+  const result = await handleStreamTextResult({
+    title: 'Approval Check Comment',
     print,
-    temperature,
-    topP,
-    topK,
+    streamTextResult: streamText({
+      model,
+      messages,
+      temperature,
+      topP,
+      topK,
+      onError: ({ error }) => {
+        throw new Error('failed to call the model', { cause: error });
+      },
+    }),
   });
 
   // return
   return {
     ...result,
+    messages: [...messages, ...result.messages],
   };
 }
 
@@ -179,21 +180,23 @@ async function generateApprovalCheckStatus(
   // options
   const { model, prevMessages, print, temperature, topP, topK } = options;
 
-  // generate approval check
-  const result = await callModel({
+  // messages
+  const messages = getApprovalCheckStatusMessages({ prevMessages });
+
+  // call model
+  const result = await handleStreamTextResult({
     title: 'Approval Check Status',
-    model,
-    messages: [
-      ...prevMessages,
-      {
-        role: 'user',
-        content: await getApprovalCheckStatusPrompt(),
-      },
-    ],
     print,
-    temperature,
-    topP,
-    topK,
+    streamTextResult: streamText({
+      model,
+      messages,
+      temperature,
+      topP,
+      topK,
+      onError: ({ error }) => {
+        throw new Error('failed to call the model', { cause: error });
+      },
+    }),
   });
 
   // parse approved flag
@@ -202,6 +205,7 @@ async function generateApprovalCheckStatus(
   // return
   return {
     ...result,
+    messages: [...messages, ...result.messages],
     approved,
   };
 }

src/code_review/model.ts

@@ -2,19 +2,18 @@ import { stdout } from 'node:process';
 import { inspect } from 'node:util';
 import chalk from 'chalk';
 import { fetch, ProxyAgent } from 'undici';
-import {
-  streamText,
-  type LanguageModel,
-  type ToolSet,
-  type ToolChoice,
-  type ModelMessage,
-} from 'ai';
+import { type StreamTextResult, type LanguageModel, type ToolSet } from 'ai';
 import { createOpenAI, type OpenAIProviderSettings } from '@ai-sdk/openai';
 import { createDeepSeek } from '@ai-sdk/deepseek';
 import { createXai } from '@ai-sdk/xai';
 import { createAnthropic } from '@ai-sdk/anthropic';
 import { createGoogleGenerativeAI } from '@ai-sdk/google';
-import type { CodeReviewOptions, CompletionStats, CompletionUsage } from '../types';
+import type {
+  CodeReviewOptions,
+  LanguageModelCallResult,
+  LanguageModelCallStats,
+  LanguageModelCallUsage,
+} from '../types';
 import { usageToString, statsToString, getHttpProxyUrl } from './utils';
 
 export function initModel(options: CodeReviewOptions): LanguageModel {
@@ -66,33 +65,13 @@ export function initModel(options: CodeReviewOptions): LanguageModel {
   return model;
 }
 
-export async function callModel<TOOLS extends ToolSet>(options: {
+export async function handleStreamTextResult<TOOLS extends ToolSet>(options: {
   title: string;
-  model: LanguageModel;
-  messages: ModelMessage[];
   print?: boolean;
-  tools?: TOOLS;
-  toolChoice?: ToolChoice<TOOLS>;
-  prepareStep?: Parameters<typeof streamText<TOOLS>>[0]['prepareStep'];
-  stopWhen?: Parameters<typeof streamText<TOOLS>>[0]['stopWhen'];
-  temperature?: number;
-  topP?: number;
-  topK?: number;
-}) {
+  streamTextResult: StreamTextResult<TOOLS, string>;
+}): Promise<LanguageModelCallResult> {
   // options
-  const {
-    title,
-    model,
-    messages,
-    print,
-    tools,
-    toolChoice,
-    prepareStep,
-    stopWhen,
-    temperature,
-    topP,
-    topK,
-  } = options;
+  const { title, print, streamTextResult } = options;
 
   // print title
   if (print) {
@@ -103,24 +82,8 @@ export async function callModel<TOOLS extends ToolSet>(options: {
     );
   }
 
-  // call model
-  const result = streamText({
-    model,
-    messages,
-    tools,
-    toolChoice,
-    prepareStep,
-    stopWhen,
-    temperature,
-    topP,
-    topK,
-    onError: ({ error }) => {
-      throw new Error('failed to call the model', { cause: error });
-    },
-  });
-
   // init stats
-  const stats: CompletionStats = {
+  const stats: LanguageModelCallStats = {
     startedAt: Date.now(),
     firstTokenReceivedAt: 0,
     finishedAt: 0,
@@ -132,7 +95,7 @@ export async function callModel<TOOLS extends ToolSet>(options: {
   let stepCnt = 0;
   let textPartCnt = 0;
   let reasoningPartCnt = 0;
-  for await (const streamPart of result.fullStream) {
+  for await (const streamPart of streamTextResult.fullStream) {
     if (streamPart.type === 'start-step') {
       if (print) {
         console.log(chalk.blue(`[STEP_${stepCnt}]\n`));
@@ -174,11 +137,11 @@ export async function callModel<TOOLS extends ToolSet>(options: {
   }
 
   // destructure result
-  const steps = await result.steps;
+  const steps = await streamTextResult.steps;
   const lastStep = steps[steps.length - 1];
   const text = lastStep.text;
   const reasoning = lastStep.reasoning;
-  const usage: CompletionUsage = await result.usage;
+  const usage: LanguageModelCallUsage = await streamTextResult.usage;
 
   // update usage
   if (typeof usage.inputTokens === 'number' && typeof usage.cachedInputTokens === 'number') {
@@ -201,14 +164,11 @@ export async function callModel<TOOLS extends ToolSet>(options: {
     console.log();
   }
 
-  // get all messages
-  const allMessages: ModelMessage[] = [...messages, ...(await result.response).messages];
-
   // return
   return {
     text,
-    reasoning: reasoning.map((r) => r.text).join('\n'),
-    messages: allMessages,
+    reasoning: reasoning.map((r) => r.text),
+    messages: (await streamTextResult.response).messages,
     usage,
     stats,
   };