refactor(jupyter): replace httpx with aiohttp

backend/open_webui/utils/code_interpreter.py

@@ -1,65 +1,119 @@
 import asyncio
 import json
+import logging
 import uuid
 from typing import Optional
 
-import httpx
+import aiohttp
 import websockets
+from pydantic import BaseModel
+from websockets import ClientConnection
+
+logger = logging.getLogger(__name__)
 
 
-async def execute_code_jupyter(
+class ResultModel(BaseModel):
-    jupyter_url: str, code: str, token: str = None, password: str = None, timeout: int = 60
-) -> Optional[dict]:
     """
-    Executes Python code in a Jupyter kernel.
+    Execute Code Result Model
-    Supports authentication with a token or password.
+    """
-    :param jupyter_url: Jupyter server URL (e.g., "http://localhost:8888")
+
+    stdout: Optional[str] = ""
+    stderr: Optional[str] = ""
+    result: Optional[str] = ""
+
+
+class JupyterCodeExecuter:
+    """
+    Execute code in jupyter notebook
+    """
+
+    def __init__(self, base_url: str, code: str, token: str = "", password: str = "", timeout: int = 60):
+        """
+        :param base_url: Jupyter server URL (e.g., "http://localhost:8888")
         :param code: Code to execute
         :param token: Jupyter authentication token (optional)
         :param password: Jupyter password (optional)
-    :param timeout: WebSocket timeout in seconds (default: 10s)
+        :param timeout: WebSocket timeout in seconds (default: 60s)
-    :return: Dictionary with stdout, stderr, and result
-             - Images are prefixed with "base64:image/png," and separated by newlines if multiple.
         """
+        self.base_url = base_url.rstrip("/")
+        self.code = code
+        self.token = token
+        self.password = password
+        self.timeout = timeout
+        self.kernel_id = ""
+        self.session = aiohttp.ClientSession(base_url=self.base_url)
+        self.params = {}
+        self.result = ResultModel()
 
-    jupyter_url = jupyter_url.rstrip("/")
+    async def __aenter__(self):
-    client = httpx.AsyncClient(base_url=jupyter_url, timeout=timeout, follow_redirects=True)
+        return self
-    headers = {}
 
-    # password authentication
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
-    if password and not token:
+        if self.kernel_id:
             try:
-            response = await client.get("/login")
+                await self.session.delete(f"/api/kernels/{self.kernel_id}", params=self.params)
+            except Exception as err:
+                logger.exception("close kernel failed, %s", err)
+        await self.session.close()
+
+    async def run(self) -> ResultModel:
+        try:
+            await self.sign_in()
+            await self.init_kernel()
+            await self.execute_code()
+        except Exception as err:
+            logger.error(err)
+            self.result.stderr = f"Error: {err}"
+        return self.result
+
+    async def sign_in(self) -> None:
+        # password authentication
+        if self.password and not self.token:
+            async with self.session.get("/login") as response:
                 response.raise_for_status()
-            xsrf_token = response.cookies.get("_xsrf")
+                xsrf_token = response.cookies["_xsrf"].value
                 if not xsrf_token:
                     raise ValueError("_xsrf token not found")
-            response = await client.post("/login", data={"_xsrf": xsrf_token, "password": password})
+                self.session.cookie_jar.update_cookies(response.cookies)
+                self.session.headers.update({"X-XSRFToken": xsrf_token})
+            async with self.session.post(
+                "/login", data={"_xsrf": xsrf_token, "password": self.password}, allow_redirects=False
+            ) as response:
                 response.raise_for_status()
-            headers["X-XSRFToken"] = xsrf_token
+                self.session.cookie_jar.update_cookies(response.cookies)
-        except Exception as e:
-            return {"stdout": "", "stderr": f"Authentication Error: {str(e)}", "result": ""}
 
         # token authentication
-    params = {"token": token} if token else {}
+        if self.token:
+            self.params.update({"token": self.token})
 
