| 3 | from app.services.scrapegraph import ScrapingService |
| 4 | |
| 5 | class SimpleScrapingAgent: |
| 6 | def __init__(self): |
| 7 | self.scraping_service = ScrapingService(settings.SCRAPEGRAPH_API_KEY) |
| 8 | self.pipelines = {} |
| 9 | |
| 10 | async def process_message(self, message: str, pipeline_id: str, context: Dict = None) -> Dict: |
| 11 | """Process a user message and return the response.""" |
| 12 | |
| 13 | # Simple command parsing |
| 14 | message_lower = message.lower() |
| 15 | |
| 16 | if context is None: |
| 17 | context = {"urls": [], "schema": {}, "generated_code": ""} |
| 18 | |
| 19 | response = "I can help you build a web scraping pipeline. Try these commands:\n" |
| 20 | response += "- 'add url <url>': Add a URL to scrape\n" |
| 21 | response += "- 'list urls': Show all URLs\n" |
| 22 | response += "- 'define schema': Define data fields to extract\n" |
| 23 | response += "- 'generate code': Generate Python code\n" |
| 24 | response += "- 'run pipeline': Execute the scraping" |
| 25 | |
| 26 | # Handle add URL |
| 27 | if "add url" in message_lower: |
| 28 | url = message.split("add url", 1)[1].strip() |
| 29 | if url: |
| 30 | context["urls"].append(url) |
| 31 | response = f"✅ Added URL: {url}\nTotal URLs: {len(context['urls'])}" |
| 32 | |
| 33 | # Handle list URLs |
| 34 | elif "list urls" in message_lower: |
| 35 | if context["urls"]: |
| 36 | response = "📋 Current URLs:\n" + "\n".join(f"- {url}" for url in context["urls"]) |
| 37 | else: |
| 38 | response = "No URLs added yet. Use 'add url <url>' to add one." |
| 39 | |
| 40 | # Handle schema definition |
| 41 | elif "define schema" in message_lower or "add field" in message_lower: |
| 42 | response = "To define schema fields, use:\n" |
| 43 | response += "'add field <name> <type>' where type is: str, int, float, bool, list\n" |
| 44 | response += f"Current fields: {list(context.get('schema', {}).keys())}" |
| 45 | |
| 46 | # Handle generate code |
| 47 | elif "generate code" in message_lower: |
| 48 | if context["urls"] and context.get("schema"): |
| 49 | code = self._generate_simple_code(context["urls"], context["schema"]) |
| 50 | context["generated_code"] = code |
| 51 | response = "✅ Generated Python code for your pipeline. Check the Code tab!" |
| 52 | else: |
| 53 | response = "❌ Please add URLs and define schema fields first." |
| 54 | |
| 55 | # Handle run pipeline |
| 56 | elif "run" in message_lower or "execute" in message_lower: |
| 57 | if context["urls"] and context.get("schema"): |
| 58 | response = "🚀 Starting pipeline execution..." |
| 59 | # In real implementation, this would trigger actual scraping |
| 60 | else: |
| 61 | response = "❌ Please add URLs and define schema fields first." |
| 62 | |