"""Upload management tools for listing, deleting, and resending uploaded files."""
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from .base import ToolResult
from .registry import folder_bot, get_service
if TYPE_CHECKING:
from ..bot import BotContext
from ..session_manager import SessionManager
[docs]
@dataclass
class UploadServices:
"""Services for upload tools."""
session_manager: SessionManager
uploads_dir: Path
send_document: Callable[[int, Path, str], Awaitable[None]]
chat_id: int = 0
[docs]
class ListUploadsRequest(BaseModel, frozen=True):
"""Request for listing uploaded files."""
[docs]
@folder_bot.tool(
name="list_uploads",
request_type=ListUploadsRequest,
response_type=ToolResult,
)
async def list_uploads(
request: ListUploadsRequest, _context: BotContext | None = None
) -> ToolResult:
"""List all files the user has uploaded.
Returns a list of uploaded files with their ID, original filename, size, and upload date.
"""
services: UploadServices | None = get_service(_context, "uploads")
if services is None:
return ToolResult(content="Upload services not available", is_error=True)
user_id = _context.user_id if _context else 0
uploads = services.session_manager.get_uploads(user_id)
if not uploads:
return ToolResult(content="No uploaded files.")
lines = []
for u in uploads:
lines.append(
f"[{u.id}] {u.original_filename} "
f"({u.file_size:,} bytes, {u.mime_type}, {u.created_at})"
)
return ToolResult(content="\n".join(lines))
[docs]
class DeleteUploadRequest(BaseModel, frozen=True):
"""Request for deleting an uploaded file."""
upload_id: int = Field(description="The ID of the upload to delete.")
[docs]
@folder_bot.tool(
name="delete_upload",
request_type=DeleteUploadRequest,
response_type=ToolResult,
)
async def delete_upload(
request: DeleteUploadRequest, _context: BotContext | None = None
) -> ToolResult:
"""Delete an uploaded file by its ID.
Removes the file from disk and the database record.
Use list_uploads first to find the upload ID.
"""
services: UploadServices | None = get_service(_context, "uploads")
if services is None:
return ToolResult(content="Upload services not available", is_error=True)
user_id = _context.user_id if _context else 0
record = services.session_manager.get_upload_by_id(request.upload_id, user_id)
if record is None:
return ToolResult(
content=f"Upload {request.upload_id} not found.", is_error=True
)
# Delete file from disk
file_path = services.uploads_dir / record.hash_filename
if file_path.exists():
file_path.unlink()
# Delete from database
services.session_manager.delete_upload(request.upload_id, user_id)
return ToolResult(content=f"Deleted: {record.original_filename}")
[docs]
class SendUploadRequest(BaseModel, frozen=True):
"""Request for sending an uploaded file back to the user."""
upload_id: int = Field(description="The ID of the upload to send back.")
[docs]
@folder_bot.tool(
name="send_upload",
request_type=SendUploadRequest,
response_type=ToolResult,
)
async def send_upload(
request: SendUploadRequest, _context: BotContext | None = None
) -> ToolResult:
"""Send an uploaded file back to the user in the chat.
Use list_uploads first to find the upload ID.
"""
services: UploadServices | None = get_service(_context, "uploads")
if services is None:
return ToolResult(content="Upload services not available", is_error=True)
user_id = _context.user_id if _context else 0
record = services.session_manager.get_upload_by_id(request.upload_id, user_id)
if record is None:
return ToolResult(
content=f"Upload {request.upload_id} not found.", is_error=True
)
file_path = services.uploads_dir / record.hash_filename
if not file_path.exists():
return ToolResult(
content=f"File missing from disk: {record.original_filename}", is_error=True
)
try:
await services.send_document(
services.chat_id, file_path, record.original_filename
)
except Exception as e:
return ToolResult(content=f"Error sending file: {e}", is_error=True)
return ToolResult(content=f"Sent: {record.original_filename}")