stonks-oracle/services/intelligence_pipeline_v3/resolution/alias_index.py

"""In-memory alias index for company name/ticker/alias lookup.

Supports case-insensitive matching with common corporate suffix stripping
(Inc., Corp., LLC, etc.) to maximize recall against varied document text.

The primary entry point for consumers is `build_alias_index(companies)` which
constructs a populated AliasIndex from company registry data.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

# Common corporate suffixes to strip for matching purposes.
_SUFFIX_PATTERN = re.compile(
    r"\s*\b("
    r"inc\.?|incorporated|"
    r"corp\.?|corporation|"
    r"co\.?|company|"
    r"ltd\.?|limited|"
    r"llc|l\.l\.c\.?|"
    r"plc|p\.l\.c\.?|"
    r"sa|s\.a\.?|"
    r"nv|n\.v\.?|"
    r"ag|"
    r"se|"
    r"group|"
    r"holdings?"
    r")\s*$",
    re.IGNORECASE,
)

# Trailing punctuation after suffix removal.
_TRAILING_PUNCT = re.compile(r"[.,;:\s]+$")


@dataclass
class IndexEntry:
    """A single entry linking a normalized key to a company."""

    company_id: str
    ticker: str
    name: str
    match_type: str  # "exact_ticker", "exact_name", "alias"


@dataclass
class AliasIndex:
    """Case-insensitive index mapping normalized strings to company entries.

    Stores tickers, full legal names, and known aliases. Returns all
    matching companies for a given query string.
    """

    _entries: dict[str, list[IndexEntry]] = field(default_factory=dict)

    @staticmethod
    def normalize(text: str) -> str:
        """Normalize a string for matching: lowercase, strip suffixes and punctuation."""
        s = text.strip().lower()
        # Strip corporate suffixes.
        s = _SUFFIX_PATTERN.sub("", s)
        # Remove trailing punctuation left behind.
        s = _TRAILING_PUNCT.sub("", s)
        # Collapse whitespace.
        s = re.sub(r"\s+", " ", s).strip()
        return s

    def add(self, key: str, entry: IndexEntry) -> None:
        """Add a lookup key mapped to an index entry."""
        normalized = self.normalize(key)
        if not normalized:
            return
        self._entries.setdefault(normalized, []).append(entry)

    def lookup(self, query: str) -> list[IndexEntry]:
        """Return all entries matching the normalized query."""
        normalized = self.normalize(query)
        if not normalized:
            return []
        return list(self._entries.get(normalized, []))

    def keys(self) -> list[str]:
        """Return all normalized keys in the index."""
        return list(self._entries.keys())

    def __len__(self) -> int:
        """Number of distinct normalized keys in the index."""
        return len(self._entries)


def build_alias_index(companies: list[dict]) -> AliasIndex:
    """Build a complete alias index from company registry data.

    This is the primary factory function for constructing an AliasIndex.
    It processes the same company dict format used by the symbol registry seed:

    Each company dict should contain:
        - id: str (company UUID)
        - ticker: str
        - legal_name: str

    Optional fields:
        - aliases: list[dict] with keys "alias" and optionally "alias_type"
          OR list[tuple[str, str]] of (alias_text, alias_type)
        - exchange: str (used for qualified ticker indexing)
        - sector: str
        - industry: str

    The function indexes:
        1. Ticker symbols (exact_ticker match type)
        2. Full legal names (exact_name match type)
        3. Legal names with suffixes stripped (exact_name match type)
        4. All known aliases (alias match type)

    Returns:
        A fully populated AliasIndex ready for lookup operations.
    """
    index = AliasIndex()

    for company in companies:
        company_id = str(company["id"])
        ticker = company["ticker"]
        name = company["legal_name"]

        # Index the ticker itself.
        entry_ticker = IndexEntry(
            company_id=company_id,
            ticker=ticker,
            name=name,
            match_type="exact_ticker",
        )
        index.add(ticker, entry_ticker)

        # Index the full legal name.
        entry_name = IndexEntry(
            company_id=company_id,
            ticker=ticker,
            name=name,
            match_type="exact_name",
        )
        index.add(name, entry_name)

        # Index known aliases.
        aliases = company.get("aliases", [])
        for alias_entry in aliases:
            if isinstance(alias_entry, dict):
                alias_text = alias_entry.get("alias", "")
            elif isinstance(alias_entry, (list, tuple)) and len(alias_entry) >= 1:
                alias_text = alias_entry[0]
            else:
                alias_text = str(alias_entry)

            if alias_text:
                entry_alias = IndexEntry(
                    company_id=company_id,
                    ticker=ticker,
                    name=name,
                    match_type="alias",
                )
                index.add(alias_text, entry_alias)

    return index