initial commit

src/commands/web-search-config.ts

@@ -0,0 +1,229 @@
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import {
+  getDefaultWebSearchConfigPath,
+  readRawWebSearchConfig,
+  writeWebSearchConfig,
+  WebSearchConfigError,
+} from "../config.ts";
+import type { WebSearchConfig, WebSearchProviderConfig } from "../schema.ts";
+
+export function createDefaultWebSearchConfig(input: { tavilyName: string; tavilyApiKey: string }): WebSearchConfig {
+  return {
+    defaultProvider: input.tavilyName,
+    providers: [
+      {
+        name: input.tavilyName,
+        type: "tavily",
+        apiKey: input.tavilyApiKey,
+      },
+    ],
+  };
+}
+
+export function setDefaultProviderOrThrow(config: WebSearchConfig, providerName: string): WebSearchConfig {
+  if (!config.providers.some((provider) => provider.name === providerName)) {
+    throw new Error(`Unknown provider: ${providerName}`);
+  }
+  return { ...config, defaultProvider: providerName };
+}
+
+export function renameProviderOrThrow(
+  config: WebSearchConfig,
+  currentName: string,
+  nextName: string,
+): WebSearchConfig {
+  if (!nextName.trim()) {
+    throw new Error("Provider name cannot be blank.");
+  }
+  if (config.providers.some((provider) => provider.name === nextName && provider.name !== currentName)) {
+    throw new Error(`Duplicate provider name: ${nextName}`);
+  }
+
+  return {
+    defaultProvider: config.defaultProvider === currentName ? nextName : config.defaultProvider,
+    providers: config.providers.map((provider) =>
+      provider.name === currentName ? { ...provider, name: nextName } : provider,
+    ),
+  };
+}
+
+export function updateProviderOrThrow(
+  config: WebSearchConfig,
+  providerName: string,
+  patch: { apiKey?: string; options?: WebSearchProviderConfig["options"] },
+): WebSearchConfig {
+  const existing = config.providers.find((provider) => provider.name === providerName);
+  if (!existing) {
+    throw new Error(`Unknown provider: ${providerName}`);
+  }
+  if (patch.apiKey !== undefined && !patch.apiKey.trim()) {
+    throw new Error("Provider apiKey cannot be blank.");
+  }
+
+  return {
+    ...config,
+    providers: config.providers.map((provider) =>
+      provider.name === providerName
+        ? {
+            ...provider,
+            apiKey: patch.apiKey ?? provider.apiKey,
+            options: patch.options ?? provider.options,
+          }
+        : provider,
+    ),
+  };
+}
+
+export function removeProviderOrThrow(config: WebSearchConfig, providerName: string): WebSearchConfig {
+  if (config.providers.length === 1) {
+    throw new Error("Cannot remove the last provider.");
+  }
+  if (config.defaultProvider === providerName) {
+    throw new Error("Cannot remove the default provider before selecting a new default.");
+  }
+  return {
+    ...config,
+    providers: config.providers.filter((provider) => provider.name !== providerName),
+  };
+}
+
+function upsertProviderOrThrow(config: WebSearchConfig, nextProvider: WebSearchProviderConfig): WebSearchConfig {
+  if (!nextProvider.name.trim()) {
+    throw new Error("Provider name cannot be blank.");
+  }
+  if (!nextProvider.apiKey.trim()) {
+    throw new Error("Provider apiKey cannot be blank.");
+  }
+
+  const withoutSameName = config.providers.filter((provider) => provider.name !== nextProvider.name);
+  return {
+    ...config,
+    providers: [...withoutSameName, nextProvider],
+  };
+}
+
+async function promptProviderOptions(ctx: any, provider: WebSearchProviderConfig) {
+  const defaultSearchLimit = await ctx.ui.input(
+    `Default search limit for ${provider.name}`,
+    provider.options?.defaultSearchLimit !== undefined ? String(provider.options.defaultSearchLimit) : "",
+  );
+  const defaultFetchTextMaxCharacters = await ctx.ui.input(
+    `Default fetch text max characters for ${provider.name}`,
+    provider.options?.defaultFetchTextMaxCharacters !== undefined
+      ? String(provider.options.defaultFetchTextMaxCharacters)
+      : "",
+  );
+
+  const options = {
+    defaultSearchLimit: defaultSearchLimit ? Number(defaultSearchLimit) : undefined,
+    defaultFetchTextMaxCharacters: defaultFetchTextMaxCharacters
+      ? Number(defaultFetchTextMaxCharacters)
+      : undefined,
+  };
+
+  return Object.values(options).some((value) => value !== undefined) ? options : undefined;
+}
+
+export function registerWebSearchConfigCommand(pi: ExtensionAPI) {
+  pi.registerCommand("web-search-config", {
+    description: "Configure Tavily/Exa providers for web_search and web_fetch",
+    handler: async (_args, ctx) => {
+      const path = getDefaultWebSearchConfigPath();
+
+      let config: WebSearchConfig;
+      try {
+        config = await readRawWebSearchConfig(path);
+      } catch (error) {
+        if (!(error instanceof WebSearchConfigError)) {
+          throw error;
+        }
+
+        const tavilyName = await ctx.ui.input("Create Tavily provider", "tavily-main");
+        const tavilyApiKey = await ctx.ui.input("Tavily API key", "tvly-...");
+        if (!tavilyName || !tavilyApiKey) {
+          return;
+        }
+        config = createDefaultWebSearchConfig({ tavilyName, tavilyApiKey });
+      }
+
+      const action = await ctx.ui.select("Web search config", [
+        "Set default provider",
+        "Add Tavily provider",
+        "Add Exa provider",
+        "Edit provider",
+        "Remove provider",
+      ]);
+      if (!action) {
+        return;
+      }
+
+      if (action === "Set default provider") {
+        const nextDefault = await ctx.ui.select(
+          "Choose default provider",
+          config.providers.map((provider) => provider.name),
+        );
+        if (!nextDefault) {
+          return;
+        }
+        config = setDefaultProviderOrThrow(config, nextDefault);
+      }
+
+      if (action === "Add Tavily provider") {
+        const name = await ctx.ui.input("Provider name", "tavily-main");
+        const apiKey = await ctx.ui.input("Tavily API key", "tvly-...");
+        if (!name || !apiKey) {
+          return;
+        }
+        config = upsertProviderOrThrow(config, { name, type: "tavily", apiKey });
+      }
+
+      if (action === "Add Exa provider") {
+        const name = await ctx.ui.input("Provider name", "exa-fallback");
+        const apiKey = await ctx.ui.input("Exa API key", "exa_...");
+        if (!name || !apiKey) {
+          return;
+        }
+        config = upsertProviderOrThrow(config, { name, type: "exa", apiKey });
+      }
+
+      if (action === "Edit provider") {
+        const providerName = await ctx.ui.select(
+          "Choose provider",
+          config.providers.map((provider) => provider.name),
+        );
+        if (!providerName) {
+          return;
+        }
+
+        const existing = config.providers.find((provider) => provider.name === providerName)!;
+        const nextName = await ctx.ui.input("Provider name", existing.name);
+        const nextApiKey = await ctx.ui.input(`API key for ${existing.name}`, existing.apiKey);
+        if (!nextName || !nextApiKey) {
+          return;
+        }
+
+        config = renameProviderOrThrow(config, existing.name, nextName);
+        const renamed = config.providers.find((provider) => provider.name === nextName)!;
+        const nextOptions = await promptProviderOptions(ctx, renamed);
+        config = updateProviderOrThrow(config, nextName, {
+          apiKey: nextApiKey,
+          options: nextOptions,
+        });
+      }
+
+      if (action === "Remove provider") {
+        const providerName = await ctx.ui.select(
+          "Choose provider to remove",
+          config.providers.map((provider) => provider.name),
+        );
+        if (!providerName) {
+          return;
+        }
+        config = removeProviderOrThrow(config, providerName);
+      }
+
+      await writeWebSearchConfig(path, config);
+      ctx.ui.notify(`Saved web-search config to ${path}`, "info");
+    },
+  });
+}