feat: add Azure AI Search integration for OpenAI models

- Add Azure AI Search configuration fields to provider settings schema
- Update OpenAICompatible UI component with Azure AI Search options
- Add data_sources field to OpenAI API requests when Azure AI Search is enabled
- Add comprehensive tests for Azure AI Search functionality
- Add translation keys for all Azure AI Search UI elements

Implements #6282

packages/types/src/provider-settings.ts

@@ -141,6 +141,17 @@ const openAiSchema = baseProviderSettingsSchema.extend({
 	openAiStreamingEnabled: z.boolean().optional(),
 	openAiHostHeader: z.string().optional(), // Keep temporarily for backward compatibility during migration.
 	openAiHeaders: z.record(z.string(), z.string()).optional(),
+	// Azure AI Search fields
+	azureAiSearchEnabled: z.boolean().optional(),
+	azureAiSearchEndpoint: z.string().optional(),
+	azureAiSearchIndexName: z.string().optional(),
+	azureAiSearchApiKey: z.string().optional(),
+	azureAiSearchSemanticConfiguration: z.string().optional(),
+	azureAiSearchQueryType: z.string().optional(),
+	azureAiSearchEmbeddingEndpoint: z.string().optional(),
+	azureAiSearchEmbeddingApiKey: z.string().optional(),
+	azureAiSearchTopNDocuments: z.number().optional(),
+	azureAiSearchStrictness: z.number().optional(),
 })
 
 const ollamaSchema = baseProviderSettingsSchema.extend({

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

@@ -11,19 +11,43 @@ const mockCreate = vitest.fn()
 
 vitest.mock("openai", () => {
 	const mockConstructor = vitest.fn()
-	return {
-		__esModule: true,
-		default: mockConstructor.mockImplementation(() => ({
-			chat: {
-				completions: {
-					create: mockCreate.mockImplementation(async (options) => {
-						if (!options.stream) {
-							return {
-								id: "test-completion",
+	const mockImplementation = () => ({
+		chat: {
+			completions: {
+				create: mockCreate.mockImplementation(async (options) => {
+					if (!options.stream) {
+						return {
+							id: "test-completion",
+							choices: [
+								{
+									message: { role: "assistant", content: "Test response", refusal: null },
+									finish_reason: "stop",
+									index: 0,
+								},
+							],
+							usage: {
+								prompt_tokens: 10,
+								completion_tokens: 5,
+								total_tokens: 15,
+							},
+						}
+					}
+
+					return {
+						[Symbol.asyncIterator]: async function* () {
+							yield {
 								choices: [
 									{
-										message: { role: "assistant", content: "Test response", refusal: null },
-										finish_reason: "stop",
+										delta: { content: "Test response" },
+										index: 0,
+									},
+								],
+								usage: null,
+							}
+							yield {
+								choices: [
+									{
+										delta: {},
 										index: 0,
 									},
 								],
@@ -33,38 +57,16 @@ vitest.mock("openai", () => {
 									total_tokens: 15,
 								},
 							}
-						}
-
-						return {
-							[Symbol.asyncIterator]: async function* () {
-								yield {
-									choices: [
-										{
-											delta: { content: "Test response" },
-											index: 0,
-										},
-									],
-									usage: null,
-								}
-								yield {
-									choices: [
-										{
-											delta: {},
-											index: 0,
-										},
-									],
-									usage: {
-										prompt_tokens: 10,
-										completion_tokens: 5,
-										total_tokens: 15,
-									},
-								}
-							},
-						}
-					}),
-				},
+						},
+					}
+				}),
 			},
-		})),
+		},
+	})
+	return {
+		__esModule: true,
+		default: mockConstructor.mockImplementation(mockImplementation),
+		AzureOpenAI: mockConstructor.mockImplementation(mockImplementation),
 	}
 })
 
@@ -775,4 +777,223 @@ describe("OpenAiHandler", () => {
 			)
 		})
 	})
+
+	describe("Azure AI Search", () => {
+		const azureSearchOptions = {
+			...mockOptions,
+			openAiUseAzure: true,
+			azureAiSearchEnabled: true,
+			azureAiSearchEndpoint: "https://test-search.search.windows.net/",
+			azureAiSearchIndexName: "test-index",
+			azureAiSearchApiKey: "test-search-api-key",
+			azureAiSearchSemanticConfiguration: "azureml-default",
+			azureAiSearchQueryType: "vector_simple_hybrid",
+			azureAiSearchEmbeddingEndpoint:
+				"https://test-embedding.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2023-07-01-preview",
+			azureAiSearchEmbeddingApiKey: "test-embedding-api-key",
+			azureAiSearchTopNDocuments: 5,
+			azureAiSearchStrictness: 3,
+		}
+
+		it("should include data_sources when Azure AI Search is enabled", async () => {
+			const azureSearchHandler = new OpenAiHandler(azureSearchOptions)
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = azureSearchHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).toHaveProperty("data_sources")
+			expect(callArgs.data_sources).toHaveLength(1)
+
+			const dataSource = callArgs.data_sources[0]
+			expect(dataSource.type).toBe("azure_search")
+			expect(dataSource.parameters).toMatchObject({
+				endpoint: azureSearchOptions.azureAiSearchEndpoint,
+				index_name: azureSearchOptions.azureAiSearchIndexName,
+				semantic_configuration: azureSearchOptions.azureAiSearchSemanticConfiguration,
+				query_type: azureSearchOptions.azureAiSearchQueryType,
+				in_scope: true,
+				role_information: "You are an AI assistant that helps people find information.",
+				strictness: azureSearchOptions.azureAiSearchStrictness,
+				top_n_documents: azureSearchOptions.azureAiSearchTopNDocuments,
+				authentication: {
+					type: "api_key",
+					key: azureSearchOptions.azureAiSearchApiKey,
+				},
+				embedding_dependency: {
+					type: "endpoint",
+					endpoint: azureSearchOptions.azureAiSearchEmbeddingEndpoint,
+					authentication: {
+						type: "api_key",
+						key: azureSearchOptions.azureAiSearchEmbeddingApiKey,
+					},
+				},
+				fields_mapping: {
+					content_fields: ["content"],
+					filepath_field: "filepath",
+					title_field: "title",
+					url_field: "url",
+					content_fields_separator: "\n",
+					vector_fields: ["contentVector"],
+				},
+			})
+		})
+
+		it("should not include data_sources when Azure AI Search is disabled", async () => {
+			const noSearchHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				azureAiSearchEnabled: false,
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = noSearchHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).not.toHaveProperty("data_sources")
+		})
+
+		it("should not include data_sources when not using Azure OpenAI", async () => {
+			const nonAzureHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				openAiUseAzure: false,
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = nonAzureHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).not.toHaveProperty("data_sources")
+		})
+
+		it("should handle Azure AI Search without embedding configuration", async () => {
+			const searchWithoutEmbeddingHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				azureAiSearchEmbeddingEndpoint: undefined,
+				azureAiSearchEmbeddingApiKey: undefined,
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = searchWithoutEmbeddingHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).toHaveProperty("data_sources")
+
+			const dataSource = callArgs.data_sources[0]
+			expect(dataSource.parameters).not.toHaveProperty("embedding_dependency")
+		})
+
+		it("should not include fields_mapping for non-vector query types", async () => {
+			const simpleSearchHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				azureAiSearchQueryType: "simple",
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = simpleSearchHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).toHaveProperty("data_sources")
+
+			const dataSource = callArgs.data_sources[0]
+			expect(dataSource.parameters).not.toHaveProperty("fields_mapping")
+		})
+
+		it("should include data_sources in non-streaming mode", async () => {
+			const nonStreamingHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				openAiStreamingEnabled: false,
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = nonStreamingHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).toHaveProperty("data_sources")
+			expect(callArgs.data_sources).toHaveLength(1)
+			expect(callArgs.data_sources[0].type).toBe("azure_search")
+		})
+
+		it("should not include data_sources when endpoint or index name is missing", async () => {
+			const incompleteHandler = new OpenAiHandler({
+				...azureSearchOptions,
+				azureAiSearchEndpoint: undefined,
+			})
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello!",
+				},
+			]
+
+			const stream = incompleteHandler.createMessage(systemPrompt, messages)
+			// Consume the stream to trigger the API call
+			for await (const _chunk of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][0]
+			expect(callArgs).not.toHaveProperty("data_sources")
+		})
+	})
 })

