diff --git a/scripts/README.md b/scripts/README.md index d9ec3f8..793bcc3 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -26,6 +26,12 @@ pnpm test:db:seed:token # against the token server pnpm test:db:seed:basic # against the basic-auth server ``` +The seed data includes: + +- `articles` and `code_snippets` with small 8-dimensional fixture embeddings. +- `huggingface_384d` with HuggingFace-style 384-dimensional fixture embeddings + for testing inferred collection dimensions and custom embedding server URLs. + Stop with `Ctrl+C`. ## Auth credentials diff --git a/scripts/seed-test-db.mjs b/scripts/seed-test-db.mjs index ad8990e..5967f11 100644 --- a/scripts/seed-test-db.mjs +++ b/scripts/seed-test-db.mjs @@ -27,13 +27,11 @@ if (authMode === 'token') { const client = new ChromaClient({ host: 'localhost', port, ssl: false, headers }) -const EMBED_DIM = 8 - -function seededVector(seed) { +function seededVector(seed, dimension) { // Deterministic pseudo-random vector so re-seeds produce stable embeddings. - const out = new Array(EMBED_DIM) + const out = new Array(dimension) let s = seed - for (let i = 0; i < EMBED_DIM; i++) { + for (let i = 0; i < dimension; i++) { s = (s * 9301 + 49297) % 233280 out[i] = s / 233280 - 0.5 } @@ -43,6 +41,7 @@ function seededVector(seed) { const collections = [ { name: 'articles', + dimension: 8, metadata: { description: 'Sample blog articles for explorer testing' }, docs: [ { id: 'a1', document: 'Getting started with vector databases.', metadata: { topic: 'intro', words: 5 } }, @@ -54,6 +53,7 @@ const collections = [ }, { name: 'code_snippets', + dimension: 8, metadata: { description: 'Code snippets across languages' }, docs: [ { id: 'c1', document: 'def hello():\n print("hi")', metadata: { language: 'python', lines: 2 } }, @@ -63,6 +63,32 @@ const collections = [ { id: 'c5', document: 'puts "hello"', metadata: { language: 'ruby', lines: 1 } }, ], }, + { + name: 'huggingface_384d', + dimension: 384, + metadata: { + description: 'HuggingFace-style 384-dimensional embeddings for URL and dimension testing', + embedding_provider: 'huggingface-server', + embedding_model: 'sentence-transformers/all-MiniLM-L6-v2', + }, + docs: [ + { + id: 'hf1', + document: 'A local HuggingFace embedding server can expose sentence-transformer vectors.', + metadata: { provider: 'huggingface-server', model: 'all-MiniLM-L6-v2', dimension: 384 }, + }, + { + id: 'hf2', + document: 'MiniLM embeddings are compact enough for fast document exploration.', + metadata: { provider: 'huggingface-server', model: 'all-MiniLM-L6-v2', dimension: 384 }, + }, + { + id: 'hf3', + document: 'This fixture helps verify inferred dimensions in the collection UI.', + metadata: { provider: 'huggingface-server', model: 'all-MiniLM-L6-v2', dimension: 384 }, + }, + ], + }, ] async function main() { @@ -76,9 +102,9 @@ async function main() { const ids = spec.docs.map((d) => d.id) const documents = spec.docs.map((d) => d.document) const metadatas = spec.docs.map((d) => d.metadata) - const embeddings = spec.docs.map((_, idx) => seededVector(seedBase + idx + 1)) + const embeddings = spec.docs.map((_, idx) => seededVector(seedBase + idx + 1, spec.dimension)) await col.upsert({ ids, documents, metadatas, embeddings }) - console.log(`[seed] upserted ${ids.length} docs into "${spec.name}"`) + console.log(`[seed] upserted ${ids.length} docs into "${spec.name}" (${spec.dimension}d)`) } console.log('[seed] done') diff --git a/src/components/collections/CollectionConfigView.tsx b/src/components/collections/CollectionConfigView.tsx index 03a91ff..6415624 100644 --- a/src/components/collections/CollectionConfigView.tsx +++ b/src/components/collections/CollectionConfigView.tsx @@ -3,7 +3,16 @@ import { ChevronDown, ChevronRight, Plus, X } from 'lucide-react' import { useDraftCollection, DraftHNSWConfig } from '../../context/DraftCollectionContext' import { useChromaDB } from '../../providers/ChromaDBProvider' import { useCollection } from '../../context/CollectionContext' -import { EMBEDDING_FUNCTIONS, EMBEDDING_FUNCTION_GROUPS, getEmbeddingFunctionById } from '../../constants/embedding-functions' +import { + buildEmbeddingFunctionOverride, + EMBEDDING_FUNCTIONS, + EMBEDDING_FUNCTION_GROUPS, + embeddingFunctionUsesUrl, + getEmbeddingFunctionById, + getEmbeddingFunctionUrlLabel, + getEmbeddingFunctionUrlPlaceholder, + validateEmbeddingFunctionUrl, +} from '../../constants/embedding-functions' import { MetadataValueType, validateMetadataValue } from '../../types/metadata' import { CopyProgressDialog } from './CopyProgressDialog' @@ -29,19 +38,19 @@ export function CollectionConfigView() { }) const selectedEf = draftCollection ? getEmbeddingFunctionById(draftCollection.embeddingFunctionId) : EMBEDDING_FUNCTIONS[0] + const embeddingFunctionUrlError = draftCollection && selectedEf + ? validateEmbeddingFunctionUrl(selectedEf, draftCollection.embeddingFunctionUrl) + : null + const dimensionsLabel = selectedEf?.dimensions ? `${selectedEf.dimensions}d` : 'Inferred from first embedding' // Handle embedding function change - auto-populate dimension const handleEfChange = useCallback( (efId: string) => { - updateDraft({ embeddingFunctionId: efId, dimensionOverride: '' }) - }, - [updateDraft] - ) - - // Handle dimension override change - const handleDimensionChange = useCallback( - (value: string) => { - updateDraft({ dimensionOverride: value }) + const nextEf = getEmbeddingFunctionById(efId) + updateDraft({ + embeddingFunctionId: efId, + embeddingFunctionUrl: nextEf?.url || '', + }) }, [updateDraft] ) @@ -70,13 +79,16 @@ export function CollectionConfigView() { if (!draftCollection?.sourceCollection) return false const source = draftCollection.sourceCollection - const sourceType = source.embeddingFunction?.name === 'openai' ? 'openai' : 'default' + const sourceType = source.embeddingFunction?.name ?? 'default' const sourceModel = source.embeddingFunction?.config?.model_name as string | undefined + const sourceUrl = source.embeddingFunction?.config?.url as string | undefined const draftEf = getEmbeddingFunctionById(draftCollection.embeddingFunctionId) + const draftUrl = draftCollection.embeddingFunctionUrl.trim() || draftEf?.url if (draftEf?.type !== sourceType) return true - if (draftEf?.type === 'openai' && draftEf?.modelName !== sourceModel) return true + if (sourceModel && draftEf?.modelName !== sourceModel) return true + if (draftUrl !== sourceUrl) return true return false }, [draftCollection]) @@ -89,6 +101,11 @@ export function CollectionConfigView() { if (!draftCollection.name.trim()) { errors.name = 'Collection name is required' } + const draftEf = getEmbeddingFunctionById(draftCollection.embeddingFunctionId) + const urlError = validateEmbeddingFunctionUrl(draftEf, draftCollection.embeddingFunctionUrl) + if (urlError) { + errors.embeddingFunctionUrl = urlError + } if (Object.keys(errors).length > 0) { return @@ -132,10 +149,9 @@ export function CollectionConfigView() { const result = await window.electronAPI.chromadb.copyCollection(currentProfile.id, { sourceCollectionName: draftCollection.sourceCollection.name, targetName: draftCollection.name.trim(), - embeddingFunction: draftEf ? { - type: draftEf.type, - modelName: draftEf.modelName, - } : undefined, + embeddingFunction: draftEf + ? buildEmbeddingFunctionOverride(draftEf, { url: draftCollection.embeddingFunctionUrl }) + : undefined, hnsw: Object.keys(hnswConfig).length > 0 ? hnswConfig : undefined, regenerateEmbeddings: needsEmbeddingRegeneration(), }) @@ -283,22 +299,40 @@ export function CollectionConfigView() { {selectedEf &&

