342 lines
11 KiB
Python
342 lines
11 KiB
Python
"""
|
|
Execute command for running pipeline steps.
|
|
"""
|
|
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Optional, List
|
|
from uuid import UUID
|
|
import re
|
|
|
|
import typer
|
|
from rich.console import Console
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
|
|
|
|
from ..services.supabase import SupabaseService
|
|
from ..services.apify import ApifyService
|
|
from ..services.firecrawl import FirecrawlService
|
|
from ..services.haiku import HaikuService
|
|
from ..services.images import ImageService
|
|
from ..services.tags import TagAggregationService
|
|
from ..models.schemas import Company, StepStatus, TagCategory
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
|
|
def execute_command(
|
|
plan_id: str,
|
|
step: Optional[str] = None,
|
|
city: Optional[str] = None,
|
|
dry_run: bool = False,
|
|
) -> None:
|
|
"""
|
|
Execute an execution plan or specific step.
|
|
|
|
Runs the pipeline to scrape, enrich, and generate content for
|
|
Odoo partners.
|
|
|
|
Args:
|
|
plan_id: UUID of the execution plan
|
|
step: Specific step to run (scrape, enrich, content, export)
|
|
city: Specific city slug to process
|
|
dry_run: If True, don't make actual API calls
|
|
"""
|
|
console.print(f"\n[bold blue]Executing Plan: {plan_id}[/bold blue]\n")
|
|
|
|
db = SupabaseService()
|
|
|
|
try:
|
|
plan_uuid = UUID(plan_id)
|
|
except ValueError:
|
|
console.print(f"[red]Invalid plan ID: {plan_id}[/red]")
|
|
raise typer.Exit(1)
|
|
|
|
plan = db.get_execution_plan(plan_uuid)
|
|
if not plan:
|
|
console.print(f"[red]Plan not found: {plan_id}[/red]")
|
|
raise typer.Exit(1)
|
|
|
|
# Get country info
|
|
country = db.get_country_by_code(plan.country_code)
|
|
if not country:
|
|
console.print(f"[red]Country not found: {plan.country_code}[/red]")
|
|
raise typer.Exit(1)
|
|
|
|
# Filter cities if specified
|
|
city_slugs = plan.cities
|
|
if city:
|
|
if city not in city_slugs:
|
|
console.print(f"[red]City {city} not in plan[/red]")
|
|
raise typer.Exit(1)
|
|
city_slugs = [city]
|
|
|
|
# Determine steps to run
|
|
if step:
|
|
steps_to_run = [step]
|
|
else:
|
|
steps_to_run = ["scrape", "enrich", "content"]
|
|
|
|
console.print(f"Country: {country.name}")
|
|
console.print(f"Cities: {len(city_slugs)}")
|
|
console.print(f"Steps: {', '.join(steps_to_run)}")
|
|
console.print(f"Dry run: {dry_run}")
|
|
console.print()
|
|
|
|
# Update plan status
|
|
if not dry_run:
|
|
db.update_execution_plan(
|
|
plan_uuid, {"status": "running", "started_at": datetime.utcnow().isoformat()}
|
|
)
|
|
|
|
# Run each step
|
|
for step_name in steps_to_run:
|
|
console.print(f"\n[bold cyan]Running step: {step_name}[/bold cyan]\n")
|
|
|
|
for city_slug in city_slugs:
|
|
city_obj = db.get_city_by_slug(city_slug, country.id)
|
|
if not city_obj:
|
|
console.print(f"[yellow]City not found: {city_slug}[/yellow]")
|
|
continue
|
|
|
|
try:
|
|
if step_name == "scrape":
|
|
_run_scrape_step(db, city_obj, country, dry_run)
|
|
elif step_name == "enrich":
|
|
_run_enrich_step(db, city_obj, dry_run)
|
|
elif step_name == "content":
|
|
_run_content_step(db, city_obj, country, dry_run)
|
|
except Exception as e:
|
|
console.print(f"[red]Error in {step_name} for {city_slug}: {e}[/red]")
|
|
logger.exception(f"Step {step_name} failed for {city_slug}")
|
|
|
|
# Update plan status
|
|
if not dry_run:
|
|
db.update_execution_plan(
|
|
plan_uuid, {"status": "completed", "completed_at": datetime.utcnow().isoformat()}
|
|
)
|
|
|
|
console.print("\n[bold green]Execution complete![/bold green]")
|
|
|
|
|
|
def _run_scrape_step(db: SupabaseService, city, country, dry_run: bool) -> None:
|
|
"""Run the scrape step for a city."""
|
|
console.print(f"Scraping {city.name}...")
|
|
|
|
if dry_run:
|
|
console.print("[yellow]Dry run - skipping actual scrape[/yellow]")
|
|
return
|
|
|
|
apify = ApifyService()
|
|
|
|
# Scrape Google Maps
|
|
results = apify.scrape_city(
|
|
city_name=city.name,
|
|
language=country.language,
|
|
country=country.code,
|
|
)
|
|
|
|
# Deduplicate
|
|
results = apify.deduplicate_results(results)
|
|
console.print(f"Found {len(results)} unique companies")
|
|
|
|
# Save to database
|
|
saved = 0
|
|
for result in results:
|
|
# Check if already exists
|
|
if result.placeId:
|
|
existing = db.get_company_by_place_id(result.placeId)
|
|
if existing:
|
|
continue
|
|
|
|
# Create slug from name
|
|
slug = _create_slug(result.title)
|
|
|
|
# Extract location
|
|
lat = lng = None
|
|
if result.location:
|
|
lat = result.location.get("lat")
|
|
lng = result.location.get("lng")
|
|
|
|
company = Company(
|
|
city_id=city.id,
|
|
name=result.title,
|
|
slug=slug,
|
|
google_place_id=result.placeId,
|
|
address=result.address,
|
|
phone=result.phone,
|
|
website=result.website,
|
|
rating=result.rating,
|
|
review_count=result.reviewsCount,
|
|
latitude=lat,
|
|
longitude=lng,
|
|
status="scraped",
|
|
last_scraped_at=datetime.utcnow(),
|
|
)
|
|
|
|
db.create_company(company)
|
|
saved += 1
|
|
|
|
console.print(f"[green]Saved {saved} new companies[/green]")
|
|
|
|
|
|
def _run_enrich_step(db: SupabaseService, city, dry_run: bool) -> None:
|
|
"""Run the enrich step for companies in a city."""
|
|
console.print(f"Enriching companies in {city.name}...")
|
|
|
|
# Get companies that need enrichment
|
|
companies = db.list_companies(city_id=city.id, status="scraped", limit=100)
|
|
console.print(f"Found {len(companies)} companies to enrich")
|
|
|
|
if not companies or dry_run:
|
|
if dry_run:
|
|
console.print("[yellow]Dry run - skipping enrichment[/yellow]")
|
|
return
|
|
|
|
firecrawl = FirecrawlService()
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task("Enriching...", total=len(companies))
|
|
|
|
for company in companies:
|
|
if not company.website:
|
|
progress.advance(task)
|
|
continue
|
|
|
|
try:
|
|
info = firecrawl.extract_company_info(company.website)
|
|
|
|
updates = {
|
|
"website_markdown": info.get("markdown"),
|
|
"logo_url": info.get("logo_url"),
|
|
"email": info.get("email") or company.email,
|
|
"status": "enriched",
|
|
"last_enriched_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
db.update_company(company.id, updates)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to enrich {company.name}: {e}")
|
|
|
|
progress.advance(task)
|
|
|
|
console.print("[green]Enrichment complete[/green]")
|
|
|
|
|
|
def _run_content_step(db: SupabaseService, city, country, dry_run: bool) -> None:
|
|
"""Run the content generation step."""
|
|
console.print(f"Generating content for {city.name}...")
|
|
|
|
# Get companies that need content
|
|
companies = db.list_companies(city_id=city.id, status="enriched", limit=100)
|
|
console.print(f"Found {len(companies)} companies for content generation")
|
|
|
|
if not companies or dry_run:
|
|
if dry_run:
|
|
console.print("[yellow]Dry run - skipping content generation[/yellow]")
|
|
return
|
|
|
|
haiku = HaikuService()
|
|
tags_service = TagAggregationService(db)
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task("Generating...", total=len(companies))
|
|
|
|
for company in companies:
|
|
try:
|
|
# Extract tags
|
|
tags_response = haiku.extract_tags(
|
|
company_name=company.name,
|
|
website_markdown=company.website_markdown,
|
|
address=company.address,
|
|
)
|
|
|
|
# Save tags
|
|
_save_company_tags(db, company.id, tags_response)
|
|
|
|
# Generate content
|
|
content = haiku.generate_company_content(
|
|
company_name=company.name,
|
|
city_name=city.name,
|
|
country_name=country.name,
|
|
website_markdown=company.website_markdown,
|
|
services=tags_response.services,
|
|
industries=tags_response.industries,
|
|
modules=tags_response.modules,
|
|
partner_level=tags_response.partner_level,
|
|
language=country.language,
|
|
)
|
|
|
|
updates = {
|
|
"description": content.description,
|
|
"content_html": content.content_html,
|
|
"meta_title": content.meta_title,
|
|
"meta_description": content.meta_description,
|
|
"status": "content_generated",
|
|
"last_content_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
db.update_company(company.id, updates)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to generate content for {company.name}: {e}")
|
|
|
|
progress.advance(task)
|
|
|
|
# Aggregate tags for city
|
|
console.print("Aggregating city tags...")
|
|
tags_service.aggregate_city_tags(city.id)
|
|
|
|
console.print("[green]Content generation complete[/green]")
|
|
|
|
|
|
def _save_company_tags(db: SupabaseService, company_id: UUID, tags_response) -> None:
|
|
"""Save extracted tags to company."""
|
|
# Clear existing tags
|
|
db.remove_company_tags(company_id)
|
|
|
|
# Add service tags
|
|
for name in tags_response.services:
|
|
tag = db.get_or_create_tag(name, TagCategory.SERVICE)
|
|
db.add_company_tag(company_id, tag.id, tags_response.confidence)
|
|
|
|
# Add industry tags
|
|
for name in tags_response.industries:
|
|
tag = db.get_or_create_tag(name, TagCategory.INDUSTRY)
|
|
db.add_company_tag(company_id, tag.id, tags_response.confidence)
|
|
|
|
# Add module tags
|
|
for name in tags_response.modules:
|
|
tag = db.get_or_create_tag(name, TagCategory.MODULE)
|
|
db.add_company_tag(company_id, tag.id, tags_response.confidence)
|
|
|
|
# Add partner level if present
|
|
if tags_response.partner_level:
|
|
tag = db.get_or_create_tag(tags_response.partner_level, TagCategory.PARTNER_LEVEL)
|
|
db.add_company_tag(company_id, tag.id, tags_response.confidence)
|
|
|
|
|
|
def _create_slug(name: str) -> str:
|
|
"""Create URL-friendly slug from company name."""
|
|
# Convert to lowercase
|
|
slug = name.lower()
|
|
# Replace spaces and special chars with hyphens
|
|
slug = re.sub(r"[^a-z0-9]+", "-", slug)
|
|
# Remove leading/trailing hyphens
|
|
slug = slug.strip("-")
|
|
# Limit length
|
|
return slug[:100]
|