feat: add Fireworks AI provider for chat agent

Routes chat-completions through a managed OpenAI-compatible inference endpoint as an alternative to local Ollama, useful when the agent needs fast multi-hop reasoning that local hardware can't sustain. - backend: rename `call_ollama_chat_messages` → `call_chat_messages`, dispatch by provider; add `call_fireworks` branch (Bearer auth, `response_format: json_object` mapped from internal `format="json"`) and `list_fireworks_models` Tauri command - settings: extend `AiProvider` enum + `AiSettings.fireworks_api_key` (serde-default for legacy config compat); Fireworks base URL hardcoded - UI: provider selector in both popover and AppSettingsSheet (only ollama+fireworks shown; legacy openai/anthropic kept for serde-compat but normalized to ollama in UI); password input + dynamic model list for Fireworks; switching provider clears stale model selection - 4 unit tests: serde round-trip, legacy settings deserialization, Fireworks chat-completions parsing, models-list parsing
2026-05-06 23:04:10 +03:00
parent 532ebf3b44
commit 96a54edcd0
10 changed files with 524 additions and 65 deletions
--- a/src/components/ai/AiSettingsFields.tsx
+++ b/src/components/ai/AiSettingsFields.tsx
@@ -7,27 +7,66 @@ import {
  SelectTrigger,
  SelectValue,
 } from "@/components/ui/select";
-import { useOllamaModels } from "@/hooks/use-ai";
+import { useFireworksModels, useOllamaModels } from "@/hooks/use-ai";
 import { RefreshCw, Loader2 } from "lucide-react";
