abtsousa commited on
Commit
79d28fc
·
1 Parent(s): 94d07d4

Add multi-language code execution functionality (all credits go to @fisherman611 )

Browse files
Files changed (2) hide show
  1. tools/__init__.py +3 -1
  2. tools/code_interpreter.py +347 -0
tools/__init__.py CHANGED
@@ -1,5 +1,6 @@
1
  from .wikipedia import wiki_search
2
  from .search import web_search
 
3
  from langchain_core.tools import BaseTool
4
 
5
  def get_all_tools() -> list[BaseTool]:
@@ -11,7 +12,8 @@ def get_all_tools() -> list[BaseTool]:
11
  """
12
  tools = [
13
  wiki_search,
14
- web_search
 
15
  ]
16
 
17
  return tools
 
1
  from .wikipedia import wiki_search
2
  from .search import web_search
3
+ from .code_interpreter import execute_code_multilang
4
  from langchain_core.tools import BaseTool
5
 
6
  def get_all_tools() -> list[BaseTool]:
 
12
  """
13
  tools = [
14
  wiki_search,
15
+ web_search,
16
+ execute_code_multilang
17
  ]
18
 
19
  return tools
tools/code_interpreter.py ADDED
@@ -0,0 +1,347 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import io
3
+ import sys
4
+ import uuid
5
+ import base64
6
+ import traceback
7
+ import contextlib
8
+ import tempfile
9
+ import subprocess
10
+ import sqlite3
11
+ from typing import Dict, List, Any, Optional, Union
12
+ import numpy as np
13
+ import pandas as pd
14
+ import matplotlib.pyplot as plt
15
+ from PIL import Image
16
+ from langchain_core.tools import tool
17
+
18
+ @tool
19
+ def execute_code_multilang(code: str, language: str = "python") -> str:
20
+ """Execute code in multiple languages (Python, Bash, SQL, C, Java) and return results.
21
+
22
+ Args:
23
+ code (str): The source code to execute.
24
+ language (str): The language of the code. Supported: "python", "bash", "sql", "c", "java".
25
+
26
+ Returns:
27
+ A string summarizing the execution results (stdout, stderr, errors, plots, dataframes if any).
28
+ """
29
+ supported_languages = ["python", "bash", "sql", "c", "java"]
30
+ language = language.lower()
31
+
32
+ if language not in supported_languages:
33
+ return f"❌ Unsupported language: {language}. Supported languages are: {', '.join(supported_languages)}"
34
+
35
+ result = interpreter_instance.execute_code(code, language=language)
36
+
37
+ response = []
38
+
39
+ if result["status"] == "success":
40
+ response.append(f"✅ Code executed successfully in **{language.upper()}**")
41
+
42
+ if result.get("stdout"):
43
+ response.append(
44
+ "\n**Standard Output:**\n```\n" + result["stdout"].strip() + "\n```"
45
+ )
46
+
47
+ if result.get("stderr"):
48
+ response.append(
49
+ "\n**Standard Error (if any):**\n```\n"
50
+ + result["stderr"].strip()
51
+ + "\n```"
52
+ )
53
+
54
+ if result.get("result") is not None:
55
+ response.append(
56
+ "\n**Execution Result:**\n```\n"
57
+ + str(result["result"]).strip()
58
+ + "\n```"
59
+ )
60
+
61
+ if result.get("dataframes"):
62
+ for df_info in result["dataframes"]:
63
+ response.append(
64
+ f"\n**DataFrame `{df_info['name']}` (Shape: {df_info['shape']})**"
65
+ )
66
+ df_preview = pd.DataFrame(df_info["head"])
67
+ response.append("First 5 rows:\n```\n" + str(df_preview) + "\n```")
68
+
69
+ if result.get("plots"):
70
+ response.append(
71
+ f"\n**Generated {len(result['plots'])} plot(s)** (Image data returned separately)"
72
+ )
73
+
74
+ else:
75
+ response.append(f"❌ Code execution failed in **{language.upper()}**")
76
+ if result.get("stderr"):
77
+ response.append(
78
+ "\n**Error Log:**\n```\n" + result["stderr"].strip() + "\n```"
79
+ )
80
+
81
+ return "\n".join(response)
82
+
83
+ class CodeInterpreter:
84
+ def __init__(self, allowed_modules=None, max_execution_time=30, working_directory=None):
85
+ """Initialize the code interpreter with safety measures."""
86
+ self.allowed_modules = allowed_modules or [
87
+ "numpy", "pandas", "matplotlib", "scipy", "sklearn",
88
+ "math", "random", "statistics", "datetime", "collections",
89
+ "itertools", "functools", "operator", "re", "json",
90
+ "sympy", "networkx", "nltk", "PIL", "pytesseract",
91
+ "cmath", "uuid", "tempfile", "requests", "urllib"
92
+ ]
93
+ self.max_execution_time = max_execution_time
94
+ self.working_directory = working_directory or os.path.join(os.getcwd())
95
+ if not os.path.exists(self.working_directory):
96
+ os.makedirs(self.working_directory)
97
+
98
+ self.globals = {
99
+ "__builtins__": __builtins__,
100
+ "np": np,
101
+ "pd": pd,
102
+ "plt": plt,
103
+ "Image": Image,
104
+ }
105
+ self.temp_sqlite_db = os.path.join(tempfile.gettempdir(), "code_exec.db")
106
+
107
+ def execute_code(self, code: str, language: str = "python") -> Dict[str, Any]:
108
+ """Execute the provided code in the selected programming language."""
109
+ language = language.lower()
110
+ execution_id = str(uuid.uuid4())
111
+
112
+ result = {
113
+ "execution_id": execution_id,
114
+ "status": "error",
115
+ "stdout": "",
116
+ "stderr": "",
117
+ "result": None,
118
+ "plots": [],
119
+ "dataframes": []
120
+ }
121
+
122
+ try:
123
+ if language == "python":
124
+ return self._execute_python(code, execution_id)
125
+ elif language == "bash":
126
+ return self._execute_bash(code, execution_id)
127
+ elif language == "sql":
128
+ return self._execute_sql(code, execution_id)
129
+ elif language == "c":
130
+ return self._execute_c(code, execution_id)
131
+ elif language == "java":
132
+ return self._execute_java(code, execution_id)
133
+ else:
134
+ result["stderr"] = f"Unsupported language: {language}"
135
+ except Exception as e:
136
+ result["stderr"] = str(e)
137
+
138
+ return result
139
+
140
+ def _execute_python(self, code: str, execution_id: str) -> dict:
141
+ output_buffer = io.StringIO()
142
+ error_buffer = io.StringIO()
143
+ result = {
144
+ "execution_id": execution_id,
145
+ "status": "error",
146
+ "stdout": "",
147
+ "stderr": "",
148
+ "result": None,
149
+ "plots": [],
150
+ "dataframes": []
151
+ }
152
+
153
+ try:
154
+ exec_dir = os.path.join(self.working_directory, execution_id)
155
+ os.makedirs(exec_dir, exist_ok=True)
156
+ plt.switch_backend('Agg')
157
+
158
+ with contextlib.redirect_stdout(output_buffer), contextlib.redirect_stderr(error_buffer):
159
+ exec_result = exec(code, self.globals)
160
+
161
+ if plt.get_fignums():
162
+ for i, fig_num in enumerate(plt.get_fignums()):
163
+ fig = plt.figure(fig_num)
164
+ img_path = os.path.join(exec_dir, f"plot_{i}.png")
165
+ fig.savefig(img_path)
166
+ with open(img_path, "rb") as img_file:
167
+ img_data = base64.b64encode(img_file.read()).decode('utf-8')
168
+ result["plots"].append({
169
+ "figure_number": fig_num,
170
+ "data": img_data
171
+ })
172
+
173
+ for var_name, var_value in self.globals.items():
174
+ if isinstance(var_value, pd.DataFrame) and len(var_value) > 0:
175
+ result["dataframes"].append({
176
+ "name": var_name,
177
+ "head": var_value.head().to_dict(),
178
+ "shape": var_value.shape,
179
+ "dtypes": str(var_value.dtypes)
180
+ })
181
+
182
+ result["status"] = "success"
183
+ result["stdout"] = output_buffer.getvalue()
184
+ result["result"] = exec_result
185
+
186
+ except Exception as e:
187
+ result["status"] = "error"
188
+ result["stderr"] = f"{error_buffer.getvalue()}\n{traceback.format_exc()}"
189
+
190
+ return result
191
+
192
+ def _execute_bash(self, code: str, execution_id: str) -> dict:
193
+ try:
194
+ completed = subprocess.run(
195
+ code, shell=True, capture_output=True, text=True, timeout=self.max_execution_time
196
+ )
197
+ return {
198
+ "execution_id": execution_id,
199
+ "status": "success" if completed.returncode == 0 else "error",
200
+ "stdout": completed.stdout,
201
+ "stderr": completed.stderr,
202
+ "result": None,
203
+ "plots": [],
204
+ "dataframes": []
205
+ }
206
+ except subprocess.TimeoutExpired:
207
+ return {
208
+ "execution_id": execution_id,
209
+ "status": "error",
210
+ "stdout": "",
211
+ "stderr": "Execution timed out.",
212
+ "result": None,
213
+ "plots": [],
214
+ "dataframes": []
215
+ }
216
+
217
+ def _execute_sql(self, code: str, execution_id: str) -> dict:
218
+ result = {
219
+ "execution_id": execution_id,
220
+ "status": "error",
221
+ "stdout": "",
222
+ "stderr": "",
223
+ "result": None,
224
+ "plots": [],
225
+ "dataframes": []
226
+ }
227
+ try:
228
+ conn = sqlite3.connect(self.temp_sqlite_db)
229
+ cur = conn.cursor()
230
+ cur.execute(code)
231
+ if code.strip().lower().startswith("select"):
232
+ columns = [description[0] for description in cur.description]
233
+ rows = cur.fetchall()
234
+ df = pd.DataFrame(rows, columns=columns)
235
+ result["dataframes"].append({
236
+ "name": "query_result",
237
+ "head": df.head().to_dict(),
238
+ "shape": df.shape,
239
+ "dtypes": str(df.dtypes)
240
+ })
241
+ else:
242
+ conn.commit()
243
+
244
+ result["status"] = "success"
245
+ result["stdout"] = "Query executed successfully."
246
+
247
+ except Exception as e:
248
+ result["stderr"] = str(e)
249
+ finally:
250
+ conn.close()
251
+
252
+ return result
253
+
254
+ def _execute_c(self, code: str, execution_id: str) -> dict:
255
+ temp_dir = tempfile.mkdtemp()
256
+ source_path = os.path.join(temp_dir, "program.c")
257
+ binary_path = os.path.join(temp_dir, "program")
258
+
259
+ try:
260
+ with open(source_path, "w") as f:
261
+ f.write(code)
262
+
263
+ compile_proc = subprocess.run(
264
+ ["gcc", source_path, "-o", binary_path],
265
+ capture_output=True, text=True, timeout=self.max_execution_time
266
+ )
267
+ if compile_proc.returncode != 0:
268
+ return {
269
+ "execution_id": execution_id,
270
+ "status": "error",
271
+ "stdout": compile_proc.stdout,
272
+ "stderr": compile_proc.stderr,
273
+ "result": None,
274
+ "plots": [],
275
+ "dataframes": []
276
+ }
277
+
278
+ run_proc = subprocess.run(
279
+ [binary_path],
280
+ capture_output=True, text=True, timeout=self.max_execution_time
281
+ )
282
+ return {
283
+ "execution_id": execution_id,
284
+ "status": "success" if run_proc.returncode == 0 else "error",
285
+ "stdout": run_proc.stdout,
286
+ "stderr": run_proc.stderr,
287
+ "result": None,
288
+ "plots": [],
289
+ "dataframes": []
290
+ }
291
+ except Exception as e:
292
+ return {
293
+ "execution_id": execution_id,
294
+ "status": "error",
295
+ "stdout": "",
296
+ "stderr": str(e),
297
+ "result": None,
298
+ "plots": [],
299
+ "dataframes": []
300
+ }
301
+
302
+ def _execute_java(self, code: str, execution_id: str) -> dict:
303
+ temp_dir = tempfile.mkdtemp()
304
+ source_path = os.path.join(temp_dir, "Main.java")
305
+
306
+ try:
307
+ with open(source_path, "w") as f:
308
+ f.write(code)
309
+
310
+ compile_proc = subprocess.run(
311
+ ["javac", source_path],
312
+ capture_output=True, text=True, timeout=self.max_execution_time
313
+ )
314
+ if compile_proc.returncode != 0:
315
+ return {
316
+ "execution_id": execution_id,
317
+ "status": "error",
318
+ "stdout": compile_proc.stdout,
319
+ "stderr": compile_proc.stderr,
320
+ "result": None,
321
+ "plots": [],
322
+ "dataframes": []
323
+ }
324
+
325
+ run_proc = subprocess.run(
326
+ ["java", "-cp", temp_dir, "Main"],
327
+ capture_output=True, text=True, timeout=self.max_execution_time
328
+ )
329
+ return {
330
+ "execution_id": execution_id,
331
+ "status": "success" if run_proc.returncode == 0 else "error",
332
+ "stdout": run_proc.stdout,
333
+ "stderr": run_proc.stderr,
334
+ "result": None,
335
+ "plots": [],
336
+ "dataframes": []
337
+ }
338
+ except Exception as e:
339
+ return {
340
+ "execution_id": execution_id,
341
+ "status": "error",
342
+ "stdout": "",
343
+ "stderr": str(e),
344
+ "result": None,
345
+ "plots": [],
346
+ "dataframes": []
347
+ }