feat: enhance token counting by extracting text from messages using VSCode LM API

Author: NaccOll

src/api/providers/vscode-lm.ts

@@ -7,7 +7,7 @@ import type { ApiHandlerOptions } from "../../shared/api"
 import { SELECTOR_SEPARATOR, stringifyVsCodeLmModelSelector } from "../../shared/vsCodeSelectorUtils"
 
 import { ApiStream } from "../transform/stream"
-import { convertToVsCodeLmMessages } from "../transform/vscode-lm-format"
+import { convertToVsCodeLmMessages, extractTextCountFromMessage } from "../transform/vscode-lm-format"
 
 import { BaseProvider } from "./base-provider"
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
@@ -231,7 +231,8 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
 					console.debug("Roo Code <Language Model API>: Empty chat message content")
 					return 0
 				}
-				tokenCount = await this.client.countTokens(text, this.currentRequestCancellation.token)
+				const countMessage = extractTextCountFromMessage(text)
+				tokenCount = await this.client.countTokens(countMessage, this.currentRequestCancellation.token)
 			} else {
 				console.warn("Roo Code <Language Model API>: Invalid input type for token counting")
 				return 0
@@ -268,15 +269,10 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
 		}
 	}
 
-	private async calculateTotalInputTokens(
-		systemPrompt: string,
-		vsCodeLmMessages: vscode.LanguageModelChatMessage[],
-	): Promise<number> {
-		const systemTokens: number = await this.internalCountTokens(systemPrompt)
-
+	private async calculateTotalInputTokens(vsCodeLmMessages: vscode.LanguageModelChatMessage[]): Promise<number> {
 		const messageTokens: number[] = await Promise.all(vsCodeLmMessages.map((msg) => this.internalCountTokens(msg)))
 
-		return systemTokens + messageTokens.reduce((sum: number, tokens: number): number => sum + tokens, 0)
+		return messageTokens.reduce((sum: number, tokens: number): number => sum + tokens, 0)
 	}
 
 	private ensureCleanState(): void {
@@ -359,7 +355,7 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
 		this.currentRequestCancellation = new vscode.CancellationTokenSource()
 
 		// Calculate input tokens before starting the stream
-		const totalInputTokens: number = await this.calculateTotalInputTokens(systemPrompt, vsCodeLmMessages)
+		const totalInputTokens: number = await this.calculateTotalInputTokens(vsCodeLmMessages)
 
 		// Accumulate the text and count at the end of the stream to reduce token counting overhead.
 		let accumulatedText: string = ""

src/api/transform/vscode-lm-format.ts

@@ -155,3 +155,32 @@ export function convertToAnthropicRole(vsCodeLmMessageRole: vscode.LanguageModel
 			return null
 	}
 }
+
+export function extractTextCountFromMessage(message: vscode.LanguageModelChatMessage): string {
+	let text = ""
+	if (Array.isArray(message.content)) {
+		for (const item of message.content) {
+			if (item instanceof vscode.LanguageModelTextPart) {
+				text += item.value
+			}
+			if (item instanceof vscode.LanguageModelToolResultPart) {
+				text += item.callId
+				for (const part of item.content) {
+					if (part instanceof vscode.LanguageModelTextPart) {
+						text += part.value
+					}
+				}
+			}
+			if (item instanceof vscode.LanguageModelToolCallPart) {
+				text += item.name
+				text += item.callId
+				if (item.input && Object.keys(item.input).length > 0) {
+					text += JSON.stringify(item.input)
+				}
+			}
+		}
+	} else if (typeof message.content === "string") {
+		text += message.content
+	}
+	return text
+}