Update routers/scraping_router.py

routers/scraping_router.py

@@ -1,13 +1,12 @@
 # routers/scraping_router.py
 from fastapi import APIRouter, HTTPException, Header
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-import os
 from typing import Optional
 import logging
 import re
 from helpers.ai_client import AIClient
 
-# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -16,7 +15,6 @@ router = APIRouter(
     tags=["Web Scraping"]
 )
 
-# Initialize AI client
 ai_client = AIClient()
 
 class CheerioScriptRequest(BaseModel):
@@ -24,16 +22,73 @@ class CheerioScriptRequest(BaseModel):
     user_input: str
     api_key: Optional[str] = None
 
-@router.post("/generate-cheerio-script")
+class CheerioScriptResponse(BaseModel):
+    cheerio_script: str
+    status: str
+    message: str
+
+@router.post("/generate-cheerio-script", response_model=CheerioScriptResponse)
 async def generate_cheerio_script(
     request: CheerioScriptRequest,
     x_api_key: str = Header(None)
 ):
     try:
-        system_prompt = f"""return cheerio js script to extract {request.user_input} from the following html, 
-        enclose the cheerio script in <cheerio_script> .. </cheerio_script> format, {request.html}"""
-        user_prompt = "return valid cherio js script enclosed in <cheerio_script> </cheerio_script> format"
-        response = ""
+        example = """
+        Input HTML:
+        <div class="product-card">
+            <h2 class="title">iPhone 14</h2>
+            <span class="price">$999</span>
+        </div>
+
+        Input Request: "extract product title and price"
+
+        Expected Output:
+        <cheerio_script>
+        let result = {};
+        try {
+            const productCard = $('.product-card');
+            result = {
+                success: true,
+                data: {
+                    title: productCard.find('.title').text().trim() || null,
+                    price: productCard.find('.price').text().trim() || null
+                },
+                error: null
+            };
+        } catch (error) {
+            result = {
+                success: false,
+                data: null,
+                error: error.message
+            };
+        }
+        return result;
+        </cheerio_script>
+        """
+
+        system_prompt = f"""You are an expert at writing Cheerio.js web scraping scripts.
+        Task: Generate a Cheerio.js script to extract {request.user_input} from the provided HTML.
+        
+        Requirements:
+        - Return a dictionary/object with the structure: {{ success: boolean, data: object|null, error: string|null }}
+        - Use modern JavaScript syntax
+        - Include try-catch error handling
+        - Make the script reusable and efficient
+        - Enclose the script in <cheerio_script> tags
+        
+        Here's an example of the expected format:
+        {example}
+        
+        HTML to process:
+        {request.html}"""
+        
+        user_prompt = f"""Generate a Cheerio.js script to extract {request.user_input}.
+        The script must:
+        1. Return a dictionary/object with success, data, and error fields
+        2. Handle missing elements gracefully
+        3. Use proper Cheerio selectors
+        4. Be enclosed in <cheerio_script> tags"""
+        
         response = ai_client.chat(
             prompt=user_prompt,
             system_message=system_prompt,
@@ -41,19 +96,36 @@ async def generate_cheerio_script(
             api_key=x_api_key
         )
         
-        # Extract Cheerio script using regex
         cheerio_pattern = r'<cheerio_script>(.*?)</cheerio_script>'
         matches = re.search(cheerio_pattern, response, re.DOTALL)
         
         if matches:
             cheerio_script = matches.group(1).strip()
-            return {"cheerio_script": cheerio_script}
+            return JSONResponse(
+                status_code=200,
+                content={
+                    "cheerio_script": cheerio_script,
+                    "status": "success",
+                    "message": "Cheerio script generated successfully"
+                }
+            )
         else:
-            raise HTTPException(
+            return JSONResponse(
                 status_code=422,
-                detail=f"No valid Cheerio script found in the response: {response}"
+                content={
+                    "cheerio_script": "",
+                    "status": "error",
+                    "message": "No valid Cheerio script found in response"
+                }
             )
             
     except Exception as e:
         logger.error(f"Error generating Cheerio script: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        return JSONResponse(
+            status_code=500,
+            content={
+                "cheerio_script": "",
+                "status": "error",
+                "message": f"Error generating script: {str(e)}"
+            }
+        )