import json import asyncio import functools from typing import Dict, Any, List, Optional import google.generativeai as genai from backend.browser_controller import BrowserController import base64 from bs4 import BeautifulSoup import pandas as pd from reportlab.lib.pagesizes import letter from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer from reportlab.lib.styles import getSampleStyleSheet from pathlib import Path import re MODEL = genai.GenerativeModel("gemini-2.5-flash-preview-05-20") UNIVERSAL_EXTRACTION_PROMPT = """ You are a universal data extraction specialist. Your task is to analyze any webpage and extract the most relevant information based on the user's specific goal. USER'S GOAL: {goal} CURRENT URL: {url} PAGE TITLE: {title} WEBSITE TYPE: {website_type} EXTRACTION GUIDELINES: **For PERSON/PROFILE information:** - Full name and professional title - Current position and company - Professional background and experience - Education and credentials - Skills and expertise areas - Contact information (if publicly available) - Notable achievements or projects - Social media profiles and professional links **For COMPANY/ORGANIZATION information:** - Company name and industry - Mission, vision, and description - Products or services offered - Leadership team and key personnel - Company size and locations - Contact information and headquarters - Recent news, funding, or updates - Key statistics or metrics **For PRODUCT/SERVICE information:** - Product/service name and category - Key features and specifications - Pricing information - User reviews and ratings - Availability and purchasing options - Technical requirements - Comparison with alternatives **For NEWS/CONTENT information:** - Article headline and summary - Publication date and source - Key facts and main points - Author information - Related topics or tags - Important quotes or statistics **For DATA/RESEARCH information:** - Main findings or conclusions - Statistical data and metrics - Methodology or sources - Publication details - Key insights and implications **For GENERAL INFORMATION:** - Extract the main facts relevant to the user's goal - Include supporting details and context - Provide sources and references when available - Focus on actionable or useful information IMPORTANT: - Only extract information that is VISIBLE and RELEVANT to the user's goal - Organize information in a clear, structured format - Include metadata about the source and extraction context - Be comprehensive but avoid irrelevant details - If the page doesn't contain the requested information, clearly state what was found instead WEBPAGE CONTENT: {content} Return a well-structured JSON object with the extracted information: """ class UniversalExtractor: def __init__(self): self.extraction_cache = {} async def extract_intelligent_content(self, browser: BrowserController, goal: str, fmt: str = "json", job_id: str = None) -> str: """Extract content intelligently from any website based on user's goal""" try: # Get comprehensive page information url = browser.page.url title = await browser.page.title() # Detect website type website_type = self._detect_website_type(url, title) # Get clean, structured content content = await self._get_structured_content(browser) # Use AI to extract relevant information extracted_data = await self._ai_extract(goal, url, title, website_type, content) # Format the output based on requested format return await self._format_output(extracted_data, fmt, goal, job_id) # Pass job_id except Exception as e: print(f"❌ Universal extraction failed: {e}") return await self._fallback_extraction(browser, fmt, goal) def _detect_website_type(self, url: str, title: str) -> str: """Detect website type for better extraction strategy""" url_lower = url.lower() title_lower = title.lower() # Professional networks if "linkedin.com" in url_lower: return "linkedin_profile" if "github.com" in url_lower: return "github_profile" # Social media if any(domain in url_lower for domain in ["twitter.com", "facebook.com", "instagram.com"]): return "social_media" # E-commerce if any(domain in url_lower for domain in ["amazon", "ebay", "shopify", "etsy"]): return "ecommerce" # News and content if any(word in title_lower for word in ["news", "article", "blog", "post"]): return "news_content" # Company websites if any(word in title_lower for word in ["company", "corp", "about", "careers"]): return "company_website" # Search results if "/search" in url_lower or "google.com" in url_lower: return "search_results" return "general_website" async def _get_structured_content(self, browser: BrowserController) -> str: """Get clean, structured content from the page""" try: # Get HTML content html = await browser.page.content() soup = BeautifulSoup(html, 'html.parser') # Remove script, style, and other non-content elements for tag in soup(['script', 'style', 'nav', 'footer', 'header', 'aside', 'advertisement']): tag.decompose() # Extract main content areas main_content = [] # Look for main content containers main_containers = soup.find_all(['main', 'article', 'section']) or [soup.find('body')] for container in main_containers[:3]: # Limit to avoid too much content if container: # Extract headings headings = container.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6']) for heading in headings: if heading.get_text(strip=True): main_content.append(f"HEADING: {heading.get_text(strip=True)}") # Extract paragraphs paragraphs = container.find_all('p') for p in paragraphs[:20]: # Limit paragraphs text = p.get_text(strip=True) if len(text) > 20: # Only meaningful paragraphs main_content.append(f"TEXT: {text}") # Extract lists lists = container.find_all(['ul', 'ol']) for list_elem in lists[:5]: # Limit lists items = list_elem.find_all('li') if items: main_content.append("LIST:") for item in items[:10]: # Limit list items text = item.get_text(strip=True) if text: main_content.append(f" - {text}") # Extract table data tables = container.find_all('table') for table in tables[:3]: # Limit tables rows = table.find_all('tr') if rows: main_content.append("TABLE:") for row in rows[:10]: # Limit rows cells = row.find_all(['td', 'th']) if cells: row_text = " | ".join([cell.get_text(strip=True) for cell in cells]) if row_text.strip(): main_content.append(f" {row_text}") # Join and limit content content = "\n".join(main_content) return content[:12000] # Limit total content to avoid token limits except Exception as e: print(f"❌ Error getting structured content: {e}") # Fallback to simple text extraction try: return await browser.page.inner_text("body")[:8000] except: return "Content extraction failed" async def _ai_extract(self, goal: str, url: str, title: str, website_type: str, content: str) -> Dict[str, Any]: """Use AI to extract relevant information based on context""" try: prompt = UNIVERSAL_EXTRACTION_PROMPT.format( goal=goal, url=url, title=title, website_type=website_type, content=content ) response = await asyncio.to_thread( functools.partial(MODEL.generate_content, prompt) ) # Parse AI response raw_text = response.text # Extract JSON from response start = raw_text.find('{') end = raw_text.rfind('}') + 1 if start != -1 and end > start: json_str = raw_text[start:end] extracted_data = json.loads(json_str) # Add metadata extracted_data["_metadata"] = { "source_url": url, "page_title": title, "website_type": website_type, "extraction_goal": goal, "extraction_timestamp": asyncio.get_event_loop().time(), "extraction_method": "ai_powered" } return extracted_data else: # Fallback: structure the raw text return { "extracted_content": raw_text, "content_type": "unstructured_text", "_metadata": { "source_url": url, "page_title": title, "website_type": website_type, "extraction_goal": goal, "extraction_timestamp": asyncio.get_event_loop().time(), "extraction_method": "text_fallback" } } except Exception as e: print(f"❌ AI extraction failed: {e}") return self._create_fallback_structure(content, url, title, website_type, goal) def _create_fallback_structure(self, content: str, url: str, title: str, website_type: str, goal: str) -> Dict[str, Any]: """Create structured fallback when AI extraction fails""" return { "extraction_status": "fallback_mode", "raw_content": content[:2000], # Truncated content "content_summary": self._create_simple_summary(content), "_metadata": { "source_url": url, "page_title": title, "website_type": website_type, "extraction_goal": goal, "extraction_method": "fallback_structure", "note": "AI extraction failed, using fallback method" } } def _create_simple_summary(self, content: str) -> Dict[str, Any]: """Create a simple summary of content without AI""" lines = content.split('\n') summary = { "headings": [], "key_text": [], "lists": [], "total_lines": len(lines) } current_list = [] for line in lines[:50]: # Limit processing line = line.strip() if not line: continue if line.startswith("HEADING:"): summary["headings"].append(line[8:].strip()) elif line.startswith("TEXT:"): text = line[5:].strip() if len(text) > 30: # Only substantial text summary["key_text"].append(text[:200]) elif line.startswith("LIST:"): if current_list: summary["lists"].append(current_list) current_list = [] elif line.startswith(" -"): current_list.append(line[4:].strip()) if current_list: summary["lists"].append(current_list) return summary async def _format_output(self, data: Dict[str, Any], fmt: str, goal: str, job_id: str = None) -> str: """Format extracted data in the requested format""" if fmt == "json": return json.dumps(data, indent=2, ensure_ascii=False) elif fmt == "txt": return self._format_as_text(data) elif fmt == "md": return self._format_as_markdown(data) elif fmt == "html": return self._format_as_html(data) elif fmt == "csv": return self._format_as_csv(data) elif fmt == "pdf": return await self._format_as_pdf(data, goal, job_id) # Pass job_id else: return json.dumps(data, indent=2, ensure_ascii=False) def _format_as_text(self, data: Dict[str, Any]) -> str: """Format as clean text""" lines = [] metadata = data.get("_metadata", {}) if metadata: lines.append(f"EXTRACTED INFORMATION") lines.append(f"Source: {metadata.get('source_url', 'Unknown')}") lines.append(f"Goal: {metadata.get('extraction_goal', 'Unknown')}") lines.append(f"Website Type: {metadata.get('website_type', 'Unknown')}") lines.append("-" * 60) lines.append("") def format_item(key: str, value, indent: int = 0): spaces = " " * indent if isinstance(value, dict): if key != "_metadata": lines.append(f"{spaces}{key.replace('_', ' ').title()}:") for k, v in value.items(): format_item(k, v, indent + 1) elif isinstance(value, list): lines.append(f"{spaces}{key.replace('_', ' ').title()}:") for item in value: if isinstance(item, str): lines.append(f"{spaces} • {item}") else: lines.append(f"{spaces} • {str(item)}") else: lines.append(f"{spaces}{key.replace('_', ' ').title()}: {value}") for key, value in data.items(): format_item(key, value) return "\n".join(lines) def _format_as_markdown(self, data: Dict[str, Any]) -> str: """Format as Markdown""" lines = [] metadata = data.get("_metadata", {}) if metadata: lines.append("# Extracted Information") lines.append("") lines.append(f"**Source:** {metadata.get('source_url', 'Unknown')}") lines.append(f"**Goal:** {metadata.get('extraction_goal', 'Unknown')}") lines.append(f"**Website Type:** {metadata.get('website_type', 'Unknown')}") lines.append("") lines.append("---") lines.append("") def format_item(key: str, value, level: int = 2): if isinstance(value, dict): if key != "_metadata": lines.append(f"{'#' * level} {key.replace('_', ' ').title()}") lines.append("") for k, v in value.items(): format_item(k, v, level + 1) elif isinstance(value, list): lines.append(f"{'#' * level} {key.replace('_', ' ').title()}") lines.append("") for item in value: lines.append(f"- {item}") lines.append("") else: lines.append(f"**{key.replace('_', ' ').title()}:** {value}") lines.append("") for key, value in data.items(): format_item(key, value) return "\n".join(lines) def _format_as_html(self, data: Dict[str, Any]) -> str: """Format as HTML""" html_parts = ["Extracted Information"] html_parts.append("") html_parts.append("") metadata = data.get("_metadata", {}) if metadata: html_parts.append("

