pi-web-search/src/runtime.ts

import { loadWebSearchConfig, type ResolvedWebSearchConfig } from "./config.ts";
import { createExaProvider } from "./providers/exa.ts";
import { createTavilyProvider } from "./providers/tavily.ts";
import type {
  NormalizedFetchRequest,
  NormalizedFetchResponse,
  NormalizedSearchRequest,
  NormalizedSearchResponse,
  WebProvider,
} from "./providers/types.ts";
import type { WebSearchProviderConfig } from "./schema.ts";

export interface ProviderExecutionMeta {
  requestedProviderName?: string;
  actualProviderName: string;
  failoverFromProviderName?: string;
  failoverReason?: string;
}

export interface RuntimeSearchResponse extends NormalizedSearchResponse {
  execution: ProviderExecutionMeta;
}

export interface RuntimeFetchResponse extends NormalizedFetchResponse {
  execution: ProviderExecutionMeta;
}

export function createWebSearchRuntime(
  deps: {
    loadConfig?: () => Promise<ResolvedWebSearchConfig>;
    createProvider?: (providerConfig: WebSearchProviderConfig) => WebProvider;
  } = {},
) {
  const loadConfig = deps.loadConfig ?? loadWebSearchConfig;
  const createProvider = deps.createProvider ?? ((providerConfig: WebSearchProviderConfig) => {
    switch (providerConfig.type) {
      case "tavily":
        return createTavilyProvider(providerConfig);
      case "exa":
        return createExaProvider(providerConfig);
    }
  });

  async function resolveConfigAndProvider(providerName?: string) {
    const config = await loadConfig();
    const selectedName = providerName ?? config.defaultProviderName;
    const selectedConfig = config.providersByName.get(selectedName);

    if (!selectedConfig) {
      throw new Error(
        `Unknown web-search provider \"${selectedName}\". Configured providers: ${[...config.providersByName.keys()].join(", ")}`,
      );
    }

    return {
      config,
      selectedName,
      selectedConfig,
      selectedProvider: createProvider(selectedConfig),
    };
  }

  async function search(request: NormalizedSearchRequest): Promise<RuntimeSearchResponse> {
    const { config, selectedName, selectedConfig, selectedProvider } = await resolveConfigAndProvider(request.provider);

    try {
      const response = await selectedProvider.search(request);
      return {
        ...response,
        execution: {
          requestedProviderName: request.provider,
          actualProviderName: selectedName,
        },
      };
    } catch (error) {
      if (selectedConfig.type !== "tavily") {
        throw error;
      }

      const fallbackConfig = [...config.providersByName.values()].find((provider) => provider.type === "exa");
      if (!fallbackConfig) {
        throw error;
      }

      const fallbackProvider = createProvider(fallbackConfig);
      const fallbackResponse = await fallbackProvider.search({ ...request, provider: fallbackConfig.name });
      return {
        ...fallbackResponse,
        execution: {
          requestedProviderName: request.provider,
          actualProviderName: fallbackConfig.name,
          failoverFromProviderName: selectedName,
          failoverReason: (error as Error).message,
        },
      };
    }
  }

  async function fetch(request: NormalizedFetchRequest): Promise<RuntimeFetchResponse> {
    const { config, selectedName, selectedConfig, selectedProvider } = await resolveConfigAndProvider(request.provider);

    try {
      const response = await selectedProvider.fetch(request);
      return {
        ...response,
        execution: {
          requestedProviderName: request.provider,
          actualProviderName: selectedName,
        },
      };
    } catch (error) {
      if (selectedConfig.type !== "tavily") {
        throw error;
      }

      const fallbackConfig = [...config.providersByName.values()].find((provider) => provider.type === "exa");
      if (!fallbackConfig) {
        throw error;
      }

      const fallbackProvider = createProvider(fallbackConfig);
      const fallbackResponse = await fallbackProvider.fetch({ ...request, provider: fallbackConfig.name });
      return {
        ...fallbackResponse,
        execution: {
          requestedProviderName: request.provider,
          actualProviderName: fallbackConfig.name,
          failoverFromProviderName: selectedName,
          failoverReason: (error as Error).message,
        },
      };
    }
  }

  return {
    search,
    fetch,
  };
}