spb/company-atlas
Public
Python 66.3%
TypeScript 22.7%
JavaScript 8.6%
HTML 1.4%
CSS 0.7%
1"""`/ask` helper: turn a natural-language question into structured filters — deterministically first (countries, industries, event2types/subtypes, windows, intents: hiring / pricing / AI / launch / leadership / expansion / legal / developer), then optionally refined by3the small model (`prompts/ask-router/v1.md`). The result is an `interpretation` object the API runs against its own tables; nothing4here ever fabricates companies, events or numbers — `build_answer()` only phrases counts the API measured.56 interp = parse_question("companies hiring AI engineers in Canada last month")7 interp = await route_question(q) # + LLM refinement when configured8 answer = build_answer(interp, companies=12, events=48)9"""10from __future__ import annotations1112import logging13import re14from dataclasses import asdict, dataclass, field15from datetime import UTC, datetime16from typing import Any1718from companyatlas.config import settings19from companyatlas.services.llm.schemas import AskRoute20from companyatlas.taxonomy import EVENT_SUBTYPES, EventType2122log = logging.getLogger(__name__)2324ASK_PARSER_VERSION = "ask-parser-v1"2526# Compact built-in geography (the API may pass the full `countries` table for better coverage).27COUNTRY_TERMS: dict[str, str] = {28 "canada": "CA", "canadian": "CA", "quebec": "CA", "québec": "CA", "ontario": "CA", "united states": "US", "usa": "US", "u.s.": "US", "us ": "US",29 "american": "US", "america": "US", "united kingdom": "GB", "uk": "GB", "britain": "GB", "british": "GB", "england": "GB", "france": "FR", "french": "FR",30 "germany": "DE", "german": "DE", "japan": "JP", "japanese": "JP", "south korea": "KR", "korea": "KR", "korean": "KR", "india": "IN", "indian": "IN",31 "australia": "AU", "australian": "AU", "brazil": "BR", "brazilian": "BR", "mexico": "MX", "mexican": "MX", "spain": "ES", "spanish": "ES",32 "italy": "IT", "italian": "IT", "netherlands": "NL", "dutch": "NL", "sweden": "SE", "swedish": "SE", "switzerland": "CH", "swiss": "CH",33 "singapore": "SG", "china": "CN", "chinese": "CN", "israel": "IL", "israeli": "IL", "ireland": "IE", "irish": "IE", "norway": "NO", "denmark": "DK",34 "finland": "FI", "poland": "PL", "portugal": "PT", "belgium": "BE", "austria": "AT", "uae": "AE", "emirates": "AE", "saudi": "SA", "south africa": "ZA",35 "nigeria": "NG", "kenya": "KE", "argentina": "AR", "chile": "CL", "colombia": "CO", "indonesia": "ID", "vietnam": "VN", "thailand": "TH", "taiwan": "TW",36 "new zealand": "NZ", "turkey": "TR", "türkiye": "TR",37}38REGION_TERMS = {"europe": "europe", "european": "europe", "asia": "asia", "asian": "asia", "latin america": "latam", "latam": "latam", "africa": "africa",39 "middle east": "middle-east", "nordic": "nordics", "nordics": "nordics", "apac": "apac", "emea": "emea", "north america": "north-america"}4041# Intent → (event types, subtypes, tags). Order matters: first match sets the primary intent.42INTENTS: list[tuple[str, re.Pattern[str], list[str], list[str], list[str]]] = [43 ("ai", re.compile(r"\b(ai|artificial intelligence|machine learning|ml|llm|generative|genai)\b", re.IGNORECASE), [], ["AI_HIRING", "AI_LAUNCH", "TECHNOLOGY_ADOPTION"], ["ai"]),44 ("hiring", re.compile(r"\b(hiring|hire[sd]?|recruit(ing|ment)?|job[s]?|positions?|openings?|careers?|headcount|talent)\b", re.IGNORECASE), [EventType.HIRING], [], []),45 ("pricing", re.compile(r"\b(pric(e|es|ing)|plans?|tiers?|cheaper|more expensive|raised prices|cost)\b", re.IGNORECASE), [EventType.PRICING], [], []),46 ("launch", re.compile(r"\b(launch(es|ed|ing)?|new products?|released?|introduc(ed|es|ing)|unveil(ed|s)?|ship(ped|ping)?)\b", re.IGNORECASE), [EventType.PRODUCT], [], []),47 ("leadership", re.compile(r"\b(ceo|cfo|cto|coo|executive[s]?|leadership|board|founder[s]?|appoint(ed|s|ment)?|management team)\b", re.IGNORECASE), [EventType.LEADERSHIP], [], []),48 ("expansion", re.compile(r"\b(expan(d|ded|ding|sion)|new (office|offices|countr(y|ies)|market[s]?|location[s]?)|open(ed|ing)? an? office|enter(ed|ing)?)\b", re.IGNORECASE), [EventType.LOCATION], [], []),49 ("legal", re.compile(r"\b(terms( of service)?|privacy( policy)?|legal|tos|gdpr|policy changes?)\b", re.IGNORECASE), [EventType.LEGAL], [], []),50 ("developer", re.compile(r"\b(api[s]?|sdk[s]?|developer[s]?|docs|documentation|changelog|open[- ]source)\b", re.IGNORECASE), [EventType.DEVELOPER], [], []),51 ("financing", re.compile(r"\b(funding|raise[sd]? (?:\$|€|£|[0-9]|an? )|series [a-f]\b|ipo|investors?|financing|venture capital)\b", re.IGNORECASE), [EventType.FINANCING], [], []),52 ("m&a", re.compile(r"\b(acqui(re|red|sition|sitions)|merger[s]?|merged|bought|takeover)\b", re.IGNORECASE), [EventType.MA], [], []),53 ("communication", re.compile(r"\b(news|press releases?|announc(e|ed|ements?)|blog)\b", re.IGNORECASE), [EventType.COMMUNICATION], [], []),54]55_WINDOW_RE = re.compile(r"\b(?:in the |over the |during the |within the )?(?:last|past|previous)\s+(\d+)?\s*(day|days|week|weeks|month|months|quarter|quarters|year|years)\b", re.IGNORECASE)56_THIS_RE = re.compile(r"\bthis (week|month|quarter|year)\b", re.IGNORECASE)57_SINCE_RE = re.compile(r"\bsince (\d{4})(?:-(\d{2}))?(?:-(\d{2}))?\b", re.IGNORECASE)58_COUNT_RE = re.compile(r"\b(how many|number of|count of)\b", re.IGNORECASE)59_COMPARE_RE = re.compile(r"\b(vs\.?|versus|compare[d]?|comparison)\b", re.IGNORECASE)60_TIMELINE_RE = re.compile(r"\b(when did|history of|timeline|over time)\b", re.IGNORECASE)61_TREND_RE = re.compile(r"\b(trend(s|ing)?|fastest|most active|which (industries|countries|sectors))\b", re.IGNORECASE)62_QUOTED_RE = re.compile(r"[\"“”']([^\"“”']{2,60})[\"“”']")63_STOP = {"the", "a", "an", "of", "in", "on", "at", "for", "to", "and", "or", "with", "which", "what", "who", "are", "is", "that", "have", "has", "had",64 "companies", "company", "show", "me", "list", "find", "all", "any", "about", "from", "by", "their", "recently", "new", "did", "does", "do",65 "how", "many", "much", "last", "past", "this", "year", "years", "month", "months", "week", "weeks", "day", "days", "since", "were", "was",66 "been", "being", "into", "than", "more", "most", "less", "some", "there", "where", "when", "why", "get", "give", "tell"}67_UNITS = {"day": 1, "days": 1, "week": 7, "weeks": 7, "month": 30, "months": 30, "quarter": 90, "quarters": 90, "year": 365, "years": 365}686970@dataclass(slots=True)71class Interpretation:72 question: str73 intent: str = "search"74 intents: list[str] = field(default_factory=list)75 countries: list[str] = field(default_factory=list)76 regions: list[str] = field(default_factory=list)77 industries: list[str] = field(default_factory=list)78 event_types: list[str] = field(default_factory=list)79 event_subtypes: list[str] = field(default_factory=list)80 tags: list[str] = field(default_factory=list)81 companies: list[str] = field(default_factory=list)82 keywords: list[str] = field(default_factory=list)83 window_days: int | None = None84 answer_style: str = "list"85 min_importance: float | None = None86 confidence: float = 0.587 source: str = "deterministic"88 parser_version: str = ASK_PARSER_VERSION89 llm_model: str | None = None90 prompt_version: str | None = None9192 def to_dict(self) -> dict[str, Any]:93 return asdict(self)9495 @property96 def filters(self) -> dict[str, Any]:97 """Exactly what the API needs to query `events` / `companies`."""98 return {k: v for k, v in {"countries": self.countries, "industries": self.industries, "event_types": self.event_types,99 "event_subtypes": self.event_subtypes, "tags": self.tags, "companies": self.companies, "keywords": self.keywords,100 "window_days": self.window_days, "min_importance": self.min_importance}.items() if v}101102103def _find_countries(q: str, extra: dict[str, str] | None) -> list[str]:104 low = f" {q.lower()} "105 found: list[str] = []106 terms = dict(COUNTRY_TERMS)107 for code, name in (extra or {}).items():108 if name:109 terms[name.lower()] = code.upper()110 for term, code in sorted(terms.items(), key=lambda kv: -len(kv[0])):111 pattern = r"(?<![a-z])" + re.escape(term.strip()) + r"(?![a-z])"112 if re.search(pattern, low) and code not in found:113 found.append(code)114 return found[:8]115116117def _find_industries(q: str, industries: dict[str, str] | None) -> list[str]:118 if not industries:119 return []120 low = q.lower()121 found: list[str] = []122 for slug, name in sorted(industries.items(), key=lambda kv: -len(kv[1] or kv[0])):123 for needle in {slug.replace("-", " "), (name or "").lower()}:124 if needle and len(needle) >= 3 and re.search(r"(?<![a-z])" + re.escape(needle) + r"(?![a-z])", low) and slug not in found:125 found.append(slug)126 return found[:8]127128129def _window(q: str) -> int | None:130 m = _WINDOW_RE.search(q)131 if m:132 n = int(m.group(1) or 1)133 return n * _UNITS[m.group(2).lower()]134 m = _THIS_RE.search(q)135 if m:136 return _UNITS[m.group(1).lower()]137 if re.search(r"\btoday\b", q, re.IGNORECASE):138 return 1139 if re.search(r"\byesterday\b", q, re.IGNORECASE):140 return 2141 m = _SINCE_RE.search(q)142 if m:143 start = datetime(int(m.group(1)), int(m.group(2) or 1), int(m.group(3) or 1), tzinfo=UTC)144 return max(1, (datetime.now(UTC) - start).days)145 return None146147148def _keywords(q: str, *, drop: set[str]) -> list[str]:149 words = re.findall(r"[a-zA-Z][a-zA-Z0-9\-\+\.]{1,}", q.lower())150 out: list[str] = []151 for w in words:152 w = w.strip(".")153 if w in _STOP or w in drop or len(w) < 3:154 continue155 if w not in out:156 out.append(w)157 return out[:8]158159160def parse_question(q: str, *, industries: dict[str, str] | None = None, countries: dict[str, str] | None = None) -> Interpretation:161 """Deterministic parser. `industries` = {slug: name} and `countries` = {code: name} from the registry (optional)."""162 q = " ".join((q or "").split())[:400]163 it = Interpretation(question=q)164 it.countries = _find_countries(q, countries)165 it.regions = [v for k, v in REGION_TERMS.items() if re.search(r"(?<![a-z])" + re.escape(k) + r"(?![a-z])", q.lower())]166 it.industries = _find_industries(q, industries)167 it.window_days = _window(q)168 drop: set[str] = set()169 for name, pattern, types, subtypes, tags in INTENTS:170 if pattern.search(q):171 it.intents.append(name)172 for t in types:173 if str(t) not in it.event_types:174 it.event_types.append(str(t))175 for s in subtypes:176 if s in EVENT_SUBTYPES and s not in it.event_subtypes:177 it.event_subtypes.append(s)178 for tag in tags:179 if tag not in it.tags:180 it.tags.append(tag)181 if it.intents:182 it.intent = it.intents[0]183 if it.intent == "ai" and "hiring" in it.intents and str(EventType.HIRING) not in it.event_types:184 it.event_types.append(str(EventType.HIRING))185 if re.search(r"\b(price increase[s]?|raised prices|more expensive|increased (their )?prices)\b", q, re.IGNORECASE):186 it.event_subtypes = [s for s in it.event_subtypes if s != "PRICE_DECREASE"] + ["PRICE_INCREASE"]187 if re.search(r"\b(price (cut|decrease)[s]?|cheaper|lowered prices)\b", q, re.IGNORECASE):188 it.event_subtypes.append("PRICE_DECREASE")189 if re.search(r"\b(new countr(y|ies)|international(ly)?|abroad|overseas)\b", q, re.IGNORECASE):190 it.event_subtypes.append("COUNTRY_EXPANSION")191 if re.search(r"\b(important|major|significant|big)\b", q, re.IGNORECASE):192 it.min_importance = 0.6193 it.companies = [m.strip() for m in _QUOTED_RE.findall(q)][:8]194 if _COUNT_RE.search(q):195 it.answer_style = "count"196 elif _COMPARE_RE.search(q):197 it.answer_style = "compare"198 it.intent = "compare" if it.intent == "search" else it.intent199 elif _TIMELINE_RE.search(q):200 it.answer_style = "timeline"201 elif _TREND_RE.search(q):202 it.answer_style = "trend"203 it.intent = "trend" if it.intent == "search" else it.intent204 for term in COUNTRY_TERMS:205 drop.update(term.split())206 drop.update({"ai", "hiring", "pricing", "launch", "launched", "executive", "executives"})207 it.keywords = _keywords(q, drop=drop | set(REGION_TERMS))208 it.event_subtypes = list(dict.fromkeys(it.event_subtypes))209 signal = sum(bool(x) for x in (it.countries, it.industries, it.event_types or it.event_subtypes, it.window_days, it.keywords))210 it.confidence = round(min(0.9, 0.35 + 0.12 * signal), 2)211 return it212213214async def route_question(q: str, *, industries: dict[str, str] | None = None, countries: dict[str, str] | None = None,215 use_llm: bool | None = None) -> Interpretation:216 """Deterministic parse, then an optional LLM refinement that may only tighten filters (never invents results)."""217 it = parse_question(q, industries=industries, countries=countries)218 enabled = settings.llm_configured if use_llm is None else (use_llm and settings.llm_configured)219 if not enabled or not q.strip():220 return it221 try:222 from companyatlas.db import jsonb223 from companyatlas.services.llm.gateway import get_provider224 from companyatlas.services.llm.prompts import load_prompt225226 prompt = load_prompt("ask-router")227 context = {"question": q, "parsed": it.filters | {"intent": it.intent, "answer_style": it.answer_style},228 "allowed_event_types": [str(t) for t in EventType], "allowed_event_subtypes": sorted(EVENT_SUBTYPES),229 "industries": [{"slug": k, "name": v} for k, v in (industries or {}).items()][:200],230 "countries": [{"code": k, "name": v} for k, v in (countries or {}).items()][:250]}231 res = await get_provider().complete_json("small", prompt.system, jsonb(context), AskRoute, max_tokens=400)232 r: AskRoute = res.data233 allowed_ind = set(industries or {})234 it.countries = list(dict.fromkeys(it.countries + [c for c in r.countries if not countries or c in countries]))[:8]235 it.industries = list(dict.fromkeys(it.industries + [s for s in r.industries if s in allowed_ind]))[:8]236 it.event_types = list(dict.fromkeys(it.event_types + [t for t in r.event_types if t in {str(x) for x in EventType}]))[:8]237 it.event_subtypes = list(dict.fromkeys(it.event_subtypes + r.event_subtypes))[:8]238 it.companies = list(dict.fromkeys(it.companies + r.companies))[:8]239 it.keywords = list(dict.fromkeys(r.keywords or it.keywords))[:8]240 it.window_days = it.window_days or r.window_days241 if it.intent == "search" and r.intent:242 it.intent = r.intent243 if r.answer_style in ("list", "count", "compare", "timeline"):244 it.answer_style = r.answer_style245 it.confidence = round(max(it.confidence, min(0.95, r.confidence)), 2)246 it.source = "llm"247 it.llm_model = res.model248 it.prompt_version = prompt.ref249 except Exception as exc: # noqa: BLE001 — the deterministic parse is always a valid fallback250 log.warning("ask-router llm refinement skipped", extra={"error": str(exc)[:200]})251 return it252253254def build_answer(it: Interpretation, *, companies: int, events: int) -> str:255 """Phrase measured counts only. The API fills `companies` / `events` from its own query."""256 bits: list[str] = []257 what = {"hiring": "hiring-related events", "pricing": "pricing events", "ai": "AI-related events", "launch": "product events",258 "leadership": "leadership events", "expansion": "location events", "legal": "legal page changes", "developer": "developer events",259 "financing": "financing mentions", "m&a": "M&A mentions", "communication": "communication events"}.get(it.intent, "events")260 scope = []261 if it.countries:262 scope.append("in " + ", ".join(it.countries))263 if it.industries:264 scope.append("within " + ", ".join(it.industries))265 if it.window_days:266 scope.append(f"over the last {it.window_days} days")267 scope_txt = (" " + " ".join(scope)) if scope else ""268 if events == 0 and companies == 0:269 return f"No monitored evidence matches this question yet{scope_txt}. Results only include events detected on public company pages."270 bits.append(f"{events} {what}{scope_txt} across {companies} monitored {'company' if companies == 1 else 'companies'}.")271 bits.append("Every item links to its public source; interpretations are labelled with a confidence level.")272 return " ".join(bits)273274275__all__ = ["ASK_PARSER_VERSION", "COUNTRY_TERMS", "Interpretation", "build_answer", "parse_question", "route_question"]276