✨ feat: Enhance AI commit generation with token tracking and streaming support

Author: canisminor1990

packages/lobe-commit/src/commands/Ai/index.tsx

@@ -16,6 +16,7 @@ const Ai = memo(() => {
     start,
     loadingInfo,
     loading,
+    tokenUsage,
     message: streamingMessage,
   } = useCommits({ setMessage });
 
@@ -34,7 +35,7 @@ const Ai = memo(() => {
         )
       }
       reverse
-      title={`🤯 AI Commit Generator ${commitConfig.stream ? '(Streaming)' : ''}`}
+      title={`🤯 AI Commit Generator ${commitConfig.stream ? '(Streaming)' : ''} ${tokenUsage > 0 ? `[Tokens: ${tokenUsage}]` : ''}`}
     >
       <AiMessageDisplay
         loading={loading}

packages/lobe-commit/src/commands/Commit/AiCommit.tsx

@@ -23,6 +23,7 @@ const AiCommit = memo(() => {
     loadingInfo,
     loading,
     restart,
+    tokenUsage,
     message: streamingMessage,
   } = useCommits({ setMessage });
 
@@ -71,7 +72,7 @@ const AiCommit = memo(() => {
   return (
     <Panel
       footer={!loading && message && <SelectInput items={items} onSelect={handleSelect} />}
-      title={`🤯 AI Commit Generator ${commitConfig.stream ? '(Streaming)' : ''}`}
+      title={`🤯 AI Commit Generator ${commitConfig.stream ? '(Streaming)' : ''} ${tokenUsage > 0 ? `[Tokens: ${tokenUsage}]` : ''}`}
     >
       {summary && (
         <SplitView direction={'bottom'}>

packages/lobe-commit/src/core/Commits.ts

@@ -20,6 +20,7 @@ import { RecursiveCharacterTextSplitter } from '../../../common/textSplitter';
 export interface GenAiCommitProps {
   cacheSummary?: string;
   onStreamMessage?: (message: string) => void;
+  onTokenUsage?: (tokenUsage: number) => void;
   setLoadingInfo: (text: string) => void;
   setSummary: (text: string) => void;
 }
@@ -52,7 +53,13 @@ export class Commits {
     this.prompt = promptCommits();
   }
 
-  async genCommit({ setLoadingInfo, setSummary, cacheSummary, onStreamMessage }: GenAiCommitProps) {
+  async genCommit({
+    setLoadingInfo,
+    setSummary,
+    cacheSummary,
+    onStreamMessage,
+    onTokenUsage,
+  }: GenAiCommitProps) {
     setLoadingInfo(' Generating...');
 
     // STEP 1
@@ -63,16 +70,21 @@ export class Commits {
       summary,
     });
 
+    // Calculate input tokens
+    const inputTokens = calcToken(JSON.stringify(messages));
+
     if (this.config.stream && onStreamMessage) {
-      return this.genCommitStream(messages, onStreamMessage);
+      return this.genCommitStream(messages, onStreamMessage, onTokenUsage, inputTokens);
     } else {
-      return this.genCommitNonStream(messages);
+      return this.genCommitNonStream(messages, onTokenUsage, inputTokens);
     }
   }
 
   private async genCommitStream(
     messages: any[],
     onStreamMessage: (message: string) => void,
+    onTokenUsage?: (tokenUsage: number) => void,
+    inputTokens?: number,
   ): Promise<string> {
     // 开始流式输出，先调用一次回调来切换UI状态
     onStreamMessage('');
@@ -85,6 +97,7 @@ export class Commits {
     });
 
     let fullMessage = '';
+    let tokenUsage = inputTokens || 0;
 
     for await (const chunk of stream) {
       const content = chunk.choices[0]?.delta?.content || '';
@@ -95,19 +108,26 @@ export class Commits {
           fullMessage.replace(/\((.*?)\):/, (match, p1) => match && `(${p1.toLowerCase()}):`),
         );
         onStreamMessage(processedMessage);
+        tokenUsage += calcToken(content);
       }
     }
 
     if (!fullMessage) {
       alert.error('Diff summary failed, please check your network or try again...', true);
     }
 
+    onTokenUsage?.(tokenUsage);
+
     return addEmojiToMessage(
       fullMessage.replace(/\((.*?)\):/, (match, p1) => match && `(${p1.toLowerCase()}):`),
     );
   }
 
-  private async genCommitNonStream(messages: any[]): Promise<string> {
+  private async genCommitNonStream(
+    messages: any[],
+    onTokenUsage?: (tokenUsage: number) => void,
+    inputTokens?: number,
+  ): Promise<string> {
     const completion = await this.client.chat.completions.create({
       messages: messages as OpenAI.Chat.Completions.ChatCompletionMessageParam[],
       model: this.config.modelName,
@@ -119,6 +139,10 @@ export class Commits {
     if (!result)
       alert.error('Diff summary failed, please check your network or try again...', true);
 
+    const outputTokens = calcToken(result!);
+    const totalTokenUsage = (inputTokens || 0) + outputTokens;
+    onTokenUsage?.(totalTokenUsage);
+
     return addEmojiToMessage(
       result!.replace(/\((.*?)\):/, (match, p1) => match && `(${p1.toLowerCase()}):`),
     );

packages/lobe-commit/src/hooks/useCommits.ts

@@ -16,6 +16,7 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
   const [isGlobalLoading, setIsGlobalLoading] = useState(true);
   const [key, setKey] = useState<string>(Date.now().toString());
   const [streamingMessage, setStreamingMessage] = useState<string>('');
+  const [tokenUsage, setTokenUsage] = useState<number>(0);
   const commitConfig = selectors.getCommitConfig();
 
   const handleStreamMessage = useCallback(
@@ -27,6 +28,10 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
     [setMessage],
   );
 
+  const handleTokenUsage = useCallback((usage: number) => {
+    setTokenUsage(usage);
+  }, []);
+
   const { data, isLoading } = useSWR(
     shouldFetch ? key : null,
     async () => {
@@ -36,6 +41,7 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
       return commits.current.genCommit({
         cacheSummary: summary,
         onStreamMessage: commitConfig.stream ? handleStreamMessage : undefined,
+        onTokenUsage: handleTokenUsage,
         setLoadingInfo,
         setSummary,
       });
@@ -60,6 +66,7 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
     setKey(Date.now().toString());
     setIsGlobalLoading(true);
     setStreamingMessage('');
+    setTokenUsage(0);
     setShouldFetch(true);
   }, []);
 
@@ -68,6 +75,7 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
     setKey(Date.now().toString());
     setIsGlobalLoading(true);
     setStreamingMessage('');
+    setTokenUsage(0);
     setShouldFetch(true);
   }, []);
 
@@ -84,5 +92,6 @@ export const useCommits = ({ setMessage, onSuccess, onError, ...config }: Commit
     start,
     stop,
     summary,
+    tokenUsage,
   };
 };