kamidso-back/kamidso/engine/app/openapi/api_v1/tools/code_interpreter.py

"""code_interpreter tool — execute Python code in a subprocess.

Execution boundary: file locations (``FILE_PATH``/``FILES_JSON``/``PLOT_DIR``)
travel only through the subprocess environment; the generated Python source
never interpolates a path literal, so adversarial display names or paths
cannot inject code.
"""

import asyncio
import json
import logging
import os
import shutil
import sys
import uuid
from typing import Any, Dict, List, Optional, Tuple

from kamidso.engine.core.agent.resource.tool.base import tool
from kamidso.engine.app.openapi.api_v1.tools.file_access import FILE_ACCESS_GUARD

logger = logging.getLogger(__name__)

EXECUTION_TIMEOUT_SECONDS = 60


async def _run_python_file(
    script_path: str,
    *,
    cwd: str,
    env: Optional[Dict[str, str]] = None,
    timeout: int = EXECUTION_TIMEOUT_SECONDS,
) -> Tuple[Optional[int], bytes, bytes]:
    """Run one Python script via ``asyncio.create_subprocess_exec``.

    The process is spawned with an argument list (never a shell string) and
    inherits the parent environment; file locations reach the child only
    through ``env``. Returns ``(returncode, stdout, stderr)``; ``returncode``
    is ``None`` when the run timed out.
    """
    proc = await asyncio.create_subprocess_exec(
        sys.executable,
        script_path,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        cwd=cwd,
        env=env,
    )
    try:
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
    except asyncio.TimeoutError:
        try:
            proc.kill()
            await proc.wait()
        except Exception:
            pass
        return None, b"", b""
    return proc.returncode, stdout, stderr


def build_execution_env(
    *,
    work_dir: str,
    file_path: Optional[str] = None,
    files_json_path: Optional[str] = None,
    extra: Optional[Dict[str, str]] = None,
    sandbox: bool = False,
) -> Dict[str, str]:
    """Build the subprocess environment for file-aware code execution."""
    if sandbox:
        from kamidso.database_runtime.host import sandboxed_env

        return sandboxed_env(
            work_dir=work_dir,
            file_path=file_path,
            files_json_path=files_json_path,
            extra=extra,
        )
    env = dict(os.environ)
    env["PLOT_DIR"] = work_dir
    if file_path:
        env["FILE_PATH"] = file_path
    if files_json_path:
        env["FILES_JSON"] = files_json_path
    for key, value in (extra or {}).items():
        if value is not None:
            env[key] = value
    return env


def _try_repair_truncated_code(raw_code: str) -> Optional[str]:
    """Attempt to fix code that was truncated by the LLM's token limit."""
    lines = raw_code.split("\n")
    for trim in range(1, min(11, len(lines))):
        candidate_lines = lines[: len(lines) - trim]
        if not candidate_lines:
            continue
        candidate = "\n".join(candidate_lines)
        open_chars = {"(": ")", "[": "]", "{": "}"}
        close_chars = set(open_chars.values())
        stack: list = []
        for ch in candidate:
            if ch in open_chars:
                stack.append(open_chars[ch])
            elif ch in close_chars:
                if stack and stack[-1] == ch:
                    stack.pop()
        if stack:
            candidate += "\n" + "".join(reversed(stack))
        try:
            compile(candidate, "<repair>", "exec")
            return candidate
        except SyntaxError:
            continue
    return None


