Update app.py

app.py

@@ -36,9 +36,14 @@ app.add_middleware(
 driver_pool = []
 executor = ThreadPoolExecutor(max_workers=2)
 
-class ThreadsResponse(BaseModel):
+class MediaItem(BaseModel):
     url: str
-    media_urls: List[str]
+
+class ThreadsResponse(BaseModel):
+    post_url: str
+    media_urls: Optional[List[str]] = None
+    url: Optional[str] = None
+    picker: Optional[List[MediaItem]] = None
     media_count: int
     post_text: Optional[str] = None
     author: Optional[str] = None
@@ -276,15 +281,24 @@ async def extract_media(url: str = Query(..., description="Threads post URL")):
         loop = asyncio.get_event_loop()
         extracted_data = await loop.run_in_executor(executor, extract_media_sync, url)
         
-        return ThreadsResponse(
-            url=url,
-            media_urls=extracted_data["media_urls"],
-            media_count=len(extracted_data["media_urls"]),
-            post_text=extracted_data["post_text"],
-            author=extracted_data["author"],
-            success=True,
-            processing_time=extracted_data.get("processing_time")
-        )
+        media_urls = extracted_data["media_urls"]
+        media_count = len(media_urls)
+        
+        response_data = {
+            "post_url": url,
+            "media_count": media_count,
+            "post_text": extracted_data["post_text"],
+            "author": extracted_data["author"],
+            "success": True,
+            "processing_time": extracted_data.get("processing_time")
+        }
+        
+        if media_count == 1:
+            response_data["url"] = media_urls[0]
+        else:
+            response_data["picker"] = [{"url": url} for url in media_urls]
+            
+        return ThreadsResponse(**response_data)
         
     except HTTPException:
         raise