|
| 1 | +/** |
| 2 | + * @vitest-environment node |
| 3 | + */ |
| 4 | +import { beforeEach, describe, expect, it, vi } from 'vitest' |
| 5 | + |
| 6 | +const { mockEnv, mockOllamaMetadata } = vi.hoisted(() => ({ |
| 7 | + mockEnv: {} as { KB_EMBEDDING_MODEL?: string; EMBEDDING_OUTPUT_DIMS?: string }, |
| 8 | + mockOllamaMetadata: vi.fn(), |
| 9 | +})) |
| 10 | + |
| 11 | +/** |
| 12 | + * `envNumber` is the real implementation, not a stub: the whole point of the |
| 13 | + * cases below is that `createEnv` runs with `skipValidation`, so every value |
| 14 | + * arrives as the raw string from the environment however its schema is declared. |
| 15 | + */ |
| 16 | +vi.mock('@/lib/core/config/env', async (importOriginal) => ({ |
| 17 | + ...(await importOriginal<typeof import('@/lib/core/config/env')>()), |
| 18 | + env: mockEnv, |
| 19 | +})) |
| 20 | + |
| 21 | +vi.mock('@/lib/embeddings/ollama-model-catalog.server', () => ({ |
| 22 | + getOllamaEmbeddingModelMetadata: mockOllamaMetadata, |
| 23 | +})) |
| 24 | + |
| 25 | +import { getConfiguredKbEmbedding } from '@/lib/knowledge/embeddings' |
| 26 | + |
| 27 | +describe('getConfiguredKbEmbedding', () => { |
| 28 | + beforeEach(() => { |
| 29 | + vi.clearAllMocks() |
| 30 | + mockEnv.KB_EMBEDDING_MODEL = undefined |
| 31 | + mockEnv.EMBEDDING_OUTPUT_DIMS = undefined |
| 32 | + mockOllamaMetadata.mockRejectedValue(new Error('no server configured')) |
| 33 | + }) |
| 34 | + |
| 35 | + it('defaults to the model and width knowledge bases were always created at', async () => { |
| 36 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 37 | + model: 'text-embedding-3-small', |
| 38 | + dimensions: 1536, |
| 39 | + }) |
| 40 | + }) |
| 41 | + |
| 42 | + it('stores at the configured width when the model can emit it', async () => { |
| 43 | + mockEnv.KB_EMBEDDING_MODEL = 'text-embedding-3-large' |
| 44 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '3072' |
| 45 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 46 | + model: 'text-embedding-3-large', |
| 47 | + dimensions: 3072, |
| 48 | + }) |
| 49 | + }) |
| 50 | + |
| 51 | + it('accepts any storable width from a model on the deployment’s own Ollama', async () => { |
| 52 | + mockEnv.KB_EMBEDDING_MODEL = 'ollama/nomic-embed-text' |
| 53 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '768' |
| 54 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 55 | + model: 'ollama/nomic-embed-text', |
| 56 | + dimensions: 768, |
| 57 | + }) |
| 58 | + }) |
| 59 | + |
| 60 | + /** |
| 61 | + * Sim can read a local model's width off the server it is installed on, so an |
| 62 | + * unstated width is resolved rather than defaulted — defaulting would create |
| 63 | + * every base at 1,536 and fail each document against a 768-wide model. |
| 64 | + */ |
| 65 | + it('reads an unstated Ollama width from the server the model is installed on', async () => { |
| 66 | + mockEnv.KB_EMBEDDING_MODEL = 'ollama/nomic-embed-text' |
| 67 | + mockOllamaMetadata.mockResolvedValue({ id: 'nomic-embed-text:latest', dimensions: 768 }) |
| 68 | + |
| 69 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 70 | + model: 'ollama/nomic-embed-text', |
| 71 | + dimensions: 768, |
| 72 | + }) |
| 73 | + expect(mockOllamaMetadata).toHaveBeenCalledWith('ollama/nomic-embed-text') |
| 74 | + }) |
| 75 | + |
| 76 | + it('prefers a stated width over the server, and never asks when one is stated', async () => { |
| 77 | + mockEnv.KB_EMBEDDING_MODEL = 'ollama/nomic-embed-text' |
| 78 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '1024' |
| 79 | + mockOllamaMetadata.mockResolvedValue({ id: 'nomic-embed-text:latest', dimensions: 768 }) |
| 80 | + |
| 81 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 82 | + model: 'ollama/nomic-embed-text', |
| 83 | + dimensions: 1024, |
| 84 | + }) |
| 85 | + expect(mockOllamaMetadata).not.toHaveBeenCalled() |
| 86 | + }) |
| 87 | + |
| 88 | + /** |
| 89 | + * There is no width to fall back to: the adapter cannot ask Ollama for a |
| 90 | + * different one, so the platform default would pin the base at a width the |
| 91 | + * model does not emit and fail every document. Refusing is recoverable; |
| 92 | + * a base created at an impossible width is not. |
| 93 | + */ |
| 94 | + it('refuses to create a base when the server cannot be asked', async () => { |
| 95 | + mockEnv.KB_EMBEDDING_MODEL = 'ollama/nomic-embed-text' |
| 96 | + mockOllamaMetadata.mockRejectedValue(new Error('ECONNREFUSED')) |
| 97 | + |
| 98 | + await expect(getConfiguredKbEmbedding()).rejects.toThrow('Set EMBEDDING_OUTPUT_DIMS') |
| 99 | + }) |
| 100 | + |
| 101 | + it('refuses to create a base when the model emits a width no column can store', async () => { |
| 102 | + mockEnv.KB_EMBEDDING_MODEL = 'ollama/odd-width' |
| 103 | + mockOllamaMetadata.mockResolvedValue({ id: 'odd-width:latest', dimensions: 1152 }) |
| 104 | + |
| 105 | + await expect(getConfiguredKbEmbedding()).rejects.toThrow( |
| 106 | + 'emits 1152-dimensional vectors, which knowledge bases cannot store' |
| 107 | + ) |
| 108 | + }) |
| 109 | + |
| 110 | + it('falls back when the width is not a number at all', async () => { |
| 111 | + mockEnv.KB_EMBEDDING_MODEL = 'text-embedding-3-large' |
| 112 | + mockEnv.EMBEDDING_OUTPUT_DIMS = 'wide' |
| 113 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 114 | + model: 'text-embedding-3-large', |
| 115 | + dimensions: 1536, |
| 116 | + }) |
| 117 | + }) |
| 118 | + |
| 119 | + it('falls back when the width has no storage column, keeping the chosen model', async () => { |
| 120 | + mockEnv.KB_EMBEDDING_MODEL = 'text-embedding-3-large' |
| 121 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '1000' |
| 122 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 123 | + model: 'text-embedding-3-large', |
| 124 | + dimensions: 1536, |
| 125 | + }) |
| 126 | + }) |
| 127 | + |
| 128 | + it('falls back when the model cannot emit the configured width', async () => { |
| 129 | + mockEnv.KB_EMBEDDING_MODEL = 'gemini-embedding-001' |
| 130 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '1024' |
| 131 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 132 | + model: 'gemini-embedding-001', |
| 133 | + dimensions: 1536, |
| 134 | + }) |
| 135 | + }) |
| 136 | + |
| 137 | + it('falls back to the default model when the configured one cannot index a knowledge base', async () => { |
| 138 | + mockEnv.KB_EMBEDDING_MODEL = 'text-embedding-ada-002' |
| 139 | + mockEnv.EMBEDDING_OUTPUT_DIMS = '768' |
| 140 | + await expect(getConfiguredKbEmbedding()).resolves.toEqual({ |
| 141 | + model: 'text-embedding-3-small', |
| 142 | + dimensions: 768, |
| 143 | + }) |
| 144 | + }) |
| 145 | +}) |
0 commit comments