main.py

import uvicorn

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, HttpUrl
from typing import Union, Dict, Any, List

from crawler import WebCrawler
import os

# Configuration
PORT = 8001

app = FastAPI(
    title="Web Crawler API",
    description="API for crawling websites and sitemaps using Crawl4AI",
    version="1.0.0"
)

class CrawlRequest(BaseModel):
    url: HttpUrl
    is_sitemap: bool = False

class CrawlResponse(BaseModel):
    url: str
    results: Union[Dict[str, Any], List[Dict[str, Any]]]
    output_paths: Union[str, List[str]]

# Use a relative path for the output directory
crawler = WebCrawler(output_dir="crawl_results")

@app.post("/crawl", response_model=CrawlResponse)
async def crawl_url(request: CrawlRequest):
    """
    Crawl a URL or sitemap.
    
    - If is_sitemap is False, crawls the single URL
    - If is_sitemap is True, extracts URLs from the sitemap and crawls them in parallel
    """
    try:
        if request.is_sitemap:
            results = await crawler.crawl_sitemap(str(request.url))
            return CrawlResponse(
                url=str(request.url),
                results=results,
                output_paths=[r["output_path"] for r in results]
            )
        else:
            result = await crawler.crawl_url(str(request.url))
            return CrawlResponse(
                url=str(request.url),
                results=result,
                output_paths=result["output_path"]
            )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=PORT, reload=True)