pim-agent 0.0.1 → 0.9.0

package/packages/core/src/extensions/web-search/ExaMcpClient.ts

@@ -0,0 +1,251 @@
+import { Json } from "../../shared/Json";
+import { McpClient, type McpFetch } from "../../shared/McpClient";
+import { RateLimiter } from "../../shared/RateLimiter";
+import {
+  normalizeSnippet,
+  type SearchResult,
+} from "./providers/SearchProvider";
+
+type ExaMcpClientOptions = {
+  readonly endpoint?: string;
+  readonly apiKey?: string;
+  readonly fetch?: McpFetch;
+  readonly rateLimiter?: RateLimiter;
+};
+
+type ExaSearchInput = {
+  readonly query: string;
+  readonly numResults: number;
+  readonly signal?: AbortSignal;
+};
+
+class ExaSearchError extends Error {
+  public constructor(message: string) {
+    super(message);
+    this.name = "ExaSearchError";
+  }
+}
+
+const defaultEndpoint = "https://mcp.exa.ai/mcp";
+const toolName = "web_search_exa";
+// Stay under the keyless endpoint's 2 req/s window, or a 429 no longer means the daily cap.
+const maxRequestsPerWindow = 2;
+const windowMs = 1100;
+
+export class ExaMcpClient {
+  private readonly client: McpClient;
+
+  public constructor(options: ExaMcpClientOptions = {}) {
+    const apiKey =
+      options.apiKey === undefined || options.apiKey.length === 0
+        ? undefined
+        : options.apiKey;
+    // Throttle only on the free tier; an API key lifts the rate limit.
+    const rateLimiter =
+      apiKey !== undefined
+        ? undefined
+        : (options.rateLimiter ??
+          new RateLimiter({
+            maxRequests: maxRequestsPerWindow,
+            windowMs: windowMs,
+          }));
+
+    this.client = new McpClient({
+      endpoint: options.endpoint ?? defaultEndpoint,
+      ...(apiKey === undefined ? {} : { headers: { "x-api-key": apiKey } }),
+      ...(rateLimiter === undefined ? {} : { rateLimiter }),
+      ...(options.fetch === undefined ? {} : { fetch: options.fetch }),
+    });
+  }
+
+  public async search(input: ExaSearchInput): Promise<readonly SearchResult[]> {
+    const result = await this.client.callTool({
+      name: toolName,
+      arguments: {
+        query: input.query,
+        numResults: input.numResults,
+      },
+      ...(input.signal === undefined ? {} : { signal: input.signal }),
+    });
+
+    return extractResults(result);
+  }
+}
+
+function extractResults(result: unknown): readonly SearchResult[] {
+  const record = Json.asRecord(result);
+  const content = record?.["content"];
+
+  if (!Array.isArray(content) || content.length === 0) {
+    throw new ExaSearchError("Exa returned malformed tool content.");
+  }
+
+  const textBlocks = content.map(readTextBlock);
+  const plainTextResults = extractPlainTextResults(textBlocks);
+
+  if (plainTextResults !== undefined) {
+    return plainTextResults;
+  }
+
+  const resultObjects = findFirstObjectArray(
+    textBlocks
+      .map((block) => Json.tryParseJson(block))
+      .filter((value) => value !== undefined)
+  );
+
+  if (resultObjects === undefined) {
+    throw new ExaSearchError("Exa returned malformed search results.");
+  }
+
+  return resultObjects.map(projectResult);
+}
+
+function readTextBlock(block: unknown): string {
+  const record = Json.asRecord(block);
+
+  if (record?.["type"] !== "text" || typeof record["text"] !== "string") {
+    throw new ExaSearchError("Exa returned malformed tool content.");
+  }
+
+  return record["text"];
+}
+
+function extractPlainTextResults(
+  textBlocks: readonly string[]
+): readonly SearchResult[] | undefined {
+  for (const textBlock of textBlocks) {
+    const blocks = textBlock
+      .split(/\n---\n/u)
+      .map((block) => block.trim())
+      .filter((block) => block.length > 0);
+    const results = blocks
+      .map((block) => parsePlainTextResult(block))
+      .filter((result) => result !== undefined);
+
+    if (results.length > 0) {
+      return results;
+    }
+  }
+
+  return undefined;
+}
+
+function parsePlainTextResult(block: string): SearchResult | undefined {
+  const lines = block
+    .split(/\r?\n/u)
+    .map((line) => line.trim())
+    .filter((line) => line.length > 0);
+  const title = readLabeledLine(lines, "Title");
+  const url = readLabeledLine(lines, "URL");
+
+  if (title === undefined || url === undefined) {
+    return undefined;
+  }
+
+  return {
+    title,
+    url,
+    snippet: normalizeSnippet(readPlainTextSnippet(lines)),
+  };
+}
+
+function readLabeledLine(
+  lines: readonly string[],
+  label: string
+): string | undefined {
+  const prefix = `${label}:`;
+  const line = lines.find((candidate) =>
+    candidate.toLowerCase().startsWith(prefix.toLowerCase())
+  );
+
+  return line?.slice(prefix.length).trim();
+}
+
+function readPlainTextSnippet(lines: readonly string[]): string {
+  const highlightsIndex = lines.findIndex((line) =>
+    line.toLowerCase().startsWith("highlights:")
+  );
+  const snippetLines =
+    highlightsIndex === -1 ? lines : lines.slice(highlightsIndex + 1);
+  const skipPrefixes = ["title:", "url:", "published:", "author:"];
+  return snippetLines
+    .filter((line) => {
+      if (line.startsWith("[...]")) {
+        return false;
+      }
+      const lower = line.toLowerCase();
+      return !skipPrefixes.some((prefix) => lower.startsWith(prefix));
+    })
+    .join(" ");
+}
+
+function findFirstObjectArray(
+  values: readonly unknown[]
+): readonly Readonly<Record<string, unknown>>[] | undefined {
+  for (const value of values) {
+    if (
+      Array.isArray(value) &&
+      value.every((item) => Json.asRecord(item) !== undefined)
+    ) {
+      return value as readonly Readonly<Record<string, unknown>>[];
+    }
+
+    const record = Json.asRecord(value);
+
+    if (record === undefined) {
+      continue;
+    }
+
+    for (const nestedValue of Object.values(record)) {
+      if (
+        Array.isArray(nestedValue) &&
+        nestedValue.every((item) => Json.asRecord(item) !== undefined)
+      ) {
+        return nestedValue as readonly Readonly<Record<string, unknown>>[];
+      }
+    }
+  }
+
+  return undefined;
+}
+
+function projectResult(
+  result: Readonly<Record<string, unknown>>
+): SearchResult {
+  return {
+    title: readResultString(result, "title"),
+    url: readResultString(result, "url"),
+    snippet: normalizeSnippet(readSnippet(result)),
+  };
+}
+
+function readSnippet(result: Readonly<Record<string, unknown>>): string {
+  return (
+    readOptionalResultString(result, "snippet") ??
+    readOptionalResultString(result, "text") ??
+    readOptionalResultString(result, "summary") ??
+    ""
+  );
+}
+
+function readResultString(
+  result: Readonly<Record<string, unknown>>,
+  name: string
+): string {
+  const value = readOptionalResultString(result, name);
+
+  if (value === undefined) {
+    throw new ExaSearchError(`Exa returned a result without ${name}.`);
+  }
+
+  return value;
+}
+
+function readOptionalResultString(
+  result: Readonly<Record<string, unknown>>,
+  name: string
+): string | undefined {
+  const value = result[name];
+
+  return typeof value === "string" ? value : undefined;
+}

