diff --git a/.gitignore b/.gitignore index 30b6b92..be7b51a 100644 --- a/.gitignore +++ b/.gitignore @@ -179,4 +179,8 @@ cython_debug/ ./mantis_sdk/main.py mantis_sdk/main.py test.csv -./test.csv \ No newline at end of file +./test.csv +main.py +./main.py +./venv +venv/ \ No newline at end of file diff --git a/README.md b/README.md index a74cc21..1cc16f2 100644 --- a/README.md +++ b/README.md @@ -35,11 +35,15 @@ The `MantisClient` object requires the parameter `cookie` to be passed in. This ***Note***: If you are using a local backend, you must run it with docker, or the space creation will NOT work. To do so, `cd docker` from the backend root, then `docker compose up -d --build`. After you run the build once, you can re-run it simply with `docker compose up`. If things dno't work, check the logs and make sure they are not empty. ```python -from client import MantisClient, SpacePrivacy, DataType, ReducerModels -from render_args import RenderArgs +from mantis_sdk.client import MantisClient, SpacePrivacy, DataType, ReducerModels +from mantis_sdk.render_args import RenderArgs +from mantis_sdk.config import ConfigurationManager import pandas as pd -mantis = MantisClient("/api/proxy/", cookie) +# You need to provide your cookie and a space_id (can be dummy for creation) +cookie = "YOUR_COOKIE_HERE" + +mantis = MantisClient("/api/proxy/", cookie=cookie, space_id="dummy") # Create DF (Real data will need more points) df = pd.DataFrame({ @@ -58,10 +62,11 @@ new_space_id = mantis.create_space("Stock data", data=df, data_types=data_types, reducer=ReducerModels.UMAP, - privacy_level=SpacePrivacy.Private)["space_id"] + privacy_level=SpacePrivacy.PRIVATE)["space_id"] # Open space -space = await mantis.open_space(space_id) +# Re-initialize client with the new space_id if needed, or just use the space object +space = await mantis.open_space(new_space_id) # Interact with space await space.select_points(100) @@ -166,4 +171,39 @@ await space.run_code (code) await space.close_panel ("bags") await space.close_panel ("quicksheet") await space.close_panel ("userlogs") +``` + +### Notebooks + +You can create and manage notebooks within a space. + +```python +# Create a notebook +nb = mantis.create_notebook(space_id, "My Analysis", "user_id") + +# Add a cell +code = """ +print("Hello from Mantis Notebook!") +""" +cell = nb.add_cell(code) + +# Execute cell +outputs = cell.execute() +print(outputs) +``` + +### Configuration + +You can configure the client using `ConfigurationManager`. + +```python +from mantis_sdk.config import ConfigurationManager + +config = ConfigurationManager() +config.update({ + "timeout": 60000, + "render_args": RenderArgs(headless=True) +}) + +mantis = MantisClient("/api/proxy/", cookie=cookie, space_id=space_id, config=config) ``` \ No newline at end of file diff --git a/examples/demo.ipynb b/examples/demo.ipynb index 789a328..df56796 100644 --- a/examples/demo.ipynb +++ b/examples/demo.ipynb @@ -9,6 +9,7 @@ "# Import main SDK\n", "from mantis_sdk.client import MantisClient, SpacePrivacy, DataType, ReducerModels\n", "from mantis_sdk.render_args import RenderArgs\n", + "from mantis_sdk.config import ConfigurationManager\n", "\n", "import nest_asyncio\n", "import asyncio\n", @@ -16,6 +17,17 @@ "nest_asyncio.apply()" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Setup authentication\n", + "# You need to get your cookie from the browser (dev tools -> network -> request headers)\n", + "cookie = \"YOUR_COOKIE_HERE\"\n" + ] + }, { "cell_type": "code", "execution_count": null, @@ -41,8 +53,10 @@ "metadata": {}, "outputs": [], "source": [ - "# Load Mantis Client\n", - "mantis = MantisClient(\"/api/proxy/\", render_args=RenderArgs(viewport={\"width\": 1920, \"height\": 1080}))" + "# Load Mantis Client with Config\n", + "config = ConfigurationManager()\n", + "config.update({\"render_args\": RenderArgs(viewport={\"width\": 1920, \"height\": 1080})})\n", + "mantis = MantisClient(\"/api/proxy/\", cookie=cookie, space_id=\"dummy\", config=config)" ] }, { @@ -68,7 +82,7 @@ "outputs": [], "source": [ "# Load Mantis\n", - "mantis = MantisClient(\"/api/proxy/\")\n", + "mantis = MantisClient(\"/api/proxy/\", cookie=cookie, space_id=\"dummy\")\n", "\n", "# Set data path + types\n", "data_path = \"./StockData.csv\"\n", @@ -94,7 +108,7 @@ "outputs": [], "source": [ "# Alternative you can create dataframes\n", - "mantis = MantisClient(\"/api/proxy/\")\n", + "mantis = MantisClient(\"/api/proxy/\", cookie=cookie, space_id=\"dummy\")\n", "\n", "# Create DF\n", "df = pd.DataFrame({\n", @@ -187,39 +201,6 @@ "imshow (await stock_data_space.capture ())" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Run Code**" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "code = \"\"\"\n", - "\n", - "computation = 6**4\n", - "print ('Hello from SDK, :P -> ' + str(computation))\n", - "\n", - "\"\"\"\n", - "\n", - "\n", - "await stock_data_space.run_code (code)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "imshow (await stock_data_space.capture ())" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -335,7 +316,7 @@ "metadata": {}, "outputs": [], "source": [ - "mantis = MantisClient(\"/api/proxy/\")" + "mantis = MantisClient(\"/api/proxy/\", cookie=cookie, space_id=\"dummy\")" ] }, { @@ -381,6 +362,46 @@ " print (\"Space:\", space_name)\n", " imshow (await space.capture ())" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Notebook Automation**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create a notebook in the space\n", + "# We need a valid space_id. Assuming 'new_space_id' from previous cells is valid.\n", + "if 'new_space_id' in locals():\n", + " client = MantisClient(\"/api/proxy/\", cookie=cookie, space_id=new_space_id)\n", + " nb = client.create_notebook(new_space_id, \"SDK Demo Notebook\", \"sdk_user\")\n", + " \n", + " # Add a cell\n", + " code = \"\"\"\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "x = np.linspace(0, 10, 100)\n", + "y = np.sin(x)\n", + "\n", + "plt.plot(x, y)\n", + "plt.title('Sine Wave')\n", + "plt.show()\n", + " \"\"\"\n", + " cell = nb.add_cell(code)\n", + " \n", + " # Execute the cell\n", + " outputs = cell.execute()\n", + " print(\"Execution outputs:\", outputs)\n", + "else:\n", + " print(\"Please run the 'Create Space' cells first to generate a space ID.\")" + ] } ], "metadata": { diff --git a/mantis_sdk/client.py b/mantis_sdk/client.py index 167012d..58b164c 100644 --- a/mantis_sdk/client.py +++ b/mantis_sdk/client.py @@ -47,12 +47,14 @@ class MantisClient: SDK for interacting with your Django API. """ - def __init__(self, base_url: str, cookie: str, config: Optional[ConfigurationManager] = None): + def __init__(self, base_url: str, cookie: str, space_id: str, config: Optional[ConfigurationManager] = None): """ Initialize the client. :param base_url: Base URL of the API. - :param token: Optional authentication token. + :param cookie: Authentication cookie. + :param space_id: The space ID to authenticate against. + :param config: Optional configuration manager. """ self.base_url = base_url.rstrip("/") @@ -62,12 +64,36 @@ def __init__(self, base_url: str, cookie: str, config: Optional[ConfigurationMan self.config = config self.cookie = cookie + self.space_id = space_id + self.vscode_token = None - if self.cookie is None: - self._authenticate () + if self.cookie: + self._authenticate() - def _authenticate (self): - raise NotImplementedError ("Authentication is not implemented yet.") + def _authenticate(self): + """ + Authenticates by creating a VSCode token. + """ + try: + # We use the cookie to get the token + # The endpoint is /api/jupyter/vscode-auth-token/ + # It requires project_id in the body + url = f"{self.base_url.lstrip('/').rstrip('/')}/api/jupyter/vscode-auth-token/" + headers = {"cookie": self.cookie} + data = {"project_id": self.space_id} + + response = requests.post(url, json=data, headers=headers) + response.raise_for_status() + result = response.json() + + if result.get("success") and result.get("token"): + self.vscode_token = result.get("token") + logger.info("Successfully authenticated with VSCode token.") + else: + logger.warning(f"Failed to get VSCode token: {result}") + + except Exception as e: + logger.error(f"Authentication failed: {e}") def _request(self, method: str, endpoint: str, rm_slash: bool = False, **kwargs) -> Any: """ @@ -80,7 +106,7 @@ def _request(self, method: str, endpoint: str, rm_slash: bool = False, **kwargs) def remove_slash (s: str): return s.lstrip('/').rstrip('/') - url = f"{self.config.host}/{remove_slash(self.base_url)}/{remove_slash(endpoint)}/" + url = f"{remove_slash(self.base_url)}/{remove_slash(endpoint)}/" # This is one of the weirdest cases I have required # some endpoints don't authenticate if there is a slash at the end @@ -90,6 +116,10 @@ def remove_slash (s: str): headers = {"cookie": self.cookie} + # add VSCode token if available + if self.vscode_token: + headers["VSCode-Token"] = self.vscode_token + if method.upper() == "GET": headers["Cache-Control"] = "no-cache" # Prevent caching for GET requests params = kwargs.get("params", {}) @@ -100,12 +130,17 @@ def remove_slash (s: str): headers.update (kwargs["headers"]) del kwargs["headers"] + reponse_content = None + try: response = requests.request(method, url, headers=headers, **kwargs) + reponse_content = response.text response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: - raise RuntimeError(f"API request failed: {e}. Text: {response.text}") + # capture response text for better debugging + response_text = reponse_content if reponse_content else getattr(e.response, 'text', '') if e.response else str(e) + raise RuntimeError(f"API request failed: {e}. Text: {response_text}\nURL: {url}\nHeaders: {headers}") def get_spaces (self) -> Dict[str, Any]: """ @@ -300,6 +335,17 @@ def create_space (self, space_name: str, return {"space_id": space_id} + def resolve_map_to_project(self, map_id: str) -> str: + """ + Resolves a map ID to a project ID. + """ + response = self._request("POST", "/api/notebook/resolve_map_to_project", json={"map_id": map_id}) + + if response.get("success"): + return response.get("project_id") + else: + raise RuntimeError(f"Failed to resolve map to project: {response.get('error')}") + async def open_space(self, space_id: str) -> "Space": """ Asynchronously open a space by ID. @@ -322,4 +368,41 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): Async context manager exit """ # Clean up if needed - pass \ No newline at end of file + pass + + def create_notebook(self, space_id: str, notebook_name: str = "Untitled", user_id: str = "sdk_user") -> "Notebook": + """ + Creates a new notebook in the specified space. + """ + from .notebook import Notebook + + payload = { + "user_id": user_id, + "notebook_name": notebook_name, + "project_id": space_id, + "notebook_code": '{"cells": [], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}' + } + + response = self._request("POST", "/api/notebook/create", json=payload) + + if response.get("success"): + nid = response.get("nid") + return Notebook(self, space_id, nid, notebook_name) + else: + raise RuntimeError(f"Failed to create notebook: {response.get('error')}") + + def get_notebook(self, space_id: str, notebook_id: str) -> "Notebook": + """ + Gets an existing notebook. + """ + from .notebook import Notebook + + # verify it exists + response = self._request("GET", "/api/notebook/get/", params={"nid": notebook_id, "project_id": space_id}) + + if response.get("success"): + notebook_data = response.get("notebook", {}) + name = notebook_data.get("notebook_name", "Untitled") + return Notebook(self, space_id, notebook_id, name) + else: + raise RuntimeError(f"Failed to get notebook: {response.get('error')}") \ No newline at end of file diff --git a/mantis_sdk/notebook.py b/mantis_sdk/notebook.py new file mode 100644 index 0000000..abe6d81 --- /dev/null +++ b/mantis_sdk/notebook.py @@ -0,0 +1,287 @@ +from typing import List, Optional, Dict, Any, TYPE_CHECKING +import time +import logging + +if TYPE_CHECKING: + from .client import MantisClient + +logger = logging.getLogger(__name__) + +class Cell: + """ + Represents a single cell in a Mantis Notebook. + """ + def __init__(self, notebook: "Notebook", index: int, cell_data: Dict[str, Any]): + self.notebook = notebook + self.index = index + self._data = cell_data + + @property + def cell_type(self) -> str: + return self._data.get("cell_type", "code") + + @property + def source(self) -> str: + source = self._data.get("source", "") + if isinstance(source, list): + return "".join(source) + return source + + @property + def outputs(self) -> List[Any]: + return self._data.get("outputs", []) + + @property + def metadata(self) -> Dict[str, Any]: + return self._data.get("metadata", {}) + + def update(self, content: str): + """ + Updates the content of this cell. + """ + self.notebook.update_cell(self.index, content) + + if isinstance(self._data["source"], list): + self._data["source"] = [content] + else: + self._data["source"] = content + + def execute(self): + """ + Executes this cell. + """ + return self.notebook.execute_cell(self.index) + + def delete(self): + """ + Deletes this cell from the notebook. + """ + self.notebook.delete_cell(self.index) + + def get_metadata(self) -> Dict[str, Any]: + """ + Returns the metadata of this cell. + """ + return self.metadata + + def __repr__(self): + return f"" + + +class Notebook: + """ + Represents a Mantis Notebook. + """ + def __init__(self, client: "MantisClient", space_id: str, nid: str, notebook_name: str = "Untitled"): + self.client = client + self.space_id = space_id + self.nid = nid + self.notebook_name = notebook_name + self.session_id: Optional[str] = None + self.cells: List[Cell] = [] + self._refresh_content() + + def _ensure_session(self): + """ + Ensures that a valid session exists. If not, creates one. + """ + if self.session_id: + # Check if session is valid + try: + response = self.client._request("POST", "/api/sessions/check", json={"session_id": self.session_id, "project_id": self.space_id}) + if response.get("success"): + return + except Exception: + logger.warning("Session check failed, creating new session.") + self.session_id = None + + def _create_session(self): + user_id = getattr(self.client, "user_id", "sdk_user") + + payload = { + "user_id": user_id, + "nid": self.nid, + "project_id": self.space_id + } + response = self.client._request("POST", "/api/sessions/create", json=payload) + if response.get("success"): + self.session_id = response.get("session_id") + else: + raise RuntimeError(f"Failed to create session: {response.get('error')}") + + def _refresh_content(self): + """ + Refreshes the notebook content from the backend. + """ + if not self.session_id: + self._ensure_session() + if not self.session_id: + self._create_session() + + response = self.client._request("GET", "/api/notebook/content", params={"session_id": self.session_id, "project_id": self.space_id}) + + if response.get("success"): + content = response.get("content", {}) + if "cells" in content: + self.cells = [Cell(self, i, cell_data) for i, cell_data in enumerate(content.get("cells", []))] + else: + logger.warning("Received successful response but content is missing 'cells'. Keeping previous state.") + else: + # If session not found, try to recreate + if "session not found" in str(response.get("error", "")).lower(): + self.session_id = None + self._create_session() + # Retry once + response = self.client._request("GET", "/api/notebook/content", params={"session_id": self.session_id, "project_id": self.space_id}) + if response.get("success"): + content = response.get("content", {}) + if "cells" in content: + self.cells = [Cell(self, i, cell_data) for i, cell_data in enumerate(content.get("cells", []))] + else: + logger.warning("Received successful response (retry) but content is missing 'cells'. Keeping previous state.") + else: + raise RuntimeError(f"Failed to get notebook content: {response.get('error')}") + else: + raise RuntimeError(f"Failed to get notebook content: {response.get('error')}") + + def refresh(self): + self._refresh_content() + + def add_cell(self, content: str = "", cell_type: str = "code", position: int = -1) -> Cell: + """ + Adds a new cell to the notebook. + """ + if not self.session_id: + self._create_session() + + payload = { + "session_id": self.session_id, + "cell_type": cell_type, + "content": content, + "position": position, + "project_id": self.space_id + } + + response = self.client._request("POST", "/api/notebook/add_cell", json=payload) + if response.get("success"): + self._refresh_content() + # Return the newly created cell. + # If position was -1, it's the last one. + # If position was specified, it's at that index (assuming 0-based index from backend matches). + if position == -1: + return self.cells[-1] + else: + # If position is within bounds, return that cell. + # Note: Backend might insert at position, shifting others. + # We should probably trust the index we passed if it's valid. + if 0 <= position < len(self.cells): + return self.cells[position] + return self.cells[-1] # Fallback + else: + raise RuntimeError(f"Failed to add cell: {response.get('error')}") + + def delete_cell(self, index: int): + """ + Deletes a cell by index. + """ + if not self.session_id: + self._create_session() + + payload = { + "session_id": self.session_id, + "position": index, + "project_id": self.space_id + } + + response = self.client._request("POST", "/api/notebook/delete_cell", json=payload) + if response.get("success"): + self._refresh_content() + else: + raise RuntimeError(f"Failed to delete cell: {response.get('error')}") + + def update_cell(self, index: int, content: str): + """ + Updates the content of a cell. + """ + if not self.session_id: + self._create_session() + + payload = { + "session_id": self.session_id, + "cell_index": index, + "content": content, + "project_id": self.space_id + } + + response = self.client._request("POST", "/api/notebook/edit_cell", json=payload) + if response.get("success"): + self._refresh_content() + else: + raise RuntimeError(f"Failed to update cell: {response.get('error')}") + + def execute_cell(self, index: int): + """ + Executes a cell by index. + """ + if not self.session_id: + self._create_session() + + payload = { + "session_id": self.session_id, + "cell_index": index, + "project_id": self.space_id + } + + response = self.client._request("POST", "/api/sessions/execute", json=payload) + if not response.get("success"): + raise RuntimeError(f"Failed to execute cell: {response.get('error')}") + + while True: + time.sleep(0.5) + self._refresh_content() + + if index >= len(self.cells): + logger.warning(f"Cell index {index} out of range (cells count: {len(self.cells)}). Waiting for refresh...") + continue + + cell = self.cells[index] + # check if executing + # note: metadata might be None + meta = cell.metadata or {} + if not meta.get("executing", False): + # execution finished + return cell.outputs + + def execute_all(self): + """ + Executes all code cells in the notebook. + """ + results = [] + for i, cell in enumerate(self.cells): + if cell.cell_type == "code": + results.append(self.execute_cell(i)) + return results + + def get_cell(self, index: int) -> Cell: + """ + Returns the cell at the specified index. + """ + if 0 <= index < len(self.cells): + return self.cells[index] + raise IndexError("Cell index out of range") + + def delete(self): + """ + Deletes the notebook and its session. + """ + if self.session_id: + payload = { + "nid": self.nid, + "session_id": self.session_id, + "project_id": self.space_id + } + self.client._request("POST", "/api/notebook/drop", json=payload) + self.session_id = None + else: + # if no session, just drop the notebook + self.client._request("POST", "/api/notebook/drop", json={"nid": self.nid, "project_id": self.space_id}) diff --git a/mantis_sdk/requirements.txt b/mantis_sdk/requirements.txt index 9a352b1..3bcd514 100644 --- a/mantis_sdk/requirements.txt +++ b/mantis_sdk/requirements.txt @@ -39,7 +39,6 @@ ptyprocess==0.7.0 pure_eval==0.2.3 pyasn1==0.6.1 pyasn1_modules==0.4.1 -pyee==11.1.1 Pygments==2.19.1 pyparsing==3.2.1 pyppeteer @@ -58,5 +57,4 @@ typing_extensions==4.12.2 tzdata==2025.1 urllib3==1.26.20 wcwidth==0.2.13 -websockets==10.4 -zipp==3.21.0 +websockets==10.4 \ No newline at end of file diff --git a/mantis_sdk/space.py b/mantis_sdk/space.py index fa95421..9932b78 100644 --- a/mantis_sdk/space.py +++ b/mantis_sdk/space.py @@ -92,9 +92,11 @@ async def _init_space(self): timeout=self.config.timeout) await self._apply_init_render_args () + + wait_for = self.config.wait_for if hasattr(self.config, 'wait_for') else "isLoaded" # Wait until the exposed loading value is true - await self.page.wait_for_function ("""() => window.isLoaded === true""", + await self.page.wait_for_function (f"""() => window.{wait_for} === true""", timeout=self.config.timeout) # Let points render after data is loaded diff --git a/output_plot.png b/output_plot.png new file mode 100644 index 0000000..2141f72 Binary files /dev/null and b/output_plot.png differ