Add the print statement to the LLM-generated python code (superagen…

…t-ai#462)

* Add print to the LLM-generated python code

* Update e2b package

* Shorten import

* Improve code based on suggestions

* Fix formatting

---------

Co-authored-by: Ismail Pelaseyed <[email protected]>

libs/superagent/app/tools/e2b.py

@@ -1,6 +1,7 @@
 # flake8: noqa
+import ast
 from decouple import config
-from e2b.templates.data_analysis import DataAnalysis
+from e2b import DataAnalysis
 from langchain.tools import BaseTool
 
 
@@ -14,6 +15,21 @@ class E2BCodeExecutor(BaseTool):
     # If the "E2B_API_KEY" env var is set, E2B automatically loads it, no need to pass it to the constructor.
     _session = DataAnalysis(api_key=config("E2B_API_KEY"))
 
+    def _add_last_line_print(self, code: str):
+        tree = ast.parse(code)
+        node = tree.body[-1]
+        if isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
+            if isinstance(node.value.func, ast.Name) and node.value.func.id == "print":
+                return tree
+        tree.body[-1] = ast.Expr(
+            value=ast.Call(
+                func=ast.Name(id="print", ctx=ast.Load()),
+                args=[node.value],
+                keywords=[],
+            )
+        )
+        return ast.unparse(tree)
+
     # TODO: Once we know the the agent is done, we need to close the E2B session.
     # You most likely want to keep the E2B session active for the whole lifecycle of an agent.
     def _close_session(self):
@@ -28,9 +44,11 @@ def _download_artifact(self, artifact):
         artifact_bytes = artifact.download()
 
     def _run(self, python_code: str) -> str:
+        code = self._add_last_line_print(python_code)
+
         # E2B offers both streaming output and artifacts or retrieving them after the code has finished running.
         stdout, err, artifacts = self._session.run_python(
-            code=python_code,
+            code=code,
             # TODO: To create more responsive UI, you might want to stream stdout, stderr, and artifacts
             on_stdout=lambda line: print("stdout", line),
             on_stderr=lambda line: print("stderr", line),

libs/superagent/pyproject.toml

@@ -49,7 +49,7 @@ langchain = "^0.0.300"
 bs4 = "^0.0.1"
 segment-analytics-python = "^2.2.3"
 replicate = "^0.15.4"
-e2b = "^0.9.0"
+e2b = "^0.9.3"
 
 
 [build-system]