mito-ai-python-tool-executor 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,15 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ from mito_ai_python_tool_executor.blacklisted_words import BlacklistResult, check_for_blacklisted_words
5
+ from mito_ai_python_tool_executor.executor import AskUserMode, PythonToolExecutor
6
+ from mito_ai_python_tool_executor.notebook_io import cells_to_notebook, save_notebook
7
+
8
+ __all__ = [
9
+ "BlacklistResult",
10
+ "AskUserMode",
11
+ "PythonToolExecutor",
12
+ "check_for_blacklisted_words",
13
+ "cells_to_notebook",
14
+ "save_notebook",
15
+ ]
@@ -0,0 +1,89 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ """Detect dangerous patterns in code before execution (parity with mito-ai blacklistedWords)."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import re
9
+ from typing import List, NamedTuple, Optional, Pattern, Tuple
10
+
11
+ __all__ = ["BlacklistResult", "check_for_blacklisted_words"]
12
+
13
+
14
+ class BlacklistResult(NamedTuple):
15
+ safe: bool
16
+ reason: Optional[str] = None
17
+
18
+
19
+ # Keep in sync with mito-ai/src/utils/blacklistedWords.tsx
20
+ _BLACKLISTED_PATTERNS: List[Tuple[Pattern[str], str]] = [
21
+ (
22
+ re.compile(r"\brm\s+-rf\b"),
23
+ "This code contains a command (rm -rf) that could recursively delete files and directories from your system",
24
+ ),
25
+ (
26
+ re.compile(r"\bfs\.rmdir\b"),
27
+ "This code contains a Node.js command (fs.rmdir) that could delete directories from your system",
28
+ ),
29
+ (
30
+ re.compile(r"\bfs\.unlink\b"),
31
+ "This code contains a Node.js command (fs.unlink) that could delete files from your system",
32
+ ),
33
+ (
34
+ re.compile(r"\bshutil\.rmtree\b"),
35
+ "This code contains a Python command (shutil.rmtree) that could recursively delete directories and their contents",
36
+ ),
37
+ (
38
+ re.compile(r"\bos\.remove\b"),
39
+ "This code contains a Python command (os.remove) that could delete files from your system",
40
+ ),
41
+ (
42
+ re.compile(r"\bos\.rmdir\b"),
43
+ "This code contains a Python command (os.rmdir) that could delete directories from your system",
44
+ ),
45
+ (
46
+ re.compile(r"\bos\.unlink\b"),
47
+ "This code contains a Python command (os.unlink) that could delete files from your system",
48
+ ),
49
+ (
50
+ re.compile(r"\brmdir\b"),
51
+ "This code contains a command (rmdir) that could delete directories from your system",
52
+ ),
53
+ (
54
+ re.compile(r"\bunlink\b"),
55
+ "This code contains a function (unlink) that could delete files from your system",
56
+ ),
57
+ (
58
+ re.compile(r"\bdelete\s+from\b", re.IGNORECASE),
59
+ "This code contains an SQL DELETE command that could remove data from your database",
60
+ ),
61
+ (
62
+ re.compile(r"\bdrop\s+table\b", re.IGNORECASE),
63
+ "This code contains an SQL DROP TABLE command that could delete entire tables from your database",
64
+ ),
65
+ (
66
+ re.compile(r"\bdrop\s+database\b", re.IGNORECASE),
67
+ "This code contains an SQL DROP DATABASE command that could delete your entire database",
68
+ ),
69
+ (
70
+ re.compile(r"\bsystem\s*\("),
71
+ "This code contains a system() call that could execute arbitrary system commands, which is a security risk",
72
+ ),
73
+ (
74
+ re.compile(r"\beval\s*\("),
75
+ "This code contains an eval() function that could execute arbitrary code, which is a security risk",
76
+ ),
77
+ (
78
+ re.compile(r"\bexec\s*\("),
79
+ "This code contains an exec() function that could execute arbitrary code, which is a security risk",
80
+ ),
81
+ ]
82
+
83
+
84
+ def check_for_blacklisted_words(code: str) -> BlacklistResult:
85
+ """Return ``safe=False`` and a *reason* if *code* matches a blocked pattern."""
86
+ for pattern, message in _BLACKLISTED_PATTERNS:
87
+ if pattern.search(code):
88
+ return BlacklistResult(safe=False, reason=message)
89
+ return BlacklistResult(safe=True, reason=None)
@@ -0,0 +1,447 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ """In-process :class:`ToolExecutor` backed by an ipykernel session."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import asyncio
9
+ import logging
10
+ import uuid
11
+ from typing import Awaitable, Callable, List, Literal, Optional, Tuple
12
+
13
+ from mito_ai_core.agent.types import AgentContext, ToolResult
14
+ from mito_ai_core.completions.models import AIOptimizedCell, CellUpdate
15
+
16
+ from mito_ai_python_tool_executor.blacklisted_words import check_for_blacklisted_words
17
+ from mito_ai_python_tool_executor.kernel_session import KernelSession
18
+
19
+ AskUserMode = Literal["cli", "mcp_elicitation", "mcp_plaintext"]
20
+ AskUserHandler = Callable[[str, Optional[List[str]]], Awaitable[Optional[str]]]
21
+ logger = logging.getLogger(__name__)
22
+
23
+ ASK_USER_QUESTION_DISABLED_MESSAGE = (
24
+ "The ask_user_question tool is disabled in this environment. "
25
+ "Please use your best judgement to assume the user's response and continue working."
26
+ )
27
+
28
+ STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE = (
29
+ "This Streamlit and app functionality is disabled in this environment. "
30
+ "Please use your best judgement on how to proceed. Either continue working in the notebook "
31
+ "or tell the user that this functionality is disabled."
32
+ )
33
+
34
+
35
+ def _blacklist_error(code: str) -> Optional[str]:
36
+ result = check_for_blacklisted_words(code)
37
+ return result.reason if not result.safe else None
38
+
39
+
40
+ def _default_cell_type(cell_update: CellUpdate) -> str:
41
+ return cell_update.cell_type or "code"
42
+
43
+
44
+ def _find_cell_index(cells: List[AIOptimizedCell], cell_id: str) -> Optional[int]:
45
+ for i, c in enumerate(cells):
46
+ if c.id == cell_id:
47
+ return i
48
+ return None
49
+
50
+
51
+ class PythonToolExecutor:
52
+ """Execute agent tools against a single long-lived ipykernel.
53
+
54
+ Mutates :class:`AgentContext` in place (``cells``, ``active_cell_id``,
55
+ ``variables``) and returns :class:`ToolResult` snapshots for the agent loop.
56
+ """
57
+
58
+ def __init__(
59
+ self,
60
+ *,
61
+ ask_user_mode: AskUserMode = "cli",
62
+ ask_user_handler: Optional[AskUserHandler] = None,
63
+ kernel_cwd: str | None = None,
64
+ ) -> None:
65
+ self._session: Optional[KernelSession] = None
66
+ self._last_cell_text: dict[str, str] = {}
67
+ self._ask_user_mode: AskUserMode = ask_user_mode
68
+ self._ask_user_handler = ask_user_handler
69
+ self._kernel_cwd = kernel_cwd
70
+
71
+ def _ensure_session(self) -> KernelSession:
72
+ if self._session is None:
73
+ self._session = KernelSession(cwd=self._kernel_cwd)
74
+ return self._session
75
+
76
+ def shutdown(self) -> None:
77
+ """Stop the kernel if it was started."""
78
+ if self._session is not None:
79
+ self._session.shutdown()
80
+ self._session = None
81
+
82
+ async def execute_cell_update(
83
+ self,
84
+ ctx: AgentContext,
85
+ cell_update: CellUpdate,
86
+ message: str,
87
+ ) -> ToolResult:
88
+ def _sync() -> ToolResult:
89
+ session = self._ensure_session()
90
+ cells = list(ctx.cells)
91
+ new_cells, active_id, err = self._apply_cell_update(cells, cell_update)
92
+ if err:
93
+ return ToolResult(
94
+ success=False,
95
+ tool_name="cell_update",
96
+ error_message=err,
97
+ )
98
+ ctx.cells = new_cells
99
+ ctx.active_cell_id = active_id
100
+
101
+ idx = _find_cell_index(new_cells, active_id)
102
+ if idx is None:
103
+ vars_ = session.fetch_variables()
104
+ ctx.variables = vars_
105
+ return ToolResult(
106
+ success=True,
107
+ tool_name="cell_update",
108
+ cells=new_cells,
109
+ variables=vars_,
110
+ )
111
+
112
+ cell = new_cells[idx]
113
+ if cell.cell_type != "code":
114
+ vars_ = session.fetch_variables()
115
+ ctx.variables = vars_
116
+ return ToolResult(
117
+ success=True,
118
+ tool_name="cell_update",
119
+ cells=new_cells,
120
+ variables=vars_,
121
+ )
122
+
123
+ blocked = _blacklist_error(cell.code)
124
+ if blocked:
125
+ vars_ = session.fetch_variables()
126
+ ctx.variables = vars_
127
+ return ToolResult(
128
+ success=False,
129
+ tool_name="cell_update",
130
+ error_message=blocked,
131
+ cells=new_cells,
132
+ variables=vars_,
133
+ )
134
+
135
+ ok, out, exec_err = session.execute(cell.code)
136
+ self._last_cell_text[cell.id] = out if out else (exec_err or "")
137
+
138
+ vars_ = session.fetch_variables()
139
+ ctx.variables = vars_
140
+ if not ok:
141
+ return ToolResult(
142
+ success=False,
143
+ tool_name="cell_update",
144
+ error_message=exec_err or "Cell execution failed",
145
+ cells=new_cells,
146
+ variables=vars_,
147
+ )
148
+ return ToolResult(
149
+ success=True,
150
+ tool_name="cell_update",
151
+ cells=new_cells,
152
+ variables=vars_,
153
+ )
154
+
155
+ return await asyncio.to_thread(_sync)
156
+
157
+ def _apply_cell_update(
158
+ self,
159
+ cells: List[AIOptimizedCell],
160
+ cell_update: CellUpdate,
161
+ ) -> Tuple[List[AIOptimizedCell], str, Optional[str]]:
162
+ if cell_update.type == "modification":
163
+ if not cell_update.id:
164
+ return cells, "", "cell_update modification requires id"
165
+ idx = _find_cell_index(cells, cell_update.id)
166
+ if idx is None:
167
+ return cells, "", f"Cell id {cell_update.id!r} not found"
168
+ ct = _default_cell_type(cell_update)
169
+ updated = AIOptimizedCell(cell_type=ct, id=cell_update.id, code=cell_update.code)
170
+ new_cells = cells.copy()
171
+ new_cells[idx] = updated
172
+ return new_cells, cell_update.id, None
173
+
174
+ # new cell
175
+ after = cell_update.after_cell_id
176
+ if after is None:
177
+ return cells, "", "cell_update new requires after_cell_id"
178
+
179
+ ct = _default_cell_type(cell_update)
180
+ new_id = str(uuid.uuid4())
181
+ new_cell = AIOptimizedCell(cell_type=ct, id=new_id, code=cell_update.code)
182
+ new_cells = cells.copy()
183
+
184
+ if after == "new cell":
185
+ new_cells.insert(0, new_cell)
186
+ return new_cells, new_id, None
187
+
188
+ idx = _find_cell_index(new_cells, after)
189
+ if idx is None:
190
+ return cells, "", f"after_cell_id {after!r} not found"
191
+ new_cells.insert(idx + 1, new_cell)
192
+ return new_cells, new_id, None
193
+
194
+ async def run_all_cells(
195
+ self,
196
+ ctx: AgentContext,
197
+ message: str,
198
+ ) -> ToolResult:
199
+ def _sync() -> ToolResult:
200
+ session = self._ensure_session()
201
+ cells = list(ctx.cells)
202
+ for cell in cells:
203
+ if cell.cell_type != "code":
204
+ continue
205
+ blocked = _blacklist_error(cell.code)
206
+ if blocked:
207
+ vars_ = session.fetch_variables()
208
+ ctx.variables = vars_
209
+ ctx.cells = cells
210
+ return ToolResult(
211
+ success=False,
212
+ tool_name="run_all_cells",
213
+ error_message=blocked,
214
+ cells=cells,
215
+ variables=vars_,
216
+ )
217
+ ok, out, err = session.execute(cell.code)
218
+ self._last_cell_text[cell.id] = out if out else (err or "")
219
+ if not ok:
220
+ vars_ = session.fetch_variables()
221
+ ctx.variables = vars_
222
+ ctx.cells = cells
223
+ return ToolResult(
224
+ success=False,
225
+ tool_name="run_all_cells",
226
+ error_message=err or "Run-all failed",
227
+ cells=cells,
228
+ variables=vars_,
229
+ )
230
+ vars_ = session.fetch_variables()
231
+ ctx.variables = vars_
232
+ ctx.cells = cells
233
+ return ToolResult(
234
+ success=True,
235
+ tool_name="run_all_cells",
236
+ cells=cells,
237
+ variables=vars_,
238
+ )
239
+
240
+ return await asyncio.to_thread(_sync)
241
+
242
+ async def get_cell_output(
243
+ self,
244
+ ctx: AgentContext,
245
+ cell_id: str,
246
+ message: str,
247
+ ) -> ToolResult:
248
+ def _sync() -> ToolResult:
249
+ text = self._last_cell_text.get(cell_id)
250
+ if text is None:
251
+ return ToolResult(
252
+ success=True,
253
+ tool_name="get_cell_output",
254
+ output=None,
255
+ error_message=(
256
+ "[CLI] No captured text output for this cell yet. "
257
+ "Run the cell or use run_all_cells / scratchpad first."
258
+ ),
259
+ )
260
+ return ToolResult(
261
+ success=True,
262
+ tool_name="get_cell_output",
263
+ output=None,
264
+ error_message=f"Cell output (plain text):\n{text}",
265
+ )
266
+
267
+ return await asyncio.to_thread(_sync)
268
+
269
+ async def execute_scratchpad(
270
+ self,
271
+ ctx: AgentContext,
272
+ code: str,
273
+ summary: str,
274
+ message: str,
275
+ ) -> ToolResult:
276
+ def _sync() -> ToolResult:
277
+ session = self._ensure_session()
278
+ blocked = _blacklist_error(code)
279
+ if blocked:
280
+ vars_ = session.fetch_variables()
281
+ ctx.variables = vars_
282
+ return ToolResult(
283
+ success=False,
284
+ tool_name="scratchpad",
285
+ error_message=blocked,
286
+ variables=vars_,
287
+ )
288
+ ok, out, err = session.execute(code)
289
+ vars_ = session.fetch_variables()
290
+ ctx.variables = vars_
291
+ if not ok:
292
+ return ToolResult(
293
+ success=False,
294
+ tool_name="scratchpad",
295
+ error_message=err or "Scratchpad execution failed",
296
+ output=out,
297
+ variables=vars_,
298
+ )
299
+ return ToolResult(
300
+ success=True,
301
+ tool_name="scratchpad",
302
+ output=out,
303
+ variables=vars_,
304
+ )
305
+
306
+ return await asyncio.to_thread(_sync)
307
+
308
+ async def ask_user_question(
309
+ self,
310
+ ctx: AgentContext,
311
+ question: str,
312
+ message: str,
313
+ answers: Optional[List[str]] = None,
314
+ ) -> ToolResult:
315
+ logger.info(
316
+ "ask_user_question invoked with mode=%s, answers_provided=%s",
317
+ self._ask_user_mode,
318
+ bool(answers),
319
+ )
320
+
321
+ if self._ask_user_mode == "mcp_elicitation":
322
+ logger.info("Routing ask_user_question via MCP elicitation handler")
323
+ return await self._ask_user_question_via_mcp(ctx, question, answers)
324
+ if self._ask_user_mode == "mcp_plaintext":
325
+ logger.info("Routing ask_user_question via MCP plaintext fallback")
326
+ return await self._ask_user_question_disabled_response(ctx, question, answers)
327
+ if self._ask_user_mode != "cli":
328
+ logger.warning(
329
+ "Unknown ask_user_mode=%s; forcing plaintext fallback",
330
+ self._ask_user_mode,
331
+ )
332
+ return await self._ask_user_question_disabled_response(ctx, question, answers)
333
+
334
+ logger.warning("Routing ask_user_question via CLI stdin/stdout prompt path")
335
+ def _sync() -> ToolResult:
336
+ session = self._ensure_session()
337
+ print(question, flush=True)
338
+ if answers:
339
+ for i, a in enumerate(answers):
340
+ print(f" [{i}] {a}", flush=True)
341
+ raw = input("Enter a number or free text: ").strip()
342
+ if raw.isdigit():
343
+ j = int(raw)
344
+ if 0 <= j < len(answers):
345
+ answer = answers[j]
346
+ else:
347
+ answer = raw
348
+ else:
349
+ answer = raw
350
+ else:
351
+ answer = input("> ").strip()
352
+
353
+ vars_ = session.fetch_variables()
354
+ ctx.variables = vars_
355
+ return ToolResult(
356
+ success=True,
357
+ tool_name="ask_user_question",
358
+ output=answer,
359
+ variables=vars_,
360
+ )
361
+
362
+ return await asyncio.to_thread(_sync)
363
+
364
+ async def _ask_user_question_via_mcp(
365
+ self,
366
+ ctx: AgentContext,
367
+ question: str,
368
+ answers: Optional[List[str]],
369
+ ) -> ToolResult:
370
+ if self._ask_user_handler is None:
371
+ logger.warning(
372
+ "ask_user_question is in `mcp_elicitation` mode but no "
373
+ "elicitation handler is configured; returning disabled "
374
+ "plaintext message."
375
+ )
376
+ return await self._ask_user_question_disabled_response(ctx, question, answers)
377
+
378
+ try:
379
+ answer = await self._ask_user_handler(question, answers)
380
+ except Exception as exc:
381
+ logger.warning(
382
+ "MCP elicitation failed for ask_user_question; returning "
383
+ "disabled plaintext message. error=%s",
384
+ exc,
385
+ )
386
+ return await self._ask_user_question_disabled_response(ctx, question, answers)
387
+
388
+ session = self._ensure_session()
389
+ vars_ = session.fetch_variables()
390
+ ctx.variables = vars_
391
+ return ToolResult(
392
+ success=True,
393
+ tool_name="ask_user_question",
394
+ output=(answer or "").strip(),
395
+ variables=vars_,
396
+ )
397
+
398
+ async def _ask_user_question_disabled_response(
399
+ self,
400
+ ctx: AgentContext,
401
+ question: str,
402
+ answers: Optional[List[str]],
403
+ ) -> ToolResult:
404
+ del question, answers # Plaintext mode intentionally disables user prompting.
405
+ session = self._ensure_session()
406
+ vars_ = session.fetch_variables()
407
+ ctx.variables = vars_
408
+ return ToolResult(
409
+ success=True,
410
+ tool_name="ask_user_question",
411
+ output=ASK_USER_QUESTION_DISABLED_MESSAGE,
412
+ variables=vars_,
413
+ )
414
+
415
+ async def create_streamlit_app(
416
+ self,
417
+ ctx: AgentContext,
418
+ message: str,
419
+ streamlit_app_prompt: Optional[str] = None,
420
+ ) -> ToolResult:
421
+ del message, streamlit_app_prompt
422
+ session = self._ensure_session()
423
+ vars_ = session.fetch_variables()
424
+ ctx.variables = vars_
425
+ return ToolResult(
426
+ success=True,
427
+ tool_name="create_streamlit_app",
428
+ output=STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE,
429
+ variables=vars_,
430
+ )
431
+
432
+ async def edit_streamlit_app(
433
+ self,
434
+ ctx: AgentContext,
435
+ streamlit_app_prompt: str,
436
+ message: str,
437
+ ) -> ToolResult:
438
+ del streamlit_app_prompt, message
439
+ session = self._ensure_session()
440
+ vars_ = session.fetch_variables()
441
+ ctx.variables = vars_
442
+ return ToolResult(
443
+ success=True,
444
+ tool_name="edit_streamlit_app",
445
+ output=STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE,
446
+ variables=vars_,
447
+ )
@@ -0,0 +1,145 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ """One ipykernel session with execute + variable introspection."""
5
+
6
+ from __future__ import annotations
7
+
8
+ import json
9
+ import re
10
+ from queue import Empty
11
+ from typing import List, Optional, Tuple
12
+
13
+ from jupyter_client import KernelManager
14
+
15
+ from mito_ai_core.completions.models import KernelVariable
16
+ from mito_ai_core.kernel_variable_inspection import get_kernel_variable_inspection_script
17
+
18
+ # Same script as JupyterLab VariableInspector / kernelVariableInspectionScript.ts
19
+ # (mito_ai_core/resources/kernel_variable_inspection.txt).
20
+ _VARIABLE_PROBE = get_kernel_variable_inspection_script()
21
+ _ANSI_ESCAPE_PATTERN = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
22
+
23
+
24
+ class KernelSession:
25
+ """Starts a single kernel and exposes blocking execute + shutdown."""
26
+
27
+ def __init__(self, *, cwd: str | None = None) -> None:
28
+ self._km = KernelManager()
29
+ start_kwargs = {"cwd": cwd} if cwd else {}
30
+ self._km.start_kernel(**start_kwargs)
31
+ self._kc = self._km.client()
32
+ self._kc.start_channels()
33
+ self._kc.wait_for_ready(timeout=120)
34
+
35
+ def shutdown(self) -> None:
36
+ try:
37
+ self._kc.stop_channels()
38
+ finally:
39
+ try:
40
+ self._km.shutdown_kernel(now=True)
41
+ except Exception:
42
+ pass
43
+
44
+ def execute(
45
+ self, code: str, *, timeout_s: float = 300.0
46
+ ) -> Tuple[bool, str, Optional[str]]:
47
+ """Run *code* in the kernel.
48
+
49
+ Returns
50
+ -------
51
+ success, combined_text_output, error_message
52
+ *error_message* is ``None`` when execution finished without a traceback.
53
+ """
54
+ msg_id = self._kc.execute(code)
55
+ text_parts: List[str] = []
56
+ error_text: Optional[str] = None
57
+
58
+ while True:
59
+ try:
60
+ msg = self._kc.get_iopub_msg(timeout=timeout_s)
61
+ except Empty:
62
+ return False, "", "Timed out waiting for kernel iopub output"
63
+
64
+ if msg.get("parent_header", {}).get("msg_id") != msg_id:
65
+ continue
66
+
67
+ msg_type = msg["msg_type"]
68
+ content = msg["content"]
69
+
70
+ if msg_type == "stream":
71
+ text_parts.append(_strip_ansi_escape_sequences(content.get("text", "")))
72
+ elif msg_type == "error":
73
+ error_text = _strip_ansi_escape_sequences(
74
+ "\n".join(content.get("traceback", []))
75
+ )
76
+ elif msg_type == "execute_result":
77
+ data = content.get("data", {})
78
+ plain = data.get("text/plain")
79
+ if plain is not None:
80
+ text_parts.append(_strip_ansi_escape_sequences(plain))
81
+ elif msg_type == "display_data":
82
+ data = content.get("data", {})
83
+ plain = data.get("text/plain")
84
+ if plain is not None:
85
+ text_parts.append(_strip_ansi_escape_sequences(plain))
86
+ elif msg_type == "status" and content.get("execution_state") == "idle":
87
+ break
88
+
89
+ # Shell reply may contain error info even when iopub missed it.
90
+ try:
91
+ reply = self._kc.get_shell_msg(timeout=60)
92
+ except Empty:
93
+ reply = None
94
+
95
+ if reply is not None:
96
+ c = reply.get("content", {})
97
+ if c.get("status") == "error":
98
+ tb = c.get("traceback")
99
+ if tb:
100
+ shell_err = _strip_ansi_escape_sequences("\n".join(tb))
101
+ else:
102
+ shell_err = _strip_ansi_escape_sequences(
103
+ c.get("evalue", "Kernel error")
104
+ )
105
+ if not error_text:
106
+ error_text = shell_err
107
+
108
+ out = "".join(text_parts)
109
+ success = error_text is None
110
+ return success, out, error_text
111
+
112
+ def fetch_variables(self, *, timeout_s: float = 120.0) -> List[KernelVariable]:
113
+ """Return a simple name / type / repr view of user globals."""
114
+ ok, out, err = self.execute(_VARIABLE_PROBE, timeout_s=timeout_s)
115
+ if not ok or err:
116
+ return []
117
+ out = out.strip()
118
+ if not out:
119
+ return []
120
+ try:
121
+ raw = json.loads(out)
122
+ except json.JSONDecodeError:
123
+ return []
124
+ result: List[KernelVariable] = []
125
+ for item in raw:
126
+ if not isinstance(item, dict):
127
+ continue
128
+ name = item.get("variable_name")
129
+ typ = item.get("type")
130
+ if not isinstance(name, str) or not isinstance(typ, str):
131
+ continue
132
+ result.append(
133
+ KernelVariable(
134
+ variable_name=name,
135
+ type=typ,
136
+ value=item.get("value"),
137
+ )
138
+ )
139
+ return result
140
+
141
+
142
+ def _strip_ansi_escape_sequences(text: str) -> str:
143
+ if not isinstance(text, str) or not text:
144
+ return text
145
+ return _ANSI_ESCAPE_PATTERN.sub("", text)
@@ -0,0 +1,31 @@
1
+ # Copyright (c) Saga Inc.
2
+ # Distributed under the terms of the GNU Affero General Public License v3.0 License.
3
+
4
+ """Serialize :class:`AIOptimizedCell` lists to Jupyter notebook files."""
5
+
6
+ from __future__ import annotations
7
+
8
+ from typing import List
9
+
10
+ import nbformat
11
+ from nbformat.v4 import new_code_cell, new_markdown_cell, new_notebook
12
+ from nbformat.notebooknode import NotebookNode
13
+
14
+ from mito_ai_core.completions.models import AIOptimizedCell
15
+
16
+
17
+ def cells_to_notebook(cells: List[AIOptimizedCell]) -> NotebookNode:
18
+ """Build an ``nbformat`` v4 notebook from agent cells (no outputs)."""
19
+ nb = new_notebook()
20
+ for cell in cells:
21
+ if cell.cell_type == "markdown":
22
+ nb.cells.append(new_markdown_cell(source=cell.code))
23
+ else:
24
+ nb.cells.append(new_code_cell(source=cell.code))
25
+ return nb
26
+
27
+
28
+ def save_notebook(nb: NotebookNode, path: str) -> None:
29
+ """Write *nb* to *path* as ``.ipynb`` (UTF-8)."""
30
+ with open(path, "w", encoding="utf-8") as f:
31
+ nbformat.write(nb, f)
@@ -0,0 +1,14 @@
1
+ Metadata-Version: 2.4
2
+ Name: mito-ai-python-tool-executor
3
+ Version: 0.1.1
4
+ Summary: In-process ipykernel-backed ToolExecutor for Mito AI (CLI and headless use).
5
+ Author: Saga Inc.
6
+ License-Expression: AGPL-3.0-only
7
+ Requires-Python: >=3.9
8
+ Requires-Dist: ipykernel>=6.0
9
+ Requires-Dist: jupyter-client>=8.0
10
+ Requires-Dist: mito-ai-core
11
+ Requires-Dist: nbformat>=5.0
12
+ Provides-Extra: test
13
+ Requires-Dist: pytest; extra == 'test'
14
+ Requires-Dist: pytest-asyncio; extra == 'test'
@@ -0,0 +1,8 @@
1
+ mito_ai_python_tool_executor/__init__.py,sha256=F99Us94U0QcW_B3gYe2hXOlMrhAKPhWmTOa81r78jIg,548
2
+ mito_ai_python_tool_executor/blacklisted_words.py,sha256=Iuwrl4y0T943vQA40OtZjqpQFkYif-0TSjGk1e1BUz8,3326
3
+ mito_ai_python_tool_executor/executor.py,sha256=3oE3WV80mV_Wx38jG8xLhOQB8z0TW3pb7B2qWj3XfhQ,15643
4
+ mito_ai_python_tool_executor/kernel_session.py,sha256=5NYIjM-7wbhd9o6lqsVvghDLGgROkz7TQlpyeUw0aPs,5072
5
+ mito_ai_python_tool_executor/notebook_io.py,sha256=w4K5oSCmF5PIENdGzmMMWzmmyCkrzztYgmN3nQavBeY,1022
6
+ mito_ai_python_tool_executor-0.1.1.dist-info/METADATA,sha256=fCKH6uoML5hs5sGMYg1Glicd1S1QkNRZX9JYB-gBDnk,462
7
+ mito_ai_python_tool_executor-0.1.1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
8
+ mito_ai_python_tool_executor-0.1.1.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: hatchling 1.29.0
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any