Easy to deploy.A cloud service for python code interpreter sandbox for Code-Interpreter.
import subprocess
from fastapi import HTTPException
from collections import OrderedDict
from typing import List, Optional
from queue import Queue
import models.api as models
from models.api import File
import configs
import os
import errno
import docker
from docker.types import Mount
import base64
import models.schemas as schemas
import zipfile
import re
import socket, time, json, asyncio
from concurrent.futures import ThreadPoolExecutor
from jupyter_api import JupyterAPI

executor = ThreadPoolExecutor()
PROJECT_NETWORK='app_network_for_codebox'

def get_network_gateway_ip():
    client=docker.from_env()
    network=client.networks.get(PROJECT_NETWORK)
    gateway=network.attrs['IPAM']['Config'][0]['Gateway']
    print('network_gateway=', gateway)
    return gateway

PROJECT_NETWORK_GATEWAY=get_network_gateway_ip()



class LRUCacheCodebox:
    def __init__(self, capacity: int):
        self.cache = OrderedDict()
        self.capacity = capacity

    def get(self, key: models.CodeboxSessionKey) -> "CodeboxSession":
        if key not in self.cache:
            return None
        value = self.cache.pop(key)
        # 把 key 放到最右边（也就是最新）
        self.cache[key] = value
        return value
    
    def remove(self, key:models.CodeboxSessionKey):
        if key in self.cache:
            self.cache.pop(key)

    async def put(self, key: models.CodeboxSessionKey, value: "CodeboxSession") -> None:
        if key in self.cache:
            # 如果 key 已经在 cache 中，先移除旧的
            self.cache.pop(key)
        elif len(self.cache) == self.capacity:
            # 如果 cache 已满，移除最左边的元素（也就是最老的）
            key,value=self.cache.popitem(last=False)
            value.stop()
        # 添加新的 key-value 到最右边（也就是最新）
        self.cache[key] = value

    def keys(self)->List[models.CodeboxSessionKey]:
        return self.cache.keys()



class GlobalCodeboxSessionManager(object):
    def __init__(self,):
        self.sessions= LRUCacheCodebox(10) 
    
    async def get_codebox_session(
            self, session_key:models.CodeboxSessionKey, 
            language:Optional[schemas.ProgrammingLanguage]=schemas.ProgrammingLanguage.JUPYTER
        )->"CodeboxSession":
        session=self.sessions.get(session_key)
        if session is None:
            session=CodeboxSession(
                session_id=session_key.session_id,
                language=language,
                )
            await self.sessions.put(key=session_key,value=session)
            return session
        else:
            return session
    
    def remove_session(self,session_key:models.CodeboxSessionKey):
        self.sessions.remove(session_key)

    
global_codebox_session_manager=GlobalCodeboxSessionManager()
PORT_RANGE=range(10000,40000)
PORTS_POOL = {}

def get_one_port():
    for port in PORT_RANGE:
        if port not in PORTS_POOL.keys():
            return port

