e2b tool - fix initialization and improve tool description (#12345)

This commit is contained in:
Vasek Mlejnsky 2023-10-26 08:47:50 -07:00 committed by GitHub
parent 8ec7aade9f
commit cdd75b687e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 16 deletions

View File

@ -33,6 +33,7 @@
] ]
}, },
{ {
"attachments": {},
"cell_type": "markdown", "cell_type": "markdown",
"metadata": {}, "metadata": {},
"source": [ "source": [

View File

@ -26,9 +26,7 @@ Script should be pure python code that can be evaluated. \
It should be in python format NOT markdown. \ It should be in python format NOT markdown. \
The code should NOT be wrapped in backticks. \ The code should NOT be wrapped in backticks. \
All python packages including requests, matplotlib, scipy, numpy, pandas, \ All python packages including requests, matplotlib, scipy, numpy, pandas, \
etc are available. \ etc are available. Create and display chart using `plt.show()`."""
If you have any files outputted write them to "/home/user" directory \
path."""
def _unparse(tree: ast.AST) -> str: def _unparse(tree: ast.AST) -> str:
@ -99,7 +97,7 @@ class E2BDataAnalysisTool(BaseTool):
name = "e2b_data_analysis" name = "e2b_data_analysis"
args_schema: Type[BaseModel] = E2BDataAnalysisToolArguments args_schema: Type[BaseModel] = E2BDataAnalysisToolArguments
session: Any session: Any
_uploaded_files: List[UploadedFile] = Field(default_factory=list) uploaded_files: List[UploadedFile] = Field(default_factory=list)
def __init__( def __init__(
self, self,
@ -130,23 +128,21 @@ class E2BDataAnalysisTool(BaseTool):
on_exit=on_exit, on_exit=on_exit,
on_artifact=on_artifact, on_artifact=on_artifact,
) )
super().__init__(session=session, **kwargs) super().__init__(session=session, description=base_description, **kwargs)
self.description = ( self.uploaded_files = []
base_description + "\n\n" + self.uploaded_files_description
).strip()
def close(self) -> None: def close(self) -> None:
"""Close the cloud sandbox.""" """Close the cloud sandbox."""
self._uploaded_files = [] self.uploaded_files = []
self.session.close() self.session.close()
@property @property
def uploaded_files_description(self) -> str: def uploaded_files_description(self) -> str:
if len(self._uploaded_files) == 0: if len(self.uploaded_files) == 0:
return "" return ""
lines = ["The following files available in the sandbox:"] lines = ["The following files available in the sandbox:"]
for f in self._uploaded_files: for f in self.uploaded_files:
if f.description == "": if f.description == "":
lines.append(f"- path: `{f.remote_path}`") lines.append(f"- path: `{f.remote_path}`")
else: else:
@ -210,16 +206,14 @@ class E2BDataAnalysisTool(BaseTool):
remote_path=remote_path, remote_path=remote_path,
description=description, description=description,
) )
self._uploaded_files.append(f) self.uploaded_files.append(f)
return f return f
def remove_uploaded_file(self, uploaded_file: UploadedFile) -> None: def remove_uploaded_file(self, uploaded_file: UploadedFile) -> None:
"""Remove uploaded file from the sandbox.""" """Remove uploaded file from the sandbox."""
self.session.filesystem.remove(uploaded_file.remote_path) self.session.filesystem.remove(uploaded_file.remote_path)
self._uploaded_files = [ self.uploaded_files = [
f f for f in self.uploaded_files if f.remote_path != uploaded_file.remote_path
for f in self._uploaded_files
if f.remote_path != uploaded_file.remote_path
] ]
def as_tool(self) -> Tool: def as_tool(self) -> Tool: