refactor(tool): convert websearch tool internals to Effect (#21810)
This commit is contained in:
@@ -0,0 +1,74 @@
|
|||||||
|
import { Duration, Effect, Schema } from "effect"
|
||||||
|
import { HttpClient, HttpClientRequest } from "effect/unstable/http"
|
||||||
|
|
||||||
|
const URL = "https://mcp.exa.ai/mcp"
|
||||||
|
|
||||||
|
const McpResult = Schema.Struct({
|
||||||
|
result: Schema.Struct({
|
||||||
|
content: Schema.Array(
|
||||||
|
Schema.Struct({
|
||||||
|
type: Schema.String,
|
||||||
|
text: Schema.String,
|
||||||
|
}),
|
||||||
|
),
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
|
||||||
|
const decode = Schema.decodeUnknownEffect(Schema.fromJsonString(McpResult))
|
||||||
|
|
||||||
|
const parseSse = Effect.fn("McpExa.parseSse")(function* (body: string) {
|
||||||
|
for (const line of body.split("\n")) {
|
||||||
|
if (!line.startsWith("data: ")) continue
|
||||||
|
const data = yield* decode(line.substring(6))
|
||||||
|
if (data.result.content[0]?.text) return data.result.content[0].text
|
||||||
|
}
|
||||||
|
return undefined
|
||||||
|
})
|
||||||
|
|
||||||
|
export const SearchArgs = Schema.Struct({
|
||||||
|
query: Schema.String,
|
||||||
|
type: Schema.String,
|
||||||
|
numResults: Schema.Number,
|
||||||
|
livecrawl: Schema.String,
|
||||||
|
contextMaxCharacters: Schema.optional(Schema.Number),
|
||||||
|
})
|
||||||
|
|
||||||
|
export const CodeArgs = Schema.Struct({
|
||||||
|
query: Schema.String,
|
||||||
|
tokensNum: Schema.Number,
|
||||||
|
})
|
||||||
|
|
||||||
|
const McpRequest = <F extends Schema.Struct.Fields>(args: Schema.Struct<F>) =>
|
||||||
|
Schema.Struct({
|
||||||
|
jsonrpc: Schema.Literal("2.0"),
|
||||||
|
id: Schema.Literal(1),
|
||||||
|
method: Schema.Literal("tools/call"),
|
||||||
|
params: Schema.Struct({
|
||||||
|
name: Schema.String,
|
||||||
|
arguments: args,
|
||||||
|
}),
|
||||||
|
})
|
||||||
|
|
||||||
|
export const call = <F extends Schema.Struct.Fields>(
|
||||||
|
http: HttpClient.HttpClient,
|
||||||
|
tool: string,
|
||||||
|
args: Schema.Struct<F>,
|
||||||
|
value: Schema.Struct.Type<F>,
|
||||||
|
timeout: Duration.Input,
|
||||||
|
) =>
|
||||||
|
Effect.gen(function* () {
|
||||||
|
const request = yield* HttpClientRequest.post(URL).pipe(
|
||||||
|
HttpClientRequest.accept("application/json, text/event-stream"),
|
||||||
|
HttpClientRequest.schemaBodyJson(McpRequest(args))({
|
||||||
|
jsonrpc: "2.0" as const,
|
||||||
|
id: 1 as const,
|
||||||
|
method: "tools/call" as const,
|
||||||
|
params: { name: tool, arguments: value },
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
const response = yield* HttpClient.filterStatusOk(http).execute(request).pipe(
|
||||||
|
Effect.timeoutOrElse({ duration: timeout, orElse: () => Effect.die(new Error(`${tool} request timed out`)) }),
|
||||||
|
)
|
||||||
|
const body = yield* response.text
|
||||||
|
return yield* parseSse(body)
|
||||||
|
})
|
||||||
@@ -101,6 +101,7 @@ export namespace ToolRegistry {
|
|||||||
const lsptool = yield* LspTool
|
const lsptool = yield* LspTool
|
||||||
const plan = yield* PlanExitTool
|
const plan = yield* PlanExitTool
|
||||||
const webfetch = yield* WebFetchTool
|
const webfetch = yield* WebFetchTool
|
||||||
|
const websearch = yield* WebSearchTool
|
||||||
|
|
||||||
const state = yield* InstanceState.make<State>(
|
const state = yield* InstanceState.make<State>(
|
||||||
Effect.fn("ToolRegistry.state")(function* (ctx) {
|
Effect.fn("ToolRegistry.state")(function* (ctx) {
|
||||||
@@ -168,7 +169,7 @@ export namespace ToolRegistry {
|
|||||||
task: Tool.init(task),
|
task: Tool.init(task),
|
||||||
fetch: Tool.init(webfetch),
|
fetch: Tool.init(webfetch),
|
||||||
todo: Tool.init(todo),
|
todo: Tool.init(todo),
|
||||||
search: Tool.init(WebSearchTool),
|
search: Tool.init(websearch),
|
||||||
code: Tool.init(CodeSearchTool),
|
code: Tool.init(CodeSearchTool),
|
||||||
skill: Tool.init(SkillTool),
|
skill: Tool.init(SkillTool),
|
||||||
patch: Tool.init(ApplyPatchTool),
|
patch: Tool.init(ApplyPatchTool),
|
||||||
|
|||||||
@@ -1,15 +1,9 @@
|
|||||||
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 "./websearch.txt"
|
import DESCRIPTION from "./websearch.txt"
|
||||||
import { abortAfterAny } from "../util/abort"
|
|
||||||
|
|
||||||
const API_CONFIG = {
|
|
||||||
BASE_URL: "https://mcp.exa.ai",
|
|
||||||
ENDPOINTS: {
|
|
||||||
SEARCH: "/mcp",
|
|
||||||
},
|
|
||||||
DEFAULT_NUM_RESULTS: 8,
|
|
||||||
} as const
|
|
||||||
|
|
||||||
const Parameters = z.object({
|
const Parameters = z.object({
|
||||||
query: z.string().describe("Websearch query"),
|
query: z.string().describe("Websearch query"),
|
||||||
@@ -30,121 +24,53 @@ const Parameters = z.object({
|
|||||||
.describe("Maximum characters for context string optimized for LLMs (default: 10000)"),
|
.describe("Maximum characters for context string optimized for LLMs (default: 10000)"),
|
||||||
})
|
})
|
||||||
|
|
||||||
interface McpSearchRequest {
|
export const WebSearchTool = Tool.defineEffect(
|
||||||
jsonrpc: string
|
"websearch",
|
||||||
id: number
|
Effect.gen(function* () {
|
||||||
method: string
|
const http = yield* HttpClient.HttpClient
|
||||||
params: {
|
|
||||||
name: string
|
|
||||||
arguments: {
|
|
||||||
query: string
|
|
||||||
numResults?: number
|
|
||||||
livecrawl?: "fallback" | "preferred"
|
|
||||||
type?: "auto" | "fast" | "deep"
|
|
||||||
contextMaxCharacters?: number
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
interface McpSearchResponse {
|
return {
|
||||||
jsonrpc: string
|
get description() {
|
||||||
result: {
|
return DESCRIPTION.replace("{{year}}", new Date().getFullYear().toString())
|
||||||
content: Array<{
|
},
|
||||||
type: string
|
parameters: Parameters,
|
||||||
text: string
|
execute: (params: z.infer<typeof Parameters>, ctx: Tool.Context) =>
|
||||||
}>
|
Effect.gen(function* () {
|
||||||
}
|
yield* Effect.promise(() =>
|
||||||
}
|
ctx.ask({
|
||||||
|
permission: "websearch",
|
||||||
|
patterns: [params.query],
|
||||||
|
always: ["*"],
|
||||||
|
metadata: {
|
||||||
|
query: params.query,
|
||||||
|
numResults: params.numResults,
|
||||||
|
livecrawl: params.livecrawl,
|
||||||
|
type: params.type,
|
||||||
|
contextMaxCharacters: params.contextMaxCharacters,
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
export const WebSearchTool = Tool.define("websearch", async () => {
|
const result = yield* McpExa.call(
|
||||||
return {
|
http,
|
||||||
get description() {
|
"web_search_exa",
|
||||||
return DESCRIPTION.replace("{{year}}", new Date().getFullYear().toString())
|
McpExa.SearchArgs,
|
||||||
},
|
{
|
||||||
parameters: Parameters,
|
query: params.query,
|
||||||
async execute(params, ctx) {
|
type: params.type || "auto",
|
||||||
await ctx.ask({
|
numResults: params.numResults || 8,
|
||||||
permission: "websearch",
|
livecrawl: params.livecrawl || "fallback",
|
||||||
patterns: [params.query],
|
contextMaxCharacters: params.contextMaxCharacters,
|
||||||
always: ["*"],
|
},
|
||||||
metadata: {
|
"25 seconds",
|
||||||
query: params.query,
|
)
|
||||||
numResults: params.numResults,
|
|
||||||
livecrawl: params.livecrawl,
|
|
||||||
type: params.type,
|
|
||||||
contextMaxCharacters: params.contextMaxCharacters,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
|
|
||||||
const searchRequest: McpSearchRequest = {
|
return {
|
||||||
jsonrpc: "2.0",
|
output: result ?? "No search results found. Please try a different query.",
|
||||||
id: 1,
|
title: `Web search: ${params.query}`,
|
||||||
method: "tools/call",
|
metadata: {},
|
||||||
params: {
|
|
||||||
name: "web_search_exa",
|
|
||||||
arguments: {
|
|
||||||
query: params.query,
|
|
||||||
type: params.type || "auto",
|
|
||||||
numResults: params.numResults || API_CONFIG.DEFAULT_NUM_RESULTS,
|
|
||||||
livecrawl: params.livecrawl || "fallback",
|
|
||||||
contextMaxCharacters: params.contextMaxCharacters,
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
const { signal, clearTimeout } = abortAfterAny(25000, 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.SEARCH}`, {
|
|
||||||
method: "POST",
|
|
||||||
headers,
|
|
||||||
body: JSON.stringify(searchRequest),
|
|
||||||
signal,
|
|
||||||
})
|
|
||||||
|
|
||||||
clearTimeout()
|
|
||||||
|
|
||||||
if (!response.ok) {
|
|
||||||
const errorText = await response.text()
|
|
||||||
throw new Error(`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: McpSearchResponse = JSON.parse(line.substring(6))
|
|
||||||
if (data.result && data.result.content && data.result.content.length > 0) {
|
|
||||||
return {
|
|
||||||
output: data.result.content[0].text,
|
|
||||||
title: `Web search: ${params.query}`,
|
|
||||||
metadata: {},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}).pipe(Effect.runPromise),
|
||||||
|
}
|
||||||
return {
|
}),
|
||||||
output: "No search results found. Please try a different query.",
|
)
|
||||||
title: `Web search: ${params.query}`,
|
|
||||||
metadata: {},
|
|
||||||
}
|
|
||||||
} catch (error) {
|
|
||||||
clearTimeout()
|
|
||||||
|
|
||||||
if (error instanceof Error && error.name === "AbortError") {
|
|
||||||
throw new Error("Search request timed out")
|
|
||||||
}
|
|
||||||
|
|
||||||
throw error
|
|
||||||
}
|
|
||||||
},
|
|
||||||
}
|
|
||||||
})
|
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
import { NodeFileSystem } from "@effect/platform-node"
|
import { NodeFileSystem } from "@effect/platform-node"
|
||||||
|
import { FetchHttpClient } from "effect/unstable/http"
|
||||||
import { expect } from "bun:test"
|
import { expect } from "bun:test"
|
||||||
import { Cause, Effect, Exit, Fiber, Layer } from "effect"
|
import { Cause, Effect, Exit, Fiber, Layer } from "effect"
|
||||||
import { FetchHttpClient } from "effect/unstable/http"
|
|
||||||
import path from "path"
|
import path from "path"
|
||||||
import z from "zod"
|
import z from "zod"
|
||||||
import { Agent as AgentSvc } from "../../src/agent/agent"
|
import { Agent as AgentSvc } from "../../src/agent/agent"
|
||||||
|
|||||||
Reference in New Issue
Block a user