switch to connect_realtime accepting a file or file ID (#133)

Author: rgbkrk

README.md

@@ -44,6 +44,7 @@ Python 3.8+
 ## Installation
 
 For stable release:
+
 ```bash
 pip install noteable-origami
 ```
@@ -53,6 +54,7 @@ poetry add noteable-origami
 ```
 
 For alpha pre-release:
+
 ```bash
 pip install noteable-origami --pre
 ```
@@ -64,8 +66,7 @@ pip install noteable-origami --pre
 ## Getting Started
 
 > **Warning**
-Developer note: this documentation is written for the 1.0 alpha release. For stable release, see [pre-1.0 README](https://github.com/noteable-io/origami/blob/release/0.0.35/README.md)
-
+> Developer note: this documentation is written for the 1.0 alpha release. For stable release, see [pre-1.0 README](https://github.com/noteable-io/origami/blob/release/0.0.35/README.md)
 
 ### API Tokens
 
@@ -76,7 +77,7 @@ The Noteable API requires an authentication token. You can manage tokens at the
 
 ### Usage
 
-The example below shows how to create a Notebook, launch a Kernel, add new cells, and execute code. 
+The example below shows how to create a Notebook, launch a Kernel, add new cells, and execute code.
 
 ```python
 # Grab a project_id from the Noteable UI, the url will look like: app.noteable.io/p/....
@@ -99,15 +100,15 @@ file = await api_client.create_notebook(project_id=project_id, path="Demo.ipynb"
 await api_client.launch_kernel(file.id)
 
 # Client for Real-time Updates (RTU), used with Notebooks
-rtu_client = await api_client.rtu_client(file.id)
+realtime_notebook = await api_client.connect_realtime(file)
 
 # Add a new cell
 from origami.models.notebook import CodeCell
 cell = CodeCell(source="print('Hello World')")
-await rtu_client.add_cell(cell)
+await realtime_notebook.add_cell(cell)
 
 # Execute the cell
-queued_execution = await rtu_client.queue_execution(cell.id)
+queued_execution = await realtime_notebook.queue_execution(cell.id)
 
 # Wait for the execution to be complete, cell is an updated instance of CodeCell with metadata/outputs
 cell = await queued_execution
@@ -117,10 +118,8 @@ output_collection = await api_client.get_output_collection(cell.output_collectio
 print(output_collection.outputs[0].content.raw) # 'Hello World\n'
 ```
 
-
 <!-- --8<-- [end:start] -->
 
-
 ## 1.0 Roadmap
 
 Origami is heading towards a 1.0 release. The alpha release candidate is on Pypi now, installable with a `--pre` flag. The 1.0 release represents a major refactor of the Origami using the best practices and lessons learned from creating multiple production API and RTU clients, including our ChatGPT plugin. It will likely come out of alpha once all of our internal applications are using the Origami 1.0 syntax.
@@ -130,6 +129,7 @@ Origami is heading towards a 1.0 release. The alpha release candidate is on Pypi
 See [CONTRIBUTING.md](./CONTRIBUTING.md).
 
 ---
+
 <p align="center">Open sourced with ❤️ by <a href="https://noteable.io">Noteable</a> for the community.</p>
 
-<img href="https://pages.noteable.io/private-beta-access" src="https://assets.noteable.io/github/2022-07-29/noteable.png" alt="Boost Data Collaboration with Notebooks">
+<img href="https://pages.noteable.io/private-beta-access" src="https://assets.noteable.io/github/2022-07-29/noteable.png" alt="Boost Data Collaboration with Notebooks">

origami/clients/api.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import uuid
-from typing import List, Literal, Optional
+from typing import List, Literal, Optional, Union
 
 import httpx
 import pydantic
@@ -284,14 +284,27 @@ async def get_output_collection(
         resp.raise_for_status()
         return KernelOutputCollection.parse_obj(resp.json())
 
-    async def rtu_client(self, file_id: uuid.UUID) -> RTUClient:
+    async def connect_realtime(self, file: Union[File, uuid.UUID, str]) -> RTUClient:
         """
         Create an RTUClient for a Notebook by file id. This will perform the following steps:
          - Check /v1/files to get the current version information and presigned download url
          - Download seed notebook and create a NotebookBuilder from it
          - Create an RTUClient, initialize the websocket connection, authenticate, and subscribe
          - Apply delts to in-memory NotebookBuilder
         """
+        file_id = None
+
+        if isinstance(file, str):
+            file_id = uuid.UUID(file)
+        elif isinstance(file, uuid.UUID):
+            file_id = file
+        elif isinstance(file, File):
+            file_id = file.id
+        else:
+            raise ValueError(f"Must provide a `file_id` or a File, not {file}")
+
+        self.add_tags_and_contextvars(file_id=str(file_id))
+
         logger.info(f"Creating RTUClient for file {file_id}")
         file = await self.get_file(file_id)
         if file.type != 'notebook':

tests/rtu/test_execution.py

@@ -14,7 +14,7 @@ async def test_single_cell(api_client: APIClient, notebook_maker):
     # TODO: remove sleep when Gate stops permission denied on newly created files (db time-travel)
     await asyncio.sleep(2)
 
-    rtu_client: RTUClient = await api_client.rtu_client(file.id)
+    rtu_client: RTUClient = await api_client.connect_realtime(file)
     assert rtu_client.builder.nb.cells == notebook.cells
 
     kernel_session: KernelSession = await api_client.launch_kernel(file.id)
@@ -47,7 +47,7 @@ async def test_run_all(api_client: APIClient, notebook_maker):
     # TODO: remove sleep when Gate stops permission denied on newly created files (db time-travel)
     await asyncio.sleep(2)
 
-    rtu_client: RTUClient = await api_client.rtu_client(file.id)
+    rtu_client: RTUClient = await api_client.connect_realtime(file)
     assert rtu_client.builder.nb.cells == notebook.cells
 
     kernel_session: KernelSession = await api_client.launch_kernel(file.id)

tests/rtu/test_notebook.py

@@ -13,7 +13,7 @@ async def test_add_and_remove_cell(api_client: APIClient, notebook_maker):
     file: File = await notebook_maker()
     # TODO: remove sleep when Gate stops permission denied on newly created files (db time-travel)
     await asyncio.sleep(2)
-    rtu_client: RTUClient = await api_client.rtu_client(file.id)
+    rtu_client: RTUClient = await api_client.connect_realtime(file)
     assert rtu_client.builder.nb.cells == []
 
     cell = CodeCell(id='cell_1', source='print("hello world")')