""" Entity resolution: raw company name → ResolvedEntity. Flow: Registry Lookup → (Companies House enrichment if UK jurisdiction) → (SEC EDGAR enrichment if US / unknown) → SQLite cache (7-day TTL) Registry Lookup is geo-blocked outside the US. When blocked the resolver falls back to EDGAR for public companies and Companies House for UK ones. """ import base64 import hashlib import logging import re from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Optional import aiosqlite import httpx from src.config import settings from src.models.entities import ResolvedEntity logger = logging.getLogger(__name__) _EDGAR_UA = f"DueDiligencePlatform/1.0 {settings.contact_email}" _CORP_SUFFIX_RE = re.compile( r"[\s,]*(Inc\.?|LLC\.?|Ltd\.?|Corp\.?|Co\.?|L\.?P\.?|LLP|PLC" r"|Limited|Incorporated|Corporation|Holdings|Group|S\.?A\.?)[\s,]*$", re.IGNORECASE, ) # ── Exceptions ──────────────────────────────────────────────────────────────── class EntityNotFoundError(Exception): pass class EntityAmbiguityError(Exception): """Raised when >3 Registry Lookup matches exist; caller must let user pick.""" def __init__(self, candidates: list[ResolvedEntity]) -> None: self.candidates = candidates super().__init__(f"{len(candidates)} candidates — user must select one") # ── Alias helpers ────────────────────────────────────────────────────────────── def _strip_suffixes(name: str) -> str: return _CORP_SUFFIX_RE.sub("", name).strip().strip(",").strip() def _generate_aliases(canonical: str, raw_input: str, tickers: list[str]) -> list[str]: candidates = [ canonical, raw_input.strip(), _strip_suffixes(canonical), _strip_suffixes(raw_input), ] + tickers seen: set[str] = set() aliases: list[str] = [] for name in candidates: n = name.strip() if n and n not in seen: seen.add(n) aliases.append(n) return aliases # ── Auth helpers ────────────────────────────────────────────────────────────── def _ch_auth_header() -> str: token = base64.b64encode(f"{settings.companies_house_api_key}:".encode()).decode() return f"Basic {token}" # ── Cache key ───────────────────────────────────────────────────────────────── def _cache_key(raw_input: str) -> str: return hashlib.sha256(raw_input.strip().casefold().encode()).hexdigest() # ── SQLite cache operations ─────────────────────────────────────────────────── async def _open_db() -> aiosqlite.Connection: db_path = Path(settings.database_path) db_path.parent.mkdir(parents=True, exist_ok=True) db = await aiosqlite.connect(str(db_path)) await db.execute(""" CREATE TABLE IF NOT EXISTS entity_cache ( cache_key TEXT PRIMARY KEY, entity_json TEXT NOT NULL, expires_at TEXT NOT NULL ) """) await db.commit() return db async def _get_cached(key: str) -> Optional[ResolvedEntity]: try: db = await _open_db() async with db.execute( "SELECT entity_json FROM entity_cache WHERE cache_key = ? AND expires_at > ?", (key, datetime.now(timezone.utc).isoformat()), ) as cur: row = await cur.fetchone() await db.close() if row: return ResolvedEntity.model_validate_json(row[0]) except Exception as exc: logger.warning("Entity cache read failed: %s", exc) return None async def _set_cached(key: str, entity: ResolvedEntity) -> None: try: db = await _open_db() expires_at = ( datetime.now(timezone.utc) + timedelta(hours=settings.cache_ttl_hours) ).isoformat() await db.execute( "INSERT OR REPLACE INTO entity_cache (cache_key, entity_json, expires_at) " "VALUES (?, ?, ?)", (key, entity.model_dump_json(), expires_at), ) await db.commit() await db.close() except Exception as exc: logger.warning("Entity cache write failed: %s", exc) # ── Registry Lookup ─────────────────────────────────────────────────────────── async def _rl_search( client: httpx.AsyncClient, query: str, jurisdiction: str = "" ) -> tuple[list[dict], bool]: """Query Registry Lookup. Returns (results, is_blocked). ``jurisdiction`` (e.g. 'us', 'gb', 'us-de') narrows the search server-side — the single most effective disambiguator for common names. """ params: dict = {"q": query, "limit": 10} if jurisdiction: params["jurisdiction_code"] = jurisdiction try: r = await client.get( "https://api.registry-lookup.com/v1/companies/search", params=params, headers={"X-API-Key": settings.registry_lookup_api_key, "Accept": "application/json"}, timeout=10, ) if r.status_code == 200: data = r.json() results = ( data.get("data") or data.get("companies") or data.get("results") or (data if isinstance(data, list) else []) ) return results, False if r.status_code in (403, 520, 521): logger.info("Registry Lookup geo-blocked — falling back to EDGAR / Companies House") return [], True logger.warning("Registry Lookup HTTP %s", r.status_code) return [], False except Exception as exc: logger.info("Registry Lookup unreachable (%s) — using fallback sources", exc) return [], True def _parse_rl_result(result: dict, raw_input: str) -> ResolvedEntity: canonical = ( result.get("legal_name") or result.get("name") or result.get("company_name") or raw_input.upper() ) jurisdiction = ( result.get("jurisdiction_code") or result.get("country_code") or result.get("jurisdiction") or result.get("country") or "" ).lower().replace("_", "-") company_type = (result.get("company_type") or "").lower() is_public = False if jurisdiction.startswith("us"): # Check if company type suggests public filing if "public" in company_type or "stock" in company_type or "plc" in company_type: is_public = True elif "plc" in company_type or "public limited" in company_type: is_public = True return ResolvedEntity( canonical_name=canonical, aliases=_generate_aliases(canonical, raw_input, []), jurisdiction=jurisdiction or None, is_public=is_public, registry_lookup_id=str(result.get("id", "")) or None, industry=result.get("industry") or result.get("sic_description"), ) def _flatten_str(value) -> str: """Coerce a registry field to a display string. Registry Lookup returns some fields (e.g. registered_address) as nested objects; the picker only needs a short string, so join dict values. """ if value is None: return "" if isinstance(value, str): return value if isinstance(value, dict): return ", ".join(_flatten_str(v) for v in value.values() if v) if isinstance(value, (list, tuple)): return ", ".join(_flatten_str(v) for v in value if v) return str(value) def _candidate_view(result: dict, raw_input: str) -> dict: """Compact, display-ready candidate for the interactive picker (design §8c).""" parsed = _parse_rl_result(result, raw_input) return { "registry_id": parsed.registry_lookup_id or "", "name": parsed.canonical_name, "jurisdiction": parsed.jurisdiction or "", "status": _flatten_str( result.get("status") or result.get("company_status") or result.get("current_status") ), "company_type": _flatten_str(result.get("company_type") or result.get("type")), "address": _flatten_str( result.get("registered_address") or result.get("address") or result.get("registered_office_address") )[:160], "is_public": parsed.is_public, } # ── SEC EDGAR ───────────────────────────────────────────────────────────────── async def _edgar_lookup(client: httpx.AsyncClient, company_name: str) -> dict: """Returns {'cik', 'canonical_name', 'tickers', 'state'} or {}. Uses company_tickers.json (fuzzy name→CIK) + submissions/{CIK}.json. The old CGI browse-edgar endpoint is Akamai-blocked; this replaces it. """ from rapidfuzz import process, fuzz headers = {"User-Agent": _EDGAR_UA} try: # Step 1: download the SEC company tickers index (name→CIK map) r = await client.get( "https://www.sec.gov/files/company_tickers.json", headers=headers, timeout=15, ) if r.status_code != 200: logger.warning("EDGAR company_tickers.json HTTP %s", r.status_code) return {} tickers_data = r.json() # {str: {cik_str, ticker, title}} def _norm(s: str) -> str: # Replace punctuation with spaces (not delete) so hyphenated names like # "Colgate-Palmolive" tokenize as two words instead of merging into # "COLGATEPALMOLIVE" (which would fail token_sort_ratio). Then collapse # whitespace. "Tesla, Inc." → "TESLA INC"; "Colgate-Palmolive Co" → # "COLGATE PALMOLIVE CO". return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", s)).strip().upper() # Build {cik_str → normalized_title} for fuzzy matching choices: dict[str, str] = { str(v["cik_str"]): _norm(v["title"]) for v in tickers_data.values() } # Step 2: fuzzy-match the query against normalized company titles match = process.extractOne( _norm(company_name), choices, scorer=fuzz.token_sort_ratio, score_cutoff=70, ) if not match: logger.info("EDGAR: no fuzzy match for '%s' (score < 70)", company_name) return {} _matched_title, _score, matched_cik = match # extractOne over a dict → (value, score, key) # Guard against near-miss fuzzy collisions where a private company is # mis-bound to a similarly-named public filer (e.g. "Stripe, Inc."→ # "Stride, Inc.", "Mercury Corp"→"Marcus Corp"). token_sort_ratio rates # these ≥70 because only one character differs, so a score cutoff can't # separate them from real hits (Boeing→"Boeing Co" scores the same). Require # the first significant token of the query and the matched filer to agree. q_tokens = _norm(company_name).split() m_tokens = (_matched_title or "").split() # choices values are already normalized if q_tokens and m_tokens and q_tokens[0] != m_tokens[0]: logger.info( "EDGAR: rejecting near-miss '%s' → '%s' (first-token mismatch)", company_name, _matched_title, ) return {} cik_padded = str(matched_cik).zfill(10) logger.info("EDGAR fuzzy match: '%s' → CIK %s (score=%d)", company_name, cik_padded, _score) # Step 3: fetch full submissions metadata r2 = await client.get( f"https://data.sec.gov/submissions/CIK{cik_padded}.json", headers=headers, timeout=10, ) if r2.status_code != 200: return {"cik": cik_padded, "canonical_name": _matched_title} data = r2.json() return { "cik": str(data.get("cik", matched_cik)).zfill(10), "canonical_name": data.get("name", _matched_title), "tickers": data.get("tickers", []), "state": data.get("stateOfIncorporation", ""), } except Exception as exc: logger.warning("EDGAR lookup failed for '%s': %s", company_name, exc) return {} # ── Companies House ─────────────────────────────────────────────────────────── async def _ch_search(client: httpx.AsyncClient, query: str) -> list[dict]: """Search Companies House. Returns [] on error.""" try: r = await client.get( "https://api.company-information.service.gov.uk/search/companies", params={"q": query, "items_per_page": 5}, headers={"Authorization": _ch_auth_header()}, timeout=10, ) if r.status_code == 200: return r.json().get("items", []) logger.warning("Companies House HTTP %s", r.status_code) except Exception as exc: logger.warning("Companies House unreachable: %s", exc) return [] # ── Entity Resolver ─────────────────────────────────────────────────────────── class EntityResolver: """Resolves a raw company name to a canonical ResolvedEntity. Usage:: resolver = EntityResolver() entity = await resolver.resolve("Tesla Inc") """ async def resolve( self, raw_input: str, *, registry_id: str = "", jurisdiction: str = "", ) -> ResolvedEntity: """Resolve *raw_input* to a ``ResolvedEntity``. When ``registry_id`` is given (the user picked a candidate via the interactive picker — design §8c), the exact matching registry entry is selected instead of auto-ranking. ``jurisdiction`` narrows the registry search server-side. Both default to "" (autonomous best-match path). Raises: EntityNotFoundError: if no entity found from any source. """ key = _cache_key(f"{raw_input}|{registry_id}|{jurisdiction}") if settings.use_cache: cached = await _get_cached(key) if cached: logger.info("Entity cache hit for '%s'", raw_input) return cached async with httpx.AsyncClient() as client: entity = await self._resolve_uncached( client, raw_input, jurisdiction=jurisdiction, registry_id=registry_id ) if settings.use_cache: await _set_cached(key, entity) return entity async def search_candidates( self, raw_input: str, jurisdiction: str = "", limit: int = 8 ) -> list[dict]: """Return up to ``limit`` registry candidates for the picker (design §8c). Lightweight: registry metadata only, no Companies House / EDGAR enrichment. Returns [] when Registry Lookup is unavailable or finds nothing — the caller then falls back to the autonomous resolve path. """ async with httpx.AsyncClient() as client: results, _blocked = await _rl_search(client, raw_input, jurisdiction) return [_candidate_view(r, raw_input) for r in results[:limit]] async def _resolve_uncached( self, client: httpx.AsyncClient, raw_input: str, jurisdiction: str = "", registry_id: str = "", ) -> ResolvedEntity: rl_results, rl_blocked = await _rl_search(client, raw_input, jurisdiction) # If the user picked a specific candidate, select it exactly by id. if registry_id and rl_results: chosen = next( (r for r in rl_results if str(r.get("id", "")) == registry_id), None ) if chosen is not None: entity = _parse_rl_result(chosen, raw_input) return await self._enrich_entity(client, entity, raw_input) # Ambiguity: >3 matches from Registry Lookup. Per design §8c the ideal UX is # an interactive picker, but the autonomous pipeline can't pause for it, so # we select the top-ranked match (consistent with the 1–3 match path below) # and log the alternatives — the choice is transparent (the resolved # canonical name is shown in the report header), not silent, and the run # never dead-ends on a common name. if len(rl_results) > 3: logger.info( "Ambiguous '%s': %d registry matches; auto-selecting top match '%s'. " "Other candidates: %s", raw_input, len(rl_results), _parse_rl_result(rl_results[0], raw_input).canonical_name, ", ".join( _parse_rl_result(r, raw_input).canonical_name for r in rl_results[1:4] ), ) if rl_results: entity = _parse_rl_result(rl_results[0], raw_input) else: # Build minimal stub — may be enriched by EDGAR / Companies House below canonical = raw_input.strip().upper() entity = ResolvedEntity( canonical_name=canonical, aliases=_generate_aliases(canonical, raw_input, []), ) entity = await self._enrich_entity(client, entity, raw_input) # Nothing found from any authoritative source if not rl_results and not entity.sec_cik: if rl_blocked: # Registry Lookup is geo-blocked; we can't confirm non-existence. # Return a best-effort stub so private companies still flow through. logger.warning( "Degraded resolution for '%s': RL geo-blocked, not in EDGAR " "(likely private company). Returning stub entity.", raw_input, ) return entity raise EntityNotFoundError( f"No entity found for '{raw_input}' in Registry Lookup or SEC EDGAR." ) return entity async def _enrich_entity( self, client: httpx.AsyncClient, entity: ResolvedEntity, raw_input: str ) -> ResolvedEntity: """Shared enrichment: Companies House (UK) + SEC EDGAR (US public).""" # UK enrichment via Companies House if entity.jurisdiction and "gb" in entity.jurisdiction: entity = await self._enrich_companies_house(client, entity, raw_input) # US / unknown jurisdiction: confirm public status + CIK via SEC EDGAR. # Attempt this whenever we lack a CIK and the jurisdiction is US or unknown. # Crucially this is NOT gated on registry_lookup_id: once the Registry # Lookup geo-block is lifted (e.g. on the US-hosted deployment) every match # carries an id, and the old `not registry_lookup_id` gate then silently # skipped EDGAR for genuine public companies (Boeing, Colgate) — leaving # is_public=False and the SEC EDGAR research source disabled. EDGAR's fuzzy # name match (cutoff 70) self-gates: a private company won't match a filer. juris = (entity.jurisdiction or "").lower() if not entity.sec_cik and (juris.startswith("us") or not juris): edgar = await _edgar_lookup(client, entity.canonical_name) if edgar.get("cik"): edg_canonical = edgar.get("canonical_name") or entity.canonical_name entity = entity.model_copy(update={ "canonical_name": edg_canonical, "aliases": _generate_aliases(edg_canonical, raw_input, edgar.get("tickers", [])), "sec_cik": edgar["cik"], "is_public": True, }) return entity async def _enrich_companies_house( self, client: httpx.AsyncClient, entity: ResolvedEntity, raw_input: str, ) -> ResolvedEntity: items = await _ch_search(client, entity.canonical_name) if not items: return entity top = items[0] ch_name = top.get("title", entity.canonical_name) return entity.model_copy(update={ "canonical_name": ch_name, "aliases": _generate_aliases(ch_name, raw_input, []), "companies_house_number": top.get("company_number"), })