new site
This commit is contained in:
oliver
@@ -0,0 +1,341 @@
|
||||
"""
|
||||
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]
|
||||
Reference in New Issue
Block a user