class CodeboxSession:
    client_socket=None
    def __init__(self,session_id: str, language: schemas.ProgrammingLanguage):
        
        self.session_id = session_id
        self.language = language
        self.language_map=self._init_language_map()
        self.init_codebox_directory(self.session_id)
        self.container= self._create_container(language=self.language)
        self.network_gateway_ip=self._get_network_gateway_ip()
        
        # Start the container
        if self.language in [schemas.ProgrammingLanguage.JUPYTER]:
            self.container.start()
            self.container_ip_address = self._container_ip_address()
            print('_container_ip_address', self.container_ip_address)
            time.sleep(1)
            self.token = 'zhejianzhang'
            self.jupyter_api = JupyterAPI(port=self.export_port,ip=self.network_gateway_ip, token=self.token)

    def _init_language_map(self):
        return {
            schemas.ProgrammingLanguage.JUPYTER:{
                "image_name":"jupyter/scipy-notebook:custom", #You can modify the image
                # "command":'start-notebook.sh --NotebookApp.token="zhejianzhang" --ip="0.0.0.0" --no-browser --allow-root --allow_origin="*" --ServerApp.allow_remote_access=True --certfile=mycert.pem --keyfile=mykey.key',
                "command":'start-notebook.sh --NotebookApp.token="zhejianzhang" --ip="0.0.0.0" --no-browser --allow-root --allow_origin="*" --ServerApp.allow_remote_access=True ',
                "file_extension":".ipynb",
                "export_port":8888
            },
        }
    
    # def _jupyter_token(self):
    #     logs = self.container.logs().decode('utf-8')
    #     print(logs)
    #     token_pattern = r"token=([a-z0-9]+)"
    #     match = re.search(token_pattern, logs)
    #     if not match:
    #         raise ValueError("Token not found in container logs.")
    #     token = match.group(1)    
    #     return token
    

    def _get_network_gateway_ip(self):
        client=docker.from_env()
        network=client.networks.get(PROJECT_NETWORK)
        gateway=network.attrs['IPAM']['Config'][0]['Gateway']
        print('network_gateway=', gateway)
        return gateway


    def _create_container(self,language:schemas.ProgrammingLanguage):
        print("__init__ codeboxsession",self.root_path)
        client = docker.DockerClient(base_url='unix://var/run/docker.sock', timeout=20)
        mounts = [
            Mount(target="/codebox", source=f"{configs.codebox_root_path}{self.session_id}", type="bind"),
            Mount(target="/logs", source=f"{configs.codebox_root_path}{self.session_id}{configs.log_file_path}", type="bind")
        ]
        self.container_export_port = self.language_map[language]['export_port']
        self.export_port=get_one_port()
        PORTS_POOL[self.export_port] = schemas.PortOccupation(
            session_id=self.session_id,
            usage=schemas.PortUsage.EXPORT
        )
        port_mapping={
            '8888/tcp':self.export_port,
        }

        container = client.containers.create(self.language_map[language]['image_name'], 
                                     mounts=mounts,
                                     command=self.language_map[language]['command'],
                                     network_mode=PROJECT_NETWORK,
                                     ports = port_mapping, 
                                     detach=True)
        
        return container
    
    def _container_ip_address(self):
        client = docker.from_env()
        container = client.containers.get(self.container.name)
        return container.attrs['NetworkSettings']['Networks']['app_network_for_codebox']['IPAddress']
    
    def _form_image_name(self):
        return self.language_map[self.language]['image_name'] + self.session_id
        
    async def install_packages(self, packages:List[str])->tuple[list[str], list[str]]:
        code = f"!pip install {' '.join(packages)}"
        return self.jupyter_api.execute_code(code)

    
    async def arun_shell(self, command):
        loop = asyncio.get_event_loop()
        client = docker.from_env()
        exec_instance = await loop.run_in_executor(executor, client.api.exec_create, self.container.id, command, True, True, False)
        result = await loop.run_in_executor(executor, client.api.exec_start, exec_instance['Id'])
        print('exec_shell result:', result.decode('utf-8'))
        return result.decode('utf-8')


    def init_codebox_directory(self,session_id: str):
        self.root_path = f"/codebox/{session_id}"
        self.input_code_path = f"{self.root_path}{configs.input_code_path}"
        # self.upload_file_path= f"{self.root_path}{configs.upload_file_path}"
        # self.output_file_path= f"{self.root_path}{configs.output_file_path}"
        self.log_file_path= f"{self.root_path}{configs.log_file_path}"

        all_dirs = [
            self.root_path,
            self.input_code_path, 
            # self.upload_file_path,
            # self.output_file_path,
            self.log_file_path
        ]

        for item in all_dirs:
            if not os.path.exists(item):
                print(item)
                os.makedirs(item)
    
    def upload(self, file: File)-> tuple[bool, str]:
        file_path = f"{self.root_path}/{file.name}"
        try:
            with open(file_path, "wb") as f:
                f.write(file.content)
            if file.name.lower().endswith('.zip'):
                # 解压ZIP文件
                try:
                    with zipfile.ZipFile(file_path, 'r') as zip_ref:
                        zip_ref.extractall(file_path[:-4])
                    return (True, f'zip file {file.name} is extracted to {file_path}')
                except zipfile.BadZipFile:
                    os.remove(file_path)  # 删除不完整或损坏的ZIP文件
                    # raise HTTPException(status_code=400, detail="Invalid ZIP file")
                    return (False, f"{file.name} seems to be an invalid ZIP file")
        except OSError as e:
            if e.errno == errno.ENOSPC:
                raise HTTPException(status_code=400, detail="Insufficient disk space")
            else:
                raise
        return (True, "")
    
    def upload_batch(self, files: List[File]):
        for file in files:
            success = self.upload(file)
            print(file.name, success)

    
    #download from mnt
    def download(self, filename: str) -> File:
        if os.path.exists(filename):
            raise HTTPException(status_code=401, detail="File already exists")
        try:
            with open(filename) as f:
                file = File(
                    name=filename,
                    content=f.read()
                )
                return file
        except Exception as e:
            return None
        
    #download from jupyter
    def download_from_jupyter(self, filename):
        self.jupyter_api.download_file(filename)

    def before_execute(self,):
        if os.path.exists(f"{self.root_path}execute_lock"):
            os.remove(f"{self.root_path}execute_lock")
        with open(f"{self.root_path}execute_lock", 'w'):
            pass
        


    def execute_code(self, code: str):
        return self.jupyter_api.execute_code(code)




    def stop(self):
        try:
            if self.container:
                self.container.stop()
                self.container.remove()
        except Exception as e:
            print(e)

    
    def __del__(self):
        self.stop()
        return
    
        
        

        
