From 8540bb9403d0a36f295f4bebfaf45bd9b3d15e70 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Mon, 28 Jul 2025 05:33:18 +0000 Subject: [PATCH] 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 | 11 + src/api/providers/__tests__/openai.spec.ts | 305 +++++++++++++++--- src/api/providers/openai.ts | 74 +++++ .../settings/providers/OpenAICompatible.tsx | 133 ++++++++ webview-ui/src/i18n/locales/en/settings.json | 28 ++ 5 files changed, 509 insertions(+), 42 deletions(-) diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index 4322b9f550..939add6ccc 100644 --- a/packages/types/src/provider-settings.ts +++ b/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({ diff --git a/src/api/providers/__tests__/openai.spec.ts b/src/api/providers/__tests__/openai.spec.ts index b4b5f29204..918aef2f39 100644 --- a/src/api/providers/__tests__/openai.spec.ts +++ b/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") + }) + }) }) diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index f5e4e4c985..33ed8ec1a9 100644 --- a/src/api/providers/openai.ts +++ b/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) { diff --git a/webview-ui/src/components/settings/providers/OpenAICompatible.tsx b/webview-ui/src/components/settings/providers/OpenAICompatible.tsx index 736b0253c4..2a8d6c2233 100644 --- a/webview-ui/src/components/settings/providers/OpenAICompatible.tsx +++ b/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 | null>(null) @@ -204,6 +205,138 @@ export const OpenAICompatible = ({ )} + {/* Azure AI Search UI */} +
+ { + setAzureAiSearchEnabled(checked) + setApiConfigurationField("azureAiSearchEnabled", checked) + }}> + {t("settings:providers.azureAiSearch.enable")} + +
+ {t("settings:providers.azureAiSearch.enableDescription")} +
+ {azureAiSearchEnabled && ( +
+ + + + + + + + + + + + +
+ + +
+ + + + + + +
+ { + const value = parseInt((e.target as HTMLInputElement).value) + return isNaN(value) ? 5 : value + })} + className="w-full"> + + +
+ {t("settings:providers.azureAiSearch.topNDocumentsDescription")} +
+
+
+ { + const value = parseInt((e.target as HTMLInputElement).value) + return isNaN(value) ? 3 : value + })} + className="w-full"> + + +
+ {t("settings:providers.azureAiSearch.strictnessDescription")} +
+
+
+ )} +
+ {/* Custom Headers UI */}
diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json index cd83190326..ff3b28e13e 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/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",