package/packages/core/src/extensions/web-search/SearchBreaker.ts

@@ -0,0 +1,147 @@
+import { join } from "node:path";
+
+import { Fs } from "../../shared/Fs";
+import { Paths } from "../../shared/Paths";
+
+type BreakerEntry = {
+  readonly until: number;
+  readonly nextProbeAt: number;
+  readonly reason: string;
+};
+
+type BreakerState = Record<string, BreakerEntry>;
+
+export type SearchBreakerOptions = {
+  readonly path?: string;
+  readonly now?: () => number;
+  readonly probeIntervalMs?: number;
+};
+
+export type TripInput = {
+  readonly provider: string;
+  readonly reason: string;
+  readonly retryAfterMs?: number;
+};
+
+const DAY_MS = 86_400_000;
+
+const defaultProbeIntervalMs = 1_800_000;
+
+/** Which providers are quota-exhausted, on disk; quota is per-IP-per-day, so trips expire at UTC midnight. */
+export class SearchBreaker {
+  private readonly filePath: string;
+  private readonly now: () => number;
+  private readonly probeIntervalMs: number;
+  private readonly writes = Fs.serialised();
+
+  public constructor(options: SearchBreakerOptions = {}) {
+    this.filePath =
+      options.path ?? join(Paths.pimHomeDir(), "web-search-breaker.json");
+    this.now = options.now ?? Date.now;
+    this.probeIntervalMs = options.probeIntervalMs ?? defaultProbeIntervalMs;
+  }
+
+  public async isOpen(provider: string): Promise<boolean> {
+    const entry = (await this.read())[provider];
+
+    if (entry === undefined) {
+      return false;
+    }
+
+    const now = this.now();
+
+    return now < entry.until && now < entry.nextProbeAt;
+  }
+
+  public async trip(input: TripInput): Promise<void> {
+    const now = this.now();
+    const ttlMs = input.retryAfterMs ?? msUntilNextUtcMidnight(now);
+    const until = now + Math.max(0, ttlMs);
+
+    await this.mutate((state) => ({
+      ...state,
+      [input.provider]: {
+        until,
+        nextProbeAt: Math.min(until, now + this.probeIntervalMs),
+        reason: input.reason,
+      },
+    }));
+  }
+
+  public async reset(provider: string): Promise<void> {
+    await this.mutate((state) => {
+      if (state[provider] === undefined) {
+        return state;
+      }
+
+      const { [provider]: _removed, ...rest } = state;
+      return rest;
+    });
+  }
+
+  private async read(): Promise<BreakerState> {
+    let raw: unknown;
+
+    try {
+      raw = await Bun.file(this.filePath).json();
+    } catch {
+      return {};
+    }
+
+    if (typeof raw !== "object" || raw === null || Array.isArray(raw)) {
+      return {};
+    }
+
+    const now = this.now();
+    const state: BreakerState = {};
+
+    for (const [provider, value] of Object.entries(raw)) {
+      const entry = parseEntry(value);
+
+      if (entry !== undefined && entry.until > now) {
+        state[provider] = entry;
+      }
+    }
+
+    return state;
+  }
+
+  private async mutate(
+    update: (state: BreakerState) => BreakerState
+  ): Promise<void> {
+    await this.writes.run(async () => {
+      const state = await this.read();
+      const next = update(state);
+
+      if (next === state) {
+        return;
+      }
+
+      await Paths.ensurePimHome();
+      await Fs.writeJson(this.filePath, next);
+    });
+  }
+}
+
+function parseEntry(value: unknown): BreakerEntry | undefined {
+  if (typeof value !== "object" || value === null || Array.isArray(value)) {
+    return undefined;
+  }
+
+  const record = value as Record<string, unknown>;
+  const { until, nextProbeAt, reason } = record;
+
+  if (typeof until !== "number" || typeof nextProbeAt !== "number") {
+    return undefined;
+  }
+
+  return {
+    until,
+    nextProbeAt,
+    reason: typeof reason === "string" ? reason : "",
+  };
+}
+
+function msUntilNextUtcMidnight(now: number): number {
+  return DAY_MS - (now % DAY_MS);
+}

