workspace-mcp 1.0.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

gdocs/docs_tools.py

@@ -14,13 +14,14 @@ from googleapiclient.http import MediaIoBaseDownload
 
 # Auth & server utilities
 from auth.service_decorator import require_google_service, require_multiple_services
-from core.utils import extract_office_xml_text
+from core.utils import extract_office_xml_text, handle_http_errors
 from core.server import server
 
 logger = logging.getLogger(__name__)
 
 @server.tool()
 @require_google_service("drive", "drive_read")
+@handle_http_errors("search_docs")
 async def search_docs(
     service,
     user_google_email: str,
@@ -35,36 +36,32 @@ async def search_docs(
     """
     logger.info(f"[search_docs] Email={user_google_email}, Query='{query}'")
 
-    try:
-        escaped_query = query.replace("'", "\\'")
-
-        response = await asyncio.to_thread(
-            service.files().list(
-                q=f"name contains '{escaped_query}' and mimeType='application/vnd.google-apps.document' and trashed=false",
-                pageSize=page_size,
-                fields="files(id, name, createdTime, modifiedTime, webViewLink)"
-            ).execute
+    escaped_query = query.replace("'", "\\'")
+
+    response = await asyncio.to_thread(
+        service.files().list(
+            q=f"name contains '{escaped_query}' and mimeType='application/vnd.google-apps.document' and trashed=false",
+            pageSize=page_size,
+            fields="files(id, name, createdTime, modifiedTime, webViewLink)"
+        ).execute
+    )
+    files = response.get('files', [])
+    if not files:
+        return f"No Google Docs found matching '{query}'."
+
+    output = [f"Found {len(files)} Google Docs matching '{query}':"]
+    for f in files:
+        output.append(
+            f"- {f['name']} (ID: {f['id']}) Modified: {f.get('modifiedTime')} Link: {f.get('webViewLink')}"
         )
-        files = response.get('files', [])
-        if not files:
-            return f"No Google Docs found matching '{query}'."
-
-        output = [f"Found {len(files)} Google Docs matching '{query}':"]
-        for f in files:
-            output.append(
-                f"- {f['name']} (ID: {f['id']}) Modified: {f.get('modifiedTime')} Link: {f.get('webViewLink')}"
-            )
-        return "\n".join(output)
-
-    except HttpError as e:
-        logger.error(f"API error in search_docs: {e}", exc_info=True)
-        raise Exception(f"API error: {e}")
+    return "\n".join(output)
 
 @server.tool()
 @require_multiple_services([
     {"service_type": "drive", "scopes": "drive_read", "param_name": "drive_service"},
     {"service_type": "docs", "scopes": "docs_read", "param_name": "docs_service"}
 ])
+@handle_http_errors("get_doc_content")
 async def get_doc_content(
     drive_service,
     docs_service,
@@ -81,97 +78,87 @@ async def get_doc_content(
     """
     logger.info(f"[get_doc_content] Invoked. Document/File ID: '{document_id}' for user '{user_google_email}'")
 
-    try:
-        # Step 2: Get file metadata from Drive
-        file_metadata = await asyncio.to_thread(
-            drive_service.files().get(
-                fileId=document_id, fields="id, name, mimeType, webViewLink"
-            ).execute
+    # Step 2: Get file metadata from Drive
+    file_metadata = await asyncio.to_thread(
+        drive_service.files().get(
+            fileId=document_id, fields="id, name, mimeType, webViewLink"
+        ).execute
+    )
+    mime_type = file_metadata.get("mimeType", "")
+    file_name = file_metadata.get("name", "Unknown File")
+    web_view_link = file_metadata.get("webViewLink", "#")
+
+    logger.info(f"[get_doc_content] File '{file_name}' (ID: {document_id}) has mimeType: '{mime_type}'")
+
+    body_text = "" # Initialize body_text
+
+    # Step 3: Process based on mimeType
+    if mime_type == "application/vnd.google-apps.document":
+        logger.info(f"[get_doc_content] Processing as native Google Doc.")
+        doc_data = await asyncio.to_thread(
+            docs_service.documents().get(documentId=document_id).execute
         )
-        mime_type = file_metadata.get("mimeType", "")
-        file_name = file_metadata.get("name", "Unknown File")
-        web_view_link = file_metadata.get("webViewLink", "#")
-
-        logger.info(f"[get_doc_content] File '{file_name}' (ID: {document_id}) has mimeType: '{mime_type}'")
-
-        body_text = "" # Initialize body_text
-
-        # Step 3: Process based on mimeType
-        if mime_type == "application/vnd.google-apps.document":
-            logger.info(f"[get_doc_content] Processing as native Google Doc.")
-            doc_data = await asyncio.to_thread(
-                docs_service.documents().get(documentId=document_id).execute
-            )
-            body_elements = doc_data.get('body', {}).get('content', [])
-
-            processed_text_lines: List[str] = []
-            for element in body_elements:
-                if 'paragraph' in element:
-                    paragraph = element.get('paragraph', {})
-                    para_elements = paragraph.get('elements', [])
-                    current_line_text = ""
-                    for pe in para_elements:
-                        text_run = pe.get('textRun', {})
-                        if text_run and 'content' in text_run:
-                            current_line_text += text_run['content']
-                    if current_line_text.strip():
-                         processed_text_lines.append(current_line_text)
-            body_text = "".join(processed_text_lines)
-        else:
-            logger.info(f"[get_doc_content] Processing as Drive file (e.g., .docx, other). MimeType: {mime_type}")
-
-            export_mime_type_map = {
-                 # Example: "application/vnd.google-apps.spreadsheet"z: "text/csv",
-                 # Native GSuite types that are not Docs would go here if this function
-                 # was intended to export them. For .docx, direct download is used.
-            }
-            effective_export_mime = export_mime_type_map.get(mime_type)
-
-            request_obj = (
-                drive_service.files().export_media(fileId=document_id, mimeType=effective_export_mime)
-                if effective_export_mime
-                else drive_service.files().get_media(fileId=document_id)
-            )
-
-            fh = io.BytesIO()
-            downloader = MediaIoBaseDownload(fh, request_obj)
-            loop = asyncio.get_event_loop()
-            done = False
-            while not done:
-                status, done = await loop.run_in_executor(None, downloader.next_chunk)
-
-            file_content_bytes = fh.getvalue()
-
-            office_text = extract_office_xml_text(file_content_bytes, mime_type)
-            if office_text:
-                body_text = office_text
-            else:
-                try:
-                    body_text = file_content_bytes.decode("utf-8")
-                except UnicodeDecodeError:
-                    body_text = (
-                        f"[Binary or unsupported text encoding for mimeType '{mime_type}' - "
-                        f"{len(file_content_bytes)} bytes]"
-                    )
-
-        header = (
-            f'File: "{file_name}" (ID: {document_id}, Type: {mime_type})\n'
-            f'Link: {web_view_link}\n\n--- CONTENT ---\n'
+        body_elements = doc_data.get('body', {}).get('content', [])
+
+        processed_text_lines: List[str] = []
+        for element in body_elements:
+            if 'paragraph' in element:
+                paragraph = element.get('paragraph', {})
+                para_elements = paragraph.get('elements', [])
+                current_line_text = ""
+                for pe in para_elements:
+                    text_run = pe.get('textRun', {})
+                    if text_run and 'content' in text_run:
+                        current_line_text += text_run['content']
+                if current_line_text.strip():
+                        processed_text_lines.append(current_line_text)
+        body_text = "".join(processed_text_lines)
+    else:
+        logger.info(f"[get_doc_content] Processing as Drive file (e.g., .docx, other). MimeType: {mime_type}")
+
+        export_mime_type_map = {
+                # Example: "application/vnd.google-apps.spreadsheet"z: "text/csv",
+                # Native GSuite types that are not Docs would go here if this function
+                # was intended to export them. For .docx, direct download is used.
+        }
+        effective_export_mime = export_mime_type_map.get(mime_type)
+
+        request_obj = (
+            drive_service.files().export_media(fileId=document_id, mimeType=effective_export_mime)
+            if effective_export_mime
+            else drive_service.files().get_media(fileId=document_id)
         )
-        return header + body_text
 
-    except HttpError as error:
-        logger.error(
-            f"[get_doc_content] API error for ID {document_id}: {error}",
-            exc_info=True,
-        )
-        raise Exception(f"API error processing document/file ID {document_id}: {error}")
-    except Exception as e:
-        logger.exception(f"[get_doc_content] Unexpected error for ID {document_id}: {e}")
-        raise Exception(f"Unexpected error processing document/file ID {document_id}: {e}")
+        fh = io.BytesIO()
+        downloader = MediaIoBaseDownload(fh, request_obj)
+        loop = asyncio.get_event_loop()
+        done = False
+        while not done:
+            status, done = await loop.run_in_executor(None, downloader.next_chunk)
+
+        file_content_bytes = fh.getvalue()
+
+        office_text = extract_office_xml_text(file_content_bytes, mime_type)
+        if office_text:
+            body_text = office_text
+        else:
+            try:
+                body_text = file_content_bytes.decode("utf-8")
+            except UnicodeDecodeError:
+                body_text = (
+                    f"[Binary or unsupported text encoding for mimeType '{mime_type}' - "
+                    f"{len(file_content_bytes)} bytes]"
+                )
+
+    header = (
+        f'File: "{file_name}" (ID: {document_id}, Type: {mime_type})\n'
+        f'Link: {web_view_link}\n\n--- CONTENT ---\n'
+    )
+    return header + body_text
 
 @server.tool()
 @require_google_service("drive", "drive_read")
+@handle_http_errors("list_docs_in_folder")
 async def list_docs_in_folder(
     service,
     user_google_email: str,
@@ -186,34 +173,27 @@ async def list_docs_in_folder(
     """
     logger.info(f"[list_docs_in_folder] Invoked. Email: '{user_google_email}', Folder ID: '{folder_id}'")
 
-    try:
-        rsp = await asyncio.to_thread(
-            service.files().list(
-                q=f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.document' and trashed=false",
-                pageSize=page_size,
-                fields="files(id, name, modifiedTime, webViewLink)"
-            ).execute
-        )
-        items = rsp.get('files', [])
-        if not items:
-            return f"No Google Docs found in folder '{folder_id}'."
-        out = [f"Found {len(items)} Docs in folder '{folder_id}':"]
-        for f in items:
-            out.append(f"- {f['name']} (ID: {f['id']}) Modified: {f.get('modifiedTime')} Link: {f.get('webViewLink')}")
-        return "\n".join(out)
-
-    except HttpError as e:
-        logger.error(f"API error in list_docs_in_folder: {e}", exc_info=True)
-        raise Exception(f"API error: {e}")
-    except Exception as e:
-        logger.exception(f"Unexpected error in list_docs_in_folder: {e}")
-        raise Exception(f"Unexpected error: {e}")
+    rsp = await asyncio.to_thread(
+        service.files().list(
+            q=f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.document' and trashed=false",
+            pageSize=page_size,
+            fields="files(id, name, modifiedTime, webViewLink)"
+        ).execute
+    )
+    items = rsp.get('files', [])
+    if not items:
+        return f"No Google Docs found in folder '{folder_id}'."
+    out = [f"Found {len(items)} Docs in folder '{folder_id}':"]
+    for f in items:
+        out.append(f"- {f['name']} (ID: {f['id']}) Modified: {f.get('modifiedTime')} Link: {f.get('webViewLink')}")
+    return "\n".join(out)
 
 @server.tool()
 @require_google_service("docs", "docs_write")
+@handle_http_errors("create_doc")
 async def create_doc(
     service,
-    user_google_email: str, # Made user_google_email required
+    user_google_email: str,
     title: str,
     content: str = '',
 ) -> str:
@@ -225,20 +205,12 @@ async def create_doc(
     """
     logger.info(f"[create_doc] Invoked. Email: '{user_google_email}', Title='{title}'")
 
-    try:
-        doc = await asyncio.to_thread(service.documents().create(body={'title': title}).execute)
-        doc_id = doc.get('documentId')
-        if content:
-            requests = [{'insertText': {'location': {'index': 1}, 'text': content}}]
-            await asyncio.to_thread(service.documents().batchUpdate(documentId=doc_id, body={'requests': requests}).execute)
-        link = f"https://docs.google.com/document/d/{doc_id}/edit"
-        msg = f"Created Google Doc '{title}' (ID: {doc_id}) for {user_google_email}. Link: {link}"
-        logger.info(f"Successfully created Google Doc '{title}' (ID: {doc_id}) for {user_google_email}. Link: {link}")
-        return msg
-
-    except HttpError as e:
-        logger.error(f"API error in create_doc: {e}", exc_info=True)
-        raise Exception(f"API error: {e}")
-    except Exception as e:
-        logger.exception(f"Unexpected error in create_doc: {e}")
-        raise Exception(f"Unexpected error: {e}")
+    doc = await asyncio.to_thread(service.documents().create(body={'title': title}).execute)
+    doc_id = doc.get('documentId')
+    if content:
+        requests = [{'insertText': {'location': {'index': 1}, 'text': content}}]
+        await asyncio.to_thread(service.documents().batchUpdate(documentId=doc_id, body={'requests': requests}).execute)
+    link = f"https://docs.google.com/document/d/{doc_id}/edit"
+    msg = f"Created Google Doc '{title}' (ID: {doc_id}) for {user_google_email}. Link: {link}"
+    logger.info(f"Successfully created Google Doc '{title}' (ID: {doc_id}) for {user_google_email}. Link: {link}")
+    return msg