-    kernel_id = ""
+    async def init_kernel(self) -> None:
-    try:
+        async with self.session.post(url="/api/kernels", params=self.params) as response:
-        response = await client.post(url="/api/kernels", params=params, headers=headers)
             response.raise_for_status()
-        kernel_id = response.json()["id"]
+            kernel_data = await response.json()
+            self.kernel_id = kernel_data["id"]
 
-        ws_base = jupyter_url.replace("http", "ws")
+    def init_ws(self) -> (str, dict):
-        websocket_url = f"{ws_base}/api/kernels/{kernel_id}/channels" + (f"?token={token}" if token else "")
+        ws_base = self.base_url.replace("http", "ws")
+        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 ''}"
         ws_headers = {}
-        if password and not token:
+        if self.password and not self.token:
             ws_headers = {
-                "X-XSRFToken": client.cookies.get("_xsrf"),
+                "Cookie": "; ".join([f"{cookie.key}={cookie.value}" for cookie in self.session.cookie_jar]),
-                "Cookie": "; ".join([f"{name}={value}" for name, value in client.cookies.items()]),
+                **self.session.headers,
             }
+        return websocket_url, ws_headers
 
+    async def execute_code(self) -> None:
+        # initialize ws
+        websocket_url, ws_headers = self.init_ws()
+        # execute
         async with websockets.connect(websocket_url, additional_headers=ws_headers) as ws:
-            msg_id = str(uuid.uuid4())
+            await self.execute_in_jupyter(ws)
+
+    async def execute_in_jupyter(self, ws: ClientConnection) -> None:
+        # send message
+        msg_id = uuid.uuid4().hex
         await ws.send(
             json.dumps(
                 {
@@ -67,14 +121,14 @@ async def execute_code_jupyter(
                         "msg_id": msg_id,
                         "msg_type": "execute_request",
                         "username": "user",
-                            "session": str(uuid.uuid4()),
+                        "session": uuid.uuid4().hex,
                         "date": "",
                         "version": "5.3",
                     },
                     "parent_header": {},
                     "metadata": {},
                     "content": {
-                            "code": code,
+                        "code": self.code,
                         "silent": False,
                         "store_history": True,
                         "user_expressions": {},
@@ -85,16 +139,17 @@ async def execute_code_jupyter(
                 }
             )
         )
-
+        # parse message
         stdout, stderr, result = "", "", []
-
         while True:
             try:
-                    message = await asyncio.wait_for(ws.recv(), timeout)
+                # wait for message
+                message = await asyncio.wait_for(ws.recv(), self.timeout)
                 message_data = json.loads(message)
+                # msg id not match, skip
                 if message_data.get("parent_header", {}).get("msg_id") != msg_id:
                     continue
-
+                # check message type
                 msg_type = message_data.get("msg_type")
                 match msg_type:
                     case "stream":
@@ -117,13 +172,14 @@ async def execute_code_jupyter(
             except asyncio.TimeoutError:
                 stderr += "\nExecution timed out."
                 break
+        self.result.stdout = stdout.strip()
+        self.result.stderr = stderr.strip()
+        self.result.result = "\n".join(result).strip() if result else ""
 
-    except Exception as e:
-        return {"stdout": "", "stderr": f"Error: {str(e)}", "result": ""}
 
-    finally:
+async def execute_code_jupyter(
-        if kernel_id:
+    base_url: str, code: str, token: str = "", password: str = "", timeout: int = 60
-            await client.delete(f"/api/kernels/{kernel_id}", headers=headers, params=params)
+) -> dict:
-        await client.aclose()
+    async with JupyterCodeExecuter(base_url, code, token, password, timeout) as executor:
-
+        result = await executor.run()
-    return {"stdout": stdout.strip(), "stderr": stderr.strip(), "result": "\n".join(result).strip() if result else ""}
+        return result.model_dump()