refactor(tool): convert codesearch tool internals to Effect (#21811)
This commit is contained in:
parent
00e39d2114
commit
bf601628db
2 changed files with 58 additions and 124 deletions
|
|
@ -1,132 +1,65 @@
|
||||||
import z from "zod"
|
import z from "zod"
|
||||||
|
import { Effect } from "effect"
|
||||||
|
import { HttpClient } from "effect/unstable/http"
|
||||||
import { Tool } from "./tool"
|
import { Tool } from "./tool"
|
||||||
|
import * as McpExa from "./mcp-exa"
|
||||||
import DESCRIPTION from "./codesearch.txt"
|
import DESCRIPTION from "./codesearch.txt"
|
||||||
import { abortAfterAny } from "../util/abort"
|
|
||||||
|
|
||||||
const API_CONFIG = {
|
export const CodeSearchTool = Tool.defineEffect(
|
||||||
BASE_URL: "https://mcp.exa.ai",
|
"codesearch",
|
||||||
ENDPOINTS: {
|
Effect.gen(function* () {
|
||||||
CONTEXT: "/mcp",
|
const http = yield* HttpClient.HttpClient
|
||||||
},
|
|
||||||
} as const
|
|
||||||
|
|
||||||
interface McpCodeRequest {
|
return {
|
||||||
jsonrpc: string
|
description: DESCRIPTION,
|
||||||
id: number
|
parameters: z.object({
|
||||||
method: string
|
query: z
|
||||||
params: {
|
.string()
|
||||||
name: string
|
.describe(
|
||||||
arguments: {
|
"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'",
|
||||||
query: string
|
),
|
||||||
tokensNum: number
|
tokensNum: z
|
||||||
}
|
.number()
|
||||||
}
|
.min(1000)
|
||||||
}
|
.max(50000)
|
||||||
|
.default(5000)
|
||||||
|
.describe(
|
||||||
|
"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.",
|
||||||
|
),
|
||||||
|
}),
|
||||||
|
execute: (params: { query: string; tokensNum: number }, ctx: Tool.Context) =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
yield* Effect.promise(() =>
|
||||||
|
ctx.ask({
|
||||||
|
permission: "codesearch",
|
||||||
|
patterns: [params.query],
|
||||||
|
always: ["*"],
|
||||||
|
metadata: {
|
||||||
|
query: params.query,
|
||||||
|
tokensNum: params.tokensNum,
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
interface McpCodeResponse {
|
const result = yield* McpExa.call(
|
||||||
jsonrpc: string
|
http,
|
||||||
result: {
|
"get_code_context_exa",
|
||||||
content: Array<{
|
McpExa.CodeArgs,
|
||||||
type: string
|
{
|
||||||
text: string
|
query: params.query,
|
||||||
}>
|
tokensNum: params.tokensNum || 5000,
|
||||||
}
|
},
|
||||||
}
|
"30 seconds",
|
||||||
|
)
|
||||||
|
|
||||||
export const CodeSearchTool = Tool.define("codesearch", {
|
return {
|
||||||
description: DESCRIPTION,
|
output:
|
||||||
parameters: z.object({
|
result ??
|
||||||
query: z
|
"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.",
|
||||||
.string()
|
title: `Code search: ${params.query}`,
|
||||||
.describe(
|
metadata: {},
|
||||||
"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'",
|
|
||||||
),
|
|
||||||
tokensNum: z
|
|
||||||
.number()
|
|
||||||
.min(1000)
|
|
||||||
.max(50000)
|
|
||||||
.default(5000)
|
|
||||||
.describe(
|
|
||||||
"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.",
|
|
||||||
),
|
|
||||||
}),
|
|
||||||
async execute(params, ctx) {
|
|
||||||
await ctx.ask({
|
|
||||||
permission: "codesearch",
|
|
||||||
patterns: [params.query],
|
|
||||||
always: ["*"],
|
|
||||||
metadata: {
|
|
||||||
query: params.query,
|
|
||||||
tokensNum: params.tokensNum,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
|
|
||||||
const codeRequest: McpCodeRequest = {
|
|
||||||
jsonrpc: "2.0",
|
|
||||||
id: 1,
|
|
||||||
method: "tools/call",
|
|
||||||
params: {
|
|
||||||
name: "get_code_context_exa",
|
|
||||||
arguments: {
|
|
||||||
query: params.query,
|
|
||||||
tokensNum: params.tokensNum || 5000,
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
const { signal, clearTimeout } = abortAfterAny(30000, ctx.abort)
|
|
||||||
|
|
||||||
try {
|
|
||||||
const headers: Record<string, string> = {
|
|
||||||
accept: "application/json, text/event-stream",
|
|
||||||
"content-type": "application/json",
|
|
||||||
}
|
|
||||||
|
|
||||||
const response = await fetch(`${API_CONFIG.BASE_URL}${API_CONFIG.ENDPOINTS.CONTEXT}`, {
|
|
||||||
method: "POST",
|
|
||||||
headers,
|
|
||||||
body: JSON.stringify(codeRequest),
|
|
||||||
signal,
|
|
||||||
})
|
|
||||||
|
|
||||||
clearTimeout()
|
|
||||||
|
|
||||||
if (!response.ok) {
|
|
||||||
const errorText = await response.text()
|
|
||||||
throw new Error(`Code search error (${response.status}): ${errorText}`)
|
|
||||||
}
|
|
||||||
|
|
||||||
const responseText = await response.text()
|
|
||||||
|
|
||||||
// Parse SSE response
|
|
||||||
const lines = responseText.split("\n")
|
|
||||||
for (const line of lines) {
|
|
||||||
if (line.startsWith("data: ")) {
|
|
||||||
const data: McpCodeResponse = JSON.parse(line.substring(6))
|
|
||||||
if (data.result && data.result.content && data.result.content.length > 0) {
|
|
||||||
return {
|
|
||||||
output: data.result.content[0].text,
|
|
||||||
title: `Code search: ${params.query}`,
|
|
||||||
metadata: {},
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}).pipe(Effect.runPromise),
|
||||||
}
|
|
||||||
|
|
||||||
return {
|
|
||||||
output:
|
|
||||||
"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.",
|
|
||||||
title: `Code search: ${params.query}`,
|
|
||||||
metadata: {},
|
|
||||||
}
|
|
||||||
} catch (error) {
|
|
||||||
clearTimeout()
|
|
||||||
|
|
||||||
if (error instanceof Error && error.name === "AbortError") {
|
|
||||||
throw new Error("Code search request timed out")
|
|
||||||
}
|
|
||||||
|
|
||||||
throw error
|
|
||||||
}
|
}
|
||||||
},
|
}),
|
||||||
})
|
)
|
||||||
|
|
|
||||||
|
|
@ -102,6 +102,7 @@ export namespace ToolRegistry {
|
||||||
const plan = yield* PlanExitTool
|
const plan = yield* PlanExitTool
|
||||||
const webfetch = yield* WebFetchTool
|
const webfetch = yield* WebFetchTool
|
||||||
const websearch = yield* WebSearchTool
|
const websearch = yield* WebSearchTool
|
||||||
|
const codesearch = yield* CodeSearchTool
|
||||||
|
|
||||||
const state = yield* InstanceState.make<State>(
|
const state = yield* InstanceState.make<State>(
|
||||||
Effect.fn("ToolRegistry.state")(function* (ctx) {
|
Effect.fn("ToolRegistry.state")(function* (ctx) {
|
||||||
|
|
@ -170,7 +171,7 @@ export namespace ToolRegistry {
|
||||||
fetch: Tool.init(webfetch),
|
fetch: Tool.init(webfetch),
|
||||||
todo: Tool.init(todo),
|
todo: Tool.init(todo),
|
||||||
search: Tool.init(websearch),
|
search: Tool.init(websearch),
|
||||||
code: Tool.init(CodeSearchTool),
|
code: Tool.init(codesearch),
|
||||||
skill: Tool.init(SkillTool),
|
skill: Tool.init(SkillTool),
|
||||||
patch: Tool.init(ApplyPatchTool),
|
patch: Tool.init(ApplyPatchTool),
|
||||||
question: Tool.init(question),
|
question: Tool.init(question),
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue