fix prompt router bug

Author: Smartsheet-JB-Brown

src/api/providers/__tests__/bedrock-invokedModelId.test.ts

@@ -103,12 +103,17 @@ describe("AwsBedrockHandler with invokedModelId", () => {
 			awsAccessKey: "test-access-key",
 			awsSecretKey: "test-secret-key",
 			awsRegion: "us-east-1",
-			awsCustomArn: "arn:aws:bedrock:us-west-2:699475926481:default-prompt-router/anthropic.claude:1",
+			awsCustomArn: "arn:aws:bedrock:us-west-2:123456789:default-prompt-router/anthropic.claude:1",
 		}
 
 		const handler = new AwsBedrockHandler(mockOptions)
 
-		// Create a spy on the getModel method before mocking it
+		// Verify that getModel returns the updated model info
+		const initialModel = handler.getModel()
+		//the default prompt router model has an input price of 3. After the stream is handled it should be updated to 8
+		expect(initialModel.info.inputPrice).toBe(3)
+
+		// Create a spy on the getModel
 		const getModelByIdSpy = jest.spyOn(handler, "getModelById")
 
 		// Mock the stream to include an event with invokedModelId and usage metadata
@@ -120,7 +125,7 @@ describe("AwsBedrockHandler with invokedModelId", () => {
 						trace: {
 							promptRouter: {
 								invokedModelId:
-									"arn:aws:bedrock:us-west-2:699475926481:inference-profile/us.anthropic.claude-3-5-sonnet-20240620-v1:0",
+									"arn:aws:bedrock:us-west-2:699475926481:inference-profile/us.anthropic.claude-2-1-v1:0",
 								usage: {
 									inputTokens: 150,
 									outputTokens: 250,
@@ -159,13 +164,13 @@ describe("AwsBedrockHandler with invokedModelId", () => {
 			events.push(event)
 		}
 
-		// Verify that getModelById was called with the full ARN
-		expect(getModelByIdSpy).toHaveBeenCalledWith("anthropic.claude-3-5-sonnet-20240620-v1:0")
+		// Verify that getModelById was called with the id, not the full arn
+		expect(getModelByIdSpy).toHaveBeenCalledWith("anthropic.claude-2-1-v1:0", "inference-profile")
 
 		// Verify that getModel returns the updated model info
 		const costModel = handler.getModel()
 		//expect(costModel.id).toBe("anthropic.claude-3-5-sonnet-20240620-v1:0")
-		expect(costModel.info.inputPrice).toBe(3)
+		expect(costModel.info.inputPrice).toBe(8)
 
 		// Verify that a usage event was emitted after updating the costModelConfig
 		const usageEvents = events.filter((event) => event.type === "usage")

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

@@ -77,13 +77,11 @@ describe("AwsBedrockHandler", () => {
 
 				const modelInfo = customArnHandler.getModel()
 
-				// Verify the ARN is preserved as the ID
 				expect(modelInfo.id).toBe(
 					"arn:aws:bedrock:ap-northeast-3:123456789012:inference-profile/apne3.anthropic.claude-3-5-sonnet-20241022-v2:0",
-				)
-
-				// Verify the model info is defined
-				expect(modelInfo.info).toBeDefined()
+				),
+					// Verify the model info is defined
+					expect(modelInfo.info).toBeDefined()
 
 				// Verify parseArn was called with the correct ARN
 				expect(parseArnMock).toHaveBeenCalledWith(
@@ -102,177 +100,21 @@ describe("AwsBedrockHandler", () => {
 			}
 		})
 
-		it("should use default model when custom-arn is selected but no ARN is provided", () => {
+		it("should use default prompt router model when prompt router arn is entered but no model can be identified from the ARN", () => {
 			const customArnHandler = new AwsBedrockHandler({
-				apiModelId: "custom-arn",
+				awsCustomArn:
+					"arn:aws:bedrock:ap-northeast-3:123456789012:default-prompt-router/my_router_arn_no_model",
 				awsAccessKey: "test-access-key",
 				awsSecretKey: "test-secret-key",
 				awsRegion: "us-east-1",
-				// No awsCustomArn provided
 			})
 			const modelInfo = customArnHandler.getModel()
-			// Should fall back to default model
-			expect(modelInfo.id).not.toBe("custom-arn")
+			// Should fall back to default prompt router model
+			expect(modelInfo.id).toBe(
+				"arn:aws:bedrock:ap-northeast-3:123456789012:default-prompt-router/my_router_arn_no_model",
+			) // bedrockDefaultPromptRouterModelId
 			expect(modelInfo.info).toBeDefined()
-		})
-	})
-
-	describe("invokedModelId handling", () => {
-		it("should update costModelConfig when invokedModelId is present in custom ARN scenario", async () => {
-			const customArnHandler = new AwsBedrockHandler({
-				apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
-				awsAccessKey: "test-access-key",
-				awsSecretKey: "test-secret-key",
-				awsRegion: "us-east-1",
-				awsCustomArn: "arn:aws:bedrock:us-east-1:123456789:foundation-model/custom-model",
-			})
-
-			const mockStreamEvent = {
-				trace: {
-					promptRouter: {
-						invokedModelId: "arn:aws:bedrock:us-east-1:123456789:foundation-model/custom-model:0",
-					},
-				},
-			}
-
-			jest.spyOn(customArnHandler, "getModel").mockReturnValue({
-				id: "custom-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-
-			await customArnHandler.createMessage("system prompt", [{ role: "user", content: "user message" }]).next()
-
-			expect(customArnHandler.getModel()).toEqual({
-				id: "custom-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-		})
-
-		it("should update costModelConfig when invokedModelId is present in default model scenario", async () => {
-			handler = new AwsBedrockHandler({
-				apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
-				awsAccessKey: "test-access-key",
-				awsSecretKey: "test-secret-key",
-				awsRegion: "us-east-1",
-			})
-
-			const mockStreamEvent = {
-				trace: {
-					promptRouter: {
-						invokedModelId: "arn:aws:bedrock:us-east-1:123456789:foundation-model/default-model:0",
-					},
-				},
-			}
-
-			jest.spyOn(handler, "getModel").mockReturnValue({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-
-			await handler.createMessage("system prompt", [{ role: "user", content: "user message" }]).next()
-
-			expect(handler.getModel()).toEqual({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-		})
-
-		it("should not update costModelConfig when invokedModelId is not present", async () => {
-			handler = new AwsBedrockHandler({
-				apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
-				awsAccessKey: "test-access-key",
-				awsSecretKey: "test-secret-key",
-				awsRegion: "us-east-1",
-			})
-
-			const mockStreamEvent = {
-				trace: {
-					promptRouter: {
-						// No invokedModelId present
-					},
-				},
-			}
-
-			jest.spyOn(handler, "getModel").mockReturnValue({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-
-			await handler.createMessage("system prompt", [{ role: "user", content: "user message" }]).next()
-
-			expect(handler.getModel()).toEqual({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-		})
-
-		it("should not update costModelConfig when invokedModelId cannot be parsed", async () => {
-			handler = new AwsBedrockHandler({
-				apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
-				awsAccessKey: "test-access-key",
-				awsSecretKey: "test-secret-key",
-				awsRegion: "us-east-1",
-			})
-
-			const mockStreamEvent = {
-				trace: {
-					promptRouter: {
-						invokedModelId: "invalid-arn",
-					},
-				},
-			}
-
-			jest.spyOn(handler, "getModel").mockReturnValue({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
-
-			await handler.createMessage("system prompt", [{ role: "user", content: "user message" }]).next()
-
-			expect(handler.getModel()).toEqual({
-				id: "default-model",
-				info: {
-					maxTokens: 4096,
-					contextWindow: 128_000,
-					supportsPromptCache: false,
-					supportsImages: true,
-				},
-			})
+			expect(modelInfo.info.maxTokens).toBe(4096)
 		})
 	})
 })

src/api/providers/bedrock.ts

@@ -296,17 +296,15 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 
 				if (streamEvent?.trace?.promptRouter?.invokedModelId) {
 					try {
-						let invokedModelArn = this.parseArn(streamEvent.trace.promptRouter.invokedModelId)
-						if (invokedModelArn?.modelId) {
-							//update the in-use model info to be based on the invoked Model Id for the router
-							//so that pricing, context window, caching etc have values that can be used
-							//However, we want to keep the id of the model to be the ID for the router for
-							//subsequent requests so they are sent back through the router
-							let invokedModel = this.getModelById(invokedModelArn.modelId as string)
-							if (invokedModel) {
-								invokedModel.id = modelConfig.id
-								this.costModelConfig = invokedModel
-							}
+						//update the in-use model info to be based on the invoked Model Id for the router
+						//so that pricing, context window, caching etc have values that can be used
+						//However, we want to keep the id of the model to be the ID for the router for
+						//subsequent requests so they are sent back through the router
+						let invokedArnInfo = this.parseArn(streamEvent.trace.promptRouter.invokedModelId)
+						let invokedModel = this.getModelById(invokedArnInfo.modelId as string, invokedArnInfo.modelType)
+						if (invokedModel) {
+							invokedModel.id = modelConfig.id
+							this.costModelConfig = invokedModel
 						}
 
 						// Handle metadata events for the promptRouter.
@@ -626,26 +624,28 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 	}
 
 	//Prompt Router responses come back in a different sequence and the model used is in the response and must be fetched by name
-	getModelById(modelId: string): { id: BedrockModelId | string; info: SharedModelInfo } {
+	getModelById(modelId: string, modelType?: string): { id: BedrockModelId | string; info: SharedModelInfo } {
 		// Try to find the model in bedrockModels
 		let baseModelId = this.parseBaseModelId(modelId)
+		const id = baseModelId as BedrockModelId
+		let model
 		if (baseModelId in bedrockModels) {
-			const id = baseModelId as BedrockModelId
-
 			//Do a deep copy of the model info so that later in the code the model id and maxTokens can be set.
 			// The bedrockModels array is a constant and updating the model ID from the returned invokedModelID value
 			// in a prompt router response isn't possible on the constant.
-			let model = JSON.parse(JSON.stringify(bedrockModels[id]))
-
-			// If modelMaxTokens is explicitly set in options, override the default
-			if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) {
-				model.maxTokens = this.options.modelMaxTokens
-			}
+			model = { id: id, info: JSON.parse(JSON.stringify(bedrockModels[id])) }
+		} else if (modelType && modelType.includes("router")) {
+			model = this.getModelById(bedrockDefaultPromptRouterModelId as string)
+		} else {
+			model = this.getModelById(bedrockDefaultModelId as string)
+		}
 
-			return { id, info: model }
+		// If modelMaxTokens is explicitly set in options, override the default
+		if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) {
+			model.info.maxTokens = this.options.modelMaxTokens
 		}
 
-		return { id: bedrockDefaultModelId, info: bedrockModels[bedrockDefaultModelId] }
+		return model
 	}
 
 	override getModel(): { id: BedrockModelId | string; info: SharedModelInfo } {
@@ -657,12 +657,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
 
 		// If custom ARN is provided, use it
 		if (this.options.awsCustomArn) {
-			modelConfig = this.getModelById(this.arnInfo.modelId)
-
-			if (!modelConfig)
-				// An ARN was used, but no model info match found, use default model values for cost calculations and context window
-				// But continue using the ARN as the identifier in the Bedrock interaction
-				modelConfig = this.getModelById(bedrockDefaultPromptRouterModelId)
+			modelConfig = this.getModelById(this.arnInfo.modelId, this.arnInfo.modelType)
 
 			//If the user entered an ARN for a foundation-model they've done the same thing as picking from our list of options.
 			//We leave the model data matching the same as if a drop-down input method was used by not overwriting the model ID with the user input ARN
@@ -864,7 +859,11 @@ Please verify:
 			messageTemplate: `Request was throttled or rate limited. Please try:
 1. Reducing the frequency of requests
 2. If using a provisioned model, check its throughput settings
-3. Contact AWS support to request a quota increase if needed`,
+3. Contact AWS support to request a quota increase if needed
+
+{formattedErrorDetails}
+
+`,
 			logLevel: "error",
 		},
 		TOO_MANY_TOKENS: {