app/codebox_session.py
zhangzhejian/codeinterpreter-codebox
import os
ENV_FLAG = os.environ.get('ENV_FLAG')
APP_ID=os.environ.get('APP_ID')
APP_SECRET= os.environ.get('APP_SECRET')

TCP_SERVICE_PORT=8000

allowed_tokens = [
    'zhejianzhang'
]

# codebox_root_path='/Users/zhejianzhang/PrivateProject/codebox_server_code/'
codebox_root_path = os.environ.get('CODEBOX_ROOT_PATH')
input_code_path = '/input_code/'
upload_file_path='/upload/'
output_file_path='/output_file/'
log_file_path='/logs/'

codebox_error_log_title = 'Codebox Error Log:\n\n'


app/configs.py
zhangzhejian/codeinterpreter-codebox
import requests
import time
import re
import json
from websocket import create_connection
import uuid
import datetime
import traceback
from models.api import File
import base64


def strip_ansi_escape_codes(s):
    ansi_escape = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]')
    return ansi_escape.sub('', s)


class JupyterAPI:

    def __init__(self, port, ip, token):
        self.NOTEBOOK_PATH = "main.ipynb"
        self.ip = ip
        self.token=token
        self.BASE_URL = f"http://{ip}:{port}/api/"
        self.WS_BASE_URL = f"ws://{ip}:{port}/api/"
        self.HEADERS = {
            "Authorization": f"token {token}"
        }
        self.session=self.create_session()
        self.kernel = self.session["kernel"]
        self.kernel_id = self.kernel["id"]
        self.ws = self.create_ws_conn()
    

    def create_session(self):
        url = self.BASE_URL + 'sessions'
        params = '{"path":\"%s\","type":"notebook","name":"","kernel":{"id":null,"name":"python3"}}' % self.NOTEBOOK_PATH
        response = requests.post(url, headers=self.HEADERS , data=params)
        session = json.loads(response.text)
        return session
    
    def create_ws_conn(self):
        ws = create_connection(f"{self.WS_BASE_URL}kernels/{self.kernel_id}/channels?session_id{self.session['id']}", header=self.HEADERS)
        return ws

    # def create_notebook(self):
    #     # 创建一个新的notebook
    #     r = requests.post(self.BASE_URL + "contents", headers=self.HEADERS, json={
    #         "type": "notebook",
    #         "path": "test_notebook.ipynb"
    #     }, timeout=10, verify=False)
    #     print(r.status_code)
    #     print(r.content)
    #     notebook_path = r.json()["path"]

    # def create_kernel(self):
    #     # 创建一个新的kernel
    #     r = requests.post(self.BASE_URL + "kernels", headers=self.HEADERS,timeout=5,verify=False)
    #     kernel_id = r.json()["id"]
    #     return kernel_id
    
    def send_execute_request(self,code, msg_id):
        msg_type = 'execute_request'
        content = { 'code' : code, 'silent':False }
        hdr = { 'msg_id' :msg_id, 
            'username': 'codebox', 
            'session': self.session['id'], 
            'data': datetime.datetime.now().isoformat(),
            'msg_type': msg_type,
            'version' : '5.0' }
        msg = { 'header': hdr, 'parent_header': hdr, 
            'metadata': {},
            'content': content }
        return msg
    
    def download_file(self, filename):
        file_url = f"http://{self.ip}:{self.port}/files/{filename}?token={self.token}"
        response = requests.get(file_url)
        with open(filename, "wb") as f:
            f.write(response.content)

    def execute_code(self, code: str)->tuple[list[str], list[str]]:
        msg_id=  uuid.uuid1().hex
        msg = self.send_execute_request(code, msg_id)
        self.ws.send(json.dumps(msg))
        try:
            img_index=0
            msg_type = ''
            stdoutputs = []
            image_datas = []
            while True:
                rsp = json.loads(self.ws.recv())
                msg_type = rsp["msg_type"]
                # print(msg_type)
                if msg_type == "status" and rsp["content"]["execution_state"] == "idle" and rsp["parent_header"]['msg_id']==msg_id:
                    print('execute end')
                    break
                if msg_type == "stream":
                    stdoutput=rsp["content"]["text"]
                    if isinstance(stdoutput,list):
                        stdoutput=[strip_ansi_escape_codes(item) for item in stdoutput]
                        stdoutputs+=stdoutput
                    else:
                        stdoutputs.append(stdoutput)
                elif msg_type == "execute_result":
                    if "image/png" in (rsp["content"]["data"].keys()):
                        image_data = rsp["content"]["data"]["image/png"]
                        image_datas.append(File(name=f"{msg_id}_{img_index}.png", content=image_data))
                        img_index+=1
                    else:
                        stdoutput=rsp["content"]["data"]["text/plain"]
                        if isinstance(stdoutput,list):
                            stdoutputs+=stdoutput
                        else:
                            stdoutputs.append(stdoutput)
                elif msg_type == "display_data":
                    image_data = rsp["content"]["data"]["image/png"]
                    image_datas.append(File(name=f"{msg_id}_{img_index}.png", content=image_data))
                    img_index+=1
                elif msg_type == "error":
                    error_output=rsp["content"]["traceback"]
                    if isinstance(error_output,list):
                        error_output=[strip_ansi_escape_codes(item) for item in error_output]
                        stdoutputs+= error_output
                    else:
                        cleaned_string = strip_ansi_escape_codes(error_output)
                        stdoutputs.append(cleaned_string)
                else:
                    pass
        except:
                traceback.print_exc()
                return (stdoutputs, image_datas)
        
        return (stdoutputs, image_datas)

    

    def __del__(self):
        if self.ws:
            print('ws closed')
            self.ws.close()
        # 清理资源
        # requests.delete(self.BASE_URL + f"kernels/{kernel_id}", headers=HEADERS)
        # requests.delete(self.BASE_URL + f"contents/{notebook_path}", headers=HEADERS)
        pass
