Function generate_code

backend/app/agents/tools.py:108–181

Generate Python code for the scraping pipeline.

(urls: List[str], extraction_schema: Dict[str, Any])

@tool
def generate_code(urls: List[str], extraction_schema: Dict[str, Any]) -> Dict[str, Any]:
    """Generate Python code for the scraping pipeline."""
    code_template = '''import asyncio
from scrapegraph_py import AsyncClient
from pydantic import BaseModel
from typing import List, Optional

# Define the schema
class DataSchema(BaseModel):
{schema_fields}

async def scrape_urls(api_key: str):
    """Execute web scraping for multiple URLs concurrently."""
    urls = {urls}
    
    async with AsyncClient(api_key=api_key) as client:
        tasks = []
        for url in urls:
            task = client.smartscraper(
                website_url=url,
                user_prompt="Extract data according to the schema",
                output_schema=DataSchema
            )
            tasks.append(task)
        
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Process results
        successful_results = []
        errors = []
        
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                errors.append({{"url": urls[i], "error": str(result)}})
            else:
                successful_results.append({{"url": urls[i], "data": result}})
        
        return {{"success": successful_results, "errors": errors}}

if __name__ == "__main__":
    # Replace with your actual API key
    API_KEY = "your_scrapegraph_api_key"
    
    # Run the scraper
    results = asyncio.run(scrape_urls(API_KEY))
    
    # Print results
    print(f"Successfully scraped {{len(results['success'])}} URLs")
    print(f"Failed to scrape {{len(results['errors'])}} URLs")
    
    # Print data
    for result in results['success']:
        print(f"\\nURL: {{result['url']}}")
        print(f"Data: {{result['data']}}")
'''
    
    # Generate schema fields
    schema_fields = []