Skip to content

Commit

Permalink
feat!: use python code tool over http
Browse files Browse the repository at this point in the history
* Modify the existing PythonTool over gRPC to go over HTTP.
* Add tests (written for gRPC but adapted for the new one).
* Requires CODE_INTERPRETER_URL env var to point to exposed HTTP port (50081).

BREAKING CHANGE: Requires exposed port and updated CODE_INTERPRETER_URL.

Signed-off-by: Mark Sturdevant <mark.sturdevant@ibm.com>
  • Loading branch information
markstur committed Dec 13, 2024
1 parent 8d45d6e commit 0143054
Showing 4 changed files with 49 additions and 262 deletions.
Original file line number Diff line number Diff line change
@@ -15,26 +15,26 @@
*/

import { describe, it, expect } from "vitest";
import { PythonHttpTool } from "@/tools/python/python_http.js";
import { PythonTool } from "@/tools/python/python.js";
import { verifyDeserialization } from "@tests/e2e/utils.js";
import { LocalPythonStorage } from "@/tools/python/storage.js";

const codeInterpreterUrl = process.env.CODE_INTERPRETER_URL || "http://localhost:50081";

const getPythonTool = () =>
new PythonHttpTool({
new PythonTool({
codeInterpreter: { url: codeInterpreterUrl },
storage: new LocalPythonStorage({
interpreterWorkingDir: "/tmp/code-interpreter-storage",
localWorkingDir: "./test_dir/",
}),
});

describe("PythonHttpTool", () => {
describe("PythonTool", () => {
it("Is the expected tool", () => {
const tool = getPythonTool();
expect(tool).toBeInstanceOf(PythonHttpTool);
expect(PythonHttpTool.isTool(tool)).toBe(true);
expect(tool).toBeInstanceOf(PythonTool);
expect(PythonTool.isTool(tool)).toBe(true);
expect(tool.name).toBe("Python");
expect(tool.description).toMatch("Run Python and/or shell code");
});
@@ -69,7 +69,7 @@ describe("PythonHttpTool", () => {
it("serializes", async () => {
const tool = getPythonTool();
const serialized = tool.serialize();
const deserializedTool = PythonHttpTool.fromSerialized(serialized);
const deserializedTool = PythonTool.fromSerialized(serialized);
verifyDeserialization(tool, deserializedTool);
});
});
70 changes: 41 additions & 29 deletions src/tools/python/python.ts
Original file line number Diff line number Diff line change
@@ -14,10 +14,14 @@
* limitations under the License.
*/

import { BaseToolOptions, BaseToolRunOptions, ToolEmitter, Tool, ToolInput } from "@/tools/base.js";
import { createGrpcTransport } from "@connectrpc/connect-node";
import { PromiseClient, createPromiseClient } from "@connectrpc/connect";
import { CodeInterpreterService } from "bee-proto/code_interpreter/v1/code_interpreter_service_connect";
import {
BaseToolOptions,
BaseToolRunOptions,
ToolEmitter,
Tool,
ToolError,
ToolInput,
} from "@/tools/base.js";
import { z } from "zod";
import { BaseLLMOutput } from "@/llms/base.js";
import { LLM } from "@/llms/llm.js";
@@ -95,7 +99,6 @@ export class PythonTool extends Tool<PythonToolOutput, PythonToolOptions> {
});
}

protected readonly client: PromiseClient<typeof CodeInterpreterService>;
protected readonly preprocess;

public constructor(options: PythonToolOptions) {
@@ -109,7 +112,6 @@ export class PythonTool extends Tool<PythonToolOutput, PythonToolOptions> {
},
]);
}
this.client = this._createClient();
this.preprocess = options.preprocess;
this.storage = options.storage;
}
@@ -118,17 +120,6 @@ export class PythonTool extends Tool<PythonToolOutput, PythonToolOptions> {
this.register();
}

protected _createClient(): PromiseClient<typeof CodeInterpreterService> {
return createPromiseClient(
CodeInterpreterService,
createGrpcTransport({
baseUrl: this.options.codeInterpreter.url,
httpVersion: "2",
nodeOptions: this.options.codeInterpreter.connectionOptions,
}),
);
}

protected async _run(
input: ToolInput<this>,
_options: Partial<BaseToolRunOptions>,
@@ -156,21 +147,42 @@ export class PythonTool extends Tool<PythonToolOutput, PythonToolOptions> {

const prefix = "/workspace/";

const result = await this.client.execute(
{
sourceCode: await getSourceCode(),
executorId: this.options.executorId ?? "default",
files: Object.fromEntries(
inputFiles.map((file) => [`${prefix}${file.filename}`, file.pythonId]),
),
},
{ signal: run.signal },
);
let response;
const httpUrl = this.options.codeInterpreter.url + "/v1/execute";
try {
response = await fetch(httpUrl, {
method: "POST",
headers: {
"Accept": "application/json",
"Content-Type": "application/json",
},
body: JSON.stringify({
source_code: await getSourceCode(),
executorId: this.options.executorId ?? "default",
files: Object.fromEntries(
inputFiles.map((file) => [`${prefix}${file.filename}`, file.pythonId]),
),
}),
});
} catch (error) {
if (error.cause.name == "HTTPParserError") {
throw new ToolError("Python tool over HTTP failed -- not using HTTP endpoint!", [error]);
} else {
throw new ToolError("Python tool over HTTP failed!", [error]);
}
}

if (!response?.ok) {
throw new ToolError("HTTP request failed!", [new Error(await response.text())]);
}

const result = await response.json();

// replace absolute paths in "files" with relative paths by removing "/workspace/"
// skip files that are not in "/workspace"
// skip entries that are also entries in filesInput
const filesOutput = await this.storage.download(
// @ts-ignore
Object.entries(result.files)
.map(([k, v]) => {
const file = { path: k, pythonId: v };
@@ -194,19 +206,19 @@ export class PythonTool extends Tool<PythonToolOutput, PythonToolOptions> {
})
.filter(isTruthy),
);
return new PythonToolOutput(result.stdout, result.stderr, result.exitCode, filesOutput);
return new PythonToolOutput(result.stdout, result.stderr, result.exit_code, filesOutput);
}

createSnapshot() {
return {
...super.createSnapshot(),
files: this.files,
storage: this.storage,
preprocess: this.preprocess,
};
}

loadSnapshot(snapshot: ReturnType<typeof this.createSnapshot>): void {
super.loadSnapshot(snapshot);
Object.assign(this, { client: this._createClient() });
}
}
225 changes: 0 additions & 225 deletions src/tools/python/python_http.ts

This file was deleted.

Loading

0 comments on commit 0143054

Please sign in to comment.