if __name__=='__main__':
    print('start')
    api = JupyterAPI(port=10001, ip='localhost', token='zhejianzhang')
    api.create_notebook()
app/jupyter_api.py
zhangzhejian/codeinterpreter-codebox
import uuid
from fastapi import FastAPI, Depends, HTTPException, UploadFile,File, Form
from pydantic import BaseModel
import docker
from docker.errors import NotFound
from starlette.requests import Request
from docker.types import Mount
import os
from typing import Optional, List
import configs
from codebox_session import global_codebox_session_manager
import models.api as models
import models.schemas as schemas
import asyncio


app = FastAPI()
def validate_token(token:str):
    return token in configs.allowed_tokens


@app.post("/api/v1/jupyter/start")
async def start_codebox(token: str, language:schemas.ProgrammingLanguage) -> str:
    valid = validate_token(token)
    if not valid:
        return None
    session_id = str(uuid.uuid4())
    await global_codebox_session_manager.get_codebox_session(
        session_key=models.CodeboxSessionKey(session_id=session_id, language=language), 
        language=language)
    return session_id

@app.post("/api/v1/jupyter/execute_code", response_model=Optional[models.CodeboxOutput])
async def execute_code(code: models.Code,  session_key:models.CodeboxSessionKey) -> models.CodeboxOutput:
    if session_key not in global_codebox_session_manager.sessions.keys():
        raise HTTPException(400, detail=f"{session_key.session_id} not started")
    session = await global_codebox_session_manager.get_codebox_session(session_key=session_key)
    result = session.execute_code(code=code.code)
    print(result)
    return models.CodeboxOutput(
        output_type="text/plain" if len(result[1])<=0 else 'files/mixed',
        content='\n'.join(result[0]),
        files = result[1] if len(result[1])>0 else None
    )

