CIAgent/apps/api/app/services/enrichment_service.py

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