From ada3d2cbd129e730af1bfcc21cc1e51be015a962 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Nov=C3=A1k?= <jakub@e2b.dev>
Date: Mon, 6 Nov 2023 16:29:47 +0100
Subject: [PATCH] Add possibility to pass on_artifacts for a specific
 conversation (#12687)

Possibility to pass on_artifacts to a conversation. It can be then
achieved by adding this way:

```python
result = agent.run(
    input=message.text,
    metadata={
        "on_artifact": CALLBACK_FUNCTION
    },
)
```
---
 .../langchain/tools/e2b_data_analysis/tool.py   | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/libs/langchain/langchain/tools/e2b_data_analysis/tool.py b/libs/langchain/langchain/tools/e2b_data_analysis/tool.py
index 64c62bc4491..fb8d0a789db 100644
--- a/libs/langchain/langchain/tools/e2b_data_analysis/tool.py
+++ b/libs/langchain/langchain/tools/e2b_data_analysis/tool.py
@@ -9,6 +9,7 @@ from typing import IO, TYPE_CHECKING, Any, Callable, List, Optional, Type
 
 from langchain.callbacks.manager import (
     AsyncCallbackManagerForToolRun,
+    CallbackManager,
     CallbackManagerForToolRun,
 )
 from langchain.pydantic_v1 import BaseModel, Field, PrivateAttr
@@ -151,14 +152,26 @@ class E2BDataAnalysisTool(BaseTool):
         return "\n".join(lines)
 
     def _run(
-        self, python_code: str, run_manager: Optional[CallbackManagerForToolRun] = None
+        self,
+        python_code: str,
+        run_manager: Optional[CallbackManagerForToolRun] = None,
+        callbacks: Optional[CallbackManager] = None,
     ) -> str:
         python_code = add_last_line_print(python_code)
-        stdout, stderr, _ = self.session.run_python(python_code)
+
+        if callbacks is not None:
+            on_artifact = getattr(callbacks.metadata, "on_artifact", None)
+        else:
+            on_artifact = None
+
+        stdout, stderr, artifacts = self.session.run_python(
+            python_code, on_artifact=on_artifact
+        )
 
         out = {
             "stdout": stdout,
             "stderr": stderr,
+            "artifacts": list(map(lambda artifact: artifact.name, artifacts)),
         }
         return json.dumps(out)