tools/wordpress_media.py

import requests
from requests.auth import HTTPBasicAuth
from typing import Dict, Optional
import os
import mimetypes
from smolagents.tools import Tool


class WordPressMediaTool(Tool):
    name = "wordpress_media"
    description = "Uploads media files to WordPress and returns the attachment URL"
    inputs = {
        'file_path': {'type': 'string', 'description': 'Local path to the media file to upload'},
        'title': {
            'type': 'string',
            'description': 'Title for the media attachment',
            'nullable': True
        },
        'alt_text': {
            'type': 'string',
            'description': 'Alt text for the media attachment',
            'nullable': True
        }
    }
    output_type = "object"

    def __init__(self, wordpress_credentials):
        super().__init__()
        self.credentials = wordpress_credentials

    def forward(self, file_path: str, title: str = None, alt_text: str = None) -> Dict:
        """Uploads a media file to WordPress
        Args:
            file_path: Path to the media file
            title: Optional title for the media attachment
            alt_text: Optional alt text for the media attachment
        Returns:
            Dict containing the response from WordPress with media details
        """
        if not self.credentials:
            return {"error": "WordPress credentials not configured"}

        if not os.path.exists(file_path):
            return {"error": f"File not found: {file_path}"}

        api_url = f"{self.credentials['site_url']}/wp-json/wp/v2/media"
        auth = HTTPBasicAuth(
            self.credentials['username'], self.credentials['app_password'])

        # Determine mime type
        mime_type, _ = mimetypes.guess_type(file_path)
        if not mime_type:
            return {"error": "Could not determine file type"}

        # Prepare headers
        headers = {
            'Content-Type': mime_type,
            'Content-Disposition': f'attachment; filename="{os.path.basename(file_path)}"'
        }

        try:
            # Upload file
            with open(file_path, 'rb') as file:
                response = requests.post(
                    api_url,
                    headers=headers,
                    auth=auth,
                    data=file
                )
                response.raise_for_status()
                response_data = response.json()

            # Update media properties if title or alt_text provided
            if title or alt_text:
                media_id = response_data['id']
                update_url = f"{api_url}/{media_id}"
                update_data = {}

                if title:
                    update_data['title'] = {'rendered': title}
                if alt_text:
                    update_data['alt_text'] = alt_text

                if update_data:
                    update_response = requests.post(
                        update_url,
                        json=update_data,
                        auth=auth
                    )
                    update_response.raise_for_status()
                    response_data = update_response.json()

            return {
                "status": "success",
                "attachment_id": response_data.get("id"),
                "url": response_data.get("source_url"),
                "title": response_data.get("title", {}).get("rendered"),
                "alt_text": response_data.get("alt_text")
            }

        except requests.exceptions.RequestException as e:
            print(f"Failed to upload media: {str(e)}")
            return {"error": f"Failed to upload media: {str(e)}"}