src/api/providers/openai.ts

@@ -158,6 +158,14 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 				...(reasoning && reasoning),
 			}
 
+			// Add Azure AI Search data sources if enabled
+			if (this.options.azureAiSearchEnabled && this.options.openAiUseAzure) {
+				const dataSources = this.buildAzureAiSearchDataSources()
+				if (dataSources) {
+					;(requestOptions as any).data_sources = dataSources
+				}
+			}
+
 			// Add max_tokens if needed
 			this.addMaxTokensIfNeeded(requestOptions, modelInfo)
 
@@ -223,6 +231,14 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 			// Add max_tokens if needed
 			this.addMaxTokensIfNeeded(requestOptions, modelInfo)
 
+			// Add Azure AI Search data sources if enabled
+			if (this.options.azureAiSearchEnabled && this.options.openAiUseAzure) {
+				const dataSources = this.buildAzureAiSearchDataSources()
+				if (dataSources) {
+					;(requestOptions as any).data_sources = dataSources
+				}
+			}
+
 			const response = await this.client.chat.completions.create(
 				requestOptions,
 				this._isAzureAiInference(modelUrl) ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
@@ -408,6 +424,64 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 			requestOptions.max_completion_tokens = this.options.modelMaxTokens || modelInfo.maxTokens
 		}
 	}
