Reports: the LLM reliably used company_enrichment for prose fields but inconsistently populated the parallel Finding-list/string-list fields from the same evidence, even with progressively more explicit prompting. Add a code-level backfill (products, recent developments, financial signals, strategic initiatives, regulatory signals, risks/opportunities mirrored from SWOT, unknowns, monitoring recommendations) that only ever fills in what the model left empty, never overwrites what it produced. Enrichment tab: reorder sections (Products/Recent updates before Customers/Competitors) and add a per-section "Refresh" button that re-fetches just one of NinjaPear's six independent per-company endpoints when it came back empty - confirmed live that a data-coverage gap (e.g. Amazon returning no products) is real provider behavior, not a bug. Auth: the first account registered on a deployment with zero existing admins is now auto-promoted to admin, closing the chicken-and-egg gap where the only path to admin access was direct DB access. Self-heals if the last admin ever deletes their account. Also bumps nginx's proxy_read_timeout for api.ciagent.org to cover the enrichment refresh's synchronous funding-endpoint call (up to 5 minutes per NinjaPear's docs). Co-Authored-By: Claude Sonnet 5 <[email protected]>
267 lines
10 KiB
Python
267 lines
10 KiB
Python
"""enrichment_service.enrich_company: per-section independent failure
|
|
handling, the leadership-lookup cap, and overall status derivation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
import pytest
|
|
|
|
from app.core.errors import NotFoundError
|
|
from app.enrichment.base import (
|
|
CompanyDetails,
|
|
CompanyFunding,
|
|
LeadershipMember,
|
|
Product,
|
|
)
|
|
from app.models.company import Company
|
|
from app.models.enums import EnrichmentStatus
|
|
from app.repositories.company_enrichment_repository import CompanyEnrichmentRepository
|
|
from app.services.enrichment_service import enrich_company, refresh_section
|
|
|
|
|
|
class _FakeProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
leadership: list[LeadershipMember] | None = None,
|
|
products: list[Product] | None = None,
|
|
fail: set[str] | None = None,
|
|
):
|
|
self._leadership = leadership or []
|
|
self._products = products if products is not None else []
|
|
self._fail = fail or set()
|
|
|
|
async def get_company_details(self, name, website):
|
|
if "details" in self._fail:
|
|
raise RuntimeError("details boom")
|
|
return CompanyDetails(description="A company.", leadership_team=self._leadership)
|
|
|
|
async def get_funding(self, name, website):
|
|
if "funding" in self._fail:
|
|
raise RuntimeError("funding boom")
|
|
return CompanyFunding(total_raised="$1M")
|
|
|
|
async def get_updates(self, name, website):
|
|
if "updates" in self._fail:
|
|
raise RuntimeError("updates boom")
|
|
return []
|
|
|
|
async def get_competitors(self, name, website):
|
|
if "competitors" in self._fail:
|
|
raise RuntimeError("competitors boom")
|
|
return []
|
|
|
|
async def get_products(self, name, website):
|
|
if "products" in self._fail:
|
|
raise RuntimeError("products boom")
|
|
return self._products
|
|
|
|
async def get_customers(self, name, website):
|
|
if "customers" in self._fail:
|
|
raise RuntimeError("customers boom")
|
|
return []
|
|
|
|
async def get_work_email(self, person_name, company_website):
|
|
if "work_email" in self._fail:
|
|
raise RuntimeError("email boom")
|
|
return f"{person_name.split()[0].lower()}@example.com"
|
|
|
|
async def get_person_profile(self, person_name, company_website):
|
|
if "person_profile" in self._fail:
|
|
raise RuntimeError("profile boom")
|
|
return f"https://example.com/{person_name}", "A bio."
|
|
|
|
|
|
async def _make_company(db_session, *, website: str | None = "https://acme.example.com") -> Company:
|
|
company = Company(
|
|
id=uuid.uuid4(),
|
|
user_id=uuid.uuid4(),
|
|
name="Acme Corp",
|
|
slug=f"acme-{uuid.uuid4().hex[:6]}",
|
|
official_website=website,
|
|
)
|
|
db_session.add(company)
|
|
await db_session.commit()
|
|
return company
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_all_sections_succeeding_yields_complete_status(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
provider = _FakeProvider()
|
|
|
|
enrichment = await enrich_company(db_session, settings, provider, company)
|
|
|
|
assert enrichment.status == EnrichmentStatus.COMPLETE
|
|
assert enrichment.errors == {}
|
|
assert enrichment.data["funding"]["total_raised"] == "$1M"
|
|
assert enrichment.credits_spent is not None and enrichment.credits_spent > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_failed_section_yields_partial_without_losing_the_rest(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
provider = _FakeProvider(fail={"funding"})
|
|
|
|
enrichment = await enrich_company(db_session, settings, provider, company)
|
|
|
|
assert enrichment.status == EnrichmentStatus.PARTIAL
|
|
assert "funding" in enrichment.errors
|
|
assert "funding" not in enrichment.data
|
|
assert enrichment.data["description"] == "A company." # the other sections still ran
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_every_section_failing_yields_failed_status(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
provider = _FakeProvider(
|
|
fail={"details", "funding", "updates", "competitors", "products", "customers"}
|
|
)
|
|
|
|
enrichment = await enrich_company(db_session, settings, provider, company)
|
|
|
|
assert enrichment.status == EnrichmentStatus.FAILED
|
|
assert enrichment.data.get("leadership_team") == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_leadership_lookups_are_capped(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
leadership = [LeadershipMember(name=f"Person {i}") for i in range(8)]
|
|
provider = _FakeProvider(leadership=leadership)
|
|
capped_settings = settings.model_copy(update={"ninjapear_max_leadership_lookups": 3})
|
|
|
|
enrichment = await enrich_company(db_session, capped_settings, provider, company)
|
|
|
|
team = enrichment.data["leadership_team"]
|
|
assert len(team) == 8 # every discovered leader is kept...
|
|
with_email = [m for m in team if m.get("work_email")]
|
|
assert len(with_email) == 3 # ...but only the first 3 get person-level lookups
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_website_fails_immediately_without_calling_the_provider(db_session, settings):
|
|
"""NinjaPear identifies a company by website only - every call would
|
|
fail identically, so this must short-circuit to FAILED before spending
|
|
any credits, rather than attempting (and paying for) doomed calls."""
|
|
company = await _make_company(db_session, website=None)
|
|
|
|
class _ExplodingProvider:
|
|
provider_name = "exploding"
|
|
|
|
async def get_company_details(self, *a, **k):
|
|
raise AssertionError("must not be called without a website")
|
|
|
|
enrichment = await enrich_company(db_session, settings, _ExplodingProvider(), company)
|
|
|
|
assert enrichment.status == EnrichmentStatus.FAILED
|
|
assert enrichment.credits_spent == 0
|
|
assert "details" in enrichment.errors
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upsert_updates_the_same_row_on_a_second_call(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
|
|
first = await enrich_company(db_session, settings, _FakeProvider(), company)
|
|
second = await enrich_company(db_session, settings, _FakeProvider(fail={"funding"}), company)
|
|
|
|
assert first.id == second.id
|
|
assert second.status == EnrichmentStatus.PARTIAL
|
|
|
|
stored = await CompanyEnrichmentRepository(db_session).get_for_company(company.id)
|
|
assert stored.id == first.id
|
|
assert stored.status == EnrichmentStatus.PARTIAL
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_replaces_only_that_section(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
await enrich_company(db_session, settings, _FakeProvider(), company) # products come back []
|
|
|
|
provider = _FakeProvider(products=[Product(name="Acme Pay", description="Payments product")])
|
|
enrichment = await refresh_section(db_session, settings, provider, company, "products")
|
|
|
|
assert [p["name"] for p in enrichment.data["products"]] == ["Acme Pay"]
|
|
assert enrichment.data["description"] == "A company." # untouched by the refresh
|
|
assert enrichment.status == EnrichmentStatus.COMPLETE
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_records_a_fresh_failure_without_losing_other_data(
|
|
db_session, settings
|
|
):
|
|
company = await _make_company(db_session)
|
|
await enrich_company(db_session, settings, _FakeProvider(), company)
|
|
|
|
failing_provider = _FakeProvider(fail={"products"})
|
|
enrichment = await refresh_section(db_session, settings, failing_provider, company, "products")
|
|
|
|
assert enrichment.errors["products"] == "products boom"
|
|
assert enrichment.data["description"] == "A company."
|
|
assert enrichment.status == EnrichmentStatus.PARTIAL
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_clears_a_previously_recorded_error(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
await enrich_company(db_session, settings, _FakeProvider(fail={"products"}), company)
|
|
|
|
provider = _FakeProvider(products=[Product(name="Acme Pay")])
|
|
enrichment = await refresh_section(db_session, settings, provider, company, "products")
|
|
|
|
assert "products" not in enrichment.errors
|
|
assert enrichment.status == EnrichmentStatus.COMPLETE
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_preserves_per_leader_lookups_on_a_details_refresh(
|
|
db_session, settings
|
|
):
|
|
company = await _make_company(db_session)
|
|
leadership = [LeadershipMember(name="Jane Doe")]
|
|
await enrich_company(db_session, settings, _FakeProvider(leadership=leadership), company)
|
|
|
|
stored = await CompanyEnrichmentRepository(db_session).get_for_company(company.id)
|
|
assert stored.data["leadership_team"][0]["work_email"] == "[email protected]"
|
|
|
|
# A details refresh re-fetches the roster but not the per-leader lookups -
|
|
# the previously-found work_email must survive, not silently disappear.
|
|
provider = _FakeProvider(leadership=[LeadershipMember(name="Jane Doe", title="CEO")])
|
|
enrichment = await refresh_section(db_session, settings, provider, company, "details")
|
|
|
|
member = enrichment.data["leadership_team"][0]
|
|
assert member["title"] == "CEO"
|
|
assert member["work_email"] == "[email protected]"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_rejects_an_unknown_section(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
await enrich_company(db_session, settings, _FakeProvider(), company)
|
|
|
|
with pytest.raises(ValueError):
|
|
await refresh_section(db_session, settings, _FakeProvider(), company, "not_a_section")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_requires_an_existing_enrichment_record(db_session, settings):
|
|
company = await _make_company(db_session) # enrichment was never run
|
|
|
|
with pytest.raises(NotFoundError):
|
|
await refresh_section(db_session, settings, _FakeProvider(), company, "products")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_section_adds_its_own_credit_cost(db_session, settings):
|
|
company = await _make_company(db_session)
|
|
baseline = await enrich_company(db_session, settings, _FakeProvider(), company)
|
|
baseline_credits = baseline.credits_spent # `upsert` mutates rows in place - snapshot first
|
|
|
|
enrichment = await refresh_section(db_session, settings, _FakeProvider(), company, "products")
|
|
|
|
assert enrichment.credits_spent == baseline_credits + 3 # products' own credit cost
|