codesearch.ts 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. import z from "zod"
  2. import { Tool } from "./tool"
  3. import DESCRIPTION from "./codesearch.txt"
  4. const API_CONFIG = {
  5. BASE_URL: "https://mcp.exa.ai",
  6. ENDPOINTS: {
  7. CONTEXT: "/mcp",
  8. },
  9. } as const
  10. interface McpCodeRequest {
  11. jsonrpc: string
  12. id: number
  13. method: string
  14. params: {
  15. name: string
  16. arguments: {
  17. query: string
  18. tokensNum: number
  19. }
  20. }
  21. }
  22. interface McpCodeResponse {
  23. jsonrpc: string
  24. result: {
  25. content: Array<{
  26. type: string
  27. text: string
  28. }>
  29. }
  30. }
  31. export const CodeSearchTool = Tool.define("codesearch", {
  32. description: DESCRIPTION,
  33. parameters: z.object({
  34. query: z
  35. .string()
  36. .describe(
  37. "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'",
  38. ),
  39. tokensNum: z
  40. .number()
  41. .min(1000)
  42. .max(50000)
  43. .default(5000)
  44. .describe(
  45. "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.",
  46. ),
  47. }),
  48. async execute(params, ctx) {
  49. await ctx.ask({
  50. permission: "codesearch",
  51. patterns: [params.query],
  52. always: ["*"],
  53. metadata: {
  54. query: params.query,
  55. tokensNum: params.tokensNum,
  56. },
  57. })
  58. const codeRequest: McpCodeRequest = {
  59. jsonrpc: "2.0",
  60. id: 1,
  61. method: "tools/call",
  62. params: {
  63. name: "get_code_context_exa",
  64. arguments: {
  65. query: params.query,
  66. tokensNum: params.tokensNum || 5000,
  67. },
  68. },
  69. }
  70. const controller = new AbortController()
  71. const timeoutId = setTimeout(() => controller.abort(), 30000)
  72. try {
  73. const headers: Record<string, string> = {
  74. accept: "application/json, text/event-stream",
  75. "content-type": "application/json",
  76. }
  77. const response = await fetch(`${API_CONFIG.BASE_URL}${API_CONFIG.ENDPOINTS.CONTEXT}`, {
  78. method: "POST",
  79. headers,
  80. body: JSON.stringify(codeRequest),
  81. signal: AbortSignal.any([controller.signal, ctx.abort]),
  82. })
  83. clearTimeout(timeoutId)
  84. if (!response.ok) {
  85. const errorText = await response.text()
  86. throw new Error(`Code search error (${response.status}): ${errorText}`)
  87. }
  88. const responseText = await response.text()
  89. // Parse SSE response
  90. const lines = responseText.split("\n")
  91. for (const line of lines) {
  92. if (line.startsWith("data: ")) {
  93. const data: McpCodeResponse = JSON.parse(line.substring(6))
  94. if (data.result && data.result.content && data.result.content.length > 0) {
  95. return {
  96. output: data.result.content[0].text,
  97. title: `Code search: ${params.query}`,
  98. metadata: {},
  99. }
  100. }
  101. }
  102. }
  103. return {
  104. output:
  105. "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.",
  106. title: `Code search: ${params.query}`,
  107. metadata: {},
  108. }
  109. } catch (error) {
  110. clearTimeout(timeoutId)
  111. if (error instanceof Error && error.name === "AbortError") {
  112. throw new Error("Code search request timed out")
  113. }
  114. throw error
  115. }
  116. },
  117. })