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', () => ({
   getKnowledgeBaseById: vi.fn(),
   updateKnowledgeBase: vi.fn(),

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

@@ -5,6 +5,7 @@ import { AuditAction, AuditResourceType, recordAudit } from '@/lib/audit/log'
 import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid'
 import { PlatformEvents } from '@/lib/core/telemetry'
 import { generateRequestId } from '@/lib/core/utils/request'
+import { dropKBEmbeddingTable } from '@/lib/knowledge/dynamic-tables'
 import {
   deleteKnowledgeBase,
   getKnowledgeBaseById,
@@ -25,8 +26,6 @@ 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({
@@ -200,6 +199,8 @@ export async function DELETE(
     }
 
     await deleteKnowledgeBase(id, requestId)
+    // Drop the per-KB embedding table if this was an Ollama KB (no-op for OpenAI KBs)
+    await dropKBEmbeddingTable(id)
 
     try {
       PlatformEvents.knowledgeBaseDeleted({

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

@@ -5,7 +5,17 @@ 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 { createKnowledgeBase, getKnowledgeBases } from '@/lib/knowledge/service'
+import {
+  createKBEmbeddingTable,
+  dropKBEmbeddingTable,
+  parseEmbeddingModel,
+} from '@/lib/knowledge/dynamic-tables'
+import { getOllamaBaseUrl, validateOllamaModel } from '@/lib/knowledge/embeddings'
+import {
+  createKnowledgeBase,
+  deleteKnowledgeBase,
+  getKnowledgeBases,
+} from '@/lib/knowledge/service'
 
 const logger = createLogger('KnowledgeBaseAPI')
 
@@ -21,8 +31,66 @@ 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, loopback, and private network ranges
+          if (
+            hostname === 'localhost' ||
+            hostname === '[::1]' ||
+            hostname.startsWith('127.') ||
+            hostname.startsWith('10.') ||
+            hostname.startsWith('192.168.')
+          ) {
+            return true
+          }
+          // Allow 172.16.0.0 – 172.31.255.255
+          if (hostname.startsWith('172.')) {
+            const second = Number.parseInt(hostname.split('.')[1], 10)
+            if (second >= 16 && second <= 31) return true
+          }
+          // 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) */
@@ -89,13 +157,79 @@ 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
+          )
+          // Clean up the orphaned KB row and any partially-created table
+          try {
+            await dropKBEmbeddingTable(newKnowledgeBase.id)
+            await deleteKnowledgeBase(newKnowledgeBase.id, requestId)
+            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

@@ -25,6 +25,9 @@ const {
   mockGetQueryStrategy,
   mockGenerateSearchEmbedding,
   mockGetDocumentNamesByIds,
+  mockParseEmbeddingModel,
+  mockSearchKBTable,
+  mockSearchKBTableTagOnly,
 } = vi.hoisted(() => ({
   mockDbChain: {
     select: vi.fn().mockReturnThis(),
@@ -47,6 +50,9 @@ const {
   mockGetQueryStrategy: vi.fn(),
   mockGenerateSearchEmbedding: vi.fn(),
   mockGetDocumentNamesByIds: vi.fn(),
+  mockParseEmbeddingModel: vi.fn(),
+  mockSearchKBTable: vi.fn(),
+  mockSearchKBTableTagOnly: vi.fn(),
 }))
 
 vi.mock('drizzle-orm', () => ({
@@ -126,6 +132,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'
@@ -163,6 +179,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()
@@ -275,6 +303,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)
@@ -946,6 +979,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)
@@ -1003,13 +1041,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',
@@ -1072,15 +1103,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' }],
@@ -1145,15 +1167,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',