handle_http_errors decorator to simplify tool error handling

This commit is contained in:
Taylor Wilsdon
2025-06-18 16:29:35 -04:00
parent b6ca27c581
commit 97edbc4ff8
9 changed files with 991 additions and 1222 deletions

View File

@@ -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("'", "\\'")
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
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", "#")
body_elements = doc_data.get('body', {}).get('content', [])
logger.info(f"[get_doc_content] File '{file_name}' (ID: {document_id}) has mimeType: '{mime_type}'")
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}")
body_text = "" # Initialize body_text
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)
# 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', [])
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)
)
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)
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:
logger.info(f"[get_doc_content] Processing as Drive file (e.g., .docx, other). MimeType: {mime_type}")
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]"
)
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'
)
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}")
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