Skip to content

Commit 6a339b3

Browse files
mishushakovmishushakov
committed
added context methods to the SDK + tests
1 parent 483830e commit 6a339b3

File tree

8 files changed

+389
-6
lines changed

8 files changed

+389
-6
lines changed

‎.changeset/open-lamps-drop.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
'@e2b/code-interpreter-python': minor
3+
'@e2b/code-interpreter': minor
4+
---
5+
6+
added context methods to the sdk

‎js/src/messaging.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -145,10 +145,7 @@ export class Result {
145145

146146
readonly raw: RawData
147147

148-
constructor(
149-
rawData: RawData,
150-
public readonly isMainResult: boolean
151-
) {
148+
constructor(rawData: RawData, public readonly isMainResult: boolean) {
152149
const data = { ...rawData }
153150
delete data['type']
154151
delete data['is_main_result']

‎js/src/sandbox.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,4 +320,94 @@ export class Sandbox extends BaseSandbox {
320320
throw formatRequestTimeoutError(error)
321321
}
322322
}
323+
324+
/**
325+
* Removes a context.
326+
*
327+
* @param context context to remove.
328+
*
329+
* @returns void.
330+
*/
331+
async removeCodeContext(context: Context | string): Promise<void> {
332+
try {
333+
const id = typeof context === 'string' ? context : context.id
334+
const res = await fetch(`${this.jupyterUrl}/contexts/${id}`, {
335+
method: 'DELETE',
336+
headers: {
337+
'Content-Type': 'application/json',
338+
...this.connectionConfig.headers,
339+
},
340+
keepalive: true,
341+
signal: this.connectionConfig.getSignal(
342+
this.connectionConfig.requestTimeoutMs
343+
),
344+
})
345+
346+
const error = await extractError(res)
347+
if (error) {
348+
throw error
349+
}
350+
351+
return await res.json()
352+
} catch (error) {
353+
throw formatRequestTimeoutError(error)
354+
}
355+
}
356+
357+
/**
358+
* List all contexts.
359+
*
360+
* @returns list of contexts.
361+
*/
362+
async listCodeContexts(): Promise<Context[]> {
363+
try {
364+
const res = await fetch(`${this.jupyterUrl}/contexts`, {
365+
method: 'GET',
366+
headers: {
367+
'Content-Type': 'application/json',
368+
...this.connectionConfig.headers,
369+
},
370+
keepalive: true,
371+
signal: this.connectionConfig.getSignal(
372+
this.connectionConfig.requestTimeoutMs
373+
),
374+
})
375+
376+
const error = await extractError(res)
377+
if (error) {
378+
throw error
379+
}
380+
381+
return await res.json()
382+
} catch (error) {
383+
throw formatRequestTimeoutError(error)
384+
}
385+
}
386+
387+
/**
388+
* Restart a context.
389+
*
390+
* @param context context to restart.
391+
*
392+
* @returns void.
393+
*/
394+
async restartCodeContext(context: Context | string): Promise<void> {
395+
try {
396+
const id = typeof context === 'string' ? context : context.id
397+
const res = await fetch(`${this.jupyterUrl}/contexts/${id}/restart`, {
398+
method: 'POST',
399+
headers: {
400+
'Content-Type': 'application/json',
401+
...this.connectionConfig.headers,
402+
},
403+
})
404+
405+
const error = await extractError(res)
406+
if (error) {
407+
throw error
408+
}
409+
} catch (error) {
410+
throw formatRequestTimeoutError(error)
411+
}
412+
}
323413
}

‎js/tests/contexts.test.ts‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import { expect } from 'vitest'
2+
3+
import { isDebug, sandboxTest } from './setup'
4+
5+
sandboxTest('create context with no options', async ({ sandbox }) => {
6+
const context = await sandbox.createCodeContext()
7+
8+
expect(context.id).toBeDefined()
9+
expect(context.language).toBe('python')
10+
expect(context.cwd).toBe('/home/user')
11+
})
12+
13+
sandboxTest('create context with options', async ({ sandbox }) => {
14+
const context = await sandbox.createCodeContext({
15+
language: 'python',
16+
cwd: '/home/user/test',
17+
})
18+
19+
expect(context.id).toBeDefined()
20+
expect(context.language).toBe('python')
21+
expect(context.cwd).toBe('/home/user/test')
22+
})
23+
24+
sandboxTest('remove context', async ({ sandbox }) => {
25+
const context = await sandbox.createCodeContext()
26+
27+
await sandbox.removeCodeContext(context.id)
28+
})
29+
30+
sandboxTest('list contexts', async ({ sandbox }) => {
31+
const contexts = await sandbox.listCodeContexts()
32+
33+
expect(contexts.length).toBeGreaterThan(0)
34+
})
35+
36+
sandboxTest('restart context', async ({ sandbox }) => {
37+
const context = await sandbox.createCodeContext()
38+
39+
await sandbox.restartCodeContext(context.id)
40+
})

‎python/e2b_code_interpreter/code_interpreter_async.py‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
import httpx
33

4-
from typing import Optional, Dict, overload, Union, Literal
4+
from typing import Optional, Dict, overload, Union, Literal, List
55
from httpx import AsyncClient
66

77
from e2b import (
@@ -273,3 +273,89 @@ async def create_code_context(
273273
return Context.from_json(data)
274274
except httpx.TimeoutException:
275275
raise format_request_timeout_error()
276+
277+
async def remove_code_context(
278+
self,
279+
context: Union[Context, str],
280+
) -> None:
281+
"""
282+
Removes a context.
283+
284+
:param context: Context to remove. Can be a Context object or a context ID string.
285+
286+
:return: None
287+
"""
288+
context_id = context.id if isinstance(context, Context) else context
289+
290+
headers: Dict[str, str] = {}
291+
if self._envd_access_token:
292+
headers = {"X-Access-Token": self._envd_access_token}
293+
294+
try:
295+
response = await self._client.delete(
296+
f"{self._jupyter_url}/contexts/{context_id}",
297+
headers=headers,
298+
timeout=self.connection_config.request_timeout,
299+
)
300+
301+
err = await aextract_exception(response)
302+
if err:
303+
raise err
304+
except httpx.TimeoutException:
305+
raise format_request_timeout_error()
306+
307+
async def list_code_contexts(self) -> List[Context]:
308+
"""
309+
List all contexts.
310+
311+
:return: List of contexts.
312+
"""
313+
headers: Dict[str, str] = {}
314+
if self._envd_access_token:
315+
headers = {"X-Access-Token": self._envd_access_token}
316+
317+
try:
318+
response = await self._client.get(
319+
f"{self._jupyter_url}/contexts",
320+
headers=headers,
321+
timeout=self.connection_config.request_timeout,
322+
)
323+
324+
err = await aextract_exception(response)
325+
if err:
326+
raise err
327+
328+
data = response.json()
329+
return [Context.from_json(context_data) for context_data in data]
330+
except httpx.TimeoutException:
331+
raise format_request_timeout_error()
332+
333+
async def restart_code_context(
334+
self,
335+
context: Union[Context, str],
336+
) -> None:
337+
"""
338+
Restart a context.
339+
340+
:param context: Context to restart. Can be a Context object or a context ID string.
341+
342+
:return: None
343+
"""
344+
context_id = context.id if isinstance(context, Context) else context
345+
346+
headers: Dict[str, str] = {}
347+
if self._envd_access_token:
348+
headers = {"X-Access-Token": self._envd_access_token}
349+
350+
try:
351+
response = await self._client.post(
352+
f"{self._jupyter_url}/contexts/{context_id}/restart",
353+
headers=headers,
354+
timeout=self.connection_config.request_timeout,
355+
)
356+
357+
err = await aextract_exception(response)
358+
if err:
359+
raise err
360+
except httpx.TimeoutException:
361+
raise format_request_timeout_error()

‎python/e2b_code_interpreter/code_interpreter_sync.py‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
import httpx
33

4-
from typing import Optional, Dict, overload, Literal, Union
4+
from typing import Optional, Dict, overload, Literal, Union, List
55
from httpx import Client
66
from e2b import Sandbox as BaseSandbox, InvalidArgumentException
77

@@ -270,3 +270,89 @@ def create_code_context(
270270
return Context.from_json(data)
271271
except httpx.TimeoutException:
272272
raise format_request_timeout_error()
273+
274+
def remove_code_context(
275+
self,
276+
context: Union[Context, str],
277+
) -> None:
278+
"""
279+
Removes a context.
280+
281+
:param context: Context to remove. Can be a Context object or a context ID string.
282+
283+
:return: None
284+
"""
285+
context_id = context.id if isinstance(context, Context) else context
286+
287+
headers: Dict[str, str] = {}
288+
if self._envd_access_token:
289+
headers = {"X-Access-Token": self._envd_access_token}
290+
291+
try:
292+
response = self._client.delete(
293+
f"{self._jupyter_url}/contexts/{context_id}",
294+
headers=headers,
295+
timeout=self.connection_config.request_timeout,
296+
)
297+
298+
err = extract_exception(response)
299+
if err:
300+
raise err
301+
except httpx.TimeoutException:
302+
raise format_request_timeout_error()
303+
304+
def list_code_contexts(self) -> List[Context]:
305+
"""
306+
List all contexts.
307+
308+
:return: List of contexts.
309+
"""
310+
headers: Dict[str, str] = {}
311+
if self._envd_access_token:
312+
headers = {"X-Access-Token": self._envd_access_token}
313+
314+
try:
315+
response = self._client.get(
316+
f"{self._jupyter_url}/contexts",
317+
headers=headers,
318+
timeout=self.connection_config.request_timeout,
319+
)
320+
321+
err = extract_exception(response)
322+
if err:
323+
raise err
324+
325+
data = response.json()
326+
return [Context.from_json(context_data) for context_data in data]
327+
except httpx.TimeoutException:
328+
raise format_request_timeout_error()
329+
330+
def restart_code_context(
331+
self,
332+
context: Union[Context, str],
333+
) -> None:
334+
"""
335+
Restart a context.
336+
337+
:param context: Context to restart. Can be a Context object or a context ID string.
338+
339+
:return: None
340+
"""
341+
context_id = context.id if isinstance(context, Context) else context
342+
343+
headers: Dict[str, str] = {}
344+
if self._envd_access_token:
345+
headers = {"X-Access-Token": self._envd_access_token}
346+
347+
try:
348+
response = self._client.post(
349+
f"{self._jupyter_url}/contexts/{context_id}/restart",
350+
headers=headers,
351+
timeout=self.connection_config.request_timeout,
352+
)
353+
354+
err = extract_exception(response)
355+
if err:
356+
raise err
357+
except httpx.TimeoutException:
358+
raise format_request_timeout_error()

‎python/tests/async/test_async_contexts.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
2+
3+
4+
async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
5+
context = await async_sandbox.create_code_context()
6+
7+
assert context.id is not None
8+
assert context.language == "python"
9+
assert context.cwd == "/home/user"
10+
11+
12+
async def test_create_context_with_options(async_sandbox: AsyncSandbox):
13+
context = await async_sandbox.create_code_context(
14+
language="python",
15+
cwd="/home/user/test",
16+
)
17+
18+
assert context.id is not None
19+
assert context.language == "python"
20+
assert context.cwd == "/home/user/test"
21+
22+
23+
async def test_remove_context(async_sandbox: AsyncSandbox):
24+
context = await async_sandbox.create_code_context()
25+
26+
await async_sandbox.remove_code_context(context.id)
27+
28+
29+
async def test_list_contexts(async_sandbox: AsyncSandbox):
30+
contexts = await async_sandbox.list_code_contexts()
31+
32+
assert len(contexts) > 0
33+
34+
35+
async def test_restart_context(async_sandbox: AsyncSandbox):
36+
context = await async_sandbox.create_code_context()
37+
38+
await async_sandbox.restart_code_context(context.id)
39+

0 commit comments

Comments
 (0)