Update main.py

main.py

@@ -1,17 +1,60 @@
+import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing import List, Optional
 import asyncio
 import nest_asyncio
 from crawl4ai import AsyncWebCrawler
 from crawl4ai.extraction_strategy import JsonCssExtractionStrategy, LLMExtractionStrategy
 import json
 import time
-from pydantic import BaseModel, Field
+from dotenv import load_dotenv
+
+load_dotenv()  # Load environment variables from .env file
+
+app = FastAPI()
 
 nest_asyncio.apply()
 
+class CrawlerInput(BaseModel):
+    url: str = Field(..., description="URL to crawl")
+    columns: List[str] = Field(..., description="List of required columns")
+    descriptions: List[str] = Field(..., description="Descriptions for each column")
+
+class CrawlerOutput(BaseModel):
+    data: List[dict]
+
 async def simple_crawl():
     async with AsyncWebCrawler(verbose=True) as crawler:
         result = await crawler.arun(url="https://www.nbcnews.com/business")
         print(len(result.markdown))
         return result
-result = await simple_crawl()
-print(result.markdown)
+
+@app.post("/crawl", response_model=CrawlerOutput)
+async def crawl(input: CrawlerInput):
+    if len(input.columns) != len(input.descriptions):
+        raise HTTPException(status_code=400, detail="Number of columns must match number of descriptions")
+    
+    async with AsyncWebCrawler(verbose=True) as crawler:
+        result = await crawler.arun(
+            url=input.url,
+            extraction_strategy=LLMExtractionStrategy(
+                provider="openai/gpt-3.5-turbo",
+                api_token=os.getenv('OPENAI_API_KEY'),
+                extraction_type="schema",
+                verbose=True,
+                instruction=f"Extract the following information: {', '.join(input.columns)}. Descriptions: {', '.join(input.descriptions)}"
+            )
+        )
+    
+    extracted_data = json.loads(result.extracted_content)
+    return CrawlerOutput(data=extracted_data)
+
+@app.get("/test")
+async def test():
+    result = await simple_crawl()
+    return {"markdown": result.markdown}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)