"""Onboarding-time company enrichment via a paid third-party provider (NinjaPear/nubela.co) - fires exactly once per company, never on a recurring schedule (see app/tasks/enrichment.py and company_service.create_company, which gates the enqueue itself on NINJAPEAR_API_KEY being set). Orchestrates several independent, per-endpoint provider calls; one bad call must never sink the others (same principle as tasks/collection.py's per-source loop) - every failure is recorded in `errors` rather than silently dropped or allowed to fail the whole enrichment. Person-level lookups (work email, profile) are capped at `settings.ninjapear_max_leadership_lookups` to bound the fan-out from a large leadership team. """ from __future__ import annotations from datetime import UTC, datetime from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import Settings from app.core.errors import NotFoundError from app.core.logging import get_logger from app.enrichment.base import EnrichmentProvider from app.models.company import Company from app.models.company_enrichment import CompanyEnrichment from app.models.enums import EnrichmentStatus from app.repositories.company_enrichment_repository import CompanyEnrichmentRepository logger = get_logger(__name__) # Approximate credit cost per call per NinjaPear's published pricing - # tracked for user-facing cost transparency (Settings page) only, not # billed or enforced by this app. _CREDIT_COSTS = { "details": 5, # 3 base + 2 for the employee-count add-on "funding": 3, # 2 base + a rough per-investor estimate "updates": 2, "competitors": 5, # NinjaPear's documented minimum per request "products": 3, "customers": 3, # 1 base + a rough per-company estimate "work_email": 2, "person_profile": 3, } _PER_COMPANY_SECTIONS = ("details", "funding", "updates", "competitors", "products", "customers") _PER_LEADER_SECTIONS = ("work_email", "person_profile") def estimate_max_credits_per_company(max_leadership_lookups: int) -> int: """Worst-case credit ceiling for one company's enrichment - every company-level section succeeding plus every leadership-lookup slot used. Shown to the user before they commit to creating a company (see the add-company wizard's Confirm step) so the real cost is never a surprise.""" base = sum(_CREDIT_COSTS[section] for section in _PER_COMPANY_SECTIONS) per_leader = sum(_CREDIT_COSTS[section] for section in _PER_LEADER_SECTIONS) return base + per_leader * max_leadership_lookups async def enrich_company( db: AsyncSession, settings: Settings, provider: EnrichmentProvider, company: Company ) -> CompanyEnrichment: website = company.official_website if not website: # NinjaPear identifies a company by website only - every call would # fail the same way, so skip straight to FAILED instead of burning # credits on N doomed requests. repo = CompanyEnrichmentRepository(db) enrichment = await repo.upsert( company.id, status=EnrichmentStatus.FAILED, data={}, errors={"details": "No official_website on file - NinjaPear requires one"}, credits_spent=0, fetched_at=datetime.now(UTC), ) await db.commit() return enrichment data: dict = {} errors: dict[str, str] = {} credits_spent = 0 attempted = 0 succeeded = 0 async def _run(section: str, coro): nonlocal credits_spent, attempted, succeeded attempted += 1 try: result = await coro except Exception as exc: # noqa: BLE001 - one bad call must never sink the rest errors[section] = str(exc) logger.warning( "enrichment_section_failed", section=section, company_id=str(company.id), error=str(exc), ) return None credits_spent += _CREDIT_COSTS.get(section, 0) succeeded += 1 return result details = await _run("details", provider.get_company_details(company.name, website)) leadership_team: list[dict] = [] if details is not None: data["employee_count"] = details.employee_count_range data["description"] = details.description data["industry"] = details.industry data["founded_year"] = details.founded_year data["specialties"] = details.specialties leadership_team = [m.model_dump() for m in details.leadership_team] funding = await _run("funding", provider.get_funding(company.name, website)) if funding is not None: data["funding"] = funding.model_dump() updates = await _run("updates", provider.get_updates(company.name, website)) if updates is not None: data["recent_updates"] = [u.model_dump() for u in updates] competitors = await _run("competitors", provider.get_competitors(company.name, website)) if competitors is not None: data["competitors"] = [c.model_dump() for c in competitors] products = await _run("products", provider.get_products(company.name, website)) if products is not None: data["products"] = [p.model_dump() for p in products] customers = await _run("customers", provider.get_customers(company.name, website)) if customers is not None: data["customers"] = [c.model_dump() for c in customers] if leadership_team and website: cap = settings.ninjapear_max_leadership_lookups for member in leadership_team[:cap]: person_name = member.get("name") if not person_name: continue email = await _run("work_email", provider.get_work_email(person_name, website)) if email: member["work_email"] = email profile = await _run( "person_profile", provider.get_person_profile(person_name, website) ) if profile is not None: profile_url, bio = profile member["profile_url"] = member.get("profile_url") or profile_url member["bio"] = member.get("bio") or bio data["leadership_team"] = leadership_team if attempted == 0 or succeeded == 0: status = EnrichmentStatus.FAILED elif succeeded == attempted: status = EnrichmentStatus.COMPLETE else: status = EnrichmentStatus.PARTIAL repo = CompanyEnrichmentRepository(db) enrichment = await repo.upsert( company.id, status=status, data=data, errors=errors, credits_spent=credits_spent, fetched_at=datetime.now(UTC), ) await db.commit() logger.info( "company_enrichment_finished", company_id=str(company.id), status=status.value, credits_spent=credits_spent, failed_sections=list(errors.keys()), ) return enrichment # Every entry here maps 1:1 onto one of NinjaPear's independent # per-company endpoints (see app/enrichment/ninjapear.py) - refreshing one # section never re-fetches or touches any of the others. REFRESHABLE_SECTIONS = ("details", "funding", "updates", "competitors", "products", "customers") async def refresh_section( db: AsyncSession, settings: Settings, provider: EnrichmentProvider, company: Company, section: str, ) -> CompanyEnrichment: """Re-runs a single NinjaPear endpoint for a company that already has an enrichment record but came back empty or failed for just this one section (e.g. `products` returning `[]` while everything else succeeded - a real, honest gap in NinjaPear's own coverage for that company, not a bug in this app). Only ever replaces this section's own slice of `data`/`errors` - every other section's stored data is left exactly as it was.""" if section not in REFRESHABLE_SECTIONS: raise ValueError(f"Unknown enrichment section: {section}") repo = CompanyEnrichmentRepository(db) existing = await repo.get_for_company(company.id) if existing is None: raise NotFoundError("Company has no enrichment record to refresh") website = company.official_website data = dict(existing.data) errors = dict(existing.errors) credits_spent = existing.credits_spent or 0 try: if section == "details": details = await provider.get_company_details(company.name, website) data["employee_count"] = details.employee_count_range data["description"] = details.description data["industry"] = details.industry data["founded_year"] = details.founded_year data["specialties"] = details.specialties new_leadership = [m.model_dump() for m in details.leadership_team] # A details refresh re-fetches the leadership roster itself, # but not the separate per-leader work_email/person_profile # lookups - preserve those by matching on name so a refresh # never regresses contact info that was already found. old_by_name = { m.get("name"): m for m in (data.get("leadership_team") or []) if m.get("name") } for member in new_leadership: old = old_by_name.get(member.get("name")) if old: member["work_email"] = member.get("work_email") or old.get("work_email") member["profile_url"] = member.get("profile_url") or old.get("profile_url") member["bio"] = member.get("bio") or old.get("bio") data["leadership_team"] = new_leadership elif section == "funding": funding = await provider.get_funding(company.name, website) data["funding"] = funding.model_dump() elif section == "updates": updates = await provider.get_updates(company.name, website) data["recent_updates"] = [u.model_dump() for u in updates] elif section == "competitors": competitors = await provider.get_competitors(company.name, website) data["competitors"] = [c.model_dump() for c in competitors] elif section == "products": products = await provider.get_products(company.name, website) data["products"] = [p.model_dump() for p in products] elif section == "customers": customers = await provider.get_customers(company.name, website) data["customers"] = [c.model_dump() for c in customers] errors.pop(section, None) credits_spent += _CREDIT_COSTS.get(section, 0) except Exception as exc: # noqa: BLE001 - surfaced via `errors`, same as the initial run errors[section] = str(exc) logger.warning( "enrichment_section_refresh_failed", section=section, company_id=str(company.id), error=str(exc), ) new_status = EnrichmentStatus.PARTIAL if errors else EnrichmentStatus.COMPLETE enrichment = await repo.upsert( company.id, status=new_status, data=data, errors=errors, credits_spent=credits_spent, fetched_at=datetime.now(UTC), ) await db.commit() logger.info( "company_enrichment_section_refreshed", company_id=str(company.id), section=section, succeeded=section not in errors, ) return enrichment