websearch.ts 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. import z from "zod"
  2. import { Tool } from "./tool"
  3. import DESCRIPTION from "./websearch.txt"
  4. import { Config } from "../config/config"
  5. import { Permission } from "../permission"
  6. const API_CONFIG = {
  7. BASE_URL: "https://mcp.exa.ai",
  8. ENDPOINTS: {
  9. SEARCH: "/mcp",
  10. },
  11. DEFAULT_NUM_RESULTS: 8,
  12. } as const
  13. interface McpSearchRequest {
  14. jsonrpc: string
  15. id: number
  16. method: string
  17. params: {
  18. name: string
  19. arguments: {
  20. query: string
  21. numResults?: number
  22. livecrawl?: "fallback" | "preferred"
  23. type?: "auto" | "fast" | "deep"
  24. contextMaxCharacters?: number
  25. }
  26. }
  27. }
  28. interface McpSearchResponse {
  29. jsonrpc: string
  30. result: {
  31. content: Array<{
  32. type: string
  33. text: string
  34. }>
  35. }
  36. }
  37. export const WebSearchTool = Tool.define("websearch", {
  38. description: DESCRIPTION,
  39. parameters: z.object({
  40. query: z.string().describe("Websearch query"),
  41. numResults: z.number().optional().describe("Number of search results to return (default: 8)"),
  42. livecrawl: z
  43. .enum(["fallback", "preferred"])
  44. .optional()
  45. .describe(
  46. "Live crawl mode - 'fallback': use live crawling as backup if cached content unavailable, 'preferred': prioritize live crawling (default: 'fallback')",
  47. ),
  48. type: z
  49. .enum(["auto", "fast", "deep"])
  50. .optional()
  51. .describe("Search type - 'auto': balanced search (default), 'fast': quick results, 'deep': comprehensive search"),
  52. contextMaxCharacters: z
  53. .number()
  54. .optional()
  55. .describe("Maximum characters for context string optimized for LLMs (default: 10000)"),
  56. }),
  57. async execute(params, ctx) {
  58. const cfg = await Config.get()
  59. if (cfg.permission?.webfetch === "ask")
  60. await Permission.ask({
  61. type: "websearch",
  62. sessionID: ctx.sessionID,
  63. messageID: ctx.messageID,
  64. callID: ctx.callID,
  65. title: "Search web for: " + params.query,
  66. metadata: {
  67. query: params.query,
  68. numResults: params.numResults,
  69. livecrawl: params.livecrawl,
  70. type: params.type,
  71. contextMaxCharacters: params.contextMaxCharacters,
  72. },
  73. })
  74. const searchRequest: McpSearchRequest = {
  75. jsonrpc: "2.0",
  76. id: 1,
  77. method: "tools/call",
  78. params: {
  79. name: "web_search_exa",
  80. arguments: {
  81. query: params.query,
  82. type: params.type || "auto",
  83. numResults: params.numResults || API_CONFIG.DEFAULT_NUM_RESULTS,
  84. livecrawl: params.livecrawl || "fallback",
  85. contextMaxCharacters: params.contextMaxCharacters,
  86. },
  87. },
  88. }
  89. const controller = new AbortController()
  90. const timeoutId = setTimeout(() => controller.abort(), 25000)
  91. try {
  92. const headers: Record<string, string> = {
  93. accept: "application/json, text/event-stream",
  94. "content-type": "application/json",
  95. }
  96. const response = await fetch(`${API_CONFIG.BASE_URL}${API_CONFIG.ENDPOINTS.SEARCH}`, {
  97. method: "POST",
  98. headers,
  99. body: JSON.stringify(searchRequest),
  100. signal: AbortSignal.any([controller.signal, ctx.abort]),
  101. })
  102. clearTimeout(timeoutId)
  103. if (!response.ok) {
  104. const errorText = await response.text()
  105. throw new Error(`Search error (${response.status}): ${errorText}`)
  106. }
  107. const responseText = await response.text()
  108. // Parse SSE response
  109. const lines = responseText.split("\n")
  110. for (const line of lines) {
  111. if (line.startsWith("data: ")) {
  112. const data: McpSearchResponse = JSON.parse(line.substring(6))
  113. if (data.result && data.result.content && data.result.content.length > 0) {
  114. return {
  115. output: data.result.content[0].text,
  116. title: `Web search: ${params.query}`,
  117. metadata: {},
  118. }
  119. }
  120. }
  121. }
  122. return {
  123. output: "No search results found. Please try a different query.",
  124. title: `Web search: ${params.query}`,
  125. metadata: {},
  126. }
  127. } catch (error) {
  128. clearTimeout(timeoutId)
  129. if (error instanceof Error && error.name === "AbortError") {
  130. throw new Error("Search request timed out")
  131. }
  132. throw error
  133. }
  134. },
  135. })