File size: 6,296 Bytes
a495d22
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3b49d6
a495d22
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3b49d6
a495d22
 
 
 
c3b49d6
a495d22
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3b49d6
a495d22
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c3b49d6
a495d22
 
 
 
 
 
 
c3b49d6
a495d22
 
 
 
c3b49d6
a495d22
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
"""
Shared namespace-exec core (ADR-0007 Phase 1).

`PythonExecutor` (in-process, the default) and the sandbox exec-kernel
(`sandbox/kernel.py`, runs code inside an isolated per-session container) both
need the *same* semantics: a persistent Jupyter-kernel-style namespace, seeded
with pandas/numpy + file-writing helpers, into which code is `exec`'d with stdout
captured and capped. Rather than duplicate that logic across process boundaries,
it lives here once as pure functions + a tiny `NamespaceKernel` holder.

`PythonExecutor` is a thin in-process wrapper over this; the sandbox kernel wraps
the SAME core behind an HTTP surface. Keeping one implementation means the two
executors can never drift in their exec/capture/truncation behavior.

Nothing here imports Docker, requests, or MCP — it is the innermost, dependency
-light layer (pandas/numpy are imported lazily inside `seed_namespace`, exactly
as `PythonExecutor.reset_environment` did, so importing this module never fails
when those heavy deps are absent).
"""

from __future__ import annotations

import io
import sys
from typing import Any

__all__ = [
    "MAX_OUTPUT_CHARS",
    "fresh_namespace",
    "seed_namespace",
    "exec_capture",
    "NamespaceKernel",
]

# Cap on stdout captured per execution. Large prints (e.g. full dataframes) get
# re-sent to the LLM on every subsequent step, so an uncapped print can multiply
# token cost across the rest of the run. (Moved verbatim from PythonExecutor.)
MAX_OUTPUT_CHARS = 4000


def fresh_namespace() -> dict[str, Any]:
    """Return a bare namespace dict with builtins available."""
    return {"__builtins__": __builtins__}


def seed_namespace(namespace: dict[str, Any]) -> None:
    """Seed a namespace with the standard libs + file-writing helpers.

    Mirrors `PythonExecutor.reset_environment` exactly. pandas/numpy are imported
    LAZILY so a namespace can still be seeded (minus those names) where the heavy
    deps are absent — the same tolerant posture the original had.
    """
    try:
        import os
        from pathlib import Path

        import numpy as np
        import pandas as pd

        namespace.update({"pd": pd, "np": np, "os": os, "Path": Path})

        def write_text_file(filepath, content):
            """Write text content to a file."""
            with open(filepath, "w") as f:
                f.write(content)
            return f"File written to {filepath}"

        def write_dataframe_to_csv(df, filepath, **kwargs):
            """Write pandas DataFrame to CSV file."""
            df.to_csv(filepath, **kwargs)
            return f"DataFrame written to {filepath}"

        def write_dataframe_to_tsv(df, filepath, **kwargs):
            """Write pandas DataFrame to TSV file."""
            df.to_csv(filepath, sep="\t", **kwargs)
            return f"DataFrame written to {filepath}"

        def create_directory(dirpath):
            """Create directory if it doesn't exist."""
            Path(dirpath).mkdir(parents=True, exist_ok=True)
            return f"Directory created: {dirpath}"

        namespace.update(
            {
                "write_text_file": write_text_file,
                "write_dataframe_to_csv": write_dataframe_to_csv,
                "write_dataframe_to_tsv": write_dataframe_to_tsv,
                "create_directory": create_directory,
            }
        )
    except ImportError as e:
        print(f"Warning: Could not import library: {e}")


def exec_capture(code: str, namespace: dict[str, Any]) -> str:
    """`exec` `code` into `namespace`, returning captured (and capped) stdout.

    Same contract as `PythonExecutor.execute`: empty output → the
    "Code executed successfully" sentinel; oversized output truncated with a
    self-describing tail; any exception → an "Error: …" string (never raised).
    """
    try:
        old_stdout = sys.stdout
        sys.stdout = captured_output = io.StringIO()
        try:
            # Executing LLM-generated code is this agent's core purpose. In-process
            # this runs with the container's privileges (the accepted prototype
            # risk); the ADR-0007 sandbox reuses THIS same function inside an
            # isolated per-session container so the untrusted code is confined.
            exec(code, namespace)  # nosec B102
            result = captured_output.getvalue().strip()
            if not result:
                return "Code executed successfully"
            if len(result) > MAX_OUTPUT_CHARS:
                omitted = len(result) - MAX_OUTPUT_CHARS
                result = (
                    result[:MAX_OUTPUT_CHARS]
                    + f"\n...[output truncated, {omitted} more characters omitted. "
                    f"Print a smaller summary (e.g. .head(), .describe(), or specific columns) "
                    f"if you need to inspect this further.]"
                )
            return result
        finally:
            sys.stdout = old_stdout
    except Exception as e:
        return f"Error: {str(e)}"


class NamespaceKernel:
    """A persistent, seeded namespace with `exec`-and-capture semantics.

    The shared engine behind BOTH executors: `PythonExecutor` delegates to an
    instance of this in-process; the sandbox `kernel.py` HTTP server holds one
    per session and drives it over HTTP. State (variables, injected functions)
    persists across `execute()` calls, matching the Jupyter-kernel model the
    agent relies on.
    """

    def __init__(self) -> None:
        self.namespace: dict[str, Any] = {}
        self.reset()

    def reset(self) -> None:
        """Reset to a clean, freshly-seeded namespace."""
        self.namespace = fresh_namespace()
        seed_namespace(self.namespace)

    def send_functions(self, functions: dict[str, Any]) -> None:
        """Inject callables (or any names) into the namespace."""
        if functions:
            self.namespace.update(functions)

    def send_variables(self, variables: dict[str, Any]) -> None:
        """Inject variables into the namespace."""
        if variables:
            self.namespace.update(variables)

    def execute(self, code: str) -> str:
        """Execute `code` in the persistent namespace; return captured stdout."""
        return exec_capture(code, self.namespace)