feat(azure-openai): allow usage of azure-openai for knowledgebase uploads and wand generation (#1056)

* feat(azure-openai): allow usage of azure-openai for knowledgebase uploads

* feat(azure-openai): added azure-openai for kb and wand

* added embeddings utils, added the ability to use mistral through Azure

* fix(oauth): gdrive picker race condition, token route cleanup

* fix test

* feat(mailer): consolidated all emailing to mailer service, added support for Azure ACS (#1054)

* feat(mailer): consolidated all emailing to mailer service, added support for Azure ACS

* fix batch invitation email template

* cleanup

* improvement(emails): add help template instead of doing it inline

* remove fallback version

---------

Co-authored-by: Vikhyath Mondreti <vikhyath@simstudio.ai>

Author: waleedlatif1

apps/sim/app/api/knowledge/search/utils.test.ts

@@ -4,15 +4,50 @@
  *
  * @vitest-environment node
  */
-import { describe, expect, it, vi } from 'vitest'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 vi.mock('drizzle-orm')
-vi.mock('@/lib/logs/console/logger')
+vi.mock('@/lib/logs/console/logger', () => ({
+  createLogger: vi.fn(() => ({
+    info: vi.fn(),
+    debug: vi.fn(),
+    warn: vi.fn(),
+    error: vi.fn(),
+  })),
+}))
 vi.mock('@/db')
+vi.mock('@/lib/documents/utils', () => ({
+  retryWithExponentialBackoff: (fn: any) => fn(),
+}))
 
-import { handleTagAndVectorSearch, handleTagOnlySearch, handleVectorOnlySearch } from './utils'
+vi.stubGlobal(
+  'fetch',
+  vi.fn().mockResolvedValue({
+    ok: true,
+    json: async () => ({
+      data: [{ embedding: [0.1, 0.2, 0.3] }],
+    }),
+  })
+)
+
+vi.mock('@/lib/env', () => ({
+  env: {},
+  isTruthy: (value: string | boolean | number | undefined) =>
+    typeof value === 'string' ? value === 'true' || value === '1' : Boolean(value),
+}))
+
+import {
+  generateSearchEmbedding,
+  handleTagAndVectorSearch,
+  handleTagOnlySearch,
+  handleVectorOnlySearch,
+} from './utils'
 
 describe('Knowledge Search Utils', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
   describe('handleTagOnlySearch', () => {
     it('should throw error when no filters provided', async () => {
       const params = {
@@ -140,4 +175,251 @@ describe('Knowledge Search Utils', () => {
       expect(params.distanceThreshold).toBe(0.8)
     })
   })
+
+  describe('generateSearchEmbedding', () => {
+    it('should use Azure OpenAI when KB-specific config is provided', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      const result = await generateSearchEmbedding('test query')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        'https://test.openai.azure.com/openai/deployments/text-embedding-ada-002/embeddings?api-version=2024-12-01-preview',
+        expect.objectContaining({
+          headers: expect.objectContaining({
+            'api-key': 'test-azure-key',
+          }),
+        })
+      )
+      expect(result).toEqual([0.1, 0.2, 0.3])
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should fallback to OpenAI when no KB Azure config provided', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      const result = await generateSearchEmbedding('test query')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        'https://api.openai.com/v1/embeddings',
+        expect.objectContaining({
+          headers: expect.objectContaining({
+            Authorization: 'Bearer test-openai-key',
+          }),
+        })
+      )
+      expect(result).toEqual([0.1, 0.2, 0.3])
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should use default API version when not provided in Azure config', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        KB_OPENAI_MODEL_NAME: 'custom-embedding-model',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      await generateSearchEmbedding('test query')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        expect.stringContaining('api-version='),
+        expect.any(Object)
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should use custom model name when provided in Azure config', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'custom-embedding-model',
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      await generateSearchEmbedding('test query', 'text-embedding-3-small')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        'https://test.openai.azure.com/openai/deployments/custom-embedding-model/embeddings?api-version=2024-12-01-preview',
+        expect.any(Object)
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should throw error when no API configuration provided', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow(
+        'Either OPENAI_API_KEY or Azure OpenAI configuration (AZURE_OPENAI_API_KEY + AZURE_OPENAI_ENDPOINT) must be configured'
+      )
+    })
+
+    it('should handle Azure OpenAI API errors properly', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: false,
+        status: 404,
+        statusText: 'Not Found',
+        text: async () => 'Deployment not found',
+      } as any)
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should handle OpenAI API errors properly', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: false,
+        status: 429,
+        statusText: 'Too Many Requests',
+        text: async () => 'Rate limit exceeded',
+      } as any)
+
+      await expect(generateSearchEmbedding('test query')).rejects.toThrow('Embedding API failed')
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should include correct request body for Azure OpenAI', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        AZURE_OPENAI_API_KEY: 'test-azure-key',
+        AZURE_OPENAI_ENDPOINT: 'https://test.openai.azure.com',
+        AZURE_OPENAI_API_VERSION: '2024-12-01-preview',
+        KB_OPENAI_MODEL_NAME: 'text-embedding-ada-002',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      await generateSearchEmbedding('test query')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          body: JSON.stringify({
+            input: ['test query'],
+            encoding_format: 'float',
+          }),
+        })
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+
+    it('should include correct request body for OpenAI', async () => {
+      const { env } = await import('@/lib/env')
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+      Object.assign(env, {
+        OPENAI_API_KEY: 'test-openai-key',
+      })
+
+      const fetchSpy = vi.mocked(fetch)
+      fetchSpy.mockResolvedValueOnce({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: [0.1, 0.2, 0.3] }],
+        }),
+      } as any)
+
+      await generateSearchEmbedding('test query', 'text-embedding-3-small')
+
+      expect(fetchSpy).toHaveBeenCalledWith(
+        expect.any(String),
+        expect.objectContaining({
+          body: JSON.stringify({
+            input: ['test query'],
+            model: 'text-embedding-3-small',
+            encoding_format: 'float',
+          }),
+        })
+      )
+
+      // Clean up
+      Object.keys(env).forEach((key) => delete (env as any)[key])
+    })
+  })
 })