+
+	private buildAzureAiSearchDataSources(): any[] | null {
+		if (!this.options.azureAiSearchEndpoint || !this.options.azureAiSearchIndexName) {
+			return null
+		}
+
+		const dataSource: any = {
+			type: "azure_search",
+			parameters: {
+				filter: null,
+				endpoint: this.options.azureAiSearchEndpoint,
+				index_name: this.options.azureAiSearchIndexName,
+				semantic_configuration: this.options.azureAiSearchSemanticConfiguration || "azureml-default",
+				query_type: this.options.azureAiSearchQueryType || "vector_simple_hybrid",
+				in_scope: true,
+				role_information: "You are an AI assistant that helps people find information.",
+				strictness: this.options.azureAiSearchStrictness || 3,
+				top_n_documents: this.options.azureAiSearchTopNDocuments || 5,
+			},
+		}
+
+		// Add authentication if API key is provided
+		if (this.options.azureAiSearchApiKey) {
+			dataSource.parameters.authentication = {
+				type: "api_key",
+				key: this.options.azureAiSearchApiKey,
+			}
+		}
+
+		// Add embedding dependency if configured
+		if (this.options.azureAiSearchEmbeddingEndpoint) {
+			dataSource.parameters.embedding_dependency = {
+				type: "endpoint",
+				endpoint: this.options.azureAiSearchEmbeddingEndpoint,
+			}
+
+			if (this.options.azureAiSearchEmbeddingApiKey) {
+				dataSource.parameters.embedding_dependency.authentication = {
+					type: "api_key",
+					key: this.options.azureAiSearchEmbeddingApiKey,
+				}
+			}
+		}
+
+		// Add fields mapping for vector search
+		if (this.options.azureAiSearchQueryType?.includes("vector")) {
+			dataSource.parameters.fields_mapping = {
+				content_fields: ["content"],
+				filepath_field: "filepath",
+				title_field: "title",
+				url_field: "url",
+				content_fields_separator: "\n",
+				vector_fields: ["contentVector"],
+			}
+		}
+
+		return [dataSource]
+	}
 }
 
 export async function getOpenAiModels(baseUrl?: string, apiKey?: string, openAiHeaders?: Record<string, string>) {

webview-ui/src/components/settings/providers/OpenAICompatible.tsx

@@ -40,6 +40,7 @@ export const OpenAICompatible = ({
 
 	const [azureApiVersionSelected, setAzureApiVersionSelected] = useState(!!apiConfiguration?.azureApiVersion)
 	const [openAiLegacyFormatSelected, setOpenAiLegacyFormatSelected] = useState(!!apiConfiguration?.openAiLegacyFormat)
+	const [azureAiSearchEnabled, setAzureAiSearchEnabled] = useState(!!apiConfiguration?.azureAiSearchEnabled)
 
 	const [openAiModels, setOpenAiModels] = useState<Record<string, ModelInfo> | null>(null)
 
@@ -204,6 +205,138 @@ export const OpenAICompatible = ({
 				)}
 			</div>
 
+			{/* Azure AI Search UI */}
+			<div>
+				<Checkbox
+					checked={azureAiSearchEnabled}
+					onChange={(checked: boolean) => {
+						setAzureAiSearchEnabled(checked)
+						setApiConfigurationField("azureAiSearchEnabled", checked)
+					}}>
+					{t("settings:providers.azureAiSearch.enable")}
+				</Checkbox>
+				<div className="text-sm text-vscode-descriptionForeground ml-6">
+					{t("settings:providers.azureAiSearch.enableDescription")}
+				</div>
+				{azureAiSearchEnabled && (
+					<div className="ml-6 mt-2 space-y-3">
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchEndpoint || ""}
+							onInput={handleInputChange("azureAiSearchEndpoint")}
+							placeholder={t("settings:providers.azureAiSearch.endpointPlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.endpoint")}
+							</label>
+						</VSCodeTextField>
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchIndexName || ""}
+							onInput={handleInputChange("azureAiSearchIndexName")}
+							placeholder={t("settings:providers.azureAiSearch.indexNamePlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.indexName")}
+							</label>
+						</VSCodeTextField>
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchApiKey || ""}
+							type="password"
+							onInput={handleInputChange("azureAiSearchApiKey")}
+							placeholder={t("settings:providers.azureAiSearch.apiKeyPlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.apiKey")}
+							</label>
+						</VSCodeTextField>
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchSemanticConfiguration || "azureml-default"}
+							onInput={handleInputChange("azureAiSearchSemanticConfiguration")}
+							placeholder={t("settings:providers.azureAiSearch.semanticConfigurationPlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.semanticConfiguration")}
+							</label>
+						</VSCodeTextField>
+						<div>
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.queryType")}
+							</label>
+							<select
+								value={apiConfiguration?.azureAiSearchQueryType || "vector_simple_hybrid"}
+								onChange={(e) => setApiConfigurationField("azureAiSearchQueryType", e.target.value)}
+								className="w-full p-2 bg-vscode-input-background text-vscode-input-foreground border border-vscode-input-border rounded">
+								<option value="simple">
+									{t("settings:providers.azureAiSearch.queryTypeOptions.simple")}
+								</option>
+								<option value="semantic">
+									{t("settings:providers.azureAiSearch.queryTypeOptions.semantic")}
+								</option>
+								<option value="vector">
+									{t("settings:providers.azureAiSearch.queryTypeOptions.vector")}
+								</option>
+								<option value="vector_simple_hybrid">
+									{t("settings:providers.azureAiSearch.queryTypeOptions.vectorSimpleHybrid")}
+								</option>
+								<option value="vector_semantic_hybrid">
+									{t("settings:providers.azureAiSearch.queryTypeOptions.vectorSemanticHybrid")}
+								</option>
+							</select>
+						</div>
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchEmbeddingEndpoint || ""}
+							onInput={handleInputChange("azureAiSearchEmbeddingEndpoint")}
+							placeholder={t("settings:providers.azureAiSearch.embeddingEndpointPlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.embeddingEndpoint")}
+							</label>
+						</VSCodeTextField>
+						<VSCodeTextField
+							value={apiConfiguration?.azureAiSearchEmbeddingApiKey || ""}
+							type="password"
+							onInput={handleInputChange("azureAiSearchEmbeddingApiKey")}
+							placeholder={t("settings:providers.azureAiSearch.embeddingApiKeyPlaceholder")}
+							className="w-full">
+							<label className="block font-medium mb-1">
+								{t("settings:providers.azureAiSearch.embeddingApiKey")}
+							</label>
+						</VSCodeTextField>
+						<div>
+							<VSCodeTextField
+								value={apiConfiguration?.azureAiSearchTopNDocuments?.toString() || "5"}
+								onInput={handleInputChange("azureAiSearchTopNDocuments", (e) => {
+									const value = parseInt((e.target as HTMLInputElement).value)
+									return isNaN(value) ? 5 : value
+								})}
+								className="w-full">
+								<label className="block font-medium mb-1">
+									{t("settings:providers.azureAiSearch.topNDocuments")}
+								</label>
+							</VSCodeTextField>
+							<div className="text-sm text-vscode-descriptionForeground">
+								{t("settings:providers.azureAiSearch.topNDocumentsDescription")}
+							</div>
+						</div>
+						<div>
+							<VSCodeTextField
+								value={apiConfiguration?.azureAiSearchStrictness?.toString() || "3"}
+								onInput={handleInputChange("azureAiSearchStrictness", (e) => {
+									const value = parseInt((e.target as HTMLInputElement).value)
+									return isNaN(value) ? 3 : value
+								})}
+								className="w-full">
+								<label className="block font-medium mb-1">
+									{t("settings:providers.azureAiSearch.strictness")}
+								</label>
+							</VSCodeTextField>
+							<div className="text-sm text-vscode-descriptionForeground">
+								{t("settings:providers.azureAiSearch.strictnessDescription")}
+							</div>
+						</div>
+					</div>
+				)}
+			</div>
+
 			{/* Custom Headers UI */}
 			<div className="mb-4">
 				<div className="flex justify-between items-center mb-2">

