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