{selectedEf.modelName}

} - {/* Dimension Input */} + {embeddingFunctionUsesUrl(selectedEf) && ( +
+ + updateDraft({ embeddingFunctionUrl: e.target.value })} + placeholder={getEmbeddingFunctionUrlPlaceholder(selectedEf)} + className={`${inputClassName} ${embeddingFunctionUrlError ? 'border-destructive' : ''}`} + style={inputStyle} + /> + {embeddingFunctionUrlError ? ( +

{embeddingFunctionUrlError}

+ ) : ( +

+ Used by the client when generating embeddings. +

+ )} +
+ )} + + {/* Dimension Info */}
- - handleDimensionChange(e.target.value)} - placeholder={selectedEf?.dimensions?.toString() || 'auto'} - className={inputClassName} - style={inputStyle} - /> + + Dimensions + +

+ {dimensionsLabel} +

- Default: {selectedEf?.dimensions ?? 'auto'} + Chroma sets the collection dimension from generated embeddings.

@@ -586,7 +620,7 @@ export function CollectionConfigView() { diff --git a/src/constants/embedding-functions.ts b/src/constants/embedding-functions.ts index 3bee213..e8d0236 100644 --- a/src/constants/embedding-functions.ts +++ b/src/constants/embedding-functions.ts @@ -1,3 +1,5 @@ +import type { EmbeddingFunctionOverride } from '../types/electron' + export type EmbeddingFunctionType = | 'default' | 'openai' @@ -25,6 +27,59 @@ export interface EmbeddingFunctionConfig { group: string // For UI grouping } +export function embeddingFunctionUsesUrl(ef: Pick | undefined): boolean { + return ef?.type === 'ollama' || ef?.type === 'huggingface-server' +} + +export function embeddingFunctionRequiresUrl(ef: Pick | undefined): boolean { + return ef?.type === 'huggingface-server' +} + +export function getEmbeddingFunctionUrlLabel(ef: Pick | undefined): string { + return ef?.type === 'huggingface-server' ? 'Server URL' : 'URL' +} + +export function getEmbeddingFunctionUrlPlaceholder(ef: Pick | undefined): string { + return ef?.type === 'huggingface-server' ? 'https://your-embedding-server.example.com' : 'http://localhost:11434' +} + +export function validateEmbeddingFunctionUrl( + ef: Pick | undefined, + value: string +): string | null { + if (!embeddingFunctionUsesUrl(ef)) return null + + const trimmed = value.trim() + if (!trimmed) { + return embeddingFunctionRequiresUrl(ef) ? `${getEmbeddingFunctionUrlLabel(ef)} is required` : null + } + + try { + const parsed = new URL(trimmed) + if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') { + return `${getEmbeddingFunctionUrlLabel(ef)} must use http or https` + } + } catch { + return `${getEmbeddingFunctionUrlLabel(ef)} must be a valid URL` + } + + return null +} + +export function buildEmbeddingFunctionOverride( + ef: EmbeddingFunctionConfig, + options: { url?: string } = {} +): EmbeddingFunctionOverride { + const url = options.url?.trim() || ef.url + + return { + type: ef.type, + modelName: ef.modelName, + url: url || undefined, + accountId: ef.accountId, + } +} + export const EMBEDDING_FUNCTIONS: EmbeddingFunctionConfig[] = [ // Default (Local) { diff --git a/src/context/DraftCollectionContext.tsx b/src/context/DraftCollectionContext.tsx index b1ac567..b98fef2 100644 --- a/src/context/DraftCollectionContext.tsx +++ b/src/context/DraftCollectionContext.tsx @@ -2,7 +2,11 @@ import { createContext, useContext, useState, useCallback, ReactNode } from 'rea import { useChromaDB } from '../providers/ChromaDBProvider' import { useCreateCollectionMutation } from '../hooks/useChromaQueries' import { useCollection } from './CollectionContext' -import { EMBEDDING_FUNCTIONS } from '../constants/embedding-functions' +import { + buildEmbeddingFunctionOverride, + EMBEDDING_FUNCTIONS, + validateEmbeddingFunctionUrl, +} from '../constants/embedding-functions' import { TypedMetadataRecord, typedMetadataToChromaFormat, validateMetadataValue } from '../types/metadata' export interface DraftHNSWConfig { @@ -14,7 +18,7 @@ export interface DraftHNSWConfig { export interface DraftCollection { name: string embeddingFunctionId: string - dimensionOverride: string // String to allow empty input for auto + embeddingFunctionUrl: string hnsw: DraftHNSWConfig // Optional first document (metadata here defines the collection's schema) firstDocument: { @@ -52,7 +56,7 @@ function createInitialDraft(): DraftCollection { return { name: '', embeddingFunctionId: DEFAULT_EF.id, - dimensionOverride: '', + embeddingFunctionUrl: DEFAULT_EF.url || '', hnsw: { space: 'l2', efConstruction: '', @@ -81,9 +85,11 @@ export function DraftCollectionProvider({ children }: DraftCollectionProviderPro const startCopyFromCollection = useCallback((collection: CollectionInfo) => { // Find matching embedding function from collection's config let embeddingFunctionId: string = DEFAULT_EF.id + let embeddingFunctionUrl = DEFAULT_EF.url || '' if (collection.embeddingFunction) { - const efType = collection.embeddingFunction.name === 'openai' ? 'openai' : 'default' + const efType = collection.embeddingFunction.name const efModelName = collection.embeddingFunction.config?.model_name as string | undefined + const efUrl = collection.embeddingFunction.config?.url as string | undefined // Try to find exact match const matchingEf = EMBEDDING_FUNCTIONS.find( @@ -98,6 +104,9 @@ export function DraftCollectionProvider({ children }: DraftCollectionProviderPro embeddingFunctionId = typeMatch.id } } + + const selectedEf = EMBEDDING_FUNCTIONS.find((ef) => ef.id === embeddingFunctionId) || DEFAULT_EF + embeddingFunctionUrl = efUrl || selectedEf.url || '' } // Extract HNSW config from collection metadata @@ -121,7 +130,7 @@ export function DraftCollectionProvider({ children }: DraftCollectionProviderPro setDraftCollection({ name: `${collection.name}-copy`, embeddingFunctionId, - dimensionOverride: '', + embeddingFunctionUrl, hnsw, firstDocument: null, sourceCollection: collection, @@ -142,6 +151,12 @@ export function DraftCollectionProvider({ children }: DraftCollectionProviderPro return rest }) } + if (updates.embeddingFunctionUrl !== undefined || updates.embeddingFunctionId !== undefined) { + setValidationErrors((prev) => { + const { embeddingFunctionUrl, ...rest } = prev + return rest + }) + } }, []) const cancelCreation = useCallback(() => { @@ -176,23 +191,18 @@ export function DraftCollectionProvider({ children }: DraftCollectionProviderPro setIsCreating(true) try { const selectedEf = EMBEDDING_FUNCTIONS.find((ef) => ef.id === draftCollection.embeddingFunctionId) || DEFAULT_EF - - // Parse dimension override - let dimension: number | undefined - if (draftCollection.dimensionOverride.trim()) { - const parsed = parseInt(draftCollection.dimensionOverride, 10) - if (!isNaN(parsed) && parsed > 0) { - dimension = parsed - } + const urlError = validateEmbeddingFunctionUrl(selectedEf, draftCollection.embeddingFunctionUrl) + if (urlError) { + setValidationErrors({ embeddingFunctionUrl: urlError }) + return } // Build params const params: Parameters[0] = { name: draftCollection.name.trim(), - embeddingFunction: { - type: selectedEf.type, - modelName: selectedEf.modelName, - }, + embeddingFunction: buildEmbeddingFunctionOverride(selectedEf, { + url: draftCollection.embeddingFunctionUrl, + }), } // Add HNSW config if any non-default values are set diff --git a/tests/unit/ipc-contract.test.ts b/tests/unit/ipc-contract.test.ts index c76a935..d5acaf0 100644 --- a/tests/unit/ipc-contract.test.ts +++ b/tests/unit/ipc-contract.test.ts @@ -1,7 +1,9 @@ import { describe, expect, it } from 'vitest' import { parseConnectionProfile, + parseCopyCollectionParams, parseCreateCollectionParams, + parseEmbeddingOverride, parseSearchDocumentsParams, validateExternalUrl, } from '../../electron/ipc-contract' @@ -41,9 +43,47 @@ describe('ipc contract validators', () => { }) }) + it('preserves embedding function URL and account fields', () => { + expect(parseEmbeddingOverride({ + type: 'huggingface-server', + modelName: 'custom', + url: 'https://embeddings.example.com', + })).toMatchObject({ + type: 'huggingface-server', + modelName: 'custom', + url: 'https://embeddings.example.com', + }) + + expect(parseCreateCollectionParams({ + name: 'docs', + embeddingFunction: { + type: 'ollama', + modelName: 'nomic-embed-text', + url: 'http://localhost:11434', + }, + }).embeddingFunction).toMatchObject({ + type: 'ollama', + modelName: 'nomic-embed-text', + url: 'http://localhost:11434', + }) + + expect(parseCopyCollectionParams({ + sourceCollectionName: 'docs', + targetName: 'docs-copy', + regenerateEmbeddings: true, + embeddingFunction: { + type: 'cloudflare-worker-ai', + modelName: '@cf/baai/bge-base-en-v1.5', + accountId: 'account-123', + }, + }).embeddingFunction).toMatchObject({ + type: 'cloudflare-worker-ai', + accountId: 'account-123', + }) + }) + it('only allows http and https external URLs', () => { expect(validateExternalUrl('https://trychroma.com/docs')).toBe('https://trychroma.com/docs') expect(() => validateExternalUrl('file:///etc/passwd')).toThrow(/http: or https:/) }) }) -