def make_code_interpreter(
    react_state: Dict[str, Any],
    *,
    sql_enabled: bool = True,
    database_runtime: bool = False,
    native_connector: bool = False,
):
    if database_runtime:
        description = (
            "Execute Python code for data analysis and computation. "
            "Supports pandas, numpy, matplotlib, openpyxl, json, os, etc. "
            "When a database is attached, query it with `database()` "
            "(already imported) or "
            "`from kamidso.database_runtime.sdk import database`. "
            "Inspect with `db = database(); print(db.describe())`. "
            "SQL: `df = db.query(\"SELECT …\")`. "
            "MongoDB is not SQL — never db.query(\"SELECT …\"). Use "
            "`df = db.sample(\"collection\", n=50)`, `print(db.count(\"collection\"))`, "
            "and `df = db.aggregate([{\"$group\": {\"_id\": \"$field\", \"n\": {\"$sum\": 1}}}], collection=\"collection\")`. "
            "$group aliases (n, count) are not collection fields; $sort by them is allowed. "
            "Use sample() for example rows and count()/$group for totals. "
            "Neo4j: `db.cypher(\"MATCH … RETURN count(n)\")`. "
            "JOIN / GROUP BY / COUNT / $lookup / $group MUST run in that query "
            "on the server. Never SELECT * just to count. Never open sqlite3, "
            "pymongo, sqlalchemy, or another database file/host. "
            "If several databases are attached, pass source= the connection name. "
            "If connecting to the database fails, retry at most once. "
            "If it still cannot connect, stop and tell the user to check the connection. "
            "If the user asked for Excel/CSV, write under PLOT_DIR "
            '(e.g. df.to_excel(os.path.join(PLOT_DIR, "best_customers.xlsx"), index=False)). '
            'Parameters: {"code": "python code string"}'
        )
    elif sql_enabled:
        description = (
            "Execute Python code for data analysis and computation. "
            "Supports pandas, numpy, matplotlib, openpyxl, json, os, etc. "
            "When a database is attached: JOIN, GROUP BY, COUNT, $lookup and "
            "$group MUST run in sql_query on the server. Python only formats the "
            "small aggregated table (Jalali labels, ratios, charts, Excel). "
            "Never load thousands of raw documents to group or join in pandas. "
            "For exploration, rely on sql_query observations. If the user asked "
            "for Excel/CSV, read PLOT_DIR/query_results/latest.json and save "
            "under PLOT_DIR "
            '(e.g. df.to_excel(os.path.join(PLOT_DIR, "best_customers.xlsx"), index=False)). '
            'Parameters: {{"code": "python code string"}}'
        )
    else:
        description = (
            "Execute Python code for data analysis and computation. "
            "Supports pandas, numpy, matplotlib, openpyxl, json, os, etc. "
            "This turn's data is an attached Excel/CSV file, not a database. "
            "Use pandas on FILE_PATH / FILES_JSON / User Uploaded Files paths. "
            "Do not call sql_query. Do not read PLOT_DIR/query_results/latest.json "
            "unless you wrote that file yourself. "
            'Parameters: {"code": "python code string"}'
        )
    if native_connector:
        description += (
            " Business-app data is fetched with connector_call, not with Python HTTP clients. "
            "When connector_call saves a result under PLOT_DIR/connector_results/*.json, "
            "load that JSON with Python/pandas for calculations, joins, charts or exports. "
            "Do not use requests/httpx to call attached business apps."
        )

    @tool(description=description)
    async def code_interpreter(code: str) -> str:
        """Execute arbitrary Python code and return stdout/stderr.

        CRITICAL: Each call is completely independent — variables do NOT
        persist between calls. Every code snippet MUST include all necessary
        data loading and processing. Always print() results you want to see.
        """
        from kamidso.engine.core.configs.model_config import PILOT_PATH, STATIC_MESSAGE_IMG_PATH
        from kamidso.engine.app.openapi.api_v1.tools.generated_exports import (
            file_chunks_and_summary,
            publish_new_exports,
            snapshot_export_paths,
        )

        if not code or not code.strip():
            return json.dumps(
                {"chunks": [{"output_type": "text", "content": "No code provided"}]},
                ensure_ascii=False,
            )

        cid = react_state.get("conv_id") or "default"
        work_dir = os.path.join(PILOT_PATH, "tmp", cid)
        os.makedirs(work_dir, exist_ok=True)

        IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp"}
        pre_existing_images: set = set()
        for root, _dirs, files in os.walk(work_dir):
            for f in files:
                ext = os.path.splitext(f)[1].lower()
                if ext in IMAGE_EXTS:
                    pre_existing_images.add(os.path.join(root, f))
        pre_existing_exports = snapshot_export_paths(work_dir)

        preamble_lines = [
            "import json",
            "import os",
            "import pandas as pd",
            "import numpy as np",
            'PLOT_DIR = os.environ["PLOT_DIR"]',
            "os.makedirs(PLOT_DIR, exist_ok=True)",
            'FILE_PATH = os.environ.get("FILE_PATH") or None',
            'FILES_JSON = os.environ.get("FILES_JSON") or None',
            FILE_ACCESS_GUARD,
        ]
        if database_runtime:
            from kamidso.database_runtime.host import DIRECT_DRIVER_GUARD

            preamble_lines.append(DIRECT_DRIVER_GUARD)
            preamble_lines.append(
                "from kamidso.database_runtime.sdk import database"
            )
        preamble = "\n".join(preamble_lines) + "\n"
        full_code = preamble + code

        try:
            compile(full_code, "<code_interpreter>", "exec")
        except SyntaxError as se:
            repaired = _try_repair_truncated_code(full_code)
            if repaired is not None:
                logger.warning(
                    "code_interpreter: auto-repaired truncated code "
                    "(original SyntaxError: %s line %s)",
                    se.msg,
                    se.lineno,
                )
                full_code = repaired
                code = full_code[len(preamble) :]
            else:
                error_msg = (
                    f"SyntaxError before execution: {se.msg} "
                    f"(line {se.lineno})\n"
                    "Please regenerate complete, syntactically valid Python code. "
                    "Keep code under 80 lines."
                )
                return json.dumps(
                    {
                        "chunks": [
                            {"output_type": "code", "content": code.strip()},
                            {"output_type": "text", "content": error_msg},
                        ]
                    },
                    ensure_ascii=False,
                )

        output_text = ""
        run_timeout = EXECUTION_TIMEOUT_SECONDS
        if database_runtime:
            from kamidso.database_runtime.flags import runtime_call_timeout_seconds

            run_timeout = max(
                EXECUTION_TIMEOUT_SECONDS, runtime_call_timeout_seconds() + 15
            )
        try:
            tmp_path = os.path.join(work_dir, "_run.py")
            with open(tmp_path, "w", encoding="utf-8") as tmp:
                tmp.write(full_code)

            extra: Dict[str, str] = {}
            from contextlib import nullcontext

            if database_runtime:
                from kamidso.database_runtime.host import runtime_extra, runtime_session

                session_cm = runtime_session(react_state)
            else:
                session_cm = nullcontext()

            with session_cm as server:
                if database_runtime:
                    extra = runtime_extra(server)
                returncode, stdout, stderr = await _run_python_file(
                    tmp_path,
                    cwd=work_dir,
                    env=build_execution_env(
                        work_dir=work_dir,
                        file_path=react_state.get("file_path"),
                        files_json_path=react_state.get("files_json_path"),
                        extra=extra or None,
                        sandbox=database_runtime,
                    ),
                    timeout=run_timeout,
                )
            output_text = stdout.decode("utf-8", errors="replace")
            error_text = stderr.decode("utf-8", errors="replace")

            if returncode is None:
                # ``_run_python_file`` returns ``None`` only on timeout.
                output_text = f"Execution timed out ({run_timeout}s limit)"
            elif returncode and error_text:
                output_text = (
                    output_text + "\n[ERROR]\n" + error_text
                    if output_text
                    else error_text
                )
        except Exception as e:
            output_text = f"Execution error: {e}"

        chunks: List[Dict[str, Any]] = [
            {"output_type": "code", "content": code.strip()},
        ]
        if output_text.strip():
            clean_output = output_text.strip()
            # Raised from 2000 to 50_000: data-analysis tasks frequently need
            # more than 2000 chars of output (e.g. DataFrame summaries). Larger
            # outputs are persisted to disk by the ToolResultStorage layer
            # rather than truncated here, so the agent can read_file to recover
            # the full content.
            max_out_len = 50_000
            if len(clean_output) > max_out_len:
                truncation_notice = (
                    f"\n\n... [Output truncated, length: {len(clean_output)} chars."
                    f" Only showing first {max_out_len} chars.]"
                )
                clean_output = clean_output[:max_out_len] + truncation_notice
            chunks.append({"output_type": "text", "content": clean_output})
        else:
            chunks.append(
                {
                    "output_type": "text",
                    "content": "(no output — add print() to see results)",
                }
            )

        # Scan for new images generated by this run
        try:
            os.makedirs(STATIC_MESSAGE_IMG_PATH, exist_ok=True)
            for root, _dirs, files in os.walk(work_dir):
                for fname in files:
                    ext = os.path.splitext(fname)[1].lower()
                    full_path = os.path.join(root, fname)
                    if ext in IMAGE_EXTS and full_path not in pre_existing_images:
                        unique_name = f"{uuid.uuid4().hex[:8]}_{fname}"
                        dest = os.path.join(STATIC_MESSAGE_IMG_PATH, unique_name)
                        shutil.copy2(full_path, dest)
                        img_url = f"/images/{unique_name}"
                        chunks.append({"output_type": "image", "content": img_url})
                        react_state.setdefault("generated_images", []).append(img_url)
        except Exception:
            pass

        # Scan for new Excel/CSV/export files and expose download URLs
        try:
            published = publish_new_exports(
                work_dir, pre_existing=pre_existing_exports
            )
            file_chunks, _file_summary = file_chunks_and_summary(published)
            for item in published:
                react_state.setdefault("generated_files", []).append(item)
            chunks.extend(file_chunks)
        except Exception:
            pass

        # Clean up temp script
        try:
            script_path = os.path.join(work_dir, "_run.py")
            if os.path.exists(script_path):
                os.remove(script_path)
        except Exception:
            pass

        all_images = react_state.get("generated_images", [])
        if all_images:
            img_summary = "Generated image URLs (use these in HTML):\n" + "\n".join(
                f"  - {url}" for url in all_images
            )
            chunks.append({"output_type": "text", "content": img_summary})

        return json.dumps({"chunks": chunks}, ensure_ascii=False)

    return code_interpreter