Limits output tokens to 8192 for all models on Unbound

Author: pugazhendhi-m

src/api/providers/__tests__/unbound.test.ts

@@ -1,6 +1,7 @@
-import { UnboundHandler } from "../unbound"
+import { UnboundHandler, getUnboundModels } from "../unbound"
 import { ApiHandlerOptions } from "../../../shared/api"
 import { Anthropic } from "@anthropic-ai/sdk"
+import axios from "axios"
 
 // Mock OpenAI client
 const mockCreate = jest.fn()
@@ -299,3 +300,49 @@ describe("UnboundHandler", () => {
 		})
 	})
 })
+
+describe("getUnboundModels", () => {
+	beforeEach(() => {
+		jest.clearAllMocks()
+	})
+
+	it("should restrict maxTokens to 8192 if it's greater", async () => {
+		// Mock axios.get to return a model with maxTokens > 8192
+		jest.spyOn(axios, "get").mockResolvedValueOnce({
+			data: {
+				"anthropic/claude-3-opus": {
+					maxTokens: "10000",
+					contextWindow: "200000",
+					supportsPromptCaching: true,
+					inputTokenPrice: "0.01",
+					outputTokenPrice: "0.02",
+				},
+			},
+		})
+
+		const models = await getUnboundModels()
+
+		expect(models["anthropic/claude-3-opus"]).toBeDefined()
+		expect(models["anthropic/claude-3-opus"].maxTokens).toBe(8192) // Should be restricted to 8192
+	})
+
+	it("should keep maxTokens unchanged if it's 4096 or less", async () => {
+		// Mock axios.get to return a model with maxTokens = 4096
+		jest.spyOn(axios, "get").mockResolvedValueOnce({
+			data: {
+				"anthropic/claude-3-haiku": {
+					maxTokens: "4096",
+					contextWindow: "200000",
+					supportsPromptCaching: true,
+					inputTokenPrice: "0.01",
+					outputTokenPrice: "0.02",
+				},
+			},
+		})
+
+		const models = await getUnboundModels()
+
+		expect(models["anthropic/claude-3-haiku"]).toBeDefined()
+		expect(models["anthropic/claude-3-haiku"].maxTokens).toBe(4096) // Should remain unchanged
+	})
+})

src/api/providers/unbound.ts

@@ -210,15 +210,9 @@ export async function getUnboundModels() {
 					cacheReadsPrice: model?.cacheReadPrice ? parseFloat(model.cacheReadPrice) : undefined,
 				}
 
-				switch (true) {
-					case modelId.startsWith("anthropic/"):
-						// Set max tokens to 8192 for supported Anthropic models
-						if (modelInfo.maxTokens !== 4096) {
-							modelInfo.maxTokens = 8192
-						}
-						break
-					default:
-						break
+				// Restrict max tokens to 8192 for all models
+				if (modelInfo.maxTokens && modelInfo.maxTokens > 8192) {
+					modelInfo.maxTokens = 8192
 				}
 
 				models[modelId] = modelInfo