Merge pull request #79 from Yashh56/Yashh56/fix

feat: enhance UI and Implemented Persistance Storage for LLMs API KEYS

Author: Yashh56

bridge/src/ai/providers/anthropic.provider.ts

@@ -14,15 +14,17 @@ const DEFAULT_MODEL = "claude-3-5-haiku-20241022";
 
 export class AnthropicProvider implements AIProvider {
   private client: Anthropic;
+  private model: string;
 
-  constructor(apiKey: string) {
+  constructor(apiKey: string, model?: string) {
     this.client = new Anthropic({ apiKey });
+    this.model = model?.trim() || DEFAULT_MODEL;
   }
 
   private async complete(system: string, user: string): Promise<string> {
     try {
       const msg = await this.client.messages.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         max_tokens: 4096,
         system,
         messages: [{ role: "user", content: user }],
@@ -53,7 +55,7 @@ export class AnthropicProvider implements AIProvider {
   async testConnection(): Promise<string> {
     try {
       await this.client.messages.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         max_tokens: 10,
         messages: [{ role: "user", content: "ping" }],
       });

bridge/src/ai/providers/gemini.provider.ts

@@ -14,15 +14,17 @@ const DEFAULT_MODEL = "gemini-1.5-flash";
 
 export class GeminiProvider implements AIProvider {
   private genAI: GoogleGenerativeAI;
+  private model: string;
 
-  constructor(apiKey: string) {
+  constructor(apiKey: string, model?: string) {
     this.genAI = new GoogleGenerativeAI(apiKey);
+    this.model = model?.trim() || DEFAULT_MODEL;
   }
 
   private async complete(system: string, user: string): Promise<string> {
     try {
       const model = this.genAI.getGenerativeModel({
-        model: DEFAULT_MODEL,
+        model: this.model,
         systemInstruction: system,
         generationConfig: { maxOutputTokens: 4096 },
       });
@@ -51,7 +53,7 @@ export class GeminiProvider implements AIProvider {
 
   async testConnection(): Promise<string> {
     try {
-      const model = this.genAI.getGenerativeModel({ model: DEFAULT_MODEL });
+      const model = this.genAI.getGenerativeModel({ model: this.model });
       await model.generateContent("ping");
       return "";
     } catch (err) {

bridge/src/ai/providers/groq.provider.ts

@@ -14,15 +14,17 @@ const DEFAULT_MODEL = "llama-3.3-70b-versatile";
 
 export class GroqProvider implements AIProvider {
   private client: Groq;
+  private model: string;
 
-  constructor(apiKey: string) {
+  constructor(apiKey: string, model?: string) {
     this.client = new Groq({ apiKey });
+    this.model = model?.trim() || DEFAULT_MODEL;
   }
 
   private async complete(system: string, user: string): Promise<string> {
     try {
       const res = await this.client.chat.completions.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [
           { role: "system", content: system },
           { role: "user", content: user },
@@ -54,7 +56,7 @@ export class GroqProvider implements AIProvider {
   async testConnection(): Promise<string> {
     try {
       await this.client.chat.completions.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [{ role: "user", content: "ping" }],
         max_tokens: 5,
       });

bridge/src/ai/providers/mistral.provider.ts

@@ -14,15 +14,17 @@ const DEFAULT_MODEL = "mistral-small-latest";
 
 export class MistralProvider implements AIProvider {
   private client: Mistral;
+  private model: string;
 
-  constructor(apiKey: string) {
+  constructor(apiKey: string, model?: string) {
     this.client = new Mistral({ apiKey });
+    this.model = model?.trim() || DEFAULT_MODEL;
   }
 
   private async complete(system: string, user: string): Promise<string> {
     try {
       const res = await this.client.chat.complete({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [
           { role: "system", content: system },
           { role: "user", content: user },
@@ -60,7 +62,7 @@ export class MistralProvider implements AIProvider {
   async testConnection(): Promise<string> {
     try {
       await this.client.chat.complete({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [{ role: "user", content: "ping" }],
         maxTokens: 5,
       });

bridge/src/ai/providers/openai.provider.ts

@@ -14,15 +14,17 @@ const DEFAULT_MODEL = "gpt-4o-mini";
 
 export class OpenAIProvider implements AIProvider {
   private client: OpenAI;
+  private model: string;
 
-  constructor(apiKey: string) {
+  constructor(apiKey: string, model?: string) {
     this.client = new OpenAI({ apiKey });
+    this.model = model?.trim() || DEFAULT_MODEL;
   }
 
   private async complete(system: string, user: string): Promise<string> {
     try {
       const res = await this.client.chat.completions.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [
           { role: "system", content: system },
           { role: "user", content: user },
@@ -54,7 +56,7 @@ export class OpenAIProvider implements AIProvider {
   async testConnection(): Promise<string> {
     try {
       await this.client.chat.completions.create({
-        model: DEFAULT_MODEL,
+        model: this.model,
         messages: [{ role: "user", content: "ping" }],
         max_tokens: 5,
       });

bridge/src/handlers/aiHandlers.ts

@@ -7,6 +7,7 @@ import {
   AIExplainQueryParams,
   AIRecommendChartParams,
   AITestConnectionParams,
+  AISettings,
 } from "../types/ai";
 import {
   getOrCall,
@@ -19,6 +20,9 @@ import { buildSchemaAnalysisPrompt } from "../ai/prompts/schema-analysis";
 import { buildQueryExplanationPrompt } from "../ai/prompts/query-explanation";
 import { buildChartRecommendationPrompt } from "../ai/prompts/chart-recommendation";
 import { parseChartRecommendation } from "../ai/prompts/chart-recommendation";
+import fs from "fs/promises";
+import fsSync from "fs";
+import { AI_SETTINGS_FILE, CONFIG_FOLDER, ensureDir } from "../utils/config";
 
 export class AIHandlers {
   private aiService: AIService;
@@ -219,4 +223,43 @@ export class AIHandlers {
       this.rpc.sendError(id, { code: "HISTORY_ERROR", message: err?.message ?? String(err) });
     }
   }
+
+  // ── Settings persistence (reads/writes ai-settings.json) ──────────────
+
+  async handleLoadSettings(_params: unknown, id: number | string) {
+    try {
+      ensureDir(CONFIG_FOLDER);
+      if (!fsSync.existsSync(AI_SETTINGS_FILE)) {
+        // Return empty object — frontend will fall back to defaults
+        this.rpc.sendResponse(id, { ok: true, data: {} });
+        return;
+      }
+      const raw = await fs.readFile(AI_SETTINGS_FILE, "utf-8");
+      const settings = JSON.parse(raw) as AISettings;
+      this.rpc.sendResponse(id, { ok: true, data: settings });
+    } catch (err: any) {
+      this.logger?.warn({ err }, "ai.loadSettings failed — returning empty");
+      // Non-fatal: return empty so the app still starts
+      this.rpc.sendResponse(id, { ok: true, data: {} });
+    }
+  }
+
+  async handleSaveSettings(params: { settings: AISettings }, id: number | string) {
+    try {
+      ensureDir(CONFIG_FOLDER);
+      await fs.writeFile(
+        AI_SETTINGS_FILE,
+        JSON.stringify(params.settings, null, 2),
+        "utf-8"
+      );
+      // On non-Windows platforms, restrict file permissions (contains API keys)
+      if (process.platform !== "win32") {
+        await fs.chmod(AI_SETTINGS_FILE, 0o600);
+      }
+      this.rpc.sendResponse(id, { ok: true, data: { saved: true } });
+    } catch (err: any) {
+      this.logger?.error({ err }, "ai.saveSettings failed");
+      this.rpc.sendError(id, { code: "SAVE_ERROR", message: err?.message ?? String(err) });
+    }
+  }
 }

bridge/src/jsonRpcHandler.ts

@@ -389,6 +389,12 @@ export function registerDbHandlers(
   rpcRegister(rpc, "ai.clearHistory", (p, id) =>
     aiHandlers.handleClearHistory(p, id)
   );
+  rpcRegister(rpc, "ai.loadSettings", (p, id) =>
+    aiHandlers.handleLoadSettings(p, id)
+  );
+  rpcRegister(rpc, "ai.saveSettings", (p, id) =>
+    aiHandlers.handleSaveSettings(p, id)
+  );
 
   logger?.info("All RPC handlers registered successfully");
 }

bridge/src/services/ai.impl.ts

@@ -25,27 +25,27 @@ export class AIServiceImpl {
       case "anthropic": {
         const key = settings.anthropicApiKey?.trim();
         if (!key) throw new AIError("MISSING_API_KEY", "anthropic", "Anthropic API key is not configured.");
-        return new AnthropicProvider(key);
+        return new AnthropicProvider(key, settings.anthropicModel);
       }
       case "openai": {
         const key = settings.openaiApiKey?.trim();
         if (!key) throw new AIError("MISSING_API_KEY", "openai", "OpenAI API key is not configured.");
-        return new OpenAIProvider(key);
+        return new OpenAIProvider(key, settings.openaiModel);
       }
       case "gemini": {
         const key = settings.geminiApiKey?.trim();
         if (!key) throw new AIError("MISSING_API_KEY", "gemini", "Gemini API key is not configured.");
-        return new GeminiProvider(key);
+        return new GeminiProvider(key, settings.geminiModel);
       }
       case "groq": {
         const key = settings.groqApiKey?.trim();
         if (!key) throw new AIError("MISSING_API_KEY", "groq", "Groq API key is not configured.");
-        return new GroqProvider(key);
+        return new GroqProvider(key, settings.groqModel);
       }
       case "mistral": {
         const key = settings.mistralApiKey?.trim();
         if (!key) throw new AIError("MISSING_API_KEY", "mistral", "Mistral API key is not configured.");
-        return new MistralProvider(key);
+        return new MistralProvider(key, settings.mistralModel);
       }
       case "ollama": {
         return new OllamaProvider(settings.ollamaBaseUrl, settings.ollamaModel);

bridge/src/services/aiCacheService.ts

@@ -98,15 +98,23 @@ export function hashChartRecommendation(input: ChartRecommendationInput, datasou
 
 /**
  * Resolve the model name from the settings based on provider.
- * This is best-effort — some providers don't expose the model in settings.
  */
 function resolveModelName(settings: AISettings): string {
-  const provider = settings.defaultProvider;
-  switch (provider) {
+  switch (settings.defaultProvider) {
+    case "anthropic":
+      return settings.anthropicModel ?? "claude-3-5-haiku-20241022";
+    case "openai":
+      return settings.openaiModel ?? "gpt-4o-mini";
+    case "gemini":
+      return settings.geminiModel ?? "gemini-1.5-flash";
+    case "groq":
+      return settings.groqModel ?? "llama-3.3-70b-versatile";
+    case "mistral":
+      return settings.mistralModel ?? "mistral-small-latest";
     case "ollama":
-      return settings.ollamaModel ?? "ollama-default";
+      return settings.ollamaModel ?? "llama3.2";
     default:
-      return provider; // For API-key providers, the model is selected by the SDK
+      return settings.defaultProvider;
   }
 }
 

bridge/src/types/ai.ts

@@ -21,6 +21,12 @@ export interface AISettings {
   mistralApiKey?: string;
   ollamaBaseUrl?: string;
   ollamaModel?: string;
+  // Per-provider selected model (overrides provider default)
+  anthropicModel?: string;
+  openaiModel?: string;
+  geminiModel?: string;
+  groqModel?: string;
+  mistralModel?: string;
 }
 
 // ── Feature input/output types ────────────────────────────────────────────