feat(tools): switch PythonTool and CustomTool to HTTP API

.env.template

@@ -50,7 +50,7 @@ BEE_FRAMEWORK_LOG_SINGLE_LINE="false"
 # GCP_VERTEXAI_LOCATION=""
 
 # Tools
-# CODE_INTERPRETER_URL="http://127.0.0.1:50051"
+# CODE_INTERPRETER_URL="http://127.0.0.1:50081"
 
 # For Google Search Tool
 # GOOGLE_API_KEY="your-google-api-key"

package.json

@@ -168,13 +168,10 @@
   },
   "dependencies": {
     "@ai-zen/node-fetch-event-source": "^2.1.4",
-    "@connectrpc/connect": "^1.6.1",
-    "@connectrpc/connect-node": "^1.6.1",
     "@opentelemetry/api": "^1.9.0",
     "@streamparser/json": "^0.0.21",
     "ajv": "^8.17.1",
     "ajv-formats": "^3.0.1",
-    "bee-proto": "0.0.2",
     "duck-duck-scrape": "^2.2.7",
     "fast-xml-parser": "^4.5.0",
     "header-generator": "^2.1.57",

src/tools/custom.test.ts

@@ -20,36 +20,29 @@ import { StringToolOutput } from "./base.js";
 
 const mocks = vi.hoisted(() => {
   return {
-    parseCustomTool: vi.fn(),
-    executeCustomTool: vi.fn(),
+    fetch: vi.fn(),
   };
 });
 
-vi.mock("@connectrpc/connect", () => ({
-  createPromiseClient: vi.fn().mockReturnValue({
-    parseCustomTool: mocks.parseCustomTool,
-    executeCustomTool: mocks.executeCustomTool,
-  }),
-}));
+vi.stubGlobal("fetch", mocks.fetch);
 
 describe("CustomTool", () => {
   it("should instantiate correctly", async () => {
-    mocks.parseCustomTool.mockResolvedValue({
-      response: {
-        case: "success",
-        value: {
-          toolName: "test",
-          toolDescription: "A test tool",
-          toolInputSchemaJson: `{
-            "$schema": "http://json-schema.org/draft-07/schema#",
-            "type": "object",
-            "properties": {
-              "a": { "type": "integer" },
-              "b": { "type": "string" }
-            }
-          }`,
-        },
-      },
+    mocks.fetch.mockResolvedValueOnce({
+      ok: true,
+      json: () =>
+        Promise.resolve({
+          tool_name: "test",
+          tool_description: "A test tool",
+          tool_input_schema_json: JSON.stringify({
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              a: { type: "integer" },
+              b: { type: "string" },
+            },
+          }),
+        }),
     });
 
     const customTool = await CustomTool.fromSourceCode({ url: "http://localhost" }, "source code");
@@ -66,53 +59,45 @@ describe("CustomTool", () => {
     });
   });
 
-  it("should throw InvalidCustomToolError on parse error", async () => {
-    mocks.parseCustomTool.mockResolvedValue({
-      response: {
-        case: "error",
-        value: {
-          errorMessages: ["Error parsing tool"],
-        },
-      },
+  it("should throw CustomToolCreateError on parse error", async () => {
+    mocks.fetch.mockRejectedValueOnce({
+      cause: { name: "HTTPParserError" },
     });
 
     await expect(
       CustomTool.fromSourceCode({ url: "http://localhost" }, "source code"),
-    ).rejects.toThrow("Error parsing tool");
+    ).rejects.toThrow(
+      "Request to bee-code-interpreter has failed -- ensure that CODE_INTERPRETER_URL points to the new HTTP endpoint",
+    );
   });
 
   it("should run the custom tool", async () => {
-    mocks.parseCustomTool.mockResolvedValue({
-      response: {
-        case: "success",
-        value: {
-          toolName: "test",
-          toolDescription: "A test tool",
-          toolInputSchemaJson: `{
-            "$schema": "http://json-schema.org/draft-07/schema#",
-            "type": "object",
-            "properties": {
-              "a": { "type": "integer" },
-              "b": { "type": "string" }
-            }
-          }`,
-        },
-      },
+    mocks.fetch.mockResolvedValueOnce({
+      ok: true,
+      json: () =>
+        Promise.resolve({
+          tool_name: "test",
+          tool_description: "A test tool",
+          tool_input_schema_json: JSON.stringify({
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              a: { type: "integer" },
+              b: { type: "string" },
+            },
+          }),
+        }),
     });
 
-    const customTool = await CustomTool.fromSourceCode(
-      { url: "http://localhost" },
-      "source code",
-      "executor-id",
-    );
+    const customTool = await CustomTool.fromSourceCode({ url: "http://localhost" }, "source code");
 
-    mocks.executeCustomTool.mockResolvedValue({
-      response: {
-        case: "success",
-        value: {
-          toolOutputJson: '{"something": "42"}',
-        },
-      },
+    mocks.fetch.mockResolvedValueOnce({
+      ok: true,
+      json: () =>
+        Promise.resolve({
+          exit_code: 0,
+          tool_output_json: '{"something": "42"}',
+        }),
     });
 
     const result = await customTool.run(
@@ -128,38 +113,32 @@ describe("CustomTool", () => {
     expect(result.getTextContent()).toEqual('{"something": "42"}');
   });
 
-  it("should throw CustomToolExecutionError on execution error", async () => {
-    mocks.parseCustomTool.mockResolvedValue({
-      response: {
-        case: "success",
-        value: {
-          toolName: "test",
-          toolDescription: "A test tool",
-          toolInputSchemaJson: `{
-            "$schema": "http://json-schema.org/draft-07/schema#",
-            "type": "object",
-            "properties": {
-              "a": { "type": "integer" },
-              "b": { "type": "string" }
-            }
-          }`,
-        },
-      },
+  it("should throw CustomToolExecuteError on execution error", async () => {
+    mocks.fetch.mockResolvedValueOnce({
+      ok: true,
+      json: () =>
+        Promise.resolve({
+          tool_name: "test",
+          tool_description: "A test tool",
+          tool_input_schema_json: JSON.stringify({
+            $schema: "http://json-schema.org/draft-07/schema#",
+            type: "object",
+            properties: {
+              a: { type: "integer" },
+              b: { type: "string" },
+            },
+          }),
+        }),
     });
 
-    const customTool = await CustomTool.fromSourceCode(
-      { url: "http://localhost" },
-      "source code",
-      "executor-id",
-    );
+    const customTool = await CustomTool.fromSourceCode({ url: "http://localhost" }, "source code");
 
-    mocks.executeCustomTool.mockResolvedValue({
-      response: {
-        case: "error",
-        value: {
-          stderr: "Error executing tool",
-        },
-      },
+    mocks.fetch.mockResolvedValueOnce({
+      ok: true,
+      json: () =>
+        Promise.resolve({
+          stderr: "Oh no, it does not work",
+        }),
     });
 
     await expect(
@@ -174,4 +153,16 @@ describe("CustomTool", () => {
       ),
     ).rejects.toThrow('Tool "test" has occurred an error!');
   });
+
+  it("should handle HTTP error responses", async () => {
+    mocks.fetch.mockResolvedValueOnce({
+      ok: false,
+      status: 500,
+      text: () => Promise.resolve("Internal Server Error"),
+    });
+
+    await expect(
+      CustomTool.fromSourceCode({ url: "http://localhost" }, "source code"),
+    ).rejects.toThrow("Request to bee-code-interpreter has failed with HTTP status code 500");
+  });
 });

src/tools/custom.ts

@@ -22,15 +22,13 @@ import {
   Tool,
   ToolInput,
 } from "@/tools/base.js";
-import { createGrpcTransport } from "@connectrpc/connect-node";
-import { PromiseClient, createPromiseClient } from "@connectrpc/connect";
 import { FrameworkError } from "@/errors.js";
 import { z } from "zod";
 import { validate } from "@/internals/helpers/general.js";
-import { CodeInterpreterService } from "bee-proto/code_interpreter/v1/code_interpreter_service_connect";
-import { CodeInterpreterOptions } from "./python/python.js";
+import { callCodeInterpreter, CodeInterpreterOptions } from "./python/python.js";
 import { RunContext } from "@/context.js";
 import { Emitter } from "@/emitter/emitter.js";
+import { ToolError } from "@/tools/base.js";
 
 export class CustomToolCreateError extends FrameworkError {}
 export class CustomToolExecuteError extends FrameworkError {}
@@ -45,23 +43,11 @@ const toolOptionsSchema = z
     name: z.string().min(1),
     description: z.string().min(1),
     inputSchema: z.any(),
-    executorId: z.string().nullable().optional(),
   })
   .passthrough();
 
 export type CustomToolOptions = z.output<typeof toolOptionsSchema> & BaseToolOptions;
 
-function createCodeInterpreterClient(codeInterpreter: CodeInterpreterOptions) {
-  return createPromiseClient(
-    CodeInterpreterService,
-    createGrpcTransport({
-      baseUrl: codeInterpreter.url,
-      httpVersion: "2",
-      nodeOptions: codeInterpreter.connectionOptions,
-    }),
-  );
-}
-
 export class CustomTool extends Tool<StringToolOutput, CustomToolOptions> {
   name: string;
   description: string;
@@ -75,19 +61,13 @@ export class CustomTool extends Tool<StringToolOutput, CustomToolOptions> {
     return this.options.inputSchema;
   }
 
-  protected client: PromiseClient<typeof CodeInterpreterService>;
-
   static {
     this.register();
   }
 
-  public constructor(
-    options: CustomToolOptions,
-    client?: PromiseClient<typeof CodeInterpreterService>,
-  ) {
+  public constructor(options: CustomToolOptions) {
     validate(options, toolOptionsSchema);
     super(options);
-    this.client = client || createCodeInterpreterClient(options.codeInterpreter);
     this.name = options.name;
     this.description = options.description;
   }
@@ -97,51 +77,60 @@ export class CustomTool extends Tool<StringToolOutput, CustomToolOptions> {
     _options: Partial<BaseToolRunOptions>,
     run: RunContext<typeof this>,
   ) {
-    const { response } = await this.client.executeCustomTool(
-      {
-        executorId: this.options.executorId || "default",
-        toolSourceCode: this.options.sourceCode,
-        toolInputJson: JSON.stringify(input),
-      },
-      { signal: run.signal },
-    );
-
-    if (response.case === "error") {
-      throw new CustomToolExecuteError(response.value.stderr);
+    try {
+      const result = await callCodeInterpreter({
+        url: `${this.options.codeInterpreter.url}/v1/execute-custom-tool`,
+        body: {
+          tool_source_code: this.options.sourceCode,
+          tool_input_json: JSON.stringify(input),
+        },
+        signal: run.signal,
+      });
+
+      if (result.stderr) {
+        throw new CustomToolExecuteError(result.stderr);
+      }
+
+      return new StringToolOutput(result.tool_output_json);
+    } catch (e) {
+      if (e instanceof ToolError) {
+        throw new CustomToolExecuteError(e.message, [e]);
+      } else {
+        throw e;
+      }
     }
-
-    return new StringToolOutput(response.value!.toolOutputJson);
   }
 
   loadSnapshot(snapshot: ReturnType<typeof this.createSnapshot>): void {
     super.loadSnapshot(snapshot);
-    this.client = createCodeInterpreterClient(this.options.codeInterpreter);
   }
 
-  static async fromSourceCode(
-    codeInterpreter: CodeInterpreterOptions,
-    sourceCode: string,
-    executorId?: string,
-  ) {
-    const client = createCodeInterpreterClient(codeInterpreter);
-    const response = await client.parseCustomTool({ toolSourceCode: sourceCode });
+  static async fromSourceCode(codeInterpreter: CodeInterpreterOptions, sourceCode: string) {
+    try {
+      const result = await callCodeInterpreter({
+        url: `${codeInterpreter.url}/v1/parse-custom-tool`,
+        body: {
+          tool_source_code: sourceCode,
+        },
+      });
 
-    if (response.response.case === "error") {
-      throw new CustomToolCreateError(response.response.value.errorMessages.join("\n"));
-    }
+      if (result.error_messages) {
+        throw new CustomToolCreateError(result.errorMessages.join("\n"));
+      }
 
-    const { toolName, toolDescription, toolInputSchemaJson } = response.response.value!;
-
-    return new CustomTool(
-      {
+      return new CustomTool({
         codeInterpreter,
         sourceCode,
-        name: toolName,
-        description: toolDescription,
-        inputSchema: JSON.parse(toolInputSchemaJson),
-        executorId,
-      },
-      client,
-    );
+        name: result.tool_name,
+        description: result.tool_description,
+        inputSchema: JSON.parse(result.tool_input_schema_json),
+      });
+    } catch (e) {
+      if (e instanceof ToolError) {
+        throw new CustomToolCreateError(e.message, [e]);
+      } else {
+        throw e;
+      }
+    }
   }
 }