@app.post("/api/v1/jupyter/install", response_model=str)
async def install(install_request:models.CodeboxInstallRequest,session_key:models.CodeboxSessionKey):
    if session_key not in global_codebox_session_manager.sessions.keys():
        raise HTTPException(400, detail=f"{session_key.session_id} not started")
    session = await global_codebox_session_manager.get_codebox_session(session_key=session_key)
    await session.install_packages(install_request.packagenames)
    return 'installed'


@app.post("/api/v1/execute_shell", response_model=Optional[models.CodeboxOutput])
async def execute_code(code: models.Code,  session_key:models.CodeboxSessionKey) -> models.CodeboxOutput:
    if session_key not in global_codebox_session_manager.sessions.keys():
        raise HTTPException(400, detail=f"{session_key.session_id} not started")
    session = await global_codebox_session_manager.get_codebox_session(session_key=session_key)
    try:
        execute_result = await asyncio.wait_for(session.arun_shell(code.code), timeout=20)
    except asyncio.TimeoutError as e:
        execute_result=f"The Execution of command '{code.code}' seems take too long.Maybe you don't need to wait for the result, please specify need re-execute or not."
    return models.CodeboxOutput(
        content=execute_result,
        output_type='shell'
    )

@app.post("/api/v1/upload")
async def upload(file: models.File, session_key:models.CodeboxSessionKey):
    session = await global_codebox_session_manager.get_codebox_session(session_key)
    return session.upload(file)


@app.post("/api/v1/upload_dev")
async def upload(session_key:Optional[str]=Form(None),file: UploadFile=File(...)):
    print(session_key)
    try:
        session_key = models.CodeboxSessionKey.parse_raw(session_key)
    except Exception as e:
        print(e)
        return HTTPException(status_code=400, detail='bad request with wrong format of session_key')
    print(session_key)
    session = await global_codebox_session_manager.get_codebox_session(session_key)
    file = models.File(
        name=file.filename,
        content=await file.read(),
    )
    print(file.name)
    return session.upload(file)


@app.post("/api/v1/upload_batch_dev")
async def upload(session_key:Optional[str]=Form(None),upload_files: List[UploadFile]=File(...)):
    try:
        if session_key:
            session_key = models.CodeboxSessionKey.parse_raw(session_key)
    except Exception as e:
        print(e)
        return HTTPException(status_code=400, detail='bad request with wrong format of session_key')
    print('uploaded files:', upload_files)
    session = await global_codebox_session_manager.get_codebox_session(session_key)
    processed_files = [models.File(
        name=file.filename,
        content=await file.read(),
    ) for file in upload_files]
    return session.upload_batch(processed_files)

@app.post('/api/v1/jupyter/download')
async def download(session_key:models.CodeboxSessionKey,filenames: List[str]) -> List[models.File]:
    session = await global_codebox_session_manager.get_codebox_session(session_key)
    file_list:List[models.File] = []
    for filename in filenames:
        try:
            file_list.append(session.download(filename))
        except Exception as e:
            print(e)
    return file_list


    
@app.post("/api/v1/codebox/stop")
async def close_codebox(session_key:models.CodeboxSessionKey):
    session = await global_codebox_session_manager.get_codebox_session(session_key)
    session.stop()
    global_codebox_session_manager.remove_session(session_key)


@app.on_event('shutdown')
async def shutdown():
    print('on shutting down')
    session_keys = global_codebox_session_manager.sessions.keys()
    for item in list(session_keys):
        global_codebox_session_manager.remove_session(item)
    print('all containers closed')

app/main.py
zhangzhejian/codeinterpreter-codebox
from pydantic import BaseModel, validator
from typing import List, Optional
import json
import datetime
from models.schemas import ProgrammingLanguage


class Code(BaseModel):
    code: str


