Browse Source

Stop making count_tokens requests (#9884)

Matt Rubens 1 month ago
parent
commit
4a5cbcba86

+ 0 - 4
src/api/providers/__tests__/minimax.spec.ts

@@ -16,12 +16,10 @@ import { MiniMaxHandler } from "../minimax"
 
 vitest.mock("@anthropic-ai/sdk", () => {
 	const mockCreate = vitest.fn()
-	const mockCountTokens = vitest.fn()
 	return {
 		Anthropic: vitest.fn(() => ({
 			messages: {
 				create: mockCreate,
-				countTokens: mockCountTokens,
 			},
 		})),
 	}
@@ -30,13 +28,11 @@ vitest.mock("@anthropic-ai/sdk", () => {
 describe("MiniMaxHandler", () => {
 	let handler: MiniMaxHandler
 	let mockCreate: any
-	let mockCountTokens: any
 
 	beforeEach(() => {
 		vitest.clearAllMocks()
 		const anthropicInstance = (Anthropic as unknown as any)()
 		mockCreate = anthropicInstance.messages.create
-		mockCountTokens = anthropicInstance.messages.countTokens
 	})
 
 	describe("International MiniMax (default)", () => {

+ 0 - 26
src/api/providers/anthropic.ts

@@ -402,30 +402,4 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
 		const content = message.content.find(({ type }) => type === "text")
 		return content?.type === "text" ? content.text : ""
 	}
-
-	/**
-	 * Counts tokens for the given content using Anthropic's API
-	 *
-	 * @param content The content blocks to count tokens for
-	 * @returns A promise resolving to the token count
-	 */
-	override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
-		try {
-			// Use the current model
-			const { id: model } = this.getModel()
-
-			const response = await this.client.messages.countTokens({
-				model,
-				messages: [{ role: "user", content: content }],
-			})
-
-			return response.input_tokens
-		} catch (error) {
-			// Log error but fallback to tiktoken estimation
-			console.warn("Anthropic token counting failed, using fallback", error)
-
-			// Use the base provider's implementation as fallback
-			return super.countTokens(content)
-		}
-	}
 }

+ 1 - 26
src/api/providers/gemini.ts

@@ -6,7 +6,6 @@ import {
 	type GenerateContentConfig,
 	type GroundingMetadata,
 	FunctionCallingConfigMode,
-	Content,
 } from "@google/genai"
 import type { JWTInput } from "google-auth-library"
 
@@ -15,7 +14,7 @@ import { type ModelInfo, type GeminiModelId, geminiDefaultModelId, geminiModels
 import type { ApiHandlerOptions } from "../../shared/api"
 import { safeJsonParse } from "../../shared/safeJsonParse"
 
-import { convertAnthropicContentToGemini, convertAnthropicMessageToGemini } from "../transform/gemini-format"
+import { convertAnthropicMessageToGemini } from "../transform/gemini-format"
 import { t } from "i18next"
 import type { ApiStream, GroundingSource } from "../transform/stream"
 import { getModelParams } from "../transform/model-params"
@@ -431,30 +430,6 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		}
 	}
 
-	override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
-		try {
-			const { id: model } = this.getModel()
-
-			const countTokensRequest = {
-				model,
-				// Token counting does not need encrypted continuation; always drop thoughtSignature.
-				contents: convertAnthropicContentToGemini(content, { includeThoughtSignatures: false }),
-			}
-
-			const response = await this.client.models.countTokens(countTokensRequest)
-
-			if (response.totalTokens === undefined) {
-				console.warn("Gemini token counting returned undefined, using fallback")
-				return super.countTokens(content)
-			}
-
-			return response.totalTokens
-		} catch (error) {
-			console.warn("Gemini token counting failed, using fallback", error)
-			return super.countTokens(content)
-		}
-	}
-
 	public getThoughtSignature(): string | undefined {
 		return this.lastThoughtSignature
 	}

+ 0 - 23
src/api/providers/minimax.ts

@@ -303,27 +303,4 @@ export class MiniMaxHandler extends BaseProvider implements SingleCompletionHand
 		const content = message.content.find(({ type }) => type === "text")
 		return content?.type === "text" ? content.text : ""
 	}
-
-	/**
-	 * Counts tokens for the given content using Anthropic's token counting
-	 * Falls back to base provider's tiktoken estimation if counting fails
-	 */
-	override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
-		try {
-			const { id: model } = this.getModel()
-
-			const response = await this.client.messages.countTokens({
-				model,
-				messages: [{ role: "user", content: content }],
-			})
-
-			return response.input_tokens
-		} catch (error) {
-			// Log error but fallback to tiktoken estimation
-			console.warn("MiniMax token counting failed, using fallback", error)
-
-			// Use the base provider's implementation as fallback
-			return super.countTokens(content)
-		}
-	}
 }