Extracted Information

") html_parts.append("
") html_parts.append(f"

Source: {metadata.get('source_url', 'Unknown')}

") html_parts.append(f"

Goal: {metadata.get('extraction_goal', 'Unknown')}

") html_parts.append(f"

Website Type: {metadata.get('website_type', 'Unknown')}

") html_parts.append("
") def format_item(key: str, value, level: int = 2): if isinstance(value, dict): if key != "_metadata": html_parts.append(f"{key.replace('_', ' ').title()}") for k, v in value.items(): format_item(k, v, min(level + 1, 6)) elif isinstance(value, list): html_parts.append(f"{key.replace('_', ' ').title()}") html_parts.append("") else: html_parts.append(f"

{key.replace('_', ' ').title()}: {value}

") for key, value in data.items(): format_item(key, value) html_parts.append("") return "\n".join(html_parts) def _format_as_csv(self, data: Dict[str, Any]) -> str: """Format as CSV""" try: # Flatten the nested structure flattened = self._flatten_dict(data) # Create DataFrame df = pd.DataFrame([flattened]) return df.to_csv(index=False) except Exception as e: print(f"❌ CSV formatting failed: {e}") # Simple fallback csv_lines = ["Field,Value"] for key, value in data.items(): if key != "_metadata": clean_value = str(value).replace('"', '""').replace('\n', ' ') csv_lines.append(f'"{key}","{clean_value}"') return "\n".join(csv_lines) async def _format_as_pdf(self, data: Dict[str, Any], goal: str, job_id: str = None) -> str: """Format as PDF and return file path""" try: from reportlab.lib.pagesizes import letter from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer from reportlab.lib.styles import getSampleStyleSheet import html output_dir = Path("outputs") output_dir.mkdir(exist_ok=True) # Use job_id if provided, otherwise use timestamp if job_id: filename = f"{job_id}.pdf" else: import time timestamp = int(time.time()) filename = f"extracted_data_{timestamp}.pdf" filepath = output_dir / filename doc = SimpleDocTemplate(str(filepath), pagesize=letter, topMargin=72, bottomMargin=72) styles = getSampleStyleSheet() story = [] # Title story.append(Paragraph("Extracted Information", styles['Title'])) story.append(Spacer(1, 20)) # Metadata metadata = data.get("_metadata", {}) if metadata: story.append(Paragraph(f"Source: {html.escape(str(metadata.get('source_url', 'Unknown')))}", styles['Normal'])) story.append(Paragraph(f"Goal: {html.escape(str(metadata.get('extraction_goal', 'Unknown')))}", styles['Normal'])) story.append(Paragraph(f"Website Type: {html.escape(str(metadata.get('website_type', 'Unknown')))}", styles['Normal'])) story.append(Spacer(1, 20)) # Content with better handling def add_content(key: str, value, level: int = 0): if isinstance(value, dict): if key != "_metadata": style = styles['Heading1'] if level == 0 else styles['Heading2'] clean_key = html.escape(key.replace('_', ' ').title()) story.append(Paragraph(clean_key, style)) story.append(Spacer(1, 10)) for k, v in value.items(): add_content(k, v, level + 1) elif isinstance(value, list): clean_key = html.escape(key.replace('_', ' ').title()) story.append(Paragraph(f"{clean_key}:", styles['Normal'])) story.append(Spacer(1, 6)) for item in value: # Handle long text items and escape HTML item_str = html.escape(str(item)) if len(item_str) > 300: item_str = item_str[:300] + "..." story.append(Paragraph(f"• {item_str}", styles['Normal'])) story.append(Spacer(1, 10)) else: # Handle long text values and escape HTML clean_key = html.escape(key.replace('_', ' ').title()) value_str = html.escape(str(value)) if len(value_str) > 800: value_str = value_str[:800] + "..." story.append(Paragraph(f"{clean_key}: {value_str}", styles['Normal'])) story.append(Spacer(1, 8)) for key, value in data.items(): add_content(key, value) # Build PDF with error handling try: doc.build(story) print(f"✅ PDF successfully generated: {filepath}") return f"PDF_DIRECT_SAVE:{filepath}" # Special indicator for direct save except Exception as build_error: print(f"❌ PDF build error: {build_error}") raise build_error except ImportError: print("❌ ReportLab not installed. Installing...") import subprocess import sys try: subprocess.check_call([sys.executable, "-m", "pip", "install", "reportlab"]) # Try again after installation return await self._format_as_pdf(data, goal, job_id) except subprocess.CalledProcessError: print("❌ Failed to install ReportLab") raise ImportError("ReportLab installation failed") except Exception as e: print(f"❌ PDF generation failed: {e}") # Return error indicator instead of fallback file raise RuntimeError(f"PDF generation failed: {str(e)}") def _flatten_dict(self, d: Dict[str, Any], parent_key: str = '', sep: str = '_') -> Dict[str, Any]: """Flatten nested dictionary for CSV export""" items = [] for k, v in d.items(): new_key = f"{parent_key}{sep}{k}" if parent_key else k if isinstance(v, dict): items.extend(self._flatten_dict(v, new_key, sep=sep).items()) elif isinstance(v, list): items.append((new_key, '; '.join(map(str, v)))) else: items.append((new_key, v)) return dict(items) async def _fallback_extraction(self, browser: BrowserController, fmt: str, goal: str) -> str: """Fallback extraction when AI fails""" try: content = await browser.page.inner_text("body") url = browser.page.url title = await browser.page.title() fallback_data = { "content": content[:3000], # Truncated "source": url, "title": title, "extraction_method": "fallback", "note": "AI extraction failed, using basic text extraction" } if fmt == "json": return json.dumps(fallback_data, indent=2) elif fmt == "txt": return f"Title: {title}\nSource: {url}\n\nContent:\n{content}" else: return content except Exception as e: return f"Extraction completely failed: {str(e)}"