class CodeModification(BaseModel):
    origin_code: str
    new_code: str
    filepath: str

    def on_duplicate(self):
        return f"Duplicate target code to replace found in {self.filepath}: {self.origin_code}"
    
    def on_not_found(self):
        return f"{self.origin_code} not found in {self.filepath}"

class CodeboxInitRequest(BaseModel):
    token: str

class CodeboxSessionKey(BaseModel):
    session_id: str
    language: ProgrammingLanguage
    def __hash__(self):
        return hash(self.session_id)

    def __eq__(self, other):
        if isinstance(other, CodeboxSessionKey):
            return self.session_id == other.session_id 
        return False
    

class File(BaseModel):
    name: str
    content: bytes


class CodeboxOutput(BaseModel):
    output_type: str
    content: str
    files: Optional[List[File]] = None
    

class CodeboxInstallRequest(BaseModel):
    packagenames: List[str]
app/models/api.py
zhangzhejian/codeinterpreter-codebox
from pydantic import BaseModel
from typing import List, Optional
import json
import datetime
from enum import Enum

class ProgrammingLanguage(str,Enum):
    JUPYTER="JUPYTER"

class PortUsage(str, Enum):
    TCP = "TCP"
    EXPORT="EXPORT"
    JUPYTER="JUPYTER"


class PortOccupation(BaseModel):
    session_id: str
    usage: PortUsage=PortUsage.TCP
    
app/models/schemas.py
zhangzhejian/codeinterpreter-codebox
import setuptools

# with open("README.md", "r") as fh:
#     long_description = fh.read()

with open("requirements.txt", "r") as fh:
    requirements = fh.read().splitlines()

setuptools.setup(
    name="",
    version="0.0.1",
    author="Zhejian Zhang",
    author_email="527284545@qq.com",
    install_requires=requirements,
    description="Backend for Code Interpreter for Feishu",
    # long_description=long_description,
    # long_description_content_type="text/markdown",
    packages=setuptools.find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: OS Independent",
    ],
    python_requires='>=3.11',
)

app/setup.py
zhangzhejian/codeinterpreter-codebox
import litellm
import requests
import json
from typing import Optional, Dict, Any, Union, List
import tokentrim as tt
from pydantic import BaseModel
import base64
import os

OPENAI_API_KEY='YOURAPI_KEY'
OPENAI_API_BASE='YOURAPI_BASE'

litellm.api_key=OPENAI_API_KEY
litellm.api_base=OPENAI_API_BASE

system_message="""You are Open Interpreter, a world-class programmer that can complete any goal by executing code.
When you send a message containing code to run_code, it will be executed **on the user's machine**. The user has given you **full and complete permission** to execute any code necessary to complete the task. You have full access to control their computer to help them. Code entered into run_code will be executed **in the users local environment**.
Never use (!) when running commands.
Only use the function you have been provided with, run_code.
You are capable of **any** task."""

function_schema = {
  "name": "run_code",
  "description":
  "Executes code on the remote jupyter server and returns the output",
  "parameters": {
    "type": "object",
    "properties": {
      "code": {
        "type": "string",
        "description": "The code to execute"
      }
    },
    "required": ["language", "code"]
  },
}

def make_request(
    method: str,
    url: str,
    headers: Optional[Dict[str, str]] = None,
    params: Optional[Dict[str, Union[str, int, float]]] = None,
    data: Optional[Dict[str, Any]] = None,
    json: Optional[Dict[str, Any]] = None,
    timeout: Optional[int] = 30
) -> requests.Response:
    """
    Make a generic HTTP request.
    
    :param method: HTTP method (e.g., GET, POST, PUT, DELETE)
    :param url: The URL for the request
    :param headers: Optional headers to include in the request
    :param params: Optional query parameters
    :param data: Data to send in the body of the request (for POST/PUT)
    :param json: JSON data to send in the body of the request (for POST/PUT)
    :param timeout: Request timeout in seconds
    :return: `requests.Response` object
    """

    supported_methods = {
        'GET': requests.get,
        'POST': requests.post,
        'PUT': requests.put,
        'DELETE': requests.delete
    }
    
    if method not in supported_methods:
        raise ValueError(f"Unsupported HTTP method: {method}")

    response = supported_methods[method](
        url,
        headers=headers,
        params=params,
        data=data,
        json=json,
        timeout=timeout
    )

    return response