package/packages/core/src/extensions/web-search/SearchChain.ts

@@ -0,0 +1,77 @@
+import { Errors } from "../../shared/Errors";
+import type { SearchBreaker } from "./SearchBreaker";
+import {
+  ProviderQuotaError,
+  type ProviderSearchInput,
+  type SearchProvider,
+  type SearchResult,
+} from "./providers/SearchProvider";
+
+export type ChainSearchOutcome = {
+  readonly provider: string;
+  readonly results: readonly SearchResult[];
+  readonly fellBack: boolean;
+};
+
+export type SearchChainOptions = {
+  readonly providers: readonly SearchProvider[];
+  readonly breaker: SearchBreaker;
+};
+
+/** Tries providers in order until one answers; an empty result set is an answer and ends the chain. */
+export class SearchChain {
+  private readonly providers: readonly SearchProvider[];
+  private readonly breaker: SearchBreaker;
+
+  public constructor(options: SearchChainOptions) {
+    this.providers = options.providers;
+    this.breaker = options.breaker;
+  }
+
+  public async search(input: ProviderSearchInput): Promise<ChainSearchOutcome> {
+    const failures: string[] = [];
+    let attempted = false;
+
+    for (const provider of this.providers) {
+      if (await this.breaker.isOpen(provider.name)) {
+        failures.push(`${provider.name}: skipped (quota exhausted)`);
+        continue;
+      }
+
+      if (input.signal?.aborted) {
+        throw new Error("Web search aborted before execution.");
+      }
+
+      try {
+        const results = await provider.search(input);
+        await this.breaker.reset(provider.name);
+
+        return { provider: provider.name, results, fellBack: attempted };
+      } catch (error) {
+        if (input.signal?.aborted || Errors.isAbort(error)) {
+          throw error;
+        }
+
+        attempted = true;
+
+        if (error instanceof ProviderQuotaError) {
+          await this.breaker.trip({
+            provider: provider.name,
+            reason: error.message,
+            ...(error.retryAfterMs === undefined
+              ? {}
+              : { retryAfterMs: error.retryAfterMs }),
+          });
+        }
+
+        failures.push(`${provider.name}: ${Errors.describe(error)}`);
+      }
+    }
+
+    throw new Error(
+      `All web search providers failed:\n${failures
+        .map((failure) => `  - ${failure}`)
+        .join("\n")}`
+    );
+  }
+}

