Skip to content

Commit 28eca64

Browse files
committed
feat: add R support
1 parent 00986ce commit 28eca64

12 files changed

+168
-92
lines changed

.cursor/rules/global.mdc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,5 @@ alwaysApply: true
66
- Project uses uv and pyproject.toml for deps management
77
- Install deps with uv sync --all-extras
88
- After making changes, run `pytest tests -v` or for specific test `pytest tests/{test} -v`
9+
- Check test logs from `logs/test.log`
910
- After making changes, verify that [endpoints.py](mdc:app/api/endpoints.py) still follows [librechat-code-interpreter-openapi.json](mdc:project/librechat-code-interpreter-openapi.json)

.cursor/rules/librechat-endpoints.mdc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ All endpoints are prefixed with `/v1/librechat`
1212
### 1. Execute Code
1313
**Endpoint:** `POST /exec`
1414

15-
**Description:** Execute Python code in a secure sandbox environment
15+
**Description:** Execute Python and R code in a sandboxed environment
1616

1717
**Request Body:**
1818
```json

README.md

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,14 @@ A FastAPI-based code interpreter service that provides code execution and file m
1010
## Features
1111

1212
- Easy deployment with single docker compose file
13-
- Code execution in Docker container sandbox
14-
- File upload and download
15-
- Multi-user support
13+
- Code is executed in a isolated Docker container sandbox, custom images supported
14+
- Supports file upload and download
15+
- Supports concurrent code execution
16+
- Supports Python and R languages (possibility to extend to other languages)
1617
- RESTful API with OpenAPI documentation
1718

18-
1919
## Usage
2020

21-
2221
### Running the project
2322

2423
Run the project with docker compose using `docker compose -f compose.prod.yml up`

app/api/base.py

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
# Initialize services
3131
file_manager = FileManager()
3232

33-
SUPPORTED_LANGUAGES = {"py"} # Only Python is supported for now
33+
SUPPORTED_LANGUAGES = {"py", "r"} # Python and R are supported
3434
MAX_RETRIES = 3
3535

3636

@@ -48,18 +48,26 @@ async def execute_code(
4848
CodeExecutionRequest,
4949
Body(
5050
openapi_examples={
51-
"Hello World": {
52-
"summary": "Hello World",
51+
"Hello World (Python)": {
52+
"summary": "Hello World in Python",
5353
"value": {"code": "print('Hello, world!')", "lang": "py"},
5454
},
55-
"Random Number": {
56-
"summary": "Random Number",
55+
"Random Number (Python)": {
56+
"summary": "Random Number in Python",
5757
"value": {"code": "import random; print(random.randint(1, 100))", "lang": "py"},
5858
},
59-
"Sleep": {
59+
"Sleep (Python)": {
6060
"summary": "Sleep",
6161
"value": {"code": "import time; time.sleep(10); print('Done sleeping')", "lang": "py"},
6262
},
63+
"Hello World (R)": {
64+
"summary": "Hello World in R",
65+
"value": {"code": "cat('Hello, world!')", "lang": "r"},
66+
},
67+
"Random Number (R)": {
68+
"summary": "Random Number in R",
69+
"value": {"code": "cat(sample(1:100, 1))", "lang": "r"},
70+
},
6371
}
6472
),
6573
],
@@ -69,7 +77,7 @@ async def execute_code(
6977

7078
if request.lang not in SUPPORTED_LANGUAGES:
7179
raise BadLanguageException( # noqa: F821
72-
message=f"Language '{request.lang}' is not supported. Please use Python ('py')."
80+
message=f"Language '{request.lang}' is not supported. Please use Python ('py') or R ('r')."
7381
)
7482

7583
try:
@@ -90,22 +98,40 @@ async def execute_code(
9098

9199
# Execute code in Docker container
92100
result = await docker_executor.execute(
93-
code=request.code, session_id=session_id, files=files, timeout=settings.SANDBOX_MAX_EXECUTION_TIME
101+
code=request.code,
102+
session_id=session_id,
103+
lang=request.lang,
104+
files=files,
105+
timeout=settings.SANDBOX_MAX_EXECUTION_TIME,
94106
)
95107

96-
# Add a more Python specific error message if the stdout is empty (i.e. the code didn't print anything)
108+
# Add a language-specific error message if the stdout is empty
97109
if not result.get("stdout"):
98-
result["stdout"] = "Empty. Make sure to explicitly print the results"
110+
if request.lang == "py":
111+
result["stdout"] = "Empty. Make sure to explicitly print the results in Python"
112+
elif request.lang == "r":
113+
result["stdout"] = "Empty. Make sure to use print() or cat() to display results in R"
114+
else:
115+
result["stdout"] = "Empty. Make sure to explicitly output the results"
99116

100117
# Convert output files to FileRef model
101118
output_files = [
102119
FileRef(id=file["id"], name=file["filename"], path=file["filepath"]) for file in result.get("files", [])
103120
]
104121

122+
# Get language-specific version information
123+
version_info = ""
124+
if request.lang == "py":
125+
version_info = f"Python {sys.version.split()[0]}"
126+
elif request.lang == "r":
127+
version_info = "R (Jupyter R-notebook)"
128+
else:
129+
version_info = f"Unknown language: {request.lang}"
130+
105131
response = CodeExecutionResponse(
106132
run=result,
107133
language=request.lang,
108-
version=f"Python {sys.version.split()[0]}",
134+
version=version_info,
109135
session_id=session_id,
110136
files=output_files,
111137
)
@@ -121,7 +147,7 @@ async def execute_code(
121147
"status": result.get("status"),
122148
},
123149
"language": request.lang,
124-
"version": f"Python {sys.version.split()[0]}",
150+
"version": version_info,
125151
"session_id": session_id,
126152
"files": [f.model_dump() for f in output_files],
127153
}

app/api/librechat.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,21 +47,21 @@ def create_error_response(status_code: int, message: str) -> JSONResponse:
4747
"/exec",
4848
responses={400: {"model": LibreChatError}, 500: {"model": LibreChatError}},
4949
response_model=LibreChatExecuteResponse,
50-
description="Execute Python code in a secure sandbox environment",
51-
summary="Execute Python code",
50+
description="Execute Python or R code in a sandboxed environment",
51+
summary="Execute code",
5252
response_description="Returns the execution results",
5353
)
5454
async def execute_code(request: CodeExecutionRequest) -> LibreChatExecuteResponse:
55-
"""Execute Python code in a secure sandbox environment.
55+
"""Execute code in a sandboxed environment.
5656
5757
This endpoint handles code execution requests from LibreChat. It processes the provided
58-
Python code in an isolated environment and returns the execution results.
58+
code in an isolated environment and returns the execution results.
5959
6060
Args:
6161
request (CodeExecutionRequest): Request object containing:
62-
- code: Python code to execute
62+
- code: Code to execute
6363
- files: Optional list of files needed for execution
64-
- language: Programming language (must be 'python')
64+
- language: Programming language ('py' for Python or 'r' for R)
6565
- stdin: Optional standard input for the code
6666
6767
Returns:

app/main.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ async def lifespan(app: FastAPI):
4545
# Create FastAPI application
4646
app = FastAPI(
4747
title="Code Interpreter API",
48-
description="API for executing Python code in a secure environment",
48+
description="API for executing Python and R code in a sandboxed environment",
4949
version="1.0.0",
5050
lifespan=lifespan,
5151
)

app/models/base.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ class CodeExecutionRequest(BaseModel):
3232
lang: str = Field(
3333
...,
3434
description="The programming language of the code",
35-
examples=["py"],
35+
examples=["py", "r"],
3636
pattern="^(c|cpp|d|f90|go|java|js|php|py|rs|ts|r)$",
3737
)
3838
args: Optional[List[str]] = Field(None, description="Optional command line arguments to pass to the program")

app/services/docker_executor.py

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import docker
22
from loguru import logger
33
from pathlib import Path
4-
from typing import Dict, List, Optional, Any
4+
from typing import Dict, List, Optional, Any, Literal
55
import time
66
from docker.errors import APIError, ImageNotFound
77
import asyncio
@@ -35,14 +35,29 @@ class DockerExecutor:
3535

3636
WORK_DIR = "/mnt/data" # Working directory will be the same as data mount point
3737
DATA_MOUNT = "/mnt/data" # Mount point for session data
38-
MAX_CONCURRENT_CONTAINERS = 10
38+
39+
# Language-specific execution commands
40+
LANGUAGE_EXECUTORS = {
41+
"py": ["python", "-c"],
42+
"r": ["Rscript", "-e"],
43+
}
44+
45+
# Language-specific messages
46+
LANGUAGE_SPECIFIC_MESSAGES = {
47+
"py": {
48+
"empty_output": "Empty. Make sure to explicitly print() the results in Python"
49+
},
50+
"r": {
51+
"empty_output": "Empty. Make sure to use print() or cat() to display results in R"
52+
}
53+
}
3954

4055
def __init__(self):
41-
self._container_semaphore = asyncio.Semaphore(self.MAX_CONCURRENT_CONTAINERS)
56+
self._container_semaphore = asyncio.Semaphore(settings.MAX_CONCURRENT_CONTAINERS)
4257
self._active_containers: Dict[str, ContainerMetrics] = {}
4358
self._lock = asyncio.Lock()
4459
self._docker = None # Will be initialized in initialize()
45-
self._image_pull_locks: Dict[str, asyncio.Lock] = {} # Locks for image pulling
60+
self._image_pull_locks: Dict[str, asyncio.Lock] = {}
4661

4762
async def initialize(self):
4863
"""Initialize the Docker client."""
@@ -55,7 +70,7 @@ async def initialize(self):
5570
# Reinitialize if there was an error
5671
await self.close()
5772
self._docker = aiodocker.Docker()
58-
73+
5974
logger.info("Docker client initialized successfully")
6075
return self
6176
except Exception as e:
@@ -155,7 +170,12 @@ def _clean_output(self, raw_output: bytes) -> str:
155170
return b"".join(output_parts).decode("utf-8").strip()
156171

157172
async def execute(
158-
self, code: str, session_id: str, files: Optional[List[Dict[str, Any]]] = None, timeout: int = 30
173+
self,
174+
code: str,
175+
session_id: str,
176+
lang: Literal["py", "r"],
177+
files: Optional[List[Dict[str, Any]]] = None,
178+
timeout: int = 30,
159179
) -> Dict[str, Any]:
160180
"""Execute code in a Docker container with file management."""
161181
container = None
@@ -183,9 +203,9 @@ async def execute(
183203
async with self._container_semaphore:
184204
try:
185205
# Ensure the image is available
186-
image_name = settings.PYTHON_CONTAINER_IMAGE
206+
image_name = settings.LANGUAGE_CONTAINERS.get(lang)
187207
logger.info(f"Using container image: {image_name}")
188-
208+
189209
try:
190210
# Check if image exists
191211
await self._docker.images.inspect(image_name)
@@ -196,15 +216,18 @@ async def execute(
196216
# Get or create a lock for this specific image
197217
if image_name not in self._image_pull_locks:
198218
self._image_pull_locks[image_name] = asyncio.Lock()
199-
219+
200220
# Acquire the lock for this image to prevent multiple pulls
201221
async with self._image_pull_locks[image_name]:
202222
# Check again if the image exists (another request might have pulled it while we were waiting)
203223
try:
204224
await self._docker.images.inspect(image_name)
205225
logger.info(f"Image {image_name} is now available (pulled by another request)")
206226
except Exception as check_again_error:
207-
if isinstance(check_again_error, aiodocker.exceptions.DockerError) and check_again_error.status == 404:
227+
if (
228+
isinstance(check_again_error, aiodocker.exceptions.DockerError)
229+
and check_again_error.status == 404
230+
):
208231
# Pull the image if not available
209232
logger.info(f"Image {image_name} not found, pulling...")
210233
try:
@@ -227,7 +250,7 @@ async def execute(
227250
# Re-raise if it's not a 404 error
228251
logger.error(f"Error checking for image {image_name}: {str(e)}")
229252
raise
230-
253+
231254
# Create container config
232255
config = {
233256
"Image": image_name,
@@ -284,8 +307,16 @@ async def execute(
284307
output = await response.read()
285308
output_text = self._clean_output(output)
286309

287-
# Execute the Python code as jovyan user
288-
exec = await container.exec(cmd=["python", "-c", code], user="jovyan", stdout=True, stderr=True)
310+
# Execute the code with the appropriate interpreter
311+
logger.info(f"Code to execute: {code}")
312+
logger.info(f"Language: {lang}")
313+
314+
# Get the execution command for the specified language
315+
exec_cmd = self.LANGUAGE_EXECUTORS.get(lang, self.LANGUAGE_EXECUTORS["py"])
316+
logger.info(f"Using execution command: {exec_cmd}")
317+
318+
# Execute the code with the appropriate interpreter
319+
exec = await container.exec(cmd=[*exec_cmd, code], user="jovyan", stdout=True, stderr=True)
289320
# Use raw API call to get output
290321
exec_url = f"exec/{exec._id}/start"
291322
async with self._docker._query(

app/shared/config.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from pydantic_settings import BaseSettings, SettingsConfigDict
33
from functools import lru_cache
44
from pathlib import Path
5-
from typing import Set
5+
from typing import Dict, Set
66

77

88
class Settings(BaseSettings):
@@ -84,7 +84,19 @@ def HOST_FILE_UPLOAD_PATH_ABS(self) -> Path:
8484
CLEANUP_RUN_INTERVAL: int = 3600 # How often to run the cleanup in seconds
8585
CLEANUP_FILE_MAX_AGE: int = 86400 # How old files can be before they are deleted in seconds
8686

87-
PYTHON_CONTAINER_IMAGE: str = "jupyter/scipy-notebook:latest"
87+
PY_CONTAINER_IMAGE: str = "jupyter/scipy-notebook:latest"
88+
R_CONTAINER_IMAGE: str = "jupyter/r-notebook:latest"
89+
90+
@property
91+
def LANGUAGE_CONTAINERS(self) -> Dict[str, str]:
92+
"""Map language codes to container images."""
93+
return {
94+
"py": self.PY_CONTAINER_IMAGE,
95+
"r": self.R_CONTAINER_IMAGE,
96+
}
97+
98+
# Docker execution settings
99+
MAX_CONCURRENT_CONTAINERS: int = 10 # Maximum number of concurrent Docker containers
88100

89101

90102
@lru_cache()

tests/conftest.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,12 @@
1111
from app.main import app
1212
from loguru import logger
1313

14-
14+
# We log to a file to avoid polluting the console with logs
1515
logger.remove()
1616
logs_path = Path("logs/test.log")
1717
logs_path.parent.mkdir(exist_ok=True, parents=True)
18-
# logs_path.unlink(missing_ok=True)
18+
# Clear logs from previous runs
19+
logs_path.unlink(missing_ok=True)
1920
logger.add(logs_path)
2021

2122

0 commit comments

Comments
 (0)