refac for the better

gdocs/docs_tools.py

@@ -30,7 +30,7 @@ from gdocs.docs_helpers import (
     create_bullet_list_request,
     create_insert_doc_tab_request,
     create_update_doc_tab_request,
-    create_delete_doc_tab_request
+    create_delete_doc_tab_request,
 )
 
 # Import document structure and table utilities
@@ -584,7 +584,9 @@ async def find_and_replace_doc(
         f"[find_and_replace_doc] Doc={document_id}, find='{find_text}', replace='{replace_text}', tab='{tab_id}'"
     )
 
-    requests = [create_find_replace_request(find_text, replace_text, match_case, tab_id)]
+    requests = [
+        create_find_replace_request(find_text, replace_text, match_case, tab_id)
+    ]
 
     result = await asyncio.to_thread(
         service.documents()
@@ -1075,6 +1077,7 @@ async def inspect_doc_structure(
 
     # Always include available tabs if no tab_id was specified
     if not tab_id:
+
         def get_tabs_summary(tabs):
             summary = []
             for tab in tabs:
@@ -1107,6 +1110,7 @@ async def create_table_with_data(
     table_data: List[List[str]],
     index: int,
     bold_headers: bool = True,
+    tab_id: Optional[str] = None,
 ) -> str:
     """
     Creates a table and populates it with data in one reliable operation.
@@ -1145,6 +1149,7 @@ async def create_table_with_data(
         table_data: 2D list of strings - EXACT format: [["col1", "col2"], ["row1col1", "row1col2"]]
         index: Document position (MANDATORY: get from inspect_doc_structure 'total_length')
         bold_headers: Whether to make first row bold (default: true)
+        tab_id: Optional tab ID to create the table in a specific tab
 
     Returns:
         str: Confirmation with table details and link
@@ -1171,7 +1176,7 @@ async def create_table_with_data(
 
     # Try to create the table, and if it fails due to index being at document end, retry with index-1
     success, message, metadata = await table_manager.create_and_populate_table(
-        document_id, table_data, index, bold_headers
+        document_id, table_data, index, bold_headers, tab_id
     )
 
     # If it failed due to index being at or beyond document end, retry with adjusted index
@@ -1180,7 +1185,7 @@ async def create_table_with_data(
             f"Index {index} is at document boundary, retrying with index {index - 1}"
         )
         success, message, metadata = await table_manager.create_and_populate_table(
-            document_id, table_data, index - 1, bold_headers
+            document_id, table_data, index - 1, bold_headers, tab_id
         )
 
     if success:
@@ -1786,14 +1791,23 @@ async def insert_doc_tab(
     logger.info(f"[insert_doc_tab] Doc={document_id}, title='{title}', index={index}")
 
     request = create_insert_doc_tab_request(title, index, parent_tab_id)
-    await asyncio.to_thread(
+    result = await asyncio.to_thread(
         service.documents()
         .batchUpdate(documentId=document_id, body={"requests": [request]})
         .execute
     )
 
+    # Extract the new tab ID from the batchUpdate response
+    tab_id = None
+    if "replies" in result and result["replies"]:
+        reply = result["replies"][0]
+        if "createDocumentTab" in reply:
+            tab_id = reply["createDocumentTab"].get("tabProperties", {}).get("tabId")
+
     link = f"https://docs.google.com/document/d/{document_id}/edit"
     msg = f"Inserted tab '{title}' at index {index} in document {document_id}."
+    if tab_id:
+        msg += f" Tab ID: {tab_id}."
     if parent_tab_id:
         msg += f" Nested under parent tab {parent_tab_id}."
     return f"{msg} Link: {link}"
@@ -1854,7 +1868,9 @@ async def update_doc_tab(
     Returns:
         str: Confirmation message with document link
     """
-    logger.info(f"[update_doc_tab] Doc={document_id}, tab_id='{tab_id}', title='{title}'")
+    logger.info(
+        f"[update_doc_tab] Doc={document_id}, tab_id='{tab_id}', title='{title}'"
+    )
 
     request = create_update_doc_tab_request(tab_id, title)
     await asyncio.to_thread(
@@ -1864,7 +1880,9 @@ async def update_doc_tab(
     )
 
     link = f"https://docs.google.com/document/d/{document_id}/edit"
-    return f"Renamed tab '{tab_id}' to '{title}' in document {document_id}. Link: {link}"
+    return (
+        f"Renamed tab '{tab_id}' to '{title}' in document {document_id}. Link: {link}"
+    )
 
 
 # Create comment management tools for documents

gdocs/managers/batch_operation_manager.py

@@ -90,13 +90,20 @@ class BatchOperationManager:
                 "operation_summary": operation_descriptions[:5],  # First 5 operations
             }
 
-            summary = self._build_operation_summary(operation_descriptions)
+            # Extract new tab IDs from insert_doc_tab replies
+            created_tabs = self._extract_created_tabs(result)
+            if created_tabs:
+                metadata["created_tabs"] = created_tabs
 
-            return (
+            summary = self._build_operation_summary(operation_descriptions)
-                True,
+            msg = f"Successfully executed {len(operations)} operations ({summary})"
-                f"Successfully executed {len(operations)} operations ({summary})",
+            if created_tabs:
-                metadata,
+                tab_info = ", ".join(
-            )
+                    f"'{t['title']}' (tab_id: {t['tab_id']})" for t in created_tabs
+                )
+                msg += f". Created tabs: {tab_info}"
+
+            return True, msg, metadata
 
         except Exception as e:
             logger.error(f"Failed to execute batch operations: {str(e)}")
@@ -349,6 +356,26 @@ class BatchOperationManager:
             .execute
         )
 
+    def _extract_created_tabs(self, result: dict[str, Any]) -> list[dict[str, str]]:
+        """
+        Extract tab IDs from insert_doc_tab replies in the batchUpdate response.
+
+        Args:
+            result: The batchUpdate API response
+
+        Returns:
+            List of dicts with tab_id and title for each created tab
+        """
+        created_tabs = []
+        for reply in result.get("replies", []):
+            if "createDocumentTab" in reply:
+                props = reply["createDocumentTab"].get("tabProperties", {})
+                tab_id = props.get("tabId")
+                title = props.get("title", "")
+                if tab_id:
+                    created_tabs.append({"tab_id": tab_id, "title": title})
+        return created_tabs
+
     def _build_operation_summary(self, operation_descriptions: list[str]) -> str:
         """
         Build a concise summary of operations performed.

gdocs/managers/table_operation_manager.py

@@ -41,6 +41,7 @@ class TableOperationManager:
         table_data: List[List[str]],
         index: int,
         bold_headers: bool = True,
+        tab_id: str = None,
     ) -> Tuple[bool, str, Dict[str, Any]]:
         """
         Creates a table and populates it with data in a reliable multi-step process.
@@ -52,6 +53,7 @@ class TableOperationManager:
             table_data: 2D list of strings for table content
             index: Position to insert the table
             bold_headers: Whether to make the first row bold
+            tab_id: Optional tab ID for targeting a specific tab
 
         Returns:
             Tuple of (success, message, metadata)
@@ -70,16 +72,16 @@ class TableOperationManager:
 
         try:
             # Step 1: Create empty table
-            await self._create_empty_table(document_id, index, rows, cols)
+            await self._create_empty_table(document_id, index, rows, cols, tab_id)
 
             # Step 2: Get fresh document structure to find actual cell positions
-            fresh_tables = await self._get_document_tables(document_id)
+            fresh_tables = await self._get_document_tables(document_id, tab_id)
             if not fresh_tables:
                 return False, "Could not find table after creation", {}
 
             # Step 3: Populate each cell with proper index refreshing
             population_count = await self._populate_table_cells(
-                document_id, table_data, bold_headers
+                document_id, table_data, bold_headers, tab_id
             )
 
             metadata = {
@@ -100,7 +102,7 @@ class TableOperationManager:
             return False, f"Table creation failed: {str(e)}", {}
 
     async def _create_empty_table(
-        self, document_id: str, index: int, rows: int, cols: int
+        self, document_id: str, index: int, rows: int, cols: int, tab_id: str = None
     ) -> None:
         """Create an empty table at the specified index."""
         logger.debug(f"Creating {rows}x{cols} table at index {index}")
@@ -109,20 +111,49 @@ class TableOperationManager:
             self.service.documents()
             .batchUpdate(
                 documentId=document_id,
-                body={"requests": [create_insert_table_request(index, rows, cols)]},
+                body={
+                    "requests": [create_insert_table_request(index, rows, cols, tab_id)]
+                },
             )
             .execute
         )
 
-    async def _get_document_tables(self, document_id: str) -> List[Dict[str, Any]]:
+    async def _get_document_tables(
+        self, document_id: str, tab_id: str = None
+    ) -> List[Dict[str, Any]]:
         """Get fresh document structure and extract table information."""
         doc = await asyncio.to_thread(
-            self.service.documents().get(documentId=document_id).execute
+            self.service.documents()
+            .get(documentId=document_id, includeTabsContent=True)
+            .execute
         )
+
+        if tab_id:
+            tab = self._find_tab(doc.get("tabs", []), tab_id)
+            if tab and "documentTab" in tab:
+                doc = doc.copy()
+                doc["body"] = tab["documentTab"].get("body", {})
+
         return find_tables(doc)
 
+    @staticmethod
+    def _find_tab(tabs: list, target_id: str):
+        """Recursively find a tab by ID."""
+        for tab in tabs:
+            if tab.get("tabProperties", {}).get("tabId") == target_id:
+                return tab
+            if "childTabs" in tab:
+                found = TableOperationManager._find_tab(tab["childTabs"], target_id)
+                if found:
+                    return found
+        return None
+
     async def _populate_table_cells(
-        self, document_id: str, table_data: List[List[str]], bold_headers: bool
+        self,
+        document_id: str,
+        table_data: List[List[str]],
+        bold_headers: bool,
+        tab_id: str = None,
     ) -> int:
         """
         Populate table cells with data, refreshing structure after each insertion.
@@ -147,6 +178,7 @@ class TableOperationManager:
                         col_idx,
                         cell_text,
                         bold_headers and row_idx == 0,
+                        tab_id,
                     )
 
                     if success:
@@ -169,6 +201,7 @@ class TableOperationManager:
         col_idx: int,
         cell_text: str,
         apply_bold: bool = False,
+        tab_id: str = None,
     ) -> bool:
         """
         Populate a single cell with text, with optional bold formatting.
@@ -177,7 +210,7 @@ class TableOperationManager:
         """
         try:
             # Get fresh table structure to avoid index shifting issues
-            tables = await self._get_document_tables(document_id)
+            tables = await self._get_document_tables(document_id, tab_id)
             if not tables:
                 return False