+import type { AiProvider, OllamaModel } from "@/types";

 interface Props {
+  provider: AiProvider;
  ollamaUrl: string;
  onOllamaUrlChange: (url: string) => void;
+  fireworksApiKey: string;
+  onFireworksApiKeyChange: (key: string) => void;
  model: string;
  onModelChange: (model: string) => void;
 }

 export function AiSettingsFields({
+  provider,
+  ollamaUrl,
+  onOllamaUrlChange,
+  fireworksApiKey,
+  onFireworksApiKeyChange,
+  model,
+  onModelChange,
+}: Props) {
+  if (provider === "fireworks") {
+    return (
+      <FireworksFields
+        apiKey={fireworksApiKey}
+        onApiKeyChange={onFireworksApiKeyChange}
+        model={model}
+        onModelChange={onModelChange}
+      />
+    );
+  }
+
+  return (
+    <OllamaFields
+      ollamaUrl={ollamaUrl}
+      onOllamaUrlChange={onOllamaUrlChange}
+      model={model}
+      onModelChange={onModelChange}
+    />
+  );
+}
+
+function OllamaFields({
  ollamaUrl,
  onOllamaUrlChange,
  model,
  onModelChange,
-}: Props) {
+}: {
+  ollamaUrl: string;
+  onOllamaUrlChange: (url: string) => void;
+  model: string;
+  onModelChange: (model: string) => void;
+}) {
  const {
    data: models,
-    isLoading: modelsLoading,
-    isError: modelsError,
-    refetch: refetchModels,
+    isLoading,
+    isError,
+    refetch,
  } = useOllamaModels(ollamaUrl);

  return (
@@ -42,41 +81,122 @@ export function AiSettingsFields({
        />
      </div>

-      <div className="flex flex-col gap-1.5">
-        <div className="flex items-center justify-between">
-          <label className="text-xs text-muted-foreground">Model</label>
-          <Button
-            size="sm"
-            variant="ghost"
-            className="h-5 w-5 p-0"
-            onClick={() => refetchModels()}
-            disabled={modelsLoading}
-            title="Refresh models"
-          >
-            {modelsLoading ? (
-              <Loader2 className="h-3 w-3 animate-spin" />
-            ) : (
-              <RefreshCw className="h-3 w-3" />
-            )}
-          </Button>
-        </div>
-        {modelsError ? (
-          <p className="text-xs text-destructive">Cannot connect to Ollama</p>
-        ) : (
-          <Select value={model} onValueChange={onModelChange}>
-            <SelectTrigger className="h-8 w-full text-xs">
-              <SelectValue placeholder="Select a model" />
-            </SelectTrigger>
-            <SelectContent>
-              {models?.map((m) => (
-                <SelectItem key={m.name} value={m.name}>
-                  {m.name}
-                </SelectItem>
-              ))}
-            </SelectContent>
-          </Select>
-        )}
-      </div>
+      <ModelDropdown
+        models={models}
+        loading={isLoading}
+        errored={isError}
+        errorText="Cannot connect to Ollama"
+        onRefresh={() => refetch()}
+        model={model}
+        onModelChange={onModelChange}
+      />
    </>
  );
 }
+
+function FireworksFields({
+  apiKey,
+  onApiKeyChange,
+  model,
+  onModelChange,
+}: {
+  apiKey: string;
+  onApiKeyChange: (key: string) => void;
+  model: string;
+  onModelChange: (model: string) => void;
+}) {
+  const {
+    data: models,
+    isLoading,
+    isError,
+    refetch,
+  } = useFireworksModels(apiKey);
+
+  return (
+    <>
+      <div className="flex flex-col gap-1.5">
+        <label className="text-xs text-muted-foreground">Fireworks API key</label>
+        <Input
+          type="password"
+          value={apiKey}
+          onChange={(e) => onApiKeyChange(e.target.value)}
+          placeholder="fw_..."
+          className="h-8 text-xs"
+          autoComplete="off"
+        />
+        <p className="text-[10px] text-muted-foreground/70">
+          Stored locally; sent only to api.fireworks.ai.
+        </p>
+      </div>
+
+      <ModelDropdown
+        models={models}
+        loading={isLoading}
+        errored={isError}
+        errorText="Cannot reach Fireworks (check API key)"
+        onRefresh={() => refetch()}
+        model={model}
+        onModelChange={onModelChange}
+        emptyHint={apiKey.trim() ? "Click ↻ to load models" : "Enter API key first"}
+      />
+    </>
+  );
+}
+
+function ModelDropdown({
+  models,
+  loading,
+  errored,
+  errorText,
+  onRefresh,
+  model,
+  onModelChange,
+  emptyHint,
+}: {
+  models: OllamaModel[] | undefined;
+  loading: boolean;
+  errored: boolean;
+  errorText: string;
+  onRefresh: () => void;
+  model: string;
+  onModelChange: (model: string) => void;
+  emptyHint?: string;
+}) {
+  return (
+    <div className="flex flex-col gap-1.5">
+      <div className="flex items-center justify-between">
+        <label className="text-xs text-muted-foreground">Model</label>
+        <Button
+          size="sm"
+          variant="ghost"
+          className="h-5 w-5 p-0"
+          onClick={onRefresh}
+          disabled={loading}
+          title="Refresh models"
+        >
+          {loading ? (
+            <Loader2 className="h-3 w-3 animate-spin" />
+          ) : (
+            <RefreshCw className="h-3 w-3" />
+          )}
+        </Button>
+      </div>
+      {errored ? (
+        <p className="text-xs text-destructive">{errorText}</p>
+      ) : (
+        <Select value={model} onValueChange={onModelChange}>
+          <SelectTrigger className="h-8 w-full text-xs">
+            <SelectValue placeholder={emptyHint ?? "Select a model"} />
+          </SelectTrigger>
+          <SelectContent>
+            {models?.map((m) => (
+              <SelectItem key={m.name} value={m.name}>
+                {m.name}
+              </SelectItem>
+            ))}
+          </SelectContent>
+        </Select>
+      )}
+    </div>
+  );
+}
--- a/src/components/ai/AiSettingsPopover.tsx
+++ b/src/components/ai/AiSettingsPopover.tsx
@@ -4,25 +4,68 @@ import {
  PopoverContent,
  PopoverTrigger,
 } from "@/components/ui/popover";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
 import { Button } from "@/components/ui/button";
 import { useAiSettings, useSaveAiSettings } from "@/hooks/use-ai";
 import { Settings } from "lucide-react";
 import { toast } from "sonner";
 import { AiSettingsFields } from "./AiSettingsFields";
+import type { AiProvider } from "@/types";
+
+const SUPPORTED_PROVIDERS: { value: AiProvider; label: string }[] = [
+  { value: "ollama", label: "Ollama (local)" },
+  { value: "fireworks", label: "Fireworks AI" },
+];

 export function AiSettingsPopover() {
  const { data: settings } = useAiSettings();
  const saveMutation = useSaveAiSettings();

+  const [provider, setProvider] = useState<AiProvider | null>(null);
  const [url, setUrl] = useState<string | null>(null);
+  const [fireworksKey, setFireworksKey] = useState<string | null>(null);
  const [model, setModel] = useState<string | null>(null);

+  const settingsProvider = settings?.provider;
+  // Hide unsupported legacy values (openai/anthropic) from the selector.
+  const normalizedSettingsProvider: AiProvider | undefined =
+    settingsProvider === "ollama" || settingsProvider === "fireworks"
+      ? settingsProvider
+      : settingsProvider
+        ? "ollama"
+        : undefined;
+
+  const currentProvider: AiProvider =
+    provider ?? normalizedSettingsProvider ?? "ollama";
  const currentUrl = url ?? settings?.ollama_url ?? "http://localhost:11434";
+  const currentFireworksKey =
+    fireworksKey ?? settings?.fireworks_api_key ?? "";
  const currentModel = model ?? settings?.model ?? "";

+  const handleProviderChange = (next: AiProvider) => {
+    if (next === currentProvider) return;
+    setProvider(next);
+    // Model lists differ between providers — drop the previous selection.
+    setModel("");
+  };
+
  const handleSave = () => {
    saveMutation.mutate(
-      { provider: "ollama", ollama_url: currentUrl, model: currentModel },
+      {
+        provider: currentProvider,
+        ollama_url: currentUrl,
+        fireworks_api_key:
+          currentProvider === "fireworks"
+            ? currentFireworksKey.trim() || undefined
+            : settings?.fireworks_api_key,
+        model: currentModel,
+      },
      {
        onSuccess: () => toast.success("AI settings saved"),
        onError: (err) =>
@@ -47,11 +90,33 @@ export function AiSettingsPopover() {
      </PopoverTrigger>
      <PopoverContent className="w-80" align="end">
        <div className="flex flex-col gap-3">
-          <h4 className="text-sm font-medium">Ollama Settings</h4>
+          <h4 className="text-sm font-medium">AI Settings</h4>
+
+          <div className="flex flex-col gap-1.5">
+            <label className="text-xs text-muted-foreground">Provider</label>
+            <Select
+              value={currentProvider}
+              onValueChange={(v) => handleProviderChange(v as AiProvider)}
+            >
+              <SelectTrigger className="h-8 w-full text-xs">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent>
+                {SUPPORTED_PROVIDERS.map((p) => (
+                  <SelectItem key={p.value} value={p.value}>
+                    {p.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>

          <AiSettingsFields
+            provider={currentProvider}
            ollamaUrl={currentUrl}
            onOllamaUrlChange={setUrl}
+            fireworksApiKey={currentFireworksKey}
+            onFireworksApiKeyChange={setFireworksKey}
            model={currentModel}
            onModelChange={setModel}
          />
--- a/src/components/settings/AppSettingsSheet.tsx
+++ b/src/components/settings/AppSettingsSheet.tsx
@@ -22,7 +22,12 @@ import { useAiSettings, useSaveAiSettings } from "@/hooks/use-ai";
 import { AiSettingsFields } from "@/components/ai/AiSettingsFields";
 import { Loader2, Copy, Check } from "lucide-react";
 import { toast } from "sonner";
-import type { AppSettings } from "@/types";
+import type { AiProvider, AppSettings } from "@/types";
+
+const SUPPORTED_AI_PROVIDERS: { value: AiProvider; label: string }[] = [
+  { value: "ollama", label: "Ollama (local)" },
+  { value: "fireworks", label: "Fireworks AI" },
+];

 interface Props {
  open: boolean;
@@ -42,7 +47,9 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
  const [mcpPort, setMcpPort] = useState(9427);

  // AI state
+  const [aiProvider, setAiProvider] = useState<AiProvider>("ollama");
  const [ollamaUrl, setOllamaUrl] = useState("http://localhost:11434");
+  const [fireworksApiKey, setFireworksApiKey] = useState("");
  const [aiModel, setAiModel] = useState("");

  const [copied, setCopied] = useState(false);
@@ -61,11 +68,23 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {
  if (aiSettings !== prevAiSettings) {
    setPrevAiSettings(aiSettings);
    if (aiSettings) {
+      // Legacy openai/anthropic values aren't user-selectable here — fall back to ollama.
+      setAiProvider(
+        aiSettings.provider === "fireworks" ? "fireworks" : "ollama"
+      );
      setOllamaUrl(aiSettings.ollama_url);
+      setFireworksApiKey(aiSettings.fireworks_api_key ?? "");
      setAiModel(aiSettings.model);
    }
  }

+  const handleAiProviderChange = (next: AiProvider) => {
+    if (next === aiProvider) return;
+    setAiProvider(next);
+    // Model lists differ per provider — clear stale selection.
+    setAiModel("");
+  };
+
  const mcpEndpoint = `http://127.0.0.1:${mcpPort}/mcp`;

  const handleCopy = async () => {
@@ -89,7 +108,15 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {

    // Save AI settings separately
    saveAiMutation.mutate(
-      { provider: "ollama", ollama_url: ollamaUrl, model: aiModel },
+      {
+        provider: aiProvider,
+        ollama_url: ollamaUrl,
+        fireworks_api_key:
+          aiProvider === "fireworks"
+            ? fireworksApiKey.trim() || undefined
+            : aiSettings?.fireworks_api_key,
+        model: aiModel,
+      },
      {
        onError: (err) =>
          toast.error("Failed to save AI settings", { description: String(err) }),
@@ -179,19 +206,29 @@ export function AppSettingsSheet({ open, onOpenChange }: Props) {

            <div className="flex flex-col gap-1.5">
              <label className="text-xs text-muted-foreground">Provider</label>
-              <Select value="ollama" disabled>
+              <Select
+                value={aiProvider}
+                onValueChange={(v) => handleAiProviderChange(v as AiProvider)}
+              >
                <SelectTrigger className="h-8 text-xs">
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
-                  <SelectItem value="ollama">Ollama</SelectItem>
+                  {SUPPORTED_AI_PROVIDERS.map((p) => (
+                    <SelectItem key={p.value} value={p.value}>
+                      {p.label}
+                    </SelectItem>
+                  ))}
                </SelectContent>
              </Select>
            </div>

            <AiSettingsFields
+              provider={aiProvider}
              ollamaUrl={ollamaUrl}
              onOllamaUrlChange={setOllamaUrl}
+              fireworksApiKey={fireworksApiKey}
+              onFireworksApiKeyChange={setFireworksApiKey}
              model={aiModel}
              onModelChange={setAiModel}
            />
--- a/src/hooks/use-ai.ts
+++ b/src/hooks/use-ai.ts
@@ -3,6 +3,7 @@ import {
  getAiSettings,
  saveAiSettings,
  listOllamaModels,
+  listFireworksModels,
  generateSql,
  explainSql,
  fixSqlError,
@@ -36,6 +37,16 @@ export function useOllamaModels(ollamaUrl: string | undefined) {
  });
 }

+export function useFireworksModels(apiKey: string | undefined) {
+  return useQuery({
+    queryKey: ["fireworks-models", apiKey],
+    queryFn: () => listFireworksModels(apiKey!),
+    enabled: !!apiKey && apiKey.trim().length > 0,
+    retry: false,
+    staleTime: 60_000,
+  });
+}
+
 export function useGenerateSql() {
  return useMutation({
    mutationFn: ({
--- a/src/lib/tauri.ts
+++ b/src/lib/tauri.ts
@@ -211,6 +211,9 @@ export const saveAiSettings = (settings: AiSettings) =>
 export const listOllamaModels = (ollamaUrl: string) =>
  invoke<OllamaModel[]>("list_ollama_models", { ollamaUrl });

+export const listFireworksModels = (apiKey: string) =>
+  invoke<OllamaModel[]>("list_fireworks_models", { apiKey });
+
 export const generateSql = (connectionId: string, prompt: string) =>
  invoke<string>("generate_sql", { connectionId, prompt });

--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -134,13 +134,14 @@ export interface SavedQuery {
  created_at: string;
 }

-export type AiProvider = "ollama" | "openai" | "anthropic";
+export type AiProvider = "ollama" | "openai" | "anthropic" | "fireworks";

 export interface AiSettings {
  provider: AiProvider;
  ollama_url: string;
  openai_api_key?: string;
  anthropic_api_key?: string;
+  fireworks_api_key?: string;
  model: string;
 }