Update main.py

main.py

@@ -4,7 +4,7 @@ import fitz  # PyMuPDF
 import logging
 import asyncio
 from uuid import uuid4
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Form
 from telegram import Bot
 from telegram.constants import ParseMode
 from starlette.middleware.cors import CORSMiddleware
@@ -39,6 +39,7 @@ app.add_middleware(
 def check_status():
     return {"status": "API is working"}
 
+# ----- Original PDF Upload Endpoint (unchanged) -----
 @app.post("/upload")
 async def upload_pdf(file: UploadFile = File(...)):
     """ Upload PDF, convert to images, send to Telegram & cleanup """
@@ -64,66 +65,65 @@ async def upload_pdf(file: UploadFile = File(...)):
 
 async def convert_pdf_to_images(pdf_path):
     """ Converts PDF pages to images & sends them to Telegram """
+    # Use a thread executor for synchronous PDF operations
+    loop = asyncio.get_running_loop()
+    doc = await loop.run_in_executor(None, fitz.open, pdf_path)
+    image_links = []
+
+    for i in range(len(doc)):
+        page = await loop.run_in_executor(None, doc.load_page, i)
+        image_path = os.path.join(IMAGES_DIR, f"{uuid4()}.png")
+        
+        # Get pixmap and save image
+        pix = await loop.run_in_executor(None, page.get_pixmap)
+        await loop.run_in_executor(None, pix.save, image_path)
+
+        # Send image to Telegram with page number caption
+        caption = f"Page {i+1}"
+        telegram_url = await send_image_to_telegram(image_path, caption)
+        image_links.append(str(telegram_url))
+        
+    return image_links
+# -------------------------------------------------------
+
+# ----- New Endpoint: Image Upload with Caption -----
+@app.post("/upload-image")
+async def upload_image(
+    file: UploadFile = File(...),
+    caption: str = Form("Uploaded image")
+):
+    """
+    Upload an image along with a caption and send it to Telegram.
+    """
+    file_id = str(uuid4())
+    extension = file.filename.split(".")[-1]
+    image_path = os.path.join(IMAGES_DIR, f"{file_id}.{extension}")
+
     try:
-        # Use a thread executor for synchronous PDF operations
-        loop = asyncio.get_running_loop()
-        doc = await loop.run_in_executor(None, fitz.open, pdf_path)
-        image_links = []
-
-        for i in range(len(doc)):
-            page = await loop.run_in_executor(None, doc.load_page, i)
-            image_path = os.path.join(IMAGES_DIR, f"{uuid4()}.png")
-            
-            # Get pixmap and save image
-            pix = await loop.run_in_executor(None, page.get_pixmap)
-            await loop.run_in_executor(None, pix.save, image_path)
-
-            # Send image to Telegram with page number caption
-            caption = f"Page {i+1}"
-            telegram_url = await send_image_to_telegram(image_path, caption)
-            
-            if telegram_url:
-                image_links.append(telegram_url)
-                os.remove(image_path)  # Cleanup only if successful
-            else:
-                logger.error(f"Failed to send page {i+1}")
-                if os.path.exists(image_path):
-                    os.remove(image_path)  # Cleanup failed image
-
-        await loop.run_in_executor(None, doc.close)
-        return image_links
+        with open(image_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
 
-    except Exception as e:
-        logger.error(f"Error converting PDF: {e}")
-        raise
-    finally:
-        # Cleanup PDF after processing
-        if os.path.exists(pdf_path):
-            os.remove(pdf_path)
+        telegram_response = await send_image_to_telegram(image_path, caption)
+        return {"status": "Image sent to Telegram", "telegram_response": str(telegram_response)}
 
-async def send_image_to_telegram(image_path, caption):
-    """ Sends an image to Telegram Bot and returns the direct image URL """
-    try:
-        with open(image_path, "rb") as img_file:
-            message = await bot.send_photo(
-                chat_id=TELEGRAM_CHAT_ID,
-                photo=img_file,
-                caption=caption
-            )
-        # Extract the direct URL of the uploaded image
-        image_url = message.photo[-1].file_id  # Get the highest quality image ID
-        
-        # Generate the direct file link
-        file_info = await bot.get_file(image_url)
-        direct_link = f"{file_info.file_path.lstrip('/')}"
-        
-        
-        return direct_link
     except Exception as e:
-        logger.error(f"Error sending image: {e}")
-        return None
+        logger.error(f"Error processing image: {e}")
+        if os.path.exists(image_path):
+            os.remove(image_path)
+        return {"status": "error", "message": str(e)}
+# -------------------------------------------------------
+
+def send_photo_sync(chat_id, image_file, caption):
+    """ Synchronous helper to send a photo via Telegram """
+    return bot.send_photo(chat_id=chat_id, photo=image_file, caption=caption, parse_mode=ParseMode.HTML)
 
+async def send_image_to_telegram(image_path, caption):
+    """ Send an image to Telegram asynchronously """
+    loop = asyncio.get_running_loop()
+    with open(image_path, "rb") as image_file:
+        response = await loop.run_in_executor(None, send_photo_sync, TELEGRAM_CHAT_ID, image_file, caption)
+    return response
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info", reload=True)
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info", reload=True)