package/packages/core/src/extensions/web-search/index.ts

@@ -0,0 +1,90 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { PimSettings } from "../../shared/PimSettings";
+import { Tools } from "../../shared/Tools";
+import { DuckDuckGoProvider } from "./providers/DuckDuckGoProvider";
+import { ExaProvider } from "./providers/ExaProvider";
+import { FirecrawlProvider } from "./providers/FirecrawlProvider";
+import { webSearchView } from "./render";
+import { SearchBreaker } from "./SearchBreaker";
+import { SearchChain } from "./SearchChain";
+import { type WebSearchInput, webSearchSchema } from "./schema";
+import { clampNumResults, formatResults } from "./search";
+
+async function createChain(): Promise<SearchChain> {
+  const [exaApiKey, firecrawlApiKey, jinaApiKey] = await Promise.all([
+    PimSettings.getExaApiKey(),
+    PimSettings.getFirecrawlApiKey(),
+    PimSettings.getJinaApiKey(),
+  ]);
+
+  return new SearchChain({
+    breaker: new SearchBreaker(),
+    providers: [
+      new ExaProvider(exaApiKey === undefined ? {} : { apiKey: exaApiKey }),
+      new FirecrawlProvider(
+        firecrawlApiKey === undefined ? {} : { apiKey: firecrawlApiKey }
+      ),
+      new DuckDuckGoProvider(
+        jinaApiKey === undefined ? {} : { apiKey: jinaApiKey }
+      ),
+    ],
+  });
+}
+
+export default function (pi: ExtensionAPI): void {
+  let chainPromise: Promise<SearchChain> | undefined;
+  const getChain = () => (chainPromise ??= createChain());
+
+  Tools.register(pi, {
+    name: "web_search",
+    label: "web_search",
+    description:
+      "Search the web. " +
+      "Returns ranked results with title, URL, and a short snippet.",
+    parameters: webSearchSchema,
+    renderShell: "self",
+    effect: { kind: "readOnly" },
+    executionMode: "parallel",
+    async execute(_id, params, signal) {
+      const { query, numResults } = params as WebSearchInput;
+
+      if (signal?.aborted) {
+        throw new Error("Web search aborted before execution.");
+      }
+
+      const trimmed = query.trim();
+      if (trimmed.length === 0) {
+        throw new Error(
+          "Web search query is empty. Provide a non-empty query."
+        );
+      }
+
+      const clamped = clampNumResults(numResults);
+      const chain = await getChain();
+      const outcome = await chain.search({
+        query: trimmed,
+        numResults: clamped,
+        ...(signal === undefined ? {} : { signal }),
+      });
+
+      if (outcome.results.length === 0) {
+        throw new Error(
+          `No web results for "${trimmed}" (via ${outcome.provider}). ` +
+            "Try broader keywords or different phrasing."
+        );
+      }
+
+      return {
+        content: [{ type: "text", text: formatResults(outcome.results) }],
+        details: {
+          query: trimmed,
+          numResults: clamped,
+          count: outcome.results.length,
+          provider: outcome.provider,
+          fellBack: outcome.fellBack,
+        },
+      };
+    },
+    toViewModel: webSearchView,
+  });
+}