|
1 | 1 | /** |
2 | 2 | * @vitest-environment node |
3 | 3 | */ |
| 4 | +import { resetEnvMock, setEnv } from '@sim/testing' |
4 | 5 | import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' |
5 | | -import { embed } from '@/lib/embeddings/client' |
| 6 | +import { |
| 7 | + EmbeddingAPIError, |
| 8 | + embed, |
| 9 | + embedKnowledgeForDeployment, |
| 10 | + isTransientEmbeddingError, |
| 11 | +} from '@/lib/embeddings/client' |
| 12 | + |
| 13 | +const { mockGetBYOKKey } = vi.hoisted(() => ({ |
| 14 | + mockGetBYOKKey: vi.fn(), |
| 15 | +})) |
| 16 | + |
| 17 | +vi.mock('@/lib/api-key/byok', () => ({ |
| 18 | + getBYOKKey: mockGetBYOKKey, |
| 19 | +})) |
6 | 20 |
|
7 | 21 | /** |
8 | 22 | * Exercises the orchestrator end-to-end against a mocked transport: batching, |
@@ -35,11 +49,25 @@ let fetchMock: ReturnType<typeof vi.fn> |
35 | 49 | beforeEach(() => { |
36 | 50 | fetchMock = vi.fn() |
37 | 51 | global.fetch = fetchMock as unknown as typeof fetch |
| 52 | + mockGetBYOKKey.mockResolvedValue(null) |
| 53 | + setEnv({ |
| 54 | + AZURE_OPENAI_API_KEY: undefined, |
| 55 | + AZURE_OPENAI_ENDPOINT: undefined, |
| 56 | + AZURE_OPENAI_API_VERSION: undefined, |
| 57 | + GEMINI_API_KEY: undefined, |
| 58 | + OPENAI_API_KEY: undefined, |
| 59 | + OPENAI_API_KEY_1: undefined, |
| 60 | + OPENAI_API_KEY_2: undefined, |
| 61 | + OPENAI_API_KEY_3: undefined, |
| 62 | + OPENROUTER_API_KEY: undefined, |
| 63 | + }) |
38 | 64 | }) |
39 | 65 |
|
40 | 66 | afterEach(() => { |
41 | 67 | global.fetch = originalFetch |
| 68 | + vi.useRealTimers() |
42 | 69 | vi.restoreAllMocks() |
| 70 | + resetEnvMock() |
43 | 71 | }) |
44 | 72 |
|
45 | 73 | describe('embed', () => { |
@@ -384,3 +412,147 @@ describe('embed', () => { |
384 | 412 | }) |
385 | 413 | }) |
386 | 414 | }) |
| 415 | + |
| 416 | +describe('knowledge embedding transport fallback', () => { |
| 417 | + const options = { |
| 418 | + model: 'text-embedding-3-small', |
| 419 | + taskType: 'document' as const, |
| 420 | + dimensions: 1536, |
| 421 | + projectInputs: null, |
| 422 | + } |
| 423 | + |
| 424 | + it('uses OpenRouter when it is the only configured self-hosted transport', async () => { |
| 425 | + setEnv({ OPENROUTER_API_KEY: 'or-test' }) |
| 426 | + fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]], 3))) |
| 427 | + |
| 428 | + const result = await embedKnowledgeForDeployment(['hello'], options, false) |
| 429 | + |
| 430 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 431 | + const [url, init] = fetchMock.mock.calls[0] |
| 432 | + expect(url).toBe('https://openrouter.ai/api/v1/embeddings') |
| 433 | + expect(JSON.parse((init as RequestInit).body as string)).toMatchObject({ |
| 434 | + model: 'openai/text-embedding-3-small', |
| 435 | + dimensions: 1536, |
| 436 | + }) |
| 437 | + expect(result).toMatchObject({ |
| 438 | + embeddings: [[1, 2]], |
| 439 | + modelName: 'text-embedding-3-small', |
| 440 | + dimensions: 1536, |
| 441 | + }) |
| 442 | + }) |
| 443 | + |
| 444 | + it('keeps the original OpenAI path when OpenRouter is not configured', async () => { |
| 445 | + setEnv({ OPENAI_API_KEY: 'openai-test' }) |
| 446 | + fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]]))) |
| 447 | + |
| 448 | + await embedKnowledgeForDeployment(['hello'], options, false) |
| 449 | + |
| 450 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 451 | + expect(fetchMock.mock.calls[0][0]).toBe('https://api.openai.com/v1/embeddings') |
| 452 | + }) |
| 453 | + |
| 454 | + it('uses Azure before OpenAI and OpenRouter when all are configured', async () => { |
| 455 | + setEnv({ |
| 456 | + AZURE_OPENAI_API_KEY: 'azure-test', |
| 457 | + AZURE_OPENAI_ENDPOINT: 'https://example.openai.azure.com', |
| 458 | + AZURE_OPENAI_API_VERSION: '2024-10-21', |
| 459 | + KB_OPENAI_MODEL_NAME: 'kb-embedding-deployment', |
| 460 | + OPENAI_API_KEY: 'openai-test', |
| 461 | + OPENROUTER_API_KEY: 'or-test', |
| 462 | + }) |
| 463 | + fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]]))) |
| 464 | + |
| 465 | + const result = await embedKnowledgeForDeployment(['hello'], options, false) |
| 466 | + |
| 467 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 468 | + expect(fetchMock.mock.calls[0][0]).toBe( |
| 469 | + 'https://example.openai.azure.com/openai/deployments/kb-embedding-deployment/embeddings?api-version=2024-10-21' |
| 470 | + ) |
| 471 | + expect(result.modelName).toBe('kb-embedding-deployment') |
| 472 | + }) |
| 473 | + |
| 474 | + it('uses a workspace OpenAI key before OpenRouter', async () => { |
| 475 | + setEnv({ OPENROUTER_API_KEY: 'or-test' }) |
| 476 | + mockGetBYOKKey.mockResolvedValue({ apiKey: 'workspace-openai-test', isBYOK: true }) |
| 477 | + fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]]))) |
| 478 | + |
| 479 | + const result = await embedKnowledgeForDeployment( |
| 480 | + ['hello'], |
| 481 | + { ...options, workspaceId: 'workspace-1' }, |
| 482 | + false |
| 483 | + ) |
| 484 | + |
| 485 | + expect(mockGetBYOKKey).toHaveBeenCalledWith('workspace-1', 'openai') |
| 486 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 487 | + const [url, init] = fetchMock.mock.calls[0] |
| 488 | + expect(url).toBe('https://api.openai.com/v1/embeddings') |
| 489 | + expect((init as RequestInit).headers).toMatchObject({ |
| 490 | + Authorization: 'Bearer workspace-openai-test', |
| 491 | + }) |
| 492 | + expect(result.isBYOK).toBe(true) |
| 493 | + }) |
| 494 | + |
| 495 | + it('does not use OpenRouter for non-OpenAI knowledge models', async () => { |
| 496 | + setEnv({ GEMINI_API_KEY: 'gemini-test', OPENROUTER_API_KEY: 'or-test' }) |
| 497 | + fetchMock.mockResolvedValue(jsonResponse({ embeddings: [{ values: [1, 2] }] })) |
| 498 | + |
| 499 | + await embedKnowledgeForDeployment( |
| 500 | + ['hello'], |
| 501 | + { ...options, model: 'gemini-embedding-001' }, |
| 502 | + false |
| 503 | + ) |
| 504 | + |
| 505 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 506 | + expect(fetchMock.mock.calls[0][0]).toContain('generativelanguage.googleapis.com') |
| 507 | + }) |
| 508 | + |
| 509 | + it('ignores OpenRouter on hosted deployments', async () => { |
| 510 | + setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' }) |
| 511 | + fetchMock.mockResolvedValue(jsonResponse(openAIBody([[1, 2]]))) |
| 512 | + |
| 513 | + await embedKnowledgeForDeployment(['hello'], options, true) |
| 514 | + |
| 515 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 516 | + expect(fetchMock.mock.calls[0][0]).toBe('https://api.openai.com/v1/embeddings') |
| 517 | + }) |
| 518 | + |
| 519 | + it('does not fall back after a fatal provider error', async () => { |
| 520 | + setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' }) |
| 521 | + fetchMock.mockResolvedValue(jsonResponse({ error: 'invalid key' }, 401)) |
| 522 | + |
| 523 | + await expect(embedKnowledgeForDeployment(['hello'], options, false)).rejects.toThrow( |
| 524 | + /Embedding API failed: 401/ |
| 525 | + ) |
| 526 | + expect(fetchMock).toHaveBeenCalledOnce() |
| 527 | + }) |
| 528 | + |
| 529 | + it('falls back after transient retries and projects inputs only once', async () => { |
| 530 | + vi.useFakeTimers() |
| 531 | + setEnv({ OPENAI_API_KEY: 'openai-test', OPENROUTER_API_KEY: 'or-test' }) |
| 532 | + const projectInputs = vi.fn(() => ['projected']) |
| 533 | + fetchMock.mockImplementation(async (url) => |
| 534 | + url === 'https://api.openai.com/v1/embeddings' |
| 535 | + ? jsonResponse({ error: 'unavailable' }, 503) |
| 536 | + : jsonResponse(openAIBody([[7, 8]], 2)) |
| 537 | + ) |
| 538 | + |
| 539 | + const pending = embedKnowledgeForDeployment(['secret'], { ...options, projectInputs }, false) |
| 540 | + await vi.runAllTimersAsync() |
| 541 | + const result = await pending |
| 542 | + |
| 543 | + expect(fetchMock).toHaveBeenCalledTimes(5) |
| 544 | + expect(fetchMock.mock.calls.slice(0, 4).every(([url]) => url.includes('api.openai.com'))).toBe( |
| 545 | + true |
| 546 | + ) |
| 547 | + expect(fetchMock.mock.calls[4][0]).toBe('https://openrouter.ai/api/v1/embeddings') |
| 548 | + expect(projectInputs).toHaveBeenCalledOnce() |
| 549 | + expect(result.embeddings).toEqual([[7, 8]]) |
| 550 | + }) |
| 551 | + |
| 552 | + it('classifies only transient embedding failures for failover', () => { |
| 553 | + expect(isTransientEmbeddingError(new EmbeddingAPIError('unavailable', 503))).toBe(true) |
| 554 | + expect(isTransientEmbeddingError(new EmbeddingAPIError('rate limited', 429))).toBe(true) |
| 555 | + expect(isTransientEmbeddingError(new EmbeddingAPIError('invalid key', 401))).toBe(false) |
| 556 | + expect(isTransientEmbeddingError(new DOMException('timed out', 'AbortError'))).toBe(true) |
| 557 | + }) |
| 558 | +}) |
0 commit comments