Adding Ollama Support

src/main/index.ts

@@ -16,6 +16,10 @@ import {
   setClaudeApiKey,
   getClaudeModel,
   setClaudeModel,
+  getOllamaBaseUrl,
+  setOllamaBaseUrl,
+  getOllamaModel,
+  setOllamaModel,
   // Connection management functions
   createConnection,
   editConnection,
@@ -119,6 +123,7 @@ app.whenReady().then(() => {
     optimizer.watchWindowShortcuts(window)
   })
 
+  // OpenAI handlers
   ipcMain.handle('getOpenAiKey', async () => {
     return (await getOpenAiKey()) ?? ''
   })
@@ -151,6 +156,7 @@ app.whenReady().then(() => {
     await setAiProvider(aiProvider)
   })
 
+  // Claude handlers
   ipcMain.handle('getClaudeApiKey', async () => {
     return (await getClaudeApiKey()) ?? ''
   })
@@ -167,6 +173,70 @@ app.whenReady().then(() => {
     await setClaudeModel(claudeModel)
   })
 
+  // Ollama handlers
+  ipcMain.handle('getOllamaBaseUrl', async () => {
+    return (await getOllamaBaseUrl()) ?? 'http://localhost:11434/v1/'
+  })
+
+  ipcMain.handle('setOllamaBaseUrl', async (_, ollamaBaseUrl) => {
+    await setOllamaBaseUrl(ollamaBaseUrl)
+  })
+
+  ipcMain.handle('getOllamaModel', async () => {
+    return (await getOllamaModel()) ?? ''
+  })
+
+  ipcMain.handle('setOllamaModel', async (_, ollamaModel) => {
+    await setOllamaModel(ollamaModel)
+  })
+
+  // Add handler to list available Ollama models
+  ipcMain.handle('getOllamaModels', async () => {
+    try {
+      const baseUrl = await getOllamaBaseUrl() || 'http://localhost:11434/v1/'
+      const response = await fetch(`${baseUrl}/api/tags`)
+      if (!response.ok) {
+        throw new Error(`Failed to fetch Ollama models: ${response.statusText}`)
+      }
+      const data = await response.json()
+      return {
+        success: true,
+        models: data.models || []
+      }
+    } catch (error: any) {
+      return {
+        success: false,
+        error: error.message,
+        models: []
+      }
+    }
+  })
+
+  // Test Ollama connection
+  ipcMain.handle('testOllamaConnection', async (_, baseUrl) => {
+    try {
+      const response = await fetch(`${baseUrl}/api/tags`, {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+      })
+
+      if (!response.ok) {
+        throw new Error(`Connection failed: ${response.statusText}`)
+      }
+
+      return { success: true }
+    } catch (error: any) {
+      return { 
+        success: false, 
+        error: error.message.includes('fetch') 
+          ? 'Cannot connect to Ollama. Make sure Ollama is running and accessible.'
+          : error.message 
+      }
+    }
+  })
+
   // Connection management handlers
   ipcMain.handle('createConnection', async (_, name, connectionMetadata) => {
     try {
@@ -327,17 +397,21 @@ app.whenReady().then(() => {
       const aiProvider = await getAiProvider()
       const promptExtension = await getConnectionPromptExtension(name)
 
-      let apiKey: string
+      let apiKey: string = '' // Initialize with empty string
       let model: string | undefined
-      let openAiBaseUrl: string | undefined
+      let baseUrl: string | undefined
 
       if (aiProvider === 'openai') {
         apiKey = (await getOpenAiKey()) ?? ''
         model = await getOpenAiModel()
-        openAiBaseUrl = await getOpenAiBaseUrl()
-      } else {
+        baseUrl = await getOpenAiBaseUrl()
+      } else if (aiProvider === 'claude') {
         apiKey = (await getClaudeApiKey()) ?? ''
         model = await getClaudeModel()
+      } else if (aiProvider === 'ollama') {
+        apiKey = '' // Ollama doesn't require API key
+        model = await getOllamaModel()
+        baseUrl = await getOllamaBaseUrl()
       }
 
       const query = await generateQuery(
@@ -347,7 +421,7 @@ app.whenReady().then(() => {
         apiKey,
         existingQuery,
         promptExtension ?? '',
-        openAiBaseUrl,
+        baseUrl,
         model
       )
       return {
@@ -416,4 +490,4 @@ app.on('window-all-closed', () => {
 })
 
 // In this file you can include the rest of your app's specific main process
-// code. You can also put them in separate files and require them here.
+// code. You can also put them in separate files and require them here.

src/main/lib/ai.ts

@@ -13,7 +13,7 @@ export type QueryResponse = {
 export async function generateQuery(
   input: string,
   connectionString: string,
-  aiProvider: 'openai' | 'claude',
+  aiProvider: 'openai' | 'claude' | 'ollama',
   apiKey: string,
   existingQuery: string,
   promptExtension: string,
@@ -34,17 +34,22 @@ export async function generateQuery(
 
     if (aiProvider === 'openai') {
       const openai = createOpenAI({
-        apiKey: apiKey,
+        apiKey,
         baseURL: openAiUrl || undefined
       })
       modelToUse = model || 'gpt-4o'
       aiModel = openai(modelToUse)
     } else if (aiProvider === 'claude') {
-      const anthropic = createAnthropic({
-        apiKey: apiKey
-      })
+      const anthropic = createAnthropic({ apiKey })
       modelToUse = model || 'claude-sonnet-4-20250514'
       aiModel = anthropic(modelToUse)
+    } else if (aiProvider === 'ollama') {
+      const openai = createOpenAI({
+        apiKey: 'ollama-placeholder', // required by SDK, even if not used
+        baseURL: openAiUrl || 'http://localhost:11434/v1'
+      })
+      modelToUse = model || 'llama3'
+      aiModel = openai(modelToUse)
     } else {
       throw new Error(`Unsupported AI provider: ${aiProvider}`)
     }
@@ -61,7 +66,7 @@ export async function generateQuery(
       ${promptExtension.length > 0 ? `Extra information: ${promptExtension}` : ``}
 
       format the query in a way that is easy to read and understand.
-      ${dbType === 'postgres' ? 'wrap table names in double quotes' : ''}
+      ${dbType === 'postgres' ? (aiProvider === 'ollama' ? 'Wrap both table names and column names in double quotes, especially if they contain uppercase letters or special characters.' : 'Wrap table names in double quotes.') : ''}
       if the query results can be effectively visualized using a graph, specify which column should be used for the x-axis (domain) and which column(s) should be used for the y-axis (range).
     `
 

src/main/lib/state.ts

@@ -31,14 +31,17 @@ const favoritesSchema = z.object({
 
 // Global settings schema (AI provider settings only)
 const globalSettingsSchema = z.object({
-  aiProvider: z.enum(['openai', 'claude']).optional(),
+  aiProvider: z.enum(['openai', 'claude', 'ollama']).optional(),
   openAiKey: z.string().optional(),
   openAiBaseUrl: z.string().optional(),
   openAiModel: z.string().optional(),
   claudeApiKey: z.string().optional(),
-  claudeModel: z.string().optional()
+  claudeModel: z.string().optional(),
+  ollamaBaseUrl: z.string().optional(),
+  ollamaModel: z.string().optional()
 })
 
+
 // Connection settings schema
 const connectionSettingsSchema = z.object({
   connectionString: z.string(),
@@ -51,9 +54,12 @@ const defaultGlobalSettings: z.infer<typeof globalSettingsSchema> = {
   openAiBaseUrl: undefined,
   openAiModel: undefined,
   claudeApiKey: undefined,
-  claudeModel: undefined
+  claudeModel: undefined,
+  ollamaBaseUrl: undefined,
+  ollamaModel: undefined
 }
 
+
 function rootDir() {
   // Allow tests to override the root directory
   if (process.env.SNAPQL_TEST_ROOT) {
@@ -194,7 +200,7 @@ export async function getAiProvider() {
   return settings.aiProvider || 'openai'
 }
 
-export async function setAiProvider(aiProvider: 'openai' | 'claude') {
+export async function setAiProvider(aiProvider: 'openai' | 'claude' | 'ollama') {
   const settings = await getGlobalSettings()
   settings.aiProvider = aiProvider
   await setGlobalSettings(settings)
@@ -221,6 +227,27 @@ export async function setClaudeModel(claudeModel: string) {
   settings.claudeModel = claudeModel
   await setGlobalSettings(settings)
 }
+export async function getOllamaBaseUrl() {
+  const settings = await getGlobalSettings()
+  return settings.ollamaBaseUrl
+}
+
+export async function setOllamaBaseUrl(ollamaBaseUrl: string) {
+  const settings = await getGlobalSettings()
+  settings.ollamaBaseUrl = ollamaBaseUrl
+  await setGlobalSettings(settings)
+}
+
+export async function getOllamaModel() {
+  const settings = await getGlobalSettings()
+  return settings.ollamaModel
+}
+
+export async function setOllamaModel(ollamaModel: string) {
+  const settings = await getGlobalSettings()
+  settings.ollamaModel = ollamaModel
+  await setGlobalSettings(settings)
+}
 
 // Connection management
 export async function createConnection(

src/preload/index.d.ts

@@ -15,12 +15,19 @@ declare global {
       setOpenAiBaseUrl: (openAiBaseUrl: string) => Promise<void>
       getOpenAiModel: () => Promise<string>
       setOpenAiModel: (openAiModel: string) => Promise<void>
-      getAiProvider: () => Promise<'openai' | 'claude'>
-      setAiProvider: (aiProvider: 'openai' | 'claude') => Promise<void>
+      getAiProvider: () => Promise<'openai' | 'claude' | 'ollama'>
+      setAiProvider: (aiProvider: 'openai' | 'claude' | 'ollama') => Promise<void>
       getClaudeApiKey: () => Promise<string>
       setClaudeApiKey: (claudeApiKey: string) => Promise<void>
       getClaudeModel: () => Promise<string>
       setClaudeModel: (claudeModel: string) => Promise<void>
+      getOllamaBaseUrl: () => Promise<string>
+      setOllamaBaseUrl: (ollamaBaseUrl: string) => Promise<void>
+      getOllamaModel: () => Promise<string>
+      setOllamaModel: (ollamaModel: string) => Promise<void>
+      // Add missing Ollama-related functions
+      getOllamaModels: () => Promise<{ success: boolean; models: any[]; error?: string }>
+      testOllamaConnection: (baseUrl: string) => Promise<{ success: boolean; error?: string }>
 
       // Connection management
       createConnection: (name: string, connectionMetadata: any) => Promise<void>

src/preload/index.ts

@@ -17,14 +17,24 @@ try {
     setOpenAiModel: async (openAiModel: string) =>
       await ipcRenderer.invoke('setOpenAiModel', openAiModel),
     getAiProvider: async () => await ipcRenderer.invoke('getAiProvider'),
-    setAiProvider: async (aiProvider: 'openai' | 'claude') =>
+    setAiProvider: async (aiProvider: 'openai' | 'claude' | 'ollama') =>
       await ipcRenderer.invoke('setAiProvider', aiProvider),
     getClaudeApiKey: async () => await ipcRenderer.invoke('getClaudeApiKey'),
     setClaudeApiKey: async (claudeApiKey: string) =>
       await ipcRenderer.invoke('setClaudeApiKey', claudeApiKey),
     getClaudeModel: async () => await ipcRenderer.invoke('getClaudeModel'),
     setClaudeModel: async (claudeModel: string) =>
       await ipcRenderer.invoke('setClaudeModel', claudeModel),
+    getOllamaBaseUrl: async () => await ipcRenderer.invoke('getOllamaBaseUrl'),
+    setOllamaBaseUrl: async (ollamaBaseUrl: string) =>
+      await ipcRenderer.invoke('setOllamaBaseUrl', ollamaBaseUrl),
+    getOllamaModel: async () => await ipcRenderer.invoke('getOllamaModel'),
+    setOllamaModel: async (ollamaModel: string) =>
+      await ipcRenderer.invoke('setOllamaModel', ollamaModel),
+    // Add missing Ollama-related functions
+    getOllamaModels: async () => await ipcRenderer.invoke('getOllamaModels'),
+    testOllamaConnection: async (baseUrl: string) =>
+      await ipcRenderer.invoke('testOllamaConnection', baseUrl),
 
     // Connection management
     createConnection: async (name: string, connectionMetadata: any) =>
@@ -65,4 +75,4 @@ try {
   })
 } catch (error) {
   console.error(error)
-}
+}