Modular wait for + Notebooks

.gitignore

@@ -179,4 +179,8 @@ cython_debug/
 ./mantis_sdk/main.py
 mantis_sdk/main.py
 test.csv
-./test.csv
+./test.csv
+main.py
+./main.py
+./venv
+venv/

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)
 ```

examples/demo.ipynb

@@ -9,13 +9,25 @@
     "# 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",
     "import pandas as pd\n",
     "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": {

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
+        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')}")