|
| 1 | +/** |
| 2 | + * Tests for v1 knowledge search API route. |
| 3 | + * Specifically guards the per-KB embedding model resolution and the |
| 4 | + * multi-model rejection so the v1 endpoint stays in lockstep with the |
| 5 | + * internal route. |
| 6 | + * |
| 7 | + * @vitest-environment node |
| 8 | + */ |
| 9 | +import { createMockRequest, knowledgeApiUtilsMock, knowledgeApiUtilsMockFns } from '@sim/testing' |
| 10 | +import { beforeEach, describe, expect, it, vi } from 'vitest' |
| 11 | + |
| 12 | +const { |
| 13 | + mockHandleVectorOnlySearch, |
| 14 | + mockHandleTagOnlySearch, |
| 15 | + mockHandleTagAndVectorSearch, |
| 16 | + mockGetQueryStrategy, |
| 17 | + mockGenerateSearchEmbedding, |
| 18 | + mockGetDocumentNamesByIds, |
| 19 | + mockAuthenticateRequest, |
| 20 | + mockValidateWorkspaceAccess, |
| 21 | +} = vi.hoisted(() => ({ |
| 22 | + mockHandleVectorOnlySearch: vi.fn(), |
| 23 | + mockHandleTagOnlySearch: vi.fn(), |
| 24 | + mockHandleTagAndVectorSearch: vi.fn(), |
| 25 | + mockGetQueryStrategy: vi.fn(), |
| 26 | + mockGenerateSearchEmbedding: vi.fn(), |
| 27 | + mockGetDocumentNamesByIds: vi.fn(), |
| 28 | + mockAuthenticateRequest: vi.fn(), |
| 29 | + mockValidateWorkspaceAccess: vi.fn(), |
| 30 | +})) |
| 31 | + |
| 32 | +vi.mock('@/app/api/knowledge/search/utils', () => ({ |
| 33 | + handleVectorOnlySearch: mockHandleVectorOnlySearch, |
| 34 | + handleTagOnlySearch: mockHandleTagOnlySearch, |
| 35 | + handleTagAndVectorSearch: mockHandleTagAndVectorSearch, |
| 36 | + getQueryStrategy: mockGetQueryStrategy, |
| 37 | + generateSearchEmbedding: mockGenerateSearchEmbedding, |
| 38 | + getDocumentNamesByIds: mockGetDocumentNamesByIds, |
| 39 | +})) |
| 40 | + |
| 41 | +vi.mock('@/app/api/knowledge/utils', () => knowledgeApiUtilsMock) |
| 42 | + |
| 43 | +vi.mock('@/app/api/v1/knowledge/utils', () => ({ |
| 44 | + authenticateRequest: mockAuthenticateRequest, |
| 45 | + validateWorkspaceAccess: mockValidateWorkspaceAccess, |
| 46 | + parseJsonBody: async (req: Request) => { |
| 47 | + try { |
| 48 | + return { success: true, data: await req.json() } |
| 49 | + } catch { |
| 50 | + return { |
| 51 | + success: false, |
| 52 | + response: new Response(JSON.stringify({ error: 'Invalid JSON' }), { status: 400 }), |
| 53 | + } |
| 54 | + } |
| 55 | + }, |
| 56 | + validateSchema: <T>( |
| 57 | + schema: { |
| 58 | + safeParse: (v: unknown) => { |
| 59 | + success: boolean |
| 60 | + data?: T |
| 61 | + error?: { issues: { message: string }[] } |
| 62 | + } |
| 63 | + }, |
| 64 | + data: unknown |
| 65 | + ) => { |
| 66 | + const result = schema.safeParse(data) |
| 67 | + if (!result.success) { |
| 68 | + return { |
| 69 | + success: false, |
| 70 | + response: new Response( |
| 71 | + JSON.stringify({ error: result.error?.issues.map((i) => i.message).join(', ') }), |
| 72 | + { status: 400 } |
| 73 | + ), |
| 74 | + } |
| 75 | + } |
| 76 | + return { success: true, data: result.data } |
| 77 | + }, |
| 78 | + handleError: (e: unknown) => |
| 79 | + new Response(JSON.stringify({ error: e instanceof Error ? e.message : 'error' }), { |
| 80 | + status: 500, |
| 81 | + }), |
| 82 | +})) |
| 83 | + |
| 84 | +vi.mock('@/lib/knowledge/tags/service', () => ({ |
| 85 | + getDocumentTagDefinitions: vi.fn().mockResolvedValue([]), |
| 86 | +})) |
| 87 | + |
| 88 | +import { POST } from '@/app/api/v1/knowledge/search/route' |
| 89 | + |
| 90 | +const mockCheckKnowledgeBaseAccess = knowledgeApiUtilsMockFns.mockCheckKnowledgeBaseAccess |
| 91 | + |
| 92 | +const baseKb = (id: string, embeddingModel: string) => ({ |
| 93 | + id, |
| 94 | + userId: 'user-1', |
| 95 | + name: `KB ${id}`, |
| 96 | + workspaceId: 'ws-1', |
| 97 | + embeddingModel, |
| 98 | + deletedAt: null, |
| 99 | +}) |
| 100 | + |
| 101 | +describe('v1 knowledge search route — per-KB embedding model', () => { |
| 102 | + beforeEach(() => { |
| 103 | + vi.clearAllMocks() |
| 104 | + mockAuthenticateRequest.mockResolvedValue({ |
| 105 | + requestId: 'req-1', |
| 106 | + userId: 'user-1', |
| 107 | + rateLimit: {}, |
| 108 | + }) |
| 109 | + mockValidateWorkspaceAccess.mockResolvedValue(null) |
| 110 | + mockGetQueryStrategy.mockReturnValue({ distanceThreshold: 0.5 }) |
| 111 | + mockGenerateSearchEmbedding.mockResolvedValue([0.1, 0.2, 0.3]) |
| 112 | + mockHandleVectorOnlySearch.mockResolvedValue([]) |
| 113 | + mockGetDocumentNamesByIds.mockResolvedValue({}) |
| 114 | + }) |
| 115 | + |
| 116 | + it('passes the KB embedding model into generateSearchEmbedding', async () => { |
| 117 | + mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ |
| 118 | + hasAccess: true, |
| 119 | + knowledgeBase: baseKb('kb-gemini', 'gemini-embedding-001'), |
| 120 | + }) |
| 121 | + |
| 122 | + const req = createMockRequest('POST', { |
| 123 | + workspaceId: 'ws-1', |
| 124 | + knowledgeBaseIds: 'kb-gemini', |
| 125 | + query: 'hello', |
| 126 | + }) |
| 127 | + const res = await POST(req) |
| 128 | + |
| 129 | + expect(res.status).toBe(200) |
| 130 | + expect(mockGenerateSearchEmbedding).toHaveBeenCalledWith( |
| 131 | + 'hello', |
| 132 | + 'gemini-embedding-001', |
| 133 | + 'ws-1' |
| 134 | + ) |
| 135 | + }) |
| 136 | + |
| 137 | + it('rejects cross-KB queries with mixed embedding models', async () => { |
| 138 | + mockCheckKnowledgeBaseAccess |
| 139 | + .mockResolvedValueOnce({ |
| 140 | + hasAccess: true, |
| 141 | + knowledgeBase: baseKb('kb-openai', 'text-embedding-3-small'), |
| 142 | + }) |
| 143 | + .mockResolvedValueOnce({ |
| 144 | + hasAccess: true, |
| 145 | + knowledgeBase: baseKb('kb-gemini', 'gemini-embedding-001'), |
| 146 | + }) |
| 147 | + |
| 148 | + const req = createMockRequest('POST', { |
| 149 | + workspaceId: 'ws-1', |
| 150 | + knowledgeBaseIds: ['kb-openai', 'kb-gemini'], |
| 151 | + query: 'hello', |
| 152 | + }) |
| 153 | + const res = await POST(req) |
| 154 | + |
| 155 | + expect(res.status).toBe(400) |
| 156 | + expect(mockGenerateSearchEmbedding).not.toHaveBeenCalled() |
| 157 | + }) |
| 158 | + |
| 159 | + it('allows tag-only search across mixed embedding models', async () => { |
| 160 | + mockHandleTagOnlySearch.mockResolvedValue([]) |
| 161 | + mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({ |
| 162 | + hasAccess: true, |
| 163 | + knowledgeBase: baseKb('kb-mixed', 'text-embedding-3-small'), |
| 164 | + }) |
| 165 | + |
| 166 | + const req = createMockRequest('POST', { |
| 167 | + workspaceId: 'ws-1', |
| 168 | + knowledgeBaseIds: 'kb-mixed', |
| 169 | + tagFilters: [{ tagName: 'category', operator: 'eq', value: 'docs' }], |
| 170 | + }) |
| 171 | + const res = await POST(req) |
| 172 | + |
| 173 | + expect(res.status).toBe(400) |
| 174 | + // tagName "category" is undefined in our empty getDocumentTagDefinitions mock, |
| 175 | + // so the route returns 400 before reaching the search handlers — but crucially |
| 176 | + // it never tries to generate an embedding. |
| 177 | + expect(mockGenerateSearchEmbedding).not.toHaveBeenCalled() |
| 178 | + }) |
| 179 | +}) |
0 commit comments