code_interpreter.py: allow non-root jupyter base_url

This commit is contained in:
Anthony Uk 2025-05-09 20:32:01 +02:00
parent 8acc1ab425
commit 337ece4d50

View File

@ -44,12 +44,14 @@ class JupyterCodeExecuter:
:param password: Jupyter password (optional) :param password: Jupyter password (optional)
:param timeout: WebSocket timeout in seconds (default: 60s) :param timeout: WebSocket timeout in seconds (default: 60s)
""" """
self.base_url = base_url.rstrip("/") self.base_url = base_url
self.code = code self.code = code
self.token = token self.token = token
self.password = password self.password = password
self.timeout = timeout self.timeout = timeout
self.kernel_id = "" self.kernel_id = ""
if self.base_url[-1] != '/':
self.base_url += '/'
self.session = aiohttp.ClientSession(trust_env=True, base_url=self.base_url) self.session = aiohttp.ClientSession(trust_env=True, base_url=self.base_url)
self.params = {} self.params = {}
self.result = ResultModel() self.result = ResultModel()
@ -61,7 +63,7 @@ class JupyterCodeExecuter:
if self.kernel_id: if self.kernel_id:
try: try:
async with self.session.delete( async with self.session.delete(
f"/api/kernels/{self.kernel_id}", params=self.params f"api/kernels/{self.kernel_id}", params=self.params
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
except Exception as err: except Exception as err:
@ -81,7 +83,7 @@ class JupyterCodeExecuter:
async def sign_in(self) -> None: async def sign_in(self) -> None:
# password authentication # password authentication
if self.password and not self.token: if self.password and not self.token:
async with self.session.get("/login") as response: async with self.session.get("login") as response:
response.raise_for_status() response.raise_for_status()
xsrf_token = response.cookies["_xsrf"].value xsrf_token = response.cookies["_xsrf"].value
if not xsrf_token: if not xsrf_token:
@ -89,7 +91,7 @@ class JupyterCodeExecuter:
self.session.cookie_jar.update_cookies(response.cookies) self.session.cookie_jar.update_cookies(response.cookies)
self.session.headers.update({"X-XSRFToken": xsrf_token}) self.session.headers.update({"X-XSRFToken": xsrf_token})
async with self.session.post( async with self.session.post(
"/login", "login",
data={"_xsrf": xsrf_token, "password": self.password}, data={"_xsrf": xsrf_token, "password": self.password},
allow_redirects=False, allow_redirects=False,
) as response: ) as response:
@ -102,16 +104,16 @@ class JupyterCodeExecuter:
async def init_kernel(self) -> None: async def init_kernel(self) -> None:
async with self.session.post( async with self.session.post(
url="/api/kernels", params=self.params url="api/kernels", params=self.params
) as response: ) as response:
response.raise_for_status() response.raise_for_status()
kernel_data = await response.json() kernel_data = await response.json()
self.kernel_id = kernel_data["id"] self.kernel_id = kernel_data["id"]
def init_ws(self) -> (str, dict): def init_ws(self) -> (str, dict):
ws_base = self.base_url.replace("http", "ws") ws_base = self.base_url.replace("http", "ws", 1)
ws_params = "?" + "&".join([f"{key}={val}" for key, val in self.params.items()]) ws_params = "?" + "&".join([f"{key}={val}" for key, val in self.params.items()])
websocket_url = f"{ws_base}/api/kernels/{self.kernel_id}/channels{ws_params if len(ws_params) > 1 else ''}" websocket_url = f"{ws_base}api/kernels/{self.kernel_id}/channels{ws_params if len(ws_params) > 1 else ''}"
ws_headers = {} ws_headers = {}
if self.password and not self.token: if self.password and not self.token:
ws_headers = { ws_headers = {