feat(knowledge): add Ollama embedding provider support

apps/sim/app/api/knowledge/[id]/route.test.ts

@@ -15,6 +15,7 @@ const { mockGetSession, mockDbChain } = vi.hoisted(() => {
     limit: vi.fn().mockReturnThis(),
     update: vi.fn().mockReturnThis(),
     set: vi.fn().mockReturnThis(),
+    execute: vi.fn().mockResolvedValue(undefined),
   }
   return { mockGetSession, mockDbChain }
 })
@@ -98,6 +99,10 @@ vi.mock('@sim/db/schema', () => ({
 
 vi.mock('@/lib/audit/log', () => auditMock)
 
+vi.mock('@/lib/knowledge/dynamic-tables', () => ({
+  dropKBEmbeddingTable: vi.fn().mockResolvedValue(undefined),
+}))
+
 vi.mock('@/lib/knowledge/service', async (importOriginal) => {
   const actual = await importOriginal<typeof import('@/lib/knowledge/service')>()
   return {

apps/sim/app/api/knowledge/[id]/route.ts

@@ -26,16 +26,14 @@ const logger = createLogger('KnowledgeBaseByIdAPI')
 const UpdateKnowledgeBaseSchema = z.object({
   name: z.string().min(1, 'Name is required').optional(),
   description: z.string().optional(),
-  embeddingModel: z.literal('text-embedding-3-small').optional(),
-  embeddingDimension: z.literal(1536).optional(),
   workspaceId: z.string().nullable().optional(),
   chunkingConfig: z
     .object({
       /** Maximum chunk size in tokens (1 token ≈ 4 characters) */
       maxSize: z.number().min(100).max(4000),
       /** Minimum chunk size in characters */
       minSize: z.number().min(1).max(2000),
-      /** Overlap between chunks in characters */
+      /** Overlap between chunks in tokens (1 token ≈ 4 characters) */
       overlap: z.number().min(0).max(500),
     })
     .refine(
@@ -205,6 +203,11 @@ export async function DELETE(
     }
 
     await deleteKnowledgeBase(id, requestId)
+    // Note: per-KB embedding tables (kb_embeddings_{id}) are intentionally NOT dropped here.
+    // deleteKnowledgeBase performs a soft delete — the KB can be restored via the restore endpoint,
+    // which requires the per-KB table to still exist. This is consistent with how OpenAI embeddings
+    // remain in the shared table after a KB soft delete. Per-KB table cleanup should occur
+    // as part of a permanent purge/hard-delete operation.
 
     try {
       PlatformEvents.knowledgeBaseDeleted({

apps/sim/app/api/knowledge/route.ts

@@ -5,9 +5,16 @@ import { AuditAction, AuditResourceType, recordAudit } from '@/lib/audit/log'
 import { getSession } from '@/lib/auth'
 import { PlatformEvents } from '@/lib/core/telemetry'
 import { generateRequestId } from '@/lib/core/utils/request'
+import {
+  createKBEmbeddingTable,
+  dropKBEmbeddingTable,
+  parseEmbeddingModel,
+} from '@/lib/knowledge/dynamic-tables'
+import { getOllamaBaseUrl, validateOllamaModel } from '@/lib/knowledge/embeddings'
 import {
   createKnowledgeBase,
   getKnowledgeBases,
+  hardDeleteKnowledgeBase,
   KnowledgeBaseConflictError,
   type KnowledgeBaseScope,
 } from '@/lib/knowledge/service'
@@ -26,8 +33,71 @@ const CreateKnowledgeBaseSchema = z.object({
   name: z.string().min(1, 'Name is required'),
   description: z.string().optional(),
   workspaceId: z.string().min(1, 'Workspace ID is required'),
-  embeddingModel: z.literal('text-embedding-3-small').default('text-embedding-3-small'),
-  embeddingDimension: z.literal(1536).default(1536),
+  embeddingModel: z
+    .union([
+      z.literal('text-embedding-3-small'),
+      z.literal('text-embedding-3-large'),
+      z.string().regex(/^ollama\/.+/, 'Ollama models must be prefixed with "ollama/"'),
+    ])
+    .default('text-embedding-3-small'),
+  embeddingDimension: z.number().int().min(64).max(8192).default(1536),
+  ollamaBaseUrl: z
+    .string()
+    .url('Ollama base URL must be a valid URL')
+    .refine(
+      (url) => {
+        try {
+          const parsed = new URL(url)
+          // Only allow http/https schemes
+          if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
+            return false
+          }
+          const hostname = parsed.hostname.toLowerCase()
+          // Block known cloud metadata endpoints
+          if (hostname === '169.254.169.254' || hostname === 'metadata.google.internal') {
+            return false
+          }
+          // Block IPv6 addresses (except loopback) — prevents IPv6-mapped IPv4 bypass
+          // URL.hostname keeps brackets for IPv6, e.g. "[::ffff:169.254.169.254]"
+          if (hostname.startsWith('[') && hostname !== '[::1]') {
+            return false
+          }
+          // Allow localhost and IPv6 loopback
+          if (hostname === 'localhost' || hostname === '[::1]') {
+            return true
+          }
+          // Allow private IPv4 ranges — only match actual IPs, not domains like "10.evil.com"
+          const ipv4 = /^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/
+          if (ipv4.test(hostname)) {
+            if (
+              hostname.startsWith('127.') ||
+              hostname.startsWith('10.') ||
+              hostname.startsWith('192.168.')
+            ) {
+              return true
+            }
+            if (hostname.startsWith('172.')) {
+              const second = Number.parseInt(hostname.split('.')[1], 10)
+              if (second >= 16 && second <= 31) return true
+            }
+            return false
+          }
+          // Allow Docker service hostnames (no dots = not a public domain)
+          // e.g. "ollama", "host.docker.internal"
+          if (!hostname.includes('.') || hostname.endsWith('.internal')) {
+            return true
+          }
+          return false
+        } catch {
+          return false
+        }
+      },
+      {
+        message:
+          'Ollama base URL must point to localhost, a private network address, or a Docker service hostname',
+      }
+    )
+    .optional(),
   chunkingConfig: z
     .object({
       /** Maximum chunk size in tokens (1 token ≈ 4 characters) */
@@ -98,13 +168,80 @@ export async function POST(req: NextRequest) {
     try {
       const validatedData = CreateKnowledgeBaseSchema.parse(body)
 
+      const { provider, modelName } = parseEmbeddingModel(validatedData.embeddingModel)
+
+      // For Ollama models, validate the model is available and auto-detect dimension
+      let effectiveDimension = validatedData.embeddingDimension
+      if (provider === 'ollama') {
+        const ollamaBaseUrl = getOllamaBaseUrl(validatedData.ollamaBaseUrl)
+        try {
+          const modelInfo = await validateOllamaModel(modelName, ollamaBaseUrl)
+
+          // Auto-correct dimension if the model reports a different one
+          if (modelInfo.embeddingLength && modelInfo.embeddingLength !== effectiveDimension) {
+            if (modelInfo.embeddingLength < 64 || modelInfo.embeddingLength > 8192) {
+              return NextResponse.json(
+                {
+                  error: `Ollama model "${modelName}" reported an unsupported embedding dimension (${modelInfo.embeddingLength}). Supported range: 64–8192.`,
+                },
+                { status: 400 }
+              )
+            }
+            logger.info(
+              `[${requestId}] Auto-correcting embedding dimension from ${effectiveDimension} ` +
+                `to ${modelInfo.embeddingLength} (reported by Ollama model ${modelName})`
+            )
+            effectiveDimension = modelInfo.embeddingLength
+          }
+        } catch {
+          return NextResponse.json(
+            {
+              error:
+                `Cannot reach Ollama at ${ollamaBaseUrl} or model "${modelName}" is not available. ` +
+                `Make sure Ollama is running and the model is pulled (ollama pull ${modelName}).`,
+            },
+            { status: 400 }
+          )
+        }
+      }
+
       const createData = {
         ...validatedData,
+        embeddingDimension: effectiveDimension,
         userId: session.user.id,
       }
 
       const newKnowledgeBase = await createKnowledgeBase(createData, requestId)
 
+      if (provider === 'ollama') {
+        try {
+          await createKBEmbeddingTable(newKnowledgeBase.id, effectiveDimension)
+        } catch (tableError) {
+          logger.error(
+            `[${requestId}] Failed to create embedding table for KB ${newKnowledgeBase.id}`,
+            tableError
+          )
+          // Hard-delete the KB row — this is a creation-time rollback, not a user-initiated
+          // deletion, so a soft delete would leave a restorable broken KB in the archive.
+          try {
+            await dropKBEmbeddingTable(newKnowledgeBase.id)
+            await hardDeleteKnowledgeBase(newKnowledgeBase.id)
+            logger.info(
+              `[${requestId}] Cleaned up orphaned KB ${newKnowledgeBase.id} after table creation failure`
+            )
+          } catch (cleanupError) {
+            logger.error(
+              `[${requestId}] Failed to clean up orphaned KB ${newKnowledgeBase.id}`,
+              cleanupError
+            )
+          }
+          return NextResponse.json(
+            { error: 'Failed to create embedding storage. Please try again.' },
+            { status: 500 }
+          )
+        }
+      }
+
       try {
         PlatformEvents.knowledgeBaseCreated({
           knowledgeBaseId: newKnowledgeBase.id,

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

@@ -20,6 +20,9 @@ const {
   mockGetQueryStrategy,
   mockGenerateSearchEmbedding,
   mockGetDocumentNamesByIds,
+  mockParseEmbeddingModel,
+  mockSearchKBTable,
+  mockSearchKBTableTagOnly,
 } = vi.hoisted(() => ({
   mockDbChain: {
     select: vi.fn().mockReturnThis(),
@@ -42,6 +45,9 @@ const {
   mockGetQueryStrategy: vi.fn(),
   mockGenerateSearchEmbedding: vi.fn(),
   mockGetDocumentNamesByIds: vi.fn(),
+  mockParseEmbeddingModel: vi.fn(),
+  mockSearchKBTable: vi.fn(),
+  mockSearchKBTableTagOnly: vi.fn(),
 }))
 
 vi.mock('drizzle-orm', () => ({
@@ -188,6 +194,16 @@ vi.mock('./utils', () => ({
   },
 }))
 
+vi.mock('@/lib/knowledge/dynamic-tables', () => ({
+  parseEmbeddingModel: mockParseEmbeddingModel,
+  searchKBTable: mockSearchKBTable,
+  searchKBTableTagOnly: mockSearchKBTableTagOnly,
+}))
+
+vi.mock('@/lib/knowledge/embeddings', () => ({
+  generateSearchEmbedding: mockGenerateSearchEmbedding,
+}))
+
 import { estimateTokenCount } from '@/lib/tokenization/estimators'
 import { POST } from '@/app/api/knowledge/search/route'
 import { calculateCost } from '@/providers/utils'
@@ -225,6 +241,18 @@ describe('Knowledge Search API Route', () => {
       }
     })
 
+    // KB config fetch: db.select().from().where() resolves to default single-KB config
+    mockDbChain.where.mockResolvedValue([
+      { id: 'kb-123', embeddingModel: 'text-embedding-3-small', chunkingConfig: {} },
+    ])
+
+    mockParseEmbeddingModel.mockReturnValue({
+      provider: 'openai',
+      modelName: 'text-embedding-3-small',
+    })
+    mockSearchKBTable.mockResolvedValue([])
+    mockSearchKBTableTagOnly.mockResolvedValue([])
+
     mockHandleTagOnlySearch.mockClear()
     mockHandleVectorOnlySearch.mockClear()
     mockHandleTagAndVectorSearch.mockClear()
@@ -337,6 +365,11 @@ describe('Knowledge Search API Route', () => {
         .mockResolvedValueOnce({ hasAccess: true, knowledgeBase: multiKbs[0] })
         .mockResolvedValueOnce({ hasAccess: true, knowledgeBase: multiKbs[1] })
 
+      mockDbChain.where.mockResolvedValue([
+        { id: 'kb-123', embeddingModel: 'text-embedding-3-small', chunkingConfig: {} },
+        { id: 'kb-456', embeddingModel: 'text-embedding-3-small', chunkingConfig: {} },
+      ])
+
       mockDbChain.limit.mockResolvedValue([])
 
       mockHandleVectorOnlySearch.mockResolvedValue(mockSearchResults)
@@ -1008,6 +1041,11 @@ describe('Knowledge Search API Route', () => {
 
       mockHandleTagOnlySearch.mockResolvedValue(mockTaggedResults)
 
+      mockDbChain.where.mockResolvedValue([
+        { id: 'kb-123', embeddingModel: 'text-embedding-3-small', chunkingConfig: {} },
+        { id: 'kb-456', embeddingModel: 'text-embedding-3-small', chunkingConfig: {} },
+      ])
+
       mockDbChain.limit.mockResolvedValueOnce(mockTagDefinitions)
 
       const req = createMockRequest('POST', multiKbTagData)
@@ -1065,13 +1103,6 @@ describe('Knowledge Search API Route', () => {
         'doc-active': 'Active Document.pdf',
       })
 
-      const mockTagDefs = {
-        select: vi.fn().mockReturnThis(),
-        from: vi.fn().mockReturnThis(),
-        where: vi.fn().mockResolvedValue([]),
-      }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
-
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
         query: 'test query',
@@ -1134,15 +1165,6 @@ describe('Knowledge Search API Route', () => {
         'doc-active-tagged': 'Active Tagged Document.pdf',
       })
 
-      const mockTagDefs = {
-        select: vi.fn().mockReturnThis(),
-        from: vi.fn().mockReturnThis(),
-        where: vi
-          .fn()
-          .mockResolvedValue([{ tagSlot: 'tag1', displayName: 'tag1', fieldType: 'text' }]),
-      }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
-
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
         tagFilters: [{ tagName: 'tag1', value: 'api', fieldType: 'text', operator: 'eq' }],
@@ -1207,15 +1229,6 @@ describe('Knowledge Search API Route', () => {
         'doc-active-combined': 'Active Combined Search.pdf',
       })
 
-      const mockTagDefs = {
-        select: vi.fn().mockReturnThis(),
-        from: vi.fn().mockReturnThis(),
-        where: vi
-          .fn()
-          .mockResolvedValue([{ tagSlot: 'tag1', displayName: 'tag1', fieldType: 'text' }]),
-      }
-      mockDbChain.select.mockReturnValueOnce(mockTagDefs)
-
       const req = createMockRequest('POST', {
         knowledgeBaseIds: ['kb-123'],
         query: 'relevant content',