Update app/routers/screenshot/screenshot.py

app/routers/screenshot/screenshot.py

@@ -23,9 +23,10 @@ async def screenshot(
 
     try:
         response = await playwright_instance.screenshot(
-            screenshot_model=screenshot_model, page_model=page_model,
+            screenshot_model=screenshot_model,
+            page_model=page_model,
         )
-        return Response(content=response, media_type="image/png")
+        return Response(content=response, media_type=f"image/{screenshot_model.image_type}")
     except TimeoutError as e:
         status_code = 504 if isinstance(e, TimeoutError) else 400
         detail = f"An error occurred while generating the screenshot: {e}"
@@ -33,4 +34,3 @@ async def screenshot(
     except Exception as e:
         logging.exception("Unhandled error in screenshot generation")
         raise HTTPException(status_code=500, detail="An unexpected error occurred") from e
-