Add URL base check for Databricks instance on AWS and GCP

Author: Duc Nguyen

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

@@ -394,43 +394,53 @@ describe("OpenAiHandler", () => {
 	})
 
 	describe("Databricks AI Provider", () => {
-		const databricksOptions = {
+		const baseDatabricksOptions = {
 			...mockOptions,
-			openAiBaseUrl: "https://adb-xxxx.azuredatabricks.net/serving-endpoints",
 			openAiModelId: "databricks-dbrx-instruct",
 		}
 
-		it("should initialize with Databricks AI configuration", () => {
-			const databricksHandler = new OpenAiHandler(databricksOptions)
+		const databricksUrls = [
+			"https://adb-xxxx.azuredatabricks.net/serving-endpoints",
+			"https://myworkspace.cloud.databricks.com/serving-endpoints/myendpoint",
+			"https://anotherworkspace.gcp.databricks.com/serving-endpoints/anotherendpoint",
+		]
+
+		it.each(databricksUrls)("should initialize with Databricks AI configuration for %s", (url) => {
+			const options = { ...baseDatabricksOptions, openAiBaseUrl: url }
+			const databricksHandler = new OpenAiHandler(options)
 			expect(databricksHandler).toBeInstanceOf(OpenAiHandler)
-			expect(databricksHandler.getModel().id).toBe(databricksOptions.openAiModelId)
+			expect(databricksHandler.getModel().id).toBe(options.openAiModelId)
 		})
 
-		it("should exclude stream_options when streaming with Databricks AI", async () => {
-			const databricksHandler = new OpenAiHandler(databricksOptions)
-			const systemPrompt = "You are a helpful assistant."
-			const messages: Anthropic.Messages.MessageParam[] = [
-				{
-					role: "user",
-					content: "Hello!",
-				},
-			]
+		it.each(databricksUrls)(
+			"should exclude stream_options when streaming with Databricks AI for %s",
+			async (url) => {
+				const options = { ...baseDatabricksOptions, openAiBaseUrl: url }
+				const databricksHandler = new OpenAiHandler(options)
+				const systemPrompt = "You are a helpful assistant."
+				const messages: Anthropic.Messages.MessageParam[] = [
+					{
+						role: "user",
+						content: "Hello!",
+					},
+				]
 
-			const stream = databricksHandler.createMessage(systemPrompt, messages)
-			await stream.next() // Consume one item to trigger the mock call
+				const stream = databricksHandler.createMessage(systemPrompt, messages)
+				await stream.next() // Consume one item to trigger the mock call
 
-			expect(mockCreate).toHaveBeenCalledWith(
-				expect.objectContaining({
-					model: databricksOptions.openAiModelId,
-					stream: true,
-				}),
-				{}, // Expecting empty options object as second argument
-			)
+				expect(mockCreate).toHaveBeenCalledWith(
+					expect.objectContaining({
+						model: options.openAiModelId,
+						stream: true,
+					}),
+					{}, // Expecting empty options object as second argument
+				)
 
-			// Verify stream_options is not present in the last call's arguments
-			const mockCalls = mockCreate.mock.calls
-			const lastCallArgs = mockCalls[mockCalls.length - 1][0]
-			expect(lastCallArgs).not.toHaveProperty("stream_options")
-		})
+				// Verify stream_options is not present in the last call's arguments
+				const mockCalls = mockCreate.mock.calls
+				const lastCallArgs = mockCalls[mockCalls.length - 1][0]
+				expect(lastCallArgs).not.toHaveProperty("stream_options")
+			},
+		)
 	})
 })

src/api/providers/openai.ts

@@ -350,7 +350,11 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 
 	private _isDatabricksAI(baseUrl?: string): boolean {
 		const urlHost = this._getUrlHost(baseUrl)
-		return urlHost.includes(".azuredatabricks.net")
+		return (
+			urlHost.endsWith(".azuredatabricks.net") ||
+			urlHost.endsWith(".cloud.databricks.com") ||
+			urlHost.endsWith(".gcp.databricks.com")
+		)
 	}
 
 	private _isAzureAiInference(baseUrl?: string): boolean {