"""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.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