codesearch.ts 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138
  1. import z from "zod"
  2. import { Tool } from "./tool"
  3. import DESCRIPTION from "./codesearch.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. CONTEXT: "/mcp",
  10. },
  11. } as const
  12. interface McpCodeRequest {
  13. jsonrpc: string
  14. id: number
  15. method: string
  16. params: {
  17. name: string
  18. arguments: {
  19. query: string
  20. tokensNum: number
  21. }
  22. }
  23. }
  24. interface McpCodeResponse {
  25. jsonrpc: string
  26. result: {
  27. content: Array<{
  28. type: string
  29. text: string
  30. }>
  31. }
  32. }
  33. export const CodeSearchTool = Tool.define("codesearch", {
  34. description: DESCRIPTION,
  35. parameters: z.object({
  36. query: z
  37. .string()
  38. .describe(
  39. "Search query to find relevant context for APIs, Libraries, and SDKs. For example, 'React useState hook examples', 'Python pandas dataframe filtering', 'Express.js middleware', 'Next js partial prerendering configuration'",
  40. ),
  41. tokensNum: z
  42. .number()
  43. .min(1000)
  44. .max(50000)
  45. .default(5000)
  46. .describe(
  47. "Number of tokens to return (1000-50000). Default is 5000 tokens. Adjust this value based on how much context you need - use lower values for focused queries and higher values for comprehensive documentation.",
  48. ),
  49. }),
  50. async execute(params, ctx) {
  51. const cfg = await Config.get()
  52. if (cfg.permission?.webfetch === "ask")
  53. await Permission.ask({
  54. type: "codesearch",
  55. sessionID: ctx.sessionID,
  56. messageID: ctx.messageID,
  57. callID: ctx.callID,
  58. title: "Search code for: " + params.query,
  59. metadata: {
  60. query: params.query,
  61. tokensNum: params.tokensNum,
  62. },
  63. })
  64. const codeRequest: McpCodeRequest = {
  65. jsonrpc: "2.0",
  66. id: 1,
  67. method: "tools/call",
  68. params: {
  69. name: "get_code_context_exa",
  70. arguments: {
  71. query: params.query,
  72. tokensNum: params.tokensNum || 5000,
  73. },
  74. },
  75. }
  76. const controller = new AbortController()
  77. const timeoutId = setTimeout(() => controller.abort(), 30000)
  78. try {
  79. const headers: Record<string, string> = {
  80. accept: "application/json, text/event-stream",
  81. "content-type": "application/json",
  82. }
  83. const response = await fetch(`${API_CONFIG.BASE_URL}${API_CONFIG.ENDPOINTS.CONTEXT}`, {
  84. method: "POST",
  85. headers,
  86. body: JSON.stringify(codeRequest),
  87. signal: AbortSignal.any([controller.signal, ctx.abort]),
  88. })
  89. clearTimeout(timeoutId)
  90. if (!response.ok) {
  91. const errorText = await response.text()
  92. throw new Error(`Code search error (${response.status}): ${errorText}`)
  93. }
  94. const responseText = await response.text()
  95. // Parse SSE response
  96. const lines = responseText.split("\n")
  97. for (const line of lines) {
  98. if (line.startsWith("data: ")) {
  99. const data: McpCodeResponse = JSON.parse(line.substring(6))
  100. if (data.result && data.result.content && data.result.content.length > 0) {
  101. return {
  102. output: data.result.content[0].text,
  103. title: `Code search: ${params.query}`,
  104. metadata: {},
  105. }
  106. }
  107. }
  108. }
  109. return {
  110. output:
  111. "No code snippets or documentation found. Please try a different query, be more specific about the library or programming concept, or check the spelling of framework names.",
  112. title: `Code search: ${params.query}`,
  113. metadata: {},
  114. }
  115. } catch (error) {
  116. clearTimeout(timeoutId)
  117. if (error instanceof Error && error.name === "AbortError") {
  118. throw new Error("Code search request timed out")
  119. }
  120. throw error
  121. }
  122. },
  123. })