webview-ui/src/i18n/locales/en/settings.json

@@ -236,6 +236,34 @@
 		"headerName": "Header name",
 		"headerValue": "Header value",
 		"noCustomHeaders": "No custom headers defined. Click the + button to add one.",
+		"azureAiSearch": {
+			"enable": "Enable Azure AI Search",
+			"enableDescription": "Use Azure AI Search to provide context from your internal documentation and knowledge base",
+			"endpoint": "Azure AI Search Endpoint",
+			"endpointPlaceholder": "https://your-resource.search.windows.net/",
+			"indexName": "Index Name",
+			"indexNamePlaceholder": "Enter your index name",
+			"apiKey": "Azure AI Search API Key",
+			"apiKeyPlaceholder": "Enter your Azure AI Search API key",
+			"semanticConfiguration": "Semantic Configuration",
+			"semanticConfigurationPlaceholder": "azureml-default",
+			"queryType": "Query Type",
+			"queryTypeOptions": {
+				"simple": "Simple",
+				"semantic": "Semantic",
+				"vector": "Vector",
+				"vectorSimpleHybrid": "Vector + Simple Hybrid",
+				"vectorSemanticHybrid": "Vector + Semantic Hybrid"
+			},
+			"embeddingEndpoint": "Embedding Model Endpoint",
+			"embeddingEndpointPlaceholder": "https://your-resource.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2023-07-01-preview",
+			"embeddingApiKey": "Embedding Model API Key",
+			"embeddingApiKeyPlaceholder": "Enter your embedding model API key",
+			"topNDocuments": "Top N Documents",
+			"topNDocumentsDescription": "Number of documents to retrieve from the search index",
+			"strictness": "Search Strictness",
+			"strictnessDescription": "Controls how strictly the search results must match (1-5, higher is stricter)"
+		},
 		"requestyApiKey": "Requesty API Key",
 		"refreshModels": {
 			"label": "Refresh Models",