You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""
|
|
Entity Resolver - resolve text mentions to ticker symbols using 4-stage matching.
|
|
"""
|
|
|
|
import re
|
|
import logging
|
|
from typing import List, Optional, Set
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy import select
|
|
|
|
from app.models.overlay_registry import CompanyAlias
|
|
from app.core.overlay_config import TOP_50_SYMBOLS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Pre-compiled common English stop-words to avoid false ticker matches
|
|
_STOP_WORDS = {
|
|
"A", "I", "IN", "ON", "AT", "IT", "IS", "BE", "AS", "OR", "AND",
|
|
"THE", "FOR", "TO", "OF", "BY", "AN", "UP", "DO", "GO", "US",
|
|
"PM", "AM", "ET", "AI", # keep AI out to avoid false positives
|
|
}
|
|
|
|
|
|
class EntityResolver:
|
|
"""Resolve text mentions to ticker symbols using 4-stage matching."""
|
|
|
|
def __init__(self):
|
|
self._alias_cache: dict = {}
|
|
|
|
async def load_aliases(self, db: AsyncSession) -> None:
|
|
"""Load company aliases from DB into memory cache."""
|
|
result = await db.execute(
|
|
select(CompanyAlias).where(CompanyAlias.active == True)
|
|
)
|
|
rows = result.scalars().all()
|
|
cache: dict = {}
|
|
for row in rows:
|
|
key = row.alias_value.lower().strip()
|
|
if key not in cache:
|
|
cache[key] = []
|
|
cache[key].append((row.symbol, row.confidence))
|
|
self._alias_cache = cache
|
|
logger.debug(f"EntityResolver: loaded {len(cache)} alias entries")
|
|
|
|
def resolve_from_title(self, title: str) -> List[str]:
|
|
"""
|
|
Extract ticker symbols from a text string.
|
|
|
|
Stage 1: $TICKER pattern (highest confidence)
|
|
Stage 2: Direct uppercase word match against TOP_50_SYMBOLS
|
|
Stage 3: Alias / company name match (case-insensitive)
|
|
"""
|
|
symbols: Set[str] = set()
|
|
|
|
# Stage 1: $TICKER pattern
|
|
dollar_tickers = re.findall(r'\$([A-Z]{1,5})\b', title)
|
|
for t in dollar_tickers:
|
|
symbols.add(t)
|
|
|
|
# Stage 2: Uppercase word match against known symbols
|
|
words = re.findall(r'\b([A-Z]{1,5})\b', title)
|
|
for w in words:
|
|
if w in TOP_50_SYMBOLS and w not in _STOP_WORDS:
|
|
symbols.add(w)
|
|
|
|
# Stage 3: Alias / company name match (case-insensitive, word-boundary).
|
|
# Plain `substring in title` produces false positives like
|
|
# "Critical Metals" → META or "Ups Guidance" → UPS.
|
|
if self._alias_cache:
|
|
title_lower = title.lower()
|
|
for alias_text, candidates in self._alias_cache.items():
|
|
pattern = r'\b' + re.escape(alias_text) + r'\b'
|
|
if re.search(pattern, title_lower):
|
|
for sym, conf in candidates:
|
|
if conf >= 0.7:
|
|
symbols.add(sym)
|
|
|
|
return list(symbols)
|
|
|
|
def resolve_symbol(self, text: str) -> Optional[str]:
|
|
"""Resolve a single best-match symbol from text."""
|
|
results = self.resolve_from_title(text)
|
|
if not results:
|
|
return None
|
|
# Prefer symbols in TOP_50 list
|
|
for s in results:
|
|
if s in TOP_50_SYMBOLS:
|
|
return s
|
|
return results[0]
|