Merge pull request #86 from rgbkrk/read-notebook

More notebook reading

chatlab/builtins/noteable.py

@@ -120,7 +120,7 @@ async def wait_for_kernel_idle(self):
     async def create_cell(
         self,
         source: str,
-        cell_type: Literal["code", "markdown", "sql"] = "code",
+        cell_type: Literal["code", "markdown", "sql"],
         cell_id: Optional[str] = None,
         and_run: Optional[bool] = False,
         after_cell_id: Optional[str] = None,
@@ -130,6 +130,11 @@ async def create_cell(
         """Create a code, markdown, or SQL cell."""
         rtu_client = await self.get_or_create_rtu_client()
 
+        # We could make this optional. However, the LLM being forced to provide
+        # this
+        if cell_type is None:
+            cell_type = "code"
+
         if after_cell_id is None:
             existing_cells = rtu_client.cell_ids
             if existing_cells and len(existing_cells) > 0:
@@ -458,15 +463,21 @@ async def get_cells_for_llm(self):
 
         return llm_cells
 
-    async def get_notebook(self):
+    async def read_notebook(self):
         """Get enough of the notebook to make follow up calls to get_cell"""
         rtu_client = await self.get_or_create_rtu_client()
 
         response = f"Notebook URL: {self.notebook_url}\n\n"
         response += f"Notebook ID: {self.file_id}\n\n"
         response += f"Notebook Metadata: {rtu_client.builder.nb.metadata}\n\n"
 
-        response += await self.get_cell_ids()
+        cells = await self.get_cells_for_llm()
+
+        for cell in cells:
+            response += await self.get_cell(cell, with_outputs=False)
+            response += "\n\n"
+
+        response += "NOTE: Output is not shown for cells to keep the response size to the chat model low."
 
         return response
 
@@ -498,7 +509,7 @@ async def shutdown(self):
 
     async def python(self, code: str):
         """Creates a python cell, runs it, and returns output."""
-        return await self.create_cell(code, and_run=True)
+        return await self.create_cell(code, cell_type="code", and_run=True)
 
     @property
     def chat_functions(self):
@@ -509,6 +520,7 @@ def chat_functions(self):
             self.run_cell,
             self.get_cell,
             self.get_cell_ids,
+            self.read_notebook,
             self.get_datasources,
         ]