page-assist/src/web/web.ts

import { getWebSearchPrompt } from "~/services/ollama"
import { webGoogleSearch } from "./search-engines/google"
import { webDuckDuckGoSearch } from "./search-engines/duckduckgo"
import { getIsVisitSpecificWebsite, getSearchProvider } from "@/services/search"
import { webSogouSearch } from "./search-engines/sogou"
import { webBraveSearch } from "./search-engines/brave"
import { getWebsiteFromQuery, processSingleWebsite } from "./website"
import { searxngSearch } from "./search-engines/searxng"
import { braveAPISearch } from "./search-engines/brave-api"
import { webBaiduSearch } from "./search-engines/baidu"
import { searchIod } from "./iod"
import type { WebSearchResult } from "~/types/web"
import type { IodRegistryEntry } from "~/types/iod"

const getHostName = (url: string) => {
  try {
    const hostname = new URL(url).hostname
    return hostname
  } catch (e) {
    return ""
  }
}

async function searchWeb(
  provider: string,
  query: string
): Promise<WebSearchResult[]> {
  let results = []
  switch (provider) {
    case "duckduckgo":
      results = await webDuckDuckGoSearch(query)
      break
    case "sogou":
      results = await webSogouSearch(query)
      break
    case "brave":
      results = await webBraveSearch(query)
      break
    case "searxng":
      results = await searxngSearch(query)
      break
    case "brave-api":
      results = await braveAPISearch(query)
      break
    case "baidu":
      results = await webBaiduSearch(query)
      break
    default:
      results = await webGoogleSearch(query)
      break
  }
  return results.map((r) => ({ ...r, name: getHostName(r.url) }))
}

export const getSystemPromptForWeb = async (
  query: string,
  keywords: string[] = [],
  webSearch = true,
  iodSearch = false
) => {
  try {
    const websiteVisit = getWebsiteFromQuery(query)
    let webSearchResults: WebSearchResult[] = []
    // let search_results_web = ""

    if (webSearch) {
      const isVisitSpecificWebsite = await getIsVisitSpecificWebsite()

      if (isVisitSpecificWebsite && websiteVisit.hasUrl) {
        const url = websiteVisit.url
        const queryWithoutUrl = websiteVisit.queryWithouUrls
        webSearchResults = await processSingleWebsite(url, queryWithoutUrl)
      } else {
        const searchProvider = await getSearchProvider()
        webSearchResults = await searchWeb(searchProvider, query)
      }

      // search_results_web = webSearchResults
      //   .map(
      //     (result, idx) =>
      //       `<result source="${result.url}" id="${idx}">${result.content}</result>`
      //   )
      //   .join("\n")
    }

    let iodSearchResults: IodRegistryEntry[] = []
    // let search_results_iod = ""

    if (iodSearch) {
      iodSearchResults = await searchIod(query, keywords)
      // search_results_iod = iodSearchResults
      //   .map(
      //     (result, idx) =>
      //       `<result source="${result.url}" id="${idx}">${result.content}</result>`
      //   )
      //   .join("\n")
    }
    const _iodSearchResults = iodSearchResults
        .map((res) => ({
          doId: res.doId,
          name: res.name,
          url: res.url,
          data_space: res.data_space,
          content: res.content || res.description,
          tokenCount: (res.content || res.description)?.length ?? 0,
        }))

    const iod_search_results = _iodSearchResults
      .map(
        (result, idx) =>
          `<result doId="${result.doId}" name="${result.name}" source="${result.url}" id="${idx + 1}">${result.content}</result>`
      )
      .join("\n")
    console.log("iod_search_result: " + iod_search_results)

    const web_search_results = webSearchResults
      .map(
        (result, idx) =>
          `<result source="${result.url}" name="${result.name}" id="${idx + 1}">${result.content}</result>`
      )
      .join("\n")
    console.log("web_search_result: " + web_search_results)

    const current_date_time = new Date().toLocaleString()

    const system = await getWebSearchPrompt()

    const prompt = system
      .replace("{current_date_time}", current_date_time)
      .replace("{iod_search_results}", iod_search_results)
      .replace("{web_search_results}", web_search_results)

    return {
      prompt,
      webSources: webSearchResults.map((result) => {
        return {
          url: result.url,
          name: result.name,
          type: "url"
        }
      }),
      iodSources: iodSearchResults,
      iodSearchResults: _iodSearchResults,
      iodTokenCount: _iodSearchResults.reduce((acc, cur) => (acc + cur.content.length), 0)
    }
  } catch (e) {
    console.error(e)
    return {
      prompt: "",
      webSources: [],
      iodSources: [],
      iodSearchResults: [],
      iodTokenCount: 0,
    }
  }
}