apps/sim/app/api/knowledge/search/utils.ts

@@ -1,22 +1,10 @@
 import { and, eq, inArray, sql } from 'drizzle-orm'
-import { retryWithExponentialBackoff } from '@/lib/documents/utils'
-import { env } from '@/lib/env'
 import { createLogger } from '@/lib/logs/console/logger'
 import { db } from '@/db'
 import { embedding } from '@/db/schema'
 
 const logger = createLogger('KnowledgeSearchUtils')
 
-export class APIError extends Error {
-  public status: number
-
-  constructor(message: string, status: number) {
-    super(message)
-    this.name = 'APIError'
-    this.status = status
-  }
-}
-
 export interface SearchResult {
   id: string
   content: string
@@ -41,61 +29,8 @@ export interface SearchParams {
   distanceThreshold?: number
 }
 
-export async function generateSearchEmbedding(query: string): Promise<number[]> {
-  const openaiApiKey = env.OPENAI_API_KEY
-  if (!openaiApiKey) {
-    throw new Error('OPENAI_API_KEY not configured')
-  }
-
-  try {
-    const embedding = await retryWithExponentialBackoff(
-      async () => {
-        const response = await fetch('https://api.openai.com/v1/embeddings', {
-          method: 'POST',
-          headers: {
-            Authorization: `Bearer ${openaiApiKey}`,
-            'Content-Type': 'application/json',
-          },
-          body: JSON.stringify({
-            input: query,
-            model: 'text-embedding-3-small',
-            encoding_format: 'float',
-          }),
-        })
-
-        if (!response.ok) {
-          const errorText = await response.text()
-          const error = new APIError(
-            `OpenAI API error: ${response.status} ${response.statusText} - ${errorText}`,
-            response.status
-          )
-          throw error
-        }
-
-        const data = await response.json()
-
-        if (!data.data || !Array.isArray(data.data) || data.data.length === 0) {
-          throw new Error('Invalid response format from OpenAI embeddings API')
-        }
-
-        return data.data[0].embedding
-      },
-      {
-        maxRetries: 5,
-        initialDelayMs: 1000,
-        maxDelayMs: 30000,
-        backoffMultiplier: 2,
-      }
-    )
-
-    return embedding
-  } catch (error) {
-    logger.error('Failed to generate search embedding:', error)
-    throw new Error(
-      `Embedding generation failed: ${error instanceof Error ? error.message : 'Unknown error'}`
-    )
-  }
-}
+// Use shared embedding utility
+export { generateSearchEmbedding } from '@/lib/embeddings/utils'
 
 function getTagFilters(filters: Record<string, string>, embedding: any) {
   return Object.entries(filters).map(([key, value]) => {