class File(BaseModel):
    name: str
    content: bytes


class CodeboxOutput(BaseModel):
    output_type: str
    content: str
    files: Optional[List[File]] = None

class CodeinterpreterSession:
    domain="http://localhost:5002"

    def __init__(self,):
        self.model='gpt-4'
        self.messages=[]
        self.uploaded_files = []
        self.session_id=self.start_session()
        self.temperature=0.5
        self.system_message = system_message
        return


    def start_session(self):
        
        url = f"{self.domain}/api/v1/jupyter/start"
        response = make_request(
            method='POST', 
            url=url,
            headers={"accept":"application/json"},
            params={
                "token":"zhejianzhang",
                "language":"JUPYTER"
            }
            )
        session_id=response.text.strip('"')
        print(session_id)
        return session_id

    
    def install(self,packages):
        url = f"{self.domain}/api/v1/jupyter/install"
        data = {
                    "install_request": {
                    "packagenames":packages 
                    },
                    "session_key": {
                        "session_id":self.session_id,
                        "language": "JUPYTER"
                    }
                }
        response = make_request(
            method='POST', 
            url=url,
            headers={"accept":"application/json",'Content-Type': 'application/json'},
            json=data
            )

        print(response)
        print('Execute Result=',response.text)



    def upload_files(self,filenames: List[str]):
        files = [(f"upload_files", (os.path.basename(filename), open(filename, 'rb'))) for i, filename in enumerate(filenames)]
        url = self.domain + '/api/v1/upload_batch_dev'
        session_key = json.dumps({"session_id":self.session_id,"language": "JUPYTER"})

        response = requests.post(
            url,  
            data={'session_key':session_key},
            files= files)
        print(response.status_code)
        for file in files:
            self.uploaded_files.append(file[1][0])
            file[1][1].close()

    
    def run_code(self,code):
        url = f"{self.domain}/api/v1/jupyter/execute_code"
        data = {
                    "code": {
                    "code": code
                    },
                    "session_key": {
                        "session_id":self.session_id,
                        "language": "JUPYTER"
                    }
                }
        response = make_request(
            method='POST', 
            url=url,
            headers={"accept":"application/json",'Content-Type': 'application/json'},
            json=data
            )
        
        rsp = CodeboxOutput.parse_raw(response.content)
        
        print('Code Execute Result=',rsp.content)
        if rsp.files:
            for item in rsp.files:
                with open(item.name, 'wb') as f:
                    f.write(base64.b64decode(item.content))
                    print(f'image saved at {item.name}')

        return response.text
    
    def chat(self, msg: str):
        self.messages.append({'role':'user','content':msg})
        system_message = self.system_message + f"\n Uploaded files in directory /codebox/:{' '.join(self.uploaded_files)}"
        response = litellm.completion(
            model=self.model,
            messages=tt.trim(self.messages, self.model, system_message=system_message),
            functions=[function_schema],
            stream=False,
            temperature=self.temperature,
        )
        completion=response.choices[0].message
        print(completion)
        # args = completion['function_call']['arguments'].split("\n")

        # print(args)
        
        if "function_call" in completion.keys() and 'run_code' == completion['function_call']['name']:
            code = json.loads(completion['function_call']['arguments'])['code']
        #     lines = completion['function_call']['arguments'].split("\n")
        #     code = '\n'.join(lines[1:]).strip("` \n")
            self.run_code(code=code)
        
        return
    


    def close(self):
        url = f"{self.domain}/api/v1/codebox/stop"
        data = {
                        "session_id":self.session_id,
                        "language": "JUPYTER"
                    }
                
        response = make_request(
            method='POST', 
            url=url,
            headers={"accept":"application/json",'Content-Type': 'application/json'},
            json=data
            )

        if response.status_code ==200:
            print('remote container stopped')
        else:
            print('close error')
    

if __name__=='__main__':
    session=CodeinterpreterSession()
    try:
        session.upload_files(['./../data/test_data.csv'])
        session.chat('上传的文件中有多少列数据')
    finally:
        session.close()


examples/client/codeinterpreter_session.py
zhangzhejian/codeinterpreter-codebox