Update routers/scraping_router.py

routers/scraping_router.py

@@ -1,10 +1,12 @@
 # routers/scraping_router.py
-from fastapi import APIRouter, HTTPException, Header
+from fastapi import APIRouter, HTTPException, Header, Depends
 from fastapi.responses import JSONResponse
+from fastapi.security import APIKeyHeader
 from pydantic import BaseModel
 from typing import Optional
 import logging
 import re
+import os
 from helpers.ai_client import AIClient
 
 logging.basicConfig(level=logging.INFO)
@@ -16,6 +18,18 @@ router = APIRouter(
 )
 
 ai_client = AIClient()
+AI_SCRAPER_API_KEY = os.getenv("AI_SCRAPER_API_KEY")
+
+# API Key security scheme
+api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True)
+
+async def verify_api_key(api_key: str = Depends(api_key_header)):
+    if api_key != AI_SCRAPER_API_KEY:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid API key"
+        )
+    return api_key
 
 class CheerioScriptRequest(BaseModel):
     html: str
@@ -30,7 +44,7 @@ class CheerioScriptResponse(BaseModel):
 @router.post("/generate-cheerio-script", response_model=CheerioScriptResponse)
 async def generate_cheerio_script(
     request: CheerioScriptRequest,
-    x_api_key: str = Header(None)
+    api_key: str = Depends(verify_api_key)
 ):
     try:
         example = """
@@ -91,9 +105,7 @@ async def generate_cheerio_script(
         
         response = ai_client.chat(
             prompt=user_prompt,
-            system_message=system_prompt,
-            model_id="openai/gpt-4o-mini",
-            api_key=x_api_key
+            system_message=system_prompt
         )
         
         cheerio_pattern = r'<cheerio_script>(.*?)</cheerio_script>'