intelligence layer (events, clustering, LLM gateway + prompts, metrics, signals, trends, alerts, retention) and API v1 (74 endpoints, SSE, rate limits, exports, admin)
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
64 changed files +10,114 −6
modified
docs/API.md
+69 −1
@@ -83,7 +83,7 @@ type Signal = { id: string; company_id: string | null; scope: string; scope_key: | ||
| 83 | 83 | ### Companies |
| 84 | 84 | | Method | Path | Notes | |
| 85 | 85 | |---|---|---| |
| 86 | −| GET | `/companies` | filters: `q, country, industry, tier, public (bool), status, has_events (bool), sort=activity\|events\|hiring\|name\|importance\|recent, sparkline=1` → page of `CompanyCard` | | |
| 86 | +| GET | `/companies` | filters: `q, country, industry, tier, public (bool), status, has_events (bool), sort=activity\|events\|hiring\|name\|importance\|recent\|relevance` (relevance is automatic with `q`), `sparkline=1` → page of `CompanyCard` | | |
| 87 | 87 | | GET | `/companies/{slug_or_id}` | `CompanyCard & { aliases: string[], domains: {domain, kind}[], relationships: {kind, company: {slug, display_name} \| null, to_name, valid_from, valid_to, confidence}[], metrics_detail: {metric, value, confidence, computed_at, inputs}[], sensors_by_surface: Record<string, number>, coverage: {historical_coverage, first_observed_at, days_observed, sensor_uptime}, signals: Signal[], sparklines: {activity_30d: number[], hiring_90d: number[]} }` | |
| 88 | 88 | | GET | `/companies/{slug}/events` | filters `event_type, event_subtype, since, until, min_importance, surface`; `sort=recent\|importance` → page of `Event` | |
| 89 | 89 | | GET | `/companies/{slug}/timeline?filter=all\|products\|jobs\|pricing\|leadership\|locations\|legal\|news\|developer&limit=200` | `{ items: (Event & {day: string})[], days: {day, count}[] }` grouped for the timeline UI | |
@@ -168,6 +168,74 @@ type Signal = { id: string; company_id: string | null; scope: string; scope_key: | ||
| 168 | 168 | | GET | `/admin/costs?days=30` | `{ items: {day, dimension, key, units, cost_estimate}[], per_1000_companies, per_million_observations, per_meaningful_event }` | |
| 169 | 169 | | POST | `/admin/cache/clear` | | |
| 170 | 170 | |
| 171 | +## Rate limits & auth | |
| 172 | + | |
| 173 | +| Tier | How | Limit (per client, per minute) | | |
| 174 | +|---|---|---| | |
| 175 | +| anonymous | no header — keyed by client IP (first hop of `X-Forwarded-For`) | 120 | | |
| 176 | +| authenticated | `X-CA-API-Key: <key>` (create with `catlas api-key create <name>`) | 600 | | |
| 177 | +| paid | `X-CA-API-Key` with `tier = paid` (`catlas api-key create <name> --tier paid`) | 3 000 | | |
| 178 | +| internal | `X-CA-API-Key` with `tier = internal` | unlimited | | |
| 179 | +| admin | valid `X-CA-Admin-Token` | bypass (also required for `/admin/*`) | | |
| 180 | + | |
| 181 | +Token bucket per process (refills continuously at `limit / 60` tokens per second). Every limited response carries `X-RateLimit-Limit`, | |
| 182 | +`X-RateLimit-Remaining` and `X-RateLimit-Tier`; when the bucket is empty the API answers `429 { "detail": "rate limit exceeded" }` with a | |
| 183 | +`Retry-After` (seconds). `/health`, `/ready`, `/api/v1/docs` and `/api/v1/openapi.json` are never limited. Keys are stored as sha256 hashes | |
| 184 | +(`api_keys.key_hash`); the raw key is printed once by the CLI. `last_used_at` / `request_count` are updated lazily (≈ every 30 s). Owner endpoints | |
| 185 | +(`/watchlist`, `/alerts*`) need `X-CA-Owner-Token` (client-generated, 24–200 chars); the owner row is created on first use and the token is only | |
| 186 | +ever stored hashed. Admin endpoints answer `401 { "detail": "admin token required" }` without a valid `X-CA-Admin-Token`. | |
| 187 | + | |
| 188 | +## Implementation notes (API v1.0 — additive details) | |
| 189 | + | |
| 190 | +- **Caching / ETags.** Cached public aggregates (`/pulse` 60 s, `/stats` 60 s, `/system` 30 s, `/rankings` 120 s, `/industries*`, `/countries*`, | |
| 191 | + `/index`, `/map`, `/trends`, `/events/types`, `/events/summary`, `/stats/history` 300 s, `/sitemap` 600 s, `/methodology` 3600 s) return | |
| 192 | + `Cache-Control: public, max-age=N, stale-while-revalidate=2N` plus a weak `ETag`; a matching `If-None-Match` yields `304`. Other public GETs | |
| 193 | + carry `public, max-age=30…600` (per resource volatility); `/live*`, owner and admin routes are `no-store`. `POST /admin/cache/clear?prefix=` | |
| 194 | + clears one key prefix (e.g. `pulse`, `industries:`) or everything. | |
| 195 | +- **Pagination totals.** `total` is exact up to **10 000** (bounded count) so deep feeds stay cheap; `pages` derives from it. | |
| 196 | +- **`/live`** also accepts `country` and `industry`; the payload is `{ items, count, cursor, server_time }` — pass `cursor` back as `since`. | |
| 197 | +- **`/live/stream`** accepts `?since=&event_type=&min_importance=&max_s=` (`max_s` ≤ 3600 bounds the connection for proxies with idle timeouts; | |
| 198 | + the stream ends with `event: end` carrying the last cursor). The first message is always an `event: heartbeat` with `{at, cursor}`; | |
| 199 | + `event: event` messages carry `id: <event id>`. Response headers include `X-Accel-Buffering: no`. | |
| 200 | +- **`/companies`** gains `sort=relevance` (automatic when `q` is present); `q` uses FTS (`companies.search`) for ≥ 3 characters with a | |
| 201 | + trigram fallback and a prefix match for shorter strings. `sparkline=1` adds 30 daily `activity_score` points. | |
| 202 | +- **`/companies/{slug}`** additionally returns `company_type, employees, wikidata_id, indexed, discovered_at, first_observed_at, last_change_at, | |
| 203 | + recent_events: Event[10]` and `sparkline`. | |
| 204 | +- **`/companies/{slug}/timeline`** accepts `filter=corporate` (FINANCING, M&A, PARTNERSHIP, STRATEGY) and `before=<iso>`; items carry `day`. | |
| 205 | +- **`/companies/{slug}/jobs`** accepts `department`, `remote`, `sort=recent|title|posted`; `q` is a trigram `ILIKE` on `title`. | |
| 206 | +- **`/companies/{slug}/history`** accepts `versions=1…20`; `/companies/{slug}/sensors?include_retired=1`; `/companies/{slug}/news?category=`; | |
| 207 | + `/companies/{slug}/locations?status=all`; `/companies/{slug}/pricing?history_limit=`. | |
| 208 | +- **`/companies/compare`** accepts `days=` (7–365) for the activity series; unknown slugs → `404 company not found: <slug>`; fewer than 2 → 422. | |
| 209 | +- **`/events`** accepts `status=active|retracted|duplicate|review|all` (default `active`). Retracted events expose `retracted_reason`. | |
| 210 | + `/events/{id}` embeds the full `CompanyCard` under `company` (a superset of `CompanyRef`), `sources` (falls back to the event's own | |
| 211 | + `source_url` when `event_sources` is empty), `change` and `cluster`. | |
| 212 | +- **`/events/types`** lists every taxonomy type/subtype (count `0` when unseen) so filter UIs can render on an empty database. | |
| 213 | +- **`/events/summary`** items also carry `previous` (count of the preceding window). | |
| 214 | +- **`/rankings`** items are `CompanyCard & { rank, value, delta }`; the payload also lists `kinds`. `hiring_growth`/`hiring_decline` read | |
| 215 | + `hiring_momentum_7d|30d|90d` according to the window; `pricing_changes` counts PRICING events in the window; `delta` is the difference with | |
| 216 | + the `metric_series` value at the start of the window (`null` when no history yet). | |
| 217 | +- **`/countries/{code}`** accepts an ISO-2 code or a name slug (`/countries/canada`); `CountryRow` includes `slug` and `subregion`. | |
| 218 | + `/industries/{slug}` adds `children`, `signals`, `companies_total`; `trending` is `[]` until per-industry trend data exists (never faked). | |
| 219 | +- **`/signals`** accepts `scope_key`, `company`, `min_strength`; items include a `company` ref when scoped to a company. | |
| 220 | +- **`/search`** people/products use trigram similarity (`migrations/pending/api.sql` adds the two supporting GIN indexes). `/search/suggest` | |
| 221 | + items add `slug` / `code` / `event_type` next to `href`. | |
| 222 | +- **`/ask`** returns `engine` (`deterministic` | `llm`) and `events_total`; the parser is `services/llm/ask.py` when present (deterministic | |
| 223 | + parse + optional LLM refinement), else `api/ask_fallback.py`. Answers only phrase measured counts. | |
| 224 | +- **`/snapshots/{id}`** accepts `include=text,blocks,extracted` and returns `text_truncated` (text capped at 200 kB), `sensor`, `object_keys`. | |
| 225 | + `/snapshots/{a}/diff/{b}` orders the pair by `fetched_at`, returns `source: "change"` (stored diff, with `change_id`) or `"computed"` | |
| 226 | + (`sdk.diff.compare`); `501` when the diff engine is unavailable, `404` when block objects are missing. | |
| 227 | +- **Exports.** `/export/events.{fmt}` also accepts `until`, `company`, `min_importance`, `limit` (≤ 10 000); `/export/companies.{fmt}` accepts | |
| 228 | + `status`, `tier`, `limit` (≤ 20 000); `/export/jobs.{json|ndjson|csv}` accepts `status=open|removed|all`, `country`, `ai`, `limit` (≤ 10 000). | |
| 229 | + CSV flattens the company ref (`company_slug`, `company_name`, `company_domain`, `country`). | |
| 230 | +- **`/sitemap`** pages hold 5 000 entries; `kind=companies` returns `total`; industries/countries only list entries with ≥ 1 active company. | |
| 231 | +- **Admin additions.** `GET /admin/sensors` accepts `surface`, `sort=recent|next_run|failures|quality|changes|created`; sensor actions return | |
| 232 | + `{ ok, action, sensor, queued? , interval_s? }` (`set_interval` is clamped to `CA_MIN_INTERVAL_S…CA_MAX_INTERVAL_S` and recomputes the tier). | |
| 233 | + `POST /admin/companies` → `201 { ok, company, queued }` (409 when the registrable domain already exists; 422 for non-http(s) URLs, unknown | |
| 234 | + country codes or industry slugs). `GET /admin/failures` returns `by_class` and `classes`; `GET /admin/queue` returns `counts` + `items`; | |
| 235 | + `POST /admin/queue/requeue-dead { kind? }`; `GET /admin/llm` includes `stats`; `POST /admin/reviews/{id}` accepts `label` (correct | | |
| 236 | + duplicate | noise | misclassified) feeding `/admin/quality.calibration`; retract/restore append an audit trail under `payload._audit` and | |
| 237 | + never delete; `GET /admin/storage` reports the object store. `/admin/overview` also returns `companies_by_onboarding`, `scheduler_heartbeat`. | |
| 238 | + | |
| 171 | 239 | ## Conventions for implementers |
| 172 | 240 | - Every list endpoint is bounded (`per_page ≤ 200`, `limit ≤ 500`), uses indexed predicates, and returns `Cache-Control: public, max-age=60` for public aggregates (`/pulse`, `/stats`, `/rankings`, `/industries`, `/countries`) and `no-store` for `/live*`, owner and admin routes. |
| 173 | 241 | - Company lookups accept slug **or** id. Unknown → 404 `{detail: "company not found"}`. |
added
docs/EVENT-TAXONOMY.md
+87 −0
@@ -0,0 +1,87 @@ | ||
| 1 | +# Event taxonomy — deterministic rules, wording policy, importance & confidence | |
| 2 | + | |
| 3 | +Owner: intelligence layer (`services/events.py`, `services/clustering.py`, rules version `rules-v1`, schema `event-v1`). | |
| 4 | +Vocabulary: `taxonomy.EVENT_SUBTYPES` (subtype → type + default importance). The LLM classifier may only pick from that list. | |
| 5 | + | |
| 6 | +## Pipeline | |
| 7 | + | |
| 8 | +``` | |
| 9 | +changes (status='pending', kind ≥ meaningful) ─▶ derive_events(change, company, sensor, baseline) pure rules, no I/O | |
| 10 | + ─▶ persist_change_events() events + event_sources + clusters + review_queue + llm_jobs | |
| 11 | + ─▶ changes.status = 'processed' · sensors.event_count · companies.last_event_at | |
| 12 | + ─▶ alerts.evaluate_alerts(new_event_ids) | |
| 13 | +``` | |
| 14 | + | |
| 15 | +- `catlas process-changes [--limit] [--loop]` · periodic `process-changes` every 20 s (SKIP LOCKED, batch 200). | |
| 16 | +- `catlas reprocess-events --since 7d [--company]` re-runs the rules on processed changes **without refetching**; dedupe keys make it idempotent, | |
| 17 | + new rules simply add the events they now produce. | |
| 18 | +- **Noise / minor changes never produce events** (they are archived if still pending). Meaningful (≥ 0.40), major (≥ 0.65), critical (≥ 0.85) do. | |
| 19 | + | |
| 20 | +## Rules → subtypes | |
| 21 | + | |
| 22 | +| Family | Input (`structured_delta`) | Subtype(s) | Title pattern | | |
| 23 | +|---|---|---|---| | |
| 24 | +| HIRING | `jobs.added/removed/open_before/open_after` | `JOB_COUNT_INCREASE` when net > 0 (entities.jobs ≤ 50) | *12 new positions detected on careers page* | | |
| 25 | +| | | `JOB_COUNT_DECREASE` when net < 0 | *7 monitored job listings no longer visible on careers page* | | |
| 26 | +| | AI keyword in added titles (`taxonomy.AI_KEYWORDS`) or `is_ai` | `AI_HIRING` | *2 AI-related positions detected on careers page* | | |
| 27 | +| | ≤ 5 jobs added | `NEW_JOB` per job | *New position listed: Senior ML Engineer (Toronto, CA)* | | |
| 28 | +| | added ≥ mean + 2σ of `baselines.jobs_new_weekly` (≥ 4 samples) — fallback ≥ 10 and ≥ 50 % of open_before | `HIRING_SURGE` | *Hiring surge signal: 25 new positions detected in one observation (baseline ≈ 3.0 new/week)* | | |
| 29 | +| | removed ≥ 10, ≥ 50 % of open_before and open_after ≤ 50 % | `HIRING_FREEZE_SIGNAL` (+ review `unexpected_activity`) | *Hiring slowdown signal: 30 of 40 monitored listings no longer visible* | | |
| 30 | +| PRICING | `plans.price_changed` | `PRICE_INCREASE` / `PRICE_DECREASE` (old/new values, `payload.pct`) | *Pro plan price observed at $59 (was $49)* | | |
| 31 | +| | `plans.added` | `NEW_PRICING_TIER` (tag `enterprise` when contact-sales) | *New pricing tier listed: Enterprise (contact sales)* | | |
| 32 | +| | `plans.removed` | `PRICING_TIER_REMOVED` | *Pricing tier no longer listed: Starter* | | |
| 33 | +| | pricing surface, no typed delta | `PRICING_CHANGE` (text-diff) | *Pricing page materially updated (3 blocks changed)* | | |
| 34 | +| LEADERSHIP | `people.added` (executive) | `NEW_EXECUTIVE` | *Jane Doe listed as Chief Financial Officer on leadership page* | | |
| 35 | +| | `people.removed` (executive) | `EXECUTIVE_NO_LONGER_LISTED` | *Jane Doe no longer listed on leadership page* | | |
| 36 | +| | `people.title_changed` | `EXECUTIVE_TITLE_CHANGE` | *Ann Lee now listed as COO (was VP Operations)* | | |
| 37 | +| | ≥ 3 people changes, or non-executive changes only | `LEADERSHIP_CHANGE` aggregate | *Leadership page updated: 2 added, 1 no longer listed* | | |
| 38 | +| PRODUCT | `products.added` / `removed` | `NEW_PRODUCT` / `PRODUCT_REMOVED` | *New product listed: Atlas Pro* · *Product no longer listed: Atlas Lite* | | |
| 39 | +| LOCATION | `locations.added` | `NEW_LOCATION` (kind-aware label) | *New office listed: Toronto, CA* | | |
| 40 | +| | `locations.new_countries` | `COUNTRY_EXPANSION` | *New country presence listed: Japan (Tokyo)* | | |
| 41 | +| | `locations.removed` | `OFFICE_REMOVED` | *Office no longer listed: Berlin, DE* | | |
| 42 | +| COMMUNICATION / IR / DEVELOPER | `news.added` (≤ 20 per change) | `NEWS_RELEASE` · `BLOG_POST` · `CHANGELOG_ENTRY` · `INVESTOR_UPDATE` · `EARNINGS_RELEASE` (title heuristics: earnings/quarter/fiscal → earnings; investor/annual report/dividend → IR; category/surface otherwise) | *News release: <title>* | | |
| 43 | +| DEVELOPER (text-diff) | docs/developer · api · changelog | `DOC_CHANGE` · `API_CHANGE` · `CHANGELOG_ENTRY` | *Documentation updated (5 sections changed)* | | |
| 44 | +| LEGAL (text-diff) | legal_terms · legal_privacy · security | `TERMS_CHANGE` · `PRIVACY_POLICY_CHANGE` · `SECURITY_UPDATE` (+ review `legal_sensitive`) | *Terms of service page materially updated (3 sections changed)* — sections from block `path` | | |
| 45 | +| WEBSITE | homepage meaningful · major+ · `meta.title_changed` | `WEBSITE_CHANGE` · `HOMEPAGE_REDESIGN` · `MESSAGING_CHANGE` | *Homepage materially redesigned (9 blocks changed, 62% of text)* | | |
| 46 | +| other surfaces (text-diff) | products/services/solutions · investor_relations · sustainability · status · careers · locations · about | `PRODUCT_UPDATE` · `INVESTOR_UPDATE` · `SUSTAINABILITY_UPDATE` · `OPERATIONS_UPDATE` · `WEBSITE_CHANGE` | *<Surface> updated (n sections changed)* | | |
| 47 | + | |
| 48 | +Surfaces with no rule (`other`, `partners`, `customers`, `support`, …) and meaningful+ significance produce no deterministic event; the change is | |
| 49 | +queued for the LLM classifier instead (`llm_jobs.kind = classify_change`), together with ambiguous surfaces (homepage/about/products text-only). | |
| 50 | + | |
| 51 | +## Wording policy (spec §167–168) | |
| 52 | + | |
| 53 | +- Verbs: **detected, observed, listed, now listed, no longer listed, no longer visible, appears, signal**. | |
| 54 | +- Never: *fired, laid off, layoffs, shut down, bankrupt, collapsed* (`taxonomy.FORBIDDEN_WORDING`). `safe_wording()` rewrites them defensively; | |
| 55 | + the LLM schemas reject them at validation time. | |
| 56 | +- Disappearance ≠ departure: `EXECUTIVE_NO_LONGER_LISTED` summaries say so explicitly; `HIRING_FREEZE_SIGNAL` / `hiring_freeze` are labelled *signal*. | |
| 57 | +- Titles ≤ 200 chars, summaries ≤ 600; entity lists bounded to 50; per-entity events capped (5 jobs, 10 people, 20 news items). | |
| 58 | + | |
| 59 | +## Importance & confidence | |
| 60 | + | |
| 61 | +- `importance = default(subtype) × (0.7 + 0.6·significance) × (1 + 0.3·magnitude)` clamped to [0.05, 1]. `magnitude` is rule-specific and 0–1: | |
| 62 | + relative job delta (`n / max(open_before, 5)`), |price pct| / 50, headquarters vs office, country expansion 0.6, launch-tagged news 0.5, text | |
| 63 | + ratio × 2 for page changes. | |
| 64 | +- `confidence` = evidence quality (`taxonomy.EVIDENCE_CONFIDENCE`): structured ATS JSON **0.95** (jobs_board surface, Greenhouse/Lever/Ashby/ | |
| 65 | + SmartRecruiters/Workday/JSON connectors), JSON-LD / feed **0.9**, HTML extraction **0.8**, text-diff-only **0.7**. `confidence_label` via | |
| 66 | + `taxonomy.confidence_label` (VERIFIED ≥ 0.95, HIGH_CONFIDENCE ≥ 0.85, LIKELY ≥ 0.7, INFERRED ≥ 0.5, else LOW_CONFIDENCE). | |
| 67 | +- Corroboration: a second surface in the same cluster adds **+0.03 per extra surface**, capped at 0.99 (clustering below). | |
| 68 | +- Review queue: `major_event` for critical changes, `legal_sensitive` for legal inferences, `low_confidence` (< 0.5), `unexpected_activity` for freeze signals. | |
| 69 | + | |
| 70 | +## Idempotency & clustering | |
| 71 | + | |
| 72 | +- `events.dedupe_key = sha(company, subtype, normalised entity key, sensor, detection day)` → inserting the same change twice is a no-op | |
| 73 | + (`on conflict do nothing`). Entity keys: job label, plan name, person name, product name, place, news title, `jobs:<before>><after>` for aggregates, | |
| 74 | + `page:<sensor>` for page-level events. | |
| 75 | +- `event_clusters.cluster_key = sha(company, subtype, normalised entity key, 7-day bucket)` — surface-independent, so the same press release seen | |
| 76 | + on the newsroom **and** the feed, or the same executive on leadership **and** about, fold into one cluster. The first event is canonical; later | |
| 77 | + ones are stored with `status='duplicate'` + `cluster_id`, the canonical event gains an `event_sources` row (`kind='corroboration'`), | |
| 78 | + `payload.sources[]`, `payload.corroborations` and the confidence bump. Clusters track `source_count`, `surfaces[]`, `confidence`. | |
| 79 | + | |
| 80 | +## LLM enrichment hooks | |
| 81 | + | |
| 82 | +- `classify_change` (small model) when no deterministic event or ambiguous surface, significance ≥ `CA_LLM_MIN_SIGNIFICANCE`, within | |
| 83 | + `CA_LLM_DAILY_BUDGET`, only when `settings.llm_configured`. Material classifications create `origin='llm'` events (dedupe | |
| 84 | + `sha(company, 'llm', change, subtype)`); when a deterministic event of the same subtype already exists for the change, the classification enriches | |
| 85 | + it instead (`origin='hybrid'`, `payload.llm_classification`). | |
| 86 | +- `summarize_event` (medium model) for legal, homepage/messaging and news/IR events that carry diff content → `events.summary`, `origin='hybrid'`, | |
| 87 | + `payload.llm` (key points or legal sections + materiality), model/prompt/schema versions recorded. See docs/LLM.md. | |
added
docs/LLM.md
+91 −0
@@ -0,0 +1,91 @@ | ||
| 1 | +# LLM enrichment — gateway, prompts, worker, budgets | |
| 2 | + | |
| 3 | +LLMs are enrichment, never the crawler (spec §2.4, §24–25). Everything deterministic happens first (`docs/EVENT-TAXONOMY.md`); the model only | |
| 4 | +sees bounded, already-diffed content, answers in strict JSON, and every output is validated, versioned and attributed. | |
| 5 | + | |
| 6 | +## Configuration (`config.Settings`) | |
| 7 | + | |
| 8 | +| Setting | Default | Meaning | | |
| 9 | +|---|---|---| | |
| 10 | +| `CA_LLM_BASE_URL` / `CA_LLM_API_KEY` | `https://www.llm-api.io/v1` (`.env`) / key in `deploy/.llm-key` (git-ignored) | OpenAI-compatible endpoint; `settings.llm_configured` = enabled **and** URL set | | |
| 11 | +| `CA_LLM_ENABLED` | true | master switch — when off, jobs stay pending and the worker logs once | | |
| 12 | +| `CA_LLM_SMALL_MODEL` / `MEDIUM` / `LARGE` / `EMBEDDING` | `qwen3-4b-instruct-2507-4bit` / `qwen3.6-35b-a3b-4bit` / `qwen3.8-27b-4bit` / `qwen3-embedding-0.6b-8bit` | tiers used by tasks | | |
| 13 | +| `CA_LLM_DAILY_BUDGET` | 1500 | max jobs per UTC day (enqueue side counts created jobs, worker side counts finished jobs) | | |
| 14 | +| `CA_LLM_MIN_SIGNIFICANCE` | 0.40 | changes below this never reach a model | | |
| 15 | +| `CA_LLM_TIMEOUT` | 300 s | per request (the server loads models on demand: first call 30–60 s) | | |
| 16 | +| `CA_LLM_MAX_TRIES` / `BACKOFF_INITIAL_S` / `BACKOFF_MAX_S` | 6 / 15 / 120 | 429/503/5xx/timeouts: 15 → 30 → 60 → 120 → 120 s (Retry-After honoured) | | |
| 17 | +| `CA_LLM_JOB_MAX_ATTEMPTS` | 3 | retryable failures re-queue the job until this | | |
| 18 | +| `CA_LLM_CONTEXT_BLOCK_BYTES` | 3072 | before/after text per block sent to the model (≤ 12 blocks) | | |
| 19 | +| `CA_PROMPTS_DIR` | `<repo>/prompts` | prompt files location | | |
| 20 | + | |
| 21 | +## Gateway (`services/llm/gateway.py`) | |
| 22 | + | |
| 23 | +`LLMProvider` (abstract) → `OpenAICompatibleProvider` (httpx, `/chat/completions`, `/models`, `/embeddings`): | |
| 24 | + | |
| 25 | +- `complete_json(tier, system, user, schema, max_tokens)` — `response_format={"type": "json_object"}`, the compact JSON schema of the pydantic model | |
| 26 | + appended to the system prompt, `<think>` blocks / code fences / prose tolerated when extracting the object, **strict pydantic validation**, and | |
| 27 | + **one repair round-trip** (the model sees its own answer plus the validation error). Still invalid → `LLMValidationError` (job failed, nothing stored). | |
| 28 | +- `complete_text(...)`, `embed(texts)` (embedding tier; reserved for later search/clustering work), `health()` (model list + latency). | |
| 29 | +- Returns `LLMResult(data, model, request_tokens, response_tokens, latency_ms, attempts, repaired)`; tokens go to `llm_jobs` and `cost_ledger` | |
| 30 | + (`dimension='llm'`, key = model, units = tokens). | |
| 31 | +- `get_provider()` / `set_provider()` for a process-wide instance or test doubles. The gateway is the only module allowed to call the endpoint. | |
| 32 | + | |
| 33 | +Sticky model per stream: the worker drains one job kind at a time (classification → small model, summaries → medium model) so the on-demand | |
| 34 | +server does not evict models between calls. Use `catlas llm-test ["prompt"] [--tier medium]` for a live health + round-trip check. | |
| 35 | + | |
| 36 | +## Prompts (`prompts/<task>/v<N>.md`) | |
| 37 | + | |
| 38 | +Loaded by `services/llm/prompts.load_prompt(task, version=None)` (latest `vN` by default); the leading `<!-- … -->` comment is metadata, the rest is | |
| 39 | +the **system** prompt. The user message is always a compact JSON context — untrusted page text is never interpolated into instructions. | |
| 40 | +`events.prompt_version` stores `task/vN`; bump the file (v2.md) when wording changes materially and keep v1 for reproducibility. | |
| 41 | + | |
| 42 | +| Prompt | Tier | Schema (`services/llm/schemas.py`, version) | Used by | | |
| 43 | +|---|---|---|---| | |
| 44 | +| `change-classifier/v1` | small | `ChangeClassification` (`classify-v1`): subtype ∈ `EVENT_SUBTYPES` else OTHER, importance, confidence, title ≤ 120, summary ≤ 400, old/new value, entities, tags, language, `is_material` | `classify_change` jobs | | |
| 45 | +| `event-summarizer/v1` | medium | `EventSummary` (`summary-v1`): summary ≤ 400, key_points ≤ 5, confidence, language | `summarize_event` (non-legal) | | |
| 46 | +| `legal-diff/v1` | medium | `LegalDiffSummary` (`legal-v1`): sections_changed[{section, change}], materiality ∈ editorial/minor/material/unclear, summary, user_impact | `summarize_event` for LEGAL events | | |
| 47 | +| `industry-tagger/v1` | small | `IndustryTags` (`industry-v1`): slugs from the allowed list only, primary, keywords | `classify_industry` jobs → `companies.source_meta.llm_industries` (suggestion only) | | |
| 48 | +| `ask-router/v1` | small | `AskRoute` (`ask-v1`): intent, countries, industries, event types/subtypes, companies, window, keywords, answer_style | `services/llm/ask.route_question` (optional refinement of the deterministic parser) | | |
| 49 | + | |
| 50 | +All schemas reject the forbidden wording (`taxonomy.FORBIDDEN_WORDING`: fired, laid off, shut down…) at validation time; prompts state the ban and | |
| 51 | +the observational language ("detected", "no longer listed"). | |
| 52 | + | |
| 53 | +## Worker (`services/llm/enrich.py`) | |
| 54 | + | |
| 55 | +`run_llm_jobs(limit)` — periodic `llm-enrich` every 15 s, `catlas enrich [--limit] [--once]`: | |
| 56 | + | |
| 57 | +1. Budget check (finished jobs today < `CA_LLM_DAILY_BUDGET`), pick the sticky kind, claim jobs with `FOR UPDATE SKIP LOCKED` (`attempts += 1`). | |
| 58 | +2. Build the bounded context: company meta (name, domain, country, industries, description ≤ 400), surface, source URL, significance/kind, block | |
| 59 | + counts, ≤ 12 diff blocks (before/after ≤ 3 kB each), structured deltas with lists trimmed to 10 items. | |
| 60 | +3. Call the gateway, validate, write: | |
| 61 | + - `classify_change` → material & non-OTHER & importance ≥ 0.25 → new event `origin='llm'` (confidence capped at 0.9, `status='review'` + | |
| 62 | + review_queue when < 0.5, dedupe `sha(company, 'llm', change, subtype)`, event_sources, clustering, alerts). If a deterministic event with the | |
| 63 | + same subtype already exists for that change, it is enriched instead (`origin='hybrid'`, `payload.llm_classification`). `changes.status='enriched'`. | |
| 64 | + - `summarize_event` → `events.summary` (previous kept in `payload.summary_prev`), `origin='hybrid'`, `model_provider/model_name/prompt_version/ | |
| 65 | + schema_version`, `payload.llm` (key points, or legal sections + materiality + tag `materiality:<x>`). | |
| 66 | + - `classify_industry` → `companies.source_meta.llm_industries` (never overwrites registry industries). | |
| 67 | +4. `llm_jobs` gets status, model, prompt_version, tokens, latency, `result` JSON or `error`; retryable failures go back to `pending` until | |
| 68 | + `CA_LLM_JOB_MAX_ATTEMPTS`, a retryable failure also stops the current batch (the server is unhealthy — the periodic task returns). | |
| 69 | + | |
| 70 | +Enqueue side (`services/events.py`): `classify_change` for meaningful+ changes with no deterministic event or an ambiguous surface; | |
| 71 | +`summarize_event` for legal / homepage / messaging / news / IR events that have diff content. Both only when `settings.llm_configured`, above | |
| 72 | +`CA_LLM_MIN_SIGNIFICANCE`, within budget, and never twice for the same ref. | |
| 73 | + | |
| 74 | +## `/ask` helper (`services/llm/ask.py`) | |
| 75 | + | |
| 76 | +`parse_question(q, industries=, countries=)` is deterministic: countries (names/demonyms + registry table), industries (registry slugs/names), | |
| 77 | +intents → event types/subtypes/tags (hiring, pricing ± increase/decrease, AI, launch, leadership, expansion, legal, developer, financing, M&A, | |
| 78 | +communication), windows ("last 3 months", "this week", "since 2024", today/yesterday), quoted company names, answer style (list / count / compare / | |
| 79 | +timeline / trend), `min_importance` for "major/important". `route_question()` adds the LLM refinement only when configured and only to tighten | |
| 80 | +filters; on any failure the deterministic interpretation is returned. `build_answer()` phrases counts the API measured — it never invents results. | |
| 81 | + | |
| 82 | +## Live check (2026-09-12) | |
| 83 | + | |
| 84 | +`catlas llm-test` against `https://www.llm-api.io/v1`: health OK (15 models listed), small model round-trip **4.5 s**, 96 prompt / 19 completion | |
| 85 | +tokens, no repair needed. Live enrichment of factory changes: `classify_change` (small) 3.9 s; three `summarize_event` (medium, incl. legal-diff) | |
| 86 | +3.1–8.2 s, 0 failures. | |
| 87 | + | |
| 88 | +## Retention & cost | |
| 89 | + | |
| 90 | +`llm_jobs` finished > `CA_RETENTION_LLM_JOBS_DAYS` (60) are summarised into `cost_ledger` (`archived:<model>` tokens, `archived-jobs:<model>` count) | |
| 91 | +and deleted (`services/retention.py`); events keep their model/prompt attribution forever. | |
added
docs/OPERATIONS.md
+72 −0
@@ -0,0 +1,72 @@ | ||
| 1 | +# Company Atlas — operations | |
| 2 | + | |
| 3 | +## API | |
| 4 | + | |
| 5 | +The API is one FastAPI process (`catlas api`, uvicorn) bound to loopback; Next.js proxies `/api/v1/*` to it. It is stateless: every cache is | |
| 6 | +in-process (`api/common.TTLCache`, cleared with `POST /api/v1/admin/cache/clear` or on restart), so several workers may run behind PM2 and | |
| 7 | +the rate limiter buckets are per process (effective limits scale with the number of workers — keep `--workers 1` unless traffic requires | |
| 8 | +more, or terminate rate limiting at the tunnel). | |
| 9 | + | |
| 10 | +### Environment variables (read by `config.Settings`, prefix `CA_`) | |
| 11 | + | |
| 12 | +| Variable | Default | Used by the API for | | |
| 13 | +|---|---|---| | |
| 14 | +| `DATABASE_URL` | `postgresql+asyncpg://companyatlas:companyatlas@127.0.0.1:5432/companyatlas` | every query (pool 8 + 8 overflow per process) | | |
| 15 | +| `CA_API_HOST` / `CA_API_PORT` | `127.0.0.1` / `8371` (prod `8361`) | bind address of `catlas api` | | |
| 16 | +| `CA_ADMIN_TOKEN` | *(empty → every `/admin/*` call is 401)* | `X-CA-Admin-Token` (also bypasses rate limits) | | |
| 17 | +| `CA_SITE_URL` | `https://www.company-atlas.co` | CORS allow-list (plus the localhost dev/prod ports) | | |
| 18 | +| `CA_DATA_DIR` | `./data` | object store read by `/snapshots/{id}` (text/blocks) and `/stats.archive` | | |
| 19 | +| `CA_SEO_MIN_EVENTS` / `CA_SEO_MIN_SENSORS` | `1` / `3` | `/sitemap` inclusion rule (`companies.indexed` or ≥ events **and** ≥ active sensors) | | |
| 20 | +| `CA_MIN_INTERVAL_S` / `CA_MAX_INTERVAL_S` | `900` / `604800` | clamp for `POST /admin/sensors/{id}/set_interval` | | |
| 21 | +| `CA_LLM_DAILY_BUDGET`, `CA_LLM_BASE_URL`, `CA_LLM_ENABLED` | `1500`, *(empty)*, `true` | `/admin/overview.llm.budget_left`, `/ask` LLM refinement availability | | |
| 22 | +| `CA_NOISE_THRESHOLD` … `CA_CRITICAL_THRESHOLD` | `0.20 / 0.40 / 0.65 / 0.85` | `/methodology.significance_bands`, `/admin/costs` meaningful-event denominator | | |
| 23 | +| `CA_LOG_JSON` | `true` | JSON logs (`service=ca-api`) | | |
| 24 | + | |
| 25 | +Run: `.venv/bin/catlas api` (dev, port 8371) · `.venv/bin/catlas api --port 8361` (prod under PM2 `company-atlas-api`). | |
| 26 | +Health: `GET /health` → `{status, db, version, api_version, uptime_s}`; `GET /ready` → `{ready}` (DB reachable). Docs: `/api/v1/docs`, | |
| 27 | +schema `/api/v1/openapi.json`. The scheduler's heartbeat (`settings_kv['scheduler:heartbeat']`) surfaces as `/system.scheduler_last_tick_at`. | |
| 28 | + | |
| 29 | +### API keys and rate-limit tiers | |
| 30 | + | |
| 31 | +Anonymous traffic is limited to 120 requests/min per client IP (first hop of `X-Forwarded-For` — Caddy on the MacLustr Tunnel sets it). | |
| 32 | +Keys raise the tier (`authenticated` 600/min, `paid` 3 000/min, `internal` unlimited). Keys are stored as sha256 hashes in `api_keys`; | |
| 33 | +the raw key is printed **once**: | |
| 34 | + | |
| 35 | +```bash | |
| 36 | +.venv/bin/catlas api-key create "acme research" --tier paid # prints ca_paid_… once; store it in the customer's vault | |
| 37 | +.venv/bin/catlas api-key list [--include-revoked] # id, prefix, tier, last used, request count | |
| 38 | +.venv/bin/catlas api-key revoke key_01… # soft revoke (row kept, `revoked_at` set); cached for ≤ 60 s in running API processes | |
| 39 | +``` | |
| 40 | + | |
| 41 | +Clients send `X-CA-API-Key: <key>`; responses carry `X-RateLimit-Limit`, `X-RateLimit-Remaining`, `X-RateLimit-Tier`; a `429` carries | |
| 42 | +`Retry-After`. Usage counters (`last_used_at`, `request_count`) are flushed every ~30 s per process. There is no self-service key issuance — | |
| 43 | +create keys from the M2U64 shell (or any host with `DATABASE_URL`). | |
| 44 | + | |
| 45 | +### Caches and freshness | |
| 46 | + | |
| 47 | +| Endpoint | TTL | Notes | | |
| 48 | +|---|---|---| | |
| 49 | +| `/pulse`, `/stats` | 60 s | homepage / counters; `archive` stats refresh every 10 min (`settings_kv['archive:stats']` if the ops job writes it, else a directory walk in a thread) | | |
| 50 | +| `/system` | 30 s | public health | | |
| 51 | +| `/rankings` | 120 s | per (kind, window, country, industry, limit) | | |
| 52 | +| `/industries*`, `/countries*`, `/index`, `/map`, `/trends`, `/events/types`, `/events/summary`, `/stats/history` | 300 s | aggregate producers in `api/aggregates.py` | | |
| 53 | +| `/sitemap` | 600 s | 5 000 entries per page | | |
| 54 | +| `/live*`, owner, admin | none (`no-store`) | | | |
| 55 | + | |
| 56 | +After bulk backfills or manual corrections: `curl -X POST -H "X-CA-Admin-Token: $CA_ADMIN_TOKEN" localhost:8371/api/v1/admin/cache/clear` | |
| 57 | +(optionally `?prefix=pulse`). Retract/restore of events clears the caches automatically. | |
| 58 | + | |
| 59 | +### Pending SQL | |
| 60 | + | |
| 61 | +`migrations/pending/api.sql` adds two trigram GIN indexes (`people.name`, `products.name`) used by `/search?types=people,products`. Apply | |
| 62 | +with `psql -f` (uses `CREATE INDEX CONCURRENTLY`, so outside a transaction) or fold into the next numbered Alembic migration. The endpoints | |
| 63 | +work without them (sequential scans on small tables). | |
| 64 | + | |
| 65 | +### Troubleshooting | |
| 66 | + | |
| 67 | +- `503`/`degraded` on `/health` → Postgres unreachable; the API keeps serving cached aggregates until they expire. | |
| 68 | +- `429` for the web app itself → the Next.js server is one IP; give it an `internal` key (`API_KEY` env in `apps/web`) or terminate limits upstream. | |
| 69 | +- `/snapshots/{id}` returns `text: null` → the object key is missing from `CA_DATA_DIR/objects` (different data dir or pruned by retention); | |
| 70 | + metadata still renders and `/snapshots/{a}/diff/{b}` falls back to the stored `changes.diff` when present. | |
| 71 | +- `/ask` engine is always `deterministic` → `CA_LLM_BASE_URL`/`CA_LLM_API_KEY` unset or the small model unreachable (the deterministic parser is | |
| 72 | + the designed fallback; nothing is fabricated either way). | |
added
docs/SCORING.md
+78 −0
@@ -0,0 +1,78 @@ | ||
| 1 | +# Scoring — metrics, indices, baselines, signals, trends | |
| 2 | + | |
| 3 | +Owner: `services/metrics.py` (`metrics-v1`, CCI `cci-v1`), `services/signals.py` (`signals-v1`), `services/trends.py` (`trends-v1`). | |
| 4 | +Parameters live in `taxonomy.METRIC_PARAMS` / `taxonomy.CCI_WEIGHTS` / `config.Settings`; **bump the formula version whenever they change** | |
| 5 | +(spec §177). Every `metrics_current` row stores `value`, `confidence`, `inputs` (the numbers the formula saw), `formula_version`, `computed_at`; | |
| 6 | +`metric_series` keeps one point per company/metric/day. **No inputs → no row** (never 0-as-measured). | |
| 7 | + | |
| 8 | +Schedules: `metrics-hourly` (`CA_METRICS_CRON`, default `7 * * * *`) for companies with activity in the last `CA_METRICS_ACTIVE_WINDOW_DAYS` (90); | |
| 9 | +`daily-aggregates` (`CA_DAILY_CRON`, `20 0 * * *`) = catch-up of missed days + all-company recompute; `signals` hourly; `trends` at :25. | |
| 10 | +CLI: `catlas metrics [--all] [--company slug]`, `catlas daily [--day|--catch-up]`, `catlas signals`, `catlas trends`. | |
| 11 | + | |
| 12 | +Notation: `sat(x, k) = 1 − e^(−x/k)` (0–1, x=k → 0.63); `decay(age) = e^(−age/τ)`; ages in days. | |
| 13 | + | |
| 14 | +## Company metrics (`metrics_current`) | |
| 15 | + | |
| 16 | +| Metric | Formula | Inputs stored | Emitted when | | |
| 17 | +|---|---|---|---| | |
| 18 | +| `activity_score` (0–100) | `raw = Σ_changes w(kind)·decay(age, τ=10) + Σ_events importance·decay(age, τ)` over 30 d, `w` = meaningful 1 / major 2 / critical 3; `density = raw / active_sensors^0.5` (coverage normalisation); `score = 100·log1p(density)/log1p(6)` capped at 100 | changes_30d, events_30d, raw_changes, raw_events, active_sensors, density, τ, exponent, D_max | ≥ 1 active sensor or any activity | | |
| 19 | +| `open_jobs` | count of `jobs.status='open'` | jobs_new_30d, jobs_removed_30d, remote_ratio, by_country, by_department, ai_open | company has a careers/jobs_board sensor or any job rows | | |
| 20 | +| `hiring_momentum_{7,30,90}d` (%) | `(open_now − open_then)/open_then × 100`, `open_then` reconstructed from `first_seen_at ≤ t−N and (removed_at is null or > t−N)` | open_now, open_then, window (+ 30 d extras) | `open_then ≥ 3` (hiring_min_listings) — no % from 0 or 2 listings | | |
| 21 | +| `ai_adoption` (0–100) | components: jobs `min(1, ai_open/open × 2)` (w 0.5), events `sat(ai_events_90d, 3)` (w 0.25: AI_HIRING/AI_LAUNCH/tag `ai`), keywords `sat(hits, 5)` (w 0.25: `AI_KEYWORDS` in product names, 90 d news titles, latest snapshot titles); weights renormalised over available components | ai_open, open_jobs, ai_events_90d, keyword_hits, components, weights_used | any component has inputs | | |
| 22 | +| `product_velocity` | `100·sat(Σ importance·decay(age, 30), 6)` over 90 d PRODUCT + changelog/docs/API/SDK subtypes | events_90d, surfaces | products/changelog/docs/api/developer/services sensor or a PRODUCT event | | |
| 23 | +| `developer_momentum` | `100·sat(Σ DEVELOPER importance·decay + 0.5·meaningful dev-surface changes 30 d, 6)` | events_90d, meaningful_changes_30d, surfaces | docs/api/developer/changelog sensor or DEVELOPER event | | |
| 24 | +| `communication_activity` | `100·sat(max(news_30d, comm_events_30d) + 0.5·min(·), 8)` | news_items_30d, events_30d | newsroom/blog/feed/IR/research sensor or news rows | | |
| 25 | +| `pricing_activity` | `100·sat(Σ PRICING importance·decay + 0.25·pricing changes 30 d, 3)` | events_90d, meaningful_changes_30d | pricing sensor or PRICING event | | |
| 26 | +| `leadership_activity` | `100·sat(Σ LEADERSHIP importance·decay, 3)` | events_90d | leadership sensor or LEADERSHIP event | | |
| 27 | +| `geo_expansion` | `100·sat(new_locations_90d + 3·|new countries| + Σ COUNTRY_EXPANSION importance, 4)`; new countries = countries whose first location **or** first job appeared within 90 d | new_locations_90d, new_countries_90d, countries_listed, location_events_90d | locations/contact/about/careers sensor, location rows or jobs | | |
| 28 | +| `corporate_change_index` (`cci-v1`) | `Σ w_i·c_i / Σ w_i` over **available** components with `CCI_WEIGHTS` (hiring 0.25 · product 0.20 · geo 0.15 · leadership 0.15 · developer 0.10 · communication 0.10 · pricing 0.05); hiring momentum (%) mapped to `50 + clamp(m, −100, 100)/2`; confidence `0.4 + 0.6·weight_coverage` | components, weights, weight_coverage | ≥ 1 component | | |
| 29 | +| `anomaly_score` (z) | `(this_week_meaningful − mean) / max(stddev, 0.5)` against `baselines.meaningful_changes_weekly` | this_week, mean, stddev, samples | baseline with ≥ 4 samples | | |
| 30 | +| `historical_coverage` (0–100) | `100·(0.5·min(1, observed_ok/expected) + 0.3·continuity + 0.2·min(1, surfaces/8))`; `expected = Σ_sensors age_s / current_interval_s`; `continuity = days_with_observations / days_since_first_observed` | observed, expected, days_with_obs, days_since_first, surfaces | `first_observed_at` set and ≥ 1 sensor | | |
| 31 | + | |
| 32 | +Confidence conventions: activity `min(0.95, 0.5 + 0.05·sensors)`; hiring 0.9 with a structured jobs_board sensor else 0.75; AI adoption 0.6 | |
| 33 | +(inferred from public signals only); saturating scores 0.7; coverage 0.8. | |
| 34 | + | |
| 35 | +## Daily aggregates | |
| 36 | + | |
| 37 | +`compute_daily(day)` (UTC day, idempotent upserts): | |
| 38 | + | |
| 39 | +- `company_daily` — observations, changes, meaningful_changes, events (+ `events_by_type`), jobs_open / jobs_new / jobs_removed / jobs_ai_open at end of | |
| 40 | + day (reconstructed from `first_seen_at` / `removed_at`), news_items, sensors_active. Only companies with activity that day get a row. | |
| 41 | +- `global_daily` — sums + `companies_active` / `sensors_active` (distinct in observations that day; falls back to active sensors when no observation | |
| 42 | + rows exist), `events_by_type`, `jobs_open`, `by_country` / `by_industry` (`{companies, meaningful_changes, events}` per key) and the | |
| 43 | + **Global Corporate Activity Index**: `ratio = meaningful_changes / sensors_active`; `baseline = mean(ratio)` over the trailing 28 days with | |
| 44 | + sensors; `activity_index = ratio / baseline × 100` (baseline = 100). Null when no baseline exists yet (never fabricated). | |
| 45 | +- `baselines` — per company, weekly buckets of `company_daily` over `CA_BASELINE_WINDOW_DAYS` (56): `meaningful_changes_weekly`, `jobs_new_weekly`, | |
| 46 | + `news_weekly` → population mean / stddev / samples. Requires ≥ 2 observed weeks. Consumed by `anomaly_score`, `HIRING_SURGE` and signals. | |
| 47 | +- `compute_daily_catch_up()` fills every day between the last computed day (or the first observation/change/event) and yesterday. | |
| 48 | + | |
| 49 | +Coverage normalisation (spec §146): activity uses per-sensor density, the global index divides by active sensors and compares to a trailing | |
| 50 | +baseline, trends divide mentions by active companies, and the CCI renormalises over available components — growth of the sensor network alone | |
| 51 | +must not move any index. | |
| 52 | + | |
| 53 | +## Signals (`signals`, `signals-v1`) | |
| 54 | + | |
| 55 | +Labelled as signals, with `strength` 0–1, `confidence`, `explanation`, `evidence` (event ids + metric values), `window_days`, `expires_at` | |
| 56 | +(`CA_SIGNALS_TTL_DAYS` = 14). One active row per (company, kind): re-detection updates it, disappearance sets `status='expired'`. | |
| 57 | + | |
| 58 | +| kind | trigger | | |
| 59 | +|---|---| | |
| 60 | +| `hiring_surge` | `hiring_momentum_30d ≥ +30 %` with ≥ 10 open, or a `HIRING_SURGE` event in the window | | |
| 61 | +| `hiring_freeze` | `hiring_momentum_30d ≤ −30 %` with ≤ 1 new listing, or a `HIRING_FREEZE_SIGNAL` event — worded "listings no longer visible" | | |
| 62 | +| `launch_buildup` | within 14 d ≥ 3 of {product, docs/API, changelog, careers, messaging} categories including product, or ≥ 4 categories — "Possible launch preparation signal" | | |
| 63 | +| `expansion` | `COUNTRY_EXPANSION` event, or new countries (locations or job countries, 90 d) with location events / new jobs | | |
| 64 | +| `pricing_migration` | ≥ 2 PRICING events in 30 d, or a tier added **and** a tier removed | | |
| 65 | +| `developer_push` | ≥ 3 DEVELOPER events in 30 d or `developer_momentum ≥ 60` | | |
| 66 | +| `enterprise_repositioning` | new contact-sales tier(s) in 30 d, or ≥ 2 events mentioning enterprise / SSO / audit logs / contact sales | | |
| 67 | +| `ai_acceleration` | ≥ 2 AI-tagged events in 30 d, or ≥ 25 % of ≥ 4 new listings are AI-related | | |
| 68 | +| `abnormal_activity` | `anomaly_score ≥ CA_ANOMALY_Z` (2.5) | | |
| 69 | + | |
| 70 | +Scope aggregates (`scope = industry | country`): ≥ 3 companies with the same active kind → one aggregate signal per (scope, key, kind) with the | |
| 71 | +company slugs as evidence; aggregates expire when the support drops below the threshold. | |
| 72 | + | |
| 73 | +## Trends (`trends`, `trends-v1`) | |
| 74 | + | |
| 75 | +Terms = lowercase 1–3-grams from event titles (prefix like "News release:" removed) and news titles of the day, company names removed, | |
| 76 | +stopwords excluded at n-gram edges, tokens ≥ 3 chars (unigrams ≥ 4). A term is kept for a day only when ≥ `CA_TRENDS_MIN_COMPANIES` (3) distinct | |
| 77 | +companies used it (`mentions`, `companies`). Momentum for a window W ∈ {7, 30, 90}: `rate = mentions / mean(companies_active)` over the window, | |
| 78 | +`momentum = (rate_now − rate_prev) / max(rate_prev, ε)`; snapshots are stored in `settings_kv['trends:momentum:<W>d']` for the API. | |
added
migrations/pending/api.sql
+4 −0
@@ -0,0 +1,4 @@ | ||
| 1 | +-- API (search): trigram indexes for people / product name search (`/search?types=people,products`). | |
| 2 | +-- Apply once, outside a transaction (CONCURRENTLY), or fold into the next numbered migration. | |
| 3 | +create index concurrently if not exists people_name_trgm_idx on people using gin (name gin_trgm_ops); | |
| 4 | +create index concurrently if not exists products_name_trgm_idx on products using gin (name gin_trgm_ops); | |
added
prompts/ask-router/v1.md
+20 −0
@@ -0,0 +1,20 @@ | ||
| 1 | +<!-- ask-router v1 · model tier: small · schema: AskRoute (ask-v1) --> | |
| 2 | +You translate a natural-language question about companies into structured search filters for the Company Atlas API. You NEVER answer the question yourself and you NEVER invent companies, events or numbers — the platform runs the filters against its own database and links every result to its source. | |
| 3 | + | |
| 4 | +You receive a JSON context: the user's question, the filters a deterministic parser already extracted (`parsed`), the allowed event types and subtypes, the known industry slugs (`industries`) and country codes with names (`countries`). | |
| 5 | + | |
| 6 | +Produce an `AskRoute` object: | |
| 7 | +- `intent`: one of search, hiring, pricing, ai, launch, leadership, expansion, legal, developer, compare, trend, count. | |
| 8 | +- `countries`: ISO-3166 alpha-2 codes mentioned or clearly implied ("Canadian companies" → ["CA"], "Europe" → leave empty and add keyword "europe"). | |
| 9 | +- `industries`: slugs from the provided list only. | |
| 10 | +- `event_types` / `event_subtypes`: from the allowed lists only. | |
| 11 | +- `companies`: company names literally mentioned (as written), max 8. | |
| 12 | +- `window_days`: the time window when the question states one ("last month" → 30, "this week" → 7, "this year" → 365, "since 2024" → days from 2024-01-01 to today); null otherwise. | |
| 13 | +- `keywords`: up to 8 lowercase free-text terms that should be matched against titles (technologies, product names, roles). | |
| 14 | +- `answer_style`: list (default), count ("how many"), compare ("X vs Y"), timeline ("when did", "history of"). | |
| 15 | +- `confidence` (0–1). | |
| 16 | + | |
| 17 | +Rules: | |
| 18 | +1. Start from `parsed` and only add or tighten filters that the question clearly supports. When in doubt, leave a filter empty rather than guessing. | |
| 19 | +2. Do not translate vague adjectives into filters ("big companies", "interesting") — put them in keywords or ignore them. | |
| 20 | +3. Output one JSON object only. | |
added
prompts/change-classifier/v1.md
+18 −0
@@ -0,0 +1,18 @@ | ||
| 1 | +<!-- change-classifier v1 · model tier: small · schema: ChangeClassification (classify-v1) --> | |
| 2 | +You are the change classifier of Company Atlas, a platform that observes public company web pages and records how companies change over time. | |
| 3 | + | |
| 4 | +You receive a JSON context describing ONE detected change on ONE public page: company metadata, the page surface (careers, pricing, homepage, legal_terms…), the blocks of text that were added, removed or modified, and any structured deltas already extracted deterministically. | |
| 5 | + | |
| 6 | +Your task: decide whether this change corresponds to a corporate event, and describe it as a structured JSON object. | |
| 7 | + | |
| 8 | +Rules — follow all of them: | |
| 9 | +1. Choose `event_subtype` ONLY from this list (or `OTHER`): PRODUCT_LAUNCH, NEW_PRODUCT, PRODUCT_REMOVED, PRODUCT_RENAME, PRODUCT_UPDATE, FEATURE_LAUNCH, PRICE_INCREASE, PRICE_DECREASE, NEW_PRICING_TIER, PRICING_TIER_REMOVED, PRICING_CHANGE, NEW_JOB, JOB_REMOVED, JOB_COUNT_INCREASE, JOB_COUNT_DECREASE, AI_HIRING, NEW_EXECUTIVE, EXECUTIVE_NO_LONGER_LISTED, EXECUTIVE_TITLE_CHANGE, LEADERSHIP_CHANGE, NEW_OFFICE, NEW_LOCATION, OFFICE_REMOVED, COUNTRY_EXPANSION, FUNDING_ROUND, IPO, ACQUISITION, DIVESTITURE, MERGER, NEW_PARTNERSHIP, PARTNERSHIP_ENDED, BRAND_REPOSITIONING, STRATEGY_UPDATE, ENTERPRISE_REPOSITIONING, TECHNOLOGY_ADOPTION, AI_LAUNCH, TERMS_CHANGE, PRIVACY_POLICY_CHANGE, LEGAL_UPDATE, REGULATORY, CAMPAIGN_LAUNCH, MESSAGING_CHANGE, API_LAUNCH, API_CHANGE, SDK_RELEASE, DOCUMENTATION_CHANGE, CHANGELOG_ENTRY, SECURITY_INCIDENT, SECURITY_UPDATE, OUTAGE, OPERATIONS_UPDATE, SUSTAINABILITY_UPDATE, EARNINGS_RELEASE, INVESTOR_UPDATE, NEWS_RELEASE, BLOG_POST, WEBSITE_CHANGE, HOMEPAGE_REDESIGN, DOC_CHANGE. | |
| 10 | +2. Only describe what the changed text itself shows. Never infer internal decisions, motives, headcount, revenue or outcomes. If the change is cosmetic (layout, dates, cookie banners, navigation, footers, typos) set `is_material` to false, `event_subtype` to `OTHER` and `importance` ≤ 0.2. | |
| 11 | +3. Wording must be careful and observational: "detected", "observed", "now listed", "no longer listed", "appears". NEVER write "fired", "laid off", "layoffs", "shut down", "bankrupt" or any statement about people losing jobs. | |
| 12 | +4. `title` ≤ 120 characters, factual, starts with the subject (e.g. "Enterprise plan now lists SSO and audit logs"). `summary` ≤ 400 characters, 1–3 sentences, no marketing tone, no speculation. | |
| 13 | +5. `old_value` / `new_value` are short literal values when the change is a replacement (price, title, name); otherwise null. | |
| 14 | +6. `importance` (0–1): how much a company analyst would care. Pricing, leadership, launches, expansion, legal terms are usually 0.5–0.9; routine blog posts 0.2–0.4; cosmetic 0–0.2. | |
| 15 | +7. `confidence` (0–1): how sure you are the subtype is right given the evidence. Below 0.5 when the text is ambiguous or partial. | |
| 16 | +8. `entities`: an object with optional arrays `products`, `people` (name, title), `locations` (city, country), `plans` (plan_name, price), `amounts`, `organizations`. Only entities literally present in the text. | |
| 17 | +9. `tags`: up to 8 lowercase keywords (e.g. "ai", "enterprise", "api", "europe"). `language`: ISO-639-1 code of the page text. | |
| 18 | +10. Output one JSON object and nothing else. | |
added
prompts/event-summarizer/v1.md
+17 −0
@@ -0,0 +1,17 @@ | ||
| 1 | +<!-- event-summarizer v1 · model tier: medium · schema: EventSummary (summary-v1) --> | |
| 2 | +You write the short neutral summary shown under a Company Atlas event card. Company Atlas records how public company web pages change; every event links back to its source page, so the summary must only restate what the source text shows. | |
| 3 | + | |
| 4 | +You receive a JSON context: company metadata, the event already detected (type, subtype, title), the page surface, the text blocks that were added / removed / modified, and any structured deltas (jobs, plans, people, locations, news). | |
| 5 | + | |
| 6 | +Write: | |
| 7 | +- `summary`: 1–3 sentences, ≤ 400 characters, plain and factual, in English even if the source is in another language (mention the source language if it is not English). Say what is now on the page versus before. Do not repeat the title verbatim; add the concrete details (names, plan names, prices, cities, product names, section names). | |
| 8 | +- `key_points`: up to 5 short bullets (≤ 160 characters each) with concrete observed details. Omit if nothing beyond the summary. | |
| 9 | +- `confidence` (0–1): how well the provided text supports the summary. | |
| 10 | +- `language`: ISO-639-1 code of the source text. | |
| 11 | + | |
| 12 | +Strict rules: | |
| 13 | +1. No speculation about causes, strategy, intent, financial impact or people's employment status. Use "detected", "observed", "now lists", "no longer lists", "appears". | |
| 14 | +2. NEVER use "fired", "laid off", "layoffs", "shut down", "bankrupt", "collapse". A person or listing disappearing from a page is described as "no longer listed". | |
| 15 | +3. Do not invent numbers, dates, names or products that are not literally in the context. If the context is too thin, say so in the summary ("The page changed materially; the visible text does not state the reason."). | |
| 16 | +4. No marketing tone, no adjectives like "exciting", "innovative", "leading". | |
| 17 | +5. Output a single JSON object only. | |
added
prompts/industry-tagger/v1.md
+16 −0
@@ -0,0 +1,16 @@ | ||
| 1 | +<!-- industry-tagger v1 · model tier: small · schema: IndustryTags (industry-v1) --> | |
| 2 | +You assign industry tags to a company for Company Atlas using only public, observable evidence. | |
| 3 | + | |
| 4 | +You receive a JSON context: company name, canonical domain, the public description (if any), the homepage title and meta description, the first product names observed, up to 10 recent job titles, and the list of allowed industry slugs with their names (`allowed_industries`). | |
| 5 | + | |
| 6 | +Produce: | |
| 7 | +- `industries`: 1–5 slugs chosen ONLY from `allowed_industries`, most specific first. Never invent a slug. | |
| 8 | +- `primary`: the single best slug from `industries`. | |
| 9 | +- `keywords`: up to 10 lowercase keywords that describe what the company visibly does (products, markets, technologies) — derived from the context, not general knowledge. | |
| 10 | +- `confidence` (0–1): 0.9 when the description and products clearly match; ≤ 0.5 when you rely mostly on the name or domain. | |
| 11 | + | |
| 12 | +Rules: | |
| 13 | +1. Prefer what the company says about itself on its own pages over what you believe you know about it. If the context contradicts your prior knowledge, follow the context and lower confidence. | |
| 14 | +2. Holding companies and conglomerates: tag the visible operating activities, then `conglomerate` if it is in the allowed list. | |
| 15 | +3. If the context is empty or unrelated (parked domain, error page), return `industries: []`, `primary: null`, `confidence: 0.1`. | |
| 16 | +4. Output one JSON object only. | |
added
prompts/legal-diff/v1.md
+18 −0
@@ -0,0 +1,18 @@ | ||
| 1 | +<!-- legal-diff v1 · model tier: medium · schema: LegalDiffSummary (legal-v1) --> | |
| 2 | +You summarise changes between two versions of a public legal page (terms of service, privacy policy, acceptable-use, security or data-processing terms) for Company Atlas. Readers are analysts; the summary is shown next to a link to both versions, so it must be precise and restrained. | |
| 3 | + | |
| 4 | +You receive a JSON context: company metadata, the surface (legal_terms / legal_privacy / security), and the blocks that were added, removed or modified, each with its heading path (e.g. "Terms > 7. Termination"). | |
| 5 | + | |
| 6 | +Produce: | |
| 7 | +- `sections_changed`: up to 12 objects `{section, change}` — `section` is the heading path or the section number/name visible in the text, `change` is one sentence (≤ 240 characters) describing what the text now says compared to before ("now requires 30 days' notice, previously 14"). | |
| 8 | +- `materiality`: one of `editorial` (wording, formatting, dates, typos), `minor` (clarifications with no change in rights or obligations), `material` (changes to rights, obligations, fees, data use, liability, jurisdiction, termination, arbitration), `unclear` (not enough text). | |
| 9 | +- `summary`: ≤ 400 characters, neutral, starting with the page and the number of sections changed ("Privacy policy: 3 sections changed. …"). | |
| 10 | +- `user_impact`: ≤ 240 characters, only if the text explicitly changes what users/customers can or must do; else null. | |
| 11 | +- `confidence` (0–1). | |
| 12 | + | |
| 13 | +Rules: | |
| 14 | +1. This is NOT legal advice. Never state whether the change is lawful, fair or compliant; never guess the reason for the change. | |
| 15 | +2. Only describe text present in the context. Missing "before" text means the section is new; missing "after" text means it is no longer present — say "no longer includes", never "removed rights" or similar conclusions. | |
| 16 | +3. Quote key phrases briefly when they matter (fees, notice periods, jurisdictions, data categories). | |
| 17 | +4. Never use "fired", "laid off", "shut down", "bankrupt". | |
| 18 | +5. Output one JSON object only. | |
added
src/companyatlas/api/aggregates.py
+417 −0
@@ -0,0 +1,417 @@ | ||
| 1 | +"""Aggregate producers shared by several routers (`/pulse` composes them): industry & country rows, rankings, trends, map buckets, | |
| 2 | +global stats and the activity index. Every producer degrades to empty lists / `null` values on an empty database — nothing is invented.""" | |
| 3 | +from __future__ import annotations | |
| 4 | + | |
| 5 | +import asyncio | |
| 6 | +from collections import defaultdict | |
| 7 | +from datetime import UTC, datetime, timedelta | |
| 8 | +from typing import Any | |
| 9 | + | |
| 10 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 11 | + | |
| 12 | +from companyatlas import archive | |
| 13 | +from companyatlas.api import queries as q | |
| 14 | +from companyatlas.api import serializers as ser | |
| 15 | +from companyatlas.api.common import cache, cached | |
| 16 | +from companyatlas.db import connection, fetch_all, fetch_one | |
| 17 | +from companyatlas.ids import slugify | |
| 18 | +from companyatlas.taxonomy import METRICS_FORMULA_VERSION, Metric | |
| 19 | + | |
| 20 | +# ------------------------------------------------------------------------------------------------ industries / countries | |
| 21 | + | |
| 22 | +_ROW_METRICS = (Metric.ACTIVITY_SCORE.value, Metric.HIRING_MOMENTUM_30D.value, Metric.AI_ADOPTION.value) | |
| 23 | + | |
| 24 | + | |
| 25 | +def _avg_map(rows: list[dict[str, Any]], key: str) -> dict[str, dict[str, float]]: | |
| 26 | + out: dict[str, dict[str, float]] = defaultdict(dict) | |
| 27 | + for r in rows: | |
| 28 | + if r[key] is not None and r["v"] is not None: | |
| 29 | + out[r[key]][r["metric"]] = float(r["v"]) | |
| 30 | + return out | |
| 31 | + | |
| 32 | + | |
| 33 | +async def industry_rows(conn: AsyncConnection, *, country: str | None = None) -> list[dict[str, Any]]: | |
| 34 | + scope = " and c.country = cast(:country as char(2))" if country else "" | |
| 35 | + params: dict[str, Any] = {"country": country} if country else {} | |
| 36 | + d7, d30 = q.days_ago(7), q.days_ago(30) | |
| 37 | + taxonomy = await fetch_all(conn, "select slug, name, parent_slug, description, sort_order from industries order by sort_order, name") | |
| 38 | + companies = await fetch_all(conn, f"select ind, count(*) as n from companies c, unnest(c.industries) ind where c.status = 'ACTIVE'{scope} " | |
| 39 | + "group by ind", **params) | |
| 40 | + events = await fetch_all(conn, "select ind, e.event_type, count(*) filter (where e.detected_at >= :d7) as n7, count(*) as n30 " | |
| 41 | + "from events e join companies c on c.id = e.company_id, unnest(c.industries) ind " | |
| 42 | + f"where e.status = 'active' and e.detected_at >= :d30{scope} group by ind, e.event_type", d7=d7, d30=d30, **params) | |
| 43 | + metrics = await fetch_all(conn, "select ind, m.metric, avg(m.value) as v from metrics_current m join companies c on c.id = m.company_id, " | |
| 44 | + f"unnest(c.industries) ind where m.metric = any(cast(:ms as text[])){scope} group by ind, m.metric", | |
| 45 | + ms=list(_ROW_METRICS), **params) | |
| 46 | + n_companies = {r["ind"]: int(r["n"]) for r in companies} | |
| 47 | + ev7: dict[str, int] = defaultdict(int) | |
| 48 | + ev30: dict[str, int] = defaultdict(int) | |
| 49 | + types: dict[str, dict[str, int]] = defaultdict(dict) | |
| 50 | + for r in events: | |
| 51 | + ev7[r["ind"]] += int(r["n7"]) | |
| 52 | + ev30[r["ind"]] += int(r["n30"]) | |
| 53 | + types[r["ind"]][r["event_type"]] = int(r["n30"]) | |
| 54 | + avg = _avg_map(metrics, "ind") | |
| 55 | + names = {t["slug"]: t for t in taxonomy} | |
| 56 | + slugs = list(names) + [s for s in n_companies if s not in names] | |
| 57 | + rows = [] | |
| 58 | + for slug in slugs: | |
| 59 | + t = names.get(slug, {}) | |
| 60 | + m = avg.get(slug, {}) | |
| 61 | + rows.append({"slug": slug, "name": t.get("name") or slug.replace("-", " ").title(), "parent_slug": t.get("parent_slug"), | |
| 62 | + "description": t.get("description"), "companies": n_companies.get(slug, 0), "events_7d": ev7.get(slug, 0), | |
| 63 | + "events_30d": ev30.get(slug, 0), "hiring_momentum_30d": ser.metric_value("hiring_momentum_30d", m.get("hiring_momentum_30d")), | |
| 64 | + "activity_score": ser.metric_value("activity_score", m.get("activity_score")), | |
| 65 | + "ai_adoption": ser.metric_value("ai_adoption", m.get("ai_adoption")), | |
| 66 | + "top_event_types": [k for k, _ in sorted(types.get(slug, {}).items(), key=lambda kv: (-kv[1], kv[0]))[:3]]}) | |
| 67 | + rows.sort(key=lambda r: (-r["companies"], -(r["activity_score"] or 0), r["name"])) | |
| 68 | + return rows | |
| 69 | + | |
| 70 | + | |
| 71 | +async def country_rows(conn: AsyncConnection, *, industry: str | None = None) -> list[dict[str, Any]]: | |
| 72 | + scope = " and cast(:industry as text) = any(c.industries)" if industry else "" | |
| 73 | + params: dict[str, Any] = {"industry": industry} if industry else {} | |
| 74 | + d7, d30 = q.days_ago(7), q.days_ago(30) | |
| 75 | + ref = await fetch_all(conn, "select code, name, region, subregion, lat, lon from countries order by name") | |
| 76 | + companies = await fetch_all(conn, f"select c.country as code, count(*) as n from companies c where c.status = 'ACTIVE' and c.country is not null{scope} " | |
| 77 | + "group by c.country", **params) | |
| 78 | + events = await fetch_all(conn, "select c.country as code, count(*) filter (where e.detected_at >= :d7) as n7, count(*) as n30 " | |
| 79 | + "from events e join companies c on c.id = e.company_id " | |
| 80 | + f"where e.status = 'active' and e.detected_at >= :d30 and c.country is not null{scope} group by c.country", | |
| 81 | + d7=d7, d30=d30, **params) | |
| 82 | + metrics = await fetch_all(conn, "select c.country as code, m.metric, avg(m.value) as v from metrics_current m join companies c on c.id = m.company_id " | |
| 83 | + f"where m.metric = any(cast(:ms as text[])) and c.country is not null{scope} group by c.country, m.metric", | |
| 84 | + ms=list(_ROW_METRICS), **params) | |
| 85 | + mix = await fetch_all(conn, "select c.country as code, ind, count(*) as n from companies c, unnest(c.industries) ind " | |
| 86 | + f"where c.status = 'ACTIVE' and c.country is not null{scope} group by c.country, ind", **params) | |
| 87 | + n_companies = {r["code"]: int(r["n"]) for r in companies} | |
| 88 | + ev = {r["code"]: (int(r["n7"]), int(r["n30"])) for r in events} | |
| 89 | + avg = _avg_map(metrics, "code") | |
| 90 | + mixes: dict[str, list[tuple[str, int]]] = defaultdict(list) | |
| 91 | + for r in mix: | |
| 92 | + mixes[r["code"]].append((r["ind"], int(r["n"]))) | |
| 93 | + names = {r["code"]: r for r in ref} | |
| 94 | + codes = list(names) + [c for c in n_companies if c not in names] | |
| 95 | + rows = [] | |
| 96 | + for code in codes: | |
| 97 | + c = names.get(code, {}) | |
| 98 | + m = avg.get(code, {}) | |
| 99 | + name = c.get("name") or code | |
| 100 | + rows.append({"code": code, "slug": slugify(name), "name": name, "region": c.get("region"), "subregion": c.get("subregion"), | |
| 101 | + "companies": n_companies.get(code, 0), "events_7d": ev.get(code, (0, 0))[0], "events_30d": ev.get(code, (0, 0))[1], | |
| 102 | + "hiring_momentum_30d": ser.metric_value("hiring_momentum_30d", m.get("hiring_momentum_30d")), | |
| 103 | + "activity_score": ser.metric_value("activity_score", m.get("activity_score")), | |
| 104 | + "ai_adoption": ser.metric_value("ai_adoption", m.get("ai_adoption")), | |
| 105 | + "industry_mix": [{"industry": i, "companies": n} for i, n in sorted(mixes.get(code, []), key=lambda x: -x[1])[:6]], | |
| 106 | + "lat": c.get("lat"), "lon": c.get("lon")}) | |
| 107 | + rows.sort(key=lambda r: (-r["companies"], -(r["activity_score"] or 0), r["name"])) | |
| 108 | + return rows | |
| 109 | + | |
| 110 | + | |
| 111 | +async def cached_industry_rows(country: str | None = None) -> list[dict[str, Any]]: | |
| 112 | + async def produce() -> list[dict[str, Any]]: | |
| 113 | + async with connection() as conn: | |
| 114 | + return await industry_rows(conn, country=country) | |
| 115 | + return await cached(f"industries:{country or ''}", 300, produce) | |
| 116 | + | |
| 117 | + | |
| 118 | +async def cached_country_rows(industry: str | None = None) -> list[dict[str, Any]]: | |
| 119 | + async def produce() -> list[dict[str, Any]]: | |
| 120 | + async with connection() as conn: | |
| 121 | + return await country_rows(conn, industry=industry) | |
| 122 | + return await cached(f"countries:{industry or ''}", 300, produce) | |
| 123 | + | |
| 124 | + | |
| 125 | +async def resolve_country(conn: AsyncConnection, key: str) -> dict[str, Any] | None: | |
| 126 | + """Accept an ISO-2 code (`CA`) or a name slug (`canada`).""" | |
| 127 | + key = (key or "").strip() | |
| 128 | + if not key: | |
| 129 | + return None | |
| 130 | + if len(key) == 2: | |
| 131 | + row = await fetch_one(conn, "select code, name, region, subregion, lat, lon from countries where code = cast(:c as char(2))", c=key.upper()) | |
| 132 | + if row: | |
| 133 | + return row | |
| 134 | + rows = await fetch_all(conn, "select code, name, region, subregion, lat, lon from countries") | |
| 135 | + k = slugify(key) | |
| 136 | + for r in rows: | |
| 137 | + if slugify(r["name"]) == k or r["code"].lower() == key.lower(): | |
| 138 | + return r | |
| 139 | + return None | |
| 140 | + | |
| 141 | + | |
| 142 | +# ------------------------------------------------------------------------------------------------ rankings | |
| 143 | + | |
| 144 | +RANKING_KINDS: dict[str, dict[str, Any]] = { | |
| 145 | + "most_active": {"metric": Metric.ACTIVITY_SCORE.value}, | |
| 146 | + "hiring_growth": {"metric": "hiring_momentum_{w}", "min": 0.0}, | |
| 147 | + "hiring_decline": {"metric": "hiring_momentum_{w}", "asc": True, "max": 0.0}, | |
| 148 | + "product_velocity": {"metric": Metric.PRODUCT_VELOCITY.value}, | |
| 149 | + "ai_active": {"metric": Metric.AI_ADOPTION.value}, | |
| 150 | + "geo_expansion": {"metric": Metric.GEO_EXPANSION.value}, | |
| 151 | + "developer_momentum": {"metric": Metric.DEVELOPER_MOMENTUM.value}, | |
| 152 | + "pricing_changes": {"events": "PRICING"}, | |
| 153 | + "unusual_activity": {"metric": Metric.ANOMALY_SCORE.value}, | |
| 154 | +} | |
| 155 | +_HIRING_WINDOW = {"24h": "7d", "7d": "7d", "30d": "30d", "90d": "90d", "1y": "90d"} | |
| 156 | + | |
| 157 | + | |
| 158 | +async def ranking(conn: AsyncConnection, kind: str, window: str, *, country: str | None = None, industry: str | None = None, | |
| 159 | + limit: int = 50) -> list[dict[str, Any]]: | |
| 160 | + """[{company_id, value, delta}] for a ranking kind. `delta` = value − series value at the start of the window (null if unknown).""" | |
| 161 | + spec = RANKING_KINDS[kind] | |
| 162 | + scope, params = [], {} | |
| 163 | + if country: | |
| 164 | + scope.append("c.country = cast(:country as char(2))") | |
| 165 | + params["country"] = country.upper()[:2] | |
| 166 | + if industry: | |
| 167 | + scope.append("cast(:industry as text) = any(c.industries)") | |
| 168 | + params["industry"] = industry[:80] | |
| 169 | + scope_sql = (" and " + " and ".join(scope)) if scope else "" | |
| 170 | + since = q.window_start(window) | |
| 171 | + if "events" in spec: | |
| 172 | + rows = await fetch_all(conn, "select e.company_id, count(*) as value from events e join companies c on c.id = e.company_id " | |
| 173 | + f"where e.status = 'active' and e.event_type = :et and e.detected_at >= :since and c.status = 'ACTIVE'{scope_sql} " | |
| 174 | + "group by e.company_id order by value desc, e.company_id limit :limit", et=spec["events"], since=since, | |
| 175 | + limit=limit, **params) | |
| 176 | + return [{"company_id": r["company_id"], "value": float(r["value"]), "delta": None} for r in rows] | |
| 177 | + metric = spec["metric"].format(w=_HIRING_WINDOW.get(window, "30d")) | |
| 178 | + bounds = "" | |
| 179 | + if "min" in spec: | |
| 180 | + bounds += " and m.value > :vmin" | |
| 181 | + params["vmin"] = spec["min"] | |
| 182 | + if "max" in spec: | |
| 183 | + bounds += " and m.value < :vmax" | |
| 184 | + params["vmax"] = spec["max"] | |
| 185 | + order = "m.value asc" if spec.get("asc") else "m.value desc" | |
| 186 | + rows = await fetch_all(conn, "select m.company_id, m.value from metrics_current m join companies c on c.id = m.company_id " | |
| 187 | + f"where m.metric = :metric and c.status = 'ACTIVE'{scope_sql}{bounds} order by {order}, m.company_id limit :limit", | |
| 188 | + metric=metric, limit=limit, **params) | |
| 189 | + ids = [r["company_id"] for r in rows] | |
| 190 | + past = await q.metric_values_at(conn, ids, metric, since.date()) | |
| 191 | + out = [] | |
| 192 | + for r in rows: | |
| 193 | + v = float(r["value"]) | |
| 194 | + p = past.get(r["company_id"]) | |
| 195 | + out.append({"company_id": r["company_id"], "value": v, "delta": round(v - p, 2) if p is not None else None}) | |
| 196 | + return out | |
| 197 | + | |
| 198 | + | |
| 199 | +async def ranking_cards(conn: AsyncConnection, kind: str, window: str, *, country: str | None = None, industry: str | None = None, | |
| 200 | + limit: int = 50, sparkline: bool = False) -> list[dict[str, Any]]: | |
| 201 | + items = await ranking(conn, kind, window, country=country, industry=industry, limit=limit) | |
| 202 | + cards = await q.fetch_cards_by_ids(conn, [i["company_id"] for i in items], sparkline=sparkline) | |
| 203 | + by_id = {c["id"]: c for c in cards} | |
| 204 | + out = [] | |
| 205 | + for rank, it in enumerate(items, start=1): | |
| 206 | + row = by_id.get(it["company_id"]) | |
| 207 | + if row is None: | |
| 208 | + continue | |
| 209 | + card = ser.company_card(row) | |
| 210 | + metric = RANKING_KINDS[kind].get("metric", "").format(w=_HIRING_WINDOW.get(window, "30d")) | |
| 211 | + card.update({"rank": rank, "value": ser.metric_value(metric, it["value"]) if metric else int(it["value"]), "delta": it["delta"]}) | |
| 212 | + out.append(card) | |
| 213 | + return out | |
| 214 | + | |
| 215 | + | |
| 216 | +# ------------------------------------------------------------------------------------------------ trends / signals | |
| 217 | + | |
| 218 | + | |
| 219 | +async def trend_rows(conn: AsyncConnection, window_days: int, limit: int) -> list[dict[str, Any]]: | |
| 220 | + d0 = q.days_ago(window_days).date() | |
| 221 | + top = await fetch_all(conn, "select term, sum(mentions) as mentions, max(companies) as companies from trends where day >= :d group by term " | |
| 222 | + "order by mentions desc, term limit :limit", d=d0, limit=limit) | |
| 223 | + if not top: | |
| 224 | + return [] | |
| 225 | + terms = [t["term"] for t in top] | |
| 226 | + series = await fetch_all(conn, "select term, day, mentions from trends where day >= :d and term = any(cast(:terms as text[])) order by term, day", | |
| 227 | + d=d0, terms=terms) | |
| 228 | + by_term: dict[str, list[tuple[Any, int]]] = defaultdict(list) | |
| 229 | + for r in series: | |
| 230 | + by_term[r["term"]].append((r["day"], int(r["mentions"]))) | |
| 231 | + mid = q.days_ago(window_days // 2 or 1).date() | |
| 232 | + out = [] | |
| 233 | + for t in top: | |
| 234 | + pts = by_term.get(t["term"], []) | |
| 235 | + first = sum(m for d, m in pts if d < mid) | |
| 236 | + second = sum(m for d, m in pts if d >= mid) | |
| 237 | + momentum = round((second - first) / first * 100, 1) if first > 0 else None | |
| 238 | + out.append({"term": t["term"], "mentions": int(t["mentions"]), "companies": int(t["companies"] or 0), "momentum": momentum, | |
| 239 | + "series": [m for _, m in pts]}) | |
| 240 | + return out | |
| 241 | + | |
| 242 | + | |
| 243 | +# ------------------------------------------------------------------------------------------------ map | |
| 244 | + | |
| 245 | +MAP_MAX_BUCKETS = 600 | |
| 246 | + | |
| 247 | + | |
| 248 | +async def map_buckets(conn: AsyncConnection, metric: str = "events_30d") -> list[dict[str, Any]]: | |
| 249 | + d30 = q.days_ago(30) | |
| 250 | + ev_by_company = {r["company_id"]: int(r["n"]) for r in await fetch_all( | |
| 251 | + conn, "select company_id, count(*) as n from events where status = 'active' and detected_at >= :d group by company_id", d=d30)} | |
| 252 | + jobs_by_company = {r["company_id"]: int(r["n"]) for r in await fetch_all( | |
| 253 | + conn, "select company_id, count(*) as n from jobs where status = 'open' group by company_id")} | |
| 254 | + countries = await fetch_all(conn, "select k.code, k.name, k.lat, k.lon, c.id, c.slug, c.display_name, c.importance from countries k " | |
| 255 | + "join companies c on c.country = k.code and c.status = 'ACTIVE' where k.lat is not null order by c.importance desc") | |
| 256 | + cities = await fetch_all(conn, "select l.country, l.city, avg(l.lat) as lat, avg(l.lon) as lon, " | |
| 257 | + "array_agg(distinct l.company_id) as company_ids from locations l join companies c on c.id = l.company_id " | |
| 258 | + "where l.status = 'listed' and l.lat is not null and l.lon is not null and l.city is not null " | |
| 259 | + "group by l.country, l.city order by count(distinct l.company_id) desc limit :lim", lim=MAP_MAX_BUCKETS) | |
| 260 | + by_country: dict[str, dict[str, Any]] = {} | |
| 261 | + for r in countries: | |
| 262 | + b = by_country.setdefault(r["code"], {"lat": r["lat"], "lon": r["lon"], "country": r["code"], "city": None, "companies": 0, "events_30d": 0, | |
| 263 | + "jobs_open": 0, "top": [], "_ids": []}) | |
| 264 | + b["companies"] += 1 | |
| 265 | + b["events_30d"] += ev_by_company.get(r["id"], 0) | |
| 266 | + b["jobs_open"] += jobs_by_company.get(r["id"], 0) | |
| 267 | + if len(b["top"]) < 3: | |
| 268 | + b["top"].append({"slug": r["slug"], "display_name": r["display_name"]}) | |
| 269 | + names: dict[str, tuple[str, str]] = {} | |
| 270 | + if cities: | |
| 271 | + ids = sorted({cid for r in cities for cid in (r["company_ids"] or [])}) | |
| 272 | + names = {r["id"]: (r["slug"], r["display_name"]) for r in await fetch_all( | |
| 273 | + conn, "select id, slug, display_name from companies where id = any(cast(:ids as text[])) order by importance desc", ids=ids[:5000])} | |
| 274 | + buckets = list(by_country.values()) | |
| 275 | + for r in cities: | |
| 276 | + ids = [cid for cid in (r["company_ids"] or []) if cid in names] | |
| 277 | + buckets.append({"lat": round(float(r["lat"]), 4), "lon": round(float(r["lon"]), 4), "country": r["country"], "city": r["city"], "companies": len(ids), | |
| 278 | + "events_30d": sum(ev_by_company.get(i, 0) for i in ids), "jobs_open": sum(jobs_by_company.get(i, 0) for i in ids), | |
| 279 | + "top": [{"slug": names[i][0], "display_name": names[i][1]} for i in ids[:3]]}) | |
| 280 | + key = {"companies": "companies", "hiring": "jobs_open"}.get(metric, "events_30d") | |
| 281 | + buckets.sort(key=lambda b: (-b[key], -b["companies"])) | |
| 282 | + for b in buckets: | |
| 283 | + b.pop("_ids", None) | |
| 284 | + return buckets[:MAP_MAX_BUCKETS] | |
| 285 | + | |
| 286 | + | |
| 287 | +# ------------------------------------------------------------------------------------------------ global stats / index | |
| 288 | + | |
| 289 | + | |
| 290 | +async def archive_stats() -> dict[str, int]: | |
| 291 | + async def produce() -> dict[str, int]: | |
| 292 | + async with connection() as conn: | |
| 293 | + kv = await q.settings_value(conn, "archive:stats") | |
| 294 | + if isinstance(kv, dict) and "objects" in kv: | |
| 295 | + return {"objects": int(kv.get("objects") or 0), "bytes": int(kv.get("bytes") or 0)} | |
| 296 | + return await asyncio.to_thread(archive.store_stats) | |
| 297 | + return await cached("archive:stats", 600, produce) | |
| 298 | + | |
| 299 | + | |
| 300 | +async def global_stats(conn: AsyncConnection) -> dict[str, Any]: | |
| 301 | + today = datetime.now(UTC).replace(hour=0, minute=0, second=0, microsecond=0) | |
| 302 | + row = await fetch_one(conn, """ | |
| 303 | + select (select count(*) from companies) as companies, | |
| 304 | + (select count(*) from companies where status = 'ACTIVE' and onboarding_status = 'active') as companies_active, | |
| 305 | + (select count(*) from sensors where status <> 'retired') as sensors, | |
| 306 | + (select count(*) from sensors where status = 'active') as sensors_active, | |
| 307 | + (select coalesce(sum(observation_count), 0) from sensors) as observations, | |
| 308 | + (select coalesce(sum(snapshot_count), 0) from sensors) as snapshots, | |
| 309 | + (select count(*) from changes) as changes, | |
| 310 | + (select count(*) from changes where kind in ('meaningful', 'major', 'critical')) as meaningful_changes, | |
| 311 | + (select count(*) from events where status = 'active') as events, | |
| 312 | + (select count(*) from jobs where status = 'open') as jobs_open, | |
| 313 | + (select count(distinct country) from companies where country is not null and status = 'ACTIVE') as countries, | |
| 314 | + (select count(distinct ind) from companies c, unnest(c.industries) ind where c.status = 'ACTIVE') as industries, | |
| 315 | + (select count(*) from observations where fetched_at >= :today) as observations_today, | |
| 316 | + (select count(*) from changes where detected_at >= :today) as changes_today, | |
| 317 | + (select count(*) from events where status = 'active' and detected_at >= :today) as events_today, | |
| 318 | + (select min(first_observed_at) from companies) as oldest_observation_at, | |
| 319 | + (select max(fetched_at) from observations) as last_observation_at | |
| 320 | + """, today=today) | |
| 321 | + row = row or {} | |
| 322 | + started = await q.settings_value(conn, "dataset_started_at") | |
| 323 | + started_dt = q.parse_iso(started) if isinstance(started, str) else None | |
| 324 | + now = datetime.now(UTC) | |
| 325 | + oldest = row.get("oldest_observation_at") | |
| 326 | + return {"companies": int(row.get("companies") or 0), "companies_active": int(row.get("companies_active") or 0), | |
| 327 | + "sensors": int(row.get("sensors") or 0), "sensors_active": int(row.get("sensors_active") or 0), | |
| 328 | + "observations": int(row.get("observations") or 0), "snapshots": int(row.get("snapshots") or 0), "changes": int(row.get("changes") or 0), | |
| 329 | + "meaningful_changes": int(row.get("meaningful_changes") or 0), "events": int(row.get("events") or 0), | |
| 330 | + "jobs_open": int(row.get("jobs_open") or 0), "countries": int(row.get("countries") or 0), "industries": int(row.get("industries") or 0), | |
| 331 | + "observations_today": int(row.get("observations_today") or 0), "changes_today": int(row.get("changes_today") or 0), | |
| 332 | + "events_today": int(row.get("events_today") or 0), "dataset_started_at": started_dt, | |
| 333 | + "dataset_age_days": (now - started_dt).days if started_dt else None, | |
| 334 | + "oldest_history_days": (now - oldest).days if oldest else None, "last_observation_at": row.get("last_observation_at"), | |
| 335 | + "archive": await archive_stats()} | |
| 336 | + | |
| 337 | + | |
| 338 | +async def cached_global_stats() -> dict[str, Any]: | |
| 339 | + async def produce() -> dict[str, Any]: | |
| 340 | + async with connection() as conn: | |
| 341 | + return await global_stats(conn) | |
| 342 | + return await cached("stats", 60, produce) | |
| 343 | + | |
| 344 | + | |
| 345 | +async def global_daily_rows(conn: AsyncConnection, days: int) -> list[dict[str, Any]]: | |
| 346 | + rows = await fetch_all(conn, "select * from global_daily where day >= :d order by day asc limit :lim", d=q.days_ago(days).date(), lim=days + 1) | |
| 347 | + return [{"day": r["day"], "companies_active": r["companies_active"], "sensors_active": r["sensors_active"], "observations": r["observations"], | |
| 348 | + "changes": r["changes"], "meaningful_changes": r["meaningful_changes"], "events": r["events"], "events_by_type": ser._dict(r["events_by_type"]), | |
| 349 | + "jobs_open": r["jobs_open"], "jobs_new": r["jobs_new"], "jobs_removed": r["jobs_removed"], | |
| 350 | + "activity_index": ser._float(r["activity_index"], 2), "by_country": ser._dict(r["by_country"]), "by_industry": ser._dict(r["by_industry"])} | |
| 351 | + for r in rows] | |
| 352 | + | |
| 353 | + | |
| 354 | +def _index_at(rows: list[dict[str, Any]], days_back: int) -> float | None: | |
| 355 | + if not rows: | |
| 356 | + return None | |
| 357 | + target = rows[-1]["day"] - timedelta(days=days_back) | |
| 358 | + candidates = [r for r in rows if r["day"] <= target and r["activity_index"] is not None] | |
| 359 | + return candidates[-1]["activity_index"] if candidates else None | |
| 360 | + | |
| 361 | + | |
| 362 | +async def activity_index(conn: AsyncConnection, days: int = 365) -> dict[str, Any]: | |
| 363 | + rows = await global_daily_rows(conn, days) | |
| 364 | + with_value = [r for r in rows if r["activity_index"] is not None] | |
| 365 | + latest = with_value[-1] if with_value else None | |
| 366 | + value = latest["activity_index"] if latest else None | |
| 367 | + v7, v30 = _index_at(with_value, 7), _index_at(with_value, 30) | |
| 368 | + formula = await q.settings_value(conn, "index:formula_version") | |
| 369 | + return {"value": value, "baseline": 100, "delta_7d": round(value - v7, 2) if value is not None and v7 is not None else None, | |
| 370 | + "delta_30d": round(value - v30, 2) if value is not None and v30 is not None else None, | |
| 371 | + "series": [{"day": r["day"], "value": r["activity_index"], "confidence": None} for r in with_value], | |
| 372 | + "by_type": (latest or {}).get("events_by_type", {}), | |
| 373 | + "by_country": [{"key": k, "value": v} for k, v in sorted((latest or {}).get("by_country", {}).items(), key=lambda kv: -float(kv[1] or 0))[:50]], | |
| 374 | + "by_industry": [{"key": k, "value": v} for k, v in sorted((latest or {}).get("by_industry", {}).items(), key=lambda kv: -float(kv[1] or 0))[:50]], | |
| 375 | + "formula_version": formula if isinstance(formula, str) else METRICS_FORMULA_VERSION, "computed_at": (latest or {}).get("day")} | |
| 376 | + | |
| 377 | + | |
| 378 | +async def system_health(conn: AsyncConnection) -> dict[str, Any]: | |
| 379 | + today = datetime.now(UTC).replace(hour=0, minute=0, second=0, microsecond=0) | |
| 380 | + row = await fetch_one(conn, """ | |
| 381 | + select (select count(*) from sensors where status = 'active') as sensors_online, | |
| 382 | + (select count(*) from sensors where status in ('failing', 'stale', 'blocked')) as sensors_failing, | |
| 383 | + (select count(*) from observations where fetched_at >= :today) as observations_today, | |
| 384 | + (select count(*) from events where status = 'active' and detected_at >= :today) as events_today, | |
| 385 | + (select count(distinct country) from companies where country is not null and status = 'ACTIVE') as countries_covered, | |
| 386 | + (select extract(epoch from (now() - min(run_at))) from queue_jobs where status = 'pending' and run_at <= now()) as queue_lag_s, | |
| 387 | + (select count(*) from observations where fetched_at >= now() - interval '10 minutes') as obs_10m, | |
| 388 | + (select count(*) from observations where fetched_at >= now() - interval '24 hours') as obs_24h, | |
| 389 | + (select count(*) from observations where fetched_at >= now() - interval '24 hours' and failure_class is null) as ok_24h | |
| 390 | + """, today=today) or {} | |
| 391 | + hb = await q.settings_value(conn, "scheduler:heartbeat") | |
| 392 | + tick = None | |
| 393 | + if isinstance(hb, dict): | |
| 394 | + tick = hb.get("at") or hb.get("ts") or hb.get("time") or hb.get("last_tick_at") | |
| 395 | + elif isinstance(hb, str): | |
| 396 | + tick = hb | |
| 397 | + obs_24h = int(row.get("obs_24h") or 0) | |
| 398 | + return {"sensors_online": int(row.get("sensors_online") or 0), "sensors_failing": int(row.get("sensors_failing") or 0), | |
| 399 | + "observations_today": int(row.get("observations_today") or 0), "events_today": int(row.get("events_today") or 0), | |
| 400 | + "countries_covered": int(row.get("countries_covered") or 0), | |
| 401 | + "queue_lag_s": round(float(row["queue_lag_s"]), 1) if row.get("queue_lag_s") is not None else 0.0, | |
| 402 | + "scheduler_last_tick_at": tick, "fetch_per_min": round(int(row.get("obs_10m") or 0) / 10.0, 2), | |
| 403 | + "success_rate_24h": round(int(row.get("ok_24h") or 0) / obs_24h, 4) if obs_24h else None} | |
| 404 | + | |
| 405 | + | |
| 406 | +async def live_events(conn: AsyncConnection, limit: int, **filters: Any) -> list[dict[str, Any]]: | |
| 407 | + where, params = q.event_filters(**filters) | |
| 408 | + return [ser.event(r) for r in await q.fetch_events(conn, where, params, sort="recent", limit=limit)] | |
| 409 | + | |
| 410 | + | |
| 411 | +def clear_aggregate_cache() -> None: | |
| 412 | + cache.clear() | |
| 413 | + | |
| 414 | + | |
| 415 | +__all__ = ["MAP_MAX_BUCKETS", "RANKING_KINDS", "activity_index", "archive_stats", "cached_country_rows", "cached_global_stats", "cached_industry_rows", | |
| 416 | + "clear_aggregate_cache", "country_rows", "global_daily_rows", "global_stats", "industry_rows", "live_events", "map_buckets", "ranking", | |
| 417 | + "ranking_cards", "resolve_country", "system_health", "trend_rows"] | |
added
src/companyatlas/api/ask_fallback.py
+180 −0
@@ -0,0 +1,180 @@ | ||
| 1 | +"""Deterministic natural-language parser for `/ask` (used when `companyatlas.services.llm.ask` is not available). | |
| 2 | + | |
| 3 | +"companies hiring AI engineers in Canada" → {event_types: [HIRING], ai: true, country: "CA", window: "30d"}. Only vocabulary from the | |
| 4 | +taxonomy and the reference tables is recognised; anything else becomes free-text `terms` used for company search. Nothing is guessed: | |
| 5 | +when no filter matches, the interpretation says so and the answer falls back to "no monitored evidence". | |
| 6 | +""" | |
| 7 | +from __future__ import annotations | |
| 8 | + | |
| 9 | +import re | |
| 10 | +from dataclasses import asdict, dataclass, field | |
| 11 | +from typing import Any | |
| 12 | + | |
| 13 | +from companyatlas.ids import normalize_alias | |
| 14 | + | |
| 15 | +WINDOW_PATTERNS: list[tuple[re.Pattern[str], str]] = [ | |
| 16 | + (re.compile(r"\b(today|last 24 ?h(ours)?|past 24 ?h(ours)?)\b", re.IGNORECASE), "24h"), | |
| 17 | + (re.compile(r"\b(yesterday)\b", re.IGNORECASE), "24h"), | |
| 18 | + (re.compile(r"\b(this week|last (7|seven) days|past (7|seven) days|last week|weekly)\b", re.IGNORECASE), "7d"), | |
| 19 | + (re.compile(r"\b(this month|last (30|thirty) days|past (30|thirty) days|last month|monthly|recently|recent)\b", re.IGNORECASE), "30d"), | |
| 20 | + (re.compile(r"\b(this quarter|last (90|ninety) days|past (90|ninety) days|last quarter|quarterly|last (3|three) months)\b", re.IGNORECASE), "90d"), | |
| 21 | + (re.compile(r"\b(this year|last (12|twelve) months|past year|last year|yearly|annual)\b", re.IGNORECASE), "1y"), | |
| 22 | +] | |
| 23 | + | |
| 24 | +EVENT_KEYWORDS: dict[str, tuple[str, ...]] = { | |
| 25 | + "HIRING": ("hiring", "hire", "hires", "jobs", "job", "recruit", "recruiting", "openings", "positions", "careers", "vacanc", "headcount", "layoff", "layoffs"), | |
| 26 | + "PRICING": ("pricing", "price", "prices", "plan", "plans", "tier", "tiers", "subscription", "cheaper", "expensive"), | |
| 27 | + "PRODUCT": ("launch", "launched", "launches", "product", "products", "release", "released", "feature", "features", "shipped", "ship"), | |
| 28 | + "LEADERSHIP": ("ceo", "cfo", "cto", "coo", "executive", "executives", "leadership", "appointed", "appoint", "board", "founder", "president", "chief"), | |
| 29 | + "LOCATION": ("office", "offices", "location", "locations", "expansion", "expanding", "expand", "opened", "opening", "headquarters", "hq", "store", "stores", "factory"), | |
| 30 | + "FINANCING": ("funding", "raised", "raise", "round", "series a", "series b", "investment", "investors", "ipo", "valuation"), | |
| 31 | + "M&A": ("acquisition", "acquired", "acquire", "acquires", "merger", "merged", "buyout", "divest", "divestiture"), | |
| 32 | + "PARTNERSHIP": ("partnership", "partner", "partners", "partnered", "alliance", "integration"), | |
| 33 | + "DEVELOPER": ("api", "apis", "sdk", "developer", "developers", "docs", "documentation", "changelog"), | |
| 34 | + "LEGAL": ("legal", "terms", "privacy", "policy", "compliance", "regulatory", "regulation", "gdpr"), | |
| 35 | + "SECURITY": ("security", "breach", "incident", "vulnerability", "cve"), | |
| 36 | + "TECHNOLOGY": ("technology", "tech stack", "adopted", "adoption", "platform"), | |
| 37 | + "SUSTAINABILITY": ("sustainability", "esg", "climate", "carbon", "net zero"), | |
| 38 | + "INVESTOR_RELATIONS": ("earnings", "investor", "quarterly results", "annual report", "guidance"), | |
| 39 | + "COMMUNICATION": ("news", "press", "announcement", "announced", "blog", "published"), | |
| 40 | +} | |
| 41 | +AI_TERMS = ("ai", "a.i.", "artificial intelligence", "machine learning", "ml", "llm", "llms", "generative", "genai", "copilot", "agentic", "deep learning") | |
| 42 | +RANKING_INTENT: list[tuple[re.Pattern[str], str]] = [ | |
| 43 | + (re.compile(r"\b(most active|fastest moving|moving fastest|busiest)\b", re.IGNORECASE), "most_active"), | |
| 44 | + (re.compile(r"\b(hiring (the )?(most|fastest)|fastest hiring|hiring growth|growing headcount)\b", re.IGNORECASE), "hiring_growth"), | |
| 45 | + (re.compile(r"\b(hiring (decline|freeze|slowdown)|cutting|fewer jobs|shrinking)\b", re.IGNORECASE), "hiring_decline"), | |
| 46 | + (re.compile(r"\b(product velocity|shipping (the )?most|most launches)\b", re.IGNORECASE), "product_velocity"), | |
| 47 | + (re.compile(r"\b(most ai|ai[- ]active|ai adoption|ai leaders)\b", re.IGNORECASE), "ai_active"), | |
| 48 | + (re.compile(r"\b(expanding (abroad|internationally|geographically)|geographic expansion|new countries)\b", re.IGNORECASE), "geo_expansion"), | |
| 49 | + (re.compile(r"\b(developer momentum|developer[- ]focused)\b", re.IGNORECASE), "developer_momentum"), | |
| 50 | + (re.compile(r"\b(pricing changes|changed (their )?pricing|price (increase|hike)s?)\b", re.IGNORECASE), "pricing_changes"), | |
| 51 | + (re.compile(r"\b(unusual|anomal|abnormal|behaving (unusually|strangely))", re.IGNORECASE), "unusual_activity"), | |
| 52 | +] | |
| 53 | +DEMONYMS: dict[str, str] = { | |
| 54 | + "canadian": "CA", "american": "US", "us": "US", "usa": "US", "u.s.": "US", "british": "GB", "uk": "GB", "u.k.": "GB", "english": "GB", "french": "FR", | |
| 55 | + "german": "DE", "japanese": "JP", "korean": "KR", "indian": "IN", "australian": "AU", "chinese": "CN", "brazilian": "BR", "mexican": "MX", | |
| 56 | + "dutch": "NL", "swedish": "SE", "swiss": "CH", "spanish": "ES", "italian": "IT", "israeli": "IL", "singaporean": "SG", "irish": "IE", | |
| 57 | + "european": None, # region, not a country | |
| 58 | +} | |
| 59 | +STOPWORDS = {"the", "a", "an", "of", "in", "on", "at", "for", "to", "and", "or", "with", "which", "what", "who", "are", "is", "was", "were", "has", | |
| 60 | + "have", "had", "do", "does", "did", "show", "me", "list", "find", "companies", "company", "that", "this", "these", "those", "any", | |
| 61 | + "all", "from", "by", "about", "their", "its", "new", "recently", "recent", "engineers", "engineer", "people", "roles", "many", "how", | |
| 62 | + "much", "top", "best", "biggest", "largest", "most", "events", "event", "changes", "change", "signal", "signals", "atlas"} | |
| 63 | + | |
| 64 | + | |
| 65 | +@dataclass | |
| 66 | +class Interpretation: | |
| 67 | + query: str | |
| 68 | + window: str = "30d" | |
| 69 | + event_types: list[str] = field(default_factory=list) | |
| 70 | + country: str | None = None | |
| 71 | + country_name: str | None = None | |
| 72 | + industry: str | None = None | |
| 73 | + industry_name: str | None = None | |
| 74 | + ai: bool = False | |
| 75 | + ranking_kind: str | None = None | |
| 76 | + company_terms: list[str] = field(default_factory=list) | |
| 77 | + terms: list[str] = field(default_factory=list) | |
| 78 | + matched: list[str] = field(default_factory=list) | |
| 79 | + | |
| 80 | + def to_json(self) -> dict[str, Any]: | |
| 81 | + d = asdict(self) | |
| 82 | + d["filters"] = {k: v for k, v in {"event_types": self.event_types or None, "country": self.country, "industry": self.industry, | |
| 83 | + "ai": self.ai or None, "window": self.window}.items() if v} | |
| 84 | + return d | |
| 85 | + | |
| 86 | + | |
| 87 | +def _contains(text: str, phrase: str) -> bool: | |
| 88 | + return re.search(rf"(?<![a-z0-9]){re.escape(phrase)}(?![a-z0-9])", text) is not None | |
| 89 | + | |
| 90 | + | |
| 91 | +def interpret(query: str, *, countries: list[dict[str, Any]], industries: list[dict[str, Any]]) -> Interpretation: | |
| 92 | + """`countries`: [{code, name}], `industries`: [{slug, name, keywords}]. Deterministic, order-independent.""" | |
| 93 | + text = " " + re.sub(r"\s+", " ", (query or "").strip().lower()) + " " | |
| 94 | + it = Interpretation(query=query.strip()) | |
| 95 | + consumed: list[str] = [] | |
| 96 | + for pat, w in WINDOW_PATTERNS: | |
| 97 | + m = pat.search(text) | |
| 98 | + if m: | |
| 99 | + it.window = w | |
| 100 | + consumed.append(m.group(0)) | |
| 101 | + it.matched.append(f"window:{w}") | |
| 102 | + break | |
| 103 | + for pat, kind in RANKING_INTENT: | |
| 104 | + m = pat.search(text) | |
| 105 | + if m: | |
| 106 | + it.ranking_kind = kind | |
| 107 | + consumed.append(m.group(0)) | |
| 108 | + it.matched.append(f"ranking:{kind}") | |
| 109 | + break | |
| 110 | + if any(_contains(text, t) for t in AI_TERMS): | |
| 111 | + it.ai = True | |
| 112 | + it.matched.append("ai") | |
| 113 | + consumed.extend(t for t in AI_TERMS if _contains(text, t)) | |
| 114 | + for etype, words in EVENT_KEYWORDS.items(): | |
| 115 | + hits = [w for w in words if _contains(text, w)] | |
| 116 | + if hits: | |
| 117 | + it.event_types.append(etype) | |
| 118 | + consumed.extend(hits) | |
| 119 | + it.matched.append(f"event_type:{etype}") | |
| 120 | + by_len = sorted(countries, key=lambda c: -len(c.get("name") or "")) | |
| 121 | + for c in by_len: | |
| 122 | + name = (c.get("name") or "").lower() | |
| 123 | + if name and _contains(text, name): | |
| 124 | + it.country, it.country_name = c["code"], c["name"] | |
| 125 | + consumed.append(name) | |
| 126 | + it.matched.append(f"country:{c['code']}") | |
| 127 | + break | |
| 128 | + if it.country is None: | |
| 129 | + codes = {c["code"].upper(): c for c in countries} | |
| 130 | + for dem, code in DEMONYMS.items(): | |
| 131 | + if code and _contains(text, dem) and code in codes: | |
| 132 | + it.country, it.country_name = code, codes[code]["name"] | |
| 133 | + consumed.append(dem) | |
| 134 | + it.matched.append(f"country:{code}") | |
| 135 | + break | |
| 136 | + for ind in sorted(industries, key=lambda i: -len(i.get("name") or "")): | |
| 137 | + candidates = [ind.get("name") or "", ind.get("slug", "").replace("-", " ")] + list(ind.get("keywords") or []) | |
| 138 | + hit = next((c for c in candidates if c and len(c) >= 3 and _contains(text, c.lower())), None) | |
| 139 | + if hit: | |
| 140 | + it.industry, it.industry_name = ind["slug"], ind.get("name") or ind["slug"] | |
| 141 | + consumed.append(hit.lower()) | |
| 142 | + it.matched.append(f"industry:{ind['slug']}") | |
| 143 | + break | |
| 144 | + residual = text | |
| 145 | + for phrase in sorted(set(consumed), key=len, reverse=True): | |
| 146 | + residual = re.sub(rf"(?<![a-z0-9]){re.escape(phrase)}(?![a-z0-9])", " ", residual) | |
| 147 | + tokens = [t for t in re.findall(r"[a-z0-9][a-z0-9.&'-]*", residual) if t not in STOPWORDS and len(t) > 1] | |
| 148 | + it.terms = tokens[:8] | |
| 149 | + it.company_terms = [t for t in tokens if len(t) >= 3][:4] | |
| 150 | + return it | |
| 151 | + | |
| 152 | + | |
| 153 | +def compose_answer(it: Interpretation, *, events_total: int, companies_count: int, sample_titles: list[str]) -> str: | |
| 154 | + """Careful, non-fabricating summary sentence (spec §162–168).""" | |
| 155 | + scope = [] | |
| 156 | + if it.event_types: | |
| 157 | + scope.append(" / ".join(it.event_types).lower() + " events") | |
| 158 | + else: | |
| 159 | + scope.append("events") | |
| 160 | + if it.ai: | |
| 161 | + scope.append("with observable AI signals") | |
| 162 | + if it.industry_name: | |
| 163 | + scope.append(f"in {it.industry_name}") | |
| 164 | + if it.country_name: | |
| 165 | + scope.append(f"in {it.country_name}") | |
| 166 | + window = {"24h": "the last 24 hours", "7d": "the last 7 days", "30d": "the last 30 days", "90d": "the last 90 days", "1y": "the last 12 months"}[it.window] | |
| 167 | + if events_total == 0 and companies_count == 0: | |
| 168 | + return f"No monitored evidence matches this question yet ({' '.join(scope)}, {window}). Coverage grows as sensors observe more pages." | |
| 169 | + parts = [f"Detected {events_total:,} {' '.join(scope)} across {companies_count:,} monitored compan{'y' if companies_count == 1 else 'ies'} in {window}."] | |
| 170 | + if sample_titles: | |
| 171 | + parts.append("Most recent: " + "; ".join(t[:90] for t in sample_titles[:3]) + ".") | |
| 172 | + parts.append("Each event links to its public source; interpretations are signals, not verified facts.") | |
| 173 | + return " ".join(parts) | |
| 174 | + | |
| 175 | + | |
| 176 | +def alias_key(term: str) -> str: | |
| 177 | + return normalize_alias(term) | |
| 178 | + | |
| 179 | + | |
| 180 | +__all__ = ["EVENT_KEYWORDS", "Interpretation", "alias_key", "compose_answer", "interpret"] | |
modified
src/companyatlas/api/common.py
+32 −3
@@ -5,10 +5,10 @@ import hashlib | ||
| 5 | 5 | import hmac |
| 6 | 6 | import time |
| 7 | 7 | from collections import OrderedDict |
| 8 | −from typing import Any | |
| 8 | +from typing import Annotated, Any | |
| 9 | 9 | |
| 10 | 10 | import orjson |
| 11 | −from fastapi import Header, HTTPException, Query, Request | |
| 11 | +from fastapi import Depends, Header, HTTPException, Query, Request, Response | |
| 12 | 12 | from fastapi.responses import JSONResponse |
| 13 | 13 | |
| 14 | 14 | from companyatlas.config import settings |
@@ -44,6 +44,9 @@ class PageParams: | ||
| 44 | 44 | return (self.page - 1) * self.per_page |
| 45 | 45 | |
| 46 | 46 | |
| 47 | +PageDep = Annotated[PageParams, Depends()] # `p: PageDep` in route signatures (avoids call-in-default lint) | |
| 48 | + | |
| 49 | + | |
| 47 | 50 | def page_payload(items: list[Any], total: int, p: PageParams) -> dict[str, Any]: |
| 48 | 51 | return {"items": items, "page": p.page, "per_page": p.per_page, "total": total, "pages": max(1, -(-total // p.per_page))} |
| 49 | 52 | |
@@ -115,4 +118,30 @@ async def cached(key: str, ttl_s: float, producer): # type: ignore[no-untyped-d | ||
| 115 | 118 | return value |
| 116 | 119 | |
| 117 | 120 | |
| 118 | −__all__ = ["AtlasJSONResponse", "PageParams", "TTLCache", "cache", "cached", "client_ip", "owner_hash", "page_payload", "require_admin"] | |
| 121 | +# ------------------------------------------------------------------------------------------------ cache headers / conditional GET | |
| 122 | + | |
| 123 | +NO_STORE = "no-store" | |
| 124 | + | |
| 125 | + | |
| 126 | +def public_cache_value(max_age_s: int) -> str: | |
| 127 | + return f"public, max-age={max_age_s}, stale-while-revalidate={max_age_s * 2}" | |
| 128 | + | |
| 129 | + | |
| 130 | +def cached_response(request: Request, payload: Any, max_age_s: int) -> Any: | |
| 131 | + """JSON response for a cached public aggregate: `Cache-Control: public, max-age=…` + weak ETag, `304` when `If-None-Match` matches.""" | |
| 132 | + body = AtlasJSONResponse(payload).body | |
| 133 | + etag = 'W/"' + hashlib.sha1(body).hexdigest()[:20] + '"' | |
| 134 | + headers = {"cache-control": public_cache_value(max_age_s), "etag": etag, "vary": "Accept-Encoding"} | |
| 135 | + inm = request.headers.get("if-none-match") | |
| 136 | + if inm and etag in [x.strip() for x in inm.split(",")]: | |
| 137 | + return Response(status_code=304, headers=headers) | |
| 138 | + return Response(content=body, media_type=AtlasJSONResponse.media_type, headers=headers) | |
| 139 | + | |
| 140 | + | |
| 141 | +def is_admin_request(request: Request) -> bool: | |
| 142 | + tok = request.headers.get("x-ca-admin-token") | |
| 143 | + return bool(settings.admin_token and tok and hmac.compare_digest(tok, settings.admin_token)) | |
| 144 | + | |
| 145 | + | |
| 146 | +__all__ = ["NO_STORE", "AtlasJSONResponse", "PageDep", "PageParams", "TTLCache", "cache", "cached", "cached_response", "client_ip", "is_admin_request", | |
| 147 | + "owner_hash", "page_payload", "public_cache_value", "require_admin"] | |
modified
src/companyatlas/api/main.py
+4 −0
@@ -128,4 +128,8 @@ def _include_routers() -> None: | ||
| 128 | 128 | |
| 129 | 129 | _include_routers() |
| 130 | 130 | |
| 131 | +from companyatlas.api.ratelimit import rate_limit_middleware # registered last on purpose: outermost app-level hook | |
| 132 | + | |
| 133 | +app.middleware("http")(rate_limit_middleware) | |
| 134 | + | |
| 131 | 135 | __all__ = ["API_VERSION", "app"] |
added
src/companyatlas/api/queries.py
+309 −0
@@ -0,0 +1,309 @@ | ||
| 1 | +"""Reusable, parameterised SQL for the public/admin API. | |
| 2 | + | |
| 3 | +Rules: user input is always bound (`:param`), never interpolated; sort keys go through whitelists; every list is bounded; the | |
| 4 | +company card is built in two phases (page of ids with an indexed ORDER BY, then one card query for those ids) so the per-company | |
| 5 | +lateral aggregates only run for the rows that are returned. | |
| 6 | +""" | |
| 7 | +from __future__ import annotations | |
| 8 | + | |
| 9 | +from datetime import UTC, date, datetime, timedelta | |
| 10 | +from typing import Any | |
| 11 | + | |
| 12 | +from fastapi import HTTPException | |
| 13 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 14 | + | |
| 15 | +from companyatlas.db import fetch_all, fetch_one, fetch_val | |
| 16 | +from companyatlas.taxonomy import EventType | |
| 17 | + | |
| 18 | +# ------------------------------------------------------------------------------------------------ windows / params | |
| 19 | + | |
| 20 | +WINDOWS: dict[str, timedelta] = {"24h": timedelta(hours=24), "7d": timedelta(days=7), "30d": timedelta(days=30), "90d": timedelta(days=90), | |
| 21 | + "1y": timedelta(days=365)} | |
| 22 | +COUNT_CAP = 10_000 # `total` is exact up to this bound (keeps deep pagination cheap on the big tables) | |
| 23 | +MAX_LIMIT = 500 | |
| 24 | + | |
| 25 | + | |
| 26 | +def now_utc() -> datetime: | |
| 27 | + return datetime.now(UTC) | |
| 28 | + | |
| 29 | + | |
| 30 | +def window_start(window: str, default: str = "7d") -> datetime: | |
| 31 | + return now_utc() - WINDOWS.get(window, WINDOWS[default]) | |
| 32 | + | |
| 33 | + | |
| 34 | +def days_ago(days: int) -> datetime: | |
| 35 | + return now_utc() - timedelta(days=days) | |
| 36 | + | |
| 37 | + | |
| 38 | +def parse_iso(value: str | None, name: str = "since") -> datetime | None: | |
| 39 | + if not value: | |
| 40 | + return None | |
| 41 | + v = value.strip() | |
| 42 | + if v.endswith(("Z", "z")): | |
| 43 | + v = v[:-1] + "+00:00" | |
| 44 | + try: | |
| 45 | + dt = datetime.fromisoformat(v) | |
| 46 | + except ValueError: | |
| 47 | + try: | |
| 48 | + dt = datetime.combine(date.fromisoformat(v), datetime.min.time()) | |
| 49 | + except ValueError as exc: | |
| 50 | + raise HTTPException(status_code=422, detail=f"{name}: invalid ISO-8601 timestamp") from exc | |
| 51 | + if dt.tzinfo is None: | |
| 52 | + dt = dt.replace(tzinfo=UTC) | |
| 53 | + return dt.astimezone(UTC) | |
| 54 | + | |
| 55 | + | |
| 56 | +def parse_bool(value: str | bool | None) -> bool | None: | |
| 57 | + if value is None or value == "": | |
| 58 | + return None | |
| 59 | + if isinstance(value, bool): | |
| 60 | + return value | |
| 61 | + return value.strip().lower() in ("1", "true", "yes", "on") | |
| 62 | + | |
| 63 | + | |
| 64 | +def clamp(value: int, lo: int, hi: int) -> int: | |
| 65 | + return max(lo, min(hi, value)) | |
| 66 | + | |
| 67 | + | |
| 68 | +def csv_list(value: str | None, *, maxlen: int = 20) -> list[str]: | |
| 69 | + if not value: | |
| 70 | + return [] | |
| 71 | + return [x.strip() for x in value.split(",") if x.strip()][:maxlen] | |
| 72 | + | |
| 73 | + | |
| 74 | +# ------------------------------------------------------------------------------------------------ companies | |
| 75 | + | |
| 76 | + | |
| 77 | +async def get_company(conn: AsyncConnection, key: str) -> dict[str, Any] | None: | |
| 78 | + key = (key or "").strip() | |
| 79 | + if not key or len(key) > 200: | |
| 80 | + return None | |
| 81 | + return await fetch_one(conn, "select * from companies where slug = :k or id = :k limit 1", k=key) | |
| 82 | + | |
| 83 | + | |
| 84 | +async def require_company(conn: AsyncConnection, key: str) -> dict[str, Any]: | |
| 85 | + row = await get_company(conn, key) | |
| 86 | + if row is None: | |
| 87 | + raise HTTPException(status_code=404, detail="company not found") | |
| 88 | + return row | |
| 89 | + | |
| 90 | + | |
| 91 | +COMPANY_CARD_SQL = """ | |
| 92 | +select c.id, c.slug, c.display_name, c.legal_name, c.canonical_domain, c.website, c.description, c.industries, c.industry_primary, | |
| 93 | + c.country, c.hq_city, c.hq_region, c.public_company, c.ticker, c.exchange, c.founded_year, c.employees_band, c.logo_url, | |
| 94 | + c.status, c.onboarding_status, c.onboarding_error, c.importance, c.tier, c.indexed, c.stats, c.last_event_at, c.last_observed_at, | |
| 95 | + c.first_observed_at, c.discovered_at, c.created_at, c.updated_at, | |
| 96 | + mx.metrics, sc.sensors, sc.observations, sc.changes, sc.events, jc.jobs_open{sparkline_col} | |
| 97 | +from companies c | |
| 98 | +left join lateral (select jsonb_object_agg(m.metric, m.value) as metrics from metrics_current m where m.company_id = c.id) mx on true | |
| 99 | +left join lateral (select count(*) filter (where s.status <> 'retired') as sensors, coalesce(sum(s.observation_count), 0) as observations, | |
| 100 | + coalesce(sum(s.change_count), 0) as changes, coalesce(sum(s.event_count), 0) as events | |
| 101 | + from sensors s where s.company_id = c.id) sc on true | |
| 102 | +left join lateral (select count(*) as jobs_open from jobs j where j.company_id = c.id and j.status = 'open') jc on true | |
| 103 | +{sparkline_join} | |
| 104 | +where c.id = any(cast(:ids as text[])) | |
| 105 | +""" | |
| 106 | +SPARKLINE_COL = ", sp.sparkline" | |
| 107 | +SPARKLINE_JOIN = ("left join lateral (select array_agg(x.value order by x.day) as sparkline from (select day, value from metric_series ms " | |
| 108 | + "where ms.company_id = c.id and ms.metric = 'activity_score' and ms.day >= :spark_from order by day desc limit 30) x) sp on true") | |
| 109 | + | |
| 110 | +COMPANY_SORTS: dict[str, str] = { | |
| 111 | + "activity": "ma.value desc nulls last, c.importance desc, c.id", | |
| 112 | + "events": "c.last_event_at desc nulls last, c.importance desc, c.id", | |
| 113 | + "hiring": "mh.value desc nulls last, c.importance desc, c.id", | |
| 114 | + "name": "c.display_name asc, c.id", | |
| 115 | + "importance": "c.importance desc, c.display_name asc, c.id", | |
| 116 | + "recent": "c.discovered_at desc, c.id", | |
| 117 | + # only valid together with a `q` filter (binds :q); list endpoints switch to it automatically when q is present | |
| 118 | + "relevance": ("greatest(coalesce(ts_rank(c.search, websearch_to_tsquery('simple', :q)), 0), similarity(c.display_name, :q), " | |
| 119 | + "similarity(c.canonical_domain, :q)) desc, c.importance desc, c.id"), | |
| 120 | +} | |
| 121 | +COMPANY_SORT_JOINS: dict[str, str] = { | |
| 122 | + "activity": "left join metrics_current ma on ma.company_id = c.id and ma.metric = 'activity_score'", | |
| 123 | + "hiring": "left join metrics_current mh on mh.company_id = c.id and mh.metric = 'hiring_momentum_30d'", | |
| 124 | +} | |
| 125 | + | |
| 126 | + | |
| 127 | +async def fetch_cards_by_ids(conn: AsyncConnection, ids: list[str], *, sparkline: bool = False) -> list[dict[str, Any]]: | |
| 128 | + """Card rows for the given ids, in the given order.""" | |
| 129 | + if not ids: | |
| 130 | + return [] | |
| 131 | + sql = COMPANY_CARD_SQL.format(sparkline_col=SPARKLINE_COL if sparkline else "", sparkline_join=SPARKLINE_JOIN if sparkline else "") | |
| 132 | + params: dict[str, Any] = {"ids": ids} | |
| 133 | + if sparkline: | |
| 134 | + params["spark_from"] = (now_utc() - timedelta(days=30)).date() | |
| 135 | + rows = await fetch_all(conn, sql, **params) | |
| 136 | + by_id = {r["id"]: r for r in rows} | |
| 137 | + return [by_id[i] for i in ids if i in by_id] | |
| 138 | + | |
| 139 | + | |
| 140 | +async def company_page_ids(conn: AsyncConnection, where: list[str], params: dict[str, Any], *, sort: str = "activity", limit: int = 25, | |
| 141 | + offset: int = 0, extra_join: str = "") -> tuple[list[str], int]: | |
| 142 | + """Phase 1 of a company list: ids for the page + bounded total.""" | |
| 143 | + sort = sort if sort in COMPANY_SORTS else "activity" | |
| 144 | + if sort == "relevance" and "q" not in params: | |
| 145 | + sort = "activity" | |
| 146 | + elif sort == "activity" and params.get("q"): | |
| 147 | + sort = "relevance" | |
| 148 | + join = " ".join(x for x in (COMPANY_SORT_JOINS.get(sort, ""), extra_join) if x) | |
| 149 | + where_sql = (" where " + " and ".join(where)) if where else "" | |
| 150 | + rows = await fetch_all(conn, f"select c.id from companies c {join}{where_sql} order by {COMPANY_SORTS[sort]} limit :limit offset :offset", | |
| 151 | + **params, limit=limit, offset=offset) | |
| 152 | + total = await bounded_count(conn, f"from companies c {extra_join}{where_sql}", params) | |
| 153 | + return [r["id"] for r in rows], total | |
| 154 | + | |
| 155 | + | |
| 156 | +async def bounded_count(conn: AsyncConnection, from_where_sql: str, params: dict[str, Any], cap: int = COUNT_CAP) -> int: | |
| 157 | + val = await fetch_val(conn, f"select count(*) from (select 1 {from_where_sql} limit :cap) t", **params, cap=cap) | |
| 158 | + return int(val or 0) | |
| 159 | + | |
| 160 | + | |
| 161 | +def company_filters(*, q: str | None = None, country: str | None = None, industry: str | None = None, tier: int | None = None, | |
| 162 | + public: bool | None = None, status: str | None = None, has_events: bool | None = None, | |
| 163 | + onboarding_status: str | None = None) -> tuple[list[str], dict[str, Any]]: | |
| 164 | + where: list[str] = [] | |
| 165 | + params: dict[str, Any] = {} | |
| 166 | + if q: | |
| 167 | + q = q.strip()[:200] | |
| 168 | + params["q"] = q | |
| 169 | + if len(q) >= 3: | |
| 170 | + where.append("(c.search @@ websearch_to_tsquery('simple', :q) or c.display_name % :q or c.canonical_domain % :q)") | |
| 171 | + else: | |
| 172 | + where.append("(c.display_name ilike :q_prefix or c.canonical_domain ilike :q_prefix)") | |
| 173 | + params["q_prefix"] = q.replace("%", "").replace("_", "") + "%" | |
| 174 | + if country: | |
| 175 | + where.append("c.country = cast(:country as char(2))") | |
| 176 | + params["country"] = country.strip().upper()[:2] | |
| 177 | + if industry: | |
| 178 | + where.append("cast(:industry as text) = any(c.industries)") | |
| 179 | + params["industry"] = industry.strip()[:80] | |
| 180 | + if tier is not None: | |
| 181 | + where.append("c.tier = :tier") | |
| 182 | + params["tier"] = tier | |
| 183 | + if public is not None: | |
| 184 | + where.append("c.public_company = :public") | |
| 185 | + params["public"] = public | |
| 186 | + if status: | |
| 187 | + where.append("c.status = cast(:status as text)") | |
| 188 | + params["status"] = status.strip().upper()[:40] | |
| 189 | + if has_events is True: | |
| 190 | + where.append("c.last_event_at is not null") | |
| 191 | + elif has_events is False: | |
| 192 | + where.append("c.last_event_at is null") | |
| 193 | + if onboarding_status: | |
| 194 | + where.append("c.onboarding_status = cast(:onboarding_status as text)") | |
| 195 | + params["onboarding_status"] = onboarding_status.strip().lower()[:40] | |
| 196 | + return where, params | |
| 197 | + | |
| 198 | + | |
| 199 | +# ------------------------------------------------------------------------------------------------ events | |
| 200 | + | |
| 201 | +EVENT_SELECT = """ | |
| 202 | +select e.id, e.company_id, e.sensor_id, e.change_id, e.cluster_id, e.surface, e.event_type, e.event_subtype, e.importance, e.confidence, | |
| 203 | + e.confidence_label, e.title, e.summary, e.old_value, e.new_value, e.payload, e.entities, e.tags, e.detected_at, e.effective_at, | |
| 204 | + e.published_at, e.source_url, e.snapshot_before, e.snapshot_after, e.language, e.origin, e.model_name, e.prompt_version, e.status, | |
| 205 | + e.retracted_reason, | |
| 206 | + c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, c.country as company_country, | |
| 207 | + c.logo_url as company_logo_url | |
| 208 | +from events e join companies c on c.id = e.company_id | |
| 209 | +""" | |
| 210 | +EVENT_SORTS: dict[str, str] = {"recent": "e.detected_at desc, e.id desc", "importance": "e.importance desc, e.detected_at desc, e.id desc"} | |
| 211 | +_EVENT_TYPES = {t.value for t in EventType} | |
| 212 | + | |
| 213 | + | |
| 214 | +def event_filters(*, company_id: str | None = None, event_type: str | None = None, event_subtype: str | None = None, country: str | None = None, | |
| 215 | + industry: str | None = None, since: datetime | None = None, until: datetime | None = None, min_importance: float | None = None, | |
| 216 | + min_confidence: float | None = None, q: str | None = None, surface: str | None = None, origin: str | None = None, | |
| 217 | + status: str | None = "active", event_types: list[str] | None = None, | |
| 218 | + event_subtypes: list[str] | None = None) -> tuple[list[str], dict[str, Any]]: | |
| 219 | + where: list[str] = [] | |
| 220 | + params: dict[str, Any] = {} | |
| 221 | + if status: | |
| 222 | + where.append("e.status = cast(:e_status as text)") | |
| 223 | + params["e_status"] = status | |
| 224 | + if company_id: | |
| 225 | + where.append("e.company_id = :e_company_id") | |
| 226 | + params["e_company_id"] = company_id | |
| 227 | + if event_type: | |
| 228 | + where.append("e.event_type = cast(:e_type as text)") | |
| 229 | + params["e_type"] = event_type.strip().upper()[:40] | |
| 230 | + if event_types: | |
| 231 | + where.append("e.event_type = any(cast(:e_types as text[]))") | |
| 232 | + params["e_types"] = [t.upper()[:40] for t in event_types][:30] | |
| 233 | + if event_subtype: | |
| 234 | + where.append("e.event_subtype = cast(:e_subtype as text)") | |
| 235 | + params["e_subtype"] = event_subtype.strip().upper()[:60] | |
| 236 | + if event_subtypes: | |
| 237 | + where.append("e.event_subtype = any(cast(:e_subtypes as text[]))") | |
| 238 | + params["e_subtypes"] = [t.upper()[:60] for t in event_subtypes][:60] | |
| 239 | + if country: | |
| 240 | + where.append("c.country = cast(:e_country as char(2))") | |
| 241 | + params["e_country"] = country.strip().upper()[:2] | |
| 242 | + if industry: | |
| 243 | + where.append("cast(:e_industry as text) = any(c.industries)") | |
| 244 | + params["e_industry"] = industry.strip()[:80] | |
| 245 | + if since is not None: | |
| 246 | + where.append("e.detected_at >= :e_since") | |
| 247 | + params["e_since"] = since | |
| 248 | + if until is not None: | |
| 249 | + where.append("e.detected_at < :e_until") | |
| 250 | + params["e_until"] = until | |
| 251 | + if min_importance is not None: | |
| 252 | + where.append("e.importance >= :e_min_imp") | |
| 253 | + params["e_min_imp"] = float(min_importance) | |
| 254 | + if min_confidence is not None: | |
| 255 | + where.append("e.confidence >= :e_min_conf") | |
| 256 | + params["e_min_conf"] = float(min_confidence) | |
| 257 | + if q: | |
| 258 | + where.append("e.search @@ websearch_to_tsquery('english', :e_q)") | |
| 259 | + params["e_q"] = q.strip()[:200] | |
| 260 | + if surface: | |
| 261 | + where.append("e.surface = cast(:e_surface as text)") | |
| 262 | + params["e_surface"] = surface.strip().lower()[:40] | |
| 263 | + if origin: | |
| 264 | + where.append("e.origin = cast(:e_origin as text)") | |
| 265 | + params["e_origin"] = origin.strip().lower()[:20] | |
| 266 | + return where, params | |
| 267 | + | |
| 268 | + | |
| 269 | +async def fetch_events(conn: AsyncConnection, where: list[str], params: dict[str, Any], *, sort: str = "recent", limit: int = 50, | |
| 270 | + offset: int = 0) -> list[dict[str, Any]]: | |
| 271 | + order = EVENT_SORTS.get(sort, EVENT_SORTS["recent"]) | |
| 272 | + where_sql = (" where " + " and ".join(where)) if where else "" | |
| 273 | + return await fetch_all(conn, f"{EVENT_SELECT}{where_sql} order by {order} limit :limit offset :offset", **params, | |
| 274 | + limit=clamp(limit, 1, MAX_LIMIT), offset=max(0, offset)) | |
| 275 | + | |
| 276 | + | |
| 277 | +async def count_events(conn: AsyncConnection, where: list[str], params: dict[str, Any]) -> int: | |
| 278 | + where_sql = (" where " + " and ".join(where)) if where else "" | |
| 279 | + return await bounded_count(conn, f"from events e join companies c on c.id = e.company_id{where_sql}", params) | |
| 280 | + | |
| 281 | + | |
| 282 | +async def fetch_event(conn: AsyncConnection, event_id: str) -> dict[str, Any] | None: | |
| 283 | + return await fetch_one(conn, f"{EVENT_SELECT} where e.id = :id", id=event_id) | |
| 284 | + | |
| 285 | + | |
| 286 | +# ------------------------------------------------------------------------------------------------ metrics helpers | |
| 287 | + | |
| 288 | + | |
| 289 | +async def metric_series(conn: AsyncConnection, company_id: str, metric: str, days: int) -> list[dict[str, Any]]: | |
| 290 | + return await fetch_all(conn, "select day, value, confidence from metric_series where company_id = :cid and metric = :m and day >= :d " | |
| 291 | + "order by day asc limit :lim", cid=company_id, m=metric, d=days_ago(days).date(), lim=clamp(days, 1, 2000)) | |
| 292 | + | |
| 293 | + | |
| 294 | +async def metric_values_at(conn: AsyncConnection, ids: list[str], metric: str, on_or_before: date) -> dict[str, float]: | |
| 295 | + """Latest series value on/before a day per company (used for window deltas).""" | |
| 296 | + if not ids: | |
| 297 | + return {} | |
| 298 | + rows = await fetch_all(conn, "select distinct on (company_id) company_id, value from metric_series where metric = :m and " | |
| 299 | + "company_id = any(cast(:ids as text[])) and day <= :d order by company_id, day desc", m=metric, ids=ids, d=on_or_before) | |
| 300 | + return {r["company_id"]: float(r["value"]) for r in rows} | |
| 301 | + | |
| 302 | + | |
| 303 | +async def settings_value(conn: AsyncConnection, key: str) -> Any: | |
| 304 | + return await fetch_val(conn, "select value from settings_kv where key = :k", k=key) | |
| 305 | + | |
| 306 | + | |
| 307 | +__all__ = ["COMPANY_SORTS", "COUNT_CAP", "EVENT_SORTS", "MAX_LIMIT", "WINDOWS", "bounded_count", "clamp", "company_filters", "company_page_ids", | |
| 308 | + "count_events", "csv_list", "days_ago", "event_filters", "fetch_cards_by_ids", "fetch_event", "fetch_events", "get_company", | |
| 309 | + "metric_series", "metric_values_at", "now_utc", "parse_bool", "parse_iso", "require_company", "settings_value", "window_start"] | |
added
src/companyatlas/api/ratelimit.py
+173 −0
@@ -0,0 +1,173 @@ | ||
| 1 | +"""In-process token-bucket rate limiting (spec §136: tiers anonymous / authenticated / paid / internal). | |
| 2 | + | |
| 3 | +Key = client IP (first hop of `X-Forwarded-For`) for anonymous traffic, or the API key id when a valid `X-CA-API-Key` is presented. | |
| 4 | +API keys are looked up by sha256 hash in `api_keys` (60 s in-process cache, negative results too); `last_used_at`/`request_count` | |
| 5 | +are flushed lazily every `FLUSH_EVERY_S`. Admin requests (valid `X-CA-Admin-Token`) and health/docs paths bypass the limiter. | |
| 6 | +`429` carries `Retry-After`; every limited response carries `X-RateLimit-Limit` / `X-RateLimit-Remaining` / `X-RateLimit-Tier`. | |
| 7 | +""" | |
| 8 | +from __future__ import annotations | |
| 9 | + | |
| 10 | +import asyncio | |
| 11 | +import hashlib | |
| 12 | +import logging | |
| 13 | +import time | |
| 14 | +from dataclasses import dataclass, field | |
| 15 | +from datetime import UTC, datetime | |
| 16 | +from typing import Any | |
| 17 | + | |
| 18 | +from fastapi import Request | |
| 19 | + | |
| 20 | +from companyatlas.api.common import AtlasJSONResponse, TTLCache, client_ip, is_admin_request | |
| 21 | +from companyatlas.db import connection, execute, fetch_one, transaction | |
| 22 | + | |
| 23 | +log = logging.getLogger("companyatlas.api.ratelimit") | |
| 24 | + | |
| 25 | +TIER_LIMITS_PER_MIN: dict[str, int | None] = {"anonymous": 120, "authenticated": 600, "paid": 3000, "internal": None} | |
| 26 | +KEY_CACHE_TTL_S = 60.0 | |
| 27 | +FLUSH_EVERY_S = 30.0 | |
| 28 | +BYPASS_PREFIXES = ("/health", "/ready", "/api/v1/health", "/api/v1/docs", "/api/v1/openapi.json") | |
| 29 | +MAX_BUCKETS = 50_000 | |
| 30 | + | |
| 31 | + | |
| 32 | +@dataclass(slots=True) | |
| 33 | +class Bucket: | |
| 34 | + tokens: float | |
| 35 | + updated: float | |
| 36 | + | |
| 37 | + | |
| 38 | +@dataclass | |
| 39 | +class RateLimiter: | |
| 40 | + limits: dict[str, int | None] = field(default_factory=lambda: dict(TIER_LIMITS_PER_MIN)) | |
| 41 | + buckets: dict[str, Bucket] = field(default_factory=dict) | |
| 42 | + | |
| 43 | + def take(self, key: str, tier: str, now: float | None = None) -> tuple[bool, int, int, float]: | |
| 44 | + """Consume one token. Returns (allowed, limit, remaining, retry_after_s).""" | |
| 45 | + limit = self.limits.get(tier, self.limits["anonymous"]) | |
| 46 | + if limit is None: | |
| 47 | + return True, 0, 0, 0.0 | |
| 48 | + now = time.monotonic() if now is None else now | |
| 49 | + rate = limit / 60.0 | |
| 50 | + b = self.buckets.get(key) | |
| 51 | + if b is None: | |
| 52 | + if len(self.buckets) >= MAX_BUCKETS: | |
| 53 | + self._evict(now) | |
| 54 | + b = self.buckets[key] = Bucket(tokens=float(limit), updated=now) | |
| 55 | + else: | |
| 56 | + b.tokens = min(float(limit), b.tokens + (now - b.updated) * rate) | |
| 57 | + b.updated = now | |
| 58 | + if b.tokens >= 1.0: | |
| 59 | + b.tokens -= 1.0 | |
| 60 | + return True, limit, int(b.tokens), 0.0 | |
| 61 | + retry = (1.0 - b.tokens) / rate | |
| 62 | + return False, limit, 0, retry | |
| 63 | + | |
| 64 | + def _evict(self, now: float) -> None: | |
| 65 | + stale = [k for k, b in self.buckets.items() if now - b.updated > 120] | |
| 66 | + for k in stale: | |
| 67 | + self.buckets.pop(k, None) | |
| 68 | + if len(self.buckets) >= MAX_BUCKETS: # still full: drop the oldest half | |
| 69 | + for k in sorted(self.buckets, key=lambda k: self.buckets[k].updated)[: MAX_BUCKETS // 2]: | |
| 70 | + self.buckets.pop(k, None) | |
| 71 | + | |
| 72 | + | |
| 73 | +limiter = RateLimiter() | |
| 74 | +_key_cache = TTLCache(max_items=10_000) | |
| 75 | +_usage: dict[str, int] = {} | |
| 76 | +_usage_lock = asyncio.Lock() | |
| 77 | +_last_flush = time.monotonic() | |
| 78 | + | |
| 79 | + | |
| 80 | +def hash_key(raw: str) -> str: | |
| 81 | + return hashlib.sha256(raw.encode("utf-8")).hexdigest() | |
| 82 | + | |
| 83 | + | |
| 84 | +async def resolve_api_key(raw: str | None) -> dict[str, Any] | None: | |
| 85 | + """`{id, tier, name}` for a valid, non-revoked key; None otherwise (cached 60 s either way).""" | |
| 86 | + if not raw or len(raw) < 16 or len(raw) > 200: | |
| 87 | + return None | |
| 88 | + h = hash_key(raw) | |
| 89 | + hit = _key_cache.get(h) | |
| 90 | + if hit is not None: | |
| 91 | + return hit or None | |
| 92 | + try: | |
| 93 | + async with connection() as conn: | |
| 94 | + row = await fetch_one(conn, "select id, tier, name from api_keys where key_hash = :h and revoked_at is null", h=h) | |
| 95 | + except Exception: | |
| 96 | + log.warning("api key lookup failed", exc_info=True) | |
| 97 | + return None | |
| 98 | + info = {"id": row["id"], "tier": row["tier"], "name": row["name"]} if row else {} | |
| 99 | + _key_cache.set(h, info, KEY_CACHE_TTL_S) | |
| 100 | + return info or None | |
| 101 | + | |
| 102 | + | |
| 103 | +async def _write_usage(batch: dict[str, int]) -> None: | |
| 104 | + if not batch: | |
| 105 | + return | |
| 106 | + try: | |
| 107 | + async with transaction() as conn: | |
| 108 | + for kid, n in batch.items(): | |
| 109 | + await execute(conn, "update api_keys set last_used_at = :t, request_count = request_count + :n where id = :id", | |
| 110 | + t=datetime.now(UTC), n=n, id=kid) | |
| 111 | + except Exception: | |
| 112 | + log.warning("api key usage flush failed", exc_info=True) | |
| 113 | + | |
| 114 | + | |
| 115 | +async def _note_usage(key_id: str) -> None: | |
| 116 | + global _last_flush | |
| 117 | + async with _usage_lock: | |
| 118 | + _usage[key_id] = _usage.get(key_id, 0) + 1 | |
| 119 | + if time.monotonic() - _last_flush < FLUSH_EVERY_S: | |
| 120 | + return | |
| 121 | + batch, _last_flush = dict(_usage), time.monotonic() | |
| 122 | + _usage.clear() | |
| 123 | + await _write_usage(batch) | |
| 124 | + | |
| 125 | + | |
| 126 | +async def flush_usage() -> None: | |
| 127 | + """Force a usage flush (tests / shutdown).""" | |
| 128 | + global _last_flush | |
| 129 | + async with _usage_lock: | |
| 130 | + batch, _last_flush = dict(_usage), time.monotonic() | |
| 131 | + _usage.clear() | |
| 132 | + await _write_usage(batch) | |
| 133 | + | |
| 134 | + | |
| 135 | +def _is_local_server_side(request: Request) -> bool: | |
| 136 | + if request.headers.get("x-forwarded-for") or request.headers.get("x-real-ip"): | |
| 137 | + return False | |
| 138 | + host = request.client.host if request.client else "" | |
| 139 | + return host in ("127.0.0.1", "::1") | |
| 140 | + | |
| 141 | + | |
| 142 | +async def rate_limit_middleware(request: Request, call_next): # type: ignore[no-untyped-def] | |
| 143 | + path = request.url.path | |
| 144 | + if request.method == "OPTIONS" or path.startswith(BYPASS_PREFIXES): | |
| 145 | + return await call_next(request) | |
| 146 | + if is_admin_request(request): | |
| 147 | + response = await call_next(request) | |
| 148 | + response.headers["x-ratelimit-tier"] = "admin" | |
| 149 | + return response | |
| 150 | + key_info = await resolve_api_key(request.headers.get("x-ca-api-key")) | |
| 151 | + if key_info: | |
| 152 | + tier, bucket_key = key_info["tier"], f"key:{key_info['id']}" | |
| 153 | + asyncio.get_running_loop().create_task(_note_usage(key_info["id"])) | |
| 154 | + elif _is_local_server_side(request): | |
| 155 | + # Next.js server components fetch the loopback API without X-Forwarded-For: that is our own SSR, not a public client. | |
| 156 | + # Browser traffic proxied through the Next rewrite (and through Caddy) always carries X-Forwarded-For and stays anonymous. | |
| 157 | + tier, bucket_key = "internal", "internal:ssr" | |
| 158 | + else: | |
| 159 | + tier, bucket_key = "anonymous", f"ip:{client_ip(request)}" | |
| 160 | + allowed, limit, remaining, retry = limiter.take(bucket_key, tier) | |
| 161 | + if not allowed: | |
| 162 | + return AtlasJSONResponse({"detail": "rate limit exceeded"}, status_code=429, | |
| 163 | + headers={"retry-after": str(max(1, int(retry + 0.999))), "x-ratelimit-limit": str(limit), | |
| 164 | + "x-ratelimit-remaining": "0", "x-ratelimit-tier": tier, "cache-control": "no-store"}) | |
| 165 | + response = await call_next(request) | |
| 166 | + response.headers["x-ratelimit-tier"] = tier | |
| 167 | + if limit: | |
| 168 | + response.headers["x-ratelimit-limit"] = str(limit) | |
| 169 | + response.headers["x-ratelimit-remaining"] = str(remaining) | |
| 170 | + return response | |
| 171 | + | |
| 172 | + | |
| 173 | +__all__ = ["TIER_LIMITS_PER_MIN", "RateLimiter", "flush_usage", "hash_key", "limiter", "rate_limit_middleware", "resolve_api_key"] | |
added
src/companyatlas/api/routers/admin.py
+561 −0
@@ -0,0 +1,561 @@ | ||
| 1 | +"""Admin API (`X-CA-Admin-Token`). Nothing is ever deleted: actions update `status` columns and append audit rows in payloads.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from fastapi import APIRouter, Depends, HTTPException, Query, Response | |
| 8 | +from pydantic import BaseModel, Field | |
| 9 | + | |
| 10 | +from companyatlas import archive | |
| 11 | +from companyatlas.api import aggregates as agg | |
| 12 | +from companyatlas.api import queries as q | |
| 13 | +from companyatlas.api import serializers as ser | |
| 14 | +from companyatlas.api.common import NO_STORE, PageDep, cache, page_payload, require_admin | |
| 15 | +from companyatlas.config import settings | |
| 16 | +from companyatlas.db import connection, execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 17 | +from companyatlas.ids import new_id, slugify | |
| 18 | +from companyatlas.taxonomy import FailureClass, SensorStatus, tier_for_interval | |
| 19 | +from companyatlas.urls import canonicalize_url, registrable_domain | |
| 20 | + | |
| 21 | +ORDER = 10 | |
| 22 | +router = APIRouter(prefix="/api/v1/admin", tags=["admin"], dependencies=[Depends(require_admin)]) | |
| 23 | +LOW_QUALITY = 30.0 | |
| 24 | +STALE_DAYS = 7 | |
| 25 | +WORKER_WINDOW_MIN = 15 | |
| 26 | + | |
| 27 | + | |
| 28 | +def _ns(response: Response) -> None: | |
| 29 | + response.headers["cache-control"] = NO_STORE | |
| 30 | + | |
| 31 | + | |
| 32 | +# ------------------------------------------------------------------------------------------------ overview | |
| 33 | + | |
| 34 | + | |
| 35 | +@router.get("/overview") | |
| 36 | +async def overview(response: Response) -> dict[str, Any]: | |
| 37 | + _ns(response) | |
| 38 | + today = datetime.now(UTC).replace(hour=0, minute=0, second=0, microsecond=0) | |
| 39 | + async with connection() as conn: | |
| 40 | + cs = await fetch_all(conn, "select status, count(*) as n from companies group by status") | |
| 41 | + cos = await fetch_all(conn, "select onboarding_status, count(*) as n from companies group by onboarding_status") | |
| 42 | + ss = await fetch_all(conn, "select status, count(*) as n from sensors group by status") | |
| 43 | + st = await fetch_all(conn, "select tier, count(*) as n from sensors where status <> 'retired' group by tier") | |
| 44 | + qs = await fetch_one(conn, "select count(*) filter (where status = 'pending') as pending, count(*) filter (where status = 'running') as running, " | |
| 45 | + "count(*) filter (where status = 'dead') as dead, count(*) filter (where status = 'failed') as failed, " | |
| 46 | + "extract(epoch from (now() - min(run_at) filter (where status = 'pending' and run_at <= now()))) as oldest_pending_s from queue_jobs") or {} | |
| 47 | + llm = await fetch_one(conn, "select count(*) filter (where status = 'pending') as pending, count(*) filter (where status = 'done' and finished_at >= :t) as done_today, " | |
| 48 | + "count(*) filter (where status = 'failed' and finished_at >= :t) as failed_today from llm_jobs", t=today) or {} | |
| 49 | + fails = await fetch_all(conn, "select failure_class, count(*) as n from failures where at >= now() - interval '24 hours' group by failure_class order by n desc") | |
| 50 | + rates = await fetch_one(conn, "select (select count(*) from observations where fetched_at >= now() - interval '1 hour') as fetch_rate_1h, " | |
| 51 | + "(select count(*) from changes where detected_at >= now() - interval '1 hour') as change_rate_1h, " | |
| 52 | + "(select count(*) from changes where detected_at >= now() - interval '1 hour' and kind in ('meaningful','major','critical')) as meaningful_rate_1h") or {} | |
| 53 | + workers = await fetch_all(conn, "select worker as name, max(fetched_at) as last_seen_at, count(*) as fetches_15m from observations " | |
| 54 | + "where fetched_at >= now() - make_interval(mins => :m) and worker is not null group by worker order by last_seen_at desc limit 50", | |
| 55 | + m=WORKER_WINDOW_MIN) | |
| 56 | + inflight = {r["claimed_by"]: int(r["n"]) for r in await fetch_all(conn, "select claimed_by, count(*) as n from sensors where claimed_by is not null group by claimed_by")} | |
| 57 | + runs = await fetch_all(conn, "select worker as name, max(started_at) as last_seen_at from crawl_runs where started_at >= now() - make_interval(mins => :m) " | |
| 58 | + "and worker is not null group by worker", m=WORKER_WINDOW_MIN) | |
| 59 | + cost = await fetch_all(conn, "select dimension, sum(cost_estimate) as cost, sum(units) as units from cost_ledger where day = current_date group by dimension") | |
| 60 | + periodic = await q.settings_value(conn, "scheduler:heartbeat") | |
| 61 | + names = {w["name"] for w in workers} | {r["name"] for r in runs} | set(inflight) | |
| 62 | + seen = {w["name"]: w["last_seen_at"] for w in workers} | |
| 63 | + for r in runs: | |
| 64 | + if r["name"] not in seen or (r["last_seen_at"] and seen[r["name"]] and r["last_seen_at"] > seen[r["name"]]): | |
| 65 | + seen[r["name"]] = r["last_seen_at"] | |
| 66 | + done_today, failed_today = int(llm.get("done_today") or 0), int(llm.get("failed_today") or 0) | |
| 67 | + costs = {r["dimension"]: round(float(r["cost"] or 0), 4) for r in cost} | |
| 68 | + return {"companies_by_status": {r["status"]: int(r["n"]) for r in cs}, "companies_by_onboarding": {r["onboarding_status"]: int(r["n"]) for r in cos}, | |
| 69 | + "sensors_by_status": {r["status"]: int(r["n"]) for r in ss}, "sensors_by_tier": {(r["tier"] or "").strip(): int(r["n"]) for r in st}, | |
| 70 | + "queue": {"pending": int(qs.get("pending") or 0), "running": int(qs.get("running") or 0), "dead": int(qs.get("dead") or 0), | |
| 71 | + "failed": int(qs.get("failed") or 0), "oldest_pending_s": round(float(qs["oldest_pending_s"]), 1) if qs.get("oldest_pending_s") is not None else 0.0}, | |
| 72 | + "llm": {"pending": int(llm.get("pending") or 0), "done_today": done_today, "failed_today": failed_today, | |
| 73 | + "budget_left": max(0, settings.llm_daily_budget - done_today - failed_today), "budget": settings.llm_daily_budget, "configured": settings.llm_configured}, | |
| 74 | + "failures_24h_by_class": {r["failure_class"]: int(r["n"]) for r in fails}, "fetch_rate_1h": int(rates.get("fetch_rate_1h") or 0), | |
| 75 | + "change_rate_1h": int(rates.get("change_rate_1h") or 0), "meaningful_rate_1h": int(rates.get("meaningful_rate_1h") or 0), | |
| 76 | + "storage": await agg.archive_stats(), | |
| 77 | + "workers": [{"name": n, "last_seen_at": seen.get(n), "inflight": inflight.get(n, 0)} for n in sorted(names)], | |
| 78 | + "cost_today": {"fetch": costs.get("fetch", 0.0), "browser": costs.get("browser", 0.0), "llm": costs.get("llm", 0.0), "total": round(sum(costs.values()), 4)}, | |
| 79 | + "scheduler_heartbeat": periodic, "time": datetime.now(UTC)} | |
| 80 | + | |
| 81 | + | |
| 82 | +# ------------------------------------------------------------------------------------------------ connectors | |
| 83 | + | |
| 84 | + | |
| 85 | +@router.get("/connectors") | |
| 86 | +async def connectors(response: Response) -> dict[str, Any]: | |
| 87 | + _ns(response) | |
| 88 | + async with connection() as conn: | |
| 89 | + rows = await fetch_all(conn, "select * from connectors order by category, id") | |
| 90 | + sensors = await fetch_all(conn, "select connector_id, count(*) filter (where status = 'active') as active, " | |
| 91 | + "count(*) filter (where status in ('failing','stale','blocked')) as failing, count(*) filter (where status <> 'retired') as total, " | |
| 92 | + "max(last_run_at) as last_run_at from sensors group by connector_id") | |
| 93 | + obs = await fetch_all(conn, "select s.connector_id, count(*) as n, count(*) filter (where o.failure_class is null) as ok, avg(o.duration_ms) as latency, " | |
| 94 | + "count(*) filter (where o.changed) as changed, count(*) filter (where o.failure_class is not null) as errors " | |
| 95 | + "from observations o join sensors s on s.id = o.sensor_id where o.fetched_at >= now() - interval '24 hours' group by s.connector_id") | |
| 96 | + sm = {r["connector_id"]: r for r in sensors} | |
| 97 | + om = {r["connector_id"]: r for r in obs} | |
| 98 | + items = [] | |
| 99 | + known = {r["id"] for r in rows} | |
| 100 | + for cid in list(known) + [k for k in sm if k not in known]: | |
| 101 | + row = next((r for r in rows if r["id"] == cid), None) | |
| 102 | + base = ser.connector(row) if row else {"id": cid, "name": cid, "version": "?", "category": "?", "enabled": True} | |
| 103 | + s, o = sm.get(cid, {}), om.get(cid, {}) | |
| 104 | + n = int(o.get("n") or 0) | |
| 105 | + base.update({"sensors_active": int(s.get("active") or 0), "sensors_failing": int(s.get("failing") or 0), "sensors_total": int(s.get("total") or 0), | |
| 106 | + "success_rate_24h": round(int(o.get("ok") or 0) / n, 4) if n else None, "avg_latency_ms": round(float(o["latency"]), 1) if o.get("latency") is not None else None, | |
| 107 | + "change_rate_24h": round(int(o.get("changed") or 0) / n, 4) if n else None, "errors_24h": int(o.get("errors") or 0), "fetches_24h": n, | |
| 108 | + "last_run_at": s.get("last_run_at")}) | |
| 109 | + items.append(base) | |
| 110 | + items.sort(key=lambda x: (x.get("category") or "", x["id"])) | |
| 111 | + return {"items": items} | |
| 112 | + | |
| 113 | + | |
| 114 | +# ------------------------------------------------------------------------------------------------ sensors | |
| 115 | + | |
| 116 | +SENSOR_FILTERS: dict[str, str] = { | |
| 117 | + "healthy": "s.status = 'active' and s.consecutive_failures = 0", | |
| 118 | + "failing": "(s.status = 'failing' or s.consecutive_failures > 0)", | |
| 119 | + "stale": "(s.status = 'stale' or (s.status in ('active','failing') and (s.last_success_at is null or s.last_success_at < now() - make_interval(days => :stale_days))))", | |
| 120 | + "blocked": "(s.status = 'blocked' or s.last_failure_class in ('BOT_CHALLENGE','ROBOTS','BLOCKED_DESTINATION','RATE_LIMIT'))", | |
| 121 | + "redirected": "(s.status = 'redirected' or s.last_failure_class = 'REDIRECT')", | |
| 122 | + "low_quality": "s.quality_score < :low_quality", | |
| 123 | + "high_activity": "s.last_change_at >= now() - interval '24 hours'", | |
| 124 | +} | |
| 125 | +SENSOR_SORTS = {"recent": "s.last_run_at desc nulls last, s.id", "next_run": "s.next_run_at asc, s.id", "failures": "s.consecutive_failures desc, s.last_run_at desc nulls last, s.id", | |
| 126 | + "quality": "s.quality_score asc, s.id", "changes": "s.change_count desc, s.id", "created": "s.created_at desc, s.id"} | |
| 127 | +SENSOR_SELECT = ("select s.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, c.country as company_country, " | |
| 128 | + "c.logo_url as company_logo_url from sensors s join companies c on c.id = s.company_id") | |
| 129 | + | |
| 130 | + | |
| 131 | +@router.get("/sensors") | |
| 132 | +async def admin_sensors(response: Response, p: PageDep, status: str | None = None, domain: str | None = None, connector: str | None = None, | |
| 133 | + company: str | None = None, surface: str | None = None, | |
| 134 | + filter: str | None = Query(None, pattern="^(healthy|failing|stale|blocked|redirected|low_quality|high_activity)$"), | |
| 135 | + sort: str = Query("recent", pattern="^(recent|next_run|failures|quality|changes|created)$")) -> dict[str, Any]: | |
| 136 | + _ns(response) | |
| 137 | + where, params = ["true"], {"stale_days": STALE_DAYS, "low_quality": LOW_QUALITY} | |
| 138 | + if status: | |
| 139 | + where.append("s.status = cast(:status as text)") | |
| 140 | + params["status"] = status.lower()[:20] | |
| 141 | + if domain: | |
| 142 | + where.append("s.domain = cast(:domain as text)") | |
| 143 | + params["domain"] = domain.lower()[:253] | |
| 144 | + if connector: | |
| 145 | + where.append("s.connector_id = cast(:connector as text)") | |
| 146 | + params["connector"] = connector[:80] | |
| 147 | + if surface: | |
| 148 | + where.append("s.surface = cast(:surface as text)") | |
| 149 | + params["surface"] = surface.lower()[:40] | |
| 150 | + if filter: | |
| 151 | + where.append(SENSOR_FILTERS[filter]) | |
| 152 | + if filter == "high_activity": | |
| 153 | + sort = "changes" | |
| 154 | + async with connection() as conn: | |
| 155 | + if company: | |
| 156 | + params["company_id"] = (await q.require_company(conn, company))["id"] | |
| 157 | + where.append("s.company_id = :company_id") | |
| 158 | + wsql = " and ".join(where) | |
| 159 | + rows = await fetch_all(conn, f"{SENSOR_SELECT} where {wsql} order by {SENSOR_SORTS[sort]} limit :limit offset :offset", **params, limit=p.per_page, offset=p.offset) | |
| 160 | + total = await q.bounded_count(conn, f"from sensors s join companies c on c.id = s.company_id where {wsql}", params) | |
| 161 | + items = [] | |
| 162 | + for r in rows: | |
| 163 | + item = ser.sensor_admin(r) | |
| 164 | + item["company"] = ser.company_ref(r) | |
| 165 | + items.append(item) | |
| 166 | + return page_payload(items, total, p) | |
| 167 | + | |
| 168 | + | |
| 169 | +class SensorActionBody(BaseModel): | |
| 170 | + interval_s: int | None = Field(None, ge=60, le=90 * 86400) | |
| 171 | + connector_id: str | None = Field(None, max_length=80) | |
| 172 | + reason: str | None = Field(None, max_length=500) | |
| 173 | + | |
| 174 | + | |
| 175 | +SENSOR_ACTIONS = ("pause", "resume", "retry", "rediscover", "retire", "run_now", "set_interval", "set_connector") | |
| 176 | + | |
| 177 | + | |
| 178 | +@router.post("/sensors/{sensor_id}/{action}") | |
| 179 | +async def sensor_action(sensor_id: str, action: str, response: Response, body: SensorActionBody | None = None) -> dict[str, Any]: | |
| 180 | + _ns(response) | |
| 181 | + if action not in SENSOR_ACTIONS: | |
| 182 | + raise HTTPException(status_code=404, detail=f"unknown action (one of {', '.join(SENSOR_ACTIONS)})") | |
| 183 | + body = body or SensorActionBody() | |
| 184 | + now = datetime.now(UTC) | |
| 185 | + async with transaction() as conn: | |
| 186 | + s = await fetch_one(conn, "select * from sensors where id = :id", id=sensor_id) | |
| 187 | + if s is None: | |
| 188 | + raise HTTPException(status_code=404, detail="sensor not found") | |
| 189 | + extra: dict[str, Any] = {} | |
| 190 | + if action == "pause": | |
| 191 | + await execute(conn, "update sensors set status = :st, claimed_by = null, claimed_at = null, updated_at = :now where id = :id", st=SensorStatus.PAUSED.value, now=now, id=sensor_id) | |
| 192 | + elif action == "resume": | |
| 193 | + await execute(conn, "update sensors set status = :st, consecutive_failures = 0, next_run_at = :now, retired_at = null, updated_at = :now where id = :id", | |
| 194 | + st=SensorStatus.ACTIVE.value, now=now, id=sensor_id) | |
| 195 | + elif action in ("retry", "run_now"): | |
| 196 | + prio = ", priority = greatest(priority, 0.95)" if action == "run_now" else "" | |
| 197 | + await execute(conn, f"update sensors set next_run_at = :now, claimed_by = null, claimed_at = null{prio}, updated_at = :now where id = :id", now=now, id=sensor_id) | |
| 198 | + elif action == "retire": | |
| 199 | + await execute(conn, "update sensors set status = :st, retired_at = :now, claimed_by = null, claimed_at = null, updated_at = :now where id = :id", | |
| 200 | + st=SensorStatus.RETIRED.value, now=now, id=sensor_id) | |
| 201 | + elif action == "rediscover": | |
| 202 | + jid = new_id("queue_job") | |
| 203 | + key = f"discover:{s['company_id']}:{int(now.timestamp())}" | |
| 204 | + await execute(conn, "insert into queue_jobs (id, kind, key, payload, priority) values (:id, 'discover', :key, cast(:p as jsonb), 0.9) on conflict (key) do nothing", | |
| 205 | + id=jid, key=key, p=jsonb({"company_id": s["company_id"], "sensor_id": sensor_id, "reason": body.reason or "admin:rediscover"})) | |
| 206 | + extra["queued"] = {"id": jid, "key": key} | |
| 207 | + elif action == "set_interval": | |
| 208 | + if body.interval_s is None: | |
| 209 | + raise HTTPException(status_code=422, detail="interval_s required") | |
| 210 | + iv = max(settings.min_interval_s, min(settings.max_interval_s, body.interval_s)) | |
| 211 | + await execute(conn, "update sensors set base_interval_s = :iv, current_interval_s = :iv, tier = :tier, next_run_at = :now, updated_at = :now where id = :id", | |
| 212 | + iv=iv, tier=tier_for_interval(iv), now=now, id=sensor_id) | |
| 213 | + extra["interval_s"] = iv | |
| 214 | + elif action == "set_connector": | |
| 215 | + if not body.connector_id: | |
| 216 | + raise HTTPException(status_code=422, detail="connector_id required") | |
| 217 | + ok = await fetch_val(conn, "select enabled from connectors where id = :c", c=body.connector_id) | |
| 218 | + if ok is None: | |
| 219 | + raise HTTPException(status_code=422, detail="unknown connector_id") | |
| 220 | + await execute(conn, "update sensors set connector_id = :c, next_run_at = :now, updated_at = :now where id = :id", c=body.connector_id, now=now, id=sensor_id) | |
| 221 | + row = await fetch_one(conn, f"{SENSOR_SELECT} where s.id = :id", id=sensor_id) | |
| 222 | + out = ser.sensor_admin(row or {}) | |
| 223 | + out["company"] = ser.company_ref(row or {}) | |
| 224 | + return {"ok": True, "action": action, "sensor": out, **extra} | |
| 225 | + | |
| 226 | + | |
| 227 | +# ------------------------------------------------------------------------------------------------ companies | |
| 228 | + | |
| 229 | + | |
| 230 | +@router.get("/companies") | |
| 231 | +async def admin_companies(response: Response, p: PageDep, onboarding_status: str | None = None, status: str | None = None, | |
| 232 | + q_: str | None = Query(None, alias="q", max_length=200), country: str | None = None, | |
| 233 | + sort: str = Query("recent", pattern="^(activity|events|hiring|name|importance|recent)$")) -> dict[str, Any]: | |
| 234 | + _ns(response) | |
| 235 | + where, params = q.company_filters(q=q_, country=country, status=status, onboarding_status=onboarding_status) | |
| 236 | + async with connection() as conn: | |
| 237 | + ids, total = await q.company_page_ids(conn, where, params, sort=sort, limit=p.per_page, offset=p.offset) | |
| 238 | + rows = await q.fetch_cards_by_ids(conn, ids) | |
| 239 | + items = [] | |
| 240 | + for r in rows: | |
| 241 | + card = ser.company_card(r) | |
| 242 | + card.update({"onboarding_error": r.get("onboarding_error"), "indexed": bool(r.get("indexed")), "discovered_at": r.get("discovered_at"), "updated_at": r.get("updated_at")}) | |
| 243 | + items.append(card) | |
| 244 | + return page_payload(items, total, p) | |
| 245 | + | |
| 246 | + | |
| 247 | +class CompanyCreate(BaseModel): | |
| 248 | + website: str = Field(min_length=4, max_length=500) | |
| 249 | + display_name: str | None = Field(None, max_length=200) | |
| 250 | + country: str | None = Field(None, min_length=2, max_length=2) | |
| 251 | + industries: list[str] | None = None | |
| 252 | + importance: float | None = Field(None, ge=0, le=1) | |
| 253 | + tier: int | None = Field(None, ge=1, le=4) | |
| 254 | + | |
| 255 | + | |
| 256 | +@router.post("/companies", status_code=201) | |
| 257 | +async def admin_create_company(body: CompanyCreate, response: Response) -> dict[str, Any]: | |
| 258 | + _ns(response) | |
| 259 | + website = body.website.strip() | |
| 260 | + if "://" not in website: | |
| 261 | + website = "https://" + website | |
| 262 | + if not website.lower().startswith(("http://", "https://")): | |
| 263 | + raise HTTPException(status_code=422, detail="website: only http(s) URLs are accepted") | |
| 264 | + domain = registrable_domain(website) | |
| 265 | + if not domain or "." not in domain: | |
| 266 | + raise HTTPException(status_code=422, detail="website: could not derive a registrable domain") | |
| 267 | + website = canonicalize_url(website) | |
| 268 | + display_name = (body.display_name or "").strip() or domain.split(".")[0].capitalize() | |
| 269 | + country = body.country.upper() if body.country else None | |
| 270 | + industries = [slugify(i) for i in (body.industries or []) if i.strip()][:10] | |
| 271 | + async with transaction() as conn: | |
| 272 | + existing = await fetch_one(conn, "select id, slug from companies where canonical_domain = :d", d=domain) | |
| 273 | + if existing: | |
| 274 | + raise HTTPException(status_code=409, detail=f"company already exists: {existing['slug']}") | |
| 275 | + if country and not await fetch_val(conn, "select 1 from countries where code = cast(:c as char(2))", c=country): | |
| 276 | + raise HTTPException(status_code=422, detail="country: unknown ISO-2 code") | |
| 277 | + if industries: | |
| 278 | + known = {r["slug"] for r in await fetch_all(conn, "select slug from industries where slug = any(cast(:s as text[]))", s=industries)} | |
| 279 | + missing = [i for i in industries if i not in known] | |
| 280 | + if missing: | |
| 281 | + raise HTTPException(status_code=422, detail=f"industries: unknown slugs {', '.join(missing)}") | |
| 282 | + base = slugify(display_name) | |
| 283 | + slug, n = base, 2 | |
| 284 | + while await fetch_val(conn, "select 1 from companies where slug = :s", s=slug): | |
| 285 | + slug = f"{base}-{n}" | |
| 286 | + n += 1 | |
| 287 | + cid = new_id("company") | |
| 288 | + await execute(conn, "insert into companies (id, slug, display_name, canonical_domain, website, country, industries, industry_primary, importance, tier, " | |
| 289 | + "source_meta) values (:id, :slug, :name, :domain, :website, :country, cast(:inds as text[]), :ip, :imp, :tier, cast(:meta as jsonb))", | |
| 290 | + id=cid, slug=slug, name=display_name, domain=domain, website=website, country=country, inds=industries, ip=industries[0] if industries else None, | |
| 291 | + imp=body.importance if body.importance is not None else 0.2, tier=body.tier or 4, meta=jsonb({"source": "admin_api", "created_at": datetime.now(UTC)})) | |
| 292 | + await execute(conn, "insert into domains (id, company_id, domain, kind) values (:id, :cid, :d, 'primary') on conflict do nothing", id=new_id("domain"), cid=cid, d=domain) | |
| 293 | + jid = new_id("queue_job") | |
| 294 | + await execute(conn, "insert into queue_jobs (id, kind, key, payload, priority) values (:id, 'discover', :key, cast(:p as jsonb), 0.8) on conflict (key) do nothing", | |
| 295 | + id=jid, key=f"discover:{cid}", p=jsonb({"company_id": cid, "reason": "admin:create"})) | |
| 296 | + cards = await q.fetch_cards_by_ids(conn, [cid]) | |
| 297 | + return {"ok": True, "company": ser.company_card(cards[0]), "queued": {"id": jid, "kind": "discover"}} | |
| 298 | + | |
| 299 | + | |
| 300 | +@router.post("/companies/{key}/rediscover") | |
| 301 | +async def admin_rediscover(key: str, response: Response) -> dict[str, Any]: | |
| 302 | + _ns(response) | |
| 303 | + now = datetime.now(UTC) | |
| 304 | + async with transaction() as conn: | |
| 305 | + c = await q.require_company(conn, key) | |
| 306 | + jid = new_id("queue_job") | |
| 307 | + jkey = f"discover:{c['id']}:{int(now.timestamp())}" | |
| 308 | + await execute(conn, "insert into queue_jobs (id, kind, key, payload, priority) values (:id, 'discover', :key, cast(:p as jsonb), 0.9)", | |
| 309 | + id=jid, key=jkey, p=jsonb({"company_id": c["id"], "reason": "admin:rediscover"})) | |
| 310 | + if c["onboarding_status"] in ("failed", "no_website"): | |
| 311 | + await execute(conn, "update companies set onboarding_status = 'pending', onboarding_error = null, updated_at = :now where id = :id", now=now, id=c["id"]) | |
| 312 | + return {"ok": True, "company": c["slug"], "queued": {"id": jid, "key": jkey}} | |
| 313 | + | |
| 314 | + | |
| 315 | +# ------------------------------------------------------------------------------------------------ failures / queue / llm / reviews | |
| 316 | + | |
| 317 | + | |
| 318 | +@router.get("/failures") | |
| 319 | +async def admin_failures(response: Response, p: PageDep, class_: str | None = Query(None, alias="class", max_length=40), since: str | None = None, | |
| 320 | + sensor: str | None = None, company: str | None = None) -> dict[str, Any]: | |
| 321 | + _ns(response) | |
| 322 | + where, params = ["true"], {} | |
| 323 | + if class_: | |
| 324 | + where.append("f.failure_class = cast(:fc as text)") | |
| 325 | + params["fc"] = class_.upper() | |
| 326 | + since_dt = q.parse_iso(since) or q.days_ago(7) | |
| 327 | + where.append("f.at >= :since") | |
| 328 | + params["since"] = since_dt | |
| 329 | + if sensor: | |
| 330 | + where.append("f.sensor_id = :sid") | |
| 331 | + params["sid"] = sensor | |
| 332 | + async with connection() as conn: | |
| 333 | + if company: | |
| 334 | + params["cid"] = (await q.require_company(conn, company))["id"] | |
| 335 | + where.append("f.company_id = :cid") | |
| 336 | + wsql = " and ".join(where) | |
| 337 | + rows = await fetch_all(conn, f"select f.*, c.slug as company_slug, s.surface from failures f left join companies c on c.id = f.company_id " | |
| 338 | + f"left join sensors s on s.id = f.sensor_id where {wsql} order by f.at desc limit :limit offset :offset", **params, limit=p.per_page, offset=p.offset) | |
| 339 | + total = await q.bounded_count(conn, f"from failures f where {wsql}", params) | |
| 340 | + by_class = await fetch_all(conn, f"select f.failure_class, count(*) as n from failures f where {wsql} group by 1 order by n desc", **params) | |
| 341 | + out = page_payload([ser.failure(r) for r in rows], total, p) | |
| 342 | + out["by_class"] = {r["failure_class"]: int(r["n"]) for r in by_class} | |
| 343 | + out["classes"] = [c.value for c in FailureClass] | |
| 344 | + return out | |
| 345 | + | |
| 346 | + | |
| 347 | +@router.get("/queue") | |
| 348 | +async def admin_queue(response: Response, kind: str | None = None, status: str | None = None, limit: int = Query(100, ge=1, le=500)) -> dict[str, Any]: | |
| 349 | + _ns(response) | |
| 350 | + where, params = ["true"], {"limit": limit} | |
| 351 | + if kind: | |
| 352 | + where.append("kind = cast(:kind as text)") | |
| 353 | + params["kind"] = kind[:40] | |
| 354 | + if status: | |
| 355 | + where.append("status = cast(:status as text)") | |
| 356 | + params["status"] = status[:20] | |
| 357 | + async with connection() as conn: | |
| 358 | + counts = await fetch_all(conn, "select kind, status, count(*) as n from queue_jobs group by kind, status order by kind, status") | |
| 359 | + rows = await fetch_all(conn, f"select * from queue_jobs where {' and '.join(where)} order by case status when 'running' then 0 when 'pending' then 1 else 2 end, " | |
| 360 | + "run_at desc limit :limit", **params) | |
| 361 | + return {"counts": [{"kind": r["kind"], "status": r["status"], "n": int(r["n"])} for r in counts], "items": [ser.queue_job(r) for r in rows]} | |
| 362 | + | |
| 363 | + | |
| 364 | +class RequeueBody(BaseModel): | |
| 365 | + kind: str | None = Field(None, max_length=40) | |
| 366 | + | |
| 367 | + | |
| 368 | +@router.post("/queue/requeue-dead") | |
| 369 | +async def admin_requeue_dead(response: Response, body: RequeueBody | None = None) -> dict[str, Any]: | |
| 370 | + _ns(response) | |
| 371 | + body = body or RequeueBody() | |
| 372 | + extra, params = "", {} | |
| 373 | + if body.kind: | |
| 374 | + extra = " and kind = cast(:kind as text)" | |
| 375 | + params["kind"] = body.kind | |
| 376 | + async with transaction() as conn: | |
| 377 | + n = await fetch_val(conn, "with u as (update queue_jobs set status = 'pending', attempts = 0, run_at = now(), locked_at = null, locked_by = null, " | |
| 378 | + f"last_error = null, finished_at = null where status = 'dead'{extra} returning 1) select count(*) from u", **params) | |
| 379 | + return {"ok": True, "requeued": int(n or 0)} | |
| 380 | + | |
| 381 | + | |
| 382 | +@router.get("/llm") | |
| 383 | +async def admin_llm(response: Response, p: PageDep, status: str | None = None, kind: str | None = None) -> dict[str, Any]: | |
| 384 | + _ns(response) | |
| 385 | + where, params = ["true"], {} | |
| 386 | + if status: | |
| 387 | + where.append("status = cast(:status as text)") | |
| 388 | + params["status"] = status[:20] | |
| 389 | + if kind: | |
| 390 | + where.append("kind = cast(:kind as text)") | |
| 391 | + params["kind"] = kind[:40] | |
| 392 | + wsql = " and ".join(where) | |
| 393 | + async with connection() as conn: | |
| 394 | + rows = await fetch_all(conn, f"select * from llm_jobs where {wsql} order by created_at desc limit :limit offset :offset", **params, limit=p.per_page, offset=p.offset) | |
| 395 | + total = await q.bounded_count(conn, f"from llm_jobs where {wsql}", params) | |
| 396 | + stats = await fetch_one(conn, "select count(*) filter (where status = 'pending') as pending, count(*) filter (where status = 'running') as running, " | |
| 397 | + "count(*) filter (where status = 'done' and finished_at >= current_date) as done_today, " | |
| 398 | + "count(*) filter (where status = 'failed' and finished_at >= current_date) as failed_today, " | |
| 399 | + "avg(latency_ms) filter (where status = 'done' and finished_at >= current_date) as avg_latency_ms, " | |
| 400 | + "sum(request_tokens + coalesce(response_tokens, 0)) filter (where finished_at >= current_date) as tokens_today from llm_jobs") or {} | |
| 401 | + out = page_payload([ser.llm_job(r) for r in rows], total, p) | |
| 402 | + out["stats"] = {k: (round(float(v), 1) if k == "avg_latency_ms" and v is not None else (int(v) if v is not None else 0)) for k, v in stats.items()} | |
| 403 | + out["stats"]["budget"] = settings.llm_daily_budget | |
| 404 | + out["stats"]["configured"] = settings.llm_configured | |
| 405 | + return out | |
| 406 | + | |
| 407 | + | |
| 408 | +@router.get("/reviews") | |
| 409 | +async def admin_reviews(response: Response, p: PageDep, kind: str | None = None, status: str = Query("open", max_length=20)) -> dict[str, Any]: | |
| 410 | + _ns(response) | |
| 411 | + where, params = ["true"], {} | |
| 412 | + if kind: | |
| 413 | + where.append("r.kind = cast(:kind as text)") | |
| 414 | + params["kind"] = kind[:40] | |
| 415 | + if status and status != "all": | |
| 416 | + where.append("r.status = cast(:status as text)") | |
| 417 | + params["status"] = status | |
| 418 | + wsql = " and ".join(where) | |
| 419 | + async with connection() as conn: | |
| 420 | + rows = await fetch_all(conn, f"select r.*, c.slug as company_slug, c.display_name as company_display_name from review_queue r left join companies c on c.id = r.company_id " | |
| 421 | + f"where {wsql} order by r.created_at asc limit :limit offset :offset", **params, limit=p.per_page, offset=p.offset) | |
| 422 | + total = await q.bounded_count(conn, f"from review_queue r where {wsql}", params) | |
| 423 | + by_kind = await fetch_all(conn, "select kind, count(*) as n from review_queue where status = 'open' group by kind order by n desc") | |
| 424 | + out = page_payload([ser.review(r) for r in rows], total, p) | |
| 425 | + out["open_by_kind"] = {r["kind"]: int(r["n"]) for r in by_kind} | |
| 426 | + return out | |
| 427 | + | |
| 428 | + | |
| 429 | +class ReviewBody(BaseModel): | |
| 430 | + resolution: str = Field(pattern="^(accepted|rejected)$") | |
| 431 | + note: str | None = Field(None, max_length=1000) | |
| 432 | + label: str | None = Field(None, pattern="^(correct|duplicate|noise|misclassified)$") | |
| 433 | + | |
| 434 | + | |
| 435 | +@router.post("/reviews/{review_id}") | |
| 436 | +async def admin_resolve_review(review_id: str, body: ReviewBody, response: Response) -> dict[str, Any]: | |
| 437 | + _ns(response) | |
| 438 | + async with transaction() as conn: | |
| 439 | + r = await fetch_one(conn, "select * from review_queue where id = :id", id=review_id) | |
| 440 | + if r is None: | |
| 441 | + raise HTTPException(status_code=404, detail="review not found") | |
| 442 | + if r["status"] != "open": | |
| 443 | + raise HTTPException(status_code=409, detail=f"review already {r['status']}") | |
| 444 | + payload = ser._dict(r["payload"]) | |
| 445 | + payload["resolution"] = {"status": body.resolution, "note": body.note, "label": body.label, "at": datetime.now(UTC)} | |
| 446 | + await execute(conn, "update review_queue set status = :st, resolution = :res, resolved_at = now(), payload = cast(:p as jsonb) where id = :id", | |
| 447 | + st=body.resolution, res=body.label or body.note or body.resolution, p=jsonb(payload), id=review_id) | |
| 448 | + row = await fetch_one(conn, "select * from review_queue where id = :id", id=review_id) | |
| 449 | + return {"ok": True, "review": ser.review(row or {})} | |
| 450 | + | |
| 451 | + | |
| 452 | +# ------------------------------------------------------------------------------------------------ events (corrections) | |
| 453 | + | |
| 454 | + | |
| 455 | +class RetractBody(BaseModel): | |
| 456 | + reason: str = Field(min_length=3, max_length=500) | |
| 457 | + | |
| 458 | + | |
| 459 | +async def _audit(conn: Any, event_id: str, action: str, reason: str | None) -> None: | |
| 460 | + entry = jsonb([{"action": action, "reason": reason, "at": datetime.now(UTC)}]) | |
| 461 | + await execute(conn, "update events set payload = jsonb_set(payload, '{_audit}', coalesce(payload->'_audit', '[]'::jsonb) || cast(:e as jsonb), true) where id = :id", | |
| 462 | + e=entry, id=event_id) | |
| 463 | + | |
| 464 | + | |
| 465 | +@router.post("/events/{event_id}/retract") | |
| 466 | +async def admin_retract_event(event_id: str, body: RetractBody, response: Response) -> dict[str, Any]: | |
| 467 | + _ns(response) | |
| 468 | + async with transaction() as conn: | |
| 469 | + ev = await fetch_one(conn, "select id, status from events where id = :id", id=event_id) | |
| 470 | + if ev is None: | |
| 471 | + raise HTTPException(status_code=404, detail="event not found") | |
| 472 | + await execute(conn, "update events set status = 'retracted', retracted_reason = :r where id = :id", r=body.reason.strip(), id=event_id) | |
| 473 | + await _audit(conn, event_id, "retract", body.reason.strip()) | |
| 474 | + row = await q.fetch_event(conn, event_id) | |
| 475 | + cache.clear() | |
| 476 | + return {"ok": True, "event": ser.event(row or {})} | |
| 477 | + | |
| 478 | + | |
| 479 | +@router.post("/events/{event_id}/restore") | |
| 480 | +async def admin_restore_event(event_id: str, response: Response) -> dict[str, Any]: | |
| 481 | + _ns(response) | |
| 482 | + async with transaction() as conn: | |
| 483 | + ev = await fetch_one(conn, "select id, status, retracted_reason from events where id = :id", id=event_id) | |
| 484 | + if ev is None: | |
| 485 | + raise HTTPException(status_code=404, detail="event not found") | |
| 486 | + await execute(conn, "update events set status = 'active', retracted_reason = null where id = :id", id=event_id) | |
| 487 | + await _audit(conn, event_id, "restore", ev.get("retracted_reason")) | |
| 488 | + row = await q.fetch_event(conn, event_id) | |
| 489 | + cache.clear() | |
| 490 | + return {"ok": True, "event": ser.event(row or {})} | |
| 491 | + | |
| 492 | + | |
| 493 | +# ------------------------------------------------------------------------------------------------ quality / costs / cache | |
| 494 | + | |
| 495 | + | |
| 496 | +@router.get("/quality") | |
| 497 | +async def admin_quality(response: Response) -> dict[str, Any]: | |
| 498 | + _ns(response) | |
| 499 | + async with connection() as conn: | |
| 500 | + r = await fetch_one(conn, """ | |
| 501 | + select (select count(*) from companies) as companies, | |
| 502 | + (select count(*) from companies where onboarding_status = 'active' and status = 'ACTIVE') as companies_active, | |
| 503 | + (select count(*) from sensors where status <> 'retired') as sensors, | |
| 504 | + (select count(*) from sensors where status = 'active') as sensors_active, | |
| 505 | + (select count(*) from sensors where status <> 'retired' and last_run_at >= now() - interval '24 hours') as checked_24h, | |
| 506 | + (select count(*) from sensors where status = 'stale') as stale, | |
| 507 | + (select count(*) from sensors where status in ('failing', 'blocked')) as failed_sensors, | |
| 508 | + (select count(*) from sensors where surface = 'other' and status <> 'retired') as unknown_surfaces, | |
| 509 | + (select count(*) from events where detected_at >= now() - interval '30 days') as events_30d, | |
| 510 | + (select count(*) from events where detected_at >= now() - interval '30 days' and status = 'duplicate') as duplicates_30d, | |
| 511 | + (select avg(confidence) from events where detected_at >= now() - interval '30 days' and status = 'active') as confidence_avg | |
| 512 | + """) or {} | |
| 513 | + cal = await fetch_all(conn, "select coalesce(payload->'resolution'->>'label', resolution) as label, count(*) as n from review_queue " | |
| 514 | + "where status in ('accepted','rejected','resolved') group by 1") | |
| 515 | + companies, sensors = int(r.get("companies") or 0), int(r.get("sensors") or 0) | |
| 516 | + ev30 = int(r.get("events_30d") or 0) | |
| 517 | + labels = {c["label"]: int(c["n"]) for c in cal if c["label"]} | |
| 518 | + return {"coverage": {"companies_active_pct": round(int(r.get("companies_active") or 0) / companies * 100, 1) if companies else None, | |
| 519 | + "sensors_active_pct": round(int(r.get("sensors_active") or 0) / sensors * 100, 1) if sensors else None, | |
| 520 | + "companies": companies, "sensors": sensors}, | |
| 521 | + "freshness": {"sensors_checked_24h_pct": round(int(r.get("checked_24h") or 0) / sensors * 100, 1) if sensors else None, "stale": int(r.get("stale") or 0)}, | |
| 522 | + "duplicate_rate": round(int(r.get("duplicates_30d") or 0) / ev30, 4) if ev30 else None, | |
| 523 | + "event_confidence_avg": round(float(r["confidence_avg"]), 3) if r.get("confidence_avg") is not None else None, | |
| 524 | + "unknown_surfaces": int(r.get("unknown_surfaces") or 0), "failed_sensors": int(r.get("failed_sensors") or 0), | |
| 525 | + "calibration": {k: labels.get(k, 0) for k in ("correct", "duplicate", "noise", "misclassified")}, "events_30d": ev30} | |
| 526 | + | |
| 527 | + | |
| 528 | +@router.get("/costs") | |
| 529 | +async def admin_costs(response: Response, days: int = Query(30, ge=1, le=365)) -> dict[str, Any]: | |
| 530 | + _ns(response) | |
| 531 | + async with connection() as conn: | |
| 532 | + items = await fetch_all(conn, "select day, dimension, key, units, cost_estimate from cost_ledger where day >= :d order by day desc, dimension, key limit 5000", | |
| 533 | + d=q.days_ago(days).date()) | |
| 534 | + denom = await fetch_one(conn, "select (select count(*) from companies where status = 'ACTIVE') as companies, " | |
| 535 | + "(select count(*) from observations where fetched_at >= :d) as observations, " | |
| 536 | + "(select count(*) from events where status = 'active' and detected_at >= :d and importance >= :imp) as meaningful_events", | |
| 537 | + d=q.days_ago(days), imp=settings.meaningful_threshold) or {} | |
| 538 | + total = sum(float(i["cost_estimate"] or 0) for i in items) | |
| 539 | + by_dim: dict[str, float] = {} | |
| 540 | + for i in items: | |
| 541 | + by_dim[i["dimension"]] = by_dim.get(i["dimension"], 0.0) + float(i["cost_estimate"] or 0) | |
| 542 | + comp, obs, mev = int(denom.get("companies") or 0), int(denom.get("observations") or 0), int(denom.get("meaningful_events") or 0) | |
| 543 | + return {"days": days, "items": [{"day": i["day"], "dimension": i["dimension"], "key": i["key"], "units": float(i["units"] or 0), "cost_estimate": float(i["cost_estimate"] or 0)} | |
| 544 | + for i in items], | |
| 545 | + "total": round(total, 4), "by_dimension": {k: round(v, 4) for k, v in by_dim.items()}, | |
| 546 | + "per_1000_companies": round(total / comp * 1000, 4) if comp else None, "per_million_observations": round(total / obs * 1_000_000, 4) if obs else None, | |
| 547 | + "per_meaningful_event": round(total / mev, 4) if mev else None} | |
| 548 | + | |
| 549 | + | |
| 550 | +@router.post("/cache/clear") | |
| 551 | +async def admin_cache_clear(response: Response, prefix: str | None = Query(None, max_length=60)) -> dict[str, Any]: | |
| 552 | + _ns(response) | |
| 553 | + cache.clear(prefix) | |
| 554 | + return {"ok": True, "cleared": prefix or "all"} | |
| 555 | + | |
| 556 | + | |
| 557 | +@router.get("/storage") | |
| 558 | +async def admin_storage(response: Response) -> dict[str, Any]: | |
| 559 | + _ns(response) | |
| 560 | + stats = await agg.archive_stats() | |
| 561 | + return {"objects_dir": str(settings.objects_dir), "exists": archive.object_path("00" * 32).parent.parent.parent.exists(), **stats} | |
added
src/companyatlas/api/routers/companies.py
+354 −0
@@ -0,0 +1,354 @@ | ||
| 1 | +"""Companies: list, compare (literal path registered before the `/{slug_or_id}` catch-all), detail and every sub-resource.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from collections import defaultdict | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from fastapi import APIRouter, HTTPException, Query, Response | |
| 8 | + | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api import serializers as ser | |
| 11 | +from companyatlas.api.common import PageDep, page_payload, public_cache_value | |
| 12 | +from companyatlas.db import connection, fetch_all, fetch_one | |
| 13 | +from companyatlas.taxonomy import Metric | |
| 14 | + | |
| 15 | +ORDER = 40 | |
| 16 | +router = APIRouter(prefix="/api/v1", tags=["companies"]) | |
| 17 | + | |
| 18 | +TIMELINE_FILTERS: dict[str, list[str]] = { | |
| 19 | + "all": [], "products": ["PRODUCT"], "jobs": ["HIRING"], "pricing": ["PRICING"], "leadership": ["LEADERSHIP"], "locations": ["LOCATION"], | |
| 20 | + "legal": ["LEGAL"], "news": ["COMMUNICATION", "INVESTOR_RELATIONS", "MARKETING"], "developer": ["DEVELOPER", "TECHNOLOGY"], | |
| 21 | + "corporate": ["FINANCING", "M&A", "PARTNERSHIP", "STRATEGY"], | |
| 22 | +} | |
| 23 | +COMPARE_METRICS = [Metric.ACTIVITY_SCORE.value, Metric.HIRING_MOMENTUM_30D.value, Metric.PRODUCT_VELOCITY.value, Metric.AI_ADOPTION.value, | |
| 24 | + Metric.GEO_EXPANSION.value, Metric.DEVELOPER_MOMENTUM.value, Metric.CORPORATE_CHANGE_INDEX.value, Metric.OPEN_JOBS.value, | |
| 25 | + Metric.ANOMALY_SCORE.value] | |
| 26 | + | |
| 27 | + | |
| 28 | +def _pub(response: Response, s: int = 60) -> None: | |
| 29 | + response.headers["cache-control"] = public_cache_value(s) | |
| 30 | + | |
| 31 | + | |
| 32 | +# ------------------------------------------------------------------------------------------------ list / compare | |
| 33 | + | |
| 34 | + | |
| 35 | +@router.get("/companies", summary="Company directory") | |
| 36 | +async def list_companies(response: Response, p: PageDep, q_: str | None = Query(None, alias="q", max_length=200), | |
| 37 | + country: str | None = None, industry: str | None = None, tier: int | None = Query(None, ge=1, le=4), | |
| 38 | + public: str | None = None, status: str | None = None, has_events: str | None = None, | |
| 39 | + sort: str = Query("activity", pattern="^(activity|events|hiring|name|importance|recent|relevance)$"), | |
| 40 | + sparkline: str | None = None) -> dict[str, Any]: | |
| 41 | + _pub(response, 60) | |
| 42 | + where, params = q.company_filters(q=q_, country=country, industry=industry, tier=tier, public=q.parse_bool(public), status=status, | |
| 43 | + has_events=q.parse_bool(has_events)) | |
| 44 | + async with connection() as conn: | |
| 45 | + ids, total = await q.company_page_ids(conn, where, params, sort=sort, limit=p.per_page, offset=p.offset) | |
| 46 | + rows = await q.fetch_cards_by_ids(conn, ids, sparkline=bool(q.parse_bool(sparkline))) | |
| 47 | + return page_payload([ser.company_card(r) for r in rows], total, p) | |
| 48 | + | |
| 49 | + | |
| 50 | +@router.get("/companies/compare", summary="Side-by-side comparison of 2–6 companies") | |
| 51 | +async def compare(response: Response, companies: str = Query(..., description="comma-separated slugs or ids"), days: int = Query(90, ge=7, le=365)) -> dict[str, Any]: | |
| 52 | + _pub(response, 120) | |
| 53 | + keys = list(dict.fromkeys(q.csv_list(companies, maxlen=6))) | |
| 54 | + if len(keys) < 2: | |
| 55 | + raise HTTPException(status_code=422, detail="companies: provide 2 to 6 comma-separated slugs") | |
| 56 | + async with connection() as conn: | |
| 57 | + found = [] | |
| 58 | + for k in keys: | |
| 59 | + c = await q.get_company(conn, k) | |
| 60 | + if c is None: | |
| 61 | + raise HTTPException(status_code=404, detail=f"company not found: {k}") | |
| 62 | + found.append(c) | |
| 63 | + ids = [c["id"] for c in found] | |
| 64 | + slug_of = {c["id"]: c["slug"] for c in found} | |
| 65 | + cards = [ser.company_card(r) for r in await q.fetch_cards_by_ids(conn, ids, sparkline=True)] | |
| 66 | + series_rows = await fetch_all(conn, "select company_id, day, value, confidence from metric_series where metric = 'activity_score' and " | |
| 67 | + "company_id = any(cast(:ids as text[])) and day >= :d order by day", ids=ids, d=q.days_ago(days).date()) | |
| 68 | + ev_rows = await fetch_all(conn, "select company_id, event_type, count(*) as n from events where status = 'active' and detected_at >= :d and " | |
| 69 | + "company_id = any(cast(:ids as text[])) group by 1, 2", ids=ids, d=q.days_ago(30)) | |
| 70 | + job_rows = await fetch_all(conn, "select company_id, count(*) filter (where status = 'open') as open, " | |
| 71 | + "count(*) filter (where status = 'open' and is_ai) as ai_open, " | |
| 72 | + "count(*) filter (where first_seen_at >= :d) as new_30d from jobs where company_id = any(cast(:ids as text[])) " | |
| 73 | + "group by company_id", ids=ids, d=q.days_ago(30)) | |
| 74 | + loc_rows = await fetch_all(conn, "select company_id, count(*) as n from locations where status = 'listed' and company_id = any(cast(:ids as text[])) " | |
| 75 | + "group by company_id", ids=ids) | |
| 76 | + metrics: dict[str, dict[str, Any]] = {m: {} for m in COMPARE_METRICS} | |
| 77 | + for card in cards: | |
| 78 | + for m, v in card["metrics"].items(): | |
| 79 | + metrics.setdefault(m, {})[card["slug"]] = v | |
| 80 | + series: dict[str, list[dict[str, Any]]] = {s: [] for s in slug_of.values()} | |
| 81 | + for r in series_rows: | |
| 82 | + series[slug_of[r["company_id"]]].append(ser.metric_point(r)) | |
| 83 | + events_30d: dict[str, dict[str, int]] = {s: {} for s in slug_of.values()} | |
| 84 | + for r in ev_rows: | |
| 85 | + events_30d[slug_of[r["company_id"]]][r["event_type"]] = int(r["n"]) | |
| 86 | + jobs = {slug_of[r["company_id"]]: {"open": int(r["open"]), "ai_open": int(r["ai_open"]), "new_30d": int(r["new_30d"])} for r in job_rows} | |
| 87 | + for s in slug_of.values(): | |
| 88 | + jobs.setdefault(s, {"open": 0, "ai_open": 0, "new_30d": 0}) | |
| 89 | + locations = {slug_of[r["company_id"]]: int(r["n"]) for r in loc_rows} | |
| 90 | + for s in slug_of.values(): | |
| 91 | + locations.setdefault(s, 0) | |
| 92 | + return {"companies": cards, "metrics": metrics, "series": series, "events_30d": events_30d, "jobs": jobs, "locations": locations} | |
| 93 | + | |
| 94 | + | |
| 95 | +# ------------------------------------------------------------------------------------------------ detail | |
| 96 | + | |
| 97 | + | |
| 98 | +@router.get("/companies/{key}", summary="Company profile") | |
| 99 | +async def company_detail(key: str, response: Response) -> dict[str, Any]: | |
| 100 | + _pub(response, 60) | |
| 101 | + async with connection() as conn: | |
| 102 | + c = await q.require_company(conn, key) | |
| 103 | + cid = c["id"] | |
| 104 | + rows = await q.fetch_cards_by_ids(conn, [cid], sparkline=True) | |
| 105 | + out = ser.company_card(rows[0]) | |
| 106 | + out["company_type"] = c.get("company_type") | |
| 107 | + out["employees"] = c.get("employees") | |
| 108 | + out["wikidata_id"] = c.get("wikidata_id") | |
| 109 | + out["indexed"] = bool(c.get("indexed")) | |
| 110 | + out["discovered_at"] = c.get("discovered_at") | |
| 111 | + out["first_observed_at"] = c.get("first_observed_at") | |
| 112 | + out["last_change_at"] = c.get("last_change_at") | |
| 113 | + out["aliases"] = [r["alias"] for r in await fetch_all(conn, "select alias from company_aliases where company_id = :id order by kind, alias limit 50", id=cid)] | |
| 114 | + out["domains"] = await fetch_all(conn, "select domain, kind, status, first_seen_at, last_seen_at from domains where company_id = :id order by kind, domain limit 100", id=cid) | |
| 115 | + rel = await fetch_all(conn, "select r.kind, r.to_name, r.valid_from, r.valid_to, r.confidence, r.source_url, o.slug, o.display_name " | |
| 116 | + "from company_relationships r left join companies o on o.id = r.to_company_id where r.from_company_id = :id " | |
| 117 | + "order by r.kind, r.last_seen_at desc limit 100", id=cid) | |
| 118 | + out["relationships"] = [{"kind": r["kind"], "company": {"slug": r["slug"], "display_name": r["display_name"]} if r["slug"] else None, | |
| 119 | + "to_name": r["to_name"] or r["display_name"], "valid_from": r["valid_from"], "valid_to": r["valid_to"], | |
| 120 | + "confidence": ser._float(r["confidence"], 3), "source_url": r["source_url"]} for r in rel] | |
| 121 | + md = await fetch_all(conn, "select metric, value, confidence, computed_at, formula_version, inputs from metrics_current where company_id = :id order by metric", id=cid) | |
| 122 | + out["metrics_detail"] = [{"metric": r["metric"], "value": ser.metric_value(r["metric"], r["value"]), "confidence": ser._float(r["confidence"], 3), | |
| 123 | + "computed_at": r["computed_at"], "formula_version": r["formula_version"], "inputs": ser._dict(r["inputs"])} for r in md] | |
| 124 | + sbs = await fetch_all(conn, "select surface, count(*) as n, count(*) filter (where status = 'active') as active from sensors where company_id = :id " | |
| 125 | + "and status <> 'retired' group by surface order by n desc", id=cid) | |
| 126 | + out["sensors_by_surface"] = {r["surface"]: int(r["n"]) for r in sbs} | |
| 127 | + total_sensors = sum(int(r["n"]) for r in sbs) | |
| 128 | + active_sensors = sum(int(r["active"]) for r in sbs) | |
| 129 | + days_observed = await fetch_one(conn, "select count(*) as n from company_daily where company_id = :id and observations > 0", id=cid) | |
| 130 | + hist = next((m["value"] for m in out["metrics_detail"] if m["metric"] == Metric.HISTORICAL_COVERAGE.value), None) | |
| 131 | + out["coverage"] = {"historical_coverage": hist, "first_observed_at": c.get("first_observed_at"), | |
| 132 | + "days_observed": int(days_observed["n"]) if days_observed else 0, | |
| 133 | + "sensor_uptime": round(active_sensors / total_sensors, 3) if total_sensors else None} | |
| 134 | + sig = await fetch_all(conn, "select * from signals where company_id = :id and status = 'active' order by detected_at desc limit 10", id=cid) | |
| 135 | + out["signals"] = [ser.signal(r) for r in sig] | |
| 136 | + act = await q.metric_series(conn, cid, Metric.ACTIVITY_SCORE.value, 30) | |
| 137 | + hir = await q.metric_series(conn, cid, Metric.OPEN_JOBS.value, 90) | |
| 138 | + out["sparklines"] = {"activity_30d": [ser._float(r["value"], 1) for r in act], "hiring_90d": [ser._float(r["value"], 1) for r in hir]} | |
| 139 | + out["recent_events"] = [ser.event(r) for r in await q.fetch_events(conn, *q.event_filters(company_id=cid), limit=10)] | |
| 140 | + return out | |
| 141 | + | |
| 142 | + | |
| 143 | +# ------------------------------------------------------------------------------------------------ sub-resources | |
| 144 | + | |
| 145 | + | |
| 146 | +@router.get("/companies/{key}/events", summary="Company events") | |
| 147 | +async def company_events(key: str, response: Response, p: PageDep, event_type: str | None = None, event_subtype: str | None = None, | |
| 148 | + since: str | None = None, until: str | None = None, min_importance: float | None = Query(None, ge=0, le=1), | |
| 149 | + surface: str | None = None, sort: str = Query("recent", pattern="^(recent|importance)$")) -> dict[str, Any]: | |
| 150 | + _pub(response, 30) | |
| 151 | + async with connection() as conn: | |
| 152 | + c = await q.require_company(conn, key) | |
| 153 | + where, params = q.event_filters(company_id=c["id"], event_type=event_type, event_subtype=event_subtype, since=q.parse_iso(since), | |
| 154 | + until=q.parse_iso(until, "until"), min_importance=min_importance, surface=surface) | |
| 155 | + rows = await q.fetch_events(conn, where, params, sort=sort, limit=p.per_page, offset=p.offset) | |
| 156 | + total = await q.count_events(conn, where, params) | |
| 157 | + return page_payload([ser.event(r) for r in rows], total, p) | |
| 158 | + | |
| 159 | + | |
| 160 | +@router.get("/companies/{key}/timeline", summary="Timeline grouped by day") | |
| 161 | +async def company_timeline(key: str, response: Response, filter: str = Query("all", pattern="^(all|products|jobs|pricing|leadership|locations|legal|news|developer|corporate)$"), | |
| 162 | + limit: int = Query(200, ge=1, le=500), before: str | None = None) -> dict[str, Any]: | |
| 163 | + _pub(response, 60) | |
| 164 | + types = TIMELINE_FILTERS[filter] | |
| 165 | + async with connection() as conn: | |
| 166 | + c = await q.require_company(conn, key) | |
| 167 | + where, params = q.event_filters(company_id=c["id"], event_types=types or None, until=q.parse_iso(before, "before")) | |
| 168 | + rows = await q.fetch_events(conn, where, params, sort="recent", limit=limit) | |
| 169 | + where_sql = " where " + " and ".join(where) | |
| 170 | + days = await fetch_all(conn, f"select date(e.detected_at at time zone 'UTC') as day, count(*) as count from events e join companies c on c.id = e.company_id" | |
| 171 | + f"{where_sql} group by 1 order by 1 desc limit 730", **params) | |
| 172 | + items = [] | |
| 173 | + for r in rows: | |
| 174 | + ev = ser.event(r) | |
| 175 | + ev["day"] = r["detected_at"].date() | |
| 176 | + items.append(ev) | |
| 177 | + return {"items": items, "days": [{"day": d["day"], "count": int(d["count"])} for d in days], "filter": filter} | |
| 178 | + | |
| 179 | + | |
| 180 | +@router.get("/companies/{key}/metrics", summary="Current metrics and time series") | |
| 181 | +async def company_metrics(key: str, response: Response, metric: str | None = None, days: int = Query(90, ge=1, le=730)) -> dict[str, Any]: | |
| 182 | + _pub(response, 300) | |
| 183 | + async with connection() as conn: | |
| 184 | + c = await q.require_company(conn, key) | |
| 185 | + cur = await fetch_all(conn, "select metric, value, confidence, computed_at, formula_version, inputs from metrics_current where company_id = :id order by metric", id=c["id"]) | |
| 186 | + params: dict[str, Any] = {"id": c["id"], "d": q.days_ago(days).date()} | |
| 187 | + extra = "" | |
| 188 | + if metric: | |
| 189 | + extra = " and metric = cast(:m as text)" | |
| 190 | + params["m"] = metric[:60] | |
| 191 | + rows = await fetch_all(conn, f"select metric, day, value, confidence from metric_series where company_id = :id and day >= :d{extra} " | |
| 192 | + "order by metric, day limit 20000", **params) | |
| 193 | + series: dict[str, list[dict[str, Any]]] = defaultdict(list) | |
| 194 | + for r in rows: | |
| 195 | + series[r["metric"]].append(ser.metric_point(r)) | |
| 196 | + return {"current": [{"metric": r["metric"], "value": ser.metric_value(r["metric"], r["value"]), "confidence": ser._float(r["confidence"], 3), | |
| 197 | + "computed_at": r["computed_at"], "formula_version": r["formula_version"], "inputs": ser._dict(r["inputs"])} for r in cur], | |
| 198 | + "series": dict(series), "days": days} | |
| 199 | + | |
| 200 | + | |
| 201 | +@router.get("/companies/{key}/jobs", summary="Publicly listed jobs with summary") | |
| 202 | +async def company_jobs(key: str, response: Response, p: PageDep, status: str = Query("open", pattern="^(open|removed|all)$"), | |
| 203 | + q_: str | None = Query(None, alias="q", max_length=120), country: str | None = None, ai: str | None = None, | |
| 204 | + department: str | None = None, remote: str | None = None, | |
| 205 | + sort: str = Query("recent", pattern="^(recent|title|posted)$")) -> dict[str, Any]: | |
| 206 | + _pub(response, 120) | |
| 207 | + order = {"recent": "j.first_seen_at desc, j.id", "title": "j.title asc, j.id", "posted": "j.posted_at desc nulls last, j.id"}[sort] | |
| 208 | + where = ["j.company_id = :id"] | |
| 209 | + async with connection() as conn: | |
| 210 | + c = await q.require_company(conn, key) | |
| 211 | + params: dict[str, Any] = {"id": c["id"]} | |
| 212 | + if status == "open": | |
| 213 | + where.append("j.status = 'open'") | |
| 214 | + elif status == "removed": | |
| 215 | + where.append("j.status <> 'open'") | |
| 216 | + if q_: | |
| 217 | + where.append("j.title ilike :jq") | |
| 218 | + params["jq"] = f"%{q_.strip()}%" | |
| 219 | + if country: | |
| 220 | + where.append("j.country = cast(:jc as char(2))") | |
| 221 | + params["jc"] = country.upper()[:2] | |
| 222 | + if q.parse_bool(ai): | |
| 223 | + where.append("j.is_ai") | |
| 224 | + if department: | |
| 225 | + where.append("j.department = cast(:jd as text)") | |
| 226 | + params["jd"] = department[:80] | |
| 227 | + if q.parse_bool(remote) is not None: | |
| 228 | + where.append("j.remote = :jr") | |
| 229 | + params["jr"] = q.parse_bool(remote) | |
| 230 | + wsql = " and ".join(where) | |
| 231 | + rows = await fetch_all(conn, f"select j.* from jobs j where {wsql} order by {order} limit :limit offset :offset", **params, limit=p.per_page, offset=p.offset) | |
| 232 | + total = await q.bounded_count(conn, f"from jobs j where {wsql}", params) | |
| 233 | + d7 = q.days_ago(7) | |
| 234 | + s = await fetch_one(conn, "select count(*) filter (where status = 'open') as open, count(*) filter (where first_seen_at >= :d7 and status = 'open') as new_7d, " | |
| 235 | + "count(*) filter (where removed_at >= :d7) as removed_7d, count(*) filter (where status = 'open' and is_ai) as ai_open, " | |
| 236 | + "count(*) filter (where status = 'open' and remote is true) as remote_open, " | |
| 237 | + "count(*) filter (where status = 'open' and remote is not null) as remote_known from jobs where company_id = :id", | |
| 238 | + id=c["id"], d7=d7) or {} | |
| 239 | + by_country = await fetch_all(conn, "select country, count(*) as n from jobs where company_id = :id and status = 'open' and country is not null " | |
| 240 | + "group by country order by n desc limit 15", id=c["id"]) | |
| 241 | + by_dept = await fetch_all(conn, "select department, count(*) as n from jobs where company_id = :id and status = 'open' and department is not null " | |
| 242 | + "group by department order by n desc limit 15", id=c["id"]) | |
| 243 | + remote_known = int(s.get("remote_known") or 0) | |
| 244 | + summary = {"open": int(s.get("open") or 0), "new_7d": int(s.get("new_7d") or 0), "removed_7d": int(s.get("removed_7d") or 0), | |
| 245 | + "ai_open": int(s.get("ai_open") or 0), "by_country": [{"country": r["country"], "n": int(r["n"])} for r in by_country], | |
| 246 | + "by_department": [{"department": r["department"], "n": int(r["n"])} for r in by_dept], | |
| 247 | + "remote_ratio": round(int(s.get("remote_open") or 0) / remote_known, 3) if remote_known else None} | |
| 248 | + payload = page_payload([ser.job(r) for r in rows], total, p) | |
| 249 | + payload["meta"] = {"summary": summary} | |
| 250 | + return payload | |
| 251 | + | |
| 252 | + | |
| 253 | +@router.get("/companies/{key}/people", summary="Leadership listed on monitored pages") | |
| 254 | +async def company_people(key: str, response: Response) -> dict[str, Any]: | |
| 255 | + _pub(response, 300) | |
| 256 | + async with connection() as conn: | |
| 257 | + c = await q.require_company(conn, key) | |
| 258 | + rows = await fetch_all(conn, "select * from people where company_id = :id order by is_executive desc, status, last_seen_at desc limit 500", id=c["id"]) | |
| 259 | + return {"listed": [ser.person(r) for r in rows if r["status"] == "listed"], | |
| 260 | + "no_longer_listed": [ser.person(r) for r in rows if r["status"] != "listed"]} | |
| 261 | + | |
| 262 | + | |
| 263 | +@router.get("/companies/{key}/products", summary="Products in the public catalog") | |
| 264 | +async def company_products(key: str, response: Response) -> dict[str, Any]: | |
| 265 | + _pub(response, 300) | |
| 266 | + async with connection() as conn: | |
| 267 | + c = await q.require_company(conn, key) | |
| 268 | + rows = await fetch_all(conn, "select * from products where company_id = :id order by status, last_seen_at desc limit 500", id=c["id"]) | |
| 269 | + return {"listed": [ser.product(r) for r in rows if r["status"] == "listed"], "removed": [ser.product(r) for r in rows if r["status"] != "listed"]} | |
| 270 | + | |
| 271 | + | |
| 272 | +@router.get("/companies/{key}/pricing", summary="Current plans and every previous version") | |
| 273 | +async def company_pricing(key: str, response: Response, history_limit: int = Query(200, ge=1, le=1000)) -> dict[str, Any]: | |
| 274 | + _pub(response, 300) | |
| 275 | + async with connection() as conn: | |
| 276 | + c = await q.require_company(conn, key) | |
| 277 | + cur = await fetch_all(conn, "select * from pricing_plans where company_id = :id and status = 'current' order by price nulls last, plan_name limit 100", id=c["id"]) | |
| 278 | + hist = await fetch_all(conn, "select * from pricing_plans where company_id = :id and status <> 'current' order by valid_from desc limit :lim", | |
| 279 | + id=c["id"], lim=history_limit) | |
| 280 | + return {"current": [ser.plan(r) for r in cur], "history": [ser.plan(r) for r in hist]} | |
| 281 | + | |
| 282 | + | |
| 283 | +@router.get("/companies/{key}/locations", summary="Locations listed on monitored pages") | |
| 284 | +async def company_locations(key: str, response: Response, status: str = Query("listed", pattern="^(listed|all)$")) -> dict[str, Any]: | |
| 285 | + _pub(response, 300) | |
| 286 | + async with connection() as conn: | |
| 287 | + c = await q.require_company(conn, key) | |
| 288 | + extra = " and status = 'listed'" if status == "listed" else "" | |
| 289 | + rows = await fetch_all(conn, f"select * from locations where company_id = :id{extra} order by kind, country, city limit 1000", id=c["id"]) | |
| 290 | + items = [ser.location(r) for r in rows] | |
| 291 | + return {"items": items, "countries": sorted({r["country"] for r in rows if r["country"] and r["status"] == "listed"})} | |
| 292 | + | |
| 293 | + | |
| 294 | +@router.get("/companies/{key}/news", summary="First-party news, blog and changelog items") | |
| 295 | +async def company_news(key: str, response: Response, limit: int = Query(50, ge=1, le=500), category: str | None = None) -> dict[str, Any]: | |
| 296 | + _pub(response, 120) | |
| 297 | + async with connection() as conn: | |
| 298 | + c = await q.require_company(conn, key) | |
| 299 | + params: dict[str, Any] = {"id": c["id"], "lim": limit} | |
| 300 | + extra = "" | |
| 301 | + if category: | |
| 302 | + extra = " and category = cast(:cat as text)" | |
| 303 | + params["cat"] = category[:40] | |
| 304 | + rows = await fetch_all(conn, f"select * from news_items where company_id = :id{extra} order by coalesce(published_at, first_seen_at) desc limit :lim", **params) | |
| 305 | + return {"items": [ser.news_item(r) for r in rows]} | |
| 306 | + | |
| 307 | + | |
| 308 | +@router.get("/companies/{key}/sensors", summary="Sensors attached to the company") | |
| 309 | +async def company_sensors(key: str, response: Response, include_retired: str | None = None) -> dict[str, Any]: | |
| 310 | + _pub(response, 120) | |
| 311 | + async with connection() as conn: | |
| 312 | + c = await q.require_company(conn, key) | |
| 313 | + extra = "" if q.parse_bool(include_retired) else " and status <> 'retired'" | |
| 314 | + rows = await fetch_all(conn, f"select * from sensors where company_id = :id{extra} order by quality_score desc, surface limit 500", id=c["id"]) | |
| 315 | + return {"items": [ser.sensor(r) for r in rows]} | |
| 316 | + | |
| 317 | + | |
| 318 | +@router.get("/companies/{key}/history", summary="Historical page viewer index (≤ 20 versions per sensor)") | |
| 319 | +async def company_history(key: str, response: Response, versions: int = Query(20, ge=1, le=20)) -> dict[str, Any]: | |
| 320 | + _pub(response, 300) | |
| 321 | + async with connection() as conn: | |
| 322 | + c = await q.require_company(conn, key) | |
| 323 | + sensors = await fetch_all(conn, "select * from sensors where company_id = :id and status <> 'retired' and snapshot_count > 0 " | |
| 324 | + "order by quality_score desc, surface limit 200", id=c["id"]) | |
| 325 | + if not sensors: | |
| 326 | + sensors = await fetch_all(conn, "select * from sensors where company_id = :id and status <> 'retired' order by quality_score desc, surface limit 200", id=c["id"]) | |
| 327 | + snaps = await fetch_all(conn, "select s.* from sensors sn cross join lateral (select * from snapshots x where x.sensor_id = sn.id " | |
| 328 | + "order by x.fetched_at desc limit :v) s where sn.company_id = :id and sn.status <> 'retired' order by s.sensor_id, s.fetched_at desc", | |
| 329 | + id=c["id"], v=versions) | |
| 330 | + by_sensor: dict[str, list[dict[str, Any]]] = defaultdict(list) | |
| 331 | + for s in snaps: | |
| 332 | + by_sensor[s["sensor_id"]].append(ser.snapshot(s)) | |
| 333 | + out = [] | |
| 334 | + for s in sensors: | |
| 335 | + item = ser.sensor(s) | |
| 336 | + item["versions"] = by_sensor.get(s["id"], []) | |
| 337 | + out.append(item) | |
| 338 | + return {"sensors": out} | |
| 339 | + | |
| 340 | + | |
| 341 | +@router.get("/companies/{key}/similar", summary="Similar companies (industry, country, importance)") | |
| 342 | +async def company_similar(key: str, response: Response, limit: int = Query(8, ge=1, le=50)) -> dict[str, Any]: | |
| 343 | + _pub(response, 600) | |
| 344 | + async with connection() as conn: | |
| 345 | + c = await q.require_company(conn, key) | |
| 346 | + rows = await fetch_all(conn, "select c.id from companies c where c.id <> :id and c.status = 'ACTIVE' and " | |
| 347 | + "(c.industry_primary = cast(:ip as text) or c.industries && cast(:inds as text[]) or c.country = cast(:country as char(2))) " | |
| 348 | + "order by (c.industry_primary is not distinct from cast(:ip as text) and :ip is not null) desc, " | |
| 349 | + "(c.industries && cast(:inds as text[])) desc, (c.country is not distinct from cast(:country as char(2))) desc, " | |
| 350 | + "abs(c.importance - :imp) asc, c.importance desc limit :lim", | |
| 351 | + id=c["id"], ip=c.get("industry_primary"), inds=list(c.get("industries") or []), country=c.get("country"), | |
| 352 | + imp=float(c.get("importance") or 0), lim=limit) | |
| 353 | + cards = await q.fetch_cards_by_ids(conn, [r["id"] for r in rows]) | |
| 354 | + return {"items": [ser.company_card(r) for r in cards]} | |
added
src/companyatlas/api/routers/countries.py
+54 −0
@@ -0,0 +1,54 @@ | ||
| 1 | +"""Country Atlas: `/countries`, `/countries/{code}` — code accepts ISO-2 (`CA`) or a name slug (`canada`). Cached 300 s.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, HTTPException, Request | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api import serializers as ser | |
| 11 | +from companyatlas.api.common import cached, cached_response | |
| 12 | +from companyatlas.db import connection, fetch_all | |
| 13 | + | |
| 14 | +ORDER = 20 | |
| 15 | +router = APIRouter(prefix="/api/v1", tags=["countries"]) | |
| 16 | + | |
| 17 | + | |
| 18 | +@router.get("/countries", summary="Living index per country") | |
| 19 | +async def countries(request: Request) -> Any: | |
| 20 | + return cached_response(request, {"items": await agg.cached_country_rows()}, 300) | |
| 21 | + | |
| 22 | + | |
| 23 | +@router.get("/countries/{code}", summary="Country detail") | |
| 24 | +async def country_detail(code: str, request: Request) -> Any: | |
| 25 | + async def produce() -> dict[str, Any] | None: | |
| 26 | + async with connection() as conn: | |
| 27 | + ref = await agg.resolve_country(conn, code) | |
| 28 | + if ref is None: | |
| 29 | + return None | |
| 30 | + cc = ref["code"] | |
| 31 | + rows = await agg.cached_country_rows() | |
| 32 | + row = next((r for r in rows if r["code"] == cc), None) or { | |
| 33 | + "code": cc, "slug": None, "name": ref["name"], "region": ref.get("region"), "subregion": ref.get("subregion"), "companies": 0, "events_7d": 0, | |
| 34 | + "events_30d": 0, "hiring_momentum_30d": None, "activity_score": None, "ai_adoption": None, "industry_mix": [], "lat": ref.get("lat"), | |
| 35 | + "lon": ref.get("lon")} | |
| 36 | + where, params = q.company_filters(country=cc, status="ACTIVE") | |
| 37 | + ids, _t = await q.company_page_ids(conn, where, params, sort="activity", limit=24) | |
| 38 | + companies = [ser.company_card(r) for r in await q.fetch_cards_by_ids(conn, ids, sparkline=True)] | |
| 39 | + events = await agg.live_events(conn, 20, country=cc) | |
| 40 | + movers = await agg.ranking_cards(conn, "most_active", "7d", country=cc, limit=10) | |
| 41 | + new_ids, _t2 = await q.company_page_ids(conn, where, params, sort="recent", limit=10) | |
| 42 | + new_entrants = [ser.company_card(r) for r in await q.fetch_cards_by_ids(conn, new_ids)] | |
| 43 | + series = await fetch_all(conn, "select ms.day, avg(ms.value) as value, avg(ms.confidence) as confidence from metric_series ms " | |
| 44 | + "join companies c on c.id = ms.company_id where ms.metric = 'activity_score' and c.country = cast(:c as char(2)) " | |
| 45 | + "and ms.day >= :d group by ms.day order by ms.day", c=cc, d=q.days_ago(90).date()) | |
| 46 | + industries = await agg.industry_rows(conn, country=cc) | |
| 47 | + signals = await fetch_all(conn, "select * from signals where scope = 'country' and scope_key = :c and status = 'active' order by detected_at desc limit 10", c=cc) | |
| 48 | + return {**row, "companies_total": row["companies"], "companies": companies, "events": events, "movers": movers, "new_entrants": new_entrants, | |
| 49 | + "series": [ser.metric_point(r) for r in series], "industries": [i for i in industries if i["companies"] > 0][:40], | |
| 50 | + "signals": [ser.signal(s) for s in signals]} | |
| 51 | + payload = await cached(f"country:{code.strip().lower()[:80]}", 300, produce) | |
| 52 | + if payload is None: | |
| 53 | + raise HTTPException(status_code=404, detail="country not found") | |
| 54 | + return cached_response(request, payload, 300) | |
added
src/companyatlas/api/routers/events.py
+104 −0
@@ -0,0 +1,104 @@ | ||
| 1 | +"""Events: `/events`, `/events/types`, `/events/summary`, `/events/{id}`.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from collections import defaultdict | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from fastapi import APIRouter, HTTPException, Query, Request, Response | |
| 8 | + | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api import serializers as ser | |
| 11 | +from companyatlas.api.common import PageDep, cached, cached_response, page_payload, public_cache_value | |
| 12 | +from companyatlas.db import connection, fetch_all, fetch_one | |
| 13 | +from companyatlas.taxonomy import EVENT_SUBTYPES, EventType | |
| 14 | + | |
| 15 | +ORDER = 20 | |
| 16 | +router = APIRouter(prefix="/api/v1", tags=["events"]) | |
| 17 | + | |
| 18 | + | |
| 19 | +@router.get("/events", summary="Event feed with filters") | |
| 20 | +async def list_events(response: Response, p: PageDep, event_type: str | None = None, event_subtype: str | None = None, | |
| 21 | + country: str | None = None, industry: str | None = None, since: str | None = None, until: str | None = None, | |
| 22 | + min_importance: float | None = Query(None, ge=0, le=1), min_confidence: float | None = Query(None, ge=0, le=1), | |
| 23 | + q_: str | None = Query(None, alias="q", max_length=200), surface: str | None = None, origin: str | None = None, | |
| 24 | + company: str | None = None, status: str = Query("active", pattern="^(active|retracted|duplicate|review|all)$"), | |
| 25 | + sort: str = Query("recent", pattern="^(recent|importance)$")) -> dict[str, Any]: | |
| 26 | + response.headers["cache-control"] = public_cache_value(30) | |
| 27 | + async with connection() as conn: | |
| 28 | + company_id = None | |
| 29 | + if company: | |
| 30 | + company_id = (await q.require_company(conn, company))["id"] | |
| 31 | + where, params = q.event_filters(company_id=company_id, event_type=event_type, event_subtype=event_subtype, country=country, industry=industry, | |
| 32 | + since=q.parse_iso(since), until=q.parse_iso(until, "until"), min_importance=min_importance, | |
| 33 | + min_confidence=min_confidence, q=q_, surface=surface, origin=origin, status=None if status == "all" else status) | |
| 34 | + rows = await q.fetch_events(conn, where, params, sort=sort, limit=p.per_page, offset=p.offset) | |
| 35 | + total = await q.count_events(conn, where, params) | |
| 36 | + return page_payload([ser.event(r) for r in rows], total, p) | |
| 37 | + | |
| 38 | + | |
| 39 | +@router.get("/events/types", summary="Event taxonomy with 30-day counts") | |
| 40 | +async def event_types(request: Request) -> Any: | |
| 41 | + async def produce() -> dict[str, Any]: | |
| 42 | + async with connection() as conn: | |
| 43 | + rows = await fetch_all(conn, "select event_type, event_subtype, count(*) as n from events where status = 'active' and detected_at >= :d " | |
| 44 | + "group by event_type, event_subtype", d=q.days_ago(30)) | |
| 45 | + counts: dict[str, dict[str, int]] = defaultdict(dict) | |
| 46 | + for r in rows: | |
| 47 | + counts[r["event_type"]][r["event_subtype"]] = int(r["n"]) | |
| 48 | + subtypes_by_type: dict[str, list[str]] = defaultdict(list) | |
| 49 | + for sub, (typ, _imp) in EVENT_SUBTYPES.items(): | |
| 50 | + subtypes_by_type[typ.value].append(sub) | |
| 51 | + types = [] | |
| 52 | + for t in [x.value for x in EventType] + [k for k in counts if k not in {x.value for x in EventType}]: | |
| 53 | + observed = counts.get(t, {}) | |
| 54 | + subs = list(dict.fromkeys(subtypes_by_type.get(t, []) + list(observed))) | |
| 55 | + items = sorted(({"event_subtype": s, "count_30d": observed.get(s, 0)} for s in subs), key=lambda x: (-x["count_30d"], x["event_subtype"])) | |
| 56 | + types.append({"event_type": t, "subtypes": items, "count_30d": sum(observed.values())}) | |
| 57 | + types.sort(key=lambda x: (-x["count_30d"], x["event_type"])) | |
| 58 | + return {"types": types} | |
| 59 | + return cached_response(request, await cached("events:types", 300, produce), 300) | |
| 60 | + | |
| 61 | + | |
| 62 | +@router.get("/events/summary", summary="Event counts by type / industry / country with delta vs previous window") | |
| 63 | +async def event_summary(request: Request, days: int = Query(7, ge=1, le=365), group: str = Query("type", pattern="^(type|industry|country)$")) -> Any: | |
| 64 | + async def produce() -> dict[str, Any]: | |
| 65 | + col = {"type": "e.event_type", "industry": "c.industry_primary", "country": "c.country"}[group] | |
| 66 | + d1, d2 = q.days_ago(days), q.days_ago(days * 2) | |
| 67 | + async with connection() as conn: | |
| 68 | + rows = await fetch_all(conn, f"select {col} as key, count(*) filter (where e.detected_at >= :d1) as cur, " | |
| 69 | + f"count(*) filter (where e.detected_at < :d1) as prev from events e join companies c on c.id = e.company_id " | |
| 70 | + f"where e.status = 'active' and e.detected_at >= :d2 and {col} is not null group by 1 order by cur desc limit 200", | |
| 71 | + d1=d1, d2=d2) | |
| 72 | + items = [] | |
| 73 | + for r in rows: | |
| 74 | + cur, prev = int(r["cur"]), int(r["prev"]) | |
| 75 | + items.append({"key": r["key"], "count": cur, "previous": prev, "delta_pct": round((cur - prev) / prev * 100, 1) if prev else None}) | |
| 76 | + return {"days": days, "group": group, "items": items} | |
| 77 | + return cached_response(request, await cached(f"events:summary:{days}:{group}", 300, produce), 300) | |
| 78 | + | |
| 79 | + | |
| 80 | +@router.get("/events/{event_id}", summary="Event detail with sources, change and company") | |
| 81 | +async def event_detail(event_id: str, response: Response) -> dict[str, Any]: | |
| 82 | + response.headers["cache-control"] = public_cache_value(60) | |
| 83 | + async with connection() as conn: | |
| 84 | + row = await q.fetch_event(conn, event_id) | |
| 85 | + if row is None: | |
| 86 | + raise HTTPException(status_code=404, detail="event not found") | |
| 87 | + out = ser.event(row) | |
| 88 | + sources = await fetch_all(conn, "select * from event_sources where event_id = :id order by detected_at asc limit 50", id=event_id) | |
| 89 | + out["sources"] = [ser.event_source(s) for s in sources] | |
| 90 | + if not out["sources"] and row.get("source_url"): | |
| 91 | + out["sources"] = [{"source_url": row["source_url"], "surface": row.get("surface"), "detected_at": row["detected_at"], "kind": "primary", | |
| 92 | + "sensor_id": row.get("sensor_id"), "snapshot_id": row.get("snapshot_after")}] | |
| 93 | + out["change"] = None | |
| 94 | + if row.get("change_id"): | |
| 95 | + ch = await fetch_one(conn, "select * from changes where id = :id", id=row["change_id"]) | |
| 96 | + out["change"] = ser.change(ch) if ch else None | |
| 97 | + cards = await q.fetch_cards_by_ids(conn, [row["company_id"]]) | |
| 98 | + if cards: | |
| 99 | + out["company"] = ser.company_card(cards[0]) | |
| 100 | + if row.get("cluster_id"): | |
| 101 | + cl = await fetch_one(conn, "select id, cluster_key, source_count, surfaces, confidence, first_detected_at, last_detected_at, canonical_event_id " | |
| 102 | + "from event_clusters where id = :id", id=row["cluster_id"]) | |
| 103 | + out["cluster"] = cl | |
| 104 | + return out | |
added
src/companyatlas/api/routers/exports.py
+183 −0
@@ -0,0 +1,183 @@ | ||
| 1 | +"""Streamed exports (json / ndjson / csv) with bounded limits and keyset iteration — derived data only, never raw page content.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import csv | |
| 5 | +import io | |
| 6 | +from collections.abc import AsyncIterator, Callable | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +import orjson | |
| 10 | +from fastapi import APIRouter, HTTPException, Query | |
| 11 | +from fastapi.responses import StreamingResponse | |
| 12 | + | |
| 13 | +from companyatlas.api import queries as q | |
| 14 | +from companyatlas.api import serializers as ser | |
| 15 | +from companyatlas.api.common import _default | |
| 16 | +from companyatlas.db import connection, fetch_all | |
| 17 | + | |
| 18 | +ORDER = 20 | |
| 19 | +router = APIRouter(prefix="/api/v1", tags=["exports"]) | |
| 20 | +BATCH = 1000 | |
| 21 | +MAX_EVENTS = 10_000 | |
| 22 | +MAX_COMPANIES = 20_000 | |
| 23 | +MAX_JOBS = 10_000 | |
| 24 | +MEDIA = {"json": "application/json", "ndjson": "application/x-ndjson", "csv": "text/csv; charset=utf-8"} | |
| 25 | +EVENT_COLUMNS = ["id", "detected_at", "company_slug", "company_name", "company_domain", "country", "event_type", "event_subtype", "importance", "confidence", | |
| 26 | + "confidence_label", "title", "summary", "old_value", "new_value", "source_url", "surface", "origin", "status"] | |
| 27 | +COMPANY_COLUMNS = ["id", "slug", "display_name", "canonical_domain", "website", "country", "hq_city", "industry_primary", "industries", "public_company", "ticker", | |
| 28 | + "tier", "importance", "activity_score", "hiring_momentum_30d", "open_jobs", "ai_adoption", "sensors", "events", "last_event_at"] | |
| 29 | +JOB_COLUMNS = ["id", "company_slug", "title", "department", "location_text", "city", "country", "remote", "employment_type", "seniority", "url", "posted_at", | |
| 30 | + "first_seen_at", "last_seen_at", "removed_at", "status", "is_ai"] | |
| 31 | + | |
| 32 | + | |
| 33 | +def _dumps(obj: Any) -> bytes: | |
| 34 | + return orjson.dumps(obj, option=orjson.OPT_UTC_Z | orjson.OPT_NON_STR_KEYS, default=_default) | |
| 35 | + | |
| 36 | + | |
| 37 | +def _csv_row(values: list[Any]) -> bytes: | |
| 38 | + buf = io.StringIO() | |
| 39 | + csv.writer(buf, lineterminator="\n").writerow(["" if v is None else (v.isoformat() if hasattr(v, "isoformat") else v) for v in values]) | |
| 40 | + return buf.getvalue().encode("utf-8") | |
| 41 | + | |
| 42 | + | |
| 43 | +def _event_flat(ev: dict[str, Any]) -> list[Any]: | |
| 44 | + c = ev["company"] | |
| 45 | + return [ev["id"], ev["detected_at"], c["slug"], c["display_name"], c["canonical_domain"], c["country"], ev["event_type"], ev["event_subtype"], ev["importance"], | |
| 46 | + ev["confidence"], ev["confidence_label"], ev["title"], ev["summary"], ev["old_value"], ev["new_value"], ev["source_url"], ev["surface"], ev["origin"], | |
| 47 | + ev["status"]] | |
| 48 | + | |
| 49 | + | |
| 50 | +def _company_flat(c: dict[str, Any]) -> list[Any]: | |
| 51 | + m, n = c["metrics"], c["counts"] | |
| 52 | + return [c["id"], c["slug"], c["display_name"], c["canonical_domain"], c["website"], c["country"], c["hq_city"], c["industry_primary"], "|".join(c["industries"]), | |
| 53 | + c["public_company"], c["ticker"], c["tier"], c["importance"], m.get("activity_score"), m.get("hiring_momentum_30d"), m.get("open_jobs"), m.get("ai_adoption"), | |
| 54 | + n["sensors"], n["events"], c["last_event_at"]] | |
| 55 | + | |
| 56 | + | |
| 57 | +def _job_flat(j: dict[str, Any]) -> list[Any]: | |
| 58 | + return [j["id"], j.get("company_slug"), j["title"], j["department"], j["location_text"], j["city"], j["country"], j["remote"], j["employment_type"], j["seniority"], | |
| 59 | + j["url"], j["posted_at"], j["first_seen_at"], j["last_seen_at"], j["removed_at"], j["status"], j["is_ai"]] | |
| 60 | + | |
| 61 | + | |
| 62 | +async def _stream(fmt: str, rows: AsyncIterator[dict[str, Any]], columns: list[str], flat: Callable[[dict[str, Any]], list[Any]]) -> AsyncIterator[bytes]: | |
| 63 | + if fmt == "csv": | |
| 64 | + yield _csv_row(columns) | |
| 65 | + async for r in rows: | |
| 66 | + yield _csv_row(flat(r)) | |
| 67 | + return | |
| 68 | + if fmt == "ndjson": | |
| 69 | + async for r in rows: | |
| 70 | + yield _dumps(r) + b"\n" | |
| 71 | + return | |
| 72 | + yield b"[" | |
| 73 | + first = True | |
| 74 | + async for r in rows: | |
| 75 | + yield (b"" if first else b",") + _dumps(r) | |
| 76 | + first = False | |
| 77 | + yield b"]" | |
| 78 | + | |
| 79 | + | |
| 80 | +def _response(fmt: str, name: str, body: AsyncIterator[bytes]) -> StreamingResponse: | |
| 81 | + if fmt not in MEDIA: | |
| 82 | + raise HTTPException(status_code=404, detail="unsupported format (json, ndjson, csv)") | |
| 83 | + return StreamingResponse(body, media_type=MEDIA[fmt], headers={"content-disposition": f'attachment; filename="company-atlas-{name}.{fmt}"', | |
| 84 | + "cache-control": "public, max-age=300", "x-accel-buffering": "no"}) | |
| 85 | + | |
| 86 | + | |
| 87 | +async def _iter_events(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]: | |
| 88 | + sent = 0 | |
| 89 | + cursor: tuple[Any, str] | None = None | |
| 90 | + async with connection() as conn: | |
| 91 | + while sent < limit: | |
| 92 | + w = list(where) | |
| 93 | + p = dict(params) | |
| 94 | + if cursor: | |
| 95 | + w.append("(e.detected_at, e.id) < (:cur_at, :cur_id)") | |
| 96 | + p.update(cur_at=cursor[0], cur_id=cursor[1]) | |
| 97 | + rows = await q.fetch_events(conn, w, p, sort="recent", limit=min(BATCH, limit - sent)) | |
| 98 | + if not rows: | |
| 99 | + return | |
| 100 | + for r in rows: | |
| 101 | + yield ser.event(r) | |
| 102 | + sent += len(rows) | |
| 103 | + cursor = (rows[-1]["detected_at"], rows[-1]["id"]) | |
| 104 | + | |
| 105 | + | |
| 106 | +@router.get("/export/events.{fmt}", summary="Export events (≤ 10 000 rows)") | |
| 107 | +async def export_events(fmt: str, since: str | None = None, until: str | None = None, event_type: str | None = None, country: str | None = None, | |
| 108 | + industry: str | None = None, company: str | None = None, min_importance: float | None = Query(None, ge=0, le=1), | |
| 109 | + limit: int = Query(MAX_EVENTS, ge=1, le=MAX_EVENTS)) -> StreamingResponse: | |
| 110 | + company_id = None | |
| 111 | + if company: | |
| 112 | + async with connection() as conn: | |
| 113 | + company_id = (await q.require_company(conn, company))["id"] | |
| 114 | + where, params = q.event_filters(company_id=company_id, event_type=event_type, country=country, industry=industry, since=q.parse_iso(since), | |
| 115 | + until=q.parse_iso(until, "until"), min_importance=min_importance) | |
| 116 | + return _response(fmt, "events", _stream(fmt, _iter_events(where, params, limit), EVENT_COLUMNS, _event_flat)) | |
| 117 | + | |
| 118 | + | |
| 119 | +async def _iter_companies(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]: | |
| 120 | + sent = 0 | |
| 121 | + last_id = "" | |
| 122 | + async with connection() as conn: | |
| 123 | + while sent < limit: | |
| 124 | + w = list(where) + ["c.id > :cur_id"] | |
| 125 | + rows = await fetch_all(conn, f"select c.id from companies c where {' and '.join(w)} order by c.id limit :lim", **params, cur_id=last_id, | |
| 126 | + lim=min(BATCH, limit - sent)) | |
| 127 | + if not rows: | |
| 128 | + return | |
| 129 | + ids = [r["id"] for r in rows] | |
| 130 | + for card in await q.fetch_cards_by_ids(conn, ids): | |
| 131 | + yield ser.company_card(card) | |
| 132 | + sent += len(ids) | |
| 133 | + last_id = ids[-1] | |
| 134 | + | |
| 135 | + | |
| 136 | +@router.get("/export/companies.{fmt}", summary="Export companies (≤ 20 000 rows)") | |
| 137 | +async def export_companies(fmt: str, country: str | None = None, industry: str | None = None, status: str | None = None, tier: int | None = Query(None, ge=1, le=4), | |
| 138 | + limit: int = Query(MAX_COMPANIES, ge=1, le=MAX_COMPANIES)) -> StreamingResponse: | |
| 139 | + where, params = q.company_filters(country=country, industry=industry, status=status, tier=tier) | |
| 140 | + return _response(fmt, "companies", _stream(fmt, _iter_companies(where or ["true"], params, limit), COMPANY_COLUMNS, _company_flat)) | |
| 141 | + | |
| 142 | + | |
| 143 | +async def _iter_jobs(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]: | |
| 144 | + sent = 0 | |
| 145 | + last_id = "" | |
| 146 | + async with connection() as conn: | |
| 147 | + while sent < limit: | |
| 148 | + w = list(where) + ["j.id > :cur_id"] | |
| 149 | + rows = await fetch_all(conn, f"select j.*, c.slug as company_slug from jobs j join companies c on c.id = j.company_id where {' and '.join(w)} " | |
| 150 | + "order by j.id limit :lim", **params, cur_id=last_id, lim=min(BATCH, limit - sent)) | |
| 151 | + if not rows: | |
| 152 | + return | |
| 153 | + for r in rows: | |
| 154 | + item = ser.job(r) | |
| 155 | + item["company_slug"] = r["company_slug"] | |
| 156 | + yield item | |
| 157 | + sent += len(rows) | |
| 158 | + last_id = rows[-1]["id"] | |
| 159 | + | |
| 160 | + | |
| 161 | +@router.get("/export/jobs.{fmt}", summary="Export jobs (≤ 10 000 rows)") | |
| 162 | +async def export_jobs(fmt: str, company: str | None = None, since: str | None = None, status: str = Query("open", pattern="^(open|removed|all)$"), | |
| 163 | + country: str | None = None, ai: str | None = None, limit: int = Query(MAX_JOBS, ge=1, le=MAX_JOBS)) -> StreamingResponse: | |
| 164 | + where: list[str] = ["true"] | |
| 165 | + params: dict[str, Any] = {} | |
| 166 | + if company: | |
| 167 | + async with connection() as conn: | |
| 168 | + params["cid"] = (await q.require_company(conn, company))["id"] | |
| 169 | + where.append("j.company_id = :cid") | |
| 170 | + since_dt = q.parse_iso(since) | |
| 171 | + if since_dt: | |
| 172 | + where.append("j.first_seen_at >= :since") | |
| 173 | + params["since"] = since_dt | |
| 174 | + if status == "open": | |
| 175 | + where.append("j.status = 'open'") | |
| 176 | + elif status == "removed": | |
| 177 | + where.append("j.status <> 'open'") | |
| 178 | + if country: | |
| 179 | + where.append("j.country = cast(:country as char(2))") | |
| 180 | + params["country"] = country.upper()[:2] | |
| 181 | + if q.parse_bool(ai): | |
| 182 | + where.append("j.is_ai") | |
| 183 | + return _response(fmt, "jobs", _stream(fmt, _iter_jobs(where, params, limit), JOB_COLUMNS, _job_flat)) | |
added
src/companyatlas/api/routers/industries.py
+60 −0
@@ -0,0 +1,60 @@ | ||
| 1 | +"""Industry Atlas: `/industries`, `/industries/{slug}` (cached 300 s).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, HTTPException, Request | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api import serializers as ser | |
| 11 | +from companyatlas.api.common import cached, cached_response | |
| 12 | +from companyatlas.db import connection, fetch_all, fetch_one | |
| 13 | + | |
| 14 | +ORDER = 20 | |
| 15 | +router = APIRouter(prefix="/api/v1", tags=["industries"]) | |
| 16 | + | |
| 17 | + | |
| 18 | +@router.get("/industries", summary="Living index per industry") | |
| 19 | +async def industries(request: Request) -> Any: | |
| 20 | + return cached_response(request, {"items": await agg.cached_industry_rows()}, 300) | |
| 21 | + | |
| 22 | + | |
| 23 | +@router.get("/industries/{slug}", summary="Industry detail") | |
| 24 | +async def industry_detail(slug: str, request: Request) -> Any: | |
| 25 | + slug = slug.strip().lower()[:80] | |
| 26 | + | |
| 27 | + async def produce() -> dict[str, Any] | None: | |
| 28 | + rows = await agg.cached_industry_rows() | |
| 29 | + row = next((r for r in rows if r["slug"] == slug), None) | |
| 30 | + async with connection() as conn: | |
| 31 | + if row is None: | |
| 32 | + tax = await fetch_one(conn, "select slug, name, parent_slug, description from industries where slug = :s", s=slug) | |
| 33 | + if tax is None: | |
| 34 | + return None | |
| 35 | + row = {**tax, "companies": 0, "events_7d": 0, "events_30d": 0, "hiring_momentum_30d": None, "activity_score": None, "ai_adoption": None, | |
| 36 | + "top_event_types": []} | |
| 37 | + where, params = q.company_filters(industry=slug, status="ACTIVE") | |
| 38 | + ids, _total = await q.company_page_ids(conn, where, params, sort="activity", limit=24) | |
| 39 | + companies = [ser.company_card(r) for r in await q.fetch_cards_by_ids(conn, ids, sparkline=True)] | |
| 40 | + events = await agg.live_events(conn, 20, industry=slug) | |
| 41 | + d30 = q.days_ago(30) | |
| 42 | + h = await fetch_one(conn, "select count(*) filter (where j.status = 'open') as open, count(*) filter (where j.first_seen_at >= :d) as new_30d, " | |
| 43 | + "count(*) filter (where j.removed_at >= :d) as removed_30d from jobs j join companies c on c.id = j.company_id " | |
| 44 | + "where cast(:s as text) = any(c.industries)", d=d30, s=slug) or {} | |
| 45 | + series = await fetch_all(conn, "select ms.day, avg(ms.value) as value, avg(ms.confidence) as confidence from metric_series ms " | |
| 46 | + "join companies c on c.id = ms.company_id where ms.metric = 'activity_score' and cast(:s as text) = any(c.industries) " | |
| 47 | + "and ms.day >= :d group by ms.day order by ms.day", s=slug, d=q.days_ago(90).date()) | |
| 48 | + countries = await fetch_all(conn, "select country, count(*) as companies from companies c where cast(:s as text) = any(c.industries) and " | |
| 49 | + "country is not null and status = 'ACTIVE' group by country order by companies desc limit 50", s=slug) | |
| 50 | + children = await fetch_all(conn, "select slug, name from industries where parent_slug = :s order by sort_order, name", s=slug) | |
| 51 | + signals = await fetch_all(conn, "select * from signals where scope = 'industry' and scope_key = :s and status = 'active' order by detected_at desc limit 10", s=slug) | |
| 52 | + return {**row, "companies_total": row["companies"], "companies": companies, "events": events, | |
| 53 | + "hiring": {"open": int(h.get("open") or 0), "new_30d": int(h.get("new_30d") or 0), "removed_30d": int(h.get("removed_30d") or 0), | |
| 54 | + "momentum_30d": row.get("hiring_momentum_30d")}, | |
| 55 | + "series": [ser.metric_point(r) for r in series], "countries": [{"country": r["country"], "companies": int(r["companies"])} for r in countries], | |
| 56 | + "trending": [], "children": children, "signals": [ser.signal(s) for s in signals]} | |
| 57 | + payload = await cached(f"industry:{slug}", 300, produce) | |
| 58 | + if payload is None: | |
| 59 | + raise HTTPException(status_code=404, detail="industry not found") | |
| 60 | + return cached_response(request, payload, 300) | |
added
src/companyatlas/api/routers/live.py
+37 −0
@@ -0,0 +1,37 @@ | ||
| 1 | +"""Live feed: `/live` (latest active events, never cached) and `/live/stream` (SSE).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request, Response | |
| 7 | + | |
| 8 | +from companyatlas.api import queries as q | |
| 9 | +from companyatlas.api import serializers as ser | |
| 10 | +from companyatlas.api.common import NO_STORE | |
| 11 | +from companyatlas.api.sse import MAX_STREAM_S, live_event_stream, sse_response | |
| 12 | +from companyatlas.db import connection | |
| 13 | + | |
| 14 | +ORDER = 10 | |
| 15 | +router = APIRouter(prefix="/api/v1", tags=["live"]) | |
| 16 | + | |
| 17 | + | |
| 18 | +@router.get("/live", summary="Latest active events") | |
| 19 | +async def live(response: Response, limit: int = Query(50, ge=1, le=200), since: str | None = None, event_type: str | None = None, | |
| 20 | + min_importance: float | None = Query(None, ge=0, le=1), country: str | None = None, industry: str | None = None) -> dict[str, Any]: | |
| 21 | + response.headers["cache-control"] = NO_STORE | |
| 22 | + since_dt = q.parse_iso(since) | |
| 23 | + where, params = q.event_filters(event_type=event_type, min_importance=min_importance, country=country, industry=industry) | |
| 24 | + if since_dt is not None: | |
| 25 | + where.append("e.detected_at > :live_since") | |
| 26 | + params["live_since"] = since_dt | |
| 27 | + async with connection() as conn: | |
| 28 | + rows = await q.fetch_events(conn, where, params, sort="recent", limit=limit) | |
| 29 | + items = [ser.event(r) for r in rows] | |
| 30 | + return {"items": items, "count": len(items), "cursor": items[0]["detected_at"] if items else (since_dt or q.now_utc()), "server_time": q.now_utc()} | |
| 31 | + | |
| 32 | + | |
| 33 | +@router.get("/live/stream", summary="Server-sent events stream of new events") | |
| 34 | +async def live_stream(request: Request, since: str | None = None, event_type: str | None = None, min_importance: float | None = Query(None, ge=0, le=1), | |
| 35 | + max_s: float = Query(MAX_STREAM_S, ge=1, le=MAX_STREAM_S)) -> Any: | |
| 36 | + since_dt = q.parse_iso(since) | |
| 37 | + return sse_response(live_event_stream(request, since=since_dt, event_type=event_type, min_importance=min_importance, max_s=max_s)) | |
added
src/companyatlas/api/routers/owner.py
+163 −0
@@ -0,0 +1,163 @@ | ||
| 1 | +"""Owner endpoints (X-CA-Owner-Token): watchlist and alerts. The owner row is created on first use; nothing here is cached.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Depends, HTTPException, Query, Response | |
| 7 | +from pydantic import BaseModel, Field | |
| 8 | + | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api import serializers as ser | |
| 11 | +from companyatlas.api.common import NO_STORE, owner_hash | |
| 12 | +from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 13 | +from companyatlas.ids import new_id | |
| 14 | +from companyatlas.taxonomy import EventType | |
| 15 | + | |
| 16 | +ORDER = 20 | |
| 17 | +router = APIRouter(prefix="/api/v1", tags=["owner"]) | |
| 18 | +MAX_WATCHLIST = 200 | |
| 19 | +MAX_ALERTS = 100 | |
| 20 | +KNOWN_TYPES = {t.value for t in EventType} | |
| 21 | + | |
| 22 | + | |
| 23 | +async def _ensure_owner(conn: Any, oh: str) -> str: | |
| 24 | + await execute(conn, "insert into owners (token_hash) values (:h) on conflict (token_hash) do update set last_seen_at = now()", h=oh) | |
| 25 | + wl = await fetch_val(conn, "select id from watchlists where owner_hash = :h order by created_at limit 1", h=oh) | |
| 26 | + if wl is None: | |
| 27 | + wl = new_id("watchlist") | |
| 28 | + await execute(conn, "insert into watchlists (id, owner_hash) values (:id, :h)", id=wl, h=oh) | |
| 29 | + return wl | |
| 30 | + | |
| 31 | + | |
| 32 | +class WatchBody(BaseModel): | |
| 33 | + company: str = Field(min_length=1, max_length=200) | |
| 34 | + | |
| 35 | + | |
| 36 | +@router.get("/watchlist", summary="Watched companies and their recent events") | |
| 37 | +async def get_watchlist(response: Response, oh: str = Depends(owner_hash), events_limit: int = Query(30, ge=1, le=100)) -> dict[str, Any]: | |
| 38 | + response.headers["cache-control"] = NO_STORE | |
| 39 | + async with transaction() as conn: | |
| 40 | + wl = await _ensure_owner(conn, oh) | |
| 41 | + rows = await fetch_all(conn, "select company_id, added_at from watchlist_items where watchlist_id = :wl order by added_at desc limit :lim", wl=wl, lim=MAX_WATCHLIST) | |
| 42 | + ids = [r["company_id"] for r in rows] | |
| 43 | + cards = [ser.company_card(c) for c in await q.fetch_cards_by_ids(conn, ids, sparkline=True)] | |
| 44 | + events: list[dict[str, Any]] = [] | |
| 45 | + if ids: | |
| 46 | + where, params = q.event_filters() | |
| 47 | + where.append("e.company_id = any(cast(:wl_ids as text[]))") | |
| 48 | + params["wl_ids"] = ids | |
| 49 | + events = [ser.event(r) for r in await q.fetch_events(conn, where, params, limit=events_limit)] | |
| 50 | + added = {r["company_id"]: r["added_at"] for r in rows} | |
| 51 | + for c in cards: | |
| 52 | + c["added_at"] = added.get(c["id"]) | |
| 53 | + return {"id": wl, "items": cards, "events": events, "count": len(cards), "max": MAX_WATCHLIST} | |
| 54 | + | |
| 55 | + | |
| 56 | +@router.post("/watchlist", status_code=201, summary="Add a company to the watchlist") | |
| 57 | +async def add_to_watchlist(body: WatchBody, response: Response, oh: str = Depends(owner_hash)) -> dict[str, Any]: | |
| 58 | + response.headers["cache-control"] = NO_STORE | |
| 59 | + async with transaction() as conn: | |
| 60 | + wl = await _ensure_owner(conn, oh) | |
| 61 | + c = await q.require_company(conn, body.company) | |
| 62 | + n = await fetch_val(conn, "select count(*) from watchlist_items where watchlist_id = :wl", wl=wl) | |
| 63 | + exists = await fetch_val(conn, "select 1 from watchlist_items where watchlist_id = :wl and company_id = :cid", wl=wl, cid=c["id"]) | |
| 64 | + if not exists and int(n or 0) >= MAX_WATCHLIST: | |
| 65 | + raise HTTPException(status_code=409, detail=f"watchlist is full ({MAX_WATCHLIST})") | |
| 66 | + await execute(conn, "insert into watchlist_items (watchlist_id, company_id) values (:wl, :cid) on conflict do nothing", wl=wl, cid=c["id"]) | |
| 67 | + cards = await q.fetch_cards_by_ids(conn, [c["id"]]) | |
| 68 | + return {"added": not bool(exists), "company": ser.company_card(cards[0]) if cards else None} | |
| 69 | + | |
| 70 | + | |
| 71 | +@router.delete("/watchlist/{key}", summary="Remove a company from the watchlist") | |
| 72 | +async def remove_from_watchlist(key: str, response: Response, oh: str = Depends(owner_hash)) -> dict[str, Any]: | |
| 73 | + response.headers["cache-control"] = NO_STORE | |
| 74 | + async with transaction() as conn: | |
| 75 | + wl = await _ensure_owner(conn, oh) | |
| 76 | + c = await q.require_company(conn, key) | |
| 77 | + removed = await fetch_val(conn, "with d as (delete from watchlist_items where watchlist_id = :wl and company_id = :cid returning 1) select count(*) from d", | |
| 78 | + wl=wl, cid=c["id"]) | |
| 79 | + return {"removed": bool(removed), "company": c["slug"]} | |
| 80 | + | |
| 81 | + | |
| 82 | +class MetricCondition(BaseModel): | |
| 83 | + gt: float | None = None | |
| 84 | + lt: float | None = None | |
| 85 | + | |
| 86 | + | |
| 87 | +class AlertCondition(BaseModel): | |
| 88 | + event_types: list[str] | None = None | |
| 89 | + event_subtypes: list[str] | None = None | |
| 90 | + min_importance: float | None = Field(None, ge=0, le=1) | |
| 91 | + metrics: dict[str, MetricCondition] | None = None | |
| 92 | + industries: list[str] | None = None | |
| 93 | + countries: list[str] | None = None | |
| 94 | + | |
| 95 | + | |
| 96 | +class AlertBody(BaseModel): | |
| 97 | + name: str = Field(min_length=1, max_length=120) | |
| 98 | + company: str | None = Field(None, max_length=200) | |
| 99 | + condition: AlertCondition = Field(default_factory=AlertCondition) | |
| 100 | + channel: str = Field("web", pattern="^(web|webhook)$") | |
| 101 | + target: str | None = Field(None, max_length=500) | |
| 102 | + enabled: bool = True | |
| 103 | + | |
| 104 | + | |
| 105 | +@router.get("/alerts", summary="Alerts of this owner") | |
| 106 | +async def list_alerts(response: Response, oh: str = Depends(owner_hash)) -> dict[str, Any]: | |
| 107 | + response.headers["cache-control"] = NO_STORE | |
| 108 | + async with transaction() as conn: | |
| 109 | + await _ensure_owner(conn, oh) | |
| 110 | + rows = await fetch_all(conn, "select a.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 111 | + "c.country as company_country, c.logo_url as company_logo_url from alerts a left join companies c on c.id = a.company_id " | |
| 112 | + "where a.owner_hash = :h order by a.created_at desc limit :lim", h=oh, lim=MAX_ALERTS) | |
| 113 | + return {"items": [ser.alert(r) for r in rows], "count": len(rows), "max": MAX_ALERTS} | |
| 114 | + | |
| 115 | + | |
| 116 | +@router.post("/alerts", status_code=201, summary="Create an alert") | |
| 117 | +async def create_alert(body: AlertBody, response: Response, oh: str = Depends(owner_hash)) -> dict[str, Any]: | |
| 118 | + response.headers["cache-control"] = NO_STORE | |
| 119 | + cond = body.condition.model_dump(exclude_none=True) | |
| 120 | + unknown = [t for t in (cond.get("event_types") or []) if t.upper() not in KNOWN_TYPES] | |
| 121 | + if unknown: | |
| 122 | + raise HTTPException(status_code=422, detail=f"unknown event_types: {', '.join(unknown[:5])}") | |
| 123 | + if cond.get("event_types"): | |
| 124 | + cond["event_types"] = [t.upper() for t in cond["event_types"]] | |
| 125 | + if body.channel == "webhook" and (not body.target or not body.target.lower().startswith(("https://", "http://"))): | |
| 126 | + raise HTTPException(status_code=422, detail="target: webhook alerts need an http(s) URL") | |
| 127 | + async with transaction() as conn: | |
| 128 | + await _ensure_owner(conn, oh) | |
| 129 | + n = await fetch_val(conn, "select count(*) from alerts where owner_hash = :h", h=oh) | |
| 130 | + if int(n or 0) >= MAX_ALERTS: | |
| 131 | + raise HTTPException(status_code=409, detail=f"too many alerts ({MAX_ALERTS})") | |
| 132 | + company_id = (await q.require_company(conn, body.company))["id"] if body.company else None | |
| 133 | + if company_id is None and not cond: | |
| 134 | + raise HTTPException(status_code=422, detail="an alert needs a company or at least one condition") | |
| 135 | + aid = new_id("alert") | |
| 136 | + await execute(conn, "insert into alerts (id, owner_hash, company_id, name, condition, channel, target, enabled) " | |
| 137 | + "values (:id, :h, :cid, :name, cast(:cond as jsonb), :channel, :target, :enabled)", | |
| 138 | + id=aid, h=oh, cid=company_id, name=body.name.strip(), cond=jsonb(cond), channel=body.channel, target=body.target, enabled=body.enabled) | |
| 139 | + row = await fetch_one(conn, "select a.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 140 | + "c.country as company_country, c.logo_url as company_logo_url from alerts a left join companies c on c.id = a.company_id " | |
| 141 | + "where a.id = :id", id=aid) | |
| 142 | + return ser.alert(row or {}) | |
| 143 | + | |
| 144 | + | |
| 145 | +@router.delete("/alerts/{alert_id}", summary="Delete an alert") | |
| 146 | +async def delete_alert(alert_id: str, response: Response, oh: str = Depends(owner_hash)) -> dict[str, Any]: | |
| 147 | + response.headers["cache-control"] = NO_STORE | |
| 148 | + async with transaction() as conn: | |
| 149 | + await _ensure_owner(conn, oh) | |
| 150 | + removed = await fetch_val(conn, "with d as (delete from alerts where id = :id and owner_hash = :h returning 1) select count(*) from d", id=alert_id, h=oh) | |
| 151 | + if not removed: | |
| 152 | + raise HTTPException(status_code=404, detail="alert not found") | |
| 153 | + return {"removed": True, "id": alert_id} | |
| 154 | + | |
| 155 | + | |
| 156 | +@router.get("/alerts/deliveries", summary="Recent alert deliveries for this owner") | |
| 157 | +async def alert_deliveries(response: Response, oh: str = Depends(owner_hash), limit: int = Query(50, ge=1, le=200)) -> dict[str, Any]: | |
| 158 | + response.headers["cache-control"] = NO_STORE | |
| 159 | + async with transaction() as conn: | |
| 160 | + await _ensure_owner(conn, oh) | |
| 161 | + rows = await fetch_all(conn, "select d.*, a.name as alert_name, e.title as event_title from alert_deliveries d join alerts a on a.id = d.alert_id " | |
| 162 | + "left join events e on e.id = d.event_id where a.owner_hash = :h order by d.delivered_at desc limit :lim", h=oh, lim=limit) | |
| 163 | + return {"items": [ser.alert_delivery(r) for r in rows]} | |
added
src/companyatlas/api/routers/provenance.py
+201 −0
@@ -0,0 +1,201 @@ | ||
| 1 | +"""Provenance: sensors, snapshots (text/blocks from the object store), on-demand diffs and changes.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import json | |
| 5 | +import logging | |
| 6 | +from typing import Any | |
| 7 | + | |
| 8 | +from fastapi import APIRouter, HTTPException, Query, Response | |
| 9 | + | |
| 10 | +from companyatlas import archive | |
| 11 | +from companyatlas.api import queries as q | |
| 12 | +from companyatlas.api import serializers as ser | |
| 13 | +from companyatlas.api.common import public_cache_value | |
| 14 | +from companyatlas.db import connection, fetch_all, fetch_one, fetch_val | |
| 15 | + | |
| 16 | +log = logging.getLogger("companyatlas.api.provenance") | |
| 17 | +ORDER = 20 | |
| 18 | +router = APIRouter(prefix="/api/v1", tags=["provenance"]) | |
| 19 | +MAX_TEXT_BYTES = 200 * 1024 | |
| 20 | +MAX_BLOCKS = 2000 | |
| 21 | + | |
| 22 | + | |
| 23 | +async def _sensor_or_404(conn: Any, sensor_id: str) -> dict[str, Any]: | |
| 24 | + row = await fetch_one(conn, "select s.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 25 | + "c.country as company_country, c.logo_url as company_logo_url from sensors s join companies c on c.id = s.company_id " | |
| 26 | + "where s.id = :id", id=sensor_id) | |
| 27 | + if row is None: | |
| 28 | + raise HTTPException(status_code=404, detail="sensor not found") | |
| 29 | + return row | |
| 30 | + | |
| 31 | + | |
| 32 | +@router.get("/sensors/{sensor_id}", summary="Sensor with company and latest snapshot") | |
| 33 | +async def sensor_detail(sensor_id: str, response: Response) -> dict[str, Any]: | |
| 34 | + response.headers["cache-control"] = public_cache_value(60) | |
| 35 | + async with connection() as conn: | |
| 36 | + row = await _sensor_or_404(conn, sensor_id) | |
| 37 | + out = ser.sensor(row) | |
| 38 | + out["company"] = ser.company_ref(row) | |
| 39 | + latest = None | |
| 40 | + if row.get("last_snapshot_id"): | |
| 41 | + latest = await fetch_one(conn, "select * from snapshots where id = :id", id=row["last_snapshot_id"]) | |
| 42 | + if latest is None: | |
| 43 | + latest = await fetch_one(conn, "select * from snapshots where sensor_id = :id order by fetched_at desc limit 1", id=sensor_id) | |
| 44 | + out["latest_snapshot"] = ser.snapshot(latest) if latest else None | |
| 45 | + out["last_meaningful_change_at"] = row.get("last_meaningful_change_at") | |
| 46 | + return out | |
| 47 | + | |
| 48 | + | |
| 49 | +@router.get("/sensors/{sensor_id}/snapshots", summary="Snapshot versions of a sensor") | |
| 50 | +async def sensor_snapshots(sensor_id: str, response: Response, limit: int = Query(50, ge=1, le=500), before: str | None = None) -> dict[str, Any]: | |
| 51 | + response.headers["cache-control"] = public_cache_value(60) | |
| 52 | + before_dt = q.parse_iso(before, "before") | |
| 53 | + async with connection() as conn: | |
| 54 | + await _sensor_or_404(conn, sensor_id) | |
| 55 | + extra = " and fetched_at < :before" if before_dt else "" | |
| 56 | + params: dict[str, Any] = {"id": sensor_id, "limit": limit} | |
| 57 | + if before_dt: | |
| 58 | + params["before"] = before_dt | |
| 59 | + rows = await fetch_all(conn, f"select * from snapshots where sensor_id = :id{extra} order by fetched_at desc limit :limit", **params) | |
| 60 | + return {"items": [ser.snapshot(r) for r in rows]} | |
| 61 | + | |
| 62 | + | |
| 63 | +@router.get("/sensors/{sensor_id}/changes", summary="Changes detected by a sensor") | |
| 64 | +async def sensor_changes(sensor_id: str, response: Response, limit: int = Query(50, ge=1, le=500), | |
| 65 | + min_significance: float | None = Query(None, ge=0, le=1)) -> dict[str, Any]: | |
| 66 | + response.headers["cache-control"] = public_cache_value(60) | |
| 67 | + async with connection() as conn: | |
| 68 | + await _sensor_or_404(conn, sensor_id) | |
| 69 | + extra = " and significance >= :ms" if min_significance is not None else "" | |
| 70 | + params: dict[str, Any] = {"id": sensor_id, "limit": limit} | |
| 71 | + if min_significance is not None: | |
| 72 | + params["ms"] = min_significance | |
| 73 | + rows = await fetch_all(conn, f"select * from changes where sensor_id = :id{extra} order by detected_at desc limit :limit", **params) | |
| 74 | + return {"items": [ser.change(r) for r in rows]} | |
| 75 | + | |
| 76 | + | |
| 77 | +def _load_text(key: str | None, limit: int = MAX_TEXT_BYTES) -> tuple[str | None, bool]: | |
| 78 | + if not key: | |
| 79 | + return None, False | |
| 80 | + try: | |
| 81 | + data = archive.get_bytes(key) | |
| 82 | + except (OSError, ValueError): | |
| 83 | + return None, False | |
| 84 | + truncated = len(data) > limit | |
| 85 | + return data[:limit].decode("utf-8", errors="replace"), truncated | |
| 86 | + | |
| 87 | + | |
| 88 | +def _load_blocks(key: str | None) -> list[dict[str, Any]] | None: | |
| 89 | + if not key: | |
| 90 | + return None | |
| 91 | + try: | |
| 92 | + parsed = json.loads(archive.get_text(key)) | |
| 93 | + except (OSError, ValueError): | |
| 94 | + return None | |
| 95 | + if isinstance(parsed, dict) and isinstance(parsed.get("blocks"), list): | |
| 96 | + parsed = parsed["blocks"] | |
| 97 | + return [b for b in parsed if isinstance(b, dict)][:MAX_BLOCKS] if isinstance(parsed, list) else None | |
| 98 | + | |
| 99 | + | |
| 100 | +@router.get("/snapshots/{snapshot_id}", summary="Snapshot with normalized text, blocks and extracted fields") | |
| 101 | +async def snapshot_detail(snapshot_id: str, response: Response, include: str = Query("text,blocks,extracted")) -> dict[str, Any]: | |
| 102 | + response.headers["cache-control"] = public_cache_value(300) | |
| 103 | + parts = set(q.csv_list(include)) | |
| 104 | + async with connection() as conn: | |
| 105 | + row = await fetch_one(conn, "select * from snapshots where id = :id", id=snapshot_id) | |
| 106 | + if row is None: | |
| 107 | + raise HTTPException(status_code=404, detail="snapshot not found") | |
| 108 | + sensor = await fetch_one(conn, "select id, url, surface, company_id from sensors where id = :id", id=row["sensor_id"]) | |
| 109 | + out = ser.snapshot(row) | |
| 110 | + out["sensor"] = sensor | |
| 111 | + out["object_keys"] = {"raw": row.get("object_key"), "text": row.get("text_key"), "blocks": row.get("blocks_key")} | |
| 112 | + if "text" in parts: | |
| 113 | + text, truncated = _load_text(row.get("text_key")) | |
| 114 | + out["text"], out["text_truncated"] = text, truncated | |
| 115 | + if "blocks" in parts: | |
| 116 | + out["blocks"] = _load_blocks(row.get("blocks_key")) or [] | |
| 117 | + if "extracted" in parts: | |
| 118 | + out["extracted"] = ser._dict(row.get("extracted")) | |
| 119 | + return out | |
| 120 | + | |
| 121 | + | |
| 122 | +def _compute_diff(before_blocks: list[dict[str, Any]], after_blocks: list[dict[str, Any]], *, surface: str, before_text: str = "", | |
| 123 | + after_text: str = "") -> dict[str, Any] | None: | |
| 124 | + try: | |
| 125 | + from companyatlas.sdk.diff import compare # crawl agent's module — optional at run time | |
| 126 | + except ImportError: | |
| 127 | + return None | |
| 128 | + from companyatlas.sdk.models import Block | |
| 129 | + | |
| 130 | + def to_blocks(items: list[dict[str, Any]]) -> list[Block]: | |
| 131 | + out = [] | |
| 132 | + for b in items: | |
| 133 | + try: | |
| 134 | + out.append(Block(**{k: v for k, v in b.items() if k in Block.__dataclass_fields__})) | |
| 135 | + except TypeError: | |
| 136 | + continue | |
| 137 | + return out | |
| 138 | + | |
| 139 | + try: | |
| 140 | + result = compare(to_blocks(before_blocks), to_blocks(after_blocks), surface=surface, before_text=before_text, after_text=after_text) | |
| 141 | + except TypeError: # older/newer signature: positional blocks only | |
| 142 | + result = compare(to_blocks(before_blocks), to_blocks(after_blocks)) | |
| 143 | + if hasattr(result, "to_json"): | |
| 144 | + payload = result.to_json() | |
| 145 | + payload["significance"] = getattr(result, "significance", None) | |
| 146 | + return payload | |
| 147 | + return dict(result) if isinstance(result, dict) else None | |
| 148 | + | |
| 149 | + | |
| 150 | +@router.get("/snapshots/{snapshot_id}/diff/{other_id}", summary="Diff between two snapshots (stored change when available, else computed)") | |
| 151 | +async def snapshot_diff(snapshot_id: str, other_id: str, response: Response) -> dict[str, Any]: | |
| 152 | + response.headers["cache-control"] = public_cache_value(300) | |
| 153 | + async with connection() as conn: | |
| 154 | + a = await fetch_one(conn, "select * from snapshots where id = :id", id=snapshot_id) | |
| 155 | + b = await fetch_one(conn, "select * from snapshots where id = :id", id=other_id) | |
| 156 | + if a is None or b is None: | |
| 157 | + raise HTTPException(status_code=404, detail="snapshot not found") | |
| 158 | + before, after = (a, b) if a["fetched_at"] <= b["fetched_at"] else (b, a) | |
| 159 | + stored = await fetch_one(conn, "select diff, significance, id from changes where snapshot_before = :b and snapshot_after = :a limit 1", | |
| 160 | + b=before["id"], a=after["id"]) | |
| 161 | + out = {"before": ser.snapshot(before), "after": ser.snapshot(after), "diff": None, "source": None} | |
| 162 | + if stored and ser._dict(stored["diff"]): | |
| 163 | + diff = ser._dict(stored["diff"]) | |
| 164 | + diff.setdefault("significance", ser._float(stored["significance"], 3)) | |
| 165 | + out.update(diff=diff, source="change", change_id=stored["id"]) | |
| 166 | + return out | |
| 167 | + before_blocks, after_blocks = _load_blocks(before.get("blocks_key")), _load_blocks(after.get("blocks_key")) | |
| 168 | + if before_blocks is None or after_blocks is None: | |
| 169 | + raise HTTPException(status_code=404, detail="block objects unavailable for one of the snapshots") | |
| 170 | + async with connection() as conn: | |
| 171 | + surface = await fetch_val(conn, "select surface from sensors where id = :id", id=after["sensor_id"]) or "other" | |
| 172 | + before_text = _load_text(before.get("text_key"))[0] or "" | |
| 173 | + after_text = _load_text(after.get("text_key"))[0] or "" | |
| 174 | + try: | |
| 175 | + diff = _compute_diff(before_blocks, after_blocks, surface=surface, before_text=before_text, after_text=after_text) | |
| 176 | + except Exception: | |
| 177 | + log.exception("on-demand diff failed", extra={"before": before["id"], "after": after["id"]}) | |
| 178 | + raise HTTPException(status_code=500, detail="diff computation failed") from None | |
| 179 | + if diff is None: | |
| 180 | + raise HTTPException(status_code=501, detail="diff engine not available (companyatlas.sdk.diff.compare)") | |
| 181 | + out.update(diff=diff, source="computed") | |
| 182 | + return out | |
| 183 | + | |
| 184 | + | |
| 185 | +@router.get("/changes/{change_id}", summary="Change with diff, structured delta and derived events") | |
| 186 | +async def change_detail(change_id: str, response: Response) -> dict[str, Any]: | |
| 187 | + response.headers["cache-control"] = public_cache_value(300) | |
| 188 | + async with connection() as conn: | |
| 189 | + row = await fetch_one(conn, "select * from changes where id = :id", id=change_id) | |
| 190 | + if row is None: | |
| 191 | + raise HTTPException(status_code=404, detail="change not found") | |
| 192 | + out = ser.change(row, with_diff=True) | |
| 193 | + where, params = q.event_filters(status=None) | |
| 194 | + where.append("e.change_id = :chg") | |
| 195 | + params["chg"] = change_id | |
| 196 | + out["events"] = [ser.event(e) for e in await q.fetch_events(conn, where, params, limit=100)] | |
| 197 | + sensor = await fetch_one(conn, "select id, url, surface, connector_id from sensors where id = :id", id=row["sensor_id"]) | |
| 198 | + out["sensor"] = sensor | |
| 199 | + cards = await q.fetch_cards_by_ids(conn, [row["company_id"]]) | |
| 200 | + out["company"] = ser.company_card(cards[0]) if cards else None | |
| 201 | + return out | |
added
src/companyatlas/api/routers/rankings.py
+29 −0
@@ -0,0 +1,29 @@ | ||
| 1 | +"""Rankings over `metrics_current` with window deltas from `metric_series` (pricing_changes counts events in the window).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api.common import cached, cached_response | |
| 10 | +from companyatlas.db import connection | |
| 11 | + | |
| 12 | +ORDER = 20 | |
| 13 | +router = APIRouter(prefix="/api/v1", tags=["rankings"]) | |
| 14 | +KINDS = "|".join(agg.RANKING_KINDS) | |
| 15 | + | |
| 16 | + | |
| 17 | +@router.get("/rankings", summary="Ranked companies by kind and window (cached 120 s)") | |
| 18 | +async def rankings(request: Request, kind: str = Query("most_active", pattern=f"^({KINDS})$"), window: str = Query("7d", pattern="^(24h|7d|30d|90d|1y)$"), | |
| 19 | + country: str | None = Query(None, max_length=2), industry: str | None = Query(None, max_length=80), | |
| 20 | + limit: int = Query(50, ge=1, le=200), sparkline: str | None = None) -> Any: | |
| 21 | + spark = sparkline in ("1", "true", "yes") | |
| 22 | + | |
| 23 | + async def produce() -> dict[str, Any]: | |
| 24 | + async with connection() as conn: | |
| 25 | + items = await agg.ranking_cards(conn, kind, window, country=(country or None), industry=(industry or None), limit=limit, sparkline=spark) | |
| 26 | + return {"kind": kind, "window": window, "country": (country or "").upper() or None, "industry": industry or None, "items": items, | |
| 27 | + "kinds": list(agg.RANKING_KINDS)} | |
| 28 | + key = f"rankings:{kind}:{window}:{(country or '').upper()}:{industry or ''}:{limit}:{int(spark)}" | |
| 29 | + return cached_response(request, await cached(key, 120, produce), 120) | |
added
src/companyatlas/api/routers/search.py
+254 −0
@@ -0,0 +1,254 @@ | ||
| 1 | +"""Search: `/search` (FTS + trigram), `/search/suggest` (fast prefix), `/ask` (deterministic parser, LLM optional).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import inspect | |
| 5 | +import logging | |
| 6 | +import time | |
| 7 | +from dataclasses import dataclass, field | |
| 8 | +from typing import Any | |
| 9 | + | |
| 10 | +from fastapi import APIRouter, Query, Response | |
| 11 | + | |
| 12 | +from companyatlas.api import aggregates as agg | |
| 13 | +from companyatlas.api import ask_fallback | |
| 14 | +from companyatlas.api import queries as q | |
| 15 | +from companyatlas.api import serializers as ser | |
| 16 | +from companyatlas.api.common import cached, public_cache_value | |
| 17 | +from companyatlas.db import connection, fetch_all | |
| 18 | +from companyatlas.ids import normalize_alias, slugify | |
| 19 | +from companyatlas.taxonomy import EventType | |
| 20 | + | |
| 21 | +log = logging.getLogger("companyatlas.api.search") | |
| 22 | +ORDER = 10 | |
| 23 | +router = APIRouter(prefix="/api/v1", tags=["search"]) | |
| 24 | +ALL_TYPES = ("companies", "events", "industries", "countries", "people", "products") | |
| 25 | +SUGGEST_MAX = 10 | |
| 26 | + | |
| 27 | + | |
| 28 | +def _clean(qs: str) -> str: | |
| 29 | + return " ".join(qs.replace("%", " ").replace("_", " ").split())[:200] | |
| 30 | + | |
| 31 | + | |
| 32 | +async def _search_companies(conn: Any, qs: str, limit: int) -> list[dict[str, Any]]: | |
| 33 | + if len(qs) >= 3: | |
| 34 | + rows = await fetch_all(conn, """ | |
| 35 | + select c.id, greatest(coalesce(ts_rank(c.search, websearch_to_tsquery('simple', :q)), 0), similarity(c.display_name, :q), | |
| 36 | + similarity(c.canonical_domain, :q), case when a.company_id is not null then 1.0 else 0 end) as score | |
| 37 | + from companies c left join (select distinct company_id from company_aliases where alias_norm = :alias) a on a.company_id = c.id | |
| 38 | + where c.search @@ websearch_to_tsquery('simple', :q) or c.display_name % :q or c.canonical_domain % :q or a.company_id is not null | |
| 39 | + order by score desc, c.importance desc limit :lim""", q=qs, alias=normalize_alias(qs), lim=limit) | |
| 40 | + else: | |
| 41 | + rows = await fetch_all(conn, "select c.id, 1.0 as score from companies c where c.display_name ilike :p or c.canonical_domain ilike :p " | |
| 42 | + "order by c.importance desc limit :lim", p=qs + "%", lim=limit) | |
| 43 | + cards = await q.fetch_cards_by_ids(conn, [r["id"] for r in rows]) | |
| 44 | + return [ser.company_card(c) for c in cards] | |
| 45 | + | |
| 46 | + | |
| 47 | +async def _search_events(conn: Any, qs: str, limit: int) -> list[dict[str, Any]]: | |
| 48 | + if len(qs) < 2: | |
| 49 | + return [] | |
| 50 | + rows = await fetch_all(conn, f"{q.EVENT_SELECT} where e.status = 'active' and e.search @@ websearch_to_tsquery('english', :q) " | |
| 51 | + "order by ts_rank(e.search, websearch_to_tsquery('english', :q)) desc, e.detected_at desc limit :lim", q=qs, lim=limit) | |
| 52 | + return [ser.event(r) for r in rows] | |
| 53 | + | |
| 54 | + | |
| 55 | +async def _search_people(conn: Any, qs: str, limit: int) -> list[dict[str, Any]]: | |
| 56 | + if len(qs) < 3: | |
| 57 | + return [] | |
| 58 | + rows = await fetch_all(conn, "select p.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 59 | + "c.country as company_country, c.logo_url as company_logo_url from people p join companies c on c.id = p.company_id " | |
| 60 | + "where p.name % :q or p.name ilike :like order by similarity(p.name, :q) desc, p.is_executive desc limit :lim", | |
| 61 | + q=qs, like=f"%{qs}%", lim=limit) | |
| 62 | + out = [] | |
| 63 | + for r in rows: | |
| 64 | + item = ser.person(r) | |
| 65 | + item["company"] = ser.company_ref(r) | |
| 66 | + out.append(item) | |
| 67 | + return out | |
| 68 | + | |
| 69 | + | |
| 70 | +async def _search_products(conn: Any, qs: str, limit: int) -> list[dict[str, Any]]: | |
| 71 | + if len(qs) < 3: | |
| 72 | + return [] | |
| 73 | + rows = await fetch_all(conn, "select p.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 74 | + "c.country as company_country, c.logo_url as company_logo_url from products p join companies c on c.id = p.company_id " | |
| 75 | + "where p.name % :q or p.name ilike :like order by similarity(p.name, :q) desc, p.last_seen_at desc limit :lim", | |
| 76 | + q=qs, like=f"%{qs}%", lim=limit) | |
| 77 | + out = [] | |
| 78 | + for r in rows: | |
| 79 | + item = ser.product(r) | |
| 80 | + item["company"] = ser.company_ref(r) | |
| 81 | + out.append(item) | |
| 82 | + return out | |
| 83 | + | |
| 84 | + | |
| 85 | +def _match_rows(rows: list[dict[str, Any]], qs: str, keys: tuple[str, ...], limit: int) -> list[dict[str, Any]]: | |
| 86 | + ql = qs.lower() | |
| 87 | + starts = [r for r in rows if any(str(r.get(k) or "").lower().startswith(ql) for k in keys)] | |
| 88 | + contains = [r for r in rows if r not in starts and any(ql in str(r.get(k) or "").lower() for k in keys)] | |
| 89 | + return (starts + contains)[:limit] | |
| 90 | + | |
| 91 | + | |
| 92 | +@router.get("/search", summary="Search companies, events, industries, countries, people and products") | |
| 93 | +async def search(response: Response, q_: str = Query(..., alias="q", min_length=1, max_length=200), types: str = Query(",".join(ALL_TYPES)), | |
| 94 | + limit: int = Query(10, ge=1, le=50)) -> dict[str, Any]: | |
| 95 | + response.headers["cache-control"] = public_cache_value(30) | |
| 96 | + t0 = time.perf_counter() | |
| 97 | + qs = _clean(q_) | |
| 98 | + wanted = {t for t in q.csv_list(types) if t in ALL_TYPES} or set(ALL_TYPES) | |
| 99 | + out: dict[str, Any] = {"query": qs, "companies": [], "events": [], "industries": [], "countries": [], "people": [], "products": []} | |
| 100 | + if qs: | |
| 101 | + async with connection() as conn: | |
| 102 | + if "companies" in wanted: | |
| 103 | + out["companies"] = await _search_companies(conn, qs, limit) | |
| 104 | + if "events" in wanted: | |
| 105 | + out["events"] = await _search_events(conn, qs, limit) | |
| 106 | + if "people" in wanted: | |
| 107 | + out["people"] = await _search_people(conn, qs, limit) | |
| 108 | + if "products" in wanted: | |
| 109 | + out["products"] = await _search_products(conn, qs, limit) | |
| 110 | + if "industries" in wanted: | |
| 111 | + out["industries"] = _match_rows(await agg.cached_industry_rows(), qs, ("name", "slug"), limit) | |
| 112 | + if "countries" in wanted: | |
| 113 | + out["countries"] = _match_rows(await agg.cached_country_rows(), qs, ("name", "code"), limit) | |
| 114 | + out["took_ms"] = round((time.perf_counter() - t0) * 1000, 1) | |
| 115 | + return out | |
| 116 | + | |
| 117 | + | |
| 118 | +@router.get("/search/suggest", summary="Typeahead suggestions (≤ 10)") | |
| 119 | +async def suggest(response: Response, q_: str = Query(..., alias="q", min_length=1, max_length=100)) -> dict[str, Any]: | |
| 120 | + response.headers["cache-control"] = public_cache_value(30) | |
| 121 | + qs = _clean(q_) | |
| 122 | + if not qs: | |
| 123 | + return {"items": []} | |
| 124 | + | |
| 125 | + async def produce() -> dict[str, Any]: | |
| 126 | + items: list[dict[str, Any]] = [] | |
| 127 | + async with connection() as conn: | |
| 128 | + rows = await fetch_all(conn, "select c.slug, c.display_name, c.canonical_domain, c.country from companies c " | |
| 129 | + "where c.display_name ilike :p or c.canonical_domain ilike :p " | |
| 130 | + "or c.id in (select company_id from company_aliases where alias_norm like :n) " | |
| 131 | + "order by c.importance desc, c.display_name limit 6", p=qs + "%", n=normalize_alias(qs) + "%") | |
| 132 | + for r in rows: | |
| 133 | + items.append({"kind": "company", "label": r["display_name"], "sublabel": r["canonical_domain"], "href": f"/company/{r['slug']}", "slug": r["slug"], | |
| 134 | + "country": r["country"]}) | |
| 135 | + for r in _match_rows(await agg.cached_industry_rows(), qs, ("name", "slug"), 3): | |
| 136 | + items.append({"kind": "industry", "label": r["name"], "sublabel": f"{r['companies']} companies", "href": f"/industry/{r['slug']}", "slug": r["slug"]}) | |
| 137 | + for r in _match_rows(await agg.cached_country_rows(), qs, ("name", "code"), 3): | |
| 138 | + items.append({"kind": "country", "label": r["name"], "sublabel": f"{r['companies']} companies", "href": f"/country/{r['slug']}", "code": r["code"]}) | |
| 139 | + ql = qs.upper().replace(" ", "_") | |
| 140 | + for t in EventType: | |
| 141 | + if t.value.startswith(ql) or ql in t.value: | |
| 142 | + items.append({"kind": "event_type", "label": t.value.replace("_", " ").title(), "sublabel": "event type", "href": f"/events?event_type={t.value}", | |
| 143 | + "event_type": t.value}) | |
| 144 | + return {"items": items[:SUGGEST_MAX]} | |
| 145 | + return await cached(f"suggest:{qs.lower()}", 30, produce) | |
| 146 | + | |
| 147 | + | |
| 148 | +@dataclass | |
| 149 | +class AskPlan: | |
| 150 | + """Engine-independent query plan for `/ask` (filled from the intelligence parser when present, else from `ask_fallback`).""" | |
| 151 | + window: str = "30d" | |
| 152 | + event_types: list[str] = field(default_factory=list) | |
| 153 | + event_subtypes: list[str] = field(default_factory=list) | |
| 154 | + country: str | None = None | |
| 155 | + industry: str | None = None | |
| 156 | + ai: bool = False | |
| 157 | + ranking_kind: str | None = None | |
| 158 | + company_terms: list[str] = field(default_factory=list) | |
| 159 | + keywords: list[str] = field(default_factory=list) | |
| 160 | + min_importance: float | None = None | |
| 161 | + interpretation: dict[str, Any] = field(default_factory=dict) | |
| 162 | + engine: str = "deterministic" | |
| 163 | + answer_fn: Any = None | |
| 164 | + | |
| 165 | + | |
| 166 | +def _window_from_days(days: int | None) -> str: | |
| 167 | + if not days: | |
| 168 | + return "30d" | |
| 169 | + for name, td in sorted(q.WINDOWS.items(), key=lambda kv: kv[1]): | |
| 170 | + if days <= td.days or (name == "24h" and days <= 1): | |
| 171 | + return name | |
| 172 | + return "1y" | |
| 173 | + | |
| 174 | + | |
| 175 | +async def _plan_with_intelligence(question: str, countries: list[dict[str, Any]], industries: list[dict[str, Any]]) -> AskPlan | None: | |
| 176 | + try: | |
| 177 | + from companyatlas.services.llm import ask as llm_ask # intelligence agent's module — optional | |
| 178 | + except ImportError: | |
| 179 | + return None | |
| 180 | + route = getattr(llm_ask, "route_question", None) | |
| 181 | + if route is None: | |
| 182 | + return None | |
| 183 | + try: | |
| 184 | + it = route(question, industries={i["slug"]: i["name"] for i in industries}, countries={c["code"]: c["name"] for c in countries}) | |
| 185 | + if inspect.isawaitable(it): | |
| 186 | + it = await it | |
| 187 | + except Exception: | |
| 188 | + log.warning("intelligence ask parser failed; using fallback", exc_info=True) | |
| 189 | + return None | |
| 190 | + kind = None | |
| 191 | + intent = getattr(it, "intent", "") | |
| 192 | + if intent == "trend" or getattr(it, "answer_style", "") == "trend": | |
| 193 | + kind = "most_active" | |
| 194 | + plan = AskPlan(window=_window_from_days(getattr(it, "window_days", None)), event_types=list(getattr(it, "event_types", []) or []), | |
| 195 | + event_subtypes=list(getattr(it, "event_subtypes", []) or []), country=(getattr(it, "countries", None) or [None])[0], | |
| 196 | + industry=(getattr(it, "industries", None) or [None])[0], ai="ai" in (getattr(it, "tags", []) or []) or intent == "ai", | |
| 197 | + ranking_kind=kind, company_terms=list(getattr(it, "companies", []) or []), keywords=list(getattr(it, "keywords", []) or []), | |
| 198 | + min_importance=getattr(it, "min_importance", None), interpretation=it.to_dict() if hasattr(it, "to_dict") else {}, | |
| 199 | + engine=getattr(it, "source", "deterministic")) | |
| 200 | + build = getattr(llm_ask, "build_answer", None) | |
| 201 | + if build is not None: | |
| 202 | + plan.answer_fn = lambda companies, events, titles: build(it, companies=companies, events=events) | |
| 203 | + plan.interpretation.setdefault("filters", getattr(it, "filters", {})) | |
| 204 | + return plan | |
| 205 | + | |
| 206 | + | |
| 207 | +def _plan_with_fallback(question: str, countries: list[dict[str, Any]], industries: list[dict[str, Any]]) -> AskPlan: | |
| 208 | + it = ask_fallback.interpret(question, countries=countries, industries=industries) | |
| 209 | + plan = AskPlan(window=it.window, event_types=list(it.event_types), country=it.country, industry=it.industry, ai=it.ai, ranking_kind=it.ranking_kind, | |
| 210 | + company_terms=it.company_terms, keywords=it.terms, interpretation=it.to_json(), engine="deterministic") | |
| 211 | + plan.answer_fn = lambda companies, events, titles: ask_fallback.compose_answer(it, events_total=events, companies_count=companies, sample_titles=titles) | |
| 212 | + return plan | |
| 213 | + | |
| 214 | + | |
| 215 | +@router.get("/ask", summary="Ask Company Atlas (deterministic parser; LLM refinement optional)") | |
| 216 | +async def ask(response: Response, q_: str = Query(..., alias="q", min_length=2, max_length=300), limit: int = Query(10, ge=1, le=50)) -> dict[str, Any]: | |
| 217 | + response.headers["cache-control"] = public_cache_value(30) | |
| 218 | + question = _clean(q_) | |
| 219 | + async with connection() as conn: | |
| 220 | + countries = await fetch_all(conn, "select code, name from countries") | |
| 221 | + industries = await fetch_all(conn, "select slug, name, keywords from industries") | |
| 222 | + plan = await _plan_with_intelligence(question, countries, industries) or _plan_with_fallback(question, countries, industries) | |
| 223 | + company_filter_id: str | None = None | |
| 224 | + company_cards: list[dict[str, Any]] = [] | |
| 225 | + if plan.company_terms: | |
| 226 | + probe = await _search_companies(conn, " ".join(plan.company_terms[:2]), 3) | |
| 227 | + key = slugify(" ".join(plan.company_terms[:2])) | |
| 228 | + if probe and key in (probe[0]["slug"], slugify(probe[0]["display_name"])): | |
| 229 | + company_filter_id = probe[0]["id"] | |
| 230 | + company_cards = probe[:1] | |
| 231 | + subtypes = plan.event_subtypes or (["AI_HIRING", "AI_LAUNCH"] if (plan.ai and not plan.event_types) else None) | |
| 232 | + where, params = q.event_filters(company_id=company_filter_id, event_types=plan.event_types or None, country=plan.country, industry=plan.industry, | |
| 233 | + since=q.window_start(plan.window), event_subtypes=subtypes if not plan.event_types else None, | |
| 234 | + min_importance=plan.min_importance) | |
| 235 | + if plan.ai and plan.event_types: | |
| 236 | + where.append("(e.event_subtype in ('AI_HIRING', 'AI_LAUNCH') or 'ai' = any(e.tags) or e.search @@ websearch_to_tsquery('english', 'ai'))") | |
| 237 | + events = [ser.event(r) for r in await q.fetch_events(conn, where, params, sort="recent", limit=limit * 2)] | |
| 238 | + total = await q.count_events(conn, where, params) | |
| 239 | + if not company_cards: | |
| 240 | + if plan.ranking_kind: | |
| 241 | + company_cards = await agg.ranking_cards(conn, plan.ranking_kind, plan.window, country=plan.country, industry=plan.industry, limit=limit) | |
| 242 | + else: | |
| 243 | + free_text = " ".join(plan.keywords) if plan.keywords and not plan.event_types else None | |
| 244 | + cw, cp = q.company_filters(country=plan.country, industry=plan.industry, status="ACTIVE", has_events=True if not free_text else None, q=free_text) | |
| 245 | + sort = "hiring" if "HIRING" in plan.event_types and not plan.ai else "activity" | |
| 246 | + ids, _t = await q.company_page_ids(conn, cw, cp, sort=sort, limit=limit) | |
| 247 | + company_cards = [ser.company_card(r) for r in await q.fetch_cards_by_ids(conn, ids, sparkline=True)] | |
| 248 | + if plan.ai: | |
| 249 | + company_cards.sort(key=lambda c: -(c["metrics"].get("ai_adoption") or 0)) | |
| 250 | + distinct_companies = len({e["company"]["id"] for e in events}) if events else 0 | |
| 251 | + answer = plan.answer_fn(distinct_companies or len(company_cards), total, [e["title"] for e in events]) | |
| 252 | + sources = list(dict.fromkeys(e["source_url"] for e in events if e.get("source_url")))[:10] | |
| 253 | + return {"interpretation": plan.interpretation, "answer": answer, "companies": company_cards[:limit], "events": events[:limit], "sources": sources, | |
| 254 | + "events_total": total, "engine": plan.engine} | |
added
src/companyatlas/api/routers/signals.py
+66 −0
@@ -0,0 +1,66 @@ | ||
| 1 | +"""Cross-company intelligence: `/signals`, `/trends`, `/map`.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request, Response | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api import serializers as ser | |
| 10 | +from companyatlas.api.common import cached, cached_response, public_cache_value | |
| 11 | +from companyatlas.db import connection, fetch_all | |
| 12 | + | |
| 13 | +ORDER = 20 | |
| 14 | +router = APIRouter(prefix="/api/v1", tags=["signals"]) | |
| 15 | +TREND_WINDOW_DAYS = {"7d": 7, "30d": 30, "90d": 90} | |
| 16 | + | |
| 17 | + | |
| 18 | +@router.get("/signals", summary="Active signals (labelled as signals, never facts)") | |
| 19 | +async def signals(response: Response, kind: str | None = Query(None, max_length=60), scope: str | None = Query(None, pattern="^(company|industry|country|global)$"), | |
| 20 | + scope_key: str | None = Query(None, max_length=80), company: str | None = Query(None, max_length=200), | |
| 21 | + min_strength: float | None = Query(None, ge=0, le=1), limit: int = Query(50, ge=1, le=500)) -> dict[str, Any]: | |
| 22 | + response.headers["cache-control"] = public_cache_value(60) | |
| 23 | + where = ["s.status = 'active'"] | |
| 24 | + params: dict[str, Any] = {"limit": limit} | |
| 25 | + if kind: | |
| 26 | + where.append("s.kind = cast(:kind as text)") | |
| 27 | + params["kind"] = kind | |
| 28 | + if scope: | |
| 29 | + where.append("s.scope = cast(:scope as text)") | |
| 30 | + params["scope"] = scope | |
| 31 | + if scope_key: | |
| 32 | + where.append("s.scope_key = cast(:scope_key as text)") | |
| 33 | + params["scope_key"] = scope_key | |
| 34 | + if company: | |
| 35 | + where.append("(c.slug = cast(:company as text) or c.id = cast(:company as text))") | |
| 36 | + params["company"] = company | |
| 37 | + if min_strength is not None: | |
| 38 | + where.append("s.strength >= :min_strength") | |
| 39 | + params["min_strength"] = min_strength | |
| 40 | + async with connection() as conn: | |
| 41 | + rows = await fetch_all(conn, "select s.*, c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, " | |
| 42 | + "c.country as company_country, c.logo_url as company_logo_url from signals s left join companies c on c.id = s.company_id " | |
| 43 | + f"where {' and '.join(where)} order by s.detected_at desc, s.strength desc limit :limit", **params) | |
| 44 | + items = [] | |
| 45 | + for r in rows: | |
| 46 | + s = ser.signal(r) | |
| 47 | + s["company"] = ser.company_ref(r) if r.get("company_slug") else None | |
| 48 | + items.append(s) | |
| 49 | + return {"items": items} | |
| 50 | + | |
| 51 | + | |
| 52 | +@router.get("/trends", summary="Trending terms with momentum (cached 300 s)") | |
| 53 | +async def trends(request: Request, window: str = Query("7d", pattern="^(7d|30d|90d)$"), limit: int = Query(30, ge=1, le=200)) -> Any: | |
| 54 | + async def produce() -> dict[str, Any]: | |
| 55 | + async with connection() as conn: | |
| 56 | + return {"window": window, "items": await agg.trend_rows(conn, TREND_WINDOW_DAYS[window], limit)} | |
| 57 | + return cached_response(request, await cached(f"trends:{window}:{limit}", 300, produce), 300) | |
| 58 | + | |
| 59 | + | |
| 60 | +@router.get("/map", summary="Clustered map buckets (≤ 600, cached 300 s)") | |
| 61 | +async def map_(request: Request, metric: str = Query("events_30d", pattern="^(events_30d|companies|hiring)$")) -> Any: | |
| 62 | + async def produce() -> dict[str, Any]: | |
| 63 | + async with connection() as conn: | |
| 64 | + buckets = await agg.map_buckets(conn, metric) | |
| 65 | + return {"metric": metric, "buckets": buckets} | |
| 66 | + return cached_response(request, await cached(f"map:{metric}", 300, produce), 300) | |
added
src/companyatlas/api/routers/sitemap.py
+48 −0
@@ -0,0 +1,48 @@ | ||
| 1 | +"""`/sitemap?kind=companies|industries|countries&page=` — only companies worth indexing (spec §84: no thin profiles).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api.common import cached, cached_response | |
| 10 | +from companyatlas.config import settings | |
| 11 | +from companyatlas.db import connection, fetch_all, fetch_val | |
| 12 | + | |
| 13 | +ORDER = 10 | |
| 14 | +router = APIRouter(prefix="/api/v1", tags=["exports"]) | |
| 15 | +PAGE_SIZE = 5000 | |
| 16 | + | |
| 17 | +INDEXED_SQL = """ | |
| 18 | +with ev as (select company_id, count(*) as n from events where status = 'active' group by company_id), | |
| 19 | + se as (select company_id, count(*) as n from sensors where status in ('active', 'failing', 'stale') group by company_id) | |
| 20 | +select c.slug, greatest(c.updated_at, coalesce(c.last_event_at, c.updated_at), coalesce(c.last_observed_at, c.updated_at)) as updated_at | |
| 21 | +from companies c left join ev on ev.company_id = c.id left join se on se.company_id = c.id | |
| 22 | +where c.status = 'ACTIVE' and (c.indexed or (coalesce(ev.n, 0) >= :min_events and coalesce(se.n, 0) >= :min_sensors)) | |
| 23 | +order by c.slug limit :lim offset :off | |
| 24 | +""" | |
| 25 | +INDEXED_COUNT_SQL = """ | |
| 26 | +with ev as (select company_id, count(*) as n from events where status = 'active' group by company_id), | |
| 27 | + se as (select company_id, count(*) as n from sensors where status in ('active', 'failing', 'stale') group by company_id) | |
| 28 | +select count(*) from companies c left join ev on ev.company_id = c.id left join se on se.company_id = c.id | |
| 29 | +where c.status = 'ACTIVE' and (c.indexed or (coalesce(ev.n, 0) >= :min_events and coalesce(se.n, 0) >= :min_sensors)) | |
| 30 | +""" | |
| 31 | + | |
| 32 | + | |
| 33 | +@router.get("/sitemap", summary="Sitemap entries") | |
| 34 | +async def sitemap(request: Request, kind: str = Query("companies", pattern="^(companies|industries|countries)$"), page: int = Query(1, ge=1)) -> Any: | |
| 35 | + async def produce() -> dict[str, Any]: | |
| 36 | + if kind == "industries": | |
| 37 | + rows = await agg.cached_industry_rows() | |
| 38 | + return {"kind": kind, "items": [{"slug": r["slug"], "updated_at": None} for r in rows if r["companies"] > 0], "pages": 1, "page": 1} | |
| 39 | + if kind == "countries": | |
| 40 | + rows = await agg.cached_country_rows() | |
| 41 | + return {"kind": kind, "items": [{"slug": r["slug"], "code": r["code"], "updated_at": None} for r in rows if r["companies"] > 0], "pages": 1, "page": 1} | |
| 42 | + async with connection() as conn: | |
| 43 | + params = {"min_events": settings.seo_min_events, "min_sensors": settings.seo_min_sensors} | |
| 44 | + total = int(await fetch_val(conn, INDEXED_COUNT_SQL, **params) or 0) | |
| 45 | + rows = await fetch_all(conn, INDEXED_SQL, **params, lim=PAGE_SIZE, off=(page - 1) * PAGE_SIZE) | |
| 46 | + return {"kind": kind, "items": [{"slug": r["slug"], "updated_at": r["updated_at"]} for r in rows], "pages": max(1, -(-total // PAGE_SIZE)), "page": page, | |
| 47 | + "total": total} | |
| 48 | + return cached_response(request, await cached(f"sitemap:{kind}:{page}", 600, produce), 600) | |
added
src/companyatlas/api/routers/stats.py
+116 −0
@@ -0,0 +1,116 @@ | ||
| 1 | +"""Platform aggregates: `/stats`, `/stats/history`, `/system`, `/pulse`, `/index`, `/methodology`.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from typing import Any | |
| 5 | + | |
| 6 | +from fastapi import APIRouter, Query, Request | |
| 7 | + | |
| 8 | +from companyatlas.api import aggregates as agg | |
| 9 | +from companyatlas.api import queries as q | |
| 10 | +from companyatlas.api.common import cached, cached_response | |
| 11 | +from companyatlas.config import settings | |
| 12 | +from companyatlas.db import connection | |
| 13 | +from companyatlas.taxonomy import CCI_FORMULA_VERSION, CCI_WEIGHTS, EVENT_SUBTYPES, METRICS_FORMULA_VERSION, EventType, Metric | |
| 14 | + | |
| 15 | +ORDER = 10 | |
| 16 | +router = APIRouter(prefix="/api/v1", tags=["platform"]) | |
| 17 | + | |
| 18 | +LAUNCH_SUBTYPES = ["PRODUCT_LAUNCH", "NEW_PRODUCT", "FEATURE_LAUNCH", "AI_LAUNCH", "API_LAUNCH", "SDK_RELEASE"] | |
| 19 | + | |
| 20 | + | |
| 21 | +@router.get("/stats", summary="Platform counters (cached 60 s)") | |
| 22 | +async def stats(request: Request) -> Any: | |
| 23 | + return cached_response(request, await agg.cached_global_stats(), 60) | |
| 24 | + | |
| 25 | + | |
| 26 | +@router.get("/stats/history", summary="Global daily history") | |
| 27 | +async def stats_history(request: Request, days: int = Query(90, ge=1, le=730)) -> Any: | |
| 28 | + async def produce() -> dict[str, Any]: | |
| 29 | + async with connection() as conn: | |
| 30 | + return {"items": await agg.global_daily_rows(conn, days)} | |
| 31 | + return cached_response(request, await cached(f"stats:history:{days}", 300, produce), 300) | |
| 32 | + | |
| 33 | + | |
| 34 | +@router.get("/system", summary="Public aggregate health") | |
| 35 | +async def system(request: Request) -> Any: | |
| 36 | + async def produce() -> dict[str, Any]: | |
| 37 | + async with connection() as conn: | |
| 38 | + return await agg.system_health(conn) | |
| 39 | + return cached_response(request, await cached("system", 30, produce), 30) | |
| 40 | + | |
| 41 | + | |
| 42 | +@router.get("/index", summary="Global Corporate Activity Index") | |
| 43 | +async def index(request: Request, days: int = Query(365, ge=30, le=1095)) -> Any: | |
| 44 | + async def produce() -> dict[str, Any]: | |
| 45 | + async with connection() as conn: | |
| 46 | + return await agg.activity_index(conn, days) | |
| 47 | + return cached_response(request, await cached(f"index:{days}", 300, produce), 300) | |
| 48 | + | |
| 49 | + | |
| 50 | +@router.get("/pulse", summary="Homepage aggregate (cached 60 s)") | |
| 51 | +async def pulse(request: Request) -> Any: | |
| 52 | + async def produce() -> dict[str, Any]: | |
| 53 | + stats_payload = await agg.cached_global_stats() | |
| 54 | + async with connection() as conn: | |
| 55 | + live = await agg.live_events(conn, 12) | |
| 56 | + movers = await agg.ranking_cards(conn, "most_active", "7d", limit=10, sparkline=True) | |
| 57 | + hiring = await agg.ranking_cards(conn, "hiring_growth", "30d", limit=8) | |
| 58 | + launches = await agg.live_events(conn, 8, event_subtypes=LAUNCH_SUBTYPES) | |
| 59 | + pricing = await agg.live_events(conn, 8, event_type="PRICING") | |
| 60 | + ai = await agg.ranking_cards(conn, "ai_active", "30d", limit=8) | |
| 61 | + trending = await agg.trend_rows(conn, 7, 10) | |
| 62 | + idx = await agg.activity_index(conn, 30) | |
| 63 | + buckets = await agg.map_buckets(conn, "events_30d") | |
| 64 | + industries = (await agg.cached_industry_rows())[:12] | |
| 65 | + countries = (await agg.cached_country_rows())[:12] | |
| 66 | + return {"stats": stats_payload, "live": live, "movers": movers, "hiring": hiring, "launches": launches, "pricing": pricing, "ai": ai, | |
| 67 | + "industries": industries, "countries": countries, "trending": trending, | |
| 68 | + "activity_index": {"value": idx["value"], "delta_7d": idx["delta_7d"], "series": idx["series"][-30:]}, "map": buckets[:300]} | |
| 69 | + return cached_response(request, await cached("pulse", 60, produce), 60) | |
| 70 | + | |
| 71 | + | |
| 72 | +METRIC_DOCS: dict[str, tuple[str, list[str]]] = { | |
| 73 | + Metric.ACTIVITY_SCORE: (("Company Activity Score 0–100: weighted volume and significance of detected changes across monitored surfaces (website, products, " | |
| 74 | + "news, jobs, leadership, documentation, pricing) over rolling windows, normalised by sensor coverage."), | |
| 75 | + ["changes", "meaningful_changes", "events", "surfaces", "sensor_coverage"]), | |
| 76 | + Metric.HIRING_MOMENTUM_7D: ("Hiring momentum over 7 days: net change of publicly listed open positions (%); negative when listings disappear.", | |
| 77 | + ["jobs_open", "jobs_new", "jobs_removed"]), | |
| 78 | + Metric.HIRING_MOMENTUM_30D: ("Hiring momentum over 30 days (%).", ["jobs_open", "jobs_new", "jobs_removed"]), | |
| 79 | + Metric.HIRING_MOMENTUM_90D: ("Hiring momentum over 90 days (%).", ["jobs_open", "jobs_new", "jobs_removed"]), | |
| 80 | + Metric.OPEN_JOBS: ("Number of publicly listed open positions currently observed on monitored career surfaces.", ["jobs_open"]), | |
| 81 | + Metric.AI_ADOPTION: (("AI Adoption Score 0–100 from observable public signals only: AI products, AI job postings, documentation, marketing, " | |
| 82 | + "partnerships, research and leadership roles. Never claims internal use without evidence."), | |
| 83 | + ["ai_jobs", "ai_products", "ai_docs", "ai_news", "ai_leadership"]), | |
| 84 | + Metric.PRODUCT_VELOCITY: ("Product Velocity 0–100: cadence of product launches, updates, renames and removals detected on product surfaces.", | |
| 85 | + ["product_events", "changelog_entries", "feature_launches"]), | |
| 86 | + Metric.GEO_EXPANSION: ("Geographic Expansion 0–100: new locations, new countries and geographic hiring spread.", ["new_locations", "new_countries", "job_countries"]), | |
| 87 | + Metric.DEVELOPER_MOMENTUM: ("Developer Momentum 0–100: API/SDK launches, documentation and changelog activity.", ["developer_events", "doc_changes", "api_changes"]), | |
| 88 | + Metric.COMMUNICATION_ACTIVITY: ("Communication Activity 0–100: newsroom, blog and feed publication cadence.", ["news_items", "blog_posts"]), | |
| 89 | + Metric.PRICING_ACTIVITY: ("Pricing Activity 0–100: detected plan and price changes.", ["pricing_events"]), | |
| 90 | + Metric.LEADERSHIP_ACTIVITY: ("Leadership Activity 0–100: additions, removals and title changes on monitored leadership pages.", ["leadership_events"]), | |
| 91 | + Metric.CORPORATE_CHANGE_INDEX: (f"Corporate Change Index: weighted composite — {', '.join(f'{k} {v:.2f}' for k, v in CCI_WEIGHTS.items())}.", | |
| 92 | + list(CCI_WEIGHTS)), | |
| 93 | + Metric.ANOMALY_SCORE: ("Unusual activity 0–100: z-score of current activity against the company's own baseline (56-day window).", | |
| 94 | + ["baseline_mean", "baseline_stddev", "current"]), | |
| 95 | + Metric.HISTORICAL_COVERAGE: ("Historical Completeness 0–100: sensor uptime, continuity and surface coverage of the record.", | |
| 96 | + ["sensor_uptime", "continuity", "surface_coverage", "failed_periods"]), | |
| 97 | +} | |
| 98 | + | |
| 99 | + | |
| 100 | +@router.get("/methodology", summary="How metrics, significance and confidence are computed") | |
| 101 | +async def methodology(request: Request) -> Any: | |
| 102 | + payload = { | |
| 103 | + "metrics": [{"metric": m.value, "formula_version": CCI_FORMULA_VERSION if m is Metric.CORPORATE_CHANGE_INDEX else METRICS_FORMULA_VERSION, | |
| 104 | + "description": METRIC_DOCS[m][0], "inputs": METRIC_DOCS[m][1]} for m in Metric], | |
| 105 | + "significance_bands": {"noise": [0.0, settings.noise_threshold], "minor": [settings.noise_threshold, settings.meaningful_threshold], | |
| 106 | + "meaningful": [settings.meaningful_threshold, settings.major_threshold], | |
| 107 | + "major": [settings.major_threshold, settings.critical_threshold], "critical": [settings.critical_threshold, 1.0]}, | |
| 108 | + "event_types": [t.value for t in EventType], | |
| 109 | + "event_subtypes": [{"event_subtype": k, "event_type": v[0].value, "default_importance": v[1]} for k, v in EVENT_SUBTYPES.items()], | |
| 110 | + "confidence_labels": [{"label": "VERIFIED", "min_confidence": 0.95}, {"label": "HIGH_CONFIDENCE", "min_confidence": 0.85}, | |
| 111 | + {"label": "LIKELY", "min_confidence": 0.7}, {"label": "INFERRED", "min_confidence": 0.5}, | |
| 112 | + {"label": "LOW_CONFIDENCE", "min_confidence": 0.0}], | |
| 113 | + "windows": list(q.WINDOWS), "cci_weights": dict(CCI_WEIGHTS), "baseline_window_days": settings.baseline_window_days, "anomaly_z": settings.anomaly_z, | |
| 114 | + "language": "Interpretive language is careful by design: events say 'no longer listed', never 'fired'; inferred facts carry a confidence label.", | |
| 115 | + } | |
| 116 | + return cached_response(request, payload, 3600) | |
added
src/companyatlas/api/serializers.py
+290 −0
@@ -0,0 +1,290 @@ | ||
| 1 | +"""Row → API shapes, exactly as documented in docs/API.md. Never invent values: missing inputs stay `null` / omitted.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import json | |
| 5 | +from typing import Any | |
| 6 | + | |
| 7 | +from companyatlas.taxonomy import Metric, confidence_label | |
| 8 | + | |
| 9 | +INT_METRICS = {Metric.OPEN_JOBS.value} | |
| 10 | +PCT_METRICS = {Metric.HIRING_MOMENTUM_7D.value, Metric.HIRING_MOMENTUM_30D.value, Metric.HIRING_MOMENTUM_90D.value} | |
| 11 | + | |
| 12 | + | |
| 13 | +def _json(value: Any) -> Any: | |
| 14 | + if isinstance(value, str): | |
| 15 | + try: | |
| 16 | + return json.loads(value) | |
| 17 | + except ValueError: | |
| 18 | + return value | |
| 19 | + return value | |
| 20 | + | |
| 21 | + | |
| 22 | +def _dict(value: Any) -> dict[str, Any]: | |
| 23 | + v = _json(value) | |
| 24 | + return v if isinstance(v, dict) else {} | |
| 25 | + | |
| 26 | + | |
| 27 | +def _list(value: Any) -> list[Any]: | |
| 28 | + v = _json(value) | |
| 29 | + return list(v) if isinstance(v, list | tuple) else [] | |
| 30 | + | |
| 31 | + | |
| 32 | +def _float(value: Any, nd: int = 1) -> float | None: | |
| 33 | + if value is None: | |
| 34 | + return None | |
| 35 | + try: | |
| 36 | + return round(float(value), nd) | |
| 37 | + except (TypeError, ValueError): | |
| 38 | + return None | |
| 39 | + | |
| 40 | + | |
| 41 | +def _int(value: Any) -> int | None: | |
| 42 | + if value is None: | |
| 43 | + return None | |
| 44 | + try: | |
| 45 | + return int(value) | |
| 46 | + except (TypeError, ValueError): | |
| 47 | + return None | |
| 48 | + | |
| 49 | + | |
| 50 | +def metric_value(metric: str, value: Any) -> float | int | None: | |
| 51 | + if value is None: | |
| 52 | + return None | |
| 53 | + if metric in INT_METRICS: | |
| 54 | + return _int(value) | |
| 55 | + return _float(value, 1) | |
| 56 | + | |
| 57 | + | |
| 58 | +def metrics_map(raw: Any) -> dict[str, float | int]: | |
| 59 | + out: dict[str, float | int] = {} | |
| 60 | + for k, v in _dict(raw).items(): | |
| 61 | + mv = metric_value(k, v) | |
| 62 | + if mv is not None: | |
| 63 | + out[k] = mv | |
| 64 | + return out | |
| 65 | + | |
| 66 | + | |
| 67 | +# ------------------------------------------------------------------------------------------------ companies | |
| 68 | + | |
| 69 | + | |
| 70 | +def company_ref(row: dict[str, Any], prefix: str = "company_") -> dict[str, Any]: | |
| 71 | + return {"id": row.get("company_id") or row.get("id"), "slug": row.get(f"{prefix}slug"), "display_name": row.get(f"{prefix}display_name"), | |
| 72 | + "canonical_domain": row.get(f"{prefix}domain") or row.get(f"{prefix}canonical_domain"), "country": row.get(f"{prefix}country"), | |
| 73 | + "logo_url": row.get(f"{prefix}logo_url")} | |
| 74 | + | |
| 75 | + | |
| 76 | +def company_ref_from_company(c: dict[str, Any]) -> dict[str, Any]: | |
| 77 | + return {"id": c["id"], "slug": c["slug"], "display_name": c["display_name"], "canonical_domain": c["canonical_domain"], | |
| 78 | + "country": c.get("country"), "logo_url": c.get("logo_url")} | |
| 79 | + | |
| 80 | + | |
| 81 | +def company_card(row: dict[str, Any]) -> dict[str, Any]: | |
| 82 | + stats = _dict(row.get("stats")) | |
| 83 | + | |
| 84 | + def count(key: str, fallback: Any) -> int: | |
| 85 | + v = stats.get(key) | |
| 86 | + if isinstance(v, int | float) and not isinstance(v, bool): | |
| 87 | + return int(v) | |
| 88 | + return int(fallback or 0) | |
| 89 | + | |
| 90 | + card: dict[str, Any] = { | |
| 91 | + "id": row["id"], "slug": row["slug"], "display_name": row["display_name"], "legal_name": row.get("legal_name"), | |
| 92 | + "canonical_domain": row["canonical_domain"], "website": row["website"], "description": row.get("description"), | |
| 93 | + "industries": list(row.get("industries") or []), "industry_primary": row.get("industry_primary"), "country": row.get("country"), | |
| 94 | + "hq_city": row.get("hq_city"), "hq_region": row.get("hq_region"), "public_company": bool(row.get("public_company")), | |
| 95 | + "ticker": row.get("ticker"), "exchange": row.get("exchange"), "founded_year": row.get("founded_year"), | |
| 96 | + "employees_band": row.get("employees_band"), "logo_url": row.get("logo_url"), "status": row.get("status"), | |
| 97 | + "onboarding_status": row.get("onboarding_status"), "importance": _float(row.get("importance"), 3) or 0.0, "tier": int(row.get("tier") or 4), | |
| 98 | + "metrics": metrics_map(row.get("metrics")), | |
| 99 | + "counts": {"sensors": count("sensors", row.get("sensors")), "observations": count("observations", row.get("observations")), | |
| 100 | + "changes": count("changes", row.get("changes")), "events": count("events", row.get("events")), | |
| 101 | + "jobs_open": count("jobs_open", row.get("jobs_open"))}, | |
| 102 | + "last_event_at": row.get("last_event_at"), "last_observed_at": row.get("last_observed_at"), | |
| 103 | + } | |
| 104 | + if "sparkline" in row: | |
| 105 | + card["sparkline"] = [round(float(x), 1) for x in (row.get("sparkline") or [])] | |
| 106 | + return card | |
| 107 | + | |
| 108 | + | |
| 109 | +# ------------------------------------------------------------------------------------------------ events | |
| 110 | + | |
| 111 | + | |
| 112 | +def event(row: dict[str, Any], *, with_company: bool = True) -> dict[str, Any]: | |
| 113 | + conf = float(row.get("confidence") or 0) | |
| 114 | + out: dict[str, Any] = { | |
| 115 | + "id": row["id"], "event_type": row["event_type"], "event_subtype": row["event_subtype"], "importance": _float(row.get("importance"), 3), | |
| 116 | + "confidence": round(conf, 3), "confidence_label": row.get("confidence_label") or confidence_label(conf), "title": row["title"], | |
| 117 | + "summary": row.get("summary"), "old_value": row.get("old_value"), "new_value": row.get("new_value"), "payload": _dict(row.get("payload")), | |
| 118 | + "entities": _dict(row.get("entities")), "tags": list(row.get("tags") or []), "detected_at": row["detected_at"], | |
| 119 | + "effective_at": row.get("effective_at"), "published_at": row.get("published_at"), "source_url": row.get("source_url"), | |
| 120 | + "surface": row.get("surface"), "sensor_id": row.get("sensor_id"), "change_id": row.get("change_id"), "cluster_id": row.get("cluster_id"), | |
| 121 | + "origin": row.get("origin") or "deterministic", "model_name": row.get("model_name"), "prompt_version": row.get("prompt_version"), | |
| 122 | + "status": row.get("status") or "active", | |
| 123 | + } | |
| 124 | + if with_company: | |
| 125 | + out["company"] = company_ref(row) | |
| 126 | + if row.get("status") == "retracted": | |
| 127 | + out["retracted_reason"] = row.get("retracted_reason") | |
| 128 | + return out | |
| 129 | + | |
| 130 | + | |
| 131 | +def event_source(row: dict[str, Any]) -> dict[str, Any]: | |
| 132 | + return {"source_url": row["source_url"], "surface": row.get("surface"), "detected_at": row["detected_at"], "kind": row.get("kind") or "primary", | |
| 133 | + "sensor_id": row.get("sensor_id"), "snapshot_id": row.get("snapshot_id")} | |
| 134 | + | |
| 135 | + | |
| 136 | +# ------------------------------------------------------------------------------------------------ provenance | |
| 137 | + | |
| 138 | + | |
| 139 | +def sensor(row: dict[str, Any]) -> dict[str, Any]: | |
| 140 | + return {"id": row["id"], "company_id": row["company_id"], "surface": row["surface"], "connector_id": row["connector_id"], "url": row["url"], | |
| 141 | + "canonical_url": row["canonical_url"], "domain": row["domain"], "status": row["status"], "tier": (row.get("tier") or "D").strip(), | |
| 142 | + "quality_score": _float(row.get("quality_score"), 1), "discovery_confidence": _float(row.get("discovery_confidence"), 3), | |
| 143 | + "discovery_method": row.get("discovery_method"), "current_interval_s": int(row.get("current_interval_s") or 0), | |
| 144 | + "next_run_at": row.get("next_run_at"), "last_run_at": row.get("last_run_at"), "last_success_at": row.get("last_success_at"), | |
| 145 | + "last_change_at": row.get("last_change_at"), "last_status": row.get("last_status"), "last_failure_class": row.get("last_failure_class"), | |
| 146 | + "consecutive_failures": int(row.get("consecutive_failures") or 0), "observation_count": int(row.get("observation_count") or 0), | |
| 147 | + "snapshot_count": int(row.get("snapshot_count") or 0), "change_count": int(row.get("change_count") or 0), | |
| 148 | + "meaningful_change_count": int(row.get("meaningful_change_count") or 0), "event_count": int(row.get("event_count") or 0), | |
| 149 | + "created_at": row.get("created_at")} | |
| 150 | + | |
| 151 | + | |
| 152 | +def sensor_admin(row: dict[str, Any]) -> dict[str, Any]: | |
| 153 | + out = sensor(row) | |
| 154 | + out.update({"base_interval_s": row.get("base_interval_s"), "last_error": row.get("last_error"), "priority": _float(row.get("priority"), 3), | |
| 155 | + "claimed_by": row.get("claimed_by"), "claimed_at": row.get("claimed_at"), "retired_at": row.get("retired_at"), | |
| 156 | + "consecutive_unchanged": row.get("consecutive_unchanged"), "config": _dict(row.get("config")), "updated_at": row.get("updated_at")}) | |
| 157 | + return out | |
| 158 | + | |
| 159 | + | |
| 160 | +def snapshot(row: dict[str, Any]) -> dict[str, Any]: | |
| 161 | + return {"id": row["id"], "sensor_id": row["sensor_id"], "company_id": row.get("company_id"), "version_no": int(row.get("version_no") or 1), | |
| 162 | + "fetched_at": row["fetched_at"], "title": row.get("title"), "language": row.get("language"), "text_length": row.get("text_length"), | |
| 163 | + "block_count": row.get("block_count"), "extracted_summary": _dict(row.get("extracted_summary")), "content_hash": row["content_hash"], | |
| 164 | + "previous_snapshot_id": row.get("previous_snapshot_id"), "observation_id": row.get("observation_id"), | |
| 165 | + "collection_method": row.get("collection_method"), "connector_version": row.get("connector_version")} | |
| 166 | + | |
| 167 | + | |
| 168 | +def change(row: dict[str, Any], *, with_diff: bool = False) -> dict[str, Any]: | |
| 169 | + out: dict[str, Any] = {"id": row["id"], "sensor_id": row["sensor_id"], "surface": row["surface"], "company_id": row["company_id"], | |
| 170 | + "detected_at": row["detected_at"], "significance": _float(row.get("significance"), 3), "kind": row["kind"], | |
| 171 | + "blocks_added": int(row.get("blocks_added") or 0), "blocks_removed": int(row.get("blocks_removed") or 0), | |
| 172 | + "blocks_modified": int(row.get("blocks_modified") or 0), "blocks_moved": int(row.get("blocks_moved") or 0), | |
| 173 | + "text_delta_ratio": _float(row.get("text_delta_ratio"), 4), "similarity": _float(row.get("similarity"), 4), | |
| 174 | + "snapshot_before": row.get("snapshot_before"), "snapshot_after": row["snapshot_after"], "status": row.get("status"), | |
| 175 | + "diff_version": row.get("diff_version")} | |
| 176 | + if with_diff: | |
| 177 | + out["diff"] = _dict(row.get("diff")) | |
| 178 | + out["structured_delta"] = _dict(row.get("structured_delta")) | |
| 179 | + return out | |
| 180 | + | |
| 181 | + | |
| 182 | +# ------------------------------------------------------------------------------------------------ entities | |
| 183 | + | |
| 184 | + | |
| 185 | +def job(row: dict[str, Any]) -> dict[str, Any]: | |
| 186 | + return {"id": row["id"], "title": row["title"], "department": row.get("department"), "location_text": row.get("location_text"), | |
| 187 | + "city": row.get("city"), "country": row.get("country"), "remote": row.get("remote"), "employment_type": row.get("employment_type"), | |
| 188 | + "seniority": row.get("seniority"), "url": row.get("url"), "posted_at": row.get("posted_at"), "first_seen_at": row["first_seen_at"], | |
| 189 | + "last_seen_at": row["last_seen_at"], "removed_at": row.get("removed_at"), "status": row.get("status") or "open", | |
| 190 | + "is_ai": bool(row.get("is_ai"))} | |
| 191 | + | |
| 192 | + | |
| 193 | +def person(row: dict[str, Any]) -> dict[str, Any]: | |
| 194 | + return {"id": row["id"], "name": row["name"], "title": row.get("title"), "role_category": row.get("role_category"), | |
| 195 | + "is_executive": bool(row.get("is_executive")), "first_seen_at": row["first_seen_at"], "last_seen_at": row["last_seen_at"], | |
| 196 | + "removed_at": row.get("removed_at"), "status": row.get("status") or "listed", "source_url": row.get("source_url")} | |
| 197 | + | |
| 198 | + | |
| 199 | +def product(row: dict[str, Any]) -> dict[str, Any]: | |
| 200 | + return {"id": row["id"], "name": row["name"], "category": row.get("category"), "description": row.get("description"), "url": row.get("url"), | |
| 201 | + "first_seen_at": row["first_seen_at"], "last_seen_at": row["last_seen_at"], "removed_at": row.get("removed_at"), | |
| 202 | + "status": row.get("status") or "listed"} | |
| 203 | + | |
| 204 | + | |
| 205 | +def plan(row: dict[str, Any]) -> dict[str, Any]: | |
| 206 | + return {"id": row["id"], "plan_name": row["plan_name"], "price": _float(row.get("price"), 2), "price_text": row.get("price_text"), | |
| 207 | + "currency": row.get("currency"), "billing_period": row.get("billing_period"), "unit": row.get("unit"), | |
| 208 | + "features": [str(x) for x in _list(row.get("features"))], "contact_sales": bool(row.get("contact_sales")), | |
| 209 | + "version_no": int(row.get("version_no") or 1), "valid_from": row["valid_from"], "valid_to": row.get("valid_to"), | |
| 210 | + "status": row.get("status") or "current", "source_url": row.get("source_url")} | |
| 211 | + | |
| 212 | + | |
| 213 | +def location(row: dict[str, Any]) -> dict[str, Any]: | |
| 214 | + return {"id": row["id"], "kind": row.get("kind") or "office", "name": row.get("name"), "city": row.get("city"), "region": row.get("region"), | |
| 215 | + "country": row.get("country"), "lat": row.get("lat"), "lon": row.get("lon"), "first_seen_at": row["first_seen_at"], | |
| 216 | + "last_seen_at": row["last_seen_at"], "removed_at": row.get("removed_at"), "status": row.get("status") or "listed", | |
| 217 | + "source_url": row.get("source_url")} | |
| 218 | + | |
| 219 | + | |
| 220 | +def news_item(row: dict[str, Any]) -> dict[str, Any]: | |
| 221 | + return {"id": row["id"], "title": row["title"], "url": row["url"], "summary": row.get("summary"), "category": row.get("category"), | |
| 222 | + "published_at": row.get("published_at"), "first_seen_at": row["first_seen_at"], "language": row.get("language")} | |
| 223 | + | |
| 224 | + | |
| 225 | +def metric_point(row: dict[str, Any]) -> dict[str, Any]: | |
| 226 | + return {"day": row["day"], "value": _float(row.get("value"), 2), "confidence": _float(row.get("confidence"), 3)} | |
| 227 | + | |
| 228 | + | |
| 229 | +def signal(row: dict[str, Any]) -> dict[str, Any]: | |
| 230 | + return {"id": row["id"], "company_id": row.get("company_id"), "scope": row.get("scope") or "company", "scope_key": row.get("scope_key"), | |
| 231 | + "kind": row["kind"], "strength": _float(row.get("strength"), 3), "confidence": _float(row.get("confidence"), 3), "title": row["title"], | |
| 232 | + "explanation": row.get("explanation"), "evidence": _dict(row.get("evidence")), "window_days": int(row.get("window_days") or 30), | |
| 233 | + "detected_at": row["detected_at"], "expires_at": row.get("expires_at"), "status": row.get("status") or "active"} | |
| 234 | + | |
| 235 | + | |
| 236 | +# ------------------------------------------------------------------------------------------------ owner / admin | |
| 237 | + | |
| 238 | + | |
| 239 | +def alert(row: dict[str, Any]) -> dict[str, Any]: | |
| 240 | + out = {"id": row["id"], "name": row["name"], "company_id": row.get("company_id"), "condition": _dict(row.get("condition")), | |
| 241 | + "channel": row.get("channel") or "web", "target": row.get("target"), "enabled": bool(row.get("enabled", True)), | |
| 242 | + "created_at": row["created_at"], "last_fired_at": row.get("last_fired_at")} | |
| 243 | + if row.get("company_slug"): | |
| 244 | + out["company"] = company_ref(row) | |
| 245 | + return out | |
| 246 | + | |
| 247 | + | |
| 248 | +def alert_delivery(row: dict[str, Any]) -> dict[str, Any]: | |
| 249 | + return {"id": row["id"], "alert_id": row["alert_id"], "alert_name": row.get("alert_name"), "event_id": row.get("event_id"), | |
| 250 | + "event_title": row.get("event_title"), "delivered_at": row["delivered_at"], "channel": row["channel"], "status": row["status"], | |
| 251 | + "detail": row.get("detail")} | |
| 252 | + | |
| 253 | + | |
| 254 | +def queue_job(row: dict[str, Any]) -> dict[str, Any]: | |
| 255 | + return {"id": row["id"], "kind": row["kind"], "key": row["key"], "payload": _dict(row.get("payload")), "priority": _float(row.get("priority"), 3), | |
| 256 | + "run_at": row["run_at"], "locked_at": row.get("locked_at"), "locked_by": row.get("locked_by"), "attempts": row.get("attempts"), | |
| 257 | + "max_attempts": row.get("max_attempts"), "status": row["status"], "last_error": row.get("last_error"), "created_at": row["created_at"], | |
| 258 | + "finished_at": row.get("finished_at")} | |
| 259 | + | |
| 260 | + | |
| 261 | +def llm_job(row: dict[str, Any]) -> dict[str, Any]: | |
| 262 | + return {"id": row["id"], "kind": row["kind"], "ref_id": row["ref_id"], "company_id": row.get("company_id"), "model": row.get("model"), | |
| 263 | + "prompt_version": row.get("prompt_version"), "status": row["status"], "attempts": row.get("attempts"), | |
| 264 | + "request_tokens": row.get("request_tokens"), "response_tokens": row.get("response_tokens"), "latency_ms": row.get("latency_ms"), | |
| 265 | + "result": _json(row.get("result")), "error": row.get("error"), "created_at": row["created_at"], "started_at": row.get("started_at"), | |
| 266 | + "finished_at": row.get("finished_at")} | |
| 267 | + | |
| 268 | + | |
| 269 | +def failure(row: dict[str, Any]) -> dict[str, Any]: | |
| 270 | + return {"id": row["id"], "sensor_id": row.get("sensor_id"), "company_id": row.get("company_id"), "at": row["at"], | |
| 271 | + "failure_class": row["failure_class"], "status_code": row.get("status_code"), "message": row.get("message"), "url": row.get("url"), | |
| 272 | + "company_slug": row.get("company_slug"), "surface": row.get("surface")} | |
| 273 | + | |
| 274 | + | |
| 275 | +def review(row: dict[str, Any]) -> dict[str, Any]: | |
| 276 | + return {"id": row["id"], "kind": row["kind"], "ref_id": row.get("ref_id"), "company_id": row.get("company_id"), "payload": _dict(row.get("payload")), | |
| 277 | + "status": row["status"], "resolution": row.get("resolution"), "created_at": row["created_at"], "resolved_at": row.get("resolved_at"), | |
| 278 | + "company_slug": row.get("company_slug"), "company_display_name": row.get("company_display_name")} | |
| 279 | + | |
| 280 | + | |
| 281 | +def connector(row: dict[str, Any]) -> dict[str, Any]: | |
| 282 | + return {"id": row["id"], "name": row["name"], "version": row["version"], "category": row["category"], "fetch_mode": row.get("fetch_mode"), | |
| 283 | + "enabled": bool(row.get("enabled")), "default_interval_s": row.get("default_interval_s"), | |
| 284 | + "supports_discovery": bool(row.get("supports_discovery")), "supports_incremental": bool(row.get("supports_incremental")), | |
| 285 | + "stats": _dict(row.get("stats")), "created_at": row.get("created_at"), "updated_at": row.get("updated_at")} | |
| 286 | + | |
| 287 | + | |
| 288 | +__all__ = ["alert", "alert_delivery", "change", "company_card", "company_ref", "company_ref_from_company", "connector", "event", "event_source", | |
| 289 | + "failure", "job", "llm_job", "location", "metric_point", "metric_value", "metrics_map", "news_item", "person", "plan", "product", | |
| 290 | + "queue_job", "review", "sensor", "sensor_admin", "signal", "snapshot"] | |
added
src/companyatlas/api/sse.py
+85 −0
@@ -0,0 +1,85 @@ | ||
| 1 | +"""Server-sent events for the live feed: poll `events` every `POLL_S` with a `(detected_at, id)` cursor, heartbeat every `HEARTBEAT_S`. | |
| 2 | + | |
| 3 | +The first message is always a heartbeat carrying the cursor so clients (and proxies) see bytes immediately; reconnecting clients | |
| 4 | +pass the last cursor back as `?since=`. `max_s` bounds the connection so upstream proxies with idle timeouts can reconnect cleanly. | |
| 5 | +""" | |
| 6 | +from __future__ import annotations | |
| 7 | + | |
| 8 | +import asyncio | |
| 9 | +import logging | |
| 10 | +from collections.abc import AsyncIterator | |
| 11 | +from datetime import datetime | |
| 12 | +from typing import Any | |
| 13 | + | |
| 14 | +import orjson | |
| 15 | +from fastapi import Request | |
| 16 | +from sse_starlette.sse import EventSourceResponse | |
| 17 | + | |
| 18 | +from companyatlas.api import queries as q | |
| 19 | +from companyatlas.api import serializers as ser | |
| 20 | +from companyatlas.db import connection | |
| 21 | + | |
| 22 | +log = logging.getLogger("companyatlas.api.sse") | |
| 23 | + | |
| 24 | +POLL_S = 5.0 | |
| 25 | +HEARTBEAT_S = 20.0 | |
| 26 | +BATCH = 50 | |
| 27 | +MAX_STREAM_S = 3600 | |
| 28 | +SSE_HEADERS = {"cache-control": "no-store", "x-accel-buffering": "no", "connection": "keep-alive"} | |
| 29 | + | |
| 30 | + | |
| 31 | +def _dump(payload: Any) -> str: | |
| 32 | + return orjson.dumps(payload, option=orjson.OPT_UTC_Z, default=str).decode() | |
| 33 | + | |
| 34 | + | |
| 35 | +async def poll_new_events(since: datetime | None, last_id: str | None, *, event_type: str | None = None, min_importance: float | None = None, | |
| 36 | + limit: int = BATCH) -> list[dict[str, Any]]: | |
| 37 | + where, params = q.event_filters(event_type=event_type, min_importance=min_importance) | |
| 38 | + if since is not None: | |
| 39 | + if last_id: | |
| 40 | + where.append("(e.detected_at, e.id) > (:cur_at, :cur_id)") | |
| 41 | + params.update(cur_at=since, cur_id=last_id) | |
| 42 | + else: | |
| 43 | + where.append("e.detected_at > :cur_at") | |
| 44 | + params["cur_at"] = since | |
| 45 | + async with connection() as conn: | |
| 46 | + rows = await q.fetch_events(conn, where, params, sort="recent", limit=limit) | |
| 47 | + return [ser.event(r) for r in reversed(rows)] # oldest first so clients append in order | |
| 48 | + | |
| 49 | + | |
| 50 | +async def live_event_stream(request: Request | None, *, since: datetime | None, event_type: str | None = None, | |
| 51 | + min_importance: float | None = None, max_s: float = MAX_STREAM_S) -> AsyncIterator[dict[str, Any]]: | |
| 52 | + loop = asyncio.get_running_loop() | |
| 53 | + started = loop.time() | |
| 54 | + cursor_at: datetime | None = since or q.now_utc() | |
| 55 | + cursor_id: str | None = None | |
| 56 | + last_sent = loop.time() | |
| 57 | + yield {"event": "heartbeat", "data": _dump({"at": q.now_utc(), "cursor": cursor_at})} | |
| 58 | + while True: | |
| 59 | + if request is not None and await request.is_disconnected(): | |
| 60 | + return | |
| 61 | + try: | |
| 62 | + events = await poll_new_events(cursor_at, cursor_id, event_type=event_type, min_importance=min_importance) | |
| 63 | + except Exception: | |
| 64 | + log.warning("live stream poll failed", exc_info=True) | |
| 65 | + events = [] | |
| 66 | + for ev in events: | |
| 67 | + cursor_at, cursor_id = ev["detected_at"], ev["id"] | |
| 68 | + last_sent = loop.time() | |
| 69 | + yield {"event": "event", "id": ev["id"], "data": _dump(ev)} | |
| 70 | + now = loop.time() | |
| 71 | + if now - last_sent >= HEARTBEAT_S: | |
| 72 | + last_sent = now | |
| 73 | + yield {"event": "heartbeat", "data": _dump({"at": q.now_utc(), "cursor": cursor_at})} | |
| 74 | + remaining = max_s - (now - started) | |
| 75 | + if remaining <= 0: | |
| 76 | + yield {"event": "end", "data": _dump({"cursor": cursor_at, "reason": "max_s reached"})} | |
| 77 | + return | |
| 78 | + await asyncio.sleep(min(POLL_S, remaining)) | |
| 79 | + | |
| 80 | + | |
| 81 | +def sse_response(generator: AsyncIterator[dict[str, Any]]) -> EventSourceResponse: | |
| 82 | + return EventSourceResponse(generator, headers=SSE_HEADERS, ping=HEARTBEAT_S * 3) | |
| 83 | + | |
| 84 | + | |
| 85 | +__all__ = ["HEARTBEAT_S", "MAX_STREAM_S", "POLL_S", "live_event_stream", "poll_new_events", "sse_response"] | |
added
src/companyatlas/commands/api_keys.py
+75 −0
@@ -0,0 +1,75 @@ | ||
| 1 | +"""`catlas api-key create|list|revoke` — API keys for rate-limit tiers (the raw key is printed once; only its sha256 hash is stored).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import hashlib | |
| 5 | +import secrets | |
| 6 | +from datetime import UTC, datetime | |
| 7 | +from typing import Annotated | |
| 8 | + | |
| 9 | +import typer | |
| 10 | +from rich.table import Table | |
| 11 | + | |
| 12 | +from companyatlas.cli import out, run_async | |
| 13 | +from companyatlas.db import execute, fetch_all, fetch_val, transaction | |
| 14 | +from companyatlas.ids import new_id | |
| 15 | + | |
| 16 | +TIERS = ("authenticated", "paid", "internal") | |
| 17 | +keys_app = typer.Typer(name="api-key", help="API keys (X-CA-API-Key) and their rate-limit tiers.", no_args_is_help=True) | |
| 18 | + | |
| 19 | + | |
| 20 | +def _new_raw_key(tier: str) -> str: | |
| 21 | + return f"ca_{tier[:4]}_{secrets.token_urlsafe(32)}" | |
| 22 | + | |
| 23 | + | |
| 24 | +@keys_app.command("create") | |
| 25 | +def create(name: Annotated[str, typer.Argument(help="Human label (customer, service…)")], | |
| 26 | + tier: Annotated[str, typer.Option("--tier", "-t", help="authenticated | paid | internal")] = "authenticated") -> None: | |
| 27 | + """Create a key and print it once.""" | |
| 28 | + if tier not in TIERS: | |
| 29 | + raise typer.BadParameter(f"tier must be one of {', '.join(TIERS)}") | |
| 30 | + raw = _new_raw_key(tier) | |
| 31 | + key_hash = hashlib.sha256(raw.encode("utf-8")).hexdigest() | |
| 32 | + kid = new_id("api_key") | |
| 33 | + | |
| 34 | + async def _insert() -> None: | |
| 35 | + async with transaction() as conn: | |
| 36 | + await execute(conn, "insert into api_keys (id, key_hash, prefix, name, tier) values (:id, :h, :p, :n, :t)", id=kid, h=key_hash, p=raw[:12], n=name, t=tier) | |
| 37 | + | |
| 38 | + run_async(_insert()) | |
| 39 | + out.print(f"[green]created[/] {kid} tier={tier} name={name!r}") | |
| 40 | + out.print("[bold]API key (shown once, store it now):[/]") | |
| 41 | + out.print(raw) | |
| 42 | + | |
| 43 | + | |
| 44 | +@keys_app.command("list") | |
| 45 | +def list_keys(include_revoked: bool = False) -> None: | |
| 46 | + async def _rows(): # type: ignore[no-untyped-def] | |
| 47 | + async with transaction() as conn: | |
| 48 | + extra = "" if include_revoked else " where revoked_at is null" | |
| 49 | + return await fetch_all(conn, f"select id, prefix, name, tier, created_at, last_used_at, request_count, revoked_at from api_keys{extra} order by created_at desc") | |
| 50 | + | |
| 51 | + rows = run_async(_rows()) | |
| 52 | + t = Table("id", "prefix", "name", "tier", "created", "last used", "requests", "revoked") | |
| 53 | + for r in rows: | |
| 54 | + t.add_row(r["id"], r["prefix"] + "…", r["name"], r["tier"], str(r["created_at"])[:19], str(r["last_used_at"] or "")[:19], str(r["request_count"]), | |
| 55 | + str(r["revoked_at"] or "")[:19]) | |
| 56 | + out.print(t) | |
| 57 | + | |
| 58 | + | |
| 59 | +@keys_app.command("revoke") | |
| 60 | +def revoke(key_id: Annotated[str, typer.Argument(help="key id (key_…)")]) -> None: | |
| 61 | + async def _revoke() -> int: | |
| 62 | + async with transaction() as conn: | |
| 63 | + n = await fetch_val(conn, "with u as (update api_keys set revoked_at = :t where id = :id and revoked_at is null returning 1) select count(*) from u", | |
| 64 | + t=datetime.now(UTC), id=key_id) | |
| 65 | + return int(n or 0) | |
| 66 | + | |
| 67 | + n = run_async(_revoke()) | |
| 68 | + out.print("[green]revoked[/]" if n else "[yellow]no active key with that id[/]") | |
| 69 | + | |
| 70 | + | |
| 71 | +def register(app: typer.Typer) -> None: | |
| 72 | + app.add_typer(keys_app, name="api-key") | |
| 73 | + | |
| 74 | + | |
| 75 | +__all__ = ["keys_app", "register"] | |
added
src/companyatlas/commands/intel.py
+252 −0
@@ -0,0 +1,252 @@ | ||
| 1 | +"""`catlas` intelligence commands: process-changes, enrich, metrics, daily, signals, trends, alerts-eval, llm-test, reprocess-events, events, digest, retention.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import asyncio | |
| 5 | +import json | |
| 6 | +import time | |
| 7 | +from datetime import UTC, date, datetime, timedelta | |
| 8 | +from typing import Annotated | |
| 9 | + | |
| 10 | +import typer | |
| 11 | +from rich.table import Table | |
| 12 | + | |
| 13 | +from companyatlas.cli import console, out, run_async | |
| 14 | + | |
| 15 | + | |
| 16 | +def register(app: typer.Typer) -> None: | |
| 17 | + @app.command("process-changes") | |
| 18 | + def process_changes(limit: int = 200, loop: Annotated[bool, typer.Option("--loop", help="keep polling every 20 s")] = False) -> None: | |
| 19 | + """Turn pending changes into deterministic events (clusters, sources, review queue, LLM jobs).""" | |
| 20 | + from companyatlas.services.events import process_pending_changes | |
| 21 | + | |
| 22 | + async def go() -> None: | |
| 23 | + while True: | |
| 24 | + stats = await process_pending_changes(limit=limit) | |
| 25 | + out.print(stats) | |
| 26 | + if not loop: | |
| 27 | + return | |
| 28 | + await asyncio.sleep(20) | |
| 29 | + | |
| 30 | + run_async(go()) | |
| 31 | + | |
| 32 | + @app.command() | |
| 33 | + def enrich(limit: int = 10, once: Annotated[bool, typer.Option("--once", help="one batch only (default: drain until idle)")] = False) -> None: | |
| 34 | + """Run the LLM enrichment worker over pending llm_jobs (sticky by kind, budgeted).""" | |
| 35 | + from companyatlas.services.llm.enrich import run_llm_jobs | |
| 36 | + | |
| 37 | + async def go() -> None: | |
| 38 | + while True: | |
| 39 | + stats = await run_llm_jobs(limit=limit) | |
| 40 | + out.print(stats) | |
| 41 | + if once or not stats["claimed"] or stats.get("skipped_budget"): | |
| 42 | + return | |
| 43 | + | |
| 44 | + run_async(go()) | |
| 45 | + | |
| 46 | + @app.command() | |
| 47 | + def metrics(all: Annotated[bool, typer.Option("--all", help="every company (default: active in the last 90 days)")] = False, | |
| 48 | + company: Annotated[str | None, typer.Option("--company", help="one company slug or id")] = None) -> None: | |
| 49 | + """Compute company metrics (activity, hiring momentum, AI adoption, velocity, CCI…) into metrics_current + metric_series.""" | |
| 50 | + from companyatlas.db import fetch_val, transaction | |
| 51 | + from companyatlas.services.metrics import compute_company_metrics | |
| 52 | + | |
| 53 | + async def go() -> None: | |
| 54 | + ids = None | |
| 55 | + if company: | |
| 56 | + async with transaction() as conn: | |
| 57 | + cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company) | |
| 58 | + if not cid: | |
| 59 | + console.print(f"[red]company not found: {company}[/]") | |
| 60 | + raise typer.Exit(1) | |
| 61 | + ids = [cid] | |
| 62 | + out.print(await compute_company_metrics(ids, all_companies=all)) | |
| 63 | + if company: | |
| 64 | + async with transaction() as conn: | |
| 65 | + from companyatlas.db import fetch_all | |
| 66 | + | |
| 67 | + rows = await fetch_all(conn, "select metric, value, confidence, formula_version, computed_at from metrics_current where company_id = :c order by metric", c=ids[0]) | |
| 68 | + t = Table(title=f"metrics · {company}") | |
| 69 | + for col in ("metric", "value", "confidence", "formula", "computed_at"): | |
| 70 | + t.add_column(col) | |
| 71 | + for r in rows: | |
| 72 | + t.add_row(r["metric"], f"{r['value']:.2f}", f"{r['confidence']:.2f}", r["formula_version"], str(r["computed_at"])[:19]) | |
| 73 | + out.print(t) | |
| 74 | + | |
| 75 | + run_async(go()) | |
| 76 | + | |
| 77 | + @app.command() | |
| 78 | + def daily(day: Annotated[str | None, typer.Option("--day", help="YYYY-MM-DD (UTC) or 'today'")] = None, | |
| 79 | + catch_up: Annotated[bool, typer.Option("--catch-up", help="fill every missing day up to yesterday")] = False, | |
| 80 | + include_today: Annotated[bool, typer.Option("--include-today")] = False) -> None: | |
| 81 | + """Daily aggregates: company_daily, global_daily (activity index, baseline 100), baselines.""" | |
| 82 | + from companyatlas.services.metrics import compute_daily, compute_daily_catch_up | |
| 83 | + | |
| 84 | + async def go() -> None: | |
| 85 | + if catch_up: | |
| 86 | + results = await compute_daily_catch_up(include_today=include_today) | |
| 87 | + out.print({"days": len(results), "last": results[-1] if results else None}) | |
| 88 | + return | |
| 89 | + d = datetime.now(UTC).date() if day in (None, "today") else date.fromisoformat(day) | |
| 90 | + out.print(await compute_daily(d)) | |
| 91 | + | |
| 92 | + run_async(go()) | |
| 93 | + | |
| 94 | + @app.command() | |
| 95 | + def signals(company: Annotated[str | None, typer.Option("--company")] = None) -> None: | |
| 96 | + """Detect company / industry / country signals (hiring surge, launch build-up, expansion, AI acceleration…).""" | |
| 97 | + from companyatlas.db import fetch_all, fetch_val, transaction | |
| 98 | + from companyatlas.services.signals import compute_signals | |
| 99 | + | |
| 100 | + async def go() -> None: | |
| 101 | + ids = None | |
| 102 | + if company: | |
| 103 | + async with transaction() as conn: | |
| 104 | + cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company) | |
| 105 | + ids = [cid] if cid else [] | |
| 106 | + out.print(await compute_signals(ids)) | |
| 107 | + async with transaction() as conn: | |
| 108 | + rows = await fetch_all(conn, """select coalesce(co.slug, s.scope || ':' || s.scope_key) as who, s.kind, s.strength, s.confidence, s.title | |
| 109 | + from signals s left join companies co on co.id = s.company_id where s.status = 'active' order by s.strength desc limit 30""") | |
| 110 | + t = Table(title="active signals") | |
| 111 | + for col in ("scope", "kind", "strength", "confidence", "title"): | |
| 112 | + t.add_column(col) | |
| 113 | + for r in rows: | |
| 114 | + t.add_row(r["who"], r["kind"], f"{r['strength']:.2f}", f"{r['confidence']:.2f}", r["title"]) | |
| 115 | + out.print(t) | |
| 116 | + | |
| 117 | + run_async(go()) | |
| 118 | + | |
| 119 | + @app.command() | |
| 120 | + def trends(days: int = 1, window: int = 7) -> None: | |
| 121 | + """Extract trending terms from event/news titles for the last N days and print momentum.""" | |
| 122 | + from companyatlas.services.trends import compute_trends_range, store_momentum_snapshots, trend_momentum | |
| 123 | + | |
| 124 | + async def go() -> None: | |
| 125 | + for r in await compute_trends_range(days): | |
| 126 | + out.print(r) | |
| 127 | + await store_momentum_snapshots() | |
| 128 | + items = await trend_momentum(window) | |
| 129 | + t = Table(title=f"trend momentum · {window}d") | |
| 130 | + for col in ("term", "mentions", "companies", "momentum"): | |
| 131 | + t.add_column(col) | |
| 132 | + for it in items[:25]: | |
| 133 | + t.add_row(it["term"], str(it["mentions"]), str(it["companies"]), f"{it['momentum']:+.2f}") | |
| 134 | + out.print(t) | |
| 135 | + | |
| 136 | + run_async(go()) | |
| 137 | + | |
| 138 | + @app.command("alerts-eval") | |
| 139 | + def alerts_eval(event: Annotated[list[str] | None, typer.Option("--event", help="event id(s); default = sweep")] = None) -> None: | |
| 140 | + """Evaluate alerts for given events, or run the catch-up sweep (+ metric alerts).""" | |
| 141 | + from companyatlas.services.alerts import evaluate_alerts, sweep | |
| 142 | + | |
| 143 | + run_async(_print(evaluate_alerts(event) if event else sweep())) | |
| 144 | + | |
| 145 | + @app.command("llm-test") | |
| 146 | + def llm_test(prompt: str = "Reply with a JSON object {\"ok\": true, \"model\": \"<your model name>\"}", tier: str = "small", json_mode: bool = True) -> None: | |
| 147 | + """Live round-trip against the configured LLM endpoint: prints health, model, latency and the answer.""" | |
| 148 | + from pydantic import BaseModel | |
| 149 | + | |
| 150 | + from companyatlas.config import settings | |
| 151 | + from companyatlas.services.llm.gateway import LLMError, get_provider | |
| 152 | + | |
| 153 | + class Probe(BaseModel): | |
| 154 | + ok: bool = True | |
| 155 | + model: str | None = None | |
| 156 | + answer: str | None = None | |
| 157 | + | |
| 158 | + async def go() -> None: | |
| 159 | + if not settings.llm_configured: | |
| 160 | + console.print("[red]LLM not configured (CA_LLM_BASE_URL / CA_LLM_ENABLED)[/]") | |
| 161 | + raise typer.Exit(1) | |
| 162 | + p = get_provider() | |
| 163 | + h = await p.health() | |
| 164 | + out.print({"health": h.ok, "base_url": h.base_url, "latency_ms": h.latency_ms, "models": h.models[:20], "error": h.error}) | |
| 165 | + t0 = time.monotonic() | |
| 166 | + try: | |
| 167 | + if json_mode: | |
| 168 | + res = await p.complete_json(tier, "You are a health probe. Answer briefly.", prompt, Probe, max_tokens=120) | |
| 169 | + answer = res.data.model_dump() | |
| 170 | + else: | |
| 171 | + res = await p.complete_text(tier, "You are a health probe. Answer briefly.", prompt, max_tokens=120) | |
| 172 | + answer = res.data | |
| 173 | + except LLMError as exc: | |
| 174 | + console.print(f"[red]LLM error: {exc}[/]") | |
| 175 | + raise typer.Exit(1) from exc | |
| 176 | + out.print({"model": res.model, "latency_ms": res.latency_ms, "wall_ms": int((time.monotonic() - t0) * 1000), "request_tokens": res.request_tokens, | |
| 177 | + "response_tokens": res.response_tokens, "attempts": res.attempts, "repaired": res.repaired, "answer": answer}) | |
| 178 | + await p.close() | |
| 179 | + | |
| 180 | + run_async(go()) | |
| 181 | + | |
| 182 | + @app.command("reprocess-events") | |
| 183 | + def reprocess(since: Annotated[str, typer.Option("--since", help="ISO date/datetime or e.g. 7d")] = "7d", | |
| 184 | + company: Annotated[str | None, typer.Option("--company")] = None, limit: int = 5000) -> None: | |
| 185 | + """Re-run deterministic rules on processed changes (no refetch; dedupe keys keep it idempotent).""" | |
| 186 | + from companyatlas.db import fetch_val, transaction | |
| 187 | + from companyatlas.services.events import reprocess_events | |
| 188 | + | |
| 189 | + async def go() -> None: | |
| 190 | + cid = None | |
| 191 | + if company: | |
| 192 | + async with transaction() as conn: | |
| 193 | + cid = await fetch_val(conn, "select id from companies where slug = :s or id = :s", s=company) | |
| 194 | + out.print(await reprocess_events(_parse_since(since), limit=limit, company_id=cid)) | |
| 195 | + | |
| 196 | + run_async(go()) | |
| 197 | + | |
| 198 | + @app.command() | |
| 199 | + def events(company: Annotated[str | None, typer.Option("--company")] = None, type: Annotated[str | None, typer.Option("--type")] = None, | |
| 200 | + limit: int = 50) -> None: | |
| 201 | + """List recent events.""" | |
| 202 | + from companyatlas.services.events import list_events | |
| 203 | + | |
| 204 | + async def go() -> None: | |
| 205 | + rows = await list_events(company=company, event_type=type, limit=limit) | |
| 206 | + t = Table(title="events") | |
| 207 | + for col in ("detected", "company", "subtype", "imp", "conf", "origin", "status", "title"): | |
| 208 | + t.add_column(col) | |
| 209 | + for r in rows: | |
| 210 | + t.add_row(str(r["detected_at"])[:16], r["slug"], r["event_subtype"], f"{r['importance']:.2f}", r["confidence_label"], r["origin"], r["status"], r["title"][:90]) | |
| 211 | + out.print(t) | |
| 212 | + | |
| 213 | + run_async(go()) | |
| 214 | + | |
| 215 | + @app.command() | |
| 216 | + def digest(company: Annotated[str | None, typer.Option("--company")] = None, scope: Annotated[str | None, typer.Option("--scope", help="industry|country")] = None, | |
| 217 | + key: Annotated[str | None, typer.Option("--key")] = None, days: int = 7) -> None: | |
| 218 | + """Print digest JSON for a company or an industry/country scope.""" | |
| 219 | + from companyatlas.services.digest import company_digest, scope_digest | |
| 220 | + | |
| 221 | + async def go() -> None: | |
| 222 | + if company: | |
| 223 | + data = await company_digest(company, days=days) | |
| 224 | + elif scope and key: | |
| 225 | + data = await scope_digest(scope, key, days=days) | |
| 226 | + else: | |
| 227 | + console.print("[red]--company or --scope + --key required[/]") | |
| 228 | + raise typer.Exit(1) | |
| 229 | + out.print_json(json.dumps(data, default=str)) | |
| 230 | + | |
| 231 | + run_async(go()) | |
| 232 | + | |
| 233 | + @app.command("retention-intel") | |
| 234 | + def retention(dry_run: Annotated[bool, typer.Option("--dry-run")] = True) -> None: | |
| 235 | + """Prune unchanged observations > 90 d, crawl_runs > 30 d, archive finished llm_jobs > 60 d (dry-run by default).""" | |
| 236 | + from companyatlas.services.retention import run_retention | |
| 237 | + | |
| 238 | + run_async(_print(run_retention(dry_run=dry_run))) | |
| 239 | + | |
| 240 | + | |
| 241 | +async def _print(coro) -> None: # type: ignore[no-untyped-def] | |
| 242 | + out.print(await coro) | |
| 243 | + | |
| 244 | + | |
| 245 | +def _parse_since(value: str) -> datetime: | |
| 246 | + v = value.strip().lower() | |
| 247 | + if v.endswith("d") and v[:-1].isdigit(): | |
| 248 | + return datetime.now(UTC) - timedelta(days=int(v[:-1])) | |
| 249 | + if v.endswith("h") and v[:-1].isdigit(): | |
| 250 | + return datetime.now(UTC) - timedelta(hours=int(v[:-1])) | |
| 251 | + dt = datetime.fromisoformat(value) | |
| 252 | + return dt if dt.tzinfo else dt.replace(tzinfo=UTC) | |
modified
src/companyatlas/config.py
+16 −0
@@ -80,6 +80,13 @@ class Settings(BaseSettings): | ||
| 80 | 80 | llm_enabled: bool = Field(True, alias="CA_LLM_ENABLED") |
| 81 | 81 | llm_daily_budget: int = Field(1500, alias="CA_LLM_DAILY_BUDGET") |
| 82 | 82 | llm_min_significance: float = Field(0.40, alias="CA_LLM_MIN_SIGNIFICANCE") |
| 83 | + llm_embedding_model: str = Field("qwen3-embedding-0.6b-8bit", alias="CA_LLM_EMBEDDING_MODEL") | |
| 84 | + llm_max_tries: int = Field(6, alias="CA_LLM_MAX_TRIES") # 429/503 while the server swaps models | |
| 85 | + llm_backoff_initial_s: float = Field(15.0, alias="CA_LLM_BACKOFF_INITIAL_S") | |
| 86 | + llm_backoff_max_s: float = Field(120.0, alias="CA_LLM_BACKOFF_MAX_S") | |
| 87 | + llm_job_max_attempts: int = Field(3, alias="CA_LLM_JOB_MAX_ATTEMPTS") | |
| 88 | + llm_context_block_bytes: int = Field(3072, alias="CA_LLM_CONTEXT_BLOCK_BYTES") # before/after blocks sent to the model (≤ 3 kB each) | |
| 89 | + prompts_dir: Path | None = Field(None, alias="CA_PROMPTS_DIR") # defaults to <repo>/prompts | |
| 83 | 90 | worker_concurrency: int = Field(1, alias="CA_WORKER_CONCURRENCY") |
| 84 | 91 | |
| 85 | 92 | # ---------------------------------------------------------------- metrics / retention |
@@ -88,6 +95,15 @@ class Settings(BaseSettings): | ||
| 88 | 95 | backup_cron: str = Field("35 4 * * *", alias="CA_BACKUP_CRON") |
| 89 | 96 | baseline_window_days: int = Field(56, alias="CA_BASELINE_WINDOW_DAYS") |
| 90 | 97 | anomaly_z: float = Field(2.5, alias="CA_ANOMALY_Z") |
| 98 | + metrics_active_window_days: int = Field(90, alias="CA_METRICS_ACTIVE_WINDOW_DAYS") # hourly pass = companies active within | |
| 99 | + trends_min_companies: int = Field(3, alias="CA_TRENDS_MIN_COMPANIES") | |
| 100 | + signals_window_days: int = Field(30, alias="CA_SIGNALS_WINDOW_DAYS") | |
| 101 | + signals_ttl_days: int = Field(14, alias="CA_SIGNALS_TTL_DAYS") | |
| 102 | + alerts_sweep_window_s: int = Field(1800, alias="CA_ALERTS_SWEEP_WINDOW_S") | |
| 103 | + webhook_timeout_s: float = Field(10.0, alias="CA_WEBHOOK_TIMEOUT_S") | |
| 104 | + retention_observations_days: int = Field(90, alias="CA_RETENTION_OBSERVATIONS_DAYS") | |
| 105 | + retention_crawl_runs_days: int = Field(30, alias="CA_RETENTION_CRAWL_RUNS_DAYS") | |
| 106 | + retention_llm_jobs_days: int = Field(60, alias="CA_RETENTION_LLM_JOBS_DAYS") | |
| 91 | 107 | seo_min_events: int = Field(1, alias="CA_SEO_MIN_EVENTS") |
| 92 | 108 | seo_min_sensors: int = Field(3, alias="CA_SEO_MIN_SENSORS") |
| 93 | 109 | |
added
src/companyatlas/services/alerts.py
+237 −0
@@ -0,0 +1,237 @@ | ||
| 1 | +"""Alert evaluation (spec §140–141): match new events against `alerts.condition` | |
| 2 | +`{event_types, event_subtypes, min_importance, min_confidence, company, industries, countries, tags, metrics: {activity_score: {gt: 80}}}` | |
| 3 | +and record `alert_deliveries`. Channels: `web` (row only, the UI reads it) and `webhook` (POST JSON through the shared `fetch.Fetcher` | |
| 4 | +client after the SSRF guard, signed `X-CompanyAtlas-Signature: sha256=<hmac(body, alert id)>`). Email is stored as queued (no sender yet). | |
| 5 | + | |
| 6 | +`evaluate_alerts(event_ids)` runs right after event creation; `alerts-sweep` catches up every 2 min via a watermark in settings_kv, and | |
| 7 | +evaluates metric-only alerts (no event) at most once per 24 h per alert. | |
| 8 | +""" | |
| 9 | +from __future__ import annotations | |
| 10 | + | |
| 11 | +import hashlib | |
| 12 | +import hmac | |
| 13 | +import json | |
| 14 | +import logging | |
| 15 | +from datetime import UTC, datetime, timedelta | |
| 16 | +from typing import Any | |
| 17 | + | |
| 18 | +from companyatlas.config import settings | |
| 19 | +from companyatlas.db import execute, fetch_all, fetch_val, jsonb, transaction | |
| 20 | +from companyatlas.fetch import BlockedDestination, Fetcher, validate_destination_async | |
| 21 | +from companyatlas.ids import new_id | |
| 22 | +from companyatlas.services.periodic import periodic | |
| 23 | + | |
| 24 | +log = logging.getLogger(__name__) | |
| 25 | + | |
| 26 | +WATERMARK_KEY = "alerts:last_event_created_at" | |
| 27 | +METRIC_ALERT_COOLDOWN_H = 24 | |
| 28 | +_OPS = {"gt": lambda v, t: v > t, "gte": lambda v, t: v >= t, "lt": lambda v, t: v < t, "lte": lambda v, t: v <= t, "eq": lambda v, t: v == t} | |
| 29 | + | |
| 30 | + | |
| 31 | +def _lower_set(values: Any) -> set[str]: | |
| 32 | + if not values: | |
| 33 | + return set() | |
| 34 | + if isinstance(values, str): | |
| 35 | + values = [values] | |
| 36 | + return {str(v).strip().lower() for v in values if str(v).strip()} | |
| 37 | + | |
| 38 | + | |
| 39 | +def metrics_match(cond: dict[str, Any], metrics: dict[str, float]) -> bool: | |
| 40 | + spec = cond.get("metrics") or {} | |
| 41 | + if not isinstance(spec, dict) or not spec: | |
| 42 | + return True | |
| 43 | + for metric, rule in spec.items(): | |
| 44 | + value = metrics.get(str(metric)) | |
| 45 | + if value is None: | |
| 46 | + return False | |
| 47 | + if isinstance(rule, dict): | |
| 48 | + for op, threshold in rule.items(): | |
| 49 | + fn = _OPS.get(str(op)) | |
| 50 | + try: | |
| 51 | + if fn is None or not fn(float(value), float(threshold)): | |
| 52 | + return False | |
| 53 | + except (TypeError, ValueError): | |
| 54 | + return False | |
| 55 | + else: | |
| 56 | + try: | |
| 57 | + if float(value) < float(rule): | |
| 58 | + return False | |
| 59 | + except (TypeError, ValueError): | |
| 60 | + return False | |
| 61 | + return True | |
| 62 | + | |
| 63 | + | |
| 64 | +def event_matches(alert: dict[str, Any], event: dict[str, Any], company: dict[str, Any], metrics: dict[str, float]) -> bool: | |
| 65 | + """Pure matcher. `alert.company_id` pins a company; the condition narrows further.""" | |
| 66 | + cond = alert.get("condition") or {} | |
| 67 | + if alert.get("company_id") and alert["company_id"] != event["company_id"]: | |
| 68 | + return False | |
| 69 | + if cond.get("company") and cond["company"] not in (company.get("slug"), company.get("id")): | |
| 70 | + return False | |
| 71 | + types = _lower_set(cond.get("event_types")) | |
| 72 | + if types and str(event["event_type"]).lower() not in types and str(event["event_subtype"]).lower() not in types: | |
| 73 | + return False | |
| 74 | + subtypes = _lower_set(cond.get("event_subtypes")) | |
| 75 | + if subtypes and str(event["event_subtype"]).lower() not in subtypes: | |
| 76 | + return False | |
| 77 | + if cond.get("min_importance") is not None and float(event.get("importance") or 0) < float(cond["min_importance"]): | |
| 78 | + return False | |
| 79 | + if cond.get("min_confidence") is not None and float(event.get("confidence") or 0) < float(cond["min_confidence"]): | |
| 80 | + return False | |
| 81 | + countries = {c.upper() for c in _lower_set(cond.get("countries"))} | |
| 82 | + if countries and str(company.get("country") or "").upper() not in countries: | |
| 83 | + return False | |
| 84 | + industries = _lower_set(cond.get("industries")) | |
| 85 | + if industries and not (industries & {str(i).lower() for i in (company.get("industries") or [])}): | |
| 86 | + return False | |
| 87 | + tags = _lower_set(cond.get("tags")) | |
| 88 | + if tags and not (tags & {str(t).lower() for t in (event.get("tags") or [])}): | |
| 89 | + return False | |
| 90 | + return metrics_match(cond, metrics) | |
| 91 | + | |
| 92 | + | |
| 93 | +def sign(body: bytes, alert_id: str) -> str: | |
| 94 | + return "sha256=" + hmac.new(alert_id.encode("utf-8"), body, hashlib.sha256).hexdigest() | |
| 95 | + | |
| 96 | + | |
| 97 | +def webhook_payload(alert: dict[str, Any], event: dict[str, Any] | None, company: dict[str, Any] | None, metrics: dict[str, float] | None = None) -> dict[str, Any]: | |
| 98 | + return { | |
| 99 | + "type": "event" if event else "metric", | |
| 100 | + "alert": {"id": alert["id"], "name": alert.get("name"), "condition": alert.get("condition")}, | |
| 101 | + "company": {k: company.get(k) for k in ("id", "slug", "display_name", "canonical_domain", "country")} if company else None, | |
| 102 | + "event": {k: event.get(k) for k in ("id", "event_type", "event_subtype", "importance", "confidence", "confidence_label", "title", "summary", "old_value", | |
| 103 | + "new_value", "detected_at", "source_url", "surface", "origin", "status")} if event else None, | |
| 104 | + "metrics": metrics or {}, | |
| 105 | + "delivered_at": datetime.now(UTC).isoformat(), | |
| 106 | + "docs": f"{settings.site_url}/api", | |
| 107 | + } | |
| 108 | + | |
| 109 | + | |
| 110 | +async def deliver_webhook(alert: dict[str, Any], payload: dict[str, Any]) -> tuple[str, str]: | |
| 111 | + target = str(alert.get("target") or "").strip() | |
| 112 | + if not target: | |
| 113 | + return "failed", "no webhook target" | |
| 114 | + try: | |
| 115 | + await validate_destination_async(target) | |
| 116 | + except BlockedDestination as exc: | |
| 117 | + return "failed", f"blocked destination: {exc}" | |
| 118 | + body = json.dumps(payload, default=str, ensure_ascii=False).encode("utf-8") | |
| 119 | + headers = {"Content-Type": "application/json", "X-CompanyAtlas-Signature": sign(body, alert["id"]), "X-CompanyAtlas-Alert": alert["id"], | |
| 120 | + "User-Agent": settings.user_agent} | |
| 121 | + try: | |
| 122 | + async with Fetcher(timeout_s=settings.webhook_timeout_s, http2=False, max_connections=4) as f: | |
| 123 | + r = await f.client.post(target, content=body, headers=headers) | |
| 124 | + if r.status_code < 300: | |
| 125 | + return "sent", f"HTTP {r.status_code}" | |
| 126 | + return "failed", f"HTTP {r.status_code}" | |
| 127 | + except Exception as exc: # noqa: BLE001 | |
| 128 | + return "failed", f"{exc.__class__.__name__}: {exc}"[:300] | |
| 129 | + | |
| 130 | + | |
| 131 | +async def _record(conn, alert: dict[str, Any], event_id: str | None, status: str, detail: str | None) -> None: # type: ignore[no-untyped-def] | |
| 132 | + await execute(conn, "insert into alert_deliveries (id, alert_id, event_id, channel, status, detail) values (:id, :a, :e, :ch, :st, :d)", | |
| 133 | + id=new_id("alert"), a=alert["id"], e=event_id, ch=alert.get("channel") or "web", st=status, d=detail) | |
| 134 | + await execute(conn, "update alerts set last_fired_at = now() where id = :id", id=alert["id"]) | |
| 135 | + | |
| 136 | + | |
| 137 | +async def evaluate_alerts(event_ids: list[str]) -> dict[str, int]: | |
| 138 | + """Match the given events against all enabled alerts; deliver once per (alert, event).""" | |
| 139 | + stats = {"events": 0, "matched": 0, "delivered": 0, "failed": 0} | |
| 140 | + if not event_ids: | |
| 141 | + return stats | |
| 142 | + async with transaction() as conn: | |
| 143 | + alerts = await fetch_all(conn, "select * from alerts where enabled") | |
| 144 | + if not alerts: | |
| 145 | + return stats | |
| 146 | + events = await fetch_all(conn, """select e.*, co.slug, co.display_name, co.canonical_domain, co.country as company_country, co.industries as company_industries | |
| 147 | + from events e join companies co on co.id = e.company_id where e.id = any(cast(:ids as text[])) and e.status in ('active', 'review')""", | |
| 148 | + ids=event_ids) | |
| 149 | + stats["events"] = len(events) | |
| 150 | + metrics_cache: dict[str, dict[str, float]] = {} | |
| 151 | + pending: list[tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, float]]] = [] | |
| 152 | + for ev in events: | |
| 153 | + company = {"id": ev["company_id"], "slug": ev["slug"], "display_name": ev["display_name"], "canonical_domain": ev["canonical_domain"], | |
| 154 | + "country": ev["company_country"], "industries": ev["company_industries"]} | |
| 155 | + for alert in alerts: | |
| 156 | + cond = alert.get("condition") or {} | |
| 157 | + needs_metrics = bool(cond.get("metrics")) | |
| 158 | + if needs_metrics and ev["company_id"] not in metrics_cache: | |
| 159 | + metrics_cache[ev["company_id"]] = {r["metric"]: float(r["value"]) for r in await fetch_all(conn, "select metric, value from metrics_current where company_id = :c", c=ev["company_id"])} | |
| 160 | + metrics = metrics_cache.get(ev["company_id"], {}) | |
| 161 | + if not event_matches(alert, ev, company, metrics): | |
| 162 | + continue | |
| 163 | + dup = await fetch_val(conn, "select 1 from alert_deliveries where alert_id = :a and event_id = :e", a=alert["id"], e=ev["id"]) | |
| 164 | + if dup: | |
| 165 | + continue | |
| 166 | + stats["matched"] += 1 | |
| 167 | + pending.append((alert, ev, company, metrics)) | |
| 168 | + if (alert.get("channel") or "web") != "webhook": | |
| 169 | + await _record(conn, alert, ev["id"], "delivered" if alert.get("channel") in (None, "web") else "queued", None) | |
| 170 | + stats["delivered"] += 1 | |
| 171 | + for alert, ev, company, metrics in pending: # network outside the transaction | |
| 172 | + if (alert.get("channel") or "web") != "webhook": | |
| 173 | + continue | |
| 174 | + status, detail = await deliver_webhook(alert, webhook_payload(alert, ev, company, metrics)) | |
| 175 | + async with transaction() as conn: | |
| 176 | + await _record(conn, alert, ev["id"], status, detail) | |
| 177 | + stats["delivered" if status == "sent" else "failed"] += 1 | |
| 178 | + return stats | |
| 179 | + | |
| 180 | + | |
| 181 | +async def evaluate_metric_alerts() -> dict[str, int]: | |
| 182 | + """Alerts with only a metrics condition (no event filter) fire on the current metrics, at most once per cooldown.""" | |
| 183 | + stats = {"checked": 0, "fired": 0} | |
| 184 | + fired: list[tuple[dict[str, Any], dict[str, Any], dict[str, float]]] = [] | |
| 185 | + async with transaction() as conn: | |
| 186 | + alerts = await fetch_all(conn, """select * from alerts where enabled and jsonb_exists(condition, 'metrics') and coalesce(condition->'event_types', 'null'::jsonb) in ('null'::jsonb, '[]'::jsonb) | |
| 187 | + and (last_fired_at is null or last_fired_at < :cutoff)""", cutoff=datetime.now(UTC) - timedelta(hours=METRIC_ALERT_COOLDOWN_H)) | |
| 188 | + for alert in alerts: | |
| 189 | + stats["checked"] += 1 | |
| 190 | + cond = alert.get("condition") or {} | |
| 191 | + if alert.get("company_id"): | |
| 192 | + companies = await fetch_all(conn, "select id, slug, display_name, canonical_domain, country, industries from companies where id = :id", id=alert["company_id"]) | |
| 193 | + elif cond.get("company"): | |
| 194 | + companies = await fetch_all(conn, "select id, slug, display_name, canonical_domain, country, industries from companies where slug = :s or id = :s", s=cond["company"]) | |
| 195 | + else: | |
| 196 | + continue # metric alerts must target one company | |
| 197 | + for company in companies: | |
| 198 | + metrics = {r["metric"]: float(r["value"]) for r in await fetch_all(conn, "select metric, value from metrics_current where company_id = :c", c=company["id"])} | |
| 199 | + if not metrics or not metrics_match(cond, metrics): | |
| 200 | + continue | |
| 201 | + detail = json.dumps({k: metrics.get(k) for k in (cond.get("metrics") or {})}) | |
| 202 | + if (alert.get("channel") or "web") == "webhook": | |
| 203 | + fired.append((alert, company, metrics)) | |
| 204 | + else: | |
| 205 | + await _record(conn, alert, None, "delivered", detail) | |
| 206 | + stats["fired"] += 1 | |
| 207 | + for alert, company, metrics in fired: | |
| 208 | + status, detail = await deliver_webhook(alert, webhook_payload(alert, None, company, metrics)) | |
| 209 | + async with transaction() as conn: | |
| 210 | + await _record(conn, alert, None, status, detail) | |
| 211 | + return stats | |
| 212 | + | |
| 213 | + | |
| 214 | +async def sweep(*, window_s: int | None = None) -> dict[str, int]: | |
| 215 | + """Catch-up: evaluate events created since the watermark (bounded by the sweep window).""" | |
| 216 | + window_s = window_s or settings.alerts_sweep_window_s | |
| 217 | + async with transaction() as conn: | |
| 218 | + raw = await fetch_val(conn, "select value from settings_kv where key = :k", k=WATERMARK_KEY) | |
| 219 | + since = datetime.fromisoformat(str(raw).strip('"')) if raw else datetime.now(UTC) - timedelta(seconds=window_s) | |
| 220 | + since = max(since, datetime.now(UTC) - timedelta(seconds=window_s)) | |
| 221 | + rows = await fetch_all(conn, "select id, created_at from events where created_at > :since and status in ('active', 'review') order by created_at limit 2000", since=since) | |
| 222 | + if rows: | |
| 223 | + await execute(conn, "insert into settings_kv (key, value, updated_at) values (:k, cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()", | |
| 224 | + k=WATERMARK_KEY, v=jsonb(rows[-1]["created_at"].isoformat())) | |
| 225 | + stats = await evaluate_alerts([r["id"] for r in rows]) if rows else {"events": 0, "matched": 0, "delivered": 0, "failed": 0} | |
| 226 | + stats.update({f"metric_{k}": v for k, v in (await evaluate_metric_alerts()).items()}) | |
| 227 | + return stats | |
| 228 | + | |
| 229 | + | |
| 230 | +@periodic("alerts-sweep", every_s=120, initial_delay_s=60) | |
| 231 | +async def alerts_sweep_task() -> None: | |
| 232 | + stats = await sweep() | |
| 233 | + if stats.get("matched") or stats.get("metric_fired"): | |
| 234 | + log.info("alerts-sweep", extra=stats) | |
| 235 | + | |
| 236 | + | |
| 237 | +__all__ = ["deliver_webhook", "evaluate_alerts", "evaluate_metric_alerts", "event_matches", "metrics_match", "sign", "sweep", "webhook_payload"] | |
added
src/companyatlas/services/clustering.py
+113 −0
@@ -0,0 +1,113 @@ | ||
| 1 | +"""Event clustering (spec §23): the same corporate event seen on several surfaces (newsroom + feed + homepage, leadership + about…) | |
| 2 | +is folded into one canonical event with aggregated sources. Corroboration raises confidence; later duplicates keep their own row | |
| 3 | +(`status='duplicate'`, `cluster_id`) so provenance is never lost. | |
| 4 | + | |
| 5 | + cluster_key = sha(company_id, event_subtype, normalised entity key, 7-day window bucket) | |
| 6 | + | |
| 7 | +Only `services/events.py` and `services/llm/enrich.py` call `attach_to_cluster`; the API reads `event_clusters` / `events.cluster_id`. | |
| 8 | +""" | |
| 9 | +from __future__ import annotations | |
| 10 | + | |
| 11 | +import re | |
| 12 | +import unicodedata | |
| 13 | +from datetime import UTC, datetime | |
| 14 | +from typing import Any | |
| 15 | + | |
| 16 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 17 | + | |
| 18 | +from companyatlas.db import execute, fetch_one, jsonb | |
| 19 | +from companyatlas.ids import new_id, stable_hash | |
| 20 | + | |
| 21 | +CLUSTER_WINDOW_DAYS = 7 | |
| 22 | +CORROBORATION_BONUS = 0.03 # confidence bump per extra corroborating surface | |
| 23 | +CONFIDENCE_CAP = 0.99 | |
| 24 | + | |
| 25 | +_non_alnum = re.compile(r"[^a-z0-9]+") | |
| 26 | + | |
| 27 | + | |
| 28 | +def normalize_entity_key(value: str) -> str: | |
| 29 | + """Lowercase ASCII, punctuation collapsed — makes "Jane Doe" / "jane doe" / "Jane DOE," the same entity across surfaces.""" | |
| 30 | + s = unicodedata.normalize("NFKD", value or "").encode("ascii", "ignore").decode("ascii").lower() | |
| 31 | + s = _non_alnum.sub(" ", s).strip() | |
| 32 | + return re.sub(r"\s+", " ", s)[:160] | |
| 33 | + | |
| 34 | + | |
| 35 | +def window_bucket(at: datetime, *, days: int = CLUSTER_WINDOW_DAYS) -> int: | |
| 36 | + if at.tzinfo is None: | |
| 37 | + at = at.replace(tzinfo=UTC) | |
| 38 | + return int(at.timestamp() // (days * 86400)) | |
| 39 | + | |
| 40 | + | |
| 41 | +def cluster_key_for(company_id: str, event_subtype: str, entity_key: str, detected_at: datetime) -> str: | |
| 42 | + return stable_hash(company_id, event_subtype, normalize_entity_key(entity_key), str(window_bucket(detected_at)), length=32) | |
| 43 | + | |
| 44 | + | |
| 45 | +async def attach_to_cluster(conn: AsyncConnection, event: dict[str, Any], *, entity_key: str) -> tuple[str, bool]: | |
| 46 | + """Attach a freshly inserted event to its cluster. Returns (cluster_id, is_duplicate). | |
| 47 | + | |
| 48 | + First event of a cluster becomes canonical. A later event from another sensor/surface is marked `duplicate`, the canonical | |
| 49 | + event gains a corroboration source and a confidence bump. Re-processing the same change (same dedupe key) never reaches here. | |
| 50 | + """ | |
| 51 | + key = cluster_key_for(event["company_id"], event["event_subtype"], entity_key, event["detected_at"]) | |
| 52 | + existing = await fetch_one(conn, "select * from event_clusters where cluster_key = :k for update", k=key) | |
| 53 | + if existing is None: | |
| 54 | + cluster_id = new_id("cluster") | |
| 55 | + await execute(conn, """ | |
| 56 | + insert into event_clusters (id, company_id, cluster_key, event_type, event_subtype, title, first_detected_at, last_detected_at, | |
| 57 | + source_count, surfaces, confidence, canonical_event_id) | |
| 58 | + values (:id, :company_id, :key, :event_type, :event_subtype, :title, :at, :at, 1, cast(:surfaces as text[]), :confidence, :event_id) | |
| 59 | + on conflict (cluster_key) do nothing""", | |
| 60 | + id=cluster_id, company_id=event["company_id"], key=key, event_type=event["event_type"], event_subtype=event["event_subtype"], | |
| 61 | + title=event["title"], at=event["detected_at"], surfaces=[event.get("surface") or "other"], confidence=float(event["confidence"]), | |
| 62 | + event_id=event["id"]) | |
| 63 | + await execute(conn, "update events set cluster_id = :c where id = :e", c=cluster_id, e=event["id"]) | |
| 64 | + return cluster_id, False | |
| 65 | + | |
| 66 | + cluster_id = existing["id"] | |
| 67 | + canonical_id = existing["canonical_event_id"] | |
| 68 | + if canonical_id == event["id"]: | |
| 69 | + return cluster_id, False | |
| 70 | + surfaces: list[str] = list(existing["surfaces"] or []) | |
| 71 | + surface = event.get("surface") or "other" | |
| 72 | + same_source = surface in surfaces and event.get("sensor_id") is not None and await fetch_one( | |
| 73 | + conn, "select 1 from events where id = :c and sensor_id = :s", c=canonical_id, s=event["sensor_id"]) is not None | |
| 74 | + if surface not in surfaces: | |
| 75 | + surfaces.append(surface) | |
| 76 | + extra_surfaces = max(0, len(surfaces) - 1) | |
| 77 | + canonical = await fetch_one(conn, "select confidence, payload from events where id = :id", id=canonical_id) | |
| 78 | + base_conf = float(canonical["confidence"]) if canonical else float(existing["confidence"]) | |
| 79 | + new_conf = min(CONFIDENCE_CAP, max(base_conf, float(event["confidence"])) + CORROBORATION_BONUS * extra_surfaces) | |
| 80 | + | |
| 81 | + await execute(conn, """ | |
| 82 | + update event_clusters set source_count = source_count + 1, surfaces = cast(:surfaces as text[]), confidence = :conf, | |
| 83 | + last_detected_at = greatest(last_detected_at, cast(:at as timestamptz)) where id = :id""", | |
| 84 | + surfaces=surfaces, conf=new_conf, at=event["detected_at"], id=cluster_id) | |
| 85 | + await execute(conn, "update events set cluster_id = :c, status = 'duplicate' where id = :e", c=cluster_id, e=event["id"]) | |
| 86 | + | |
| 87 | + payload = dict((canonical or {}).get("payload") or {}) | |
| 88 | + sources = list(payload.get("sources") or []) | |
| 89 | + src = {"event_id": event["id"], "sensor_id": event.get("sensor_id"), "surface": surface, "source_url": event.get("source_url"), | |
| 90 | + "detected_at": event["detected_at"].isoformat() if isinstance(event["detected_at"], datetime) else str(event["detected_at"])} | |
| 91 | + if src not in sources: | |
| 92 | + sources.append(src) | |
| 93 | + payload["sources"] = sources[:50] | |
| 94 | + payload["corroborations"] = len(sources) | |
| 95 | + await execute(conn, """ | |
| 96 | + update events set confidence = :conf, confidence_label = :label, payload = cast(:payload as jsonb) where id = :id""", | |
| 97 | + conf=new_conf, label=_label(new_conf), payload=jsonb(payload), id=canonical_id) | |
| 98 | + if event.get("source_url") and not same_source: | |
| 99 | + await execute(conn, """ | |
| 100 | + insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, detected_at, kind) | |
| 101 | + values (:event_id, :sensor_id, :url, :snap, :surface, :at, 'corroboration') on conflict do nothing""", | |
| 102 | + event_id=canonical_id, sensor_id=event.get("sensor_id"), url=event["source_url"], snap=event.get("snapshot_after"), | |
| 103 | + surface=surface, at=event["detected_at"]) | |
| 104 | + return cluster_id, True | |
| 105 | + | |
| 106 | + | |
| 107 | +def _label(confidence: float) -> str: | |
| 108 | + from companyatlas.taxonomy import confidence_label | |
| 109 | + | |
| 110 | + return confidence_label(confidence) | |
| 111 | + | |
| 112 | + | |
| 113 | +__all__ = ["CLUSTER_WINDOW_DAYS", "CORROBORATION_BONUS", "attach_to_cluster", "cluster_key_for", "normalize_entity_key", "window_bucket"] | |
added
src/companyatlas/services/digest.py
+76 −0
@@ -0,0 +1,76 @@ | ||
| 1 | +"""Digest data (spec §142): weekly company digest and industry / country market digests as JSON. No e-mail sending here — the API or | |
| 2 | +a future mailer renders these. Everything comes from measured tables (events, metrics, jobs, signals); empty sections stay empty. | |
| 3 | +""" | |
| 4 | +from __future__ import annotations | |
| 5 | + | |
| 6 | +from datetime import UTC, datetime, timedelta | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +from companyatlas.db import fetch_all, fetch_one, transaction | |
| 10 | + | |
| 11 | +EVENT_COLS = "id, event_type, event_subtype, importance, confidence, confidence_label, title, summary, detected_at, source_url, surface, origin" | |
| 12 | + | |
| 13 | + | |
| 14 | +async def company_digest(slug_or_id: str, *, days: int = 7, now: datetime | None = None) -> dict[str, Any] | None: | |
| 15 | + now = now or datetime.now(UTC) | |
| 16 | + since = now - timedelta(days=days) | |
| 17 | + async with transaction() as conn: | |
| 18 | + co = await fetch_one(conn, "select id, slug, display_name, canonical_domain, country, industries, last_event_at from companies where slug = :s or id = :s", s=slug_or_id) | |
| 19 | + if co is None: | |
| 20 | + return None | |
| 21 | + events = await fetch_all(conn, f"select {EVENT_COLS} from events where company_id = :c and detected_at >= :since and status = 'active' order by importance desc, detected_at desc limit 50", | |
| 22 | + c=co["id"], since=since) | |
| 23 | + by_type: dict[str, int] = {} | |
| 24 | + for e in events: | |
| 25 | + by_type[e["event_type"]] = by_type.get(e["event_type"], 0) + 1 | |
| 26 | + metrics = {r["metric"]: {"value": r["value"], "confidence": r["confidence"], "computed_at": r["computed_at"]} for r in | |
| 27 | + await fetch_all(conn, "select metric, value, confidence, computed_at from metrics_current where company_id = :c", c=co["id"])} | |
| 28 | + series = await fetch_all(conn, "select metric, day, value from metric_series where company_id = :c and day >= :d and metric in ('activity_score', 'open_jobs') order by day", | |
| 29 | + c=co["id"], d=(now - timedelta(days=days * 2)).date()) | |
| 30 | + jobs = await fetch_one(conn, """select count(*) filter (where status = 'open') as open, count(*) filter (where first_seen_at >= :since) as new, | |
| 31 | + count(*) filter (where removed_at >= :since) as no_longer_listed, count(*) filter (where status = 'open' and is_ai) as ai_open | |
| 32 | + from jobs where company_id = :c""", c=co["id"], since=since) | |
| 33 | + signals = await fetch_all(conn, "select kind, strength, confidence, title, explanation, detected_at from signals where company_id = :c and status = 'active' order by strength desc", c=co["id"]) | |
| 34 | + sensors = await fetch_one(conn, "select count(*) filter (where status = 'active') as active, count(*) as total, max(last_success_at) as last_checked from sensors where company_id = :c", c=co["id"]) | |
| 35 | + return {"company": co, "window": {"days": days, "since": since.isoformat(), "until": now.isoformat()}, "highlights": events[:8], "events": events, | |
| 36 | + "events_by_type": by_type, "metrics": metrics, "series": {m: [{"day": r["day"].isoformat(), "value": r["value"]} for r in series if r["metric"] == m] for m in ("activity_score", "open_jobs")}, | |
| 37 | + "jobs": jobs, "signals": signals, "coverage": sensors, "generated_at": now.isoformat()} | |
| 38 | + | |
| 39 | + | |
| 40 | +async def scope_digest(scope: str, key: str, *, days: int = 7, now: datetime | None = None, limit: int = 10) -> dict[str, Any]: | |
| 41 | + """scope ∈ industry | country. Movers = highest activity, top events = most important, hiring = aggregate momentum.""" | |
| 42 | + now = now or datetime.now(UTC) | |
| 43 | + since = now - timedelta(days=days) | |
| 44 | + if scope == "country": | |
| 45 | + where = "co.country = :key" | |
| 46 | + elif scope == "industry": | |
| 47 | + where = "(:key = any(co.industries) or co.industry_primary = :key)" | |
| 48 | + else: | |
| 49 | + raise ValueError("scope must be industry or country") | |
| 50 | + async with transaction() as conn: | |
| 51 | + companies_n = await fetch_val_int(conn, f"select count(*) from companies co where {where}", key=key) | |
| 52 | + events = await fetch_all(conn, f"""select e.{EVENT_COLS.replace(', ', ', e.')}, co.slug, co.display_name from events e join companies co on co.id = e.company_id | |
| 53 | + where {where} and e.detected_at >= :since and e.status = 'active' order by e.importance desc, e.detected_at desc limit :limit""", | |
| 54 | + key=key, since=since, limit=limit * 2) | |
| 55 | + by_type = await fetch_all(conn, f"""select e.event_type, count(*) as n from events e join companies co on co.id = e.company_id | |
| 56 | + where {where} and e.detected_at >= :since and e.status = 'active' group by 1 order by n desc""", key=key, since=since) | |
| 57 | + movers = await fetch_all(conn, f"""select co.slug, co.display_name, m.value as activity_score from metrics_current m join companies co on co.id = m.company_id | |
| 58 | + where {where} and m.metric = 'activity_score' order by m.value desc limit :limit""", key=key, limit=limit) | |
| 59 | + hiring = await fetch_one(conn, f"""select count(*) filter (where j.status = 'open') as open, count(*) filter (where j.first_seen_at >= :since) as new, | |
| 60 | + count(*) filter (where j.removed_at >= :since) as no_longer_listed | |
| 61 | + from jobs j join companies co on co.id = j.company_id where {where}""", key=key, since=since) | |
| 62 | + momentum = await fetch_one(conn, f"""select avg(m.value) as avg_momentum_30d, count(*) as companies from metrics_current m join companies co on co.id = m.company_id | |
| 63 | + where {where} and m.metric = 'hiring_momentum_30d'""", key=key) | |
| 64 | + signals = await fetch_all(conn, "select kind, strength, title, explanation, evidence, detected_at from signals where scope = :s and scope_key = :k and status = 'active' order by strength desc", | |
| 65 | + s=scope, k=key) | |
| 66 | + return {"scope": scope, "key": key, "window": {"days": days, "since": since.isoformat(), "until": now.isoformat()}, "companies": companies_n, | |
| 67 | + "top_events": events, "events_by_type": {r["event_type"]: r["n"] for r in by_type}, "movers": movers, "hiring": {**(hiring or {}), **(momentum or {})}, | |
| 68 | + "signals": signals, "generated_at": now.isoformat()} | |
| 69 | + | |
| 70 | + | |
| 71 | +async def fetch_val_int(conn, sql: str, **params: Any) -> int: # type: ignore[no-untyped-def] | |
| 72 | + row = await fetch_one(conn, sql, **params) | |
| 73 | + return int(next(iter(row.values()))) if row else 0 | |
| 74 | + | |
| 75 | + | |
| 76 | +__all__ = ["company_digest", "scope_digest"] | |
added
src/companyatlas/services/events.py
+904 −0
@@ -0,0 +1,904 @@ | ||
| 1 | +"""Deterministic event generation (spec §21–23, §158, §167–168). | |
| 2 | + | |
| 3 | + changes (status='pending') ──▶ derive_events(): typed rules over `structured_delta` + block `diff` + surface | |
| 4 | + ──▶ events (+ event_sources, clusters, dedupe keys, review_queue) ──▶ llm_jobs when useful | |
| 5 | + | |
| 6 | +Design: | |
| 7 | +- `derive_events()` is a pure function (no I/O) so rules are unit-testable from fixtures. `process_pending_changes()` is the DB layer. | |
| 8 | +- Wording is careful by construction: "detected", "listed", "no longer listed", "observed at". Never "fired", "laid off", "shut down". | |
| 9 | +- Idempotent: `events.dedupe_key = sha(company, subtype, entity key, sensor, detection day)` → re-running a change is a no-op. | |
| 10 | +- Importance = subtype default × f(significance, magnitude); confidence = evidence quality (ATS JSON 0.95 · JSON-LD 0.9 · HTML 0.8 · | |
| 11 | + text-diff-only 0.7), corroboration handled by `services/clustering.py`. | |
| 12 | +- Noise / minor changes never produce events. LLM enrichment is queued only for meaningful+ changes, within the daily budget. | |
| 13 | +""" | |
| 14 | +from __future__ import annotations | |
| 15 | + | |
| 16 | +import logging | |
| 17 | +import re | |
| 18 | +from dataclasses import dataclass, field | |
| 19 | +from datetime import UTC, date, datetime | |
| 20 | +from typing import Any | |
| 21 | + | |
| 22 | +from companyatlas.config import settings | |
| 23 | +from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 24 | +from companyatlas.ids import new_id, stable_hash | |
| 25 | +from companyatlas.services.clustering import attach_to_cluster, normalize_entity_key | |
| 26 | +from companyatlas.services.periodic import periodic | |
| 27 | +from companyatlas.taxonomy import ( | |
| 28 | + AI_KEYWORDS, | |
| 29 | + EVENT_SUBTYPES, | |
| 30 | + EVIDENCE_CONFIDENCE, | |
| 31 | + FORBIDDEN_WORDING, | |
| 32 | + ChangeKind, | |
| 33 | + EventType, | |
| 34 | + Surface, | |
| 35 | + confidence_label, | |
| 36 | +) | |
| 37 | + | |
| 38 | +log = logging.getLogger(__name__) | |
| 39 | + | |
| 40 | +RULES_VERSION = "rules-v1" | |
| 41 | +SCHEMA_VERSION = "event-v1" | |
| 42 | +MAX_ENTITY_ITEMS = 50 # bounded entity lists on aggregate events | |
| 43 | +PER_JOB_EVENT_MAX = 5 # NEW_JOB per job only when ≤ 5 jobs added | |
| 44 | +PER_PERSON_EVENT_MAX = 10 | |
| 45 | +NEWS_ITEMS_MAX = 20 | |
| 46 | +LEADERSHIP_AGGREGATE_MIN = 3 | |
| 47 | +SURGE_MIN_JOBS = 10 # fallback thresholds when no baseline is available yet | |
| 48 | +SURGE_MIN_RATIO = 0.5 | |
| 49 | +FREEZE_MIN_RATIO = 0.5 | |
| 50 | +BASELINE_Z = 2.0 # surge/freeze when the delta exceeds mean + 2σ of the company's weekly baseline | |
| 51 | +IMPORTANCE_FLOOR = 0.05 | |
| 52 | +LEGAL_SURFACES = {Surface.LEGAL_TERMS, Surface.LEGAL_PRIVACY, Surface.SECURITY} | |
| 53 | +DEVELOPER_SURFACES = {Surface.DOCS, Surface.DEVELOPER, Surface.API, Surface.CHANGELOG} | |
| 54 | +NEWS_SURFACES = {Surface.NEWSROOM, Surface.BLOG, Surface.FEED, Surface.INVESTOR_RELATIONS, Surface.CHANGELOG, Surface.RESEARCH} | |
| 55 | +AMBIGUOUS_SURFACES = {Surface.OTHER, Surface.PARTNERS, Surface.CUSTOMERS, Surface.SOLUTIONS, Surface.SERVICES, Surface.INDUSTRIES, | |
| 56 | + Surface.SUPPORT, Surface.CONTACT, Surface.SITEMAP} | |
| 57 | +SUMMARY_SUBTYPES = {"TERMS_CHANGE", "PRIVACY_POLICY_CHANGE", "SECURITY_UPDATE", "WEBSITE_CHANGE", "HOMEPAGE_REDESIGN", "NEWS_RELEASE", | |
| 58 | + "INVESTOR_UPDATE", "EARNINGS_RELEASE", "MESSAGING_CHANGE"} | |
| 59 | +ATS_CONNECTOR_HINTS = ("greenhouse", "lever", "ashby", "smartrecruiters", "workday", "workable", "json", "ats") | |
| 60 | + | |
| 61 | +SURFACE_LABEL: dict[str, str] = { | |
| 62 | + Surface.CAREERS: "careers page", Surface.JOBS_BOARD: "job board", Surface.PRICING: "pricing page", Surface.LEADERSHIP: "leadership page", | |
| 63 | + Surface.ABOUT: "about page", Surface.PRODUCTS: "products page", Surface.SERVICES: "services page", Surface.SOLUTIONS: "solutions page", | |
| 64 | + Surface.LOCATIONS: "locations page", Surface.CONTACT: "contact page", Surface.NEWSROOM: "newsroom", Surface.BLOG: "blog", | |
| 65 | + Surface.FEED: "feed", Surface.DOCS: "documentation", Surface.DEVELOPER: "developer portal", Surface.API: "API reference", | |
| 66 | + Surface.CHANGELOG: "changelog", Surface.INVESTOR_RELATIONS: "investor relations page", Surface.LEGAL_TERMS: "terms of service page", | |
| 67 | + Surface.LEGAL_PRIVACY: "privacy policy", Surface.SECURITY: "security page", Surface.HOMEPAGE: "homepage", | |
| 68 | + Surface.SUSTAINABILITY: "sustainability page", Surface.STATUS: "status page", Surface.PARTNERS: "partners page", | |
| 69 | + Surface.CUSTOMERS: "customers page", Surface.INDUSTRIES: "industries page", Surface.RESEARCH: "research page", | |
| 70 | + Surface.SUPPORT: "support page", Surface.SITEMAP: "sitemap", Surface.OTHER: "monitored page", | |
| 71 | +} | |
| 72 | +CURRENCY_SYMBOL = {"USD": "$", "EUR": "€", "GBP": "£", "CAD": "CA$", "AUD": "A$", "JPY": "¥", "CHF": "CHF ", "INR": "₹", "BRL": "R$"} | |
| 73 | +PERIOD_LABEL = {"month": "per month", "year": "per year", "one_time": "one-time", "usage": "usage-based", "contact": "contact sales"} | |
| 74 | +EXEC_ROLES = {"ceo", "cfo", "cto", "coo", "founder", "president", "chair", "board", "cmo", "cro", "cpo", "ciso", "cio", "chro", "gm"} | |
| 75 | +_EARNINGS_RE = re.compile(r"\b(earnings|quarterly results|q[1-4]\s*(fy)?\s*20\d\d|fiscal (year|quarter)|financial results|results for the (quarter|year))\b", re.IGNORECASE) | |
| 76 | +_INVESTOR_RE = re.compile(r"\b(investor|shareholder|annual report|annual meeting|dividend|10-k|10-q|8-k|proxy statement|guidance)\b", re.IGNORECASE) | |
| 77 | +_PRESS_RE = re.compile(r"\b(announces|announced|unveils|introduces|launches|partners with|acquires|appoints|names|expands|opens)\b", re.IGNORECASE) | |
| 78 | +_LAUNCH_RE = re.compile(r"\b(launch(es|ed|ing)?|introduc(es|ed|ing)|unveil(s|ed)|now available|general availability)\b", re.IGNORECASE) | |
| 79 | +_ACQ_RE = re.compile(r"\b(acquires|acquired|acquisition of|to acquire|merger|merges with)\b", re.IGNORECASE) | |
| 80 | +_FUNDING_RE = re.compile(r"\b(raises|raised|series [a-f]\b|seed round|funding round|closes \$|financing of)\b", re.IGNORECASE) | |
| 81 | + | |
| 82 | + | |
| 83 | +# ============================================================================================================== drafts | |
| 84 | + | |
| 85 | + | |
| 86 | +@dataclass(slots=True) | |
| 87 | +class EventDraft: | |
| 88 | + subtype: str | |
| 89 | + title: str | |
| 90 | + entity_key: str | |
| 91 | + summary: str | None = None | |
| 92 | + old_value: str | None = None | |
| 93 | + new_value: str | None = None | |
| 94 | + entities: dict[str, Any] = field(default_factory=dict) | |
| 95 | + payload: dict[str, Any] = field(default_factory=dict) | |
| 96 | + tags: list[str] = field(default_factory=list) | |
| 97 | + magnitude: float = 0.0 # 0–1 rule-specific magnitude (relative job delta, |price pct| …) → importance bonus | |
| 98 | + evidence: str = "html" # ats_json | jsonld | html | text_diff | |
| 99 | + effective_at: datetime | None = None | |
| 100 | + published_at: datetime | None = None | |
| 101 | + review: str | None = None # review_queue kind when the rule itself wants a human look | |
| 102 | + importance: float = 0.0 # filled by finalize() | |
| 103 | + confidence: float = 0.0 | |
| 104 | + | |
| 105 | + @property | |
| 106 | + def event_type(self) -> str: | |
| 107 | + return str(EVENT_SUBTYPES.get(self.subtype, (EventType.OTHER, 0.3))[0]) | |
| 108 | + | |
| 109 | + | |
| 110 | +@dataclass(slots=True) | |
| 111 | +class Derived: | |
| 112 | + events: list[EventDraft] | |
| 113 | + needs_classification: bool = False | |
| 114 | + classification_reason: str | None = None | |
| 115 | + summarize: list[str] = field(default_factory=list) # subtypes whose events deserve an LLM summary | |
| 116 | + | |
| 117 | + | |
| 118 | +# ============================================================================================================== helpers | |
| 119 | + | |
| 120 | + | |
| 121 | +def safe_wording(text: str) -> str: | |
| 122 | + """Defensive: rewrite forbidden phrasing (rules never produce it, LLM output might).""" | |
| 123 | + out = text | |
| 124 | + for bad in FORBIDDEN_WORDING: | |
| 125 | + if bad in out.lower(): | |
| 126 | + out = re.sub(re.escape(bad), "no longer listed", out, flags=re.IGNORECASE) | |
| 127 | + return out | |
| 128 | + | |
| 129 | + | |
| 130 | +def scale_importance(default: float, significance: float, magnitude: float = 0.0) -> float: | |
| 131 | + """importance = default × (0.7 + 0.6·significance) × (1 + 0.3·magnitude), clamped to [0.05, 1].""" | |
| 132 | + sig = min(1.0, max(0.0, float(significance or 0.0))) | |
| 133 | + mag = min(1.0, max(0.0, float(magnitude or 0.0))) | |
| 134 | + return round(min(1.0, max(IMPORTANCE_FLOOR, default * (0.7 + 0.6 * sig) * (1.0 + 0.3 * mag))), 4) | |
| 135 | + | |
| 136 | + | |
| 137 | +def evidence_kind(sensor: dict[str, Any], delta: dict[str, Any]) -> str: | |
| 138 | + connector = (sensor.get("connector_id") or "").lower() | |
| 139 | + surface = sensor.get("surface") or "" | |
| 140 | + meta = delta.get("meta") or {} | |
| 141 | + hinted = meta.get("evidence") | |
| 142 | + if hinted in EVIDENCE_CONFIDENCE: | |
| 143 | + return str(hinted) | |
| 144 | + if surface == Surface.JOBS_BOARD or any(h in connector for h in ATS_CONNECTOR_HINTS) or sensor.get("fetch_mode") == "json": | |
| 145 | + return "ats_json" | |
| 146 | + if meta.get("jsonld") or "jsonld" in connector or "feed" in connector or surface == Surface.FEED: | |
| 147 | + return "jsonld" | |
| 148 | + if any(delta.get(k) for k in ("jobs", "people", "products", "plans", "locations", "news")): | |
| 149 | + return "html" | |
| 150 | + return "text_diff" | |
| 151 | + | |
| 152 | + | |
| 153 | +def _label(surface: str) -> str: | |
| 154 | + return SURFACE_LABEL.get(surface, "monitored page") | |
| 155 | + | |
| 156 | + | |
| 157 | +def _plural(n: int, one: str, many: str | None = None) -> str: | |
| 158 | + return one if n == 1 else (many or one + "s") | |
| 159 | + | |
| 160 | + | |
| 161 | +def _money(amount: Any, currency: str | None) -> str: | |
| 162 | + try: | |
| 163 | + value = float(amount) | |
| 164 | + except (TypeError, ValueError): | |
| 165 | + return str(amount) | |
| 166 | + text = f"{value:,.0f}" if value.is_integer() else f"{value:,.2f}" | |
| 167 | + cur = (currency or "").upper() | |
| 168 | + sym = CURRENCY_SYMBOL.get(cur) | |
| 169 | + if sym: | |
| 170 | + return f"{sym}{text}" | |
| 171 | + return f"{text} {cur}".strip() | |
| 172 | + | |
| 173 | + | |
| 174 | +def _place(item: dict[str, Any]) -> str: | |
| 175 | + parts = [p for p in (item.get("city"), item.get("region")) if p] | |
| 176 | + country = item.get("country") | |
| 177 | + if country: | |
| 178 | + parts.append(str(country).upper()) | |
| 179 | + if parts: | |
| 180 | + return ", ".join(parts) | |
| 181 | + return item.get("name") or item.get("location_text") or "" | |
| 182 | + | |
| 183 | + | |
| 184 | +def _is_ai(text: str | None) -> bool: | |
| 185 | + if not text: | |
| 186 | + return False | |
| 187 | + hay = f" {text.lower()} " | |
| 188 | + return any(k in hay for k in AI_KEYWORDS) | |
| 189 | + | |
| 190 | + | |
| 191 | +def _job_is_ai(job: dict[str, Any]) -> bool: | |
| 192 | + return bool(job.get("is_ai")) or _is_ai(job.get("title")) | |
| 193 | + | |
| 194 | + | |
| 195 | +def _job_label(job: dict[str, Any]) -> str: | |
| 196 | + title = (job.get("title") or "position").strip() | |
| 197 | + loc = job.get("location_text") or _place(job) | |
| 198 | + if job.get("remote") and not loc: | |
| 199 | + loc = "Remote" | |
| 200 | + return f"{title} ({loc})" if loc else title | |
| 201 | + | |
| 202 | + | |
| 203 | +def _sections(diff: dict[str, Any]) -> list[str]: | |
| 204 | + seen: list[str] = [] | |
| 205 | + for bucket in ("modified", "added", "removed"): | |
| 206 | + for d in diff.get(bucket) or []: | |
| 207 | + path = (d.get("path") or "").strip() | |
| 208 | + name = path.split(">")[-1].strip() if path else "" | |
| 209 | + if not name: | |
| 210 | + text = (d.get("after") or d.get("before") or "").strip() | |
| 211 | + name = text.split("\n")[0][:80] if text else "" | |
| 212 | + if name and name not in seen: | |
| 213 | + seen.append(name) | |
| 214 | + return seen[:20] | |
| 215 | + | |
| 216 | + | |
| 217 | +def _blocks_changed(diff: dict[str, Any], change: dict[str, Any]) -> int: | |
| 218 | + counts = diff.get("counts") or {} | |
| 219 | + n = int(counts.get("added") or 0) + int(counts.get("removed") or 0) + int(counts.get("modified") or 0) | |
| 220 | + if n == 0: | |
| 221 | + n = int(change.get("blocks_added") or 0) + int(change.get("blocks_removed") or 0) + int(change.get("blocks_modified") or 0) | |
| 222 | + return n | |
| 223 | + | |
| 224 | + | |
| 225 | +def _dt(value: Any) -> datetime | None: | |
| 226 | + if value is None: | |
| 227 | + return None | |
| 228 | + if isinstance(value, datetime): | |
| 229 | + return value if value.tzinfo else value.replace(tzinfo=UTC) | |
| 230 | + try: | |
| 231 | + return datetime.fromisoformat(str(value)) | |
| 232 | + except ValueError: | |
| 233 | + return None | |
| 234 | + | |
| 235 | + | |
| 236 | +# ============================================================================================================== rule families | |
| 237 | + | |
| 238 | + | |
| 239 | +def _hiring_rules(delta: dict[str, Any], surface: str, evidence: str, baseline: dict[str, Any] | None) -> list[EventDraft]: | |
| 240 | + jobs = delta.get("jobs") or {} | |
| 241 | + added: list[dict[str, Any]] = list(jobs.get("added") or []) | |
| 242 | + removed: list[dict[str, Any]] = list(jobs.get("removed") or []) | |
| 243 | + open_before = jobs.get("open_before") | |
| 244 | + open_after = jobs.get("open_after") | |
| 245 | + n_add, n_rem = len(added), len(removed) | |
| 246 | + if not (n_add or n_rem): | |
| 247 | + return [] | |
| 248 | + if isinstance(open_before, int) and isinstance(open_after, int): | |
| 249 | + net = open_after - open_before | |
| 250 | + else: | |
| 251 | + net = n_add - n_rem | |
| 252 | + label = _label(surface) | |
| 253 | + out: list[EventDraft] = [] | |
| 254 | + counts_key = f"{open_before}>{open_after}" if open_before is not None else f"+{n_add}-{n_rem}" | |
| 255 | + base_denominator = max(int(open_before or 0), 5) | |
| 256 | + countries = sorted({str(j.get("country")).upper() for j in added if j.get("country")}) | |
| 257 | + departments = sorted({str(j.get("department")) for j in added if j.get("department")})[:20] | |
| 258 | + ai_added = [j for j in added if _job_is_ai(j)] | |
| 259 | + common_payload = {"added": n_add, "removed": n_rem, "open_before": open_before, "open_after": open_after, "net": net, | |
| 260 | + "ai_added": len(ai_added), "countries": countries, "departments": departments} | |
| 261 | + | |
| 262 | + if net > 0 and n_add: | |
| 263 | + title = f"{n_add} new {_plural(n_add, 'position')} detected on {label}" | |
| 264 | + if n_rem: | |
| 265 | + title += f" ({n_rem} no longer visible)" | |
| 266 | + out.append(EventDraft( | |
| 267 | + subtype="JOB_COUNT_INCREASE", title=title, entity_key=f"jobs:{counts_key}", evidence=evidence, | |
| 268 | + summary=_open_summary(open_before, open_after, n_add, n_rem), old_value=_s(open_before), new_value=_s(open_after), | |
| 269 | + entities={"jobs": [_job_entity(j) for j in added[:MAX_ENTITY_ITEMS]]}, payload=common_payload, | |
| 270 | + tags=_hiring_tags(countries, ai_added), magnitude=min(1.0, n_add / base_denominator))) | |
| 271 | + elif net < 0 and n_rem: | |
| 272 | + title = f"{n_rem} monitored job {_plural(n_rem, 'listing')} no longer visible on {label}" | |
| 273 | + if n_add: | |
| 274 | + title += f" ({n_add} new)" | |
| 275 | + out.append(EventDraft( | |
| 276 | + subtype="JOB_COUNT_DECREASE", title=title, entity_key=f"jobs:{counts_key}", evidence=evidence, | |
| 277 | + summary=_open_summary(open_before, open_after, n_add, n_rem), old_value=_s(open_before), new_value=_s(open_after), | |
| 278 | + entities={"jobs": [_job_entity(j) for j in removed[:MAX_ENTITY_ITEMS]]}, payload=common_payload, | |
| 279 | + tags=["hiring"], magnitude=min(1.0, n_rem / base_denominator))) | |
| 280 | + | |
| 281 | + if ai_added: | |
| 282 | + k = len(ai_added) | |
| 283 | + out.append(EventDraft( | |
| 284 | + subtype="AI_HIRING", title=f"{k} AI-related {_plural(k, 'position')} detected on {label}", entity_key=f"ai_jobs:{counts_key}", | |
| 285 | + evidence=evidence, summary="AI-related roles identified from listing titles: " + "; ".join(_job_label(j) for j in ai_added[:5]), | |
| 286 | + entities={"jobs": [_job_entity(j) for j in ai_added[:MAX_ENTITY_ITEMS]]}, payload={"ai_added": k, "added": n_add}, | |
| 287 | + tags=["hiring", "ai"], magnitude=min(1.0, k / 5))) | |
| 288 | + | |
| 289 | + if 1 <= n_add <= PER_JOB_EVENT_MAX: | |
| 290 | + for j in added: | |
| 291 | + out.append(EventDraft( | |
| 292 | + subtype="NEW_JOB", title=f"New position listed: {_job_label(j)}", entity_key="job:" + normalize_entity_key(_job_label(j)), | |
| 293 | + evidence=evidence, new_value=j.get("title"), entities={"jobs": [_job_entity(j)]}, | |
| 294 | + payload={"url": j.get("url"), "department": j.get("department"), "country": j.get("country"), "remote": j.get("remote")}, | |
| 295 | + tags=["hiring"] + (["ai"] if _job_is_ai(j) else []), published_at=_dt(j.get("posted_at")))) | |
| 296 | + | |
| 297 | + surge, freeze = _surge_freeze(n_add, n_rem, open_before, open_after, baseline) | |
| 298 | + if surge is not None: | |
| 299 | + out.append(EventDraft( | |
| 300 | + subtype="HIRING_SURGE", title=f"Hiring surge signal: {n_add} new positions detected in one observation" + surge, | |
| 301 | + entity_key=f"surge:{counts_key}", evidence=evidence, payload={**common_payload, "baseline": baseline}, | |
| 302 | + summary="Signal, not a fact: the number of new listings exceeds this company's usual weekly volume.", | |
| 303 | + tags=["hiring", "signal"], magnitude=min(1.0, n_add / max(base_denominator, SURGE_MIN_JOBS)))) | |
| 304 | + if freeze is not None: | |
| 305 | + out.append(EventDraft( | |
| 306 | + subtype="HIRING_FREEZE_SIGNAL", | |
| 307 | + title=f"Hiring slowdown signal: {n_rem} of {open_before if open_before is not None else n_rem + (open_after or 0)} monitored listings no longer visible" + freeze, | |
| 308 | + entity_key=f"freeze:{counts_key}", evidence=evidence, payload={**common_payload, "baseline": baseline}, | |
| 309 | + summary="Signal, not a fact: listings disappearing from a public careers page can reflect closed roles, ATS migrations or page changes.", | |
| 310 | + tags=["hiring", "signal"], review="unexpected_activity", magnitude=min(1.0, n_rem / base_denominator))) | |
| 311 | + return out | |
| 312 | + | |
| 313 | + | |
| 314 | +def _surge_freeze(n_add: int, n_rem: int, open_before: Any, open_after: Any, baseline: dict[str, Any] | None) -> tuple[str | None, str | None]: | |
| 315 | + surge = freeze = None | |
| 316 | + b = (baseline or {}).get("jobs_new_weekly") | |
| 317 | + before = int(open_before) if isinstance(open_before, int) else None | |
| 318 | + if b and b.get("samples", 0) >= 4 and b.get("stddev") is not None: | |
| 319 | + threshold = float(b["mean"]) + BASELINE_Z * max(float(b["stddev"]), 1.0) | |
| 320 | + if n_add >= max(5, threshold): | |
| 321 | + surge = f" (baseline ≈ {float(b['mean']):.1f} new/week)" | |
| 322 | + elif n_add >= SURGE_MIN_JOBS and (before is None or n_add >= SURGE_MIN_RATIO * before): | |
| 323 | + surge = "" | |
| 324 | + if before and n_rem >= SURGE_MIN_JOBS and n_rem >= FREEZE_MIN_RATIO * before and (open_after is None or int(open_after) <= (1 - FREEZE_MIN_RATIO) * before): | |
| 325 | + freeze = "" | |
| 326 | + return surge, freeze | |
| 327 | + | |
| 328 | + | |
| 329 | +def _open_summary(before: Any, after: Any, n_add: int, n_rem: int) -> str: | |
| 330 | + parts = [f"{n_add} added" if n_add else "", f"{n_rem} no longer visible" if n_rem else ""] | |
| 331 | + s = ", ".join(p for p in parts if p) | |
| 332 | + if before is not None and after is not None: | |
| 333 | + s += f"; open listings observed: {before} → {after}" | |
| 334 | + return s + "." | |
| 335 | + | |
| 336 | + | |
| 337 | +def _hiring_tags(countries: list[str], ai_added: list[dict[str, Any]]) -> list[str]: | |
| 338 | + tags = ["hiring"] + [f"country:{c}" for c in countries[:5]] | |
| 339 | + if ai_added: | |
| 340 | + tags.append("ai") | |
| 341 | + return tags | |
| 342 | + | |
| 343 | + | |
| 344 | +def _job_entity(j: dict[str, Any]) -> dict[str, Any]: | |
| 345 | + return {k: j.get(k) for k in ("title", "url", "location_text", "country", "remote", "department", "is_ai") if j.get(k) is not None} | |
| 346 | + | |
| 347 | + | |
| 348 | +def _s(v: Any) -> str | None: | |
| 349 | + return None if v is None else str(v) | |
| 350 | + | |
| 351 | + | |
| 352 | +def _pricing_rules(delta: dict[str, Any], surface: str, evidence: str, diff: dict[str, Any], change: dict[str, Any]) -> list[EventDraft]: | |
| 353 | + plans = delta.get("plans") or {} | |
| 354 | + out: list[EventDraft] = [] | |
| 355 | + for p in plans.get("price_changed") or []: | |
| 356 | + name = p.get("plan_name") or "Plan" | |
| 357 | + before, after = p.get("before"), p.get("after") | |
| 358 | + try: | |
| 359 | + b, a = float(before), float(after) | |
| 360 | + except (TypeError, ValueError): | |
| 361 | + continue | |
| 362 | + if a == b: | |
| 363 | + continue | |
| 364 | + pct = p.get("pct") | |
| 365 | + if pct is None and b: | |
| 366 | + pct = round((a - b) / b * 100, 1) | |
| 367 | + cur = p.get("currency") | |
| 368 | + period = PERIOD_LABEL.get(p.get("billing_period") or "", "") | |
| 369 | + subtype = "PRICE_INCREASE" if a > b else "PRICE_DECREASE" | |
| 370 | + title = f"{name} plan price observed at {_money(a, cur)} (was {_money(b, cur)})" | |
| 371 | + summary = f"{'Increase' if a > b else 'Decrease'} of {abs(pct):.1f}%" if pct is not None else None | |
| 372 | + if summary and period: | |
| 373 | + summary += f", billed {period}" | |
| 374 | + out.append(EventDraft( | |
| 375 | + subtype=subtype, title=title, entity_key="plan:" + normalize_entity_key(name), evidence=evidence, summary=(summary + "." if summary else None), | |
| 376 | + old_value=_money(b, cur), new_value=_money(a, cur), | |
| 377 | + payload={"plan_name": name, "before": b, "after": a, "pct": pct, "currency": cur, "billing_period": p.get("billing_period")}, | |
| 378 | + entities={"plans": [{"plan_name": name}]}, tags=["pricing"], magnitude=min(1.0, abs(float(pct or 0)) / 50.0))) | |
| 379 | + for p in plans.get("added") or []: | |
| 380 | + name = p.get("plan_name") or "New plan" | |
| 381 | + price = p.get("price") | |
| 382 | + cur = p.get("currency") | |
| 383 | + if p.get("contact_sales") or (price is None and (p.get("billing_period") == "contact")): | |
| 384 | + price_txt = "contact sales" | |
| 385 | + tags = ["pricing", "enterprise"] | |
| 386 | + elif price is not None: | |
| 387 | + price_txt = _money(price, cur) + (f" {PERIOD_LABEL.get(p.get('billing_period') or '', '')}".rstrip()) | |
| 388 | + tags = ["pricing"] | |
| 389 | + else: | |
| 390 | + price_txt = p.get("price_text") or "price not stated" | |
| 391 | + tags = ["pricing"] | |
| 392 | + out.append(EventDraft( | |
| 393 | + subtype="NEW_PRICING_TIER", title=f"New pricing tier listed: {name} ({price_txt})", entity_key="plan:" + normalize_entity_key(name), | |
| 394 | + evidence=evidence, new_value=price_txt, payload={"plan_name": name, "price": price, "currency": cur, "billing_period": p.get("billing_period"), | |
| 395 | + "contact_sales": bool(p.get("contact_sales"))}, | |
| 396 | + entities={"plans": [{"plan_name": name}]}, tags=tags, magnitude=0.3)) | |
| 397 | + for p in plans.get("removed") or []: | |
| 398 | + name = p.get("plan_name") or "Plan" | |
| 399 | + out.append(EventDraft( | |
| 400 | + subtype="PRICING_TIER_REMOVED", title=f"Pricing tier no longer listed: {name}", entity_key="plan:" + normalize_entity_key(name), | |
| 401 | + evidence=evidence, old_value=name, payload={"plan_name": name, "price": p.get("price"), "currency": p.get("currency")}, | |
| 402 | + entities={"plans": [{"plan_name": name}]}, tags=["pricing"], magnitude=0.3)) | |
| 403 | + if not out and surface == Surface.PRICING: | |
| 404 | + n = _blocks_changed(diff, change) | |
| 405 | + sections = _sections(diff) | |
| 406 | + out.append(EventDraft( | |
| 407 | + subtype="PRICING_CHANGE", title=f"Pricing page materially updated ({n} {_plural(n, 'block')} changed)", entity_key="pricing_page", | |
| 408 | + evidence="text_diff", payload={"blocks_changed": n, "sections": sections, "text_delta_ratio": diff.get("text_delta_ratio")}, | |
| 409 | + summary=("Sections affected: " + ", ".join(sections[:6]) + ".") if sections else None, tags=["pricing"], | |
| 410 | + magnitude=min(1.0, float(diff.get("text_delta_ratio") or 0) * 2))) | |
| 411 | + return out | |
| 412 | + | |
| 413 | + | |
| 414 | +def _leadership_rules(delta: dict[str, Any], surface: str, evidence: str) -> list[EventDraft]: | |
| 415 | + people = delta.get("people") or {} | |
| 416 | + added = list(people.get("added") or []) | |
| 417 | + removed = list(people.get("removed") or []) | |
| 418 | + changed = list(people.get("title_changed") or []) | |
| 419 | + if not (added or removed or changed): | |
| 420 | + return [] | |
| 421 | + label = _label(surface) if surface in (Surface.LEADERSHIP, Surface.ABOUT) else "monitored leadership page" | |
| 422 | + out: list[EventDraft] = [] | |
| 423 | + | |
| 424 | + def is_exec(p: dict[str, Any]) -> bool: | |
| 425 | + return bool(p.get("is_executive")) or (p.get("role_category") or "").lower() in EXEC_ROLES | |
| 426 | + | |
| 427 | + for p in [x for x in added if is_exec(x)][:PER_PERSON_EVENT_MAX]: | |
| 428 | + name, title = p.get("name") or "Unnamed", p.get("title") | |
| 429 | + out.append(EventDraft( | |
| 430 | + subtype="NEW_EXECUTIVE", title=f"{name} listed as {title} on {label}" if title else f"{name} newly listed on {label}", | |
| 431 | + entity_key="person:" + normalize_entity_key(name), evidence=evidence, new_value=title, | |
| 432 | + entities={"people": [{"name": name, "title": title, "role_category": p.get("role_category")}]}, | |
| 433 | + payload={"role_category": p.get("role_category")}, tags=["leadership"], magnitude=0.5 if (p.get("role_category") or "").lower() in {"ceo", "cfo", "cto", "coo", "president"} else 0.2)) | |
| 434 | + for p in [x for x in removed if is_exec(x)][:PER_PERSON_EVENT_MAX]: | |
| 435 | + name, title = p.get("name") or "Unnamed", p.get("title") | |
| 436 | + out.append(EventDraft( | |
| 437 | + subtype="EXECUTIVE_NO_LONGER_LISTED", title=f"{name} no longer listed on {label}", entity_key="person:" + normalize_entity_key(name), | |
| 438 | + evidence=evidence, old_value=title, summary=f"Previously listed as {title}. Disappearance from a public page is not evidence of departure." if title else None, | |
| 439 | + entities={"people": [{"name": name, "title": title, "role_category": p.get("role_category")}]}, | |
| 440 | + payload={"role_category": p.get("role_category")}, tags=["leadership"], magnitude=0.5 if (p.get("role_category") or "").lower() in {"ceo", "cfo", "cto", "coo", "president"} else 0.2)) | |
| 441 | + for p in changed[:PER_PERSON_EVENT_MAX]: | |
| 442 | + name = p.get("name") or "Unnamed" | |
| 443 | + before, after = p.get("before"), p.get("after") | |
| 444 | + out.append(EventDraft( | |
| 445 | + subtype="EXECUTIVE_TITLE_CHANGE", title=f"{name} now listed as {after} (was {before})", entity_key="person:" + normalize_entity_key(name), | |
| 446 | + evidence=evidence, old_value=before, new_value=after, entities={"people": [{"name": name, "title": after}]}, tags=["leadership"], magnitude=0.3)) | |
| 447 | + total = len(added) + len(removed) + len(changed) | |
| 448 | + if total >= LEADERSHIP_AGGREGATE_MIN or (total and not out): | |
| 449 | + bits = [f"{len(added)} added" if added else "", f"{len(removed)} no longer listed" if removed else "", f"{len(changed)} title {_plural(len(changed), 'change')}" if changed else ""] | |
| 450 | + out.append(EventDraft( | |
| 451 | + subtype="LEADERSHIP_CHANGE", title=f"{label[0].upper()}{label[1:]} updated: " + ", ".join(b for b in bits if b), | |
| 452 | + entity_key=f"leadership:{len(added)}:{len(removed)}:{len(changed)}", evidence=evidence, | |
| 453 | + entities={"people": [{"name": p.get("name"), "title": p.get("title"), "status": s} for s, lst in (("listed", added), ("no_longer_listed", removed)) for p in lst][:MAX_ENTITY_ITEMS]}, | |
| 454 | + payload={"added": len(added), "removed": len(removed), "title_changed": len(changed)}, tags=["leadership"], magnitude=min(1.0, total / 6))) | |
| 455 | + return out | |
| 456 | + | |
| 457 | + | |
| 458 | +def _product_rules(delta: dict[str, Any], evidence: str) -> list[EventDraft]: | |
| 459 | + products = delta.get("products") or {} | |
| 460 | + out: list[EventDraft] = [] | |
| 461 | + for p in list(products.get("added") or [])[:MAX_ENTITY_ITEMS]: | |
| 462 | + name = p.get("name") or "Unnamed product" | |
| 463 | + out.append(EventDraft( | |
| 464 | + subtype="NEW_PRODUCT", title=f"New product listed: {name}", entity_key="product:" + normalize_entity_key(name), evidence=evidence, | |
| 465 | + new_value=name, entities={"products": [{"name": name, "url": p.get("url")}]}, payload={"url": p.get("url"), "category": p.get("category")}, | |
| 466 | + tags=["product"] + (["ai"] if _is_ai(name) else []), magnitude=0.3)) | |
| 467 | + for p in list(products.get("removed") or [])[:MAX_ENTITY_ITEMS]: | |
| 468 | + name = p.get("name") or "Unnamed product" | |
| 469 | + out.append(EventDraft( | |
| 470 | + subtype="PRODUCT_REMOVED", title=f"Product no longer listed: {name}", entity_key="product:" + normalize_entity_key(name), evidence=evidence, | |
| 471 | + old_value=name, entities={"products": [{"name": name, "url": p.get("url")}]}, payload={"url": p.get("url")}, tags=["product"], magnitude=0.3)) | |
| 472 | + return out | |
| 473 | + | |
| 474 | + | |
| 475 | +def _location_rules(delta: dict[str, Any], evidence: str, country_names: dict[str, str]) -> list[EventDraft]: | |
| 476 | + locations = delta.get("locations") or {} | |
| 477 | + out: list[EventDraft] = [] | |
| 478 | + kind_label = {"headquarters": "headquarters", "office": "office", "store": "store", "factory": "factory", "warehouse": "warehouse", | |
| 479 | + "lab": "lab", "data_center": "data center"} | |
| 480 | + for loc in list(locations.get("added") or [])[:MAX_ENTITY_ITEMS]: | |
| 481 | + place = _place(loc) or "unnamed location" | |
| 482 | + kind = kind_label.get(loc.get("kind") or "", "location") | |
| 483 | + out.append(EventDraft( | |
| 484 | + subtype="NEW_LOCATION", title=f"New {kind} listed: {place}", entity_key="location:" + normalize_entity_key(place), evidence=evidence, | |
| 485 | + new_value=place, entities={"locations": [{"name": loc.get("name"), "city": loc.get("city"), "country": loc.get("country"), "kind": loc.get("kind")}]}, | |
| 486 | + payload={"kind": loc.get("kind"), "country": loc.get("country")}, tags=["location"] + ([f"country:{str(loc['country']).upper()}"] if loc.get("country") else []), | |
| 487 | + magnitude=0.4 if loc.get("kind") == "headquarters" else 0.2)) | |
| 488 | + for loc in list(locations.get("removed") or [])[:MAX_ENTITY_ITEMS]: | |
| 489 | + place = _place(loc) or "unnamed location" | |
| 490 | + kind = kind_label.get(loc.get("kind") or "", "location") | |
| 491 | + out.append(EventDraft( | |
| 492 | + subtype="OFFICE_REMOVED", title=f"{kind[0].upper()}{kind[1:]} no longer listed: {place}", entity_key="location:" + normalize_entity_key(place), | |
| 493 | + evidence=evidence, old_value=place, entities={"locations": [{"name": loc.get("name"), "city": loc.get("city"), "country": loc.get("country"), "kind": loc.get("kind")}]}, | |
| 494 | + payload={"kind": loc.get("kind"), "country": loc.get("country")}, tags=["location"], magnitude=0.2)) | |
| 495 | + for code in locations.get("new_countries") or []: | |
| 496 | + code = str(code).upper() | |
| 497 | + name = country_names.get(code, code) | |
| 498 | + cities = [loc.get("city") for loc in locations.get("added") or [] if str(loc.get("country") or "").upper() == code and loc.get("city")] | |
| 499 | + detail = f" ({', '.join(cities[:3])})" if cities else "" | |
| 500 | + out.append(EventDraft( | |
| 501 | + subtype="COUNTRY_EXPANSION", title=f"New country presence listed: {name}{detail}", entity_key=f"country:{code}", evidence=evidence, | |
| 502 | + new_value=code, entities={"locations": [{"country": code, "city": c} for c in cities[:10]] or [{"country": code}]}, | |
| 503 | + payload={"country": code, "cities": cities[:10]}, tags=["location", "expansion", f"country:{code}"], magnitude=0.6)) | |
| 504 | + return out | |
| 505 | + | |
| 506 | + | |
| 507 | +def _news_subtype(item: dict[str, Any], surface: str) -> str: | |
| 508 | + title = item.get("title") or "" | |
| 509 | + category = (item.get("category") or "").lower() | |
| 510 | + if _EARNINGS_RE.search(title): | |
| 511 | + return "EARNINGS_RELEASE" | |
| 512 | + if category == "ir" or surface == Surface.INVESTOR_RELATIONS or _INVESTOR_RE.search(title): | |
| 513 | + return "INVESTOR_UPDATE" | |
| 514 | + if category == "changelog" or surface == Surface.CHANGELOG: | |
| 515 | + return "CHANGELOG_ENTRY" | |
| 516 | + if category == "press" or surface == Surface.NEWSROOM: | |
| 517 | + return "NEWS_RELEASE" | |
| 518 | + if category in ("blog", "research") or surface in (Surface.BLOG, Surface.RESEARCH): | |
| 519 | + return "BLOG_POST" | |
| 520 | + return "NEWS_RELEASE" if _PRESS_RE.search(title) else "BLOG_POST" | |
| 521 | + | |
| 522 | + | |
| 523 | +def _news_rules(delta: dict[str, Any], surface: str, evidence: str) -> list[EventDraft]: | |
| 524 | + news = delta.get("news") or {} | |
| 525 | + prefix = {"NEWS_RELEASE": "News release", "BLOG_POST": "Blog post", "CHANGELOG_ENTRY": "Changelog entry", "INVESTOR_UPDATE": "Investor update", | |
| 526 | + "EARNINGS_RELEASE": "Earnings release"} | |
| 527 | + out: list[EventDraft] = [] | |
| 528 | + for item in list(news.get("added") or [])[:NEWS_ITEMS_MAX]: | |
| 529 | + title = (item.get("title") or "").strip() | |
| 530 | + if not title: | |
| 531 | + continue | |
| 532 | + subtype = _news_subtype(item, surface) | |
| 533 | + tags = ["news", subtype.lower()] | |
| 534 | + if _is_ai(title): | |
| 535 | + tags.append("ai") | |
| 536 | + if _LAUNCH_RE.search(title): | |
| 537 | + tags.append("launch") | |
| 538 | + if _ACQ_RE.search(title): | |
| 539 | + tags.append("m&a-mention") | |
| 540 | + if _FUNDING_RE.search(title): | |
| 541 | + tags.append("financing-mention") | |
| 542 | + published = _dt(item.get("published_at")) | |
| 543 | + out.append(EventDraft( | |
| 544 | + subtype=subtype, title=f"{prefix[subtype]}: {title[:110]}", entity_key="news:" + normalize_entity_key(title), evidence=evidence, | |
| 545 | + summary=(item.get("summary") or None), new_value=item.get("url"), | |
| 546 | + entities={"news": [{"title": title, "url": item.get("url"), "published_at": item.get("published_at")}]}, | |
| 547 | + payload={"url": item.get("url"), "category": item.get("category"), "published_at": item.get("published_at")}, tags=tags, | |
| 548 | + published_at=published, effective_at=published, magnitude=0.5 if "launch" in tags else 0.1)) | |
| 549 | + return out | |
| 550 | + | |
| 551 | + | |
| 552 | +def _text_diff_rules(surface: str, change: dict[str, Any], diff: dict[str, Any], delta: dict[str, Any]) -> list[EventDraft]: | |
| 553 | + """Meaningful+ diffs on surfaces without typed extractions.""" | |
| 554 | + n = _blocks_changed(diff, change) | |
| 555 | + sections = _sections(diff) | |
| 556 | + ratio = float(diff.get("text_delta_ratio") or change.get("text_delta_ratio") or 0) | |
| 557 | + kind = change.get("kind") or "" | |
| 558 | + sec_txt = f"{len(sections)} {_plural(len(sections), 'section')} changed" if sections else f"{n} {_plural(n, 'block')} changed" | |
| 559 | + payload = {"blocks_changed": n, "sections": sections, "text_delta_ratio": round(ratio, 4), "similarity": diff.get("similarity")} | |
| 560 | + summary = ("Sections affected: " + ", ".join(sections[:8]) + ".") if sections else None | |
| 561 | + mag = min(1.0, ratio * 2) | |
| 562 | + ek = f"page:{change.get('sensor_id')}" | |
| 563 | + if surface == Surface.LEGAL_TERMS: | |
| 564 | + return [EventDraft("TERMS_CHANGE", f"Terms of service page materially updated ({sec_txt})", ek, summary=summary, payload=payload, | |
| 565 | + tags=["legal", "terms"], magnitude=mag, evidence="text_diff", review="legal_sensitive")] | |
| 566 | + if surface == Surface.LEGAL_PRIVACY: | |
| 567 | + return [EventDraft("PRIVACY_POLICY_CHANGE", f"Privacy policy materially updated ({sec_txt})", ek, summary=summary, payload=payload, | |
| 568 | + tags=["legal", "privacy"], magnitude=mag, evidence="text_diff", review="legal_sensitive")] | |
| 569 | + if surface == Surface.SECURITY: | |
| 570 | + return [EventDraft("SECURITY_UPDATE", f"Security page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["security"], | |
| 571 | + magnitude=mag, evidence="text_diff")] | |
| 572 | + if surface == Surface.HOMEPAGE: | |
| 573 | + out: list[EventDraft] = [] | |
| 574 | + meta = delta.get("meta") or {} | |
| 575 | + tc = meta.get("title_changed") | |
| 576 | + if isinstance(tc, dict) and tc.get("before") and tc.get("after") and tc["before"] != tc["after"]: | |
| 577 | + out.append(EventDraft("MESSAGING_CHANGE", f"Homepage title observed as “{str(tc['after'])[:60]}” (was “{str(tc['before'])[:60]}”)", | |
| 578 | + "homepage:title", old_value=tc["before"], new_value=tc["after"], payload={"field": "title"}, tags=["messaging"], | |
| 579 | + magnitude=0.4, evidence="html")) | |
| 580 | + if kind in (ChangeKind.MAJOR, ChangeKind.CRITICAL): | |
| 581 | + out.append(EventDraft("HOMEPAGE_REDESIGN", f"Homepage materially redesigned ({n} {_plural(n, 'block')} changed, {ratio:.0%} of text)", ek, | |
| 582 | + summary=summary, payload=payload, tags=["website", "homepage"], magnitude=mag, evidence="text_diff")) | |
| 583 | + else: | |
| 584 | + out.append(EventDraft("WEBSITE_CHANGE", f"Homepage content updated ({sec_txt})", ek, summary=summary, payload=payload, | |
| 585 | + tags=["website", "homepage"], magnitude=mag, evidence="text_diff")) | |
| 586 | + return out | |
| 587 | + if surface == Surface.ABOUT: | |
| 588 | + return [EventDraft("WEBSITE_CHANGE", f"About page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["website", "about"], | |
| 589 | + magnitude=mag, evidence="text_diff")] | |
| 590 | + if surface == Surface.API: | |
| 591 | + return [EventDraft("API_CHANGE", f"API reference updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["developer", "api"], | |
| 592 | + magnitude=mag, evidence="text_diff")] | |
| 593 | + if surface in (Surface.DOCS, Surface.DEVELOPER): | |
| 594 | + return [EventDraft("DOC_CHANGE", f"Documentation updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["developer", "docs"], | |
| 595 | + magnitude=mag, evidence="text_diff")] | |
| 596 | + if surface == Surface.CHANGELOG: | |
| 597 | + first = next((d for d in diff.get("added") or [] if (d.get("after") or "").strip()), None) | |
| 598 | + head = (first["after"].strip().split("\n")[0][:90]) if first else None | |
| 599 | + title = f"Changelog entry detected: {head}" if head else f"Changelog updated ({sec_txt})" | |
| 600 | + return [EventDraft("CHANGELOG_ENTRY", title, "changelog:" + normalize_entity_key(head or ek), summary=summary, payload=payload, | |
| 601 | + tags=["developer", "changelog"], magnitude=mag, evidence="text_diff")] | |
| 602 | + if surface == Surface.PRICING: | |
| 603 | + return [] # handled by _pricing_rules (generic PRICING_CHANGE) | |
| 604 | + if surface in (Surface.PRODUCTS, Surface.SERVICES, Surface.SOLUTIONS): | |
| 605 | + return [EventDraft("PRODUCT_UPDATE", f"{_label(surface)[0].upper()}{_label(surface)[1:]} updated ({sec_txt})", ek, summary=summary, | |
| 606 | + payload=payload, tags=["product"], magnitude=mag, evidence="text_diff")] | |
| 607 | + if surface == Surface.INVESTOR_RELATIONS: | |
| 608 | + return [EventDraft("INVESTOR_UPDATE", f"Investor relations page updated ({sec_txt})", ek, summary=summary, payload=payload, | |
| 609 | + tags=["investor-relations"], magnitude=mag, evidence="text_diff")] | |
| 610 | + if surface == Surface.SUSTAINABILITY: | |
| 611 | + return [EventDraft("SUSTAINABILITY_UPDATE", f"Sustainability page updated ({sec_txt})", ek, summary=summary, payload=payload, | |
| 612 | + tags=["sustainability"], magnitude=mag, evidence="text_diff")] | |
| 613 | + if surface == Surface.STATUS: | |
| 614 | + return [EventDraft("OPERATIONS_UPDATE", f"Status page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["operations"], | |
| 615 | + magnitude=mag, evidence="text_diff")] | |
| 616 | + if surface in (Surface.CAREERS, Surface.JOBS_BOARD): | |
| 617 | + return [EventDraft("WEBSITE_CHANGE", f"Careers page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["website", "careers"], | |
| 618 | + magnitude=mag, evidence="text_diff")] | |
| 619 | + if surface == Surface.LEADERSHIP: | |
| 620 | + return [EventDraft("LEADERSHIP_CHANGE", f"Leadership page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["leadership"], | |
| 621 | + magnitude=mag, evidence="text_diff", review="low_confidence")] | |
| 622 | + if surface == Surface.LOCATIONS: | |
| 623 | + return [EventDraft("WEBSITE_CHANGE", f"Locations page updated ({sec_txt})", ek, summary=summary, payload=payload, tags=["website", "locations"], | |
| 624 | + magnitude=mag, evidence="text_diff")] | |
| 625 | + if surface in (Surface.NEWSROOM, Surface.BLOG, Surface.FEED, Surface.RESEARCH): | |
| 626 | + first = next((d for d in diff.get("added") or [] if (d.get("after") or "").strip()), None) | |
| 627 | + head = (first["after"].strip().split("\n")[0][:100]) if first else None | |
| 628 | + if head: | |
| 629 | + subtype = "NEWS_RELEASE" if surface == Surface.NEWSROOM or _PRESS_RE.search(head) else "BLOG_POST" | |
| 630 | + pre = "News release" if subtype == "NEWS_RELEASE" else "Blog post" | |
| 631 | + return [EventDraft(subtype, f"{pre}: {head}", "news:" + normalize_entity_key(head), payload=payload, tags=["news"], magnitude=0.1, | |
| 632 | + evidence="text_diff")] | |
| 633 | + return [EventDraft("WEBSITE_CHANGE", f"{_label(surface)[0].upper()}{_label(surface)[1:]} updated ({sec_txt})", ek, summary=summary, | |
| 634 | + payload=payload, tags=["website"], magnitude=mag, evidence="text_diff")] | |
| 635 | + return [] | |
| 636 | + | |
| 637 | + | |
| 638 | +# ============================================================================================================== derive | |
| 639 | + | |
| 640 | + | |
| 641 | +def derive_events(change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any], *, baseline: dict[str, Any] | None = None, | |
| 642 | + country_names: dict[str, str] | None = None) -> Derived: | |
| 643 | + """Pure rule evaluation for one change. Returns drafts with importance/confidence finalised, plus LLM hints.""" | |
| 644 | + kind = str(change.get("kind") or "") | |
| 645 | + significance = float(change.get("significance") or 0.0) | |
| 646 | + if kind in (ChangeKind.NOISE, ChangeKind.MINOR) or significance < settings.meaningful_threshold and kind not in (ChangeKind.MEANINGFUL, ChangeKind.MAJOR, ChangeKind.CRITICAL): | |
| 647 | + return Derived(events=[]) | |
| 648 | + delta: dict[str, Any] = dict(change.get("structured_delta") or {}) | |
| 649 | + diff: dict[str, Any] = dict(change.get("diff") or {}) | |
| 650 | + surface = str(change.get("surface") or sensor.get("surface") or Surface.OTHER) | |
| 651 | + evidence = evidence_kind({**sensor, "surface": surface}, delta) | |
| 652 | + names = country_names or {} | |
| 653 | + | |
| 654 | + drafts: list[EventDraft] = [] | |
| 655 | + drafts += _hiring_rules(delta, surface, evidence, baseline) | |
| 656 | + drafts += _pricing_rules(delta, surface, evidence, diff, change) | |
| 657 | + drafts += _leadership_rules(delta, surface, evidence) | |
| 658 | + drafts += _product_rules(delta, evidence) | |
| 659 | + drafts += _location_rules(delta, evidence, names) | |
| 660 | + drafts += _news_rules(delta, surface, evidence) | |
| 661 | + structured_hit = bool(drafts) | |
| 662 | + if not structured_hit: | |
| 663 | + drafts += _text_diff_rules(surface, change, diff, delta) | |
| 664 | + | |
| 665 | + for d in drafts: | |
| 666 | + default = EVENT_SUBTYPES.get(d.subtype, (EventType.OTHER, 0.3))[1] | |
| 667 | + d.importance = scale_importance(default, significance, d.magnitude) | |
| 668 | + d.confidence = EVIDENCE_CONFIDENCE.get(d.evidence, EVIDENCE_CONFIDENCE["html"]) | |
| 669 | + d.title = safe_wording(d.title)[:200] | |
| 670 | + if d.summary: | |
| 671 | + d.summary = safe_wording(d.summary)[:600] | |
| 672 | + if kind == ChangeKind.CRITICAL and not d.review: | |
| 673 | + d.review = "major_event" | |
| 674 | + elif d.confidence < 0.5 and not d.review: | |
| 675 | + d.review = "low_confidence" | |
| 676 | + | |
| 677 | + needs = False | |
| 678 | + reason = None | |
| 679 | + if not drafts: | |
| 680 | + needs, reason = True, "no_deterministic_event" | |
| 681 | + elif surface in AMBIGUOUS_SURFACES or (not structured_hit and surface in (Surface.HOMEPAGE, Surface.ABOUT, Surface.PRODUCTS, Surface.SERVICES, Surface.SOLUTIONS)): | |
| 682 | + needs, reason = True, "ambiguous_surface" | |
| 683 | + summarize = sorted({d.subtype for d in drafts if d.subtype in SUMMARY_SUBTYPES and (diff.get("added") or diff.get("modified") or d.summary)}) | |
| 684 | + return Derived(events=drafts, needs_classification=needs, classification_reason=reason, summarize=summarize) | |
| 685 | + | |
| 686 | + | |
| 687 | +def dedupe_key_for(company_id: str, subtype: str, entity_key: str, sensor_id: str | None, day: date) -> str: | |
| 688 | + return stable_hash(company_id, subtype, normalize_entity_key(entity_key), sensor_id or "", day.isoformat(), length=40) | |
| 689 | + | |
| 690 | + | |
| 691 | +# ============================================================================================================== persistence | |
| 692 | + | |
| 693 | + | |
| 694 | +async def _country_names(conn) -> dict[str, str]: # type: ignore[no-untyped-def] | |
| 695 | + rows = await fetch_all(conn, "select code, name from countries") | |
| 696 | + return {str(r["code"]).upper(): r["name"] for r in rows} | |
| 697 | + | |
| 698 | + | |
| 699 | +async def _baseline(conn, company_id: str) -> dict[str, Any]: # type: ignore[no-untyped-def] | |
| 700 | + rows = await fetch_all(conn, "select metric, mean, stddev, samples from baselines where company_id = :c", c=company_id) | |
| 701 | + return {r["metric"]: {"mean": r["mean"], "stddev": r["stddev"], "samples": r["samples"]} for r in rows} | |
| 702 | + | |
| 703 | + | |
| 704 | +async def _llm_budget_left(conn) -> int: # type: ignore[no-untyped-def] | |
| 705 | + used = await fetch_val(conn, "select count(*) from llm_jobs where created_at >= date_trunc('day', now() at time zone 'utc')") | |
| 706 | + return max(0, int(settings.llm_daily_budget) - int(used or 0)) | |
| 707 | + | |
| 708 | + | |
| 709 | +async def _enqueue_llm(conn, *, kind: str, ref_id: str, company_id: str, budget: dict[str, int]) -> bool: # type: ignore[no-untyped-def] | |
| 710 | + if not settings.llm_configured or budget["left"] <= 0: | |
| 711 | + return False | |
| 712 | + exists = await fetch_val(conn, "select 1 from llm_jobs where kind = :k and ref_id = :r and status in ('pending', 'running', 'done')", k=kind, r=ref_id) | |
| 713 | + if exists: | |
| 714 | + return False | |
| 715 | + await execute(conn, "insert into llm_jobs (id, kind, ref_id, company_id, status) values (:id, :k, :r, :c, 'pending')", | |
| 716 | + id=new_id("llm_job"), k=kind, r=ref_id, c=company_id) | |
| 717 | + budget["left"] -= 1 | |
| 718 | + return True | |
| 719 | + | |
| 720 | + | |
| 721 | +async def persist_change_events(conn, change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any], derived: Derived, # type: ignore[no-untyped-def] | |
| 722 | + *, budget: dict[str, int] | None = None, enqueue_llm: bool = True) -> dict[str, Any]: | |
| 723 | + """Insert the drafts for one change (idempotent), cluster them, queue reviews and LLM jobs. Caller owns the transaction.""" | |
| 724 | + detected_at: datetime = change.get("detected_at") or datetime.now(UTC) | |
| 725 | + if detected_at.tzinfo is None: | |
| 726 | + detected_at = detected_at.replace(tzinfo=UTC) | |
| 727 | + day = detected_at.astimezone(UTC).date() | |
| 728 | + created: list[str] = [] | |
| 729 | + duplicates = 0 | |
| 730 | + budget = budget if budget is not None else {"left": await _llm_budget_left(conn)} | |
| 731 | + for d in derived.events: | |
| 732 | + key = dedupe_key_for(company["id"], d.subtype, d.entity_key, change.get("sensor_id"), day) | |
| 733 | + event_id = new_id("event") | |
| 734 | + payload = {**d.payload, "rules_version": RULES_VERSION, "evidence": d.evidence, "significance": change.get("significance"), | |
| 735 | + "change_kind": change.get("kind"), "entity_key": d.entity_key} | |
| 736 | + row = await fetch_one(conn, """ | |
| 737 | + insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, | |
| 738 | + title, summary, old_value, new_value, payload, entities, tags, detected_at, effective_at, published_at, source_url, | |
| 739 | + snapshot_before, snapshot_after, language, origin, schema_version, status, dedupe_key) | |
| 740 | + values (:id, :company_id, :sensor_id, :change_id, :surface, :event_type, :subtype, :importance, :confidence, :label, :title, :summary, | |
| 741 | + :old_value, :new_value, cast(:payload as jsonb), cast(:entities as jsonb), cast(:tags as text[]), :detected_at, :effective_at, | |
| 742 | + :published_at, :source_url, :snap_before, :snap_after, :language, 'deterministic', :schema_version, 'active', :dedupe_key) | |
| 743 | + on conflict (dedupe_key) do nothing | |
| 744 | + returning id""", | |
| 745 | + id=event_id, company_id=company["id"], sensor_id=change.get("sensor_id"), change_id=change["id"], surface=change.get("surface"), | |
| 746 | + event_type=d.event_type, subtype=d.subtype, importance=d.importance, confidence=d.confidence, label=confidence_label(d.confidence), | |
| 747 | + title=d.title, summary=d.summary, old_value=(d.old_value[:400] if d.old_value else None), new_value=(d.new_value[:400] if d.new_value else None), | |
| 748 | + payload=jsonb(payload), entities=jsonb(d.entities), tags=list(dict.fromkeys(d.tags)), detected_at=detected_at, effective_at=d.effective_at, | |
| 749 | + published_at=d.published_at, source_url=sensor.get("url"), snap_before=change.get("snapshot_before"), snap_after=change.get("snapshot_after"), | |
| 750 | + language=((change.get("structured_delta") or {}).get("meta") or {}).get("language"), schema_version=SCHEMA_VERSION, dedupe_key=key) | |
| 751 | + if row is None: | |
| 752 | + continue | |
| 753 | + created.append(event_id) | |
| 754 | + if sensor.get("url"): | |
| 755 | + await execute(conn, """ | |
| 756 | + insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, detected_at, kind) | |
| 757 | + values (:e, :s, :url, :snap, :surface, :at, 'primary') on conflict do nothing""", | |
| 758 | + e=event_id, s=change.get("sensor_id"), url=sensor["url"], snap=change.get("snapshot_after"), surface=change.get("surface"), at=detected_at) | |
| 759 | + ev = {"id": event_id, "company_id": company["id"], "sensor_id": change.get("sensor_id"), "surface": change.get("surface"), "event_type": d.event_type, | |
| 760 | + "event_subtype": d.subtype, "title": d.title, "confidence": d.confidence, "detected_at": detected_at, "source_url": sensor.get("url"), | |
| 761 | + "snapshot_after": change.get("snapshot_after")} | |
| 762 | + _, dup = await attach_to_cluster(conn, ev, entity_key=d.entity_key) | |
| 763 | + duplicates += int(dup) | |
| 764 | + if d.review and not dup: | |
| 765 | + await execute(conn, """ | |
| 766 | + insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, :kind, :ref, :c, cast(:p as jsonb))""", | |
| 767 | + id=new_id("review"), kind=d.review, ref=event_id, c=company["id"], | |
| 768 | + p=jsonb({"event_subtype": d.subtype, "title": d.title, "confidence": d.confidence, "significance": change.get("significance")})) | |
| 769 | + if enqueue_llm and not dup and d.subtype in derived.summarize and float(change.get("significance") or 0) >= settings.llm_min_significance: | |
| 770 | + await _enqueue_llm(conn, kind="summarize_event", ref_id=event_id, company_id=company["id"], budget=budget) | |
| 771 | + | |
| 772 | + if enqueue_llm and derived.needs_classification and float(change.get("significance") or 0) >= settings.llm_min_significance: | |
| 773 | + await _enqueue_llm(conn, kind="classify_change", ref_id=change["id"], company_id=company["id"], budget=budget) | |
| 774 | + | |
| 775 | + await execute(conn, "update changes set status = 'processed', processed_at = now() where id = :id and status <> 'enriched'", id=change["id"]) | |
| 776 | + if created: | |
| 777 | + active = len(created) - duplicates | |
| 778 | + if change.get("sensor_id"): | |
| 779 | + await execute(conn, "update sensors set event_count = event_count + :n where id = :s", n=len(created), s=change["sensor_id"]) | |
| 780 | + await execute(conn, "update companies set last_event_at = greatest(coalesce(last_event_at, cast(:at as timestamptz)), cast(:at as timestamptz)) where id = :c", | |
| 781 | + at=detected_at, c=company["id"]) | |
| 782 | + log.info("events created", extra={"change_id": change["id"], "company": company.get("slug"), "events": len(created), "duplicates": duplicates, | |
| 783 | + "active": active}) | |
| 784 | + return {"created": created, "duplicates": duplicates} | |
| 785 | + | |
| 786 | + | |
| 787 | +_CHANGE_SQL = """ | |
| 788 | + select c.*, s.url as sensor_url, s.connector_id, s.surface as sensor_surface, s.config as sensor_config, | |
| 789 | + co.slug as company_slug, co.display_name, co.country as company_country, co.industries | |
| 790 | + from changes c | |
| 791 | + join sensors s on s.id = c.sensor_id | |
| 792 | + join companies co on co.id = c.company_id | |
| 793 | + where {where} | |
| 794 | + order by c.detected_at | |
| 795 | + limit :limit | |
| 796 | + for update of c skip locked""" | |
| 797 | + | |
| 798 | + | |
| 799 | +def _split(row: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: | |
| 800 | + company = {"id": row["company_id"], "slug": row["company_slug"], "display_name": row["display_name"], "country": row["company_country"], | |
| 801 | + "industries": row.get("industries") or []} | |
| 802 | + sensor = {"id": row["sensor_id"], "url": row["sensor_url"], "connector_id": row["connector_id"], "surface": row["sensor_surface"], | |
| 803 | + "fetch_mode": (row.get("sensor_config") or {}).get("fetch_mode")} | |
| 804 | + change = {k: v for k, v in row.items() if k not in ("sensor_url", "connector_id", "sensor_surface", "sensor_config", "company_slug", "display_name", | |
| 805 | + "company_country", "industries")} | |
| 806 | + return change, company, sensor | |
| 807 | + | |
| 808 | + | |
| 809 | +async def process_pending_changes(limit: int = 200) -> dict[str, int]: | |
| 810 | + """Claim pending changes (SKIP LOCKED), derive events, persist. Returns counters. Noise/minor pending rows are archived.""" | |
| 811 | + stats = {"changes": 0, "events": 0, "duplicates": 0, "archived": 0, "llm_jobs": 0} | |
| 812 | + new_event_ids: list[str] = [] | |
| 813 | + async with transaction() as conn: | |
| 814 | + rows = await fetch_all(conn, _CHANGE_SQL.format(where="c.status = 'pending'"), limit=limit) | |
| 815 | + if not rows: | |
| 816 | + return stats | |
| 817 | + names = await _country_names(conn) | |
| 818 | + budget = {"left": await _llm_budget_left(conn)} | |
| 819 | + start_budget = budget["left"] | |
| 820 | + baselines: dict[str, dict[str, Any]] = {} | |
| 821 | + for row in rows: | |
| 822 | + change, company, sensor = _split(row) | |
| 823 | + if change.get("kind") in (ChangeKind.NOISE, ChangeKind.MINOR): | |
| 824 | + await execute(conn, "update changes set status = 'archived', processed_at = now() where id = :id", id=change["id"]) | |
| 825 | + stats["archived"] += 1 | |
| 826 | + continue | |
| 827 | + if company["id"] not in baselines: | |
| 828 | + baselines[company["id"]] = await _baseline(conn, company["id"]) | |
| 829 | + derived = derive_events(change, company, sensor, baseline=baselines[company["id"]], country_names=names) | |
| 830 | + res = await persist_change_events(conn, change, company, sensor, derived, budget=budget) | |
| 831 | + stats["changes"] += 1 | |
| 832 | + stats["events"] += len(res["created"]) | |
| 833 | + stats["duplicates"] += res["duplicates"] | |
| 834 | + new_event_ids += res["created"] | |
| 835 | + stats["llm_jobs"] = start_budget - budget["left"] | |
| 836 | + if new_event_ids: | |
| 837 | + try: | |
| 838 | + from companyatlas.services.alerts import evaluate_alerts | |
| 839 | + | |
| 840 | + await evaluate_alerts(new_event_ids) | |
| 841 | + except Exception: | |
| 842 | + log.exception("alert evaluation failed") | |
| 843 | + return stats | |
| 844 | + | |
| 845 | + | |
| 846 | +async def reprocess_events(since: datetime, *, limit: int = 5000, company_id: str | None = None) -> dict[str, int]: | |
| 847 | + """Re-run the deterministic rules over already processed changes (no refetch). Dedupe keys make this idempotent; new rules add events.""" | |
| 848 | + stats = {"changes": 0, "events": 0, "duplicates": 0} | |
| 849 | + where = "c.status in ('processed', 'enriched') and c.detected_at >= :since and c.kind in ('meaningful', 'major', 'critical')" | |
| 850 | + params: dict[str, Any] = {"since": since, "limit": limit} | |
| 851 | + if company_id: | |
| 852 | + where += " and c.company_id = :company_id" | |
| 853 | + params["company_id"] = company_id | |
| 854 | + async with transaction() as conn: | |
| 855 | + rows = await fetch_all(conn, _CHANGE_SQL.format(where=where), **params) | |
| 856 | + names = await _country_names(conn) | |
| 857 | + for row in rows: | |
| 858 | + change, company, sensor = _split(row) | |
| 859 | + derived = derive_events(change, company, sensor, baseline=await _baseline(conn, company["id"]), country_names=names) | |
| 860 | + res = await persist_change_events(conn, change, company, sensor, derived, enqueue_llm=False) | |
| 861 | + stats["changes"] += 1 | |
| 862 | + stats["events"] += len(res["created"]) | |
| 863 | + stats["duplicates"] += res["duplicates"] | |
| 864 | + return stats | |
| 865 | + | |
| 866 | + | |
| 867 | +async def list_events(*, company: str | None = None, event_type: str | None = None, limit: int = 50) -> list[dict[str, Any]]: | |
| 868 | + where = ["e.status in ('active', 'review')"] | |
| 869 | + params: dict[str, Any] = {"limit": limit} | |
| 870 | + if company: | |
| 871 | + where.append("(co.slug = :company or co.id = :company)") | |
| 872 | + params["company"] = company | |
| 873 | + if event_type: | |
| 874 | + where.append("(e.event_type = :t or e.event_subtype = :t)") | |
| 875 | + params["t"] = event_type.upper() | |
| 876 | + async with transaction() as conn: | |
| 877 | + return await fetch_all(conn, f""" | |
| 878 | + select e.id, co.slug, e.event_type, e.event_subtype, e.importance, e.confidence, e.confidence_label, e.title, e.detected_at, e.origin, | |
| 879 | + e.surface, e.cluster_id, e.status | |
| 880 | + from events e join companies co on co.id = e.company_id | |
| 881 | + where {' and '.join(where)} order by e.detected_at desc limit :limit""", **params) | |
| 882 | + | |
| 883 | + | |
| 884 | +@periodic("process-changes", every_s=20) | |
| 885 | +async def process_changes_task() -> None: | |
| 886 | + stats = await process_pending_changes(limit=200) | |
| 887 | + if stats["changes"] or stats["archived"]: | |
| 888 | + log.info("process-changes", extra=stats) | |
| 889 | + | |
| 890 | + | |
| 891 | +__all__ = [ | |
| 892 | + "RULES_VERSION", | |
| 893 | + "Derived", | |
| 894 | + "EventDraft", | |
| 895 | + "dedupe_key_for", | |
| 896 | + "derive_events", | |
| 897 | + "evidence_kind", | |
| 898 | + "list_events", | |
| 899 | + "persist_change_events", | |
| 900 | + "process_pending_changes", | |
| 901 | + "reprocess_events", | |
| 902 | + "safe_wording", | |
| 903 | + "scale_importance", | |
| 904 | +] | |
added
src/companyatlas/services/llm/__init__.py
+1 −0
@@ -0,0 +1 @@ | ||
| 1 | +"""LLM enrichment (spec §24–25): provider abstraction, strict JSON schemas, versioned prompts, budgeted worker. Deterministic first.""" | |
added
src/companyatlas/services/llm/ask.py
+275 −0
@@ -0,0 +1,275 @@ | ||
| 1 | +"""`/ask` helper: turn a natural-language question into structured filters — deterministically first (countries, industries, event | |
| 2 | +types/subtypes, windows, intents: hiring / pricing / AI / launch / leadership / expansion / legal / developer), then optionally refined by | |
| 3 | +the small model (`prompts/ask-router/v1.md`). The result is an `interpretation` object the API runs against its own tables; nothing | |
| 4 | +here ever fabricates companies, events or numbers — `build_answer()` only phrases counts the API measured. | |
| 5 | + | |
| 6 | + interp = parse_question("companies hiring AI engineers in Canada last month") | |
| 7 | + interp = await route_question(q) # + LLM refinement when configured | |
| 8 | + answer = build_answer(interp, companies=12, events=48) | |
| 9 | +""" | |
| 10 | +from __future__ import annotations | |
| 11 | + | |
| 12 | +import logging | |
| 13 | +import re | |
| 14 | +from dataclasses import asdict, dataclass, field | |
| 15 | +from datetime import UTC, datetime | |
| 16 | +from typing import Any | |
| 17 | + | |
| 18 | +from companyatlas.config import settings | |
| 19 | +from companyatlas.services.llm.schemas import AskRoute | |
| 20 | +from companyatlas.taxonomy import EVENT_SUBTYPES, EventType | |
| 21 | + | |
| 22 | +log = logging.getLogger(__name__) | |
| 23 | + | |
| 24 | +ASK_PARSER_VERSION = "ask-parser-v1" | |
| 25 | + | |
| 26 | +# Compact built-in geography (the API may pass the full `countries` table for better coverage). | |
| 27 | +COUNTRY_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 | +} | |
| 38 | +REGION_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"} | |
| 40 | + | |
| 41 | +# Intent → (event types, subtypes, tags). Order matters: first match sets the primary intent. | |
| 42 | +INTENTS: 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} | |
| 68 | + | |
| 69 | + | |
| 70 | +@dataclass(slots=True) | |
| 71 | +class Interpretation: | |
| 72 | + question: str | |
| 73 | + 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 = None | |
| 84 | + answer_style: str = "list" | |
| 85 | + min_importance: float | None = None | |
| 86 | + confidence: float = 0.5 | |
| 87 | + source: str = "deterministic" | |
| 88 | + parser_version: str = ASK_PARSER_VERSION | |
| 89 | + llm_model: str | None = None | |
| 90 | + prompt_version: str | None = None | |
| 91 | + | |
| 92 | + def to_dict(self) -> dict[str, Any]: | |
| 93 | + return asdict(self) | |
| 94 | + | |
| 95 | + @property | |
| 96 | + 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} | |
| 101 | + | |
| 102 | + | |
| 103 | +def _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] | |
| 115 | + | |
| 116 | + | |
| 117 | +def _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] | |
| 127 | + | |
| 128 | + | |
| 129 | +def _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 1 | |
| 139 | + if re.search(r"\byesterday\b", q, re.IGNORECASE): | |
| 140 | + return 2 | |
| 141 | + 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 None | |
| 146 | + | |
| 147 | + | |
| 148 | +def _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 | + continue | |
| 155 | + if w not in out: | |
| 156 | + out.append(w) | |
| 157 | + return out[:8] | |
| 158 | + | |
| 159 | + | |
| 160 | +def 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.6 | |
| 193 | + 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.intent | |
| 199 | + 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.intent | |
| 204 | + 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 it | |
| 212 | + | |
| 213 | + | |
| 214 | +async 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 it | |
| 221 | + try: | |
| 222 | + from companyatlas.db import jsonb | |
| 223 | + from companyatlas.services.llm.gateway import get_provider | |
| 224 | + from companyatlas.services.llm.prompts import load_prompt | |
| 225 | + | |
| 226 | + 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.data | |
| 233 | + 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_days | |
| 241 | + if it.intent == "search" and r.intent: | |
| 242 | + it.intent = r.intent | |
| 243 | + if r.answer_style in ("list", "count", "compare", "timeline"): | |
| 244 | + it.answer_style = r.answer_style | |
| 245 | + it.confidence = round(max(it.confidence, min(0.95, r.confidence)), 2) | |
| 246 | + it.source = "llm" | |
| 247 | + it.llm_model = res.model | |
| 248 | + it.prompt_version = prompt.ref | |
| 249 | + except Exception as exc: # noqa: BLE001 — the deterministic parse is always a valid fallback | |
| 250 | + log.warning("ask-router llm refinement skipped", extra={"error": str(exc)[:200]}) | |
| 251 | + return it | |
| 252 | + | |
| 253 | + | |
| 254 | +def 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) | |
| 273 | + | |
| 274 | + | |
| 275 | +__all__ = ["ASK_PARSER_VERSION", "COUNTRY_TERMS", "Interpretation", "build_answer", "parse_question", "route_question"] | |
added
src/companyatlas/services/llm/enrich.py
+357 −0
@@ -0,0 +1,357 @@ | ||
| 1 | +"""LLM enrichment worker (spec §24, §2.4): drains `llm_jobs` with SKIP LOCKED, sticky by kind (one model per stream so the on-demand | |
| 2 | +server does not thrash), builds a bounded context from the change diff and structured deltas, calls the gateway, validates the JSON | |
| 3 | +and writes events / summaries with full provenance (`origin`, model, prompt version, schema version, tokens, latency). | |
| 4 | + | |
| 5 | +Job kinds | |
| 6 | +- classify_change (ref = change id, small model) → new event `origin='llm'` when material, else recorded as non-material. | |
| 7 | +- summarize_event (ref = event id, medium model) → `events.summary` (+ payload.llm), `origin='hybrid'`; LEGAL events use the legal-diff prompt. | |
| 8 | +- classify_industry (ref = company id, small model) → `companies.source_meta.llm_industries` suggestion only (never overwrites registry data). | |
| 9 | + | |
| 10 | +Graceful when `settings.llm_configured` is False: jobs stay pending, the worker logs once and returns. | |
| 11 | +""" | |
| 12 | +from __future__ import annotations | |
| 13 | + | |
| 14 | +import logging | |
| 15 | +from datetime import UTC, datetime | |
| 16 | +from typing import Any | |
| 17 | + | |
| 18 | +from pydantic import BaseModel | |
| 19 | + | |
| 20 | +from companyatlas.config import settings | |
| 21 | +from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 22 | +from companyatlas.ids import new_id, stable_hash | |
| 23 | +from companyatlas.services.clustering import attach_to_cluster | |
| 24 | +from companyatlas.services.llm.gateway import LLMError, LLMNotConfigured, LLMValidationError, get_provider | |
| 25 | +from companyatlas.services.llm.prompts import load_prompt | |
| 26 | +from companyatlas.services.llm.schemas import SCHEMA_VERSIONS, ChangeClassification, EventSummary, IndustryTags, LegalDiffSummary | |
| 27 | +from companyatlas.services.periodic import periodic | |
| 28 | +from companyatlas.taxonomy import EVENT_SUBTYPES, EventType, confidence_label | |
| 29 | + | |
| 30 | +log = logging.getLogger(__name__) | |
| 31 | + | |
| 32 | +PROVIDER_NAME = "openai-compatible" | |
| 33 | +KIND_ORDER = ("classify_change", "summarize_event", "classify_industry") | |
| 34 | +MAX_BLOCKS = 12 | |
| 35 | +MATERIAL_MIN_IMPORTANCE = 0.25 | |
| 36 | +_sticky: dict[str, str | None] = {"kind": None} | |
| 37 | +_warned = {"unconfigured": False} | |
| 38 | + | |
| 39 | + | |
| 40 | +# ---------------------------------------------------------------------------------------------------------------- context | |
| 41 | + | |
| 42 | + | |
| 43 | +def _cut(text: str | None, limit: int) -> str | None: | |
| 44 | + if not text: | |
| 45 | + return None | |
| 46 | + b = text.encode("utf-8") | |
| 47 | + return text if len(b) <= limit else b[:limit].decode("utf-8", "ignore") + "…" | |
| 48 | + | |
| 49 | + | |
| 50 | +def build_change_context(change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any]) -> dict[str, Any]: | |
| 51 | + """Bounded, structured context: company meta, surface, ≤ 12 blocks (≤ 3 kB each side), structured deltas (trimmed lists).""" | |
| 52 | + limit = settings.llm_context_block_bytes | |
| 53 | + diff = change.get("diff") or {} | |
| 54 | + blocks: list[dict[str, Any]] = [] | |
| 55 | + for bucket in ("added", "removed", "modified"): | |
| 56 | + for d in (diff.get(bucket) or [])[:MAX_BLOCKS]: | |
| 57 | + if len(blocks) >= MAX_BLOCKS: | |
| 58 | + break | |
| 59 | + blocks.append({"op": bucket, "kind": d.get("kind"), "path": d.get("path"), "before": _cut(d.get("before"), limit), "after": _cut(d.get("after"), limit)}) | |
| 60 | + delta = change.get("structured_delta") or {} | |
| 61 | + trimmed: dict[str, Any] = {} | |
| 62 | + for key, value in delta.items(): | |
| 63 | + if isinstance(value, dict): | |
| 64 | + trimmed[key] = {k: (v[:10] if isinstance(v, list) else v) for k, v in value.items()} | |
| 65 | + else: | |
| 66 | + trimmed[key] = value | |
| 67 | + return { | |
| 68 | + "company": {"name": company.get("display_name"), "domain": company.get("canonical_domain"), "country": company.get("country"), | |
| 69 | + "industries": list(company.get("industries") or [])[:5], "description": _cut(company.get("description"), 400)}, | |
| 70 | + "surface": change.get("surface") or sensor.get("surface"), "source_url": sensor.get("url"), "detected_at": str(change.get("detected_at")), | |
| 71 | + "significance": change.get("significance"), "change_kind": change.get("kind"), | |
| 72 | + "counts": diff.get("counts") or {"added": change.get("blocks_added"), "removed": change.get("blocks_removed"), "modified": change.get("blocks_modified")}, | |
| 73 | + "blocks": blocks, "structured_delta": trimmed, | |
| 74 | + } | |
| 75 | + | |
| 76 | + | |
| 77 | +async def _load_change_bundle(conn, change_id: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]] | None: # type: ignore[no-untyped-def] | |
| 78 | + change = await fetch_one(conn, "select * from changes where id = :id", id=change_id) | |
| 79 | + if change is None: | |
| 80 | + return None | |
| 81 | + sensor = await fetch_one(conn, "select id, url, surface, connector_id from sensors where id = :id", id=change["sensor_id"]) or {} | |
| 82 | + company = await fetch_one(conn, "select id, slug, display_name, canonical_domain, country, industries, description from companies where id = :id", | |
| 83 | + id=change["company_id"]) or {} | |
| 84 | + return change, company, sensor | |
| 85 | + | |
| 86 | + | |
| 87 | +# ---------------------------------------------------------------------------------------------------------------- claiming | |
| 88 | + | |
| 89 | + | |
| 90 | +async def _budget_left(conn) -> int: # type: ignore[no-untyped-def] | |
| 91 | + used = await fetch_val(conn, "select count(*) from llm_jobs where finished_at >= date_trunc('day', now() at time zone 'utc') and status in ('done', 'failed')") | |
| 92 | + return max(0, settings.llm_daily_budget - int(used or 0)) | |
| 93 | + | |
| 94 | + | |
| 95 | +async def _pick_kind(conn) -> str | None: # type: ignore[no-untyped-def] | |
| 96 | + rows = await fetch_all(conn, "select kind, count(*) as n, min(created_at) as oldest from llm_jobs where status = 'pending' group by kind") | |
| 97 | + if not rows: | |
| 98 | + return None | |
| 99 | + pending = {r["kind"]: r for r in rows} | |
| 100 | + if _sticky["kind"] in pending: | |
| 101 | + return _sticky["kind"] | |
| 102 | + ordered = sorted(rows, key=lambda r: (KIND_ORDER.index(r["kind"]) if r["kind"] in KIND_ORDER else 99, r["oldest"])) | |
| 103 | + _sticky["kind"] = ordered[0]["kind"] | |
| 104 | + return _sticky["kind"] | |
| 105 | + | |
| 106 | + | |
| 107 | +async def claim_jobs(conn, kind: str, limit: int) -> list[dict[str, Any]]: # type: ignore[no-untyped-def] | |
| 108 | + return await fetch_all(conn, """ | |
| 109 | + update llm_jobs set status = 'running', started_at = now(), attempts = attempts + 1 | |
| 110 | + where id in (select id from llm_jobs where status = 'pending' and kind = :kind order by created_at limit :limit for update skip locked) | |
| 111 | + returning *""", kind=kind, limit=limit) | |
| 112 | + | |
| 113 | + | |
| 114 | +async def _finish(conn, job: dict[str, Any], *, status: str, model: str | None, prompt_version: str | None, result: dict[str, Any] | None, # type: ignore[no-untyped-def] | |
| 115 | + error: str | None, req: int = 0, resp: int = 0, latency: int = 0) -> None: | |
| 116 | + await execute(conn, """ | |
| 117 | + update llm_jobs set status = :status, model = :model, prompt_version = :pv, result = cast(:result as jsonb), error = :error, | |
| 118 | + request_tokens = coalesce(request_tokens, 0) + :req, response_tokens = coalesce(response_tokens, 0) + :resp, latency_ms = :latency, | |
| 119 | + finished_at = case when :status in ('done', 'failed') then now() else finished_at end | |
| 120 | + where id = :id""", status=status, model=model, pv=prompt_version, result=jsonb(result) if result is not None else None, error=error, | |
| 121 | + req=req, resp=resp, latency=latency, id=job["id"]) | |
| 122 | + if model and (req or resp): | |
| 123 | + await execute(conn, """ | |
| 124 | + insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0) | |
| 125 | + on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""", key=model, units=float(req + resp)) | |
| 126 | + | |
| 127 | + | |
| 128 | +# ---------------------------------------------------------------------------------------------------------------- handlers | |
| 129 | + | |
| 130 | + | |
| 131 | +async def _handle_classify_change(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] | |
| 132 | + bundle = await _load_change_bundle(conn, job["ref_id"]) | |
| 133 | + if bundle is None: | |
| 134 | + raise LLMError("change not found") | |
| 135 | + change, company, sensor = bundle | |
| 136 | + prompt = load_prompt("change-classifier") | |
| 137 | + context = build_change_context(change, company, sensor) | |
| 138 | + res = await get_provider().complete_json("small", prompt.system, jsonb(context), ChangeClassification, max_tokens=700) | |
| 139 | + c: ChangeClassification = res.data | |
| 140 | + result: dict[str, Any] = {"schema_version": SCHEMA_VERSIONS["ChangeClassification"], "classification": c.model_dump(), "repaired": res.repaired} | |
| 141 | + material = c.is_material and c.event_subtype != "OTHER" and c.importance >= MATERIAL_MIN_IMPORTANCE | |
| 142 | + if material: | |
| 143 | + event_id = await _insert_llm_event(conn, change, company, sensor, c, model=res.model, prompt_version=prompt.ref) | |
| 144 | + result["event_id"] = event_id | |
| 145 | + await execute(conn, "update changes set status = 'enriched' where id = :id", id=change["id"]) | |
| 146 | + return result, res.model, res.request_tokens, res.response_tokens, res.latency_ms | |
| 147 | + | |
| 148 | + | |
| 149 | +async def _insert_llm_event(conn, change: dict[str, Any], company: dict[str, Any], sensor: dict[str, Any], c: ChangeClassification, *, # type: ignore[no-untyped-def] | |
| 150 | + model: str, prompt_version: str) -> str | None: | |
| 151 | + event_type = str(EVENT_SUBTYPES.get(c.event_subtype, (EventType.OTHER, 0.3))[0]) | |
| 152 | + # A deterministic event of the same subtype already exists for this change → enrich it instead of adding a near-duplicate. | |
| 153 | + existing = await fetch_one(conn, """select id, summary from events where change_id = :c and event_subtype = :st and origin in ('deterministic', 'hybrid') | |
| 154 | + and status in ('active', 'review') order by created_at limit 1""", c=change["id"], st=c.event_subtype) | |
| 155 | + if existing is not None: | |
| 156 | + await execute(conn, """update events set summary = coalesce(summary, :summary), origin = 'hybrid', model_provider = :provider, model_name = :model, | |
| 157 | + prompt_version = :pv, payload = payload || cast(:extra as jsonb) where id = :id""", | |
| 158 | + summary=c.summary, provider=PROVIDER_NAME, model=model, pv=prompt_version, id=existing["id"], | |
| 159 | + extra=jsonb({"llm_classification": {"title": c.title, "importance": c.importance, "confidence": c.confidence, "tags": c.tags}})) | |
| 160 | + return None | |
| 161 | + dedupe = stable_hash(company["id"], "llm", change["id"], c.event_subtype, length=40) | |
| 162 | + event_id = new_id("event") | |
| 163 | + detected_at = change.get("detected_at") or datetime.now(UTC) | |
| 164 | + payload = {"llm": {"importance": c.importance, "confidence": c.confidence, "is_material": c.is_material}, "significance": change.get("significance"), | |
| 165 | + "change_kind": change.get("kind"), "entity_key": c.title} | |
| 166 | + row = await fetch_one(conn, """ | |
| 167 | + insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, summary, | |
| 168 | + old_value, new_value, payload, entities, tags, detected_at, source_url, snapshot_before, snapshot_after, language, origin, | |
| 169 | + model_provider, model_name, prompt_version, schema_version, status, dedupe_key) | |
| 170 | + values (:id, :company_id, :sensor_id, :change_id, :surface, :event_type, :subtype, :importance, :confidence, :label, :title, :summary, :old_value, | |
| 171 | + :new_value, cast(:payload as jsonb), cast(:entities as jsonb), cast(:tags as text[]), :detected_at, :source_url, :snap_before, :snap_after, | |
| 172 | + :language, 'llm', :provider, :model, :pv, :sv, :status, :dedupe) | |
| 173 | + on conflict (dedupe_key) do nothing returning id""", | |
| 174 | + id=event_id, company_id=company["id"], sensor_id=change.get("sensor_id"), change_id=change["id"], surface=change.get("surface"), event_type=event_type, | |
| 175 | + subtype=c.event_subtype, importance=round(min(1.0, c.importance), 4), confidence=round(min(0.9, c.confidence), 4), label=confidence_label(min(0.9, c.confidence)), | |
| 176 | + title=c.title, summary=c.summary, old_value=c.old_value, new_value=c.new_value, payload=jsonb(payload), entities=jsonb(c.entities), | |
| 177 | + tags=list(dict.fromkeys(c.tags + ["llm"])), detected_at=detected_at, source_url=sensor.get("url"), snap_before=change.get("snapshot_before"), | |
| 178 | + snap_after=change.get("snapshot_after"), language=c.language, provider=PROVIDER_NAME, model=model, pv=prompt_version, | |
| 179 | + sv=SCHEMA_VERSIONS["ChangeClassification"], status="review" if c.confidence < 0.5 else "active", dedupe=dedupe) | |
| 180 | + if row is None: | |
| 181 | + return None | |
| 182 | + if sensor.get("url"): | |
| 183 | + await execute(conn, """insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, detected_at, kind) | |
| 184 | + values (:e, :s, :u, :snap, :surface, :at, 'primary') on conflict do nothing""", | |
| 185 | + e=event_id, s=change.get("sensor_id"), u=sensor["url"], snap=change.get("snapshot_after"), surface=change.get("surface"), at=detected_at) | |
| 186 | + await attach_to_cluster(conn, {"id": event_id, "company_id": company["id"], "sensor_id": change.get("sensor_id"), "surface": change.get("surface"), | |
| 187 | + "event_type": event_type, "event_subtype": c.event_subtype, "title": c.title, "confidence": min(0.9, c.confidence), | |
| 188 | + "detected_at": detected_at, "source_url": sensor.get("url"), "snapshot_after": change.get("snapshot_after")}, | |
| 189 | + entity_key=c.title) | |
| 190 | + if c.confidence < 0.5: | |
| 191 | + await execute(conn, "insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, 'low_confidence', :ref, :c, cast(:p as jsonb))", | |
| 192 | + id=new_id("review"), ref=event_id, c=company["id"], p=jsonb({"origin": "llm", "title": c.title, "confidence": c.confidence})) | |
| 193 | + await execute(conn, "update sensors set event_count = event_count + 1 where id = :s", s=change.get("sensor_id")) | |
| 194 | + await execute(conn, "update companies set last_event_at = greatest(coalesce(last_event_at, cast(:at as timestamptz)), cast(:at as timestamptz)) where id = :c", | |
| 195 | + at=detected_at, c=company["id"]) | |
| 196 | + return event_id | |
| 197 | + | |
| 198 | + | |
| 199 | +async def _handle_summarize_event(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] | |
| 200 | + event = await fetch_one(conn, "select * from events where id = :id", id=job["ref_id"]) | |
| 201 | + if event is None: | |
| 202 | + raise LLMError("event not found") | |
| 203 | + bundle = await _load_change_bundle(conn, event["change_id"]) if event.get("change_id") else None | |
| 204 | + if bundle is None: | |
| 205 | + raise LLMError("event has no change context") | |
| 206 | + change, company, sensor = bundle | |
| 207 | + context = build_change_context(change, company, sensor) | |
| 208 | + context["event"] = {"type": event["event_type"], "subtype": event["event_subtype"], "title": event["title"], "old_value": event.get("old_value"), | |
| 209 | + "new_value": event.get("new_value"), "entities": event.get("entities") or {}} | |
| 210 | + legal = event["event_type"] == EventType.LEGAL | |
| 211 | + prompt = load_prompt("legal-diff" if legal else "event-summarizer") | |
| 212 | + schema: type[BaseModel] = LegalDiffSummary if legal else EventSummary | |
| 213 | + res = await get_provider().complete_json("medium", prompt.system, jsonb(context), schema, max_tokens=700) | |
| 214 | + data = res.data | |
| 215 | + payload = dict(event.get("payload") or {}) | |
| 216 | + if event.get("summary"): | |
| 217 | + payload.setdefault("summary_prev", event["summary"]) | |
| 218 | + llm_block: dict[str, Any] = {"model": res.model, "prompt_version": prompt.ref, "confidence": getattr(data, "confidence", None), "repaired": res.repaired} | |
| 219 | + if legal: | |
| 220 | + d: LegalDiffSummary = data # type: ignore[assignment] | |
| 221 | + llm_block.update({"materiality": d.materiality, "sections_changed": [s.model_dump() for s in d.sections_changed], "user_impact": d.user_impact}) | |
| 222 | + schema_version = SCHEMA_VERSIONS["LegalDiffSummary"] | |
| 223 | + summary = d.summary | |
| 224 | + tags = list(event.get("tags") or []) + [f"materiality:{d.materiality}"] | |
| 225 | + else: | |
| 226 | + e: EventSummary = data # type: ignore[assignment] | |
| 227 | + llm_block.update({"key_points": e.key_points}) | |
| 228 | + schema_version = SCHEMA_VERSIONS["EventSummary"] | |
| 229 | + summary = e.summary | |
| 230 | + tags = list(event.get("tags") or []) | |
| 231 | + payload["llm"] = llm_block | |
| 232 | + await execute(conn, """ | |
| 233 | + update events set summary = :summary, origin = case when origin = 'deterministic' then 'hybrid' else origin end, model_provider = :provider, | |
| 234 | + model_name = :model, prompt_version = :pv, schema_version = :sv, payload = cast(:payload as jsonb), tags = cast(:tags as text[]), | |
| 235 | + language = coalesce(language, :lang) | |
| 236 | + where id = :id""", summary=summary, provider=PROVIDER_NAME, model=res.model, pv=prompt.ref, sv=schema_version, payload=jsonb(payload), | |
| 237 | + tags=list(dict.fromkeys(tags)), lang=getattr(data, "language", None), id=event["id"]) | |
| 238 | + await execute(conn, "update changes set status = 'enriched' where id = :id and status = 'processed'", id=change["id"]) | |
| 239 | + return {"schema_version": schema_version, "summary": data.model_dump(), "repaired": res.repaired}, res.model, res.request_tokens, res.response_tokens, res.latency_ms | |
| 240 | + | |
| 241 | + | |
| 242 | +async def _handle_classify_industry(conn, job: dict[str, Any]) -> tuple[dict[str, Any], str, int, int, int]: # type: ignore[no-untyped-def] | |
| 243 | + company = await fetch_one(conn, "select * from companies where id = :id", id=job["ref_id"]) | |
| 244 | + if company is None: | |
| 245 | + raise LLMError("company not found") | |
| 246 | + allowed = await fetch_all(conn, "select slug, name from industries order by sort_order, slug") | |
| 247 | + products = await fetch_all(conn, "select name from products where company_id = :c and status = 'listed' order by first_seen_at desc limit 10", c=company["id"]) | |
| 248 | + jobs = await fetch_all(conn, "select title from jobs where company_id = :c and status = 'open' order by first_seen_at desc limit 10", c=company["id"]) | |
| 249 | + home = await fetch_one(conn, """select s.title, s.extracted->'meta' as meta from snapshots s join sensors se on se.id = s.sensor_id | |
| 250 | + where se.company_id = :c and se.surface = 'homepage' order by s.fetched_at desc limit 1""", c=company["id"]) | |
| 251 | + context = {"company": {"name": company["display_name"], "domain": company["canonical_domain"], "description": _cut(company.get("description"), 600)}, | |
| 252 | + "homepage": {"title": (home or {}).get("title"), "meta": (home or {}).get("meta")}, "products": [p["name"] for p in products], | |
| 253 | + "job_titles": [j["title"] for j in jobs], "allowed_industries": [{"slug": r["slug"], "name": r["name"]} for r in allowed]} | |
| 254 | + prompt = load_prompt("industry-tagger") | |
| 255 | + res = await get_provider().complete_json("small", prompt.system, jsonb(context), IndustryTags, max_tokens=400) | |
| 256 | + allowed_slugs = {r["slug"] for r in allowed} | |
| 257 | + tags: IndustryTags = res.data | |
| 258 | + industries = [s for s in tags.industries if s in allowed_slugs] | |
| 259 | + meta = dict(company.get("source_meta") or {}) | |
| 260 | + meta["llm_industries"] = {"industries": industries, "primary": tags.primary if tags.primary in allowed_slugs else None, "keywords": tags.keywords, | |
| 261 | + "confidence": tags.confidence, "model": res.model, "prompt_version": prompt.ref, "at": datetime.now(UTC).isoformat()} | |
| 262 | + await execute(conn, "update companies set source_meta = cast(:m as jsonb), updated_at = now() where id = :id", m=jsonb(meta), id=company["id"]) | |
| 263 | + return {"schema_version": SCHEMA_VERSIONS["IndustryTags"], "industries": meta["llm_industries"]}, res.model, res.request_tokens, res.response_tokens, res.latency_ms | |
| 264 | + | |
| 265 | + | |
| 266 | +HANDLERS = {"classify_change": _handle_classify_change, "summarize_event": _handle_summarize_event, "classify_industry": _handle_classify_industry} | |
| 267 | + | |
| 268 | + | |
| 269 | +# ---------------------------------------------------------------------------------------------------------------- worker | |
| 270 | + | |
| 271 | + | |
| 272 | +async def run_llm_jobs(limit: int = 10) -> dict[str, int]: | |
| 273 | + """Process up to `limit` jobs of one kind. Returns counters. Safe to call concurrently (SKIP LOCKED).""" | |
| 274 | + stats = {"claimed": 0, "done": 0, "failed": 0, "retried": 0, "skipped_budget": 0} | |
| 275 | + if not settings.llm_configured: | |
| 276 | + if not _warned["unconfigured"]: | |
| 277 | + log.info("llm enrichment disabled (CA_LLM_BASE_URL unset or CA_LLM_ENABLED=0); jobs stay pending") | |
| 278 | + _warned["unconfigured"] = True | |
| 279 | + return stats | |
| 280 | + async with transaction() as conn: | |
| 281 | + left = await _budget_left(conn) | |
| 282 | + if left <= 0: | |
| 283 | + stats["skipped_budget"] = 1 | |
| 284 | + return stats | |
| 285 | + kind = await _pick_kind(conn) | |
| 286 | + if kind is None: | |
| 287 | + return stats | |
| 288 | + jobs = await claim_jobs(conn, kind, min(limit, left)) | |
| 289 | + stats["claimed"] = len(jobs) | |
| 290 | + new_events: list[str] = [] | |
| 291 | + for job in jobs: | |
| 292 | + handler = HANDLERS.get(job["kind"]) | |
| 293 | + async with transaction() as conn: | |
| 294 | + if handler is None: | |
| 295 | + await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"unknown kind {job['kind']}") | |
| 296 | + stats["failed"] += 1 | |
| 297 | + continue | |
| 298 | + try: | |
| 299 | + result, model, req, resp, latency = await handler(conn, job) | |
| 300 | + except LLMNotConfigured: | |
| 301 | + await _finish(conn, job, status="pending", model=None, prompt_version=None, result=None, error="not configured") | |
| 302 | + return stats | |
| 303 | + except LLMValidationError as exc: | |
| 304 | + await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=str(exc)[:500]) | |
| 305 | + stats["failed"] += 1 | |
| 306 | + continue | |
| 307 | + except LLMError as exc: | |
| 308 | + retry = exc.retryable and job["attempts"] < settings.llm_job_max_attempts | |
| 309 | + await _finish(conn, job, status="pending" if retry else "failed", model=None, prompt_version=None, result=None, error=str(exc)[:500]) | |
| 310 | + stats["retried" if retry else "failed"] += 1 | |
| 311 | + if exc.retryable and not retry: | |
| 312 | + log.warning("llm job exhausted", extra={"job": job["id"], "error": str(exc)[:200]}) | |
| 313 | + if exc.retryable: | |
| 314 | + break # server is unhealthy: stop the batch, the periodic task returns | |
| 315 | + continue | |
| 316 | + except Exception as exc: | |
| 317 | + log.exception("llm job crashed", extra={"job": job["id"]}) | |
| 318 | + await _finish(conn, job, status="failed", model=None, prompt_version=None, result=None, error=f"{exc.__class__.__name__}: {exc}"[:500]) | |
| 319 | + stats["failed"] += 1 | |
| 320 | + continue | |
| 321 | + pv = None | |
| 322 | + if job["kind"] == "classify_change": | |
| 323 | + pv = load_prompt("change-classifier").ref | |
| 324 | + elif job["kind"] == "summarize_event": | |
| 325 | + pv = None | |
| 326 | + await _finish(conn, job, status="done", model=model, prompt_version=pv, result=result, error=None, req=req, resp=resp, latency=latency) | |
| 327 | + stats["done"] += 1 | |
| 328 | + if result.get("event_id"): | |
| 329 | + new_events.append(result["event_id"]) | |
| 330 | + if new_events: | |
| 331 | + try: | |
| 332 | + from companyatlas.services.alerts import evaluate_alerts | |
| 333 | + | |
| 334 | + await evaluate_alerts(new_events) | |
| 335 | + except Exception: | |
| 336 | + log.exception("alert evaluation failed after llm events") | |
| 337 | + return stats | |
| 338 | + | |
| 339 | + | |
| 340 | +async def enqueue_llm_job(kind: str, ref_id: str, company_id: str | None) -> bool: | |
| 341 | + async with transaction() as conn: | |
| 342 | + exists = await fetch_val(conn, "select 1 from llm_jobs where kind = :k and ref_id = :r and status in ('pending', 'running')", k=kind, r=ref_id) | |
| 343 | + if exists: | |
| 344 | + return False | |
| 345 | + await execute(conn, "insert into llm_jobs (id, kind, ref_id, company_id, status) values (:id, :k, :r, :c, 'pending')", | |
| 346 | + id=new_id("llm_job"), k=kind, r=ref_id, c=company_id) | |
| 347 | + return True | |
| 348 | + | |
| 349 | + | |
| 350 | +@periodic("llm-enrich", every_s=15, initial_delay_s=20) | |
| 351 | +async def llm_enrich_task() -> None: | |
| 352 | + stats = await run_llm_jobs(limit=10) | |
| 353 | + if stats["claimed"]: | |
| 354 | + log.info("llm-enrich", extra=stats) | |
| 355 | + | |
| 356 | + | |
| 357 | +__all__ = ["HANDLERS", "build_change_context", "claim_jobs", "enqueue_llm_job", "run_llm_jobs"] | |
added
src/companyatlas/services/llm/gateway.py
+301 −0
@@ -0,0 +1,301 @@ | ||
| 1 | +"""LLM provider abstraction (spec §25): OpenAI-compatible chat completions with JSON mode, strict pydantic validation, one repair | |
| 2 | +retry, backoff on 429/503 (the MacLustr llm-api.io server loads models on demand and answers 503/429 while swapping), token | |
| 3 | +accounting and health. Keep ONE model per job stream (sticky) to avoid evictions — see `enrich.py`. | |
| 4 | + | |
| 5 | + provider = get_provider() | |
| 6 | + result = await provider.complete_json("small", system, user_json, ChangeClassification, max_tokens=600) | |
| 7 | + result.data → validated pydantic model · result.request_tokens / response_tokens / latency_ms / model / attempts | |
| 8 | +""" | |
| 9 | +from __future__ import annotations | |
| 10 | + | |
| 11 | +import asyncio | |
| 12 | +import json | |
| 13 | +import logging | |
| 14 | +import re | |
| 15 | +import time | |
| 16 | +from abc import ABC, abstractmethod | |
| 17 | +from collections.abc import Awaitable, Callable | |
| 18 | +from dataclasses import dataclass, field | |
| 19 | +from typing import Any | |
| 20 | + | |
| 21 | +import httpx | |
| 22 | +from pydantic import BaseModel, ValidationError | |
| 23 | + | |
| 24 | +from companyatlas.config import settings | |
| 25 | + | |
| 26 | +log = logging.getLogger(__name__) | |
| 27 | + | |
| 28 | +Tier = str # "small" | "medium" | "large" | |
| 29 | +RETRY_STATUS = {408, 425, 429, 500, 502, 503, 504} | |
| 30 | +_JSON_BLOCK = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL) | |
| 31 | +_THINK = re.compile(r"<think>.*?</think>", re.DOTALL) | |
| 32 | + | |
| 33 | + | |
| 34 | +class LLMError(Exception): | |
| 35 | + def __init__(self, message: str, *, status: int | None = None, retryable: bool = False): | |
| 36 | + super().__init__(message) | |
| 37 | + self.status = status | |
| 38 | + self.retryable = retryable | |
| 39 | + | |
| 40 | + | |
| 41 | +class LLMNotConfigured(LLMError): | |
| 42 | + pass | |
| 43 | + | |
| 44 | + | |
| 45 | +class LLMValidationError(LLMError): | |
| 46 | + """Model output could not be parsed/validated even after the repair retry.""" | |
| 47 | + | |
| 48 | + | |
| 49 | +@dataclass(slots=True) | |
| 50 | +class LLMResult[T: BaseModel]: | |
| 51 | + data: T | |
| 52 | + model: str | |
| 53 | + request_tokens: int = 0 | |
| 54 | + response_tokens: int = 0 | |
| 55 | + latency_ms: int = 0 | |
| 56 | + attempts: int = 1 | |
| 57 | + repaired: bool = False | |
| 58 | + raw: str = "" | |
| 59 | + | |
| 60 | + | |
| 61 | +@dataclass(slots=True) | |
| 62 | +class Health: | |
| 63 | + ok: bool | |
| 64 | + base_url: str | |
| 65 | + models: list[str] = field(default_factory=list) | |
| 66 | + latency_ms: int = 0 | |
| 67 | + error: str | None = None | |
| 68 | + | |
| 69 | + | |
| 70 | +class LLMProvider(ABC): | |
| 71 | + """Vendor-neutral interface. Implementations must be safe to share across tasks in one process.""" | |
| 72 | + | |
| 73 | + @abstractmethod | |
| 74 | + def model_for(self, tier: Tier) -> str: ... | |
| 75 | + | |
| 76 | + @abstractmethod | |
| 77 | + async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800, | |
| 78 | + temperature: float = 0.1) -> LLMResult[T]: ... | |
| 79 | + | |
| 80 | + @abstractmethod | |
| 81 | + async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]: ... | |
| 82 | + | |
| 83 | + @abstractmethod | |
| 84 | + async def embed(self, texts: list[str]) -> list[list[float]]: ... | |
| 85 | + | |
| 86 | + @abstractmethod | |
| 87 | + async def health(self) -> Health: ... | |
| 88 | + | |
| 89 | + async def close(self) -> None: | |
| 90 | + return None | |
| 91 | + | |
| 92 | + | |
| 93 | +class OpenAICompatibleProvider(LLMProvider): | |
| 94 | + def __init__(self, *, base_url: str | None = None, api_key: str | None = None, timeout_s: float | None = None, | |
| 95 | + max_tries: int | None = None, backoff_initial_s: float | None = None, backoff_max_s: float | None = None, | |
| 96 | + sleep: Callable[[float], Awaitable[None]] | None = None, transport: httpx.AsyncBaseTransport | None = None): | |
| 97 | + self.base_url = (base_url if base_url is not None else settings.llm_base_url).rstrip("/") | |
| 98 | + self.api_key = api_key if api_key is not None else settings.llm_api_key | |
| 99 | + self.timeout_s = timeout_s or settings.llm_timeout_s | |
| 100 | + self.max_tries = max_tries or settings.llm_max_tries | |
| 101 | + self.backoff_initial_s = backoff_initial_s if backoff_initial_s is not None else settings.llm_backoff_initial_s | |
| 102 | + self.backoff_max_s = backoff_max_s if backoff_max_s is not None else settings.llm_backoff_max_s | |
| 103 | + self._sleep = sleep or asyncio.sleep | |
| 104 | + self._transport = transport | |
| 105 | + self._client: httpx.AsyncClient | None = None | |
| 106 | + self._models = {"small": settings.llm_small_model, "medium": settings.llm_medium_model, "large": settings.llm_large_model, | |
| 107 | + "embedding": settings.llm_embedding_model} | |
| 108 | + | |
| 109 | + # ------------------------------------------------------------------ plumbing | |
| 110 | + @property | |
| 111 | + def configured(self) -> bool: | |
| 112 | + return bool(self.base_url) | |
| 113 | + | |
| 114 | + def model_for(self, tier: Tier) -> str: | |
| 115 | + return self._models.get(tier) or tier # a literal model name is accepted too | |
| 116 | + | |
| 117 | + def _headers(self) -> dict[str, str]: | |
| 118 | + h = {"Content-Type": "application/json", "Accept": "application/json", "User-Agent": "CompanyAtlas-LLM/0.1"} | |
| 119 | + if self.api_key: | |
| 120 | + h["Authorization"] = f"Bearer {self.api_key}" | |
| 121 | + return h | |
| 122 | + | |
| 123 | + def client(self) -> httpx.AsyncClient: | |
| 124 | + if self._client is None: | |
| 125 | + self._client = httpx.AsyncClient(base_url=self.base_url, headers=self._headers(), timeout=httpx.Timeout(self.timeout_s, connect=20), | |
| 126 | + transport=self._transport) | |
| 127 | + return self._client | |
| 128 | + | |
| 129 | + async def close(self) -> None: | |
| 130 | + if self._client is not None: | |
| 131 | + await self._client.aclose() | |
| 132 | + self._client = None | |
| 133 | + | |
| 134 | + def _backoff(self, attempt: int) -> float: | |
| 135 | + return min(self.backoff_max_s, self.backoff_initial_s * (2 ** attempt)) | |
| 136 | + | |
| 137 | + async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]: | |
| 138 | + if not self.configured: | |
| 139 | + raise LLMNotConfigured("CA_LLM_BASE_URL is not set") | |
| 140 | + last: LLMError | None = None | |
| 141 | + for attempt in range(self.max_tries): | |
| 142 | + try: | |
| 143 | + r = await self.client().post(path, json=payload) | |
| 144 | + except httpx.TimeoutException as exc: | |
| 145 | + last = LLMError(f"timeout: {exc}", retryable=True) | |
| 146 | + except httpx.HTTPError as exc: | |
| 147 | + last = LLMError(f"transport: {exc}", retryable=True) | |
| 148 | + else: | |
| 149 | + if r.status_code < 300: | |
| 150 | + try: | |
| 151 | + return r.json() | |
| 152 | + except ValueError as exc: | |
| 153 | + raise LLMError(f"non-JSON response: {exc}", status=r.status_code) from exc | |
| 154 | + body = r.text[:300] | |
| 155 | + if r.status_code in RETRY_STATUS: | |
| 156 | + last = LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code, retryable=True) | |
| 157 | + retry_after = r.headers.get("Retry-After") | |
| 158 | + if retry_after and retry_after.isdigit(): | |
| 159 | + await self._sleep(min(self.backoff_max_s, float(retry_after))) | |
| 160 | + continue | |
| 161 | + else: | |
| 162 | + raise LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code) | |
| 163 | + if attempt + 1 < self.max_tries: | |
| 164 | + delay = self._backoff(attempt) | |
| 165 | + log.warning("llm retry", extra={"attempt": attempt + 1, "delay_s": delay, "error": str(last)}) | |
| 166 | + await self._sleep(delay) | |
| 167 | + raise last or LLMError("unknown LLM failure", retryable=True) | |
| 168 | + | |
| 169 | + # ------------------------------------------------------------------ completions | |
| 170 | + async def _chat(self, model: str, messages: list[dict[str, str]], *, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, int], int]: | |
| 171 | + payload: dict[str, Any] = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature, "stream": False} | |
| 172 | + if json_mode: | |
| 173 | + payload["response_format"] = {"type": "json_object"} | |
| 174 | + t0 = time.monotonic() | |
| 175 | + data = await self._post("/chat/completions", payload) | |
| 176 | + latency = int((time.monotonic() - t0) * 1000) | |
| 177 | + try: | |
| 178 | + content = data["choices"][0]["message"]["content"] or "" | |
| 179 | + except (KeyError, IndexError, TypeError) as exc: | |
| 180 | + raise LLMError(f"malformed completion: {json.dumps(data)[:200]}") from exc | |
| 181 | + usage = data.get("usage") or {} | |
| 182 | + tokens = {"request": int(usage.get("prompt_tokens") or 0), "response": int(usage.get("completion_tokens") or 0)} | |
| 183 | + return str(content), tokens, latency | |
| 184 | + | |
| 185 | + async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]: | |
| 186 | + model = self.model_for(model_tier) | |
| 187 | + content, tokens, latency = await self._chat(model, [{"role": "system", "content": system}, {"role": "user", "content": user}], | |
| 188 | + max_tokens=max_tokens, temperature=temperature, json_mode=False) | |
| 189 | + text = _THINK.sub("", content).strip() | |
| 190 | + return LLMResult(data=text, model=model, request_tokens=tokens["request"], response_tokens=tokens["response"], latency_ms=latency, raw=content) | |
| 191 | + | |
| 192 | + async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800, | |
| 193 | + temperature: float = 0.1) -> LLMResult[T]: | |
| 194 | + model = self.model_for(model_tier) | |
| 195 | + sys_prompt = system.rstrip() + "\n\nRespond with a single JSON object only. JSON schema:\n" + json.dumps(_compact_schema(schema), ensure_ascii=False) | |
| 196 | + messages = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user}] | |
| 197 | + content, tokens, latency = await self._chat(model, messages, max_tokens=max_tokens, temperature=temperature, json_mode=True) | |
| 198 | + req, resp = tokens["request"], tokens["response"] | |
| 199 | + try: | |
| 200 | + return LLMResult(data=parse_json_object(content, schema), model=model, request_tokens=req, response_tokens=resp, latency_ms=latency, raw=content) | |
| 201 | + except (ValueError, ValidationError) as exc: | |
| 202 | + error = str(exc)[:800] | |
| 203 | + log.info("llm output invalid, repairing", extra={"model": model, "schema": schema.__name__, "error": error[:160]}) | |
| 204 | + # one repair retry: show the model its own output and the validation error | |
| 205 | + messages += [{"role": "assistant", "content": content[:4000]}, | |
| 206 | + {"role": "user", "content": "Your previous answer was not valid for the schema. Error:\n" + error + | |
| 207 | + "\nReturn ONLY the corrected JSON object with the same meaning. No prose, no markdown."}] | |
| 208 | + content2, tokens2, latency2 = await self._chat(model, messages, max_tokens=max_tokens, temperature=0.0, json_mode=True) | |
| 209 | + req += tokens2["request"] | |
| 210 | + resp += tokens2["response"] | |
| 211 | + try: | |
| 212 | + data = parse_json_object(content2, schema) | |
| 213 | + except (ValueError, ValidationError) as exc: | |
| 214 | + raise LLMValidationError(f"invalid after repair: {str(exc)[:300]}") from exc | |
| 215 | + return LLMResult(data=data, model=model, request_tokens=req, response_tokens=resp, latency_ms=latency + latency2, attempts=2, repaired=True, raw=content2) | |
| 216 | + | |
| 217 | + async def embed(self, texts: list[str]) -> list[list[float]]: | |
| 218 | + if not texts: | |
| 219 | + return [] | |
| 220 | + data = await self._post("/embeddings", {"model": self.model_for("embedding"), "input": texts}) | |
| 221 | + items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0)) | |
| 222 | + return [list(map(float, d["embedding"])) for d in items] | |
| 223 | + | |
| 224 | + async def health(self) -> Health: | |
| 225 | + if not self.configured: | |
| 226 | + return Health(ok=False, base_url="", error="not configured") | |
| 227 | + t0 = time.monotonic() | |
| 228 | + try: | |
| 229 | + r = await self.client().get("/models") | |
| 230 | + latency = int((time.monotonic() - t0) * 1000) | |
| 231 | + if r.status_code >= 300: | |
| 232 | + return Health(ok=False, base_url=self.base_url, latency_ms=latency, error=f"HTTP {r.status_code}") | |
| 233 | + body = r.json() | |
| 234 | + models = [m.get("id") for m in (body.get("data") or []) if isinstance(m, dict) and m.get("id")] | |
| 235 | + return Health(ok=True, base_url=self.base_url, models=models, latency_ms=latency) | |
| 236 | + except (httpx.HTTPError, ValueError) as exc: | |
| 237 | + return Health(ok=False, base_url=self.base_url, latency_ms=int((time.monotonic() - t0) * 1000), error=str(exc)[:200]) | |
| 238 | + | |
| 239 | + | |
| 240 | +# ---------------------------------------------------------------------------------------------------------------- parsing | |
| 241 | + | |
| 242 | + | |
| 243 | +def extract_json_text(content: str) -> str: | |
| 244 | + """Tolerate reasoning tags, code fences and prose around the object; return the outermost {...} span.""" | |
| 245 | + text = _THINK.sub("", content or "").strip() | |
| 246 | + m = _JSON_BLOCK.search(text) | |
| 247 | + if m: | |
| 248 | + return m.group(1) | |
| 249 | + start = text.find("{") | |
| 250 | + end = text.rfind("}") | |
| 251 | + if start == -1 or end == -1 or end < start: | |
| 252 | + raise ValueError("no JSON object in model output") | |
| 253 | + return text[start:end + 1] | |
| 254 | + | |
| 255 | + | |
| 256 | +def parse_json_object[T: BaseModel](content: str, schema: type[T]) -> T: | |
| 257 | + raw = extract_json_text(content) | |
| 258 | + try: | |
| 259 | + obj = json.loads(raw) | |
| 260 | + except json.JSONDecodeError as exc: | |
| 261 | + raise ValueError(f"invalid JSON: {exc.msg} at {exc.pos}") from exc | |
| 262 | + if isinstance(obj, dict): | |
| 263 | + return schema.model_validate(obj) | |
| 264 | + raise ValueError("top-level JSON value must be an object") | |
| 265 | + | |
| 266 | + | |
| 267 | +def _compact_schema(schema: type[BaseModel]) -> dict[str, Any]: | |
| 268 | + full = schema.model_json_schema() | |
| 269 | + props = {} | |
| 270 | + for name, spec in (full.get("properties") or {}).items(): | |
| 271 | + t = spec.get("type") or ("/".join(x.get("type", "?") for x in spec.get("anyOf", [])) if spec.get("anyOf") else "any") | |
| 272 | + entry: dict[str, Any] = {"type": t} | |
| 273 | + if spec.get("description"): | |
| 274 | + entry["description"] = spec["description"] | |
| 275 | + for k in ("maxLength", "minimum", "maximum", "maxItems"): | |
| 276 | + if k in spec: | |
| 277 | + entry[k] = spec[k] | |
| 278 | + props[name] = entry | |
| 279 | + return {"type": "object", "properties": props, "required": full.get("required", [])} | |
| 280 | + | |
| 281 | + | |
| 282 | +# ---------------------------------------------------------------------------------------------------------------- singleton | |
| 283 | + | |
| 284 | +_provider: LLMProvider | None = None | |
| 285 | + | |
| 286 | + | |
| 287 | +def get_provider() -> LLMProvider: | |
| 288 | + global _provider | |
| 289 | + if _provider is None: | |
| 290 | + _provider = OpenAICompatibleProvider() | |
| 291 | + return _provider | |
| 292 | + | |
| 293 | + | |
| 294 | +def set_provider(provider: LLMProvider | None) -> None: | |
| 295 | + """Tests / alternative vendors.""" | |
| 296 | + global _provider | |
| 297 | + _provider = provider | |
| 298 | + | |
| 299 | + | |
| 300 | +__all__ = ["Health", "LLMError", "LLMNotConfigured", "LLMProvider", "LLMResult", "LLMValidationError", "OpenAICompatibleProvider", "extract_json_text", | |
| 301 | + "get_provider", "parse_json_object", "set_provider"] | |
added
src/companyatlas/services/llm/prompts.py
+58 −0
@@ -0,0 +1,58 @@ | ||
| 1 | +"""Versioned prompt files: `prompts/<task>/v<N>.md` (spec §125 — prompts live in code, versions are recorded on every event/job). | |
| 2 | + | |
| 3 | +A prompt file is Markdown; everything after the optional front-matter comment is the *system* prompt. The user message is always a | |
| 4 | +compact JSON context built by the caller, so prompts never interpolate untrusted text into instructions. | |
| 5 | +""" | |
| 6 | +from __future__ import annotations | |
| 7 | + | |
| 8 | +import re | |
| 9 | +from dataclasses import dataclass | |
| 10 | +from functools import lru_cache | |
| 11 | +from pathlib import Path | |
| 12 | + | |
| 13 | +from companyatlas.config import settings | |
| 14 | + | |
| 15 | +_REPO_ROOT = Path(__file__).resolve().parents[4] | |
| 16 | +_VERSION_RE = re.compile(r"^v(\d+)\.md$") | |
| 17 | + | |
| 18 | + | |
| 19 | +def prompts_dir() -> Path: | |
| 20 | + return Path(settings.prompts_dir) if settings.prompts_dir else _REPO_ROOT / "prompts" | |
| 21 | + | |
| 22 | + | |
| 23 | +@dataclass(frozen=True, slots=True) | |
| 24 | +class Prompt: | |
| 25 | + task: str | |
| 26 | + version: str | |
| 27 | + system: str | |
| 28 | + | |
| 29 | + @property | |
| 30 | + def ref(self) -> str: | |
| 31 | + return f"{self.task}/{self.version}" | |
| 32 | + | |
| 33 | + | |
| 34 | +@lru_cache(maxsize=64) | |
| 35 | +def load_prompt(task: str, version: str | None = None) -> Prompt: | |
| 36 | + """Load `prompts/<task>/<version>.md` (latest `vN` when version is None). Raises FileNotFoundError when missing.""" | |
| 37 | + folder = prompts_dir() / task | |
| 38 | + if version is None: | |
| 39 | + candidates = sorted(((int(m.group(1)), p) for p in folder.glob("v*.md") if (m := _VERSION_RE.match(p.name))), reverse=True) | |
| 40 | + if not candidates: | |
| 41 | + raise FileNotFoundError(f"no prompt versions in {folder}") | |
| 42 | + path = candidates[0][1] | |
| 43 | + version = path.stem | |
| 44 | + else: | |
| 45 | + path = folder / f"{version}.md" | |
| 46 | + text = path.read_text(encoding="utf-8") | |
| 47 | + text = re.sub(r"\A<!--.*?-->\s*", "", text, flags=re.DOTALL) # optional front-matter comment | |
| 48 | + return Prompt(task=task, version=version, system=text.strip()) | |
| 49 | + | |
| 50 | + | |
| 51 | +def available_prompts() -> dict[str, list[str]]: | |
| 52 | + root = prompts_dir() | |
| 53 | + if not root.exists(): | |
| 54 | + return {} | |
| 55 | + return {d.name: sorted(p.stem for p in d.glob("v*.md")) for d in root.iterdir() if d.is_dir()} | |
| 56 | + | |
| 57 | + | |
| 58 | +__all__ = ["Prompt", "available_prompts", "load_prompt", "prompts_dir"] | |
added
src/companyatlas/services/llm/schemas.py
+155 −0
@@ -0,0 +1,155 @@ | ||
| 1 | +"""Strict output schemas for every LLM task (spec §24: structured JSON validated against schemas). The gateway validates model output | |
| 2 | +against these models and retries once with the validation error; anything that still fails is recorded as a failed job — never stored | |
| 3 | +as an event. Schema versions are written to `events.schema_version` / `llm_jobs.result`. | |
| 4 | +""" | |
| 5 | +from __future__ import annotations | |
| 6 | + | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +from pydantic import BaseModel, ConfigDict, Field, field_validator | |
| 10 | + | |
| 11 | +from companyatlas.taxonomy import EVENT_SUBTYPES, FORBIDDEN_WORDING | |
| 12 | + | |
| 13 | +SCHEMA_VERSIONS = {"ChangeClassification": "classify-v1", "EventSummary": "summary-v1", "LegalDiffSummary": "legal-v1", "IndustryTags": "industry-v1", | |
| 14 | + "AskRoute": "ask-v1"} | |
| 15 | + | |
| 16 | + | |
| 17 | +def _clean_text(value: str | None, limit: int) -> str | None: | |
| 18 | + if value is None: | |
| 19 | + return None | |
| 20 | + text = " ".join(str(value).split()) | |
| 21 | + for bad in FORBIDDEN_WORDING: | |
| 22 | + if bad in text.lower(): | |
| 23 | + raise ValueError(f"forbidden wording: {bad!r} — use 'no longer listed' / 'observed' phrasing") | |
| 24 | + return text[:limit] or None | |
| 25 | + | |
| 26 | + | |
| 27 | +class _Strict(BaseModel): | |
| 28 | + model_config = ConfigDict(extra="ignore", str_strip_whitespace=True) | |
| 29 | + | |
| 30 | + | |
| 31 | +class ChangeClassification(_Strict): | |
| 32 | + event_subtype: str = Field(description="One of the Company Atlas event subtypes, or OTHER") | |
| 33 | + importance: float = Field(ge=0, le=1) | |
| 34 | + confidence: float = Field(ge=0, le=1) | |
| 35 | + title: str = Field(min_length=3, max_length=120) | |
| 36 | + summary: str | None = Field(default=None, max_length=400) | |
| 37 | + old_value: str | None = Field(default=None, max_length=200) | |
| 38 | + new_value: str | None = Field(default=None, max_length=200) | |
| 39 | + entities: dict[str, Any] = Field(default_factory=dict) | |
| 40 | + tags: list[str] = Field(default_factory=list, max_length=12) | |
| 41 | + language: str | None = Field(default=None, max_length=8) | |
| 42 | + is_material: bool = True | |
| 43 | + | |
| 44 | + @field_validator("event_subtype", mode="before") | |
| 45 | + @classmethod | |
| 46 | + def _subtype(cls, v: Any) -> str: | |
| 47 | + s = str(v or "OTHER").strip().upper().replace(" ", "_").replace("-", "_") | |
| 48 | + return s if s in EVENT_SUBTYPES else "OTHER" | |
| 49 | + | |
| 50 | + @field_validator("title", mode="before") | |
| 51 | + @classmethod | |
| 52 | + def _title(cls, v: Any) -> str: | |
| 53 | + return _clean_text(str(v or ""), 120) or "" | |
| 54 | + | |
| 55 | + @field_validator("summary", "old_value", "new_value", mode="before") | |
| 56 | + @classmethod | |
| 57 | + def _texts(cls, v: Any) -> str | None: | |
| 58 | + return _clean_text(None if v is None else str(v), 400) | |
| 59 | + | |
| 60 | + @field_validator("tags", mode="before") | |
| 61 | + @classmethod | |
| 62 | + def _tags(cls, v: Any) -> list[str]: | |
| 63 | + if not v: | |
| 64 | + return [] | |
| 65 | + out = [] | |
| 66 | + for t in list(v)[:12]: | |
| 67 | + s = str(t).strip().lower()[:40] | |
| 68 | + if s and s not in out: | |
| 69 | + out.append(s) | |
| 70 | + return out | |
| 71 | + | |
| 72 | + | |
| 73 | +class EventSummary(_Strict): | |
| 74 | + summary: str = Field(min_length=3, max_length=400) | |
| 75 | + key_points: list[str] = Field(default_factory=list, max_length=5) | |
| 76 | + confidence: float = Field(default=0.7, ge=0, le=1) | |
| 77 | + language: str | None = Field(default=None, max_length=8) | |
| 78 | + | |
| 79 | + @field_validator("summary", mode="before") | |
| 80 | + @classmethod | |
| 81 | + def _summary(cls, v: Any) -> str: | |
| 82 | + return _clean_text(str(v or ""), 400) or "" | |
| 83 | + | |
| 84 | + @field_validator("key_points", mode="before") | |
| 85 | + @classmethod | |
| 86 | + def _points(cls, v: Any) -> list[str]: | |
| 87 | + return [p for p in (_clean_text(str(x), 160) for x in (v or [])[:5]) if p] | |
| 88 | + | |
| 89 | + | |
| 90 | +class LegalSection(_Strict): | |
| 91 | + section: str = Field(max_length=120) | |
| 92 | + change: str = Field(max_length=240) | |
| 93 | + | |
| 94 | + | |
| 95 | +class LegalDiffSummary(_Strict): | |
| 96 | + sections_changed: list[LegalSection] = Field(default_factory=list, max_length=12) | |
| 97 | + materiality: str = Field(default="unclear") # editorial | minor | material | unclear | |
| 98 | + summary: str = Field(min_length=3, max_length=400) | |
| 99 | + user_impact: str | None = Field(default=None, max_length=240) | |
| 100 | + confidence: float = Field(default=0.6, ge=0, le=1) | |
| 101 | + | |
| 102 | + @field_validator("materiality", mode="before") | |
| 103 | + @classmethod | |
| 104 | + def _mat(cls, v: Any) -> str: | |
| 105 | + s = str(v or "unclear").strip().lower() | |
| 106 | + return s if s in ("editorial", "minor", "material", "unclear") else "unclear" | |
| 107 | + | |
| 108 | + @field_validator("summary", "user_impact", mode="before") | |
| 109 | + @classmethod | |
| 110 | + def _texts(cls, v: Any) -> str | None: | |
| 111 | + return _clean_text(None if v is None else str(v), 400) | |
| 112 | + | |
| 113 | + | |
| 114 | +class IndustryTags(_Strict): | |
| 115 | + industries: list[str] = Field(default_factory=list, max_length=5) # industry slugs from the registry | |
| 116 | + primary: str | None = None | |
| 117 | + keywords: list[str] = Field(default_factory=list, max_length=10) | |
| 118 | + confidence: float = Field(default=0.6, ge=0, le=1) | |
| 119 | + | |
| 120 | + @field_validator("industries", "keywords", mode="before") | |
| 121 | + @classmethod | |
| 122 | + def _slugs(cls, v: Any) -> list[str]: | |
| 123 | + out: list[str] = [] | |
| 124 | + for x in (v or [])[:10]: | |
| 125 | + s = str(x).strip().lower().replace(" ", "-")[:60] | |
| 126 | + if s and s not in out: | |
| 127 | + out.append(s) | |
| 128 | + return out | |
| 129 | + | |
| 130 | + | |
| 131 | +class AskRoute(_Strict): | |
| 132 | + """LLM refinement of the deterministic question parser (services/llm/ask.py). It may only tighten filters, never invent results.""" | |
| 133 | + intent: str = Field(default="search") # search | hiring | pricing | ai | launch | leadership | expansion | compare | trend | |
| 134 | + countries: list[str] = Field(default_factory=list, max_length=8) | |
| 135 | + industries: list[str] = Field(default_factory=list, max_length=8) | |
| 136 | + event_types: list[str] = Field(default_factory=list, max_length=8) | |
| 137 | + event_subtypes: list[str] = Field(default_factory=list, max_length=8) | |
| 138 | + companies: list[str] = Field(default_factory=list, max_length=8) | |
| 139 | + window_days: int | None = Field(default=None, ge=1, le=3650) | |
| 140 | + keywords: list[str] = Field(default_factory=list, max_length=8) | |
| 141 | + answer_style: str = Field(default="list") # list | count | compare | timeline | |
| 142 | + confidence: float = Field(default=0.6, ge=0, le=1) | |
| 143 | + | |
| 144 | + @field_validator("event_subtypes", mode="before") | |
| 145 | + @classmethod | |
| 146 | + def _subtypes(cls, v: Any) -> list[str]: | |
| 147 | + return [s for s in (str(x).strip().upper() for x in (v or [])[:8]) if s in EVENT_SUBTYPES] | |
| 148 | + | |
| 149 | + @field_validator("countries", mode="before") | |
| 150 | + @classmethod | |
| 151 | + def _countries(cls, v: Any) -> list[str]: | |
| 152 | + return [s for s in (str(x).strip().upper()[:2] for x in (v or [])[:8]) if len(s) == 2 and s.isalpha()] | |
| 153 | + | |
| 154 | + | |
| 155 | +__all__ = ["SCHEMA_VERSIONS", "AskRoute", "ChangeClassification", "EventSummary", "IndustryTags", "LegalDiffSummary", "LegalSection"] | |
added
src/companyatlas/services/metrics.py
+591 −0
@@ -0,0 +1,591 @@ | ||
| 1 | +"""Proprietary metrics (spec §29–37, §91, §146–147, §177): reproducible, coverage-normalised, never fabricated. | |
| 2 | + | |
| 3 | +Every metric row carries `formula_version`, `inputs` (the numbers the formula saw) and `computed_at`. A metric without inputs is not | |
| 4 | +written (no row ≠ 0). Formulas and parameters: `taxonomy.METRIC_PARAMS`, `taxonomy.CCI_WEIGHTS`, documented in docs/SCORING.md. | |
| 5 | + | |
| 6 | + compute_company_metrics(ids | None) hourly (@periodic metrics-hourly): companies active in the last 90 d; nightly: all | |
| 7 | + compute_daily(day) daily aggregates: company_daily, global_daily (activity_index, baseline 100), baselines | |
| 8 | + compute_daily_catch_up() fills missed days since the last computed day | |
| 9 | +""" | |
| 10 | +from __future__ import annotations | |
| 11 | + | |
| 12 | +import logging | |
| 13 | +import math | |
| 14 | +import statistics | |
| 15 | +from dataclasses import dataclass, field | |
| 16 | +from datetime import UTC, date, datetime, timedelta | |
| 17 | +from typing import Any | |
| 18 | + | |
| 19 | +from companyatlas.config import settings | |
| 20 | +from companyatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb, transaction | |
| 21 | +from companyatlas.services.periodic import periodic | |
| 22 | +from companyatlas.taxonomy import ( | |
| 23 | + AI_KEYWORDS, | |
| 24 | + CCI_FORMULA_VERSION, | |
| 25 | + CCI_WEIGHTS, | |
| 26 | + METRIC_PARAMS, | |
| 27 | + METRICS_FORMULA_VERSION, | |
| 28 | + ChangeKind, | |
| 29 | + EventType, | |
| 30 | + Metric, | |
| 31 | + Surface, | |
| 32 | +) | |
| 33 | + | |
| 34 | +log = logging.getLogger(__name__) | |
| 35 | + | |
| 36 | +P = METRIC_PARAMS | |
| 37 | +MEANINGFUL_KINDS = (ChangeKind.MEANINGFUL, ChangeKind.MAJOR, ChangeKind.CRITICAL) | |
| 38 | +CHANGE_WEIGHT = {ChangeKind.MEANINGFUL: P["change_weight_meaningful"], ChangeKind.MAJOR: P["change_weight_major"], ChangeKind.CRITICAL: P["change_weight_critical"]} | |
| 39 | +PRODUCT_SUBTYPES = {"PRODUCT_LAUNCH", "NEW_PRODUCT", "PRODUCT_REMOVED", "PRODUCT_RENAME", "PRODUCT_UPDATE", "FEATURE_LAUNCH", "CHANGELOG_ENTRY", "DOC_CHANGE", | |
| 40 | + "DOCUMENTATION_CHANGE", "API_CHANGE", "API_LAUNCH", "SDK_RELEASE", "AI_LAUNCH"} | |
| 41 | +AI_SUBTYPES = {"AI_HIRING", "AI_LAUNCH"} | |
| 42 | +PRODUCT_SURFACES = {Surface.PRODUCTS, Surface.CHANGELOG, Surface.DOCS, Surface.API, Surface.DEVELOPER, Surface.SERVICES, Surface.SOLUTIONS} | |
| 43 | +DEVELOPER_SURFACES = {Surface.DOCS, Surface.API, Surface.DEVELOPER, Surface.CHANGELOG} | |
| 44 | +COMM_SURFACES = {Surface.NEWSROOM, Surface.BLOG, Surface.FEED, Surface.INVESTOR_RELATIONS, Surface.RESEARCH} | |
| 45 | +GEO_SURFACES = {Surface.LOCATIONS, Surface.CONTACT, Surface.ABOUT, Surface.CAREERS, Surface.JOBS_BOARD} | |
| 46 | +HIRING_SURFACES = {Surface.CAREERS, Surface.JOBS_BOARD} | |
| 47 | +BASELINE_METRICS = ("meaningful_changes_weekly", "jobs_new_weekly", "news_weekly") | |
| 48 | + | |
| 49 | + | |
| 50 | +# ================================================================================================================ pure formulas | |
| 51 | + | |
| 52 | + | |
| 53 | +def saturate(x: float, k: float) -> float: | |
| 54 | + """0–1 saturation: 1 − exp(−x/k). x=k → 0.63, x=3k → 0.95.""" | |
| 55 | + return 0.0 if x <= 0 else 1.0 - math.exp(-x / k) | |
| 56 | + | |
| 57 | + | |
| 58 | +def decay(age_days: float, tau: float) -> float: | |
| 59 | + return math.exp(-max(0.0, age_days) / tau) | |
| 60 | + | |
| 61 | + | |
| 62 | +@dataclass(slots=True) | |
| 63 | +class MetricValue: | |
| 64 | + metric: str | |
| 65 | + value: float | |
| 66 | + confidence: float | |
| 67 | + inputs: dict[str, Any] = field(default_factory=dict) | |
| 68 | + formula_version: str = METRICS_FORMULA_VERSION | |
| 69 | + | |
| 70 | + | |
| 71 | +def activity_score(changes: list[tuple[float, str]], events: list[tuple[float, float]], active_sensors: int) -> MetricValue | None: | |
| 72 | + """changes = [(age_days, kind)], events = [(age_days, importance)] within the 30-day window. | |
| 73 | + raw = Σ w(kind)·e^(−age/τ) + Σ importance·e^(−age/τ); density = raw / sensors^0.5; score = 100·log1p(density)/log1p(D_MAX).""" | |
| 74 | + if active_sensors <= 0 and not changes and not events: | |
| 75 | + return None | |
| 76 | + tau = P["activity_decay_tau_days"] | |
| 77 | + raw_changes = sum(CHANGE_WEIGHT.get(k, 1.0) * decay(a, tau) for a, k in changes) | |
| 78 | + raw_events = sum(float(i) * decay(a, tau) for a, i in events) | |
| 79 | + raw = raw_changes + raw_events | |
| 80 | + denom = max(1.0, float(active_sensors)) ** P["activity_coverage_exp"] | |
| 81 | + density = raw / denom | |
| 82 | + score = min(100.0, 100.0 * math.log1p(density) / math.log1p(P["activity_density_max"])) | |
| 83 | + conf = min(0.95, 0.5 + 0.05 * active_sensors) | |
| 84 | + return MetricValue(Metric.ACTIVITY_SCORE, round(score, 2), round(conf, 2), | |
| 85 | + {"changes_30d": len(changes), "events_30d": len(events), "raw_changes": round(raw_changes, 4), "raw_events": round(raw_events, 4), | |
| 86 | + "active_sensors": active_sensors, "density": round(density, 4), "tau_days": tau, "coverage_exp": P["activity_coverage_exp"], | |
| 87 | + "density_max": P["activity_density_max"]}) | |
| 88 | + | |
| 89 | + | |
| 90 | +def hiring_momentum(open_now: int, open_then: int, *, window: int, extra: dict[str, Any] | None = None, confidence: float = 0.8) -> MetricValue | None: | |
| 91 | + """% change of open listings vs the reconstructed count `window` days ago. Requires ≥ `hiring_min_listings` at the reference point.""" | |
| 92 | + if open_then < P["hiring_min_listings"]: | |
| 93 | + return None | |
| 94 | + pct = (open_now - open_then) / open_then * 100.0 | |
| 95 | + metric = {7: Metric.HIRING_MOMENTUM_7D, 30: Metric.HIRING_MOMENTUM_30D, 90: Metric.HIRING_MOMENTUM_90D}[window] | |
| 96 | + return MetricValue(metric, round(pct, 2), confidence, {"open_now": open_now, "open_then": open_then, "window_days": window, **(extra or {})}) | |
| 97 | + | |
| 98 | + | |
| 99 | +def ai_adoption(*, ai_open: int | None, open_jobs: int | None, ai_events_90d: int, keyword_hits: int, has_text_inputs: bool) -> MetricValue | None: | |
| 100 | + """Observable signals only: share of AI jobs (w 0.5), AI events (w 0.25), AI keyword hits in product names / news / page titles (w 0.25). | |
| 101 | + Weights are renormalised over the inputs that exist.""" | |
| 102 | + comps: dict[str, tuple[float, float]] = {} | |
| 103 | + if open_jobs is not None and open_jobs > 0: | |
| 104 | + comps["jobs"] = (P["ai_jobs_weight"], min(1.0, (ai_open or 0) / open_jobs * 2)) # 50 % AI roles → saturates | |
| 105 | + if ai_events_90d or has_text_inputs or open_jobs: | |
| 106 | + comps["events"] = (P["ai_events_weight"], saturate(ai_events_90d, P["ai_events_saturation"])) | |
| 107 | + if has_text_inputs: | |
| 108 | + comps["keywords"] = (P["ai_keywords_weight"], saturate(keyword_hits, P["ai_keywords_saturation"])) | |
| 109 | + if not comps: | |
| 110 | + return None | |
| 111 | + total_w = sum(w for w, _ in comps.values()) | |
| 112 | + score = 100.0 * sum(w * v for w, v in comps.values()) / total_w | |
| 113 | + return MetricValue(Metric.AI_ADOPTION, round(score, 2), 0.6, {"ai_open": ai_open, "open_jobs": open_jobs, "ai_events_90d": ai_events_90d, | |
| 114 | + "keyword_hits": keyword_hits, "components": {k: round(v, 4) for k, (_, v) in comps.items()}, | |
| 115 | + "weights_used": {k: w for k, (w, _) in comps.items()}}) | |
| 116 | + | |
| 117 | + | |
| 118 | +def saturating_score(metric: str, weighted_sum: float, k: float, inputs: dict[str, Any], confidence: float = 0.7) -> MetricValue: | |
| 119 | + return MetricValue(metric, round(100.0 * saturate(weighted_sum, k), 2), confidence, {**inputs, "weighted_sum": round(weighted_sum, 4), "saturation_k": k}) | |
| 120 | + | |
| 121 | + | |
| 122 | +def corporate_change_index(values: dict[str, float]) -> MetricValue | None: | |
| 123 | + """Σ w·component over available components, renormalised. Hiring momentum (%) is mapped to 0–100 as 50 + clamp(m, −100, 100)/2.""" | |
| 124 | + comps: dict[str, float] = {} | |
| 125 | + for metric in CCI_WEIGHTS: | |
| 126 | + if metric not in values: | |
| 127 | + continue | |
| 128 | + v = values[metric] | |
| 129 | + if metric == Metric.HIRING_MOMENTUM_30D: | |
| 130 | + v = 50.0 + max(-100.0, min(100.0, v)) / 2.0 | |
| 131 | + comps[str(metric)] = max(0.0, min(100.0, v)) | |
| 132 | + if not comps: | |
| 133 | + return None | |
| 134 | + total_w = sum(CCI_WEIGHTS[m] for m in CCI_WEIGHTS if str(m) in comps) | |
| 135 | + value = sum(CCI_WEIGHTS[m] * comps[str(m)] for m in CCI_WEIGHTS if str(m) in comps) / total_w | |
| 136 | + return MetricValue(Metric.CORPORATE_CHANGE_INDEX, round(value, 2), round(0.4 + 0.6 * total_w, 2), | |
| 137 | + {"components": comps, "weights": {str(m): CCI_WEIGHTS[m] for m in CCI_WEIGHTS if str(m) in comps}, "weight_coverage": round(total_w, 3)}, | |
| 138 | + formula_version=CCI_FORMULA_VERSION) | |
| 139 | + | |
| 140 | + | |
| 141 | +def anomaly_score(this_week: int, mean: float, stddev: float, samples: int) -> MetricValue | None: | |
| 142 | + if samples < P["anomaly_min_samples"]: | |
| 143 | + return None | |
| 144 | + sd = max(float(stddev), 0.5) # floor avoids infinite z on flat baselines | |
| 145 | + z = (this_week - mean) / sd | |
| 146 | + return MetricValue(Metric.ANOMALY_SCORE, round(z, 3), min(0.9, 0.4 + 0.05 * samples), {"this_week": this_week, "mean": mean, "stddev": stddev, "samples": samples, | |
| 147 | + "stddev_floor": 0.5}) | |
| 148 | + | |
| 149 | + | |
| 150 | +def historical_coverage(*, observed: int, expected: float, days_with_obs: int, days_since_first: int, surfaces: int) -> MetricValue | None: | |
| 151 | + if expected <= 0 or days_since_first <= 0: | |
| 152 | + return None | |
| 153 | + obs_cov = min(1.0, observed / expected) | |
| 154 | + continuity = min(1.0, days_with_obs / max(1, days_since_first)) | |
| 155 | + source_cov = min(1.0, surfaces / P["coverage_expected_surfaces"]) | |
| 156 | + score = 100.0 * (P["coverage_obs_weight"] * obs_cov + P["coverage_continuity_weight"] * continuity + P["coverage_sources_weight"] * source_cov) | |
| 157 | + return MetricValue(Metric.HISTORICAL_COVERAGE, round(score, 2), 0.8, {"observed": observed, "expected": round(expected, 1), "days_with_obs": days_with_obs, | |
| 158 | + "days_since_first": days_since_first, "surfaces": surfaces}) | |
| 159 | + | |
| 160 | + | |
| 161 | +# ================================================================================================================ per-company loader | |
| 162 | + | |
| 163 | + | |
| 164 | +@dataclass(slots=True) | |
| 165 | +class CompanyContext: | |
| 166 | + company: dict[str, Any] | |
| 167 | + sensors: list[dict[str, Any]] | |
| 168 | + changes_30d: list[dict[str, Any]] | |
| 169 | + events_90d: list[dict[str, Any]] | |
| 170 | + jobs: dict[str, Any] | |
| 171 | + products: dict[str, Any] | |
| 172 | + news: dict[str, Any] | |
| 173 | + locations: dict[str, Any] | |
| 174 | + titles_ai_hits: int | |
| 175 | + baselines: dict[str, dict[str, Any]] | |
| 176 | + observations: dict[str, Any] | |
| 177 | + now: datetime | |
| 178 | + | |
| 179 | + | |
| 180 | +async def load_context(conn, company_id: str, now: datetime) -> CompanyContext | None: # type: ignore[no-untyped-def] | |
| 181 | + company = await fetch_one(conn, "select id, slug, first_observed_at, industries, country from companies where id = :id", id=company_id) | |
| 182 | + if company is None: | |
| 183 | + return None | |
| 184 | + d30, d90 = now - timedelta(days=30), now - timedelta(days=90) | |
| 185 | + sensors = await fetch_all(conn, """select id, surface, status, base_interval_s, current_interval_s, created_at, observation_count, last_success_at, meaningful_change_count | |
| 186 | + from sensors where company_id = :c and retired_at is null""", c=company_id) | |
| 187 | + changes = await fetch_all(conn, """select detected_at, kind, surface from changes where company_id = :c and detected_at >= :since | |
| 188 | + and kind in ('meaningful', 'major', 'critical')""", c=company_id, since=d30) | |
| 189 | + events = await fetch_all(conn, """select detected_at, importance, event_type, event_subtype, tags, surface from events where company_id = :c | |
| 190 | + and detected_at >= :since and status in ('active', 'review') order by detected_at desc limit 3000""", c=company_id, since=d90) | |
| 191 | + jobs_row = await fetch_one(conn, """ | |
| 192 | + select count(*) filter (where status = 'open') as open_now, | |
| 193 | + count(*) filter (where status = 'open' and is_ai) as ai_open, | |
| 194 | + count(*) filter (where status = 'open' and remote) as remote_open, | |
| 195 | + count(*) filter (where first_seen_at <= :t7 and (removed_at is null or removed_at > :t7)) as open_7, | |
| 196 | + count(*) filter (where first_seen_at <= :t30 and (removed_at is null or removed_at > :t30)) as open_30, | |
| 197 | + count(*) filter (where first_seen_at <= :t90 and (removed_at is null or removed_at > :t90)) as open_90, | |
| 198 | + count(*) filter (where first_seen_at > :t30) as new_30d, | |
| 199 | + count(*) filter (where removed_at > :t30) as removed_30d, | |
| 200 | + count(*) filter (where first_seen_at > :t30 and is_ai) as ai_new_30d, | |
| 201 | + count(*) as total | |
| 202 | + from jobs where company_id = :c""", c=company_id, t7=now - timedelta(days=7), t30=d30, t90=d90) | |
| 203 | + by_country = await fetch_all(conn, "select country, count(*) as n from jobs where company_id = :c and status = 'open' and country is not null group by country order by n desc limit 20", c=company_id) | |
| 204 | + by_department = await fetch_all(conn, "select department, count(*) as n from jobs where company_id = :c and status = 'open' and department is not null group by department order by n desc limit 20", c=company_id) | |
| 205 | + new_job_countries = await fetch_all(conn, """ | |
| 206 | + select country from jobs where company_id = :c and country is not null group by country having min(first_seen_at) > :since""", c=company_id, since=d90) | |
| 207 | + products = await fetch_one(conn, """select count(*) filter (where status = 'listed') as listed, count(*) as total, | |
| 208 | + array_agg(name) filter (where status = 'listed') as names from products where company_id = :c""", c=company_id) | |
| 209 | + news = await fetch_one(conn, """select count(*) filter (where coalesce(published_at, first_seen_at) >= :d30) as n_30d, | |
| 210 | + count(*) as total, array_agg(title) filter (where coalesce(published_at, first_seen_at) >= :d90) as titles_90d | |
| 211 | + from news_items where company_id = :c""", c=company_id, d30=d30, d90=d90) | |
| 212 | + locations = await fetch_one(conn, """ | |
| 213 | + select count(*) filter (where status = 'listed') as listed, count(*) filter (where first_seen_at > :since) as new_90d, count(*) as total, | |
| 214 | + count(distinct country) filter (where status = 'listed' and country is not null) as countries | |
| 215 | + from locations where company_id = :c""", c=company_id, since=d90) | |
| 216 | + new_loc_countries = await fetch_all(conn, "select country from locations where company_id = :c and country is not null group by country having min(first_seen_at) > :since", | |
| 217 | + c=company_id, since=d90) | |
| 218 | + titles = await fetch_all(conn, """select distinct on (s.sensor_id) s.title from snapshots s join sensors se on se.id = s.sensor_id | |
| 219 | + where se.company_id = :c order by s.sensor_id, s.fetched_at desc""", c=company_id) | |
| 220 | + baselines = await fetch_all(conn, "select metric, mean, stddev, samples from baselines where company_id = :c", c=company_id) | |
| 221 | + obs = await fetch_one(conn, """select count(*) filter (where failure_class is null) as ok, count(*) as total, count(distinct date(fetched_at)) as days_with_obs, | |
| 222 | + count(distinct sensor_id) as sensors_observed from observations where company_id = :c""", c=company_id) | |
| 223 | + this_week = await fetch_val(conn, "select count(*) from changes where company_id = :c and detected_at >= :since and kind in ('meaningful', 'major', 'critical')", | |
| 224 | + c=company_id, since=now - timedelta(days=7)) | |
| 225 | + jobs = {**(jobs_row or {}), "by_country": [{"country": r["country"], "n": r["n"]} for r in by_country], | |
| 226 | + "by_department": [{"department": r["department"], "n": r["n"]} for r in by_department], "new_countries_90d": [r["country"] for r in new_job_countries]} | |
| 227 | + return CompanyContext( | |
| 228 | + company=company, sensors=sensors, changes_30d=changes, events_90d=events, jobs=jobs, | |
| 229 | + products={**(products or {}), "names": list((products or {}).get("names") or [])}, | |
| 230 | + news={**(news or {}), "titles_90d": list((news or {}).get("titles_90d") or [])}, | |
| 231 | + locations={**(locations or {}), "new_countries_90d": [r["country"] for r in new_loc_countries]}, | |
| 232 | + titles_ai_hits=sum(1 for t in titles if _ai_hit(t.get("title"))), | |
| 233 | + baselines={r["metric"]: r for r in baselines}, observations={**(obs or {}), "this_week_meaningful": int(this_week or 0)}, now=now) | |
| 234 | + | |
| 235 | + | |
| 236 | +def _ai_hit(text: str | None) -> bool: | |
| 237 | + if not text: | |
| 238 | + return False | |
| 239 | + hay = f" {text.lower()} " | |
| 240 | + return any(k in hay for k in AI_KEYWORDS) | |
| 241 | + | |
| 242 | + | |
| 243 | +def _age(now: datetime, at: datetime) -> float: | |
| 244 | + if at.tzinfo is None: | |
| 245 | + at = at.replace(tzinfo=UTC) | |
| 246 | + return max(0.0, (now - at).total_seconds() / 86400.0) | |
| 247 | + | |
| 248 | + | |
| 249 | +# ================================================================================================================ compute | |
| 250 | + | |
| 251 | + | |
| 252 | +def compute_from_context(ctx: CompanyContext) -> list[MetricValue]: | |
| 253 | + now = ctx.now | |
| 254 | + out: list[MetricValue] = [] | |
| 255 | + active_sensors = [s for s in ctx.sensors if s["status"] in ("active", "failing", "stale")] | |
| 256 | + surfaces = {s["surface"] for s in ctx.sensors} | |
| 257 | + n_active = len(active_sensors) | |
| 258 | + events_30 = [e for e in ctx.events_90d if _age(now, e["detected_at"]) <= P["activity_window_days"]] | |
| 259 | + | |
| 260 | + # activity ---------------------------------------------------------------------------------------------------- | |
| 261 | + act = activity_score([(_age(now, c["detected_at"]), str(c["kind"])) for c in ctx.changes_30d], [(_age(now, e["detected_at"]), float(e["importance"])) for e in events_30], n_active) | |
| 262 | + if act and (ctx.changes_30d or events_30 or n_active): | |
| 263 | + out.append(act) | |
| 264 | + | |
| 265 | + # hiring ------------------------------------------------------------------------------------------------------ | |
| 266 | + j = ctx.jobs | |
| 267 | + has_jobs_source = bool(surfaces & HIRING_SURFACES) or int(j.get("total") or 0) > 0 | |
| 268 | + jobs_conf = 0.9 if Surface.JOBS_BOARD in surfaces else 0.75 | |
| 269 | + if has_jobs_source: | |
| 270 | + open_now = int(j.get("open_now") or 0) | |
| 271 | + remote_ratio = round(int(j.get("remote_open") or 0) / open_now, 4) if open_now else None | |
| 272 | + extra = {"jobs_new_30d": int(j.get("new_30d") or 0), "jobs_removed_30d": int(j.get("removed_30d") or 0), "remote_ratio": remote_ratio, | |
| 273 | + "by_country": j.get("by_country"), "by_department": j.get("by_department"), "ai_open": int(j.get("ai_open") or 0)} | |
| 274 | + out.append(MetricValue(Metric.OPEN_JOBS, float(open_now), jobs_conf, extra)) | |
| 275 | + for window, key in ((7, "open_7"), (30, "open_30"), (90, "open_90")): | |
| 276 | + m = hiring_momentum(open_now, int(j.get(key) or 0), window=window, extra=extra if window == 30 else None, confidence=jobs_conf) | |
| 277 | + if m: | |
| 278 | + out.append(m) | |
| 279 | + | |
| 280 | + # AI adoption ------------------------------------------------------------------------------------------------- | |
| 281 | + ai_events = sum(1 for e in ctx.events_90d if e["event_subtype"] in AI_SUBTYPES or "ai" in (e.get("tags") or [])) | |
| 282 | + kw_hits = sum(1 for n in ctx.products["names"] if _ai_hit(n)) + sum(1 for t in ctx.news["titles_90d"] if _ai_hit(t)) + ctx.titles_ai_hits | |
| 283 | + has_text = bool(ctx.products["names"] or ctx.news["titles_90d"] or ctx.sensors) | |
| 284 | + ai = ai_adoption(ai_open=int(j.get("ai_open") or 0) if has_jobs_source else None, open_jobs=int(j.get("open_now") or 0) if has_jobs_source else None, | |
| 285 | + ai_events_90d=ai_events, keyword_hits=kw_hits, has_text_inputs=has_text) | |
| 286 | + if ai: | |
| 287 | + out.append(ai) | |
| 288 | + | |
| 289 | + # product velocity / developer / communication / pricing / leadership ------------------------------------------ | |
| 290 | + tau = P["activity_decay_tau_days"] * 3 # 90-day windows decay slower | |
| 291 | + if surfaces & PRODUCT_SURFACES or any(e["event_type"] == EventType.PRODUCT for e in ctx.events_90d): | |
| 292 | + prod_events = [e for e in ctx.events_90d if e["event_subtype"] in PRODUCT_SUBTYPES] | |
| 293 | + s = sum(float(e["importance"]) * decay(_age(now, e["detected_at"]), tau) for e in prod_events) | |
| 294 | + out.append(saturating_score(Metric.PRODUCT_VELOCITY, s, P["velocity_saturation"], {"events_90d": len(prod_events), "surfaces": sorted(surfaces & PRODUCT_SURFACES)})) | |
| 295 | + if surfaces & DEVELOPER_SURFACES or any(e["event_type"] == EventType.DEVELOPER for e in ctx.events_90d): | |
| 296 | + dev_events = [e for e in ctx.events_90d if e["event_type"] == EventType.DEVELOPER] | |
| 297 | + dev_changes = [c for c in ctx.changes_30d if c["surface"] in DEVELOPER_SURFACES] | |
| 298 | + s = sum(float(e["importance"]) * decay(_age(now, e["detected_at"]), tau) for e in dev_events) + 0.5 * len(dev_changes) | |
| 299 | + out.append(saturating_score(Metric.DEVELOPER_MOMENTUM, s, P["developer_saturation"], {"events_90d": len(dev_events), "meaningful_changes_30d": len(dev_changes), | |
| 300 | + "surfaces": sorted(surfaces & DEVELOPER_SURFACES)})) | |
| 301 | + if surfaces & COMM_SURFACES or int(ctx.news.get("total") or 0) > 0: | |
| 302 | + comm_events = [e for e in events_30 if e["event_type"] in (EventType.COMMUNICATION, EventType.INVESTOR_RELATIONS)] | |
| 303 | + n_news = int(ctx.news.get("n_30d") or 0) | |
| 304 | + s = float(max(n_news, len(comm_events))) + 0.5 * min(n_news, len(comm_events)) | |
| 305 | + out.append(saturating_score(Metric.COMMUNICATION_ACTIVITY, s, P["communication_saturation"], {"news_items_30d": n_news, "events_30d": len(comm_events)})) | |
| 306 | + if Surface.PRICING in surfaces or any(e["event_type"] == EventType.PRICING for e in ctx.events_90d): | |
| 307 | + pr_events = [e for e in ctx.events_90d if e["event_type"] == EventType.PRICING] | |
| 308 | + pr_changes = [c for c in ctx.changes_30d if c["surface"] == Surface.PRICING] | |
| 309 | + s = sum(float(e["importance"]) * decay(_age(now, e["detected_at"]), tau) for e in pr_events) + 0.25 * len(pr_changes) | |
| 310 | + out.append(saturating_score(Metric.PRICING_ACTIVITY, s, P["pricing_saturation"], {"events_90d": len(pr_events), "meaningful_changes_30d": len(pr_changes)})) | |
| 311 | + if Surface.LEADERSHIP in surfaces or any(e["event_type"] == EventType.LEADERSHIP for e in ctx.events_90d): | |
| 312 | + ld_events = [e for e in ctx.events_90d if e["event_type"] == EventType.LEADERSHIP] | |
| 313 | + s = sum(float(e["importance"]) * decay(_age(now, e["detected_at"]), tau) for e in ld_events) | |
| 314 | + out.append(saturating_score(Metric.LEADERSHIP_ACTIVITY, s, P["leadership_saturation"], {"events_90d": len(ld_events)})) | |
| 315 | + | |
| 316 | + # geographic expansion ---------------------------------------------------------------------------------------- | |
| 317 | + loc = ctx.locations | |
| 318 | + if surfaces & GEO_SURFACES or int(loc.get("total") or 0) > 0 or has_jobs_source: | |
| 319 | + new_countries = set(loc.get("new_countries_90d") or []) | set(j.get("new_countries_90d") or []) | |
| 320 | + geo_events = [e for e in ctx.events_90d if e["event_type"] == EventType.LOCATION] | |
| 321 | + s = float(int(loc.get("new_90d") or 0)) + P["geo_country_weight"] * len(new_countries) + sum(float(e["importance"]) for e in geo_events if e["event_subtype"] == "COUNTRY_EXPANSION") | |
| 322 | + out.append(saturating_score(Metric.GEO_EXPANSION, s, P["geo_saturation"], {"new_locations_90d": int(loc.get("new_90d") or 0), "new_countries_90d": sorted(new_countries), | |
| 323 | + "countries_listed": int(loc.get("countries") or 0), "location_events_90d": len(geo_events)})) | |
| 324 | + | |
| 325 | + # CCI --------------------------------------------------------------------------------------------------------- | |
| 326 | + cci = corporate_change_index({m.metric: m.value for m in out}) | |
| 327 | + if cci: | |
| 328 | + out.append(cci) | |
| 329 | + | |
| 330 | + # anomaly ----------------------------------------------------------------------------------------------------- | |
| 331 | + b = ctx.baselines.get("meaningful_changes_weekly") | |
| 332 | + if b: | |
| 333 | + an = anomaly_score(int(ctx.observations.get("this_week_meaningful") or 0), float(b["mean"]), float(b["stddev"]), int(b["samples"])) | |
| 334 | + if an: | |
| 335 | + out.append(an) | |
| 336 | + | |
| 337 | + # historical coverage ----------------------------------------------------------------------------------------- | |
| 338 | + first = ctx.company.get("first_observed_at") | |
| 339 | + if first and ctx.sensors: | |
| 340 | + days_since_first = max(1, int(_age(now, first))) | |
| 341 | + expected = 0.0 | |
| 342 | + for s in ctx.sensors: | |
| 343 | + created = s.get("created_at") or first | |
| 344 | + span_s = max(0.0, (now - (created if created.tzinfo else created.replace(tzinfo=UTC))).total_seconds()) | |
| 345 | + expected += span_s / max(900, int(s.get("current_interval_s") or s.get("base_interval_s") or 86400)) | |
| 346 | + hc = historical_coverage(observed=int(ctx.observations.get("ok") or 0), expected=expected, days_with_obs=int(ctx.observations.get("days_with_obs") or 0), | |
| 347 | + days_since_first=days_since_first, surfaces=len(surfaces)) | |
| 348 | + if hc: | |
| 349 | + out.append(hc) | |
| 350 | + return out | |
| 351 | + | |
| 352 | + | |
| 353 | +async def write_metrics(conn, company_id: str, values: list[MetricValue], now: datetime) -> None: # type: ignore[no-untyped-def] | |
| 354 | + day = now.astimezone(UTC).date() | |
| 355 | + for m in values: | |
| 356 | + await execute(conn, """ | |
| 357 | + insert into metrics_current (company_id, metric, value, confidence, inputs, formula_version, computed_at) | |
| 358 | + values (:c, :m, :v, :conf, cast(:inputs as jsonb), :fv, :at) | |
| 359 | + on conflict (company_id, metric) do update set value = excluded.value, confidence = excluded.confidence, inputs = excluded.inputs, | |
| 360 | + formula_version = excluded.formula_version, computed_at = excluded.computed_at""", | |
| 361 | + c=company_id, m=str(m.metric), v=float(m.value), conf=float(m.confidence), inputs=jsonb(m.inputs), fv=m.formula_version, at=now) | |
| 362 | + await execute(conn, """ | |
| 363 | + insert into metric_series (company_id, metric, day, value, confidence, formula_version) values (:c, :m, :d, :v, :conf, :fv) | |
| 364 | + on conflict (company_id, metric, day) do update set value = excluded.value, confidence = excluded.confidence, formula_version = excluded.formula_version""", | |
| 365 | + c=company_id, m=str(m.metric), d=day, v=float(m.value), conf=float(m.confidence), fv=m.formula_version) | |
| 366 | + | |
| 367 | + | |
| 368 | +async def compute_company_metrics(company_ids: list[str] | None = None, *, all_companies: bool = False, now: datetime | None = None) -> dict[str, int]: | |
| 369 | + """Compute and store metrics. `company_ids=None` → companies with activity in the last `metrics_active_window_days` (or all).""" | |
| 370 | + now = now or datetime.now(UTC) | |
| 371 | + stats = {"companies": 0, "metrics": 0, "skipped": 0} | |
| 372 | + async with transaction() as conn: | |
| 373 | + if company_ids is None: | |
| 374 | + if all_companies: | |
| 375 | + rows = await fetch_all(conn, "select id from companies where onboarding_status <> 'no_website' order by importance desc") | |
| 376 | + else: | |
| 377 | + rows = await fetch_all(conn, """ | |
| 378 | + select id from companies where greatest(coalesce(last_observed_at, 'epoch'), coalesce(last_change_at, 'epoch'), coalesce(last_event_at, 'epoch')) >= :since | |
| 379 | + or id in (select distinct company_id from changes where detected_at >= :since) | |
| 380 | + or id in (select distinct company_id from events where detected_at >= :since) | |
| 381 | + order by importance desc""", since=now - timedelta(days=settings.metrics_active_window_days)) | |
| 382 | + company_ids = [r["id"] for r in rows] | |
| 383 | + for cid in company_ids: | |
| 384 | + try: | |
| 385 | + async with transaction() as conn: | |
| 386 | + ctx = await load_context(conn, cid, now) | |
| 387 | + if ctx is None: | |
| 388 | + stats["skipped"] += 1 | |
| 389 | + continue | |
| 390 | + values = compute_from_context(ctx) | |
| 391 | + if not values: | |
| 392 | + stats["skipped"] += 1 | |
| 393 | + continue | |
| 394 | + await write_metrics(conn, cid, values, now) | |
| 395 | + stats["companies"] += 1 | |
| 396 | + stats["metrics"] += len(values) | |
| 397 | + except Exception: | |
| 398 | + log.exception("metrics failed", extra={"company_id": cid}) | |
| 399 | + return stats | |
| 400 | + | |
| 401 | + | |
| 402 | +# ================================================================================================================ daily aggregates | |
| 403 | + | |
| 404 | + | |
| 405 | +def _day_bounds(day: date) -> tuple[datetime, datetime]: | |
| 406 | + d0 = datetime(day.year, day.month, day.day, tzinfo=UTC) | |
| 407 | + return d0, d0 + timedelta(days=1) | |
| 408 | + | |
| 409 | + | |
| 410 | +async def compute_daily(day: date) -> dict[str, Any]: | |
| 411 | + """company_daily + global_daily (+ baselines) for one UTC day. Idempotent (upserts).""" | |
| 412 | + d0, d1 = _day_bounds(day) | |
| 413 | + async with transaction() as conn: | |
| 414 | + per: dict[str, dict[str, Any]] = {} | |
| 415 | + | |
| 416 | + def bucket(cid: str) -> dict[str, Any]: | |
| 417 | + return per.setdefault(cid, {"observations": 0, "changes": 0, "meaningful_changes": 0, "events": 0, "events_by_type": {}, "jobs_open": None, | |
| 418 | + "jobs_new": 0, "jobs_removed": 0, "jobs_ai_open": None, "news_items": 0, "sensors_active": None}) | |
| 419 | + | |
| 420 | + for r in await fetch_all(conn, "select company_id, count(*) as n from observations where fetched_at >= :d0 and fetched_at < :d1 group by 1", d0=d0, d1=d1): | |
| 421 | + bucket(r["company_id"])["observations"] = int(r["n"]) | |
| 422 | + for r in await fetch_all(conn, """select company_id, count(*) as n, count(*) filter (where kind in ('meaningful','major','critical')) as m | |
| 423 | + from changes where detected_at >= :d0 and detected_at < :d1 group by 1""", d0=d0, d1=d1): | |
| 424 | + b = bucket(r["company_id"]) | |
| 425 | + b["changes"], b["meaningful_changes"] = int(r["n"]), int(r["m"]) | |
| 426 | + for r in await fetch_all(conn, """select company_id, event_type, count(*) as n from events where detected_at >= :d0 and detected_at < :d1 | |
| 427 | + and status in ('active', 'review') group by 1, 2""", d0=d0, d1=d1): | |
| 428 | + b = bucket(r["company_id"]) | |
| 429 | + b["events"] += int(r["n"]) | |
| 430 | + b["events_by_type"][r["event_type"]] = int(r["n"]) | |
| 431 | + for r in await fetch_all(conn, """select company_id, count(*) filter (where first_seen_at >= :d0 and first_seen_at < :d1) as new_n, | |
| 432 | + count(*) filter (where removed_at >= :d0 and removed_at < :d1) as rem_n, | |
| 433 | + count(*) filter (where first_seen_at < :d1 and (removed_at is null or removed_at >= :d1)) as open_n, | |
| 434 | + count(*) filter (where first_seen_at < :d1 and (removed_at is null or removed_at >= :d1) and is_ai) as ai_n | |
| 435 | + from jobs group by 1""", d0=d0, d1=d1): | |
| 436 | + if not (r["new_n"] or r["rem_n"] or r["open_n"]): | |
| 437 | + continue | |
| 438 | + b = bucket(r["company_id"]) | |
| 439 | + b["jobs_new"], b["jobs_removed"], b["jobs_open"], b["jobs_ai_open"] = int(r["new_n"]), int(r["rem_n"]), int(r["open_n"]), int(r["ai_n"]) | |
| 440 | + for r in await fetch_all(conn, "select company_id, count(*) as n from news_items where first_seen_at >= :d0 and first_seen_at < :d1 group by 1", d0=d0, d1=d1): | |
| 441 | + bucket(r["company_id"])["news_items"] = int(r["n"]) | |
| 442 | + if per: | |
| 443 | + ids = list(per) | |
| 444 | + for r in await fetch_all(conn, """select company_id, count(*) as n from sensors where company_id = any(cast(:ids as text[])) and status in ('active','failing','stale') | |
| 445 | + and created_at < :d1 group by 1""", ids=ids, d1=d1): | |
| 446 | + per[r["company_id"]]["sensors_active"] = int(r["n"]) | |
| 447 | + for cid, b in per.items(): | |
| 448 | + await execute(conn, """ | |
| 449 | + insert into company_daily (company_id, day, observations, changes, meaningful_changes, events, events_by_type, jobs_open, jobs_new, jobs_removed, | |
| 450 | + jobs_ai_open, news_items, sensors_active) | |
| 451 | + values (:c, :d, :o, :ch, :m, :e, cast(:ebt as jsonb), :jo, :jn, :jr, :jai, :news, :sa) | |
| 452 | + on conflict (company_id, day) do update set observations = excluded.observations, changes = excluded.changes, | |
| 453 | + meaningful_changes = excluded.meaningful_changes, events = excluded.events, events_by_type = excluded.events_by_type, | |
| 454 | + jobs_open = excluded.jobs_open, jobs_new = excluded.jobs_new, jobs_removed = excluded.jobs_removed, jobs_ai_open = excluded.jobs_ai_open, | |
| 455 | + news_items = excluded.news_items, sensors_active = excluded.sensors_active""", | |
| 456 | + c=cid, d=day, o=b["observations"], ch=b["changes"], m=b["meaningful_changes"], e=b["events"], ebt=jsonb(b["events_by_type"]), jo=b["jobs_open"], | |
| 457 | + jn=b["jobs_new"], jr=b["jobs_removed"], jai=b["jobs_ai_open"], news=b["news_items"], sa=b["sensors_active"]) | |
| 458 | + | |
| 459 | + # global ----------------------------------------------------------------------------------------------------- | |
| 460 | + g = await fetch_one(conn, """select count(distinct company_id) as companies_active, count(distinct sensor_id) as sensors_active, count(*) as observations | |
| 461 | + from observations where fetched_at >= :d0 and fetched_at < :d1""", d0=d0, d1=d1) or {} | |
| 462 | + changes_total = sum(b["changes"] for b in per.values()) | |
| 463 | + meaningful = sum(b["meaningful_changes"] for b in per.values()) | |
| 464 | + events_total = sum(b["events"] for b in per.values()) | |
| 465 | + ebt: dict[str, int] = {} | |
| 466 | + for b in per.values(): | |
| 467 | + for k, v in b["events_by_type"].items(): | |
| 468 | + ebt[k] = ebt.get(k, 0) + v | |
| 469 | + jobs_open = await fetch_val(conn, "select count(*) from jobs where first_seen_at < :d1 and (removed_at is null or removed_at >= :d1)", d1=d1) | |
| 470 | + sensors_active = int(g.get("sensors_active") or 0) | |
| 471 | + companies_active = int(g.get("companies_active") or 0) | |
| 472 | + if sensors_active == 0 and per: # no observation rows (e.g. backfilled changes) → fall back to active sensors | |
| 473 | + sensors_active = int(await fetch_val(conn, "select count(*) from sensors where status in ('active','failing','stale') and created_at < :d1", d1=d1) or 0) | |
| 474 | + companies_active = len(per) | |
| 475 | + ratio = meaningful / sensors_active if sensors_active else None | |
| 476 | + baseline_rows = await fetch_all(conn, """select meaningful_changes, sensors_active from global_daily where day < :d and day >= :d_from and sensors_active > 0""", | |
| 477 | + d=day, d_from=day - timedelta(days=int(P["index_trailing_days"]))) | |
| 478 | + baseline_ratios = [r["meaningful_changes"] / r["sensors_active"] for r in baseline_rows if r["sensors_active"]] | |
| 479 | + baseline = statistics.fmean(baseline_ratios) if baseline_ratios else None | |
| 480 | + activity_index = round(ratio / baseline * 100.0, 2) if (ratio is not None and baseline) else None | |
| 481 | + by_country, by_industry = await _breakdowns(conn, per) | |
| 482 | + await execute(conn, """ | |
| 483 | + insert into global_daily (day, companies_active, sensors_active, observations, changes, meaningful_changes, events, events_by_type, jobs_open, jobs_new, | |
| 484 | + jobs_removed, activity_index, by_country, by_industry, computed_at) | |
| 485 | + values (:d, :ca, :sa, :o, :ch, :m, :e, cast(:ebt as jsonb), :jo, :jn, :jr, :ai, cast(:bc as jsonb), cast(:bi as jsonb), now()) | |
| 486 | + on conflict (day) do update set companies_active = excluded.companies_active, sensors_active = excluded.sensors_active, observations = excluded.observations, | |
| 487 | + changes = excluded.changes, meaningful_changes = excluded.meaningful_changes, events = excluded.events, events_by_type = excluded.events_by_type, | |
| 488 | + jobs_open = excluded.jobs_open, jobs_new = excluded.jobs_new, jobs_removed = excluded.jobs_removed, activity_index = excluded.activity_index, | |
| 489 | + by_country = excluded.by_country, by_industry = excluded.by_industry, computed_at = now()""", | |
| 490 | + d=day, ca=companies_active, sa=sensors_active, o=int(g.get("observations") or 0), ch=changes_total, m=meaningful, e=events_total, ebt=jsonb(ebt), | |
| 491 | + jo=int(jobs_open or 0), jn=sum(b["jobs_new"] for b in per.values()), jr=sum(b["jobs_removed"] for b in per.values()), ai=activity_index, | |
| 492 | + bc=jsonb(by_country), bi=jsonb(by_industry)) | |
| 493 | + baselines_n = await _compute_baselines(conn, day) | |
| 494 | + return {"day": day.isoformat(), "companies": len(per), "meaningful_changes": meaningful, "events": events_total, "sensors_active": sensors_active, | |
| 495 | + "activity_index": activity_index, "baseline_days": len(baseline_ratios), "baselines": baselines_n} | |
| 496 | + | |
| 497 | + | |
| 498 | +async def _breakdowns(conn, per: dict[str, dict[str, Any]]) -> tuple[dict[str, Any], dict[str, Any]]: # type: ignore[no-untyped-def] | |
| 499 | + if not per: | |
| 500 | + return {}, {} | |
| 501 | + rows = await fetch_all(conn, "select id, country, industry_primary, industries from companies where id = any(cast(:ids as text[]))", ids=list(per)) | |
| 502 | + by_country: dict[str, dict[str, int]] = {} | |
| 503 | + by_industry: dict[str, dict[str, int]] = {} | |
| 504 | + for r in rows: | |
| 505 | + b = per[r["id"]] | |
| 506 | + if r["country"]: | |
| 507 | + c = by_country.setdefault(str(r["country"]), {"companies": 0, "meaningful_changes": 0, "events": 0}) | |
| 508 | + c["companies"] += 1 | |
| 509 | + c["meaningful_changes"] += b["meaningful_changes"] | |
| 510 | + c["events"] += b["events"] | |
| 511 | + ind = r["industry_primary"] or (r["industries"][0] if r["industries"] else None) | |
| 512 | + if ind: | |
| 513 | + i = by_industry.setdefault(str(ind), {"companies": 0, "meaningful_changes": 0, "events": 0}) | |
| 514 | + i["companies"] += 1 | |
| 515 | + i["meaningful_changes"] += b["meaningful_changes"] | |
| 516 | + i["events"] += b["events"] | |
| 517 | + return by_country, by_industry | |
| 518 | + | |
| 519 | + | |
| 520 | +async def _compute_baselines(conn, day: date) -> int: # type: ignore[no-untyped-def] | |
| 521 | + """Per-company mean/stddev of weekly meaningful changes, new jobs and news items over `settings.baseline_window_days` ending at `day`.""" | |
| 522 | + window = settings.baseline_window_days | |
| 523 | + start = day - timedelta(days=window) | |
| 524 | + rows = await fetch_all(conn, """ | |
| 525 | + select company_id, floor((day - cast(:start as date)) / 7.0)::int as week, | |
| 526 | + sum(meaningful_changes) as m, sum(jobs_new) as jn, sum(news_items) as news | |
| 527 | + from company_daily where day > :start and day <= :day group by 1, 2""", start=start, day=day) | |
| 528 | + per: dict[str, dict[int, dict[str, int]]] = {} | |
| 529 | + for r in rows: | |
| 530 | + per.setdefault(r["company_id"], {})[int(r["week"])] = {"m": int(r["m"]), "jn": int(r["jn"]), "news": int(r["news"])} | |
| 531 | + first_days = {r["company_id"]: r["first"] for r in await fetch_all(conn, "select company_id, min(day) as first from company_daily where company_id = any(cast(:ids as text[])) group by 1", | |
| 532 | + ids=list(per))} if per else {} | |
| 533 | + n = 0 | |
| 534 | + weeks_total = max(1, window // 7) | |
| 535 | + for cid, weeks in per.items(): | |
| 536 | + first = first_days.get(cid) | |
| 537 | + if first is None: | |
| 538 | + continue | |
| 539 | + observed_weeks = min(weeks_total, max(1, ((day - max(first, start)).days // 7) + 1)) | |
| 540 | + if observed_weeks < 2: | |
| 541 | + continue | |
| 542 | + for metric, key in zip(BASELINE_METRICS, ("m", "jn", "news"), strict=True): | |
| 543 | + series = [weeks.get(w, {}).get(key, 0) for w in range(weeks_total - observed_weeks, weeks_total)] | |
| 544 | + mean = statistics.fmean(series) | |
| 545 | + sd = statistics.pstdev(series) if len(series) > 1 else 0.0 | |
| 546 | + await execute(conn, """ | |
| 547 | + insert into baselines (company_id, metric, mean, stddev, samples, window_days, computed_at) values (:c, :m, :mean, :sd, :n, :w, now()) | |
| 548 | + on conflict (company_id, metric) do update set mean = excluded.mean, stddev = excluded.stddev, samples = excluded.samples, | |
| 549 | + window_days = excluded.window_days, computed_at = now()""", c=cid, m=metric, mean=mean, sd=sd, n=len(series), w=window) | |
| 550 | + n += 1 | |
| 551 | + return n | |
| 552 | + | |
| 553 | + | |
| 554 | +async def compute_daily_catch_up(*, include_today: bool = False, max_days: int = 400) -> list[dict[str, Any]]: | |
| 555 | + """Compute every missing day between the last computed day (or the dataset start) and yesterday.""" | |
| 556 | + today = datetime.now(UTC).date() | |
| 557 | + end = today if include_today else today - timedelta(days=1) | |
| 558 | + async with transaction() as conn: | |
| 559 | + last = await fetch_val(conn, "select max(day) from global_daily") | |
| 560 | + first = await fetch_val(conn, """select least(coalesce((select min(date(fetched_at)) from observations), cast(:t as date)), | |
| 561 | + coalesce((select min(date(detected_at)) from changes), cast(:t as date)), | |
| 562 | + coalesce((select min(date(detected_at)) from events), cast(:t as date)))""", t=today) | |
| 563 | + start = (last + timedelta(days=1)) if last else (first or end) | |
| 564 | + results: list[dict[str, Any]] = [] | |
| 565 | + day = max(start, end - timedelta(days=max_days)) | |
| 566 | + while day <= end: | |
| 567 | + results.append(await compute_daily(day)) | |
| 568 | + day += timedelta(days=1) | |
| 569 | + return results | |
| 570 | + | |
| 571 | + | |
| 572 | +# ================================================================================================================ periodic | |
| 573 | + | |
| 574 | + | |
| 575 | +@periodic("metrics-hourly", cron=settings.metrics_cron) | |
| 576 | +async def metrics_hourly_task() -> None: | |
| 577 | + stats = await compute_company_metrics() | |
| 578 | + log.info("metrics-hourly", extra=stats) | |
| 579 | + | |
| 580 | + | |
| 581 | +@periodic("daily-aggregates", cron=settings.daily_cron) | |
| 582 | +async def daily_task() -> None: | |
| 583 | + results = await compute_daily_catch_up() | |
| 584 | + log.info("daily-aggregates", extra={"days": len(results), "last": results[-1] if results else None}) | |
| 585 | + stats = await compute_company_metrics(all_companies=True) | |
| 586 | + log.info("metrics-nightly", extra=stats) | |
| 587 | + | |
| 588 | + | |
| 589 | +__all__ = ["CompanyContext", "MetricValue", "activity_score", "ai_adoption", "anomaly_score", "compute_company_metrics", "compute_daily", | |
| 590 | + "compute_daily_catch_up", "compute_from_context", "corporate_change_index", "hiring_momentum", "historical_coverage", "load_context", | |
| 591 | + "saturate", "saturating_score", "write_metrics"] | |
added
src/companyatlas/services/retention.py
+91 −0
@@ -0,0 +1,91 @@ | ||
| 1 | +"""Retention (spec §2.1, §95–96): history is never deleted. The only prunable rows are bookkeeping: | |
| 2 | +- `observations` older than `retention_observations_days` that were `not_modified` / unchanged, are not referenced by a snapshot and are not | |
| 3 | + the sensor's most recent observation (aggregate counts stay in `sensors.observation_count`; pruned totals are kept in settings_kv); | |
| 4 | +- `crawl_runs` older than `retention_crawl_runs_days`; | |
| 5 | +- finished `llm_jobs` older than `retention_llm_jobs_days`, summarised into `cost_ledger` (dimension `llm`, key `archived:<model>`) first. | |
| 6 | +Snapshots, changes, events, entities, metrics and series are never touched. | |
| 7 | +""" | |
| 8 | +from __future__ import annotations | |
| 9 | + | |
| 10 | +import logging | |
| 11 | +from datetime import UTC, datetime, timedelta | |
| 12 | +from typing import Any | |
| 13 | + | |
| 14 | +from companyatlas.config import settings | |
| 15 | +from companyatlas.db import execute, fetch_all, fetch_val, jsonb, transaction | |
| 16 | +from companyatlas.services.periodic import periodic | |
| 17 | + | |
| 18 | +log = logging.getLogger(__name__) | |
| 19 | + | |
| 20 | +PRUNED_KEY = "retention:pruned" | |
| 21 | +BATCH = 5000 | |
| 22 | + | |
| 23 | + | |
| 24 | +async def prune_observations(*, dry_run: bool = False, batch: int = BATCH) -> int: | |
| 25 | + cutoff = datetime.now(UTC) - timedelta(days=settings.retention_observations_days) | |
| 26 | + total = 0 | |
| 27 | + async with transaction() as conn: | |
| 28 | + while True: | |
| 29 | + ids = [r["id"] for r in await fetch_all(conn, """ | |
| 30 | + select o.id from observations o | |
| 31 | + where o.fetched_at < :cutoff and (o.not_modified or (o.changed = false and o.failure_class is null)) | |
| 32 | + and not exists (select 1 from snapshots s where s.observation_id = o.id) | |
| 33 | + and o.id <> (select o2.id from observations o2 where o2.sensor_id = o.sensor_id order by o2.fetched_at desc limit 1) | |
| 34 | + limit :batch""", cutoff=cutoff, batch=batch)] | |
| 35 | + if not ids or dry_run: | |
| 36 | + total += len(ids) | |
| 37 | + break | |
| 38 | + await execute(conn, "delete from observations where id = any(cast(:ids as text[]))", ids=ids) | |
| 39 | + total += len(ids) | |
| 40 | + if len(ids) < batch: | |
| 41 | + break | |
| 42 | + if total and not dry_run: | |
| 43 | + prev = await fetch_val(conn, "select value from settings_kv where key = :k", k=PRUNED_KEY) | |
| 44 | + state = dict(prev) if isinstance(prev, dict) else {} | |
| 45 | + state["observations"] = int(state.get("observations", 0)) + total | |
| 46 | + state["last_run_at"] = datetime.now(UTC).isoformat() | |
| 47 | + await execute(conn, "insert into settings_kv (key, value, updated_at) values (:k, cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()", | |
| 48 | + k=PRUNED_KEY, v=jsonb(state)) | |
| 49 | + return total | |
| 50 | + | |
| 51 | + | |
| 52 | +async def prune_crawl_runs(*, dry_run: bool = False) -> int: | |
| 53 | + cutoff = datetime.now(UTC) - timedelta(days=settings.retention_crawl_runs_days) | |
| 54 | + async with transaction() as conn: | |
| 55 | + n = int(await fetch_val(conn, "select count(*) from crawl_runs where started_at < :c", c=cutoff) or 0) | |
| 56 | + if n and not dry_run: | |
| 57 | + await execute(conn, "delete from crawl_runs where started_at < :c", c=cutoff) | |
| 58 | + return n | |
| 59 | + | |
| 60 | + | |
| 61 | +async def archive_llm_jobs(*, dry_run: bool = False) -> int: | |
| 62 | + cutoff = datetime.now(UTC) - timedelta(days=settings.retention_llm_jobs_days) | |
| 63 | + async with transaction() as conn: | |
| 64 | + rows = await fetch_all(conn, """select coalesce(model, 'unknown') as model, count(*) as n, coalesce(sum(request_tokens), 0) + coalesce(sum(response_tokens), 0) as tokens | |
| 65 | + from llm_jobs where status in ('done', 'failed') and finished_at < :c group by 1""", c=cutoff) | |
| 66 | + n = sum(int(r["n"]) for r in rows) | |
| 67 | + if n and not dry_run: | |
| 68 | + for r in rows: | |
| 69 | + await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0) | |
| 70 | + on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""", | |
| 71 | + key=f"archived:{r['model']}", units=float(r["tokens"])) | |
| 72 | + await execute(conn, """insert into cost_ledger (day, dimension, key, units, cost_estimate) values (current_date, 'llm', :key, :units, 0) | |
| 73 | + on conflict (day, dimension, key) do update set units = cost_ledger.units + excluded.units""", | |
| 74 | + key=f"archived-jobs:{r['model']}", units=float(r["n"])) | |
| 75 | + await execute(conn, "delete from llm_jobs where status in ('done', 'failed') and finished_at < :c", c=cutoff) | |
| 76 | + return n | |
| 77 | + | |
| 78 | + | |
| 79 | +async def run_retention(*, dry_run: bool = False) -> dict[str, Any]: | |
| 80 | + stats = {"observations": await prune_observations(dry_run=dry_run), "crawl_runs": await prune_crawl_runs(dry_run=dry_run), | |
| 81 | + "llm_jobs": await archive_llm_jobs(dry_run=dry_run), "dry_run": dry_run} | |
| 82 | + log.info("retention", extra=stats) | |
| 83 | + return stats | |
| 84 | + | |
| 85 | + | |
| 86 | +@periodic("retention", cron="50 3 * * *") | |
| 87 | +async def retention_task() -> None: | |
| 88 | + await run_retention() | |
| 89 | + | |
| 90 | + | |
| 91 | +__all__ = ["archive_llm_jobs", "prune_crawl_runs", "prune_observations", "run_retention"] | |
added
src/companyatlas/services/signals.py
+273 −0
@@ -0,0 +1,273 @@ | ||
| 1 | +"""Cross-signal detection (spec §92–93): hiring surge / freeze, launch build-up, international expansion, pricing migration, developer | |
| 2 | +push, enterprise repositioning, AI acceleration, abnormal activity. Signals are *labelled as signals*, carry a strength (0–1), a | |
| 3 | +confidence, an explanation and the evidence (event ids, metric values) that produced them, and expire. They never assert facts. | |
| 4 | + | |
| 5 | +Company-scope signals are recomputed hourly for companies with recent activity; industry/country aggregates are derived from them. | |
| 6 | +""" | |
| 7 | +from __future__ import annotations | |
| 8 | + | |
| 9 | +import logging | |
| 10 | +from dataclasses import dataclass, field | |
| 11 | +from datetime import UTC, datetime, timedelta | |
| 12 | +from typing import Any | |
| 13 | + | |
| 14 | +from companyatlas.config import settings | |
| 15 | +from companyatlas.db import execute, fetch_all, fetch_one, jsonb, transaction | |
| 16 | +from companyatlas.ids import new_id | |
| 17 | +from companyatlas.services.metrics import saturate | |
| 18 | +from companyatlas.services.periodic import periodic | |
| 19 | +from companyatlas.taxonomy import SIGNAL_KINDS, SIGNALS_FORMULA_VERSION, EventType, Metric | |
| 20 | + | |
| 21 | +log = logging.getLogger(__name__) | |
| 22 | + | |
| 23 | +SCOPE_MIN_COMPANIES = 3 | |
| 24 | +LAUNCH_WINDOW_DAYS = 14 | |
| 25 | +SURGE_MOMENTUM_PCT = 30.0 | |
| 26 | +FREEZE_MOMENTUM_PCT = -30.0 | |
| 27 | +MIN_OPEN_FOR_MOMENTUM = 10 | |
| 28 | +ENTERPRISE_RE = ("enterprise", "contact sales", "sso", "saml", "audit log", "dedicated", "custom pricing", "procurement", "soc 2", "soc2") | |
| 29 | + | |
| 30 | + | |
| 31 | +@dataclass(slots=True) | |
| 32 | +class SignalDraft: | |
| 33 | + kind: str | |
| 34 | + strength: float | |
| 35 | + confidence: float | |
| 36 | + title: str | |
| 37 | + explanation: str | |
| 38 | + evidence: dict[str, Any] = field(default_factory=dict) | |
| 39 | + window_days: int = 30 | |
| 40 | + | |
| 41 | + | |
| 42 | +@dataclass(slots=True) | |
| 43 | +class SignalInputs: | |
| 44 | + """Everything the detectors look at for one company (pure, testable).""" | |
| 45 | + metrics: dict[str, float] | |
| 46 | + events: list[dict[str, Any]] # {id, event_type, event_subtype, importance, tags, detected_at, title} | |
| 47 | + jobs: dict[str, Any] # open_now, new_30d, ai_new_30d, new_countries_90d | |
| 48 | + plans_contact_sales_new: int = 0 | |
| 49 | + locations_new_countries: list[str] = field(default_factory=list) | |
| 50 | + now: datetime = field(default_factory=lambda: datetime.now(UTC)) | |
| 51 | + | |
| 52 | + | |
| 53 | +def _age(now: datetime, at: datetime) -> float: | |
| 54 | + if at.tzinfo is None: | |
| 55 | + at = at.replace(tzinfo=UTC) | |
| 56 | + return (now - at).total_seconds() / 86400.0 | |
| 57 | + | |
| 58 | + | |
| 59 | +def _ids(events: list[dict[str, Any]]) -> list[str]: | |
| 60 | + return [e["id"] for e in events[:25]] | |
| 61 | + | |
| 62 | + | |
| 63 | +def detect(inp: SignalInputs) -> list[SignalDraft]: | |
| 64 | + m, ev, j, now = inp.metrics, inp.events, inp.jobs, inp.now | |
| 65 | + out: list[SignalDraft] = [] | |
| 66 | + window = settings.signals_window_days | |
| 67 | + recent = [e for e in ev if _age(now, e["detected_at"]) <= window] | |
| 68 | + open_now = int(j.get("open_now") or 0) | |
| 69 | + mom30 = m.get(Metric.HIRING_MOMENTUM_30D) | |
| 70 | + | |
| 71 | + # hiring surge -------------------------------------------------------------------------------------------------- | |
| 72 | + surge_events = [e for e in recent if e["event_subtype"] in ("HIRING_SURGE", "JOB_COUNT_INCREASE")] | |
| 73 | + if (mom30 is not None and mom30 >= SURGE_MOMENTUM_PCT and open_now >= MIN_OPEN_FOR_MOMENTUM) or any(e["event_subtype"] == "HIRING_SURGE" for e in recent): | |
| 74 | + strength = saturate(max(mom30 or 0, 30.0) / 50.0, 1.5) | |
| 75 | + out.append(SignalDraft("hiring_surge", round(strength, 3), 0.7, "Hiring surge signal", | |
| 76 | + f"Open listings observed up {mom30:.0f}% over 30 days ({open_now} open)." if mom30 is not None else | |
| 77 | + f"{len(surge_events)} hiring increase events detected in the last {window} days.", | |
| 78 | + {"event_ids": _ids(surge_events), "hiring_momentum_30d": mom30, "open_jobs": open_now})) | |
| 79 | + # hiring freeze — careful wording: listings no longer visible -------------------------------------------------- | |
| 80 | + freeze_events = [e for e in recent if e["event_subtype"] in ("HIRING_FREEZE_SIGNAL", "JOB_COUNT_DECREASE")] | |
| 81 | + if (mom30 is not None and mom30 <= FREEZE_MOMENTUM_PCT and int(j.get("new_30d") or 0) <= 1) or any(e["event_subtype"] == "HIRING_FREEZE_SIGNAL" for e in recent): | |
| 82 | + strength = saturate(abs(min(mom30 or -30.0, -30.0)) / 50.0, 1.5) | |
| 83 | + out.append(SignalDraft("hiring_freeze", round(strength, 3), 0.6, "Hiring slowdown signal: listings no longer visible", | |
| 84 | + (f"Monitored open listings are down {abs(mom30):.0f}% over 30 days with {int(j.get('new_30d') or 0)} new listings; " | |
| 85 | + "listings no longer visible on public pages are not evidence of workforce decisions.") if mom30 is not None else | |
| 86 | + f"{len(freeze_events)} events of listings no longer visible in the last {window} days.", | |
| 87 | + {"event_ids": _ids(freeze_events), "hiring_momentum_30d": mom30, "jobs_new_30d": j.get("new_30d")})) | |
| 88 | + # launch build-up ----------------------------------------------------------------------------------------------- | |
| 89 | + win14 = [e for e in ev if _age(now, e["detected_at"]) <= LAUNCH_WINDOW_DAYS] | |
| 90 | + cats = { | |
| 91 | + "product": [e for e in win14 if e["event_type"] == EventType.PRODUCT], | |
| 92 | + "docs": [e for e in win14 if e["event_subtype"] in ("DOC_CHANGE", "DOCUMENTATION_CHANGE", "API_CHANGE", "API_LAUNCH", "SDK_RELEASE")], | |
| 93 | + "changelog": [e for e in win14 if e["event_subtype"] == "CHANGELOG_ENTRY"], | |
| 94 | + "careers": [e for e in win14 if e["event_type"] == EventType.HIRING and e["event_subtype"] != "JOB_COUNT_DECREASE"], | |
| 95 | + "messaging": [e for e in win14 if e["event_subtype"] in ("MESSAGING_CHANGE", "HOMEPAGE_REDESIGN", "WEBSITE_CHANGE")], | |
| 96 | + } | |
| 97 | + present = {k: v for k, v in cats.items() if v} | |
| 98 | + if len(present) >= 3 and cats["product"] or len(present) >= 4: | |
| 99 | + strength = min(1.0, 0.3 + 0.15 * len(present) + 0.05 * min(6, sum(len(v) for v in present.values()))) | |
| 100 | + out.append(SignalDraft("launch_buildup", round(strength, 3), round(0.45 + 0.08 * len(present), 2), "Possible launch preparation signal", | |
| 101 | + "Within 14 days: " + ", ".join(f"{len(v)} {k}" for k, v in present.items()) + " events detected on public pages. " | |
| 102 | + "This pattern often precedes announcements; it is a probabilistic signal, not a confirmation.", | |
| 103 | + {"event_ids": _ids([e for v in present.values() for e in v]), "categories": {k: len(v) for k, v in present.items()}}, LAUNCH_WINDOW_DAYS)) | |
| 104 | + # expansion ----------------------------------------------------------------------------------------------------- | |
| 105 | + geo_events = [e for e in recent if e["event_subtype"] in ("COUNTRY_EXPANSION", "NEW_LOCATION", "NEW_OFFICE")] | |
| 106 | + new_countries = sorted(set(inp.locations_new_countries) | set(j.get("new_countries_90d") or [])) | |
| 107 | + if any(e["event_subtype"] == "COUNTRY_EXPANSION" for e in geo_events) or (new_countries and (geo_events or int(j.get("new_30d") or 0) > 0)): | |
| 108 | + strength = min(1.0, 0.4 + 0.2 * len(new_countries) + 0.05 * len(geo_events)) | |
| 109 | + out.append(SignalDraft("expansion", round(strength, 3), 0.65, "International expansion signal", | |
| 110 | + (f"New country presence listed: {', '.join(new_countries[:5])}. " if new_countries else "") + | |
| 111 | + f"{len(geo_events)} location events detected in the last {window} days" + (", with job listings in new countries." if j.get("new_countries_90d") else "."), | |
| 112 | + {"event_ids": _ids(geo_events), "new_countries": new_countries, "geo_expansion": m.get(Metric.GEO_EXPANSION)})) | |
| 113 | + # pricing migration --------------------------------------------------------------------------------------------- | |
| 114 | + pricing = [e for e in recent if e["event_type"] == EventType.PRICING] | |
| 115 | + subtypes = {e["event_subtype"] for e in pricing} | |
| 116 | + if len(pricing) >= 2 or ({"NEW_PRICING_TIER", "PRICING_TIER_REMOVED"} <= subtypes): | |
| 117 | + strength = min(1.0, 0.3 + 0.15 * len(pricing) + (0.2 if {"NEW_PRICING_TIER", "PRICING_TIER_REMOVED"} <= subtypes else 0)) | |
| 118 | + out.append(SignalDraft("pricing_migration", round(strength, 3), 0.7, "Pricing migration signal", | |
| 119 | + f"{len(pricing)} pricing events in {window} days: " + ", ".join(sorted(subtypes)).lower().replace("_", " ") + ".", | |
| 120 | + {"event_ids": _ids(pricing), "subtypes": sorted(subtypes), "pricing_activity": m.get(Metric.PRICING_ACTIVITY)})) | |
| 121 | + # developer push ------------------------------------------------------------------------------------------------ | |
| 122 | + dev = [e for e in recent if e["event_type"] == EventType.DEVELOPER] | |
| 123 | + dev_m = m.get(Metric.DEVELOPER_MOMENTUM) | |
| 124 | + if len(dev) >= 3 or (dev_m is not None and dev_m >= 60): | |
| 125 | + strength = saturate(len(dev) + (dev_m or 0) / 40.0, 4.0) | |
| 126 | + out.append(SignalDraft("developer_push", round(strength, 3), 0.7, "Developer push signal", | |
| 127 | + f"{len(dev)} developer-surface events (docs, API, changelog, SDK) in {window} days" + (f"; developer momentum {dev_m:.0f}/100." if dev_m is not None else "."), | |
| 128 | + {"event_ids": _ids(dev), "developer_momentum": dev_m})) | |
| 129 | + # enterprise repositioning -------------------------------------------------------------------------------------- | |
| 130 | + ent_events = [e for e in recent if e["event_subtype"] == "ENTERPRISE_REPOSITIONING" or "enterprise" in (e.get("tags") or []) | |
| 131 | + or any(k in (e.get("title") or "").lower() for k in ENTERPRISE_RE)] | |
| 132 | + if inp.plans_contact_sales_new or len(ent_events) >= 2: | |
| 133 | + strength = min(1.0, 0.35 + 0.25 * min(2, inp.plans_contact_sales_new) + 0.1 * len(ent_events)) | |
| 134 | + out.append(SignalDraft("enterprise_repositioning", round(strength, 3), 0.55, "Enterprise repositioning signal", | |
| 135 | + (f"{inp.plans_contact_sales_new} new contact-sales / enterprise pricing tier(s) listed. " if inp.plans_contact_sales_new else "") + | |
| 136 | + (f"{len(ent_events)} events mention enterprise features or messaging." if ent_events else ""), | |
| 137 | + {"event_ids": _ids(ent_events), "contact_sales_tiers_new": inp.plans_contact_sales_new})) | |
| 138 | + # AI acceleration ----------------------------------------------------------------------------------------------- | |
| 139 | + ai_events = [e for e in recent if e["event_subtype"] in ("AI_HIRING", "AI_LAUNCH") or "ai" in (e.get("tags") or [])] | |
| 140 | + ai_new = int(j.get("ai_new_30d") or 0) | |
| 141 | + new30 = int(j.get("new_30d") or 0) | |
| 142 | + ai_share = (ai_new / new30) if new30 else 0.0 | |
| 143 | + if len(ai_events) >= 2 or (new30 >= 4 and ai_share >= 0.25): | |
| 144 | + strength = min(1.0, 0.3 + 0.1 * len(ai_events) + ai_share) | |
| 145 | + out.append(SignalDraft("ai_acceleration", round(strength, 3), 0.6, "AI acceleration signal", | |
| 146 | + f"{len(ai_events)} AI-related events in {window} days" + (f"; {ai_new} of {new30} new listings are AI-related ({ai_share:.0%})." if new30 else ".") + | |
| 147 | + " Based on public job titles, product names and announcements only.", | |
| 148 | + {"event_ids": _ids(ai_events), "ai_new_30d": ai_new, "jobs_new_30d": new30, "ai_adoption": m.get(Metric.AI_ADOPTION)})) | |
| 149 | + # abnormal activity --------------------------------------------------------------------------------------------- | |
| 150 | + z = m.get(Metric.ANOMALY_SCORE) | |
| 151 | + if z is not None and z >= settings.anomaly_z: | |
| 152 | + out.append(SignalDraft("abnormal_activity", round(min(1.0, z / (settings.anomaly_z * 2)), 3), 0.7, "Unusual activity signal", | |
| 153 | + f"This week's meaningful changes are {z:.1f} standard deviations above this company's baseline.", | |
| 154 | + {"anomaly_z": z, "event_ids": _ids(recent[:10])}, 7)) | |
| 155 | + return out | |
| 156 | + | |
| 157 | + | |
| 158 | +# ================================================================================================================ persistence | |
| 159 | + | |
| 160 | + | |
| 161 | +async def _inputs(conn, company_id: str, now: datetime) -> SignalInputs: # type: ignore[no-untyped-def] | |
| 162 | + metrics = {r["metric"]: float(r["value"]) for r in await fetch_all(conn, "select metric, value from metrics_current where company_id = :c", c=company_id)} | |
| 163 | + events = await fetch_all(conn, """select id, event_type, event_subtype, importance, tags, detected_at, title from events where company_id = :c | |
| 164 | + and detected_at >= :since and status in ('active', 'review') order by detected_at desc limit 500""", | |
| 165 | + c=company_id, since=now - timedelta(days=90)) | |
| 166 | + jobs = await fetch_one(conn, """select count(*) filter (where status = 'open') as open_now, count(*) filter (where first_seen_at > :d30) as new_30d, | |
| 167 | + count(*) filter (where first_seen_at > :d30 and is_ai) as ai_new_30d from jobs where company_id = :c""", | |
| 168 | + c=company_id, d30=now - timedelta(days=30)) or {} | |
| 169 | + new_job_countries = [r["country"] for r in await fetch_all(conn, "select country from jobs where company_id = :c and country is not null group by country having min(first_seen_at) > :s", | |
| 170 | + c=company_id, s=now - timedelta(days=90))] | |
| 171 | + plans_new = await fetch_one(conn, "select count(*) as n from pricing_plans where company_id = :c and contact_sales and first_seen_at > :s", c=company_id, s=now - timedelta(days=30)) or {} | |
| 172 | + loc_countries = [r["country"] for r in await fetch_all(conn, "select country from locations where company_id = :c and country is not null group by country having min(first_seen_at) > :s", | |
| 173 | + c=company_id, s=now - timedelta(days=90))] | |
| 174 | + return SignalInputs(metrics=metrics, events=events, jobs={**jobs, "new_countries_90d": new_job_countries}, plans_contact_sales_new=int(plans_new.get("n") or 0), | |
| 175 | + locations_new_countries=loc_countries, now=now) | |
| 176 | + | |
| 177 | + | |
| 178 | +async def upsert_company_signals(conn, company_id: str, drafts: list[SignalDraft], now: datetime) -> dict[str, int]: # type: ignore[no-untyped-def] | |
| 179 | + stats = {"inserted": 0, "updated": 0, "expired": 0} | |
| 180 | + active = {r["kind"]: r for r in await fetch_all(conn, "select id, kind from signals where company_id = :c and scope = 'company' and status = 'active'", c=company_id)} | |
| 181 | + seen: set[str] = set() | |
| 182 | + for d in drafts: | |
| 183 | + seen.add(d.kind) | |
| 184 | + evidence = {**d.evidence, "formula_version": SIGNALS_FORMULA_VERSION} | |
| 185 | + expires = now + timedelta(days=settings.signals_ttl_days) | |
| 186 | + if d.kind in active: | |
| 187 | + await execute(conn, """update signals set strength = :s, confidence = :c, title = :t, explanation = :x, evidence = cast(:e as jsonb), window_days = :w, | |
| 188 | + expires_at = :exp where id = :id""", s=d.strength, c=d.confidence, t=d.title, x=d.explanation, e=jsonb(evidence), w=d.window_days, | |
| 189 | + exp=expires, id=active[d.kind]["id"]) | |
| 190 | + stats["updated"] += 1 | |
| 191 | + else: | |
| 192 | + await execute(conn, """insert into signals (id, company_id, scope, scope_key, kind, strength, confidence, title, explanation, evidence, window_days, detected_at, expires_at, status) | |
| 193 | + values (:id, :c, 'company', null, :k, :s, :conf, :t, :x, cast(:e as jsonb), :w, :now, :exp, 'active')""", | |
| 194 | + id=new_id("signal"), c=company_id, k=d.kind, s=d.strength, conf=d.confidence, t=d.title, x=d.explanation, e=jsonb(evidence), w=d.window_days, | |
| 195 | + now=now, exp=expires) | |
| 196 | + stats["inserted"] += 1 | |
| 197 | + for kind, row in active.items(): | |
| 198 | + if kind not in seen: | |
| 199 | + await execute(conn, "update signals set status = 'expired', expires_at = least(coalesce(expires_at, :now), :now) where id = :id", now=now, id=row["id"]) | |
| 200 | + stats["expired"] += 1 | |
| 201 | + return stats | |
| 202 | + | |
| 203 | + | |
| 204 | +async def compute_scope_signals(conn, now: datetime) -> int: # type: ignore[no-untyped-def] | |
| 205 | + """Industry / country aggregates: ≥ SCOPE_MIN_COMPANIES companies sharing an active signal kind.""" | |
| 206 | + rows = await fetch_all(conn, """ | |
| 207 | + select s.kind, co.country, coalesce(co.industry_primary, co.industries[1]) as industry, co.slug, s.strength | |
| 208 | + from signals s join companies co on co.id = s.company_id where s.scope = 'company' and s.status = 'active' and s.expires_at > :now""", now=now) | |
| 209 | + groups: dict[tuple[str, str, str], list[dict[str, Any]]] = {} | |
| 210 | + for r in rows: | |
| 211 | + if r["country"]: | |
| 212 | + groups.setdefault(("country", str(r["country"]), r["kind"]), []).append(r) | |
| 213 | + if r["industry"]: | |
| 214 | + groups.setdefault(("industry", str(r["industry"]), r["kind"]), []).append(r) | |
| 215 | + n = 0 | |
| 216 | + kept: list[str] = [] | |
| 217 | + await execute(conn, "update signals set status = 'expired' where scope in ('industry', 'country') and status = 'active' and expires_at <= :now", now=now) | |
| 218 | + for (scope, key, kind), items in groups.items(): | |
| 219 | + if len(items) < SCOPE_MIN_COMPANIES: | |
| 220 | + continue | |
| 221 | + kept.append(f"{scope}:{key}:{kind}") | |
| 222 | + label = kind.replace("_", " ") | |
| 223 | + strength = round(min(1.0, saturate(len(items), 6.0) * (sum(float(i["strength"]) for i in items) / len(items) + 0.3)), 3) | |
| 224 | + title = f"{label.capitalize()} signal across {len(items)} companies in {key}" | |
| 225 | + evidence = {"companies": [i["slug"] for i in items[:50]], "count": len(items), "formula_version": SIGNALS_FORMULA_VERSION} | |
| 226 | + existing = await fetch_one(conn, "select id from signals where scope = :s and scope_key = :k and kind = :kind and status = 'active'", s=scope, k=key, kind=kind) | |
| 227 | + expires = now + timedelta(days=settings.signals_ttl_days) | |
| 228 | + if existing: | |
| 229 | + await execute(conn, "update signals set strength = :st, title = :t, evidence = cast(:e as jsonb), expires_at = :exp where id = :id", | |
| 230 | + st=strength, t=title, e=jsonb(evidence), exp=expires, id=existing["id"]) | |
| 231 | + else: | |
| 232 | + await execute(conn, """insert into signals (id, company_id, scope, scope_key, kind, strength, confidence, title, explanation, evidence, window_days, detected_at, expires_at) | |
| 233 | + values (:id, null, :s, :k, :kind, :st, 0.6, :t, :x, cast(:e as jsonb), :w, :now, :exp)""", | |
| 234 | + id=new_id("signal"), s=scope, k=key, kind=kind, st=strength, t=title, e=jsonb(evidence), w=settings.signals_window_days, now=now, exp=expires, | |
| 235 | + x=f"{len(items)} monitored companies in this {scope} currently carry an active {label} signal.") | |
| 236 | + n += 1 | |
| 237 | + # aggregates whose supporting companies dropped below the threshold are no longer signals | |
| 238 | + await execute(conn, """update signals set status = 'expired' where scope in ('industry', 'country') and status = 'active' | |
| 239 | + and (scope || ':' || scope_key || ':' || kind) <> all(cast(:kept as text[]))""", kept=kept) | |
| 240 | + return n | |
| 241 | + | |
| 242 | + | |
| 243 | +async def compute_signals(company_ids: list[str] | None = None, *, now: datetime | None = None) -> dict[str, int]: | |
| 244 | + now = now or datetime.now(UTC) | |
| 245 | + stats = {"companies": 0, "inserted": 0, "updated": 0, "expired": 0, "scope": 0} | |
| 246 | + async with transaction() as conn: | |
| 247 | + if company_ids is None: | |
| 248 | + rows = await fetch_all(conn, """select distinct company_id as id from events where detected_at >= :since and status in ('active', 'review') | |
| 249 | + union select company_id from signals where scope = 'company' and status = 'active' | |
| 250 | + union select company_id from metrics_current where metric = 'anomaly_score' and value >= :z""", | |
| 251 | + since=now - timedelta(days=90), z=settings.anomaly_z) | |
| 252 | + company_ids = [r["id"] for r in rows] | |
| 253 | + for cid in company_ids: | |
| 254 | + try: | |
| 255 | + async with transaction() as conn: | |
| 256 | + drafts = detect(await _inputs(conn, cid, now)) | |
| 257 | + s = await upsert_company_signals(conn, cid, drafts, now) | |
| 258 | + stats["companies"] += 1 | |
| 259 | + for k in ("inserted", "updated", "expired"): | |
| 260 | + stats[k] += s[k] | |
| 261 | + except Exception: | |
| 262 | + log.exception("signals failed", extra={"company_id": cid}) | |
| 263 | + async with transaction() as conn: | |
| 264 | + stats["scope"] = await compute_scope_signals(conn, now) | |
| 265 | + return stats | |
| 266 | + | |
| 267 | + | |
| 268 | +@periodic("signals", every_s=3600, initial_delay_s=120) | |
| 269 | +async def signals_task() -> None: | |
| 270 | + log.info("signals", extra=await compute_signals()) | |
| 271 | + | |
| 272 | + | |
| 273 | +__all__ = ["SIGNAL_KINDS", "SignalDraft", "SignalInputs", "compute_scope_signals", "compute_signals", "detect", "upsert_company_signals"] | |
added
src/companyatlas/services/trends.py
+136 −0
@@ -0,0 +1,136 @@ | ||
| 1 | +"""Trend engine (spec §144–145): terms extracted from event and news titles (1–3-grams, stopwords, company names removed, lowercase), | |
| 2 | +aggregated per day into `trends(term, day, mentions, companies)` and only kept when ≥ `settings.trends_min_companies` distinct companies | |
| 3 | +use them. Momentum (7/30/90 d) is normalised by the number of active companies so that coverage growth alone does not create trends. | |
| 4 | +""" | |
| 5 | +from __future__ import annotations | |
| 6 | + | |
| 7 | +import logging | |
| 8 | +import re | |
| 9 | +from collections import defaultdict | |
| 10 | +from datetime import UTC, date, datetime, timedelta | |
| 11 | +from typing import Any | |
| 12 | + | |
| 13 | +from companyatlas.config import settings | |
| 14 | +from companyatlas.db import execute, fetch_all, fetch_val, jsonb, transaction | |
| 15 | +from companyatlas.services.clustering import normalize_entity_key | |
| 16 | +from companyatlas.services.periodic import periodic | |
| 17 | + | |
| 18 | +log = logging.getLogger(__name__) | |
| 19 | + | |
| 20 | +TRENDS_FORMULA_VERSION = "trends-v1" | |
| 21 | +MAX_TERMS_PER_TITLE = 40 | |
| 22 | +STOPWORDS = frozenset(["a", "an", "the", "and", "or", "of", "to", "in", "on", "at", "for", "by", "with", "from", "as", "is", "are", "was", "were", "be", "been", "being", "this", "that", "these", "those", "it", "its", "into", "over", "under", "about", "after", "before", "than", "then", "there", "their", "they", "them", "we", "our", "you", "your", "he", "she", "his", "her", "not", "no", "yes", "new", "now", "more", "most", "less", "least", "very", "up", "down", "out", "off", "via", "per", "vs", "detected", "listed", "observed", "monitored", "longer", "visible", "no", "page", "pages", "updated", "update", "updates", "changed", "change", "changes", "section", "sections", "block", "blocks", "content", "position", "positions", "job", "jobs", "listing", "listings", "careers", "career", "title", "plan", "plans", "price", "prices", "pricing", "tier", "tiers", "product", "products", "office", "offices", "location", "locations", "country", "presence", "leadership", "executive", "team", "news", "release", "blog", "post", "entry", "changelog", "investor", "earnings", "homepage", "website", "site", "material", "materially", "redesigned", "documentation", "docs", "api", "reference", "terms", "service", "privacy", "policy", "security", "added", "removed", "increase", "decrease", "signal", "one", "two", "three", "announces", "announced", "announce", "introduces", "introducing", "launches", "launched", "launch", "today", "year", "years", "month", "months", "week", "weeks", "day", "days", "q1", "q2", "q3", "q4", "fy", "inc", "corp", "ltd", "llc", "co", "company", "companies", "group", "plc", "ag", "sa", "nv", "se", "gmbh"]) | |
| 23 | +_TOKEN_RE = re.compile(r"[a-z0-9][a-z0-9\-\+\.]{1,}") | |
| 24 | + | |
| 25 | + | |
| 26 | +def extract_terms(title: str, *, company_name: str | None = None) -> set[str]: | |
| 27 | + """Lowercased 1–3-grams over alphabetic tokens, stopwords dropped at n-gram edges, company name removed first.""" | |
| 28 | + text = (title or "").lower() | |
| 29 | + if ":" in text and text.split(":", 1)[0].strip() in {"news release", "blog post", "changelog entry", "investor update", "earnings release"}: | |
| 30 | + text = text.split(":", 1)[1] | |
| 31 | + if company_name: | |
| 32 | + norm = normalize_entity_key(company_name) | |
| 33 | + if norm: | |
| 34 | + text = re.sub(r"(?<![a-z0-9])" + re.escape(norm) + r"(?![a-z0-9])", " ", normalize_entity_key(text) if norm in normalize_entity_key(text) else text) | |
| 35 | + tokens = [t.strip(".-+") for t in _TOKEN_RE.findall(text)] | |
| 36 | + tokens = [t for t in tokens if len(t) >= 3 and not t.isdigit()] | |
| 37 | + terms: set[str] = set() | |
| 38 | + for n in (1, 2, 3): | |
| 39 | + for i in range(len(tokens) - n + 1): | |
| 40 | + gram = tokens[i:i + n] | |
| 41 | + if gram[0] in STOPWORDS or gram[-1] in STOPWORDS: | |
| 42 | + continue | |
| 43 | + if n == 1 and (len(gram[0]) < 4 or gram[0] in STOPWORDS): | |
| 44 | + continue | |
| 45 | + if all(g in STOPWORDS for g in gram): | |
| 46 | + continue | |
| 47 | + terms.add(" ".join(gram)) | |
| 48 | + if len(terms) >= MAX_TERMS_PER_TITLE: | |
| 49 | + return terms | |
| 50 | + return terms | |
| 51 | + | |
| 52 | + | |
| 53 | +async def compute_trends(day: date | None = None) -> dict[str, Any]: | |
| 54 | + """Aggregate terms for one UTC day (default: today) into `trends`. Idempotent.""" | |
| 55 | + day = day or datetime.now(UTC).date() | |
| 56 | + d0 = datetime(day.year, day.month, day.day, tzinfo=UTC) | |
| 57 | + d1 = d0 + timedelta(days=1) | |
| 58 | + mentions: dict[str, int] = defaultdict(int) | |
| 59 | + companies: dict[str, set[str]] = defaultdict(set) | |
| 60 | + async with transaction() as conn: | |
| 61 | + rows = await fetch_all(conn, """ | |
| 62 | + select e.title, e.company_id, co.display_name from events e join companies co on co.id = e.company_id | |
| 63 | + where e.detected_at >= :d0 and e.detected_at < :d1 and e.status in ('active', 'review') | |
| 64 | + union all | |
| 65 | + select n.title, n.company_id, co.display_name from news_items n join companies co on co.id = n.company_id | |
| 66 | + where n.first_seen_at >= :d0 and n.first_seen_at < :d1""", d0=d0, d1=d1) | |
| 67 | + for r in rows: | |
| 68 | + for term in extract_terms(r["title"], company_name=r["display_name"]): | |
| 69 | + mentions[term] += 1 | |
| 70 | + companies[term].add(r["company_id"]) | |
| 71 | + kept = {t: (mentions[t], len(companies[t])) for t in mentions if len(companies[t]) >= settings.trends_min_companies} | |
| 72 | + await execute(conn, "delete from trends where day = :d", d=day) | |
| 73 | + for term, (m, c) in kept.items(): | |
| 74 | + await execute(conn, "insert into trends (term, day, mentions, companies) values (:t, :d, :m, :c) on conflict (term, day) do update set mentions = excluded.mentions, companies = excluded.companies", | |
| 75 | + t=term[:120], d=day, m=m, c=c) | |
| 76 | + return {"day": day.isoformat(), "titles": len(rows), "terms_seen": len(mentions), "terms_kept": len(kept)} | |
| 77 | + | |
| 78 | + | |
| 79 | +async def compute_trends_range(days: int) -> list[dict[str, Any]]: | |
| 80 | + today = datetime.now(UTC).date() | |
| 81 | + return [await compute_trends(today - timedelta(days=i)) for i in range(days - 1, -1, -1)] | |
| 82 | + | |
| 83 | + | |
| 84 | +async def trend_momentum(window_days: int = 7, *, limit: int = 30) -> list[dict[str, Any]]: | |
| 85 | + """Terms ranked by coverage-normalised momentum: rate = mentions / active companies; momentum = (rate_now − rate_prev) / max(rate_prev, ε).""" | |
| 86 | + today = datetime.now(UTC).date() | |
| 87 | + cur_from, prev_from = today - timedelta(days=window_days - 1), today - timedelta(days=2 * window_days - 1) | |
| 88 | + async with transaction() as conn: | |
| 89 | + rows = await fetch_all(conn, """ | |
| 90 | + select term, day, mentions, companies from trends where day >= :prev_from order by term, day""", prev_from=prev_from) | |
| 91 | + cov_cur = await fetch_val(conn, "select coalesce(avg(companies_active), 0) from global_daily where day >= :f and day <= :t", f=cur_from, t=today) | |
| 92 | + cov_prev = await fetch_val(conn, "select coalesce(avg(companies_active), 0) from global_daily where day >= :f and day < :t", f=prev_from, t=cur_from) | |
| 93 | + cov_cur = float(cov_cur or 0) or 1.0 | |
| 94 | + cov_prev = float(cov_prev or 0) or cov_cur | |
| 95 | + per: dict[str, dict[str, Any]] = {} | |
| 96 | + for r in rows: | |
| 97 | + t = per.setdefault(r["term"], {"cur": 0, "prev": 0, "companies": set(), "series": {}}) | |
| 98 | + if r["day"] >= cur_from: | |
| 99 | + t["cur"] += r["mentions"] | |
| 100 | + t["companies"].add(r["companies"]) | |
| 101 | + else: | |
| 102 | + t["prev"] += r["mentions"] | |
| 103 | + t["series"][r["day"]] = r["mentions"] | |
| 104 | + out = [] | |
| 105 | + eps = 0.5 / cov_prev | |
| 106 | + for term, t in per.items(): | |
| 107 | + if t["cur"] == 0: | |
| 108 | + continue | |
| 109 | + rate_now, rate_prev = t["cur"] / cov_cur, t["prev"] / cov_prev | |
| 110 | + momentum = (rate_now - rate_prev) / max(rate_prev, eps) | |
| 111 | + series = [t["series"].get(cur_from + timedelta(days=i), 0) for i in range(window_days)] | |
| 112 | + out.append({"term": term, "mentions": t["cur"], "companies": max(t["companies"]) if t["companies"] else 0, "momentum": round(momentum, 3), "series": series, | |
| 113 | + "window_days": window_days, "formula_version": TRENDS_FORMULA_VERSION}) | |
| 114 | + out.sort(key=lambda x: (x["momentum"], x["mentions"]), reverse=True) | |
| 115 | + return out[:limit] | |
| 116 | + | |
| 117 | + | |
| 118 | +async def store_momentum_snapshots() -> None: | |
| 119 | + for w in (7, 30, 90): | |
| 120 | + items = await trend_momentum(w, limit=50) | |
| 121 | + async with transaction() as conn: | |
| 122 | + await execute(conn, "insert into settings_kv (key, value, updated_at) values (:k, cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()", | |
| 123 | + k=f"trends:momentum:{w}d", v=jsonb({"computed_at": datetime.now(UTC).isoformat(), "items": items})) | |
| 124 | + | |
| 125 | + | |
| 126 | +@periodic("trends", cron="25 * * * *") | |
| 127 | +async def trends_task() -> None: | |
| 128 | + today = datetime.now(UTC).date() | |
| 129 | + stats = await compute_trends(today) | |
| 130 | + if datetime.now(UTC).hour == 0: | |
| 131 | + await compute_trends(today - timedelta(days=1)) | |
| 132 | + await store_momentum_snapshots() | |
| 133 | + log.info("trends", extra=stats) | |
| 134 | + | |
| 135 | + | |
| 136 | +__all__ = ["STOPWORDS", "TRENDS_FORMULA_VERSION", "compute_trends", "compute_trends_range", "extract_terms", "store_momentum_snapshots", "trend_momentum"] | |
modified
src/companyatlas/taxonomy.py
+25 −2
@@ -351,6 +351,28 @@ AI_KEYWORDS = ("machine learning", "artificial intelligence", "deep learning", " | ||
| 351 | 351 | "generative", "genai", "gen ai", "nlp", "computer vision", "data scientist", "mlops", "agentic", "copilot", |
| 352 | 352 | "foundation model", "neural", "inference", "rag ", "retrieval-augmented", "prompt engineer", "ai engineer", "ai product") |
| 353 | 353 | |
| 354 | +# Metric formula parameters (intelligence layer, docs/SCORING.md). Bump METRICS_FORMULA_VERSION when any of these change. | |
| 355 | +METRIC_PARAMS: dict[str, float] = { | |
| 356 | + "activity_window_days": 30, "activity_decay_tau_days": 10.0, "activity_coverage_exp": 0.5, "activity_density_max": 6.0, | |
| 357 | + "change_weight_meaningful": 1.0, "change_weight_major": 2.0, "change_weight_critical": 3.0, | |
| 358 | + "hiring_min_listings": 3, "velocity_window_days": 90, "velocity_saturation": 6.0, "geo_saturation": 4.0, "geo_country_weight": 3.0, | |
| 359 | + "developer_saturation": 6.0, "communication_window_days": 30, "communication_saturation": 8.0, "pricing_saturation": 3.0, | |
| 360 | + "leadership_saturation": 3.0, "ai_jobs_weight": 0.5, "ai_events_weight": 0.25, "ai_keywords_weight": 0.25, "ai_events_saturation": 3.0, | |
| 361 | + "ai_keywords_saturation": 5.0, "coverage_expected_surfaces": 8, "coverage_obs_weight": 0.5, "coverage_continuity_weight": 0.3, | |
| 362 | + "coverage_sources_weight": 0.2, "anomaly_min_samples": 4, "index_trailing_days": 28, | |
| 363 | +} | |
| 364 | + | |
| 365 | +# Confidence assigned to deterministic events by evidence quality (spec §49–50). | |
| 366 | +EVIDENCE_CONFIDENCE: dict[str, float] = {"ats_json": 0.95, "jsonld": 0.9, "html": 0.8, "text_diff": 0.7} | |
| 367 | + | |
| 368 | +# Signal kinds (spec §92) — labelled as signals, never facts. | |
| 369 | +SIGNAL_KINDS = ("hiring_surge", "hiring_freeze", "launch_buildup", "expansion", "pricing_migration", "developer_push", | |
| 370 | + "enterprise_repositioning", "ai_acceleration", "abnormal_activity") | |
| 371 | +SIGNALS_FORMULA_VERSION = "signals-v1" | |
| 372 | + | |
| 373 | +# Words the intelligence layer must never emit in generated titles/summaries (spec §167–168). | |
| 374 | +FORBIDDEN_WORDING = ("laid off", "layoff", "fired", "shut down", "shutdown", "bankrupt", "collapsed") | |
| 375 | + | |
| 354 | 376 | # ------------------------------------------------------------------------------------------------------------ companies |
| 355 | 377 | |
| 356 | 378 | COMPANY_IMPORTANCE_TIERS = {1: "global", 2: "major", 3: "notable", 4: "long_tail"} |
@@ -362,8 +384,9 @@ class CollectionMethod(StrEnum): | ||
| 362 | 384 | |
| 363 | 385 | |
| 364 | 386 | __all__ = [ |
| 365 | − "AI_KEYWORDS", "CCI_FORMULA_VERSION", "CCI_WEIGHTS", "COMPANY_IMPORTANCE_TIERS", "EVENT_SUBTYPES", "FAILURE_POLICY", | |
| 366 | − "METRICS_FORMULA_VERSION", "MVP_EVENT_SUBTYPES", "SURFACE_BASE_INTERVAL_S", "SURFACE_IMPORTANCE", "ChangeKind", "CollectionMethod", | |
| 387 | + "AI_KEYWORDS", "CCI_FORMULA_VERSION", "CCI_WEIGHTS", "COMPANY_IMPORTANCE_TIERS", "EVENT_SUBTYPES", "EVIDENCE_CONFIDENCE", "FAILURE_POLICY", | |
| 388 | + "FORBIDDEN_WORDING", "METRICS_FORMULA_VERSION", "METRIC_PARAMS", "MVP_EVENT_SUBTYPES", "SIGNALS_FORMULA_VERSION", "SIGNAL_KINDS", | |
| 389 | + "SURFACE_BASE_INTERVAL_S", "SURFACE_IMPORTANCE", "ChangeKind", "CollectionMethod", | |
| 367 | 390 | "CompanyStatus", "EventStatus", "EventType", "FailureClass", "FetchMode", "Metric", "OnboardingStatus", "SensorStatus", "Surface", |
| 368 | 391 | "change_kind", "confidence_label", "tier_for_interval", |
| 369 | 392 | ] |
added
tests/factories.py
+168 −0
@@ -0,0 +1,168 @@ | ||
| 1 | +"""Test data factories for the intelligence layer (slug prefix `ztest-`; `cleanup()` removes everything by cascade).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import contextlib | |
| 5 | +import uuid | |
| 6 | +from datetime import UTC, datetime, timedelta | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +import pytest | |
| 10 | + | |
| 11 | +from companyatlas.db import dispose, execute, fetch_one, jsonb, transaction | |
| 12 | +from companyatlas.ids import new_id, normalize_alias, stable_hash | |
| 13 | + | |
| 14 | +PREFIX = "ztest-" | |
| 15 | +CONNECTOR_HTML = "ztest-generic-html-v1" | |
| 16 | +CONNECTOR_ATS = "ztest-greenhouse-v1" | |
| 17 | + | |
| 18 | + | |
| 19 | +def _uid() -> str: | |
| 20 | + return uuid.uuid4().hex[:8] | |
| 21 | + | |
| 22 | + | |
| 23 | +async def ensure_reference(conn) -> None: # type: ignore[no-untyped-def] | |
| 24 | + for code, name in (("CA", "Canada"), ("US", "United States"), ("JP", "Japan"), ("DE", "Germany"), ("GB", "United Kingdom")): | |
| 25 | + await execute(conn, "insert into countries (code, name) values (:c, :n) on conflict (code) do nothing", c=code, n=name) | |
| 26 | + for cid, cat, mode in ((CONNECTOR_HTML, "homepage", "http"), (CONNECTOR_ATS, "jobs_board", "json")): | |
| 27 | + await execute(conn, "insert into connectors (id, name, version, category, fetch_mode) values (:id, :id, 'v1', :cat, :mode) on conflict (id) do nothing", | |
| 28 | + id=cid, cat=cat, mode=mode) | |
| 29 | + | |
| 30 | + | |
| 31 | +async def make_company(conn, *, name: str | None = None, country: str = "CA", industries: list[str] | None = None, # type: ignore[no-untyped-def] | |
| 32 | + first_observed_days_ago: int | None = 30) -> dict[str, Any]: | |
| 33 | + await ensure_reference(conn) | |
| 34 | + uid = _uid() | |
| 35 | + slug = f"{PREFIX}{uid}" | |
| 36 | + display = name or f"ZTest Corp {uid}" | |
| 37 | + cid = new_id("company") | |
| 38 | + first = datetime.now(UTC) - timedelta(days=first_observed_days_ago) if first_observed_days_ago is not None else None | |
| 39 | + await execute(conn, """ | |
| 40 | + insert into companies (id, slug, display_name, canonical_domain, website, industries, country, onboarding_status, first_observed_at, last_observed_at) | |
| 41 | + values (:id, :slug, :name, :domain, :web, cast(:ind as text[]), :country, 'active', :first, :first)""", | |
| 42 | + id=cid, slug=slug, name=display, domain=f"{slug}.example", web=f"https://{slug}.example", ind=industries or [], country=country, first=first) | |
| 43 | + await execute(conn, "insert into company_aliases (company_id, alias, alias_norm) values (:c, :a, :n) on conflict do nothing", c=cid, a=display, n=normalize_alias(display)) | |
| 44 | + return {"id": cid, "slug": slug, "display_name": display, "country": country, "industries": industries or [], "canonical_domain": f"{slug}.example"} | |
| 45 | + | |
| 46 | + | |
| 47 | +async def make_sensor(conn, company: dict[str, Any], surface: str, *, connector_id: str = CONNECTOR_HTML, path: str | None = None, # type: ignore[no-untyped-def] | |
| 48 | + status: str = "active", created_days_ago: int = 30, interval_s: int = 86400) -> dict[str, Any]: | |
| 49 | + sid = new_id("sensor") | |
| 50 | + url = f"https://{company['canonical_domain']}/{path or surface}" | |
| 51 | + await execute(conn, """ | |
| 52 | + insert into sensors (id, company_id, surface, connector_id, url, canonical_url, domain, status, base_interval_s, current_interval_s, created_at, last_success_at) | |
| 53 | + values (:id, :c, :surface, :conn, :url, :url, :domain, :status, :iv, :iv, :created, now())""", | |
| 54 | + id=sid, c=company["id"], surface=surface, conn=connector_id, url=url, domain=company["canonical_domain"], status=status, iv=interval_s, | |
| 55 | + created=datetime.now(UTC) - timedelta(days=created_days_ago)) | |
| 56 | + return {"id": sid, "company_id": company["id"], "surface": surface, "connector_id": connector_id, "url": url} | |
| 57 | + | |
| 58 | + | |
| 59 | +async def make_snapshot(conn, sensor: dict[str, Any], *, version_no: int = 1, title: str | None = None, fetched_at: datetime | None = None) -> str: # type: ignore[no-untyped-def] | |
| 60 | + snap = new_id("snapshot") | |
| 61 | + h = stable_hash(snap) | |
| 62 | + await execute(conn, """ | |
| 63 | + insert into snapshots (id, sensor_id, company_id, version_no, fetched_at, content_hash, normalized_hash, structural_hash, title) | |
| 64 | + values (:id, :s, :c, :v, :at, :h, :h, :h, :title)""", id=snap, s=sensor["id"], c=sensor["company_id"], v=version_no, at=fetched_at or datetime.now(UTC), h=h, title=title) | |
| 65 | + return snap | |
| 66 | + | |
| 67 | + | |
| 68 | +async def make_change(conn, sensor: dict[str, Any], *, significance: float = 0.6, kind: str | None = None, structured_delta: dict[str, Any] | None = None, # type: ignore[no-untyped-def] | |
| 69 | + diff: dict[str, Any] | None = None, detected_at: datetime | None = None, status: str = "pending") -> dict[str, Any]: | |
| 70 | + from companyatlas.taxonomy import change_kind | |
| 71 | + | |
| 72 | + kind = kind or str(change_kind(significance)) | |
| 73 | + before = await make_snapshot(conn, sensor, version_no=1, fetched_at=(detected_at or datetime.now(UTC)) - timedelta(days=1)) | |
| 74 | + after = await make_snapshot(conn, sensor, version_no=2, fetched_at=detected_at) | |
| 75 | + diff = diff or {"added": [], "removed": [], "modified": [], "moved": [], "counts": {"added": 0, "removed": 0, "modified": 0, "moved": 0}, | |
| 76 | + "text_delta_ratio": 0.1, "similarity": 0.9, "reasons": []} | |
| 77 | + cid = new_id("change") | |
| 78 | + await execute(conn, """ | |
| 79 | + insert into changes (id, sensor_id, company_id, surface, snapshot_before, snapshot_after, detected_at, significance, kind, blocks_added, blocks_removed, | |
| 80 | + blocks_modified, text_delta_ratio, similarity, diff, structured_delta, status) | |
| 81 | + values (:id, :s, :c, :surface, :before, :after, :at, :sig, :kind, :ba, :br, :bm, :ratio, :sim, cast(:diff as jsonb), cast(:delta as jsonb), :status)""", | |
| 82 | + id=cid, s=sensor["id"], c=sensor["company_id"], surface=sensor["surface"], before=before, after=after, at=detected_at or datetime.now(UTC), sig=significance, | |
| 83 | + kind=kind, ba=len(diff.get("added") or []), br=len(diff.get("removed") or []), bm=len(diff.get("modified") or []), ratio=diff.get("text_delta_ratio") or 0, | |
| 84 | + sim=diff.get("similarity"), diff=jsonb(diff), delta=jsonb(structured_delta or {}), status=status) | |
| 85 | + return await fetch_one(conn, "select * from changes where id = :id", id=cid) or {"id": cid} | |
| 86 | + | |
| 87 | + | |
| 88 | +async def make_job(conn, company: dict[str, Any], *, title: str = "Software Engineer", first_seen_days_ago: float = 10, removed_days_ago: float | None = None, # type: ignore[no-untyped-def] | |
| 89 | + country: str | None = "CA", is_ai: bool = False, remote: bool = False, department: str | None = "Engineering", sensor_id: str | None = None) -> str: | |
| 90 | + jid = new_id("job") | |
| 91 | + now = datetime.now(UTC) | |
| 92 | + first = now - timedelta(days=first_seen_days_ago) | |
| 93 | + removed = now - timedelta(days=removed_days_ago) if removed_days_ago is not None else None | |
| 94 | + await execute(conn, """ | |
| 95 | + insert into jobs (id, company_id, sensor_id, fingerprint, title, department, country, remote, first_seen_at, last_seen_at, removed_at, status, is_ai) | |
| 96 | + values (:id, :c, :s, :fp, :title, :dep, :country, :remote, :first, :last, :removed, :status, :ai)""", | |
| 97 | + id=jid, c=company["id"], s=sensor_id, fp=stable_hash(jid), title=title, dep=department, country=country, remote=remote, first=first, | |
| 98 | + last=removed or now, removed=removed, status="no_longer_listed" if removed else "open", ai=is_ai) | |
| 99 | + return jid | |
| 100 | + | |
| 101 | + | |
| 102 | +async def make_event(conn, company: dict[str, Any], *, subtype: str, importance: float = 0.6, days_ago: float = 1, title: str | None = None, # type: ignore[no-untyped-def] | |
| 103 | + tags: list[str] | None = None, surface: str | None = None, sensor_id: str | None = None) -> str: | |
| 104 | + from companyatlas.taxonomy import EVENT_SUBTYPES, EventType, confidence_label | |
| 105 | + | |
| 106 | + eid = new_id("event") | |
| 107 | + etype = str(EVENT_SUBTYPES.get(subtype, (EventType.OTHER, 0.3))[0]) | |
| 108 | + await execute(conn, """ | |
| 109 | + insert into events (id, company_id, sensor_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, tags, detected_at, dedupe_key) | |
| 110 | + values (:id, :c, :s, :surface, :t, :st, :imp, 0.8, :label, :title, cast(:tags as text[]), :at, :dk)""", | |
| 111 | + id=eid, c=company["id"], s=sensor_id, surface=surface, t=etype, st=subtype, imp=importance, label=confidence_label(0.8), | |
| 112 | + title=title or f"{subtype} test event", tags=tags or [], at=datetime.now(UTC) - timedelta(days=days_ago), dk=stable_hash(eid)) | |
| 113 | + return eid | |
| 114 | + | |
| 115 | + | |
| 116 | +async def make_location(conn, company: dict[str, Any], *, name: str, country: str, city: str | None = None, first_seen_days_ago: float = 5) -> str: # type: ignore[no-untyped-def] | |
| 117 | + lid = new_id("location") | |
| 118 | + await execute(conn, """insert into locations (id, company_id, kind, name, name_norm, city, country, first_seen_at, last_seen_at) | |
| 119 | + values (:id, :c, 'office', :name, :norm, :city, :country, :first, now())""", | |
| 120 | + id=lid, c=company["id"], name=name, norm=normalize_alias(name) + _uid(), city=city, country=country, first=datetime.now(UTC) - timedelta(days=first_seen_days_ago)) | |
| 121 | + return lid | |
| 122 | + | |
| 123 | + | |
| 124 | +async def make_observation(conn, sensor: dict[str, Any], *, days_ago: float, changed: bool = False, not_modified: bool = False) -> str: # type: ignore[no-untyped-def] | |
| 125 | + oid = new_id("observation") | |
| 126 | + await execute(conn, """insert into observations (id, sensor_id, company_id, fetched_at, status_code, changed, not_modified) | |
| 127 | + values (:id, :s, :c, :at, 200, :changed, :nm)""", id=oid, s=sensor["id"], c=sensor["company_id"], | |
| 128 | + at=datetime.now(UTC) - timedelta(days=days_ago), changed=changed, nm=not_modified) | |
| 129 | + return oid | |
| 130 | + | |
| 131 | + | |
| 132 | +@pytest.fixture | |
| 133 | +async def intel_db(): | |
| 134 | + """DB fixture for the intelligence tests: a fresh engine bound to this test's event loop (other modules run on a session-scoped loop), | |
| 135 | + skip when Postgres is unreachable, dispose afterwards. Import it into a test module to register it.""" | |
| 136 | + from companyatlas.config import settings | |
| 137 | + from companyatlas.db import fetch_val | |
| 138 | + | |
| 139 | + with contextlib.suppress(Exception): | |
| 140 | + await dispose() | |
| 141 | + try: | |
| 142 | + async with transaction() as conn: | |
| 143 | + ok = (await fetch_val(conn, "select 1")) == 1 | |
| 144 | + except Exception: # noqa: BLE001 | |
| 145 | + ok = False | |
| 146 | + with contextlib.suppress(Exception): | |
| 147 | + await dispose() | |
| 148 | + if not ok: | |
| 149 | + pytest.skip(f"database not reachable: {settings.database_url.split('@')[-1]}") | |
| 150 | + try: | |
| 151 | + yield | |
| 152 | + finally: | |
| 153 | + with contextlib.suppress(Exception): | |
| 154 | + await dispose() | |
| 155 | + | |
| 156 | + | |
| 157 | +async def cleanup() -> None: | |
| 158 | + async with transaction() as conn: | |
| 159 | + await execute(conn, "delete from companies where slug like :p", p=f"{PREFIX}%") | |
| 160 | + await execute(conn, "delete from connectors where id like :p", p=f"{PREFIX}%") | |
| 161 | + await execute(conn, "delete from alerts where name like :p", p=f"{PREFIX}%") | |
| 162 | + await execute(conn, "delete from owners where token_hash like :p", p=f"{PREFIX}%") | |
| 163 | + await execute(conn, "delete from global_daily where day < '2002-01-01'") | |
| 164 | + await execute(conn, "delete from trends where day < '2002-01-01'") | |
| 165 | + | |
| 166 | + | |
| 167 | +__all__ = ["CONNECTOR_ATS", "CONNECTOR_HTML", "PREFIX", "cleanup", "ensure_reference", "intel_db", "make_change", "make_company", "make_event", "make_job", "make_location", | |
| 168 | + "make_observation", "make_sensor", "make_snapshot"] | |
added
tests/test_api_admin.py
+139 −0
@@ -0,0 +1,139 @@ | ||
| 1 | +"""Admin API: auth, overview/quality/costs shapes, sensor actions, company creation, queue, reviews, event corrections, cache.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | +import test_api_support as support | |
| 6 | + | |
| 7 | +client = support.client | |
| 8 | +fixture_data = support.fixture_data | |
| 9 | +ADMIN = support.ADMIN | |
| 10 | + | |
| 11 | +from companyatlas.api.common import cache | |
| 12 | +from companyatlas.db import execute, fetch_one, fetch_val, transaction | |
| 13 | + | |
| 14 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 15 | +V = "/api/v1/admin" | |
| 16 | + | |
| 17 | + | |
| 18 | +async def test_admin_requires_token(client): # type: ignore[no-untyped-def] | |
| 19 | + for path in ("/overview", "/connectors", "/sensors", "/companies", "/failures", "/queue", "/llm", "/reviews", "/quality", "/costs"): | |
| 20 | + r = await client.get(f"{V}{path}") | |
| 21 | + assert r.status_code == 401 and r.json() == {"detail": "admin token required"}, path | |
| 22 | + assert (await client.get(f"{V}/overview", headers={"X-CA-Admin-Token": "wrong"})).status_code == 401 | |
| 23 | + assert (await client.post(f"{V}/cache/clear")).status_code == 401 | |
| 24 | + | |
| 25 | + | |
| 26 | +async def test_overview_connectors_quality_costs(client, fixture_data): # type: ignore[no-untyped-def] | |
| 27 | + r = await client.get(f"{V}/overview", headers=ADMIN) | |
| 28 | + assert r.status_code == 200 and r.headers["cache-control"] == "no-store" and r.headers["x-ratelimit-tier"] == "admin" | |
| 29 | + body = r.json() | |
| 30 | + for key in ("companies_by_status", "sensors_by_status", "sensors_by_tier", "queue", "llm", "failures_24h_by_class", "fetch_rate_1h", "change_rate_1h", | |
| 31 | + "meaningful_rate_1h", "storage", "workers", "cost_today"): | |
| 32 | + assert key in body, key | |
| 33 | + assert body["queue"]["dead"] >= 1 and body["failures_24h_by_class"].get("HTTP_5XX", 0) >= 1 and set(body["cost_today"]) >= {"fetch", "browser", "llm"} | |
| 34 | + conns = (await client.get(f"{V}/connectors", headers=ADMIN)).json()["items"] | |
| 35 | + mine = next(c for c in conns if c["id"] == fixture_data["connector"]) | |
| 36 | + assert mine["sensors_active"] == 3 and {"success_rate_24h", "avg_latency_ms", "change_rate_24h", "errors_24h", "last_run_at"} <= set(mine) | |
| 37 | + q = (await client.get(f"{V}/quality", headers=ADMIN)).json() | |
| 38 | + for key in ("coverage", "freshness", "duplicate_rate", "event_confidence_avg", "unknown_surfaces", "failed_sensors", "calibration"): | |
| 39 | + assert key in q | |
| 40 | + assert set(q["calibration"]) == {"correct", "duplicate", "noise", "misclassified"} | |
| 41 | + c = (await client.get(f"{V}/costs?days=7", headers=ADMIN)).json() | |
| 42 | + assert {"items", "per_1000_companies", "per_million_observations", "per_meaningful_event"} <= set(c) | |
| 43 | + | |
| 44 | + | |
| 45 | +async def test_sensors_list_filters_and_actions(client, fixture_data): # type: ignore[no-untyped-def] | |
| 46 | + sid = fixture_data["sensor_beta"] | |
| 47 | + r = await client.get(f"{V}/sensors", headers=ADMIN, params={"company": fixture_data["beta_slug"]}) | |
| 48 | + body = r.json() | |
| 49 | + assert r.status_code == 200 and body["total"] == 1 and body["items"][0]["company"]["slug"] == fixture_data["beta_slug"] | |
| 50 | + assert (await client.get(f"{V}/sensors", headers=ADMIN, params={"filter": "healthy", "connector": fixture_data["connector"]})).json()["total"] == 3 | |
| 51 | + assert (await client.get(f"{V}/sensors", headers=ADMIN, params={"filter": "high_activity", "connector": fixture_data["connector"]})).json()["total"] == 3 | |
| 52 | + assert (await client.get(f"{V}/sensors", headers=ADMIN, params={"filter": "nope"})).status_code == 422 | |
| 53 | + r = await client.post(f"{V}/sensors/{sid}/pause", headers=ADMIN) | |
| 54 | + assert r.status_code == 200 and r.json()["sensor"]["status"] == "paused" | |
| 55 | + assert (await client.get(f"{V}/sensors", headers=ADMIN, params={"status": "paused", "domain": f"{fixture_data['beta_slug']}.example"})).json()["total"] == 1 | |
| 56 | + r = await client.post(f"{V}/sensors/{sid}/resume", headers=ADMIN) | |
| 57 | + assert r.json()["sensor"]["status"] == "active" and r.json()["sensor"]["consecutive_failures"] == 0 | |
| 58 | + r = await client.post(f"{V}/sensors/{sid}/set_interval", headers=ADMIN, json={"interval_s": 3600}) | |
| 59 | + assert r.json()["interval_s"] == 3600 and r.json()["sensor"]["current_interval_s"] == 3600 and r.json()["sensor"]["tier"] == "B" | |
| 60 | + assert (await client.post(f"{V}/sensors/{sid}/set_interval", headers=ADMIN, json={})).status_code == 422 | |
| 61 | + assert (await client.post(f"{V}/sensors/{sid}/set_connector", headers=ADMIN, json={"connector_id": "nope-v9"})).status_code == 422 | |
| 62 | + assert (await client.post(f"{V}/sensors/{sid}/set_connector", headers=ADMIN, json={"connector_id": fixture_data["connector"]})).status_code == 200 | |
| 63 | + r = await client.post(f"{V}/sensors/{sid}/run_now", headers=ADMIN) | |
| 64 | + assert r.status_code == 200 | |
| 65 | + r = await client.post(f"{V}/sensors/{sid}/rediscover", headers=ADMIN, json={"reason": "test"}) | |
| 66 | + assert r.status_code == 200 and r.json()["queued"]["key"].startswith(f"discover:{fixture_data['beta']}:") | |
| 67 | + async with transaction() as conn: | |
| 68 | + job = await fetch_one(conn, "select kind, status, payload from queue_jobs where id = :id", id=r.json()["queued"]["id"]) | |
| 69 | + assert job["kind"] == "discover" and job["status"] == "pending" | |
| 70 | + r = await client.post(f"{V}/sensors/{sid}/retire", headers=ADMIN) | |
| 71 | + assert r.json()["sensor"]["status"] == "retired" and r.json()["sensor"]["retired_at"] | |
| 72 | + assert (await client.post(f"{V}/sensors/{sid}/resume", headers=ADMIN)).json()["sensor"]["status"] == "active" | |
| 73 | + assert (await client.post(f"{V}/sensors/{sid}/explode", headers=ADMIN)).status_code == 404 | |
| 74 | + assert (await client.post(f"{V}/sensors/sen_nope/pause", headers=ADMIN)).status_code == 404 | |
| 75 | + | |
| 76 | + | |
| 77 | +async def test_companies_list_create_rediscover(client, fixture_data): # type: ignore[no-untyped-def] | |
| 78 | + r = await client.get(f"{V}/companies", headers=ADMIN, params={"onboarding_status": "active", "country": "ZZ"}) | |
| 79 | + assert r.status_code == 200 and r.json()["total"] == 2 and "onboarding_error" in r.json()["items"][0] | |
| 80 | + suffix = fixture_data["suffix"] | |
| 81 | + r = await client.post(f"{V}/companies", headers=ADMIN, json={"website": f"https://www.ztest-created-{suffix}.example/home?utm_source=x", "display_name": f"Ztest Created {suffix}", | |
| 82 | + "country": "ZZ", "industries": [fixture_data["industry"]]}) | |
| 83 | + assert r.status_code == 201, r.text | |
| 84 | + body = r.json() | |
| 85 | + card = body["company"] | |
| 86 | + assert card["canonical_domain"] == f"ztest-created-{suffix}.example" and card["website"] == f"https://www.ztest-created-{suffix}.example/home" | |
| 87 | + assert card["slug"] == f"ztest-created-{suffix}" and card["onboarding_status"] == "pending" and body["queued"]["kind"] == "discover" | |
| 88 | + async with transaction() as conn: | |
| 89 | + assert await fetch_val(conn, "select status from queue_jobs where key = :k", k=f"discover:{card['id']}") == "pending" | |
| 90 | + dup = await client.post(f"{V}/companies", headers=ADMIN, json={"website": f"http://ztest-created-{suffix}.example"}) | |
| 91 | + assert dup.status_code == 409 | |
| 92 | + assert (await client.post(f"{V}/companies", headers=ADMIN, json={"website": "ftp://ztest.example"})).status_code == 422 | |
| 93 | + assert (await client.post(f"{V}/companies", headers=ADMIN, json={"website": "https://localhost"})).status_code == 422 | |
| 94 | + assert (await client.post(f"{V}/companies", headers=ADMIN, json={"website": f"https://ztest-created2-{suffix}.example", "country": "QQ"})).status_code == 422 | |
| 95 | + assert (await client.post(f"{V}/companies", headers=ADMIN, json={"website": f"https://ztest-created2-{suffix}.example", "industries": ["nope-ind"]})).status_code == 422 | |
| 96 | + r = await client.post(f"{V}/companies/{card['slug']}/rediscover", headers=ADMIN) | |
| 97 | + assert r.status_code == 200 and r.json()["queued"]["key"].startswith(f"discover:{card['id']}:") | |
| 98 | + assert (await client.post(f"{V}/companies/nope/rediscover", headers=ADMIN)).status_code == 404 | |
| 99 | + async with transaction() as conn: # remove the created company so the ZZ counts used by the other modules stay stable | |
| 100 | + await execute(conn, "delete from queue_jobs where key like :k", k=f"discover:{card['id']}%") | |
| 101 | + await execute(conn, "delete from companies where id = :id", id=card["id"]) | |
| 102 | + cache.clear() | |
| 103 | + | |
| 104 | + | |
| 105 | +async def test_failures_queue_llm_reviews(client, fixture_data): # type: ignore[no-untyped-def] | |
| 106 | + r = await client.get(f"{V}/failures", headers=ADMIN, params={"class": "http_5xx", "company": fixture_data["beta_slug"]}) | |
| 107 | + assert r.status_code == 200 and r.json()["total"] == 1 and r.json()["by_class"] == {"HTTP_5XX": 1} and "HTTP_5XX" in r.json()["classes"] | |
| 108 | + q = (await client.get(f"{V}/queue", headers=ADMIN, params={"status": "dead"})).json() | |
| 109 | + assert any(i["id"] == fixture_data["queue_dead"] for i in q["items"]) and any(c["status"] == "dead" for c in q["counts"]) | |
| 110 | + r = await client.post(f"{V}/queue/requeue-dead", headers=ADMIN, json={"kind": "run_sensor"}) | |
| 111 | + assert r.status_code == 200 and r.json()["requeued"] >= 1 | |
| 112 | + async with transaction() as conn: | |
| 113 | + assert await fetch_val(conn, "select status from queue_jobs where id = :id", id=fixture_data["queue_dead"]) == "pending" | |
| 114 | + llm = (await client.get(f"{V}/llm", headers=ADMIN)).json() | |
| 115 | + assert {"items", "stats"} <= set(llm) and "budget" in llm["stats"] | |
| 116 | + reviews = (await client.get(f"{V}/reviews", headers=ADMIN, params={"kind": "major_event"})).json() | |
| 117 | + assert any(x["id"] == fixture_data["review"] for x in reviews["items"]) and reviews["open_by_kind"].get("major_event", 0) >= 1 | |
| 118 | + r = await client.post(f"{V}/reviews/{fixture_data['review']}", headers=ADMIN, json={"resolution": "accepted", "label": "correct", "note": "looks right"}) | |
| 119 | + assert r.status_code == 200 and r.json()["review"]["status"] == "accepted" and r.json()["review"]["payload"]["resolution"]["label"] == "correct" | |
| 120 | + assert (await client.post(f"{V}/reviews/{fixture_data['review']}", headers=ADMIN, json={"resolution": "rejected"})).status_code == 409 | |
| 121 | + assert (await client.post(f"{V}/reviews/rev_nope", headers=ADMIN, json={"resolution": "rejected"})).status_code == 404 | |
| 122 | + assert (await client.get(f"{V}/quality", headers=ADMIN)).json()["calibration"]["correct"] >= 1 | |
| 123 | + | |
| 124 | + | |
| 125 | +async def test_event_retract_restore_and_cache_clear(client, fixture_data): # type: ignore[no-untyped-def] | |
| 126 | + eid = fixture_data["ev_product"] | |
| 127 | + r = await client.post(f"{V}/events/{eid}/retract", headers=ADMIN, json={"reason": "synthetic correction"}) | |
| 128 | + assert r.status_code == 200 and r.json()["event"]["status"] == "retracted" and r.json()["event"]["retracted_reason"] == "synthetic correction" | |
| 129 | + assert (await client.get(f"/api/v1/events/{eid}")).json()["status"] == "retracted" | |
| 130 | + assert all(e["id"] != eid for e in (await client.get("/api/v1/live?limit=50")).json()["items"]) | |
| 131 | + assert (await client.post(f"{V}/events/{eid}/retract", headers=ADMIN, json={"reason": "x"})).status_code == 422 | |
| 132 | + r = await client.post(f"{V}/events/{eid}/restore", headers=ADMIN) | |
| 133 | + assert r.status_code == 200 and r.json()["event"]["status"] == "active" | |
| 134 | + audit = r.json()["event"]["payload"]["_audit"] | |
| 135 | + assert [a["action"] for a in audit] == ["retract", "restore"] and audit[0]["reason"] == "synthetic correction" | |
| 136 | + assert (await client.post(f"{V}/events/evt_nope/restore", headers=ADMIN)).status_code == 404 | |
| 137 | + r = await client.post(f"{V}/cache/clear", headers=ADMIN) | |
| 138 | + assert r.status_code == 200 and r.json()["ok"] is True | |
| 139 | + assert (await client.post(f"{V}/cache/clear?prefix=pulse", headers=ADMIN)).json()["cleared"] == "pulse" | |
added
tests/test_api_companies.py
+156 −0
@@ -0,0 +1,156 @@ | ||
| 1 | +"""Companies: list/filters/pagination, detail, sub-resources, compare, events feed and event detail.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | +import test_api_support as support | |
| 6 | + | |
| 7 | +client = support.client | |
| 8 | +fixture_data = support.fixture_data | |
| 9 | + | |
| 10 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 11 | +V = "/api/v1" | |
| 12 | +CARD_KEYS = {"id", "slug", "display_name", "legal_name", "canonical_domain", "website", "description", "industries", "industry_primary", "country", "hq_city", "hq_region", | |
| 13 | + "public_company", "ticker", "exchange", "founded_year", "employees_band", "logo_url", "status", "onboarding_status", "importance", "tier", "metrics", "counts", | |
| 14 | + "last_event_at", "last_observed_at"} | |
| 15 | + | |
| 16 | + | |
| 17 | +async def test_list_filters_sort_pagination(client, fixture_data): # type: ignore[no-untyped-def] | |
| 18 | + r = await client.get(f"{V}/companies", params={"country": "ZZ", "per_page": 1, "page": 1, "sort": "activity", "sparkline": 1}) | |
| 19 | + body = r.json() | |
| 20 | + assert r.status_code == 200 and set(body) == {"items", "page", "per_page", "total", "pages"} | |
| 21 | + assert body["total"] == 2 and body["pages"] == 2 and len(body["items"]) == 1 | |
| 22 | + card = body["items"][0] | |
| 23 | + assert CARD_KEYS <= set(card) and card["slug"] == fixture_data["alpha_slug"] | |
| 24 | + assert card["metrics"]["activity_score"] == 72.3 and card["metrics"]["open_jobs"] == 2 and card["metrics"]["hiring_momentum_30d"] == 12.5 | |
| 25 | + assert card["counts"] == {"sensors": 2, "observations": 11, "changes": 1, "events": 3, "jobs_open": 2} | |
| 26 | + assert len(card["sparkline"]) == 10 | |
| 27 | + page2 = (await client.get(f"{V}/companies", params={"country": "ZZ", "per_page": 1, "page": 2})).json() | |
| 28 | + assert page2["items"][0]["slug"] == fixture_data["beta_slug"] | |
| 29 | + hiring = (await client.get(f"{V}/companies", params={"country": "ZZ", "sort": "hiring"})).json()["items"] | |
| 30 | + assert [c["slug"] for c in hiring] == [fixture_data["alpha_slug"], fixture_data["beta_slug"]] | |
| 31 | + name = (await client.get(f"{V}/companies", params={"industry": fixture_data["industry"], "sort": "name"})).json()["items"] | |
| 32 | + assert [c["slug"] for c in name] == [fixture_data["alpha_slug"], fixture_data["beta_slug"]] | |
| 33 | + q = (await client.get(f"{V}/companies", params={"q": f"ztest beta {fixture_data['suffix']}"})).json()["items"] | |
| 34 | + assert q and q[0]["slug"] == fixture_data["beta_slug"] | |
| 35 | + assert (await client.get(f"{V}/companies", params={"country": "ZZ", "public": "true"})).json()["total"] == 0 | |
| 36 | + assert (await client.get(f"{V}/companies", params={"country": "ZZ", "has_events": "1", "tier": 2})).json()["total"] == 2 | |
| 37 | + assert (await client.get(f"{V}/companies", params={"per_page": 999})).status_code == 422 | |
| 38 | + | |
| 39 | + | |
| 40 | +async def test_detail_by_slug_and_id(client, fixture_data): # type: ignore[no-untyped-def] | |
| 41 | + for key in (fixture_data["alpha_slug"], fixture_data["alpha"]): | |
| 42 | + r = await client.get(f"{V}/companies/{key}") | |
| 43 | + assert r.status_code == 200, key | |
| 44 | + body = r.json() | |
| 45 | + for k in ("aliases", "domains", "relationships", "metrics_detail", "sensors_by_surface", "coverage", "signals", "sparklines", "sparkline"): | |
| 46 | + assert k in body, k | |
| 47 | + assert body["aliases"] == [f"Alphaz {fixture_data['suffix']}"] | |
| 48 | + assert body["relationships"][0]["kind"] == "PARTNER_OF" and body["relationships"][0]["company"]["slug"] == fixture_data["beta_slug"] | |
| 49 | + assert body["sensors_by_surface"] == {"pricing": 1, "careers": 1} | |
| 50 | + assert body["coverage"]["days_observed"] == 1 and body["coverage"]["sensor_uptime"] == 1.0 and body["coverage"]["historical_coverage"] is None | |
| 51 | + assert len(body["sparklines"]["activity_30d"]) == 10 and body["sparklines"]["hiring_90d"] == [] | |
| 52 | + assert body["signals"][0]["kind"] == "hiring_surge" | |
| 53 | + md = {m["metric"]: m for m in body["metrics_detail"]} | |
| 54 | + assert md["activity_score"]["value"] == 72.3 and md["activity_score"]["formula_version"] == "metrics-v1" and md["activity_score"]["inputs"] == {"events": 3} | |
| 55 | + r = await client.get(f"{V}/companies/does-not-exist") | |
| 56 | + assert r.status_code == 404 and r.json() == {"detail": "company not found"} | |
| 57 | + | |
| 58 | + | |
| 59 | +async def test_events_timeline_metrics(client, fixture_data): # type: ignore[no-untyped-def] | |
| 60 | + s = fixture_data["alpha_slug"] | |
| 61 | + r = await client.get(f"{V}/companies/{s}/events") | |
| 62 | + body = r.json() | |
| 63 | + assert r.status_code == 200 and body["total"] == 3 and [e["event_type"] for e in body["items"]] == ["PRICING", "HIRING", "PRODUCT"] | |
| 64 | + imp = (await client.get(f"{V}/companies/{s}/events?sort=importance")).json()["items"] | |
| 65 | + assert imp[0]["event_type"] == "PRICING" | |
| 66 | + assert (await client.get(f"{V}/companies/{s}/events?event_type=HIRING")).json()["total"] == 1 | |
| 67 | + assert (await client.get(f"{V}/companies/{s}/events?surface=careers")).json()["total"] == 1 | |
| 68 | + assert (await client.get(f"{V}/companies/{s}/events?min_importance=0.75")).json()["total"] == 1 | |
| 69 | + tl = (await client.get(f"{V}/companies/{s}/timeline?filter=pricing")).json() | |
| 70 | + assert [e["event_type"] for e in tl["items"]] == ["PRICING"] and tl["items"][0]["day"] == tl["days"][0]["day"] and tl["days"][0]["count"] == 1 | |
| 71 | + tl_all = (await client.get(f"{V}/companies/{s}/timeline")).json() | |
| 72 | + assert len(tl_all["items"]) == 3 and sum(d["count"] for d in tl_all["days"]) == 3 | |
| 73 | + m = (await client.get(f"{V}/companies/{s}/metrics?days=30")).json() | |
| 74 | + assert {c["metric"] for c in m["current"]} >= {"activity_score", "open_jobs"} and len(m["series"]["activity_score"]) == 10 | |
| 75 | + assert set(m["series"]["activity_score"][0]) == {"day", "value", "confidence"} | |
| 76 | + only = (await client.get(f"{V}/companies/{s}/metrics?metric=activity_score")).json() | |
| 77 | + assert list(only["series"]) == ["activity_score"] | |
| 78 | + | |
| 79 | + | |
| 80 | +async def test_jobs_people_products_pricing_locations_news(client, fixture_data): # type: ignore[no-untyped-def] | |
| 81 | + s = fixture_data["alpha_slug"] | |
| 82 | + r = await client.get(f"{V}/companies/{s}/jobs") | |
| 83 | + body = r.json() | |
| 84 | + assert r.status_code == 200 and body["total"] == 2 and body["meta"]["summary"]["open"] == 2 and body["meta"]["summary"]["ai_open"] == 1 | |
| 85 | + assert body["meta"]["summary"]["remote_ratio"] == 0.5 and body["meta"]["summary"]["by_country"] == [{"country": "ZZ", "n": 2}] | |
| 86 | + job = body["items"][0] | |
| 87 | + assert {"id", "title", "department", "location_text", "city", "country", "remote", "employment_type", "seniority", "url", "posted_at", "first_seen_at", "last_seen_at", | |
| 88 | + "removed_at", "status", "is_ai"} <= set(job) | |
| 89 | + assert (await client.get(f"{V}/companies/{s}/jobs?status=removed")).json()["items"][0]["status"] == "no_longer_listed" | |
| 90 | + assert (await client.get(f"{V}/companies/{s}/jobs?status=all")).json()["total"] == 3 | |
| 91 | + assert (await client.get(f"{V}/companies/{s}/jobs?ai=1")).json()["items"][0]["title"] == "Senior ML Engineer" | |
| 92 | + assert (await client.get(f"{V}/companies/{s}/jobs?q=account")).json()["total"] == 1 | |
| 93 | + p = (await client.get(f"{V}/companies/{s}/people")).json() | |
| 94 | + assert p["listed"][0]["name"] == "Jane Ztest" and p["no_longer_listed"][0]["name"] == "John Former" | |
| 95 | + pr = (await client.get(f"{V}/companies/{s}/products")).json() | |
| 96 | + assert pr["listed"][0]["name"] == "Ztest Widget" and pr["removed"] == [] | |
| 97 | + plans = (await client.get(f"{V}/companies/{s}/pricing")).json() | |
| 98 | + assert plans["current"][0]["price"] == 12.0 and plans["current"][0]["features"] == ["1 seat"] and plans["history"][0]["price"] == 10.0 | |
| 99 | + loc = (await client.get(f"{V}/companies/{s}/locations")).json() | |
| 100 | + assert loc["countries"] == ["ZZ"] and loc["items"][0]["kind"] == "headquarters" and loc["items"][0]["lat"] == 45.5 | |
| 101 | + news = (await client.get(f"{V}/companies/{s}/news")).json() | |
| 102 | + assert news["items"][0]["title"].startswith("Ztest Alpha announces") | |
| 103 | + | |
| 104 | + | |
| 105 | +async def test_sensors_history_similar(client, fixture_data): # type: ignore[no-untyped-def] | |
| 106 | + s = fixture_data["alpha_slug"] | |
| 107 | + sensors = (await client.get(f"{V}/companies/{s}/sensors")).json()["items"] | |
| 108 | + assert {x["surface"] for x in sensors} == {"pricing", "careers"} | |
| 109 | + assert {"id", "company_id", "surface", "connector_id", "url", "canonical_url", "domain", "status", "tier", "quality_score", "discovery_confidence", "discovery_method", | |
| 110 | + "current_interval_s", "next_run_at", "last_run_at", "last_success_at", "last_change_at", "last_status", "last_failure_class", "consecutive_failures", | |
| 111 | + "observation_count", "snapshot_count", "change_count", "meaningful_change_count", "event_count", "created_at"} <= set(sensors[0]) | |
| 112 | + hist = (await client.get(f"{V}/companies/{s}/history")).json()["sensors"] | |
| 113 | + pricing = next(x for x in hist if x["surface"] == "pricing") | |
| 114 | + assert [v["version_no"] for v in pricing["versions"]] == [3, 2, 1] | |
| 115 | + assert (await client.get(f"{V}/companies/{s}/history?versions=1")).json()["sensors"][0]["versions"].__len__() <= 1 | |
| 116 | + similar = (await client.get(f"{V}/companies/{s}/similar")).json()["items"] | |
| 117 | + assert [c["slug"] for c in similar] == [fixture_data["beta_slug"]] | |
| 118 | + | |
| 119 | + | |
| 120 | +async def test_compare(client, fixture_data): # type: ignore[no-untyped-def] | |
| 121 | + a, b = fixture_data["alpha_slug"], fixture_data["beta_slug"] | |
| 122 | + r = await client.get(f"{V}/companies/compare", params={"companies": f"{a},{b}"}) | |
| 123 | + body = r.json() | |
| 124 | + assert r.status_code == 200 | |
| 125 | + assert [c["slug"] for c in body["companies"]] == [a, b] | |
| 126 | + assert body["metrics"]["activity_score"] == {a: 72.3, b: 30.0} | |
| 127 | + assert len(body["series"][a]) == 10 and body["series"][b] == [] | |
| 128 | + assert body["events_30d"][a] == {"PRICING": 1, "HIRING": 1, "PRODUCT": 1} and body["events_30d"][b] == {"LEADERSHIP": 1} | |
| 129 | + assert body["jobs"][a] == {"open": 2, "ai_open": 1, "new_30d": 3} and body["jobs"][b] == {"open": 0, "ai_open": 0, "new_30d": 0} | |
| 130 | + assert body["locations"] == {a: 1, b: 0} | |
| 131 | + assert (await client.get(f"{V}/companies/compare", params={"companies": a})).status_code == 422 | |
| 132 | + assert (await client.get(f"{V}/companies/compare", params={"companies": f"{a},nope"})).status_code == 404 | |
| 133 | + | |
| 134 | + | |
| 135 | +async def test_events_feed_types_summary_detail(client, fixture_data): # type: ignore[no-untyped-def] | |
| 136 | + r = await client.get(f"{V}/events", params={"country": "ZZ"}) | |
| 137 | + body = r.json() | |
| 138 | + assert r.status_code == 200 and body["total"] == 4 | |
| 139 | + assert (await client.get(f"{V}/events", params={"country": "ZZ", "status": "retracted"})).json()["total"] == 1 | |
| 140 | + assert (await client.get(f"{V}/events", params={"industry": fixture_data["industry"], "event_type": "LEADERSHIP"})).json()["items"][0]["company"]["slug"] == fixture_data["beta_slug"] | |
| 141 | + assert (await client.get(f"{V}/events", params={"company": fixture_data["alpha_slug"], "min_confidence": 0.9})).json()["total"] == 1 | |
| 142 | + assert (await client.get(f"{V}/events", params={"q": "executive listed"})).json()["items"][0]["id"] == fixture_data["ev_leader"] | |
| 143 | + assert (await client.get(f"{V}/events", params={"country": "ZZ", "origin": "llm"})).json()["total"] == 0 | |
| 144 | + assert (await client.get(f"{V}/events", params={"company": "nope"})).status_code == 404 | |
| 145 | + types = (await client.get(f"{V}/events/types")).json()["types"] | |
| 146 | + pricing = next(t for t in types if t["event_type"] == "PRICING") | |
| 147 | + assert pricing["count_30d"] >= 1 and any(s["event_subtype"] == "PRICE_INCREASE" and s["count_30d"] >= 1 for s in pricing["subtypes"]) | |
| 148 | + summ = (await client.get(f"{V}/events/summary?days=7&group=country")).json()["items"] | |
| 149 | + zz = next(i for i in summ if i["key"] == "ZZ") | |
| 150 | + assert zz["count"] == 4 and "delta_pct" in zz | |
| 151 | + assert (await client.get(f"{V}/events/summary?group=nope")).status_code == 422 | |
| 152 | + r = await client.get(f"{V}/events/{fixture_data['ev_pricing']}") | |
| 153 | + ev = r.json() | |
| 154 | + assert r.status_code == 200 and ev["sources"][0]["source_url"].endswith("/pricing") and ev["change"]["id"] == fixture_data["change"] | |
| 155 | + assert ev["company"]["slug"] == fixture_data["alpha_slug"] and CARD_KEYS <= set(ev["company"]) | |
| 156 | + assert (await client.get(f"{V}/events/evt_nope")).status_code == 404 | |
added
tests/test_api_limits_exports.py
+102 −0
@@ -0,0 +1,102 @@ | ||
| 1 | +"""Rate limiting (token bucket, tiers, headers), API keys and streamed exports.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import csv | |
| 5 | +import hashlib | |
| 6 | +import io | |
| 7 | +import json | |
| 8 | + | |
| 9 | +import pytest | |
| 10 | +import test_api_support as support | |
| 11 | + | |
| 12 | +from companyatlas.api.ratelimit import RateLimiter, flush_usage, resolve_api_key | |
| 13 | +from companyatlas.db import fetch_one, transaction | |
| 14 | + | |
| 15 | +client = support.client | |
| 16 | +fixture_data = support.fixture_data | |
| 17 | +ADMIN = support.ADMIN | |
| 18 | +RAW_API_KEY = support.RAW_API_KEY | |
| 19 | + | |
| 20 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 21 | +V = "/api/v1" | |
| 22 | + | |
| 23 | + | |
| 24 | +async def test_token_bucket_semantics() -> None: | |
| 25 | + rl = RateLimiter(limits={"anonymous": 3, "paid": 6, "internal": None}) | |
| 26 | + t = 1000.0 | |
| 27 | + results = [rl.take("ip:1", "anonymous", now=t) for _ in range(4)] | |
| 28 | + assert [r[0] for r in results] == [True, True, True, False] | |
| 29 | + assert results[0][1] == 3 and results[0][2] == 2 and results[3][2] == 0 and results[3][3] > 0 | |
| 30 | + allowed, *_ = rl.take("ip:1", "anonymous", now=t + 20) # 3/min → one token back after 20 s | |
| 31 | + assert allowed is True | |
| 32 | + assert rl.take("ip:2", "anonymous", now=t)[0] is True # independent bucket | |
| 33 | + assert rl.take("key:x", "internal", now=t) == (True, 0, 0, 0.0) | |
| 34 | + assert rl.take("key:y", "unknown-tier", now=t)[1] == 3 # unknown tier falls back to anonymous | |
| 35 | + | |
| 36 | + | |
| 37 | +async def test_rate_limit_headers_and_tiers(client): # type: ignore[no-untyped-def] | |
| 38 | + from companyatlas.api import ratelimit | |
| 39 | + | |
| 40 | + r = await client.get(f"{V}/methodology") | |
| 41 | + assert r.headers["x-ratelimit-tier"] == "internal" # loopback without X-Forwarded-For = our own SSR | |
| 42 | + public = {"X-Forwarded-For": "203.0.113.7"} | |
| 43 | + r = await client.get(f"{V}/methodology", headers=public) | |
| 44 | + anon = ratelimit.limiter.limits["anonymous"] | |
| 45 | + assert r.headers["x-ratelimit-tier"] == "anonymous" and r.headers["x-ratelimit-limit"] == str(anon) and int(r.headers["x-ratelimit-remaining"]) < anon | |
| 46 | + assert ratelimit.TIER_LIMITS_PER_MIN == {"anonymous": 120, "authenticated": 600, "paid": 3000, "internal": None} | |
| 47 | + r = await client.get(f"{V}/methodology", headers={"X-CA-API-Key": RAW_API_KEY}) | |
| 48 | + assert r.headers["x-ratelimit-tier"] == "paid" and r.headers["x-ratelimit-limit"] == "3000" | |
| 49 | + r = await client.get(f"{V}/methodology", headers={"X-CA-API-Key": "ca_paid_not_a_real_key_000000", **public}) | |
| 50 | + assert r.headers["x-ratelimit-tier"] == "anonymous" | |
| 51 | + r = await client.get(f"{V}/methodology", headers=ADMIN) | |
| 52 | + assert r.headers["x-ratelimit-tier"] == "admin" and "x-ratelimit-remaining" not in r.headers | |
| 53 | + r = await client.get("/health") | |
| 54 | + assert "x-ratelimit-tier" not in r.headers # bypassed path | |
| 55 | + info = await resolve_api_key(RAW_API_KEY) | |
| 56 | + assert info and info["tier"] == "paid" | |
| 57 | + await flush_usage() | |
| 58 | + async with transaction() as conn: | |
| 59 | + row = await fetch_one(conn, "select request_count, last_used_at from api_keys where key_hash = :h", h=hashlib.sha256(RAW_API_KEY.encode()).hexdigest()) | |
| 60 | + assert row and row["request_count"] >= 1 and row["last_used_at"] is not None | |
| 61 | + | |
| 62 | + | |
| 63 | +async def test_429_from_exhausted_bucket(client, monkeypatch): # type: ignore[no-untyped-def] | |
| 64 | + from companyatlas.api import ratelimit | |
| 65 | + | |
| 66 | + tiny = RateLimiter(limits={"anonymous": 2, "authenticated": 600, "paid": 3000, "internal": None}) | |
| 67 | + monkeypatch.setattr(ratelimit, "limiter", tiny) | |
| 68 | + headers = {"X-Forwarded-For": "203.0.113.9, 10.0.0.1"} | |
| 69 | + codes = [(await client.get(f"{V}/methodology", headers=headers)).status_code for _ in range(3)] | |
| 70 | + assert codes == [200, 200, 429] | |
| 71 | + r = await client.get(f"{V}/methodology", headers=headers) | |
| 72 | + assert r.status_code == 429 and int(r.headers["retry-after"]) >= 1 and r.headers["x-ratelimit-remaining"] == "0" and r.json() == {"detail": "rate limit exceeded"} | |
| 73 | + assert (await client.get(f"{V}/methodology", headers={"X-Forwarded-For": "203.0.113.10"})).status_code == 200 | |
| 74 | + | |
| 75 | + | |
| 76 | +async def test_export_events_formats(client, fixture_data): # type: ignore[no-untyped-def] | |
| 77 | + r = await client.get(f"{V}/export/events.csv", params={"country": "ZZ"}) | |
| 78 | + assert r.status_code == 200 and r.headers["content-type"].startswith("text/csv") and "attachment" in r.headers["content-disposition"] | |
| 79 | + rows = list(csv.reader(io.StringIO(r.text))) | |
| 80 | + assert rows[0][:4] == ["id", "detected_at", "company_slug", "company_name"] and len(rows) == 5 | |
| 81 | + r = await client.get(f"{V}/export/events.ndjson", params={"country": "ZZ", "event_type": "PRICING"}) | |
| 82 | + lines = [json.loads(x) for x in r.text.splitlines() if x] | |
| 83 | + assert r.headers["content-type"].startswith("application/x-ndjson") and len(lines) == 1 and lines[0]["id"] == fixture_data["ev_pricing"] | |
| 84 | + r = await client.get(f"{V}/export/events.json", params={"company": fixture_data["alpha_slug"], "limit": 2}) | |
| 85 | + data = r.json() | |
| 86 | + assert isinstance(data, list) and len(data) == 2 and data[0]["company"]["slug"] == fixture_data["alpha_slug"] | |
| 87 | + assert (await client.get(f"{V}/export/events.xml")).status_code == 404 | |
| 88 | + assert (await client.get(f"{V}/export/events.json", params={"company": "nope"})).status_code == 404 | |
| 89 | + | |
| 90 | + | |
| 91 | +async def test_export_companies_and_jobs(client, fixture_data): # type: ignore[no-untyped-def] | |
| 92 | + r = await client.get(f"{V}/export/companies.csv", params={"country": "ZZ"}) | |
| 93 | + rows = list(csv.reader(io.StringIO(r.text))) | |
| 94 | + assert rows[0][0] == "id" and "activity_score" in rows[0] and len(rows) == 3 | |
| 95 | + r = await client.get(f"{V}/export/companies.json", params={"industry": fixture_data["industry"]}) | |
| 96 | + assert {c["slug"] for c in r.json()} == {fixture_data["alpha_slug"], fixture_data["beta_slug"]} | |
| 97 | + r = await client.get(f"{V}/export/jobs.ndjson", params={"company": fixture_data["alpha_slug"], "status": "all"}) | |
| 98 | + lines = [json.loads(x) for x in r.text.splitlines() if x] | |
| 99 | + assert len(lines) == 3 and all(j["company_slug"] == fixture_data["alpha_slug"] for j in lines) | |
| 100 | + r = await client.get(f"{V}/export/jobs.csv", params={"company": fixture_data["alpha_slug"], "ai": 1}) | |
| 101 | + rows = list(csv.reader(io.StringIO(r.text))) | |
| 102 | + assert len(rows) == 2 and rows[1][2] == "Senior ML Engineer" | |
added
tests/test_api_owner.py
+58 −0
@@ -0,0 +1,58 @@ | ||
| 1 | +"""Owner endpoints: watchlist and alerts flow with X-CA-Owner-Token (auto-created owner).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | +import test_api_support as support | |
| 6 | + | |
| 7 | +client = support.client | |
| 8 | +fixture_data = support.fixture_data | |
| 9 | +OWNER = support.OWNER | |
| 10 | + | |
| 11 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 12 | +V = "/api/v1" | |
| 13 | + | |
| 14 | + | |
| 15 | +async def test_owner_token_required(client): # type: ignore[no-untyped-def] | |
| 16 | + assert (await client.get(f"{V}/watchlist")).status_code == 401 | |
| 17 | + assert (await client.get(f"{V}/watchlist", headers={"X-CA-Owner-Token": "short"})).status_code == 401 | |
| 18 | + assert (await client.get(f"{V}/alerts")).status_code == 401 | |
| 19 | + | |
| 20 | + | |
| 21 | +async def test_watchlist_flow(client, fixture_data): # type: ignore[no-untyped-def] | |
| 22 | + r = await client.get(f"{V}/watchlist", headers=OWNER) | |
| 23 | + assert r.status_code == 200 and r.headers["cache-control"] == "no-store" | |
| 24 | + assert r.json()["items"] == [] and r.json()["events"] == [] | |
| 25 | + r = await client.post(f"{V}/watchlist", headers=OWNER, json={"company": fixture_data["alpha_slug"]}) | |
| 26 | + assert r.status_code == 201 and r.json()["added"] is True and r.json()["company"]["slug"] == fixture_data["alpha_slug"] | |
| 27 | + again = await client.post(f"{V}/watchlist", headers=OWNER, json={"company": fixture_data["alpha"]}) | |
| 28 | + assert again.status_code == 201 and again.json()["added"] is False | |
| 29 | + assert (await client.post(f"{V}/watchlist", headers=OWNER, json={"company": "nope"})).status_code == 404 | |
| 30 | + body = (await client.get(f"{V}/watchlist", headers=OWNER)).json() | |
| 31 | + assert [c["slug"] for c in body["items"]] == [fixture_data["alpha_slug"]] and body["items"][0]["added_at"] | |
| 32 | + assert {e["company"]["slug"] for e in body["events"]} == {fixture_data["alpha_slug"]} and len(body["events"]) == 3 | |
| 33 | + r = await client.delete(f"{V}/watchlist/{fixture_data['alpha_slug']}", headers=OWNER) | |
| 34 | + assert r.status_code == 200 and r.json()["removed"] is True | |
| 35 | + assert (await client.delete(f"{V}/watchlist/{fixture_data['alpha_slug']}", headers=OWNER)).json()["removed"] is False | |
| 36 | + assert (await client.get(f"{V}/watchlist", headers=OWNER)).json()["items"] == [] | |
| 37 | + | |
| 38 | + | |
| 39 | +async def test_alerts_flow(client, fixture_data): # type: ignore[no-untyped-def] | |
| 40 | + r = await client.post(f"{V}/alerts", headers=OWNER, json={"name": "Pricing moves", "company": fixture_data["alpha_slug"], | |
| 41 | + "condition": {"event_types": ["pricing"], "min_importance": 0.6}, "channel": "web"}) | |
| 42 | + assert r.status_code == 201 | |
| 43 | + alert = r.json() | |
| 44 | + assert alert["condition"] == {"event_types": ["PRICING"], "min_importance": 0.6} and alert["company"]["slug"] == fixture_data["alpha_slug"] | |
| 45 | + assert (await client.post(f"{V}/alerts", headers=OWNER, json={"name": "bad", "condition": {"event_types": ["NOPE"]}})).status_code == 422 | |
| 46 | + assert (await client.post(f"{V}/alerts", headers=OWNER, json={"name": "hook", "channel": "webhook", "condition": {"min_importance": 0.9}})).status_code == 422 | |
| 47 | + assert (await client.post(f"{V}/alerts", headers=OWNER, json={"name": "empty"})).status_code == 422 | |
| 48 | + ok = await client.post(f"{V}/alerts", headers=OWNER, json={"name": "hook", "channel": "webhook", "target": "https://example.org/hook", | |
| 49 | + "condition": {"metrics": {"activity_score": {"gt": 80}}}}) | |
| 50 | + assert ok.status_code == 201 and ok.json()["condition"]["metrics"]["activity_score"] == {"gt": 80.0} | |
| 51 | + items = (await client.get(f"{V}/alerts", headers=OWNER)).json()["items"] | |
| 52 | + assert {a["name"] for a in items} == {"Pricing moves", "hook"} | |
| 53 | + deliveries = await client.get(f"{V}/alerts/deliveries", headers=OWNER) | |
| 54 | + assert deliveries.status_code == 200 and deliveries.json()["items"] == [] | |
| 55 | + for a in items: | |
| 56 | + assert (await client.delete(f"{V}/alerts/{a['id']}", headers=OWNER)).json()["removed"] is True | |
| 57 | + assert (await client.delete(f"{V}/alerts/{items[0]['id']}", headers=OWNER)).status_code == 404 | |
| 58 | + assert (await client.get(f"{V}/alerts", headers=OWNER)).json()["items"] == [] | |
added
tests/test_api_provenance.py
+72 −0
@@ -0,0 +1,72 @@ | ||
| 1 | +"""Provenance chain: sensors → snapshots (text/blocks from the archive) → diffs → changes → events.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | +import test_api_support as support | |
| 6 | + | |
| 7 | +client = support.client | |
| 8 | +fixture_data = support.fixture_data | |
| 9 | + | |
| 10 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 11 | +V = "/api/v1" | |
| 12 | + | |
| 13 | + | |
| 14 | +async def test_sensor_detail_snapshots_changes(client, fixture_data): # type: ignore[no-untyped-def] | |
| 15 | + sid = fixture_data["sensor_home"] | |
| 16 | + r = await client.get(f"{V}/sensors/{sid}") | |
| 17 | + body = r.json() | |
| 18 | + assert r.status_code == 200 and body["company"]["slug"] == fixture_data["alpha_slug"] and body["latest_snapshot"]["id"] == fixture_data["snap3"] | |
| 19 | + assert body["tier"] == "C" and body["surface"] == "pricing" | |
| 20 | + snaps = (await client.get(f"{V}/sensors/{sid}/snapshots")).json()["items"] | |
| 21 | + assert [s["version_no"] for s in snaps] == [3, 2, 1] | |
| 22 | + assert {"id", "sensor_id", "version_no", "fetched_at", "title", "language", "text_length", "block_count", "extracted_summary", "content_hash", "previous_snapshot_id"} <= set(snaps[0]) | |
| 23 | + assert snaps[1]["previous_snapshot_id"] == fixture_data["snap1"] | |
| 24 | + changes = (await client.get(f"{V}/sensors/{sid}/changes")).json()["items"] | |
| 25 | + assert changes[0]["id"] == fixture_data["change"] and changes[0]["kind"] == "major" and "diff" not in changes[0] | |
| 26 | + assert (await client.get(f"{V}/sensors/{sid}/changes?min_significance=0.9")).json()["items"] == [] | |
| 27 | + assert (await client.get(f"{V}/sensors/sen_nope")).status_code == 404 | |
| 28 | + | |
| 29 | + | |
| 30 | +async def test_snapshot_detail_reads_archive(client, fixture_data): # type: ignore[no-untyped-def] | |
| 31 | + r = await client.get(f"{V}/snapshots/{fixture_data['snap2']}") | |
| 32 | + body = r.json() | |
| 33 | + assert r.status_code == 200 | |
| 34 | + assert "Starter plan $12 per month" in body["text"] and body["text_truncated"] is False | |
| 35 | + assert [b["key"] for b in body["blocks"]] == ["b1", "b2", "b3", "b4"] and body["blocks"][1]["kind"] == "pricing_plan" | |
| 36 | + assert body["extracted"] == {"plans": [{"plan_name": "Starter"}]} and body["extracted_summary"] == {"plan_count": 3} | |
| 37 | + assert body["sensor"]["surface"] == "pricing" | |
| 38 | + lite = (await client.get(f"{V}/snapshots/{fixture_data['snap2']}?include=extracted")).json() | |
| 39 | + assert "text" not in lite and "blocks" not in lite | |
| 40 | + assert (await client.get(f"{V}/snapshots/snap_nope")).status_code == 404 | |
| 41 | + | |
| 42 | + | |
| 43 | +async def test_snapshot_diff_stored_and_computed(client, fixture_data): # type: ignore[no-untyped-def] | |
| 44 | + r = await client.get(f"{V}/snapshots/{fixture_data['snap1']}/diff/{fixture_data['snap2']}") | |
| 45 | + body = r.json() | |
| 46 | + assert r.status_code == 200 and body["source"] == "change" and body["change_id"] == fixture_data["change"] | |
| 47 | + assert body["before"]["id"] == fixture_data["snap1"] and body["after"]["id"] == fixture_data["snap2"] | |
| 48 | + assert body["diff"]["counts"] == {"added": 1, "removed": 0, "modified": 1, "moved": 0} and body["diff"]["significance"] == 0.72 | |
| 49 | + reverse = (await client.get(f"{V}/snapshots/{fixture_data['snap2']}/diff/{fixture_data['snap1']}")).json() | |
| 50 | + assert reverse["before"]["id"] == fixture_data["snap1"] | |
| 51 | + r = await client.get(f"{V}/snapshots/{fixture_data['snap2']}/diff/{fixture_data['snap3']}") | |
| 52 | + assert r.status_code in (200, 501) | |
| 53 | + if r.status_code == 200: | |
| 54 | + d = r.json() | |
| 55 | + assert d["source"] == "computed" | |
| 56 | + for key in ("added", "removed", "modified", "moved", "counts", "text_delta_ratio", "similarity", "reasons"): | |
| 57 | + assert key in d["diff"], key | |
| 58 | + assert any(x["key"] == "b5" for x in d["diff"]["added"]) | |
| 59 | + assert (await client.get(f"{V}/snapshots/{fixture_data['snap1']}/diff/snap_nope")).status_code == 404 | |
| 60 | + | |
| 61 | + | |
| 62 | +async def test_change_detail(client, fixture_data): # type: ignore[no-untyped-def] | |
| 63 | + r = await client.get(f"{V}/changes/{fixture_data['change']}") | |
| 64 | + body = r.json() | |
| 65 | + assert r.status_code == 200 | |
| 66 | + for key in ("id", "sensor_id", "surface", "company_id", "detected_at", "significance", "kind", "blocks_added", "blocks_removed", "blocks_modified", "text_delta_ratio", | |
| 67 | + "similarity", "snapshot_before", "snapshot_after", "diff", "structured_delta", "events"): | |
| 68 | + assert key in body, key | |
| 69 | + assert body["structured_delta"]["plans"]["price_changed"][0]["pct"] == 20 | |
| 70 | + assert [e["id"] for e in body["events"]] == [fixture_data["ev_pricing"]] | |
| 71 | + assert body["company"]["slug"] == fixture_data["alpha_slug"] | |
| 72 | + assert (await client.get(f"{V}/changes/chg_nope")).status_code == 404 | |
added
tests/test_api_public.py
+198 −0
@@ -0,0 +1,198 @@ | ||
| 1 | +"""Public aggregates, live feed, search, rankings, industries/countries, signals, sitemap — shapes, caching and graceful degradation.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | +import test_api_support as support | |
| 6 | + | |
| 7 | +client = support.client | |
| 8 | +fixture_data = support.fixture_data | |
| 9 | + | |
| 10 | +pytestmark = pytest.mark.asyncio(loop_scope="session") | |
| 11 | +V = "/api/v1" | |
| 12 | + | |
| 13 | + | |
| 14 | +async def test_health(client): # type: ignore[no-untyped-def] | |
| 15 | + r = await client.get("/health") | |
| 16 | + assert r.status_code == 200 and r.json()["db"] is True | |
| 17 | + assert (await client.get(f"{V}/health")).status_code == 200 | |
| 18 | + | |
| 19 | + | |
| 20 | +async def test_stats_shape_and_cache_headers(client): # type: ignore[no-untyped-def] | |
| 21 | + r = await client.get(f"{V}/stats") | |
| 22 | + assert r.status_code == 200 | |
| 23 | + body = r.json() | |
| 24 | + for key in ("companies", "companies_active", "sensors", "sensors_active", "observations", "snapshots", "changes", "meaningful_changes", "events", "jobs_open", | |
| 25 | + "countries", "industries", "observations_today", "changes_today", "events_today", "dataset_started_at", "dataset_age_days", "oldest_history_days", | |
| 26 | + "last_observation_at", "archive"): | |
| 27 | + assert key in body, key | |
| 28 | + assert body["companies"] >= 2 and body["events"] >= 4 | |
| 29 | + assert set(body["archive"]) == {"objects", "bytes"} | |
| 30 | + assert r.headers["cache-control"].startswith("public, max-age=60") | |
| 31 | + assert r.headers["etag"].startswith('W/"') | |
| 32 | + r2 = await client.get(f"{V}/stats", headers={"If-None-Match": r.headers["etag"]}) | |
| 33 | + assert r2.status_code == 304 | |
| 34 | + | |
| 35 | + | |
| 36 | +async def test_stats_history_and_index(client): # type: ignore[no-untyped-def] | |
| 37 | + r = await client.get(f"{V}/stats/history?days=30") | |
| 38 | + assert r.status_code == 200 and isinstance(r.json()["items"], list) | |
| 39 | + r = await client.get(f"{V}/index") | |
| 40 | + body = r.json() | |
| 41 | + assert r.status_code == 200 | |
| 42 | + for key in ("value", "baseline", "delta_7d", "delta_30d", "series", "by_type", "by_country", "by_industry", "formula_version"): | |
| 43 | + assert key in body | |
| 44 | + assert body["baseline"] == 100 | |
| 45 | + | |
| 46 | + | |
| 47 | +async def test_system(client): # type: ignore[no-untyped-def] | |
| 48 | + r = await client.get(f"{V}/system") | |
| 49 | + assert r.status_code == 200 | |
| 50 | + body = r.json() | |
| 51 | + for key in ("sensors_online", "sensors_failing", "observations_today", "events_today", "countries_covered", "queue_lag_s", "scheduler_last_tick_at", "fetch_per_min", | |
| 52 | + "success_rate_24h"): | |
| 53 | + assert key in body | |
| 54 | + assert body["sensors_online"] >= 3 | |
| 55 | + | |
| 56 | + | |
| 57 | +async def test_pulse_full_shape(client, fixture_data): # type: ignore[no-untyped-def] | |
| 58 | + r = await client.get(f"{V}/pulse") | |
| 59 | + assert r.status_code == 200 | |
| 60 | + body = r.json() | |
| 61 | + for key in ("stats", "live", "movers", "hiring", "launches", "pricing", "ai", "industries", "countries", "trending", "activity_index", "map"): | |
| 62 | + assert key in body, key | |
| 63 | + assert any(e["id"] == fixture_data["ev_pricing"] for e in body["pricing"]) | |
| 64 | + assert any(c["slug"] == fixture_data["alpha_slug"] for c in body["movers"]) | |
| 65 | + mover = next(c for c in body["movers"] if c["slug"] == fixture_data["alpha_slug"]) | |
| 66 | + assert mover["rank"] >= 1 and mover["value"] == 72.3 and "sparkline" in mover | |
| 67 | + assert set(body["activity_index"]) == {"value", "delta_7d", "series"} | |
| 68 | + assert any(b["country"] == "ZZ" for b in body["map"]) | |
| 69 | + | |
| 70 | + | |
| 71 | +async def test_methodology(client): # type: ignore[no-untyped-def] | |
| 72 | + body = (await client.get(f"{V}/methodology")).json() | |
| 73 | + assert {m["metric"] for m in body["metrics"]} >= {"activity_score", "hiring_momentum_30d", "ai_adoption", "corporate_change_index"} | |
| 74 | + assert "PRICING" in body["event_types"] and "significance_bands" in body and body["confidence_labels"][0]["label"] == "VERIFIED" | |
| 75 | + | |
| 76 | + | |
| 77 | +async def test_live_and_since(client, fixture_data): # type: ignore[no-untyped-def] | |
| 78 | + r = await client.get(f"{V}/live?limit=10&country=ZZ") | |
| 79 | + assert r.status_code == 200 and r.headers["cache-control"] == "no-store" | |
| 80 | + body = r.json() | |
| 81 | + ids = [e["id"] for e in body["items"]] | |
| 82 | + assert fixture_data["ev_pricing"] in ids and fixture_data["ev_retracted"] not in ids | |
| 83 | + ev = body["items"][0] | |
| 84 | + for key in ("id", "company", "event_type", "event_subtype", "importance", "confidence", "confidence_label", "title", "payload", "entities", "tags", "detected_at", | |
| 85 | + "source_url", "origin", "status"): | |
| 86 | + assert key in ev | |
| 87 | + assert set(ev["company"]) == {"id", "slug", "display_name", "canonical_domain", "country", "logo_url"} | |
| 88 | + assert ev["detected_at"].endswith("Z") | |
| 89 | + r = await client.get(f"{V}/live", params={"since": body["cursor"], "country": "ZZ"}) | |
| 90 | + assert r.status_code == 200 and r.json()["items"] == [] | |
| 91 | + r = await client.get(f"{V}/live", params={"event_type": "PRICING", "min_importance": 0.7, "country": "ZZ"}) | |
| 92 | + assert all(e["event_type"] == "PRICING" for e in r.json()["items"]) and r.json()["items"] | |
| 93 | + assert (await client.get(f"{V}/live?since=not-a-date")).status_code == 422 | |
| 94 | + | |
| 95 | + | |
| 96 | +async def test_live_stream_first_bytes(client): # type: ignore[no-untyped-def] | |
| 97 | + r = await client.get(f"{V}/live/stream?max_s=1") | |
| 98 | + assert r.status_code == 200 | |
| 99 | + assert r.headers["content-type"].startswith("text/event-stream") | |
| 100 | + assert r.headers.get("x-accel-buffering") == "no" and r.headers["cache-control"] == "no-store" | |
| 101 | + assert "event: heartbeat" in r.text and '"cursor"' in r.text and "event: end" in r.text | |
| 102 | + | |
| 103 | + | |
| 104 | +async def test_rankings(client, fixture_data): # type: ignore[no-untyped-def] | |
| 105 | + r = await client.get(f"{V}/rankings?kind=most_active&window=7d&country=ZZ") | |
| 106 | + assert r.status_code == 200 and r.headers["cache-control"].startswith("public, max-age=120") | |
| 107 | + body = r.json() | |
| 108 | + assert body["kind"] == "most_active" and body["window"] == "7d" | |
| 109 | + slugs = [i["slug"] for i in body["items"]] | |
| 110 | + assert slugs[:2] == [fixture_data["alpha_slug"], fixture_data["beta_slug"]] | |
| 111 | + top = body["items"][0] | |
| 112 | + assert top["rank"] == 1 and top["value"] == 72.3 and "delta" in top | |
| 113 | + r = await client.get(f"{V}/rankings?kind=hiring_decline&window=30d&country=ZZ") | |
| 114 | + assert [i["slug"] for i in r.json()["items"]] == [fixture_data["beta_slug"]] | |
| 115 | + r = await client.get(f"{V}/rankings?kind=pricing_changes&window=7d") | |
| 116 | + assert any(i["slug"] == fixture_data["alpha_slug"] and i["value"] == 1 for i in r.json()["items"]) | |
| 117 | + assert (await client.get(f"{V}/rankings?kind=bogus")).status_code == 422 | |
| 118 | + | |
| 119 | + | |
| 120 | +async def test_industries(client, fixture_data): # type: ignore[no-untyped-def] | |
| 121 | + r = await client.get(f"{V}/industries") | |
| 122 | + assert r.status_code == 200 | |
| 123 | + row = next(i for i in r.json()["items"] if i["slug"] == fixture_data["industry"]) | |
| 124 | + for key in ("slug", "name", "parent_slug", "companies", "events_7d", "events_30d", "hiring_momentum_30d", "activity_score", "ai_adoption", "top_event_types"): | |
| 125 | + assert key in row | |
| 126 | + assert row["companies"] == 2 and row["events_30d"] == 4 and row["activity_score"] == 51.2 and "PRICING" in row["top_event_types"] | |
| 127 | + r = await client.get(f"{V}/industries/{fixture_data['industry']}") | |
| 128 | + body = r.json() | |
| 129 | + assert r.status_code == 200 | |
| 130 | + for key in ("description", "companies", "events", "hiring", "series", "countries", "trending"): | |
| 131 | + assert key in body | |
| 132 | + assert body["hiring"]["open"] == 2 and body["countries"][0] == {"country": "ZZ", "companies": 2} | |
| 133 | + assert len(body["series"]) == 10 and body["companies"][0]["slug"] == fixture_data["alpha_slug"] | |
| 134 | + assert (await client.get(f"{V}/industries/does-not-exist")).status_code == 404 | |
| 135 | + | |
| 136 | + | |
| 137 | +async def test_countries(client, fixture_data): # type: ignore[no-untyped-def] | |
| 138 | + r = await client.get(f"{V}/countries") | |
| 139 | + row = next(c for c in r.json()["items"] if c["code"] == "ZZ") | |
| 140 | + for key in ("code", "name", "region", "companies", "events_7d", "events_30d", "hiring_momentum_30d", "activity_score", "industry_mix", "lat", "lon"): | |
| 141 | + assert key in row | |
| 142 | + assert row["companies"] == 2 and row["industry_mix"][0]["industry"] == fixture_data["industry"] and row["slug"] == "ztestland" | |
| 143 | + for key in ("ZZ", "zz", "ztestland"): | |
| 144 | + r = await client.get(f"{V}/countries/{key}") | |
| 145 | + assert r.status_code == 200, key | |
| 146 | + body = r.json() | |
| 147 | + for key in ("companies", "events", "movers", "new_entrants", "series", "industries"): | |
| 148 | + assert key in body | |
| 149 | + assert body["movers"][0]["slug"] == fixture_data["alpha_slug"] and len(body["new_entrants"]) == 2 | |
| 150 | + assert (await client.get(f"{V}/countries/atlantis")).status_code == 404 | |
| 151 | + | |
| 152 | + | |
| 153 | +async def test_signals_trends_map(client, fixture_data): # type: ignore[no-untyped-def] | |
| 154 | + r = await client.get(f"{V}/signals?scope=company&company={fixture_data['alpha_slug']}") | |
| 155 | + items = r.json()["items"] | |
| 156 | + assert r.status_code == 200 and items[0]["kind"] == "hiring_surge" and items[0]["company"]["slug"] == fixture_data["alpha_slug"] | |
| 157 | + r = await client.get(f"{V}/trends?window=30d") | |
| 158 | + assert r.status_code == 200 and isinstance(r.json()["items"], list) | |
| 159 | + r = await client.get(f"{V}/map?metric=companies") | |
| 160 | + assert r.status_code == 200 | |
| 161 | + buckets = r.json()["buckets"] | |
| 162 | + zz = [b for b in buckets if b["country"] == "ZZ"] | |
| 163 | + assert zz and {"lat", "lon", "country", "city", "companies", "events_30d", "jobs_open", "top"} <= set(zz[0]) | |
| 164 | + assert any(b["city"] == "Testville" and b["jobs_open"] == 2 for b in zz) | |
| 165 | + | |
| 166 | + | |
| 167 | +async def test_search_suggest_ask(client, fixture_data): # type: ignore[no-untyped-def] | |
| 168 | + r = await client.get(f"{V}/search", params={"q": f"ztest alpha {fixture_data['suffix']}"}) | |
| 169 | + body = r.json() | |
| 170 | + assert r.status_code == 200 and body["companies"][0]["slug"] == fixture_data["alpha_slug"] and "took_ms" in body | |
| 171 | + assert any(e["id"] == fixture_data["ev_pricing"] for e in (await client.get(f"{V}/search", params={"q": "starter price increased", "types": "events"})).json()["events"]) | |
| 172 | + people = (await client.get(f"{V}/search", params={"q": "jane ztest", "types": "people"})).json()["people"] | |
| 173 | + assert people and people[0]["company"]["slug"] == fixture_data["alpha_slug"] | |
| 174 | + r = await client.get(f"{V}/search", params={"q": f"alphaz {fixture_data['suffix']}"}) # alias | |
| 175 | + assert r.json()["companies"][0]["slug"] == fixture_data["alpha_slug"] | |
| 176 | + assert (await client.get(f"{V}/search", params={"q": "zt"})).status_code == 200 | |
| 177 | + r = await client.get(f"{V}/search/suggest", params={"q": "ztest al"}) | |
| 178 | + items = r.json()["items"] | |
| 179 | + assert r.status_code == 200 and len(items) <= 10 and items[0]["kind"] == "company" and items[0]["href"].startswith("/company/") | |
| 180 | + assert any(i["kind"] == "event_type" for i in (await client.get(f"{V}/search/suggest", params={"q": "pric"})).json()["items"]) | |
| 181 | + r = await client.get(f"{V}/ask", params={"q": "pricing changes in Ztestland this week"}) | |
| 182 | + body = r.json() | |
| 183 | + assert r.status_code == 200 | |
| 184 | + for key in ("interpretation", "answer", "companies", "events", "sources"): | |
| 185 | + assert key in body | |
| 186 | + assert any(e["id"] == fixture_data["ev_pricing"] for e in body["events"]) and body["sources"] | |
| 187 | + assert "fired" not in body["answer"].lower() | |
| 188 | + | |
| 189 | + | |
| 190 | +async def test_sitemap(client, fixture_data): # type: ignore[no-untyped-def] | |
| 191 | + r = await client.get(f"{V}/sitemap?kind=companies") | |
| 192 | + body = r.json() | |
| 193 | + assert r.status_code == 200 and body["pages"] >= 1 | |
| 194 | + slugs = {i["slug"] for i in body["items"]} | |
| 195 | + assert fixture_data["alpha_slug"] in slugs # indexed = true | |
| 196 | + assert fixture_data["beta_slug"] not in slugs # 1 sensor < seo_min_sensors | |
| 197 | + assert fixture_data["industry"] in {i["slug"] for i in (await client.get(f"{V}/sitemap?kind=industries")).json()["items"]} | |
| 198 | + assert "ztestland" in {i["slug"] for i in (await client.get(f"{V}/sitemap?kind=countries")).json()["items"]} | |
added
tests/test_api_support.py
+220 −0
@@ -0,0 +1,220 @@ | ||
| 1 | +"""Shared fixtures for the API tests: a `ztest-api-*` company family with sensors, snapshots (real objects in the archive), a change, | |
| 2 | +events, jobs, entities, metrics, a signal, queue/failure/review rows and a paid API key. Everything is removed in the finaliser. | |
| 3 | + | |
| 4 | +Import into a test module with `from test_api_support import client, fixture_data # noqa: F401` and mark the module with | |
| 5 | +`pytestmark = pytest.mark.asyncio(loop_scope="session")` so the shared asyncpg engine stays on one loop. | |
| 6 | +""" | |
| 7 | +from __future__ import annotations | |
| 8 | + | |
| 9 | +import hashlib | |
| 10 | +import json | |
| 11 | +import os | |
| 12 | +import secrets | |
| 13 | +from datetime import UTC, datetime, timedelta | |
| 14 | +from typing import Any | |
| 15 | + | |
| 16 | +import httpx | |
| 17 | +import pytest_asyncio | |
| 18 | + | |
| 19 | +from companyatlas import archive | |
| 20 | +from companyatlas.api.common import cache | |
| 21 | +from companyatlas.db import dispose, execute, fetch_val, jsonb, transaction | |
| 22 | + | |
| 23 | +ADMIN = {"X-CA-Admin-Token": "dev-admin-token"} | |
| 24 | +OWNER_TOKEN = "ztest-owner-token-" + secrets.token_hex(12) | |
| 25 | +OWNER = {"X-CA-Owner-Token": OWNER_TOKEN} | |
| 26 | +RAW_API_KEY = "ca_paid_ztest_" + secrets.token_urlsafe(24) | |
| 27 | +SUFFIX = secrets.token_hex(3) | |
| 28 | + | |
| 29 | + | |
| 30 | +def _now() -> datetime: | |
| 31 | + return datetime.now(UTC) | |
| 32 | + | |
| 33 | + | |
| 34 | +async def _insert_fixture() -> dict[str, Any]: | |
| 35 | + now = _now() | |
| 36 | + d: dict[str, Any] = {"suffix": SUFFIX, "objects": []} | |
| 37 | + ids = {"alpha": f"co_ztestapi{SUFFIX}a", "beta": f"co_ztestapi{SUFFIX}b", "sensor_home": f"sen_ztestapi{SUFFIX}h", "sensor_careers": f"sen_ztestapi{SUFFIX}c", | |
| 38 | + "sensor_beta": f"sen_ztestapi{SUFFIX}x", "snap1": f"snap_ztestapi{SUFFIX}1", "snap2": f"snap_ztestapi{SUFFIX}2", "snap3": f"snap_ztestapi{SUFFIX}3", | |
| 39 | + "change": f"chg_ztestapi{SUFFIX}1", "ev_hiring": f"evt_ztestapi{SUFFIX}h", "ev_pricing": f"evt_ztestapi{SUFFIX}p", "ev_product": f"evt_ztestapi{SUFFIX}n", | |
| 40 | + "ev_leader": f"evt_ztestapi{SUFFIX}l", "ev_retracted": f"evt_ztestapi{SUFFIX}r", "queue_dead": f"qj_ztestapi{SUFFIX}d", "failure": f"fail_ztestapi{SUFFIX}1", | |
| 41 | + "review": f"rev_ztestapi{SUFFIX}1", "signal": f"sig_ztestapi{SUFFIX}1", "api_key": f"key_ztestapi{SUFFIX}1", "connector": f"ztest-generic-{SUFFIX}"} | |
| 42 | + d.update(ids) | |
| 43 | + d["alpha_slug"], d["beta_slug"] = f"ztest-api-alpha-{SUFFIX}", f"ztest-api-beta-{SUFFIX}" | |
| 44 | + d["country"], d["industry"] = "ZZ", f"ztest-industry-{SUFFIX}" | |
| 45 | + | |
| 46 | + text1 = "Ztest Alpha — pricing\nStarter plan $10 per month\nPro plan $20 per month\nContact sales for Enterprise" | |
| 47 | + text2 = "Ztest Alpha — pricing\nStarter plan $12 per month\nPro plan $20 per month\nContact sales for Enterprise\nNew: Team plan" | |
| 48 | + text3 = text2 + "\nFooter updated" | |
| 49 | + blocks1 = [{"key": "b1", "kind": "heading", "text": "Ztest Alpha — pricing", "path": "", "hash": "h1", "simhash": 0, "weight": 1.0, "order": 0, "attrs": {}}, | |
| 50 | + {"key": "b2", "kind": "pricing_plan", "text": "Starter plan $10 per month", "path": "Pricing > Starter", "hash": "h2", "simhash": 0, "weight": 1.5, "order": 1, "attrs": {}}, | |
| 51 | + {"key": "b3", "kind": "pricing_plan", "text": "Pro plan $20 per month", "path": "Pricing > Pro", "hash": "h3", "simhash": 0, "weight": 1.5, "order": 2, "attrs": {}}] | |
| 52 | + blocks2 = [dict(blocks1[0]), {**blocks1[1], "text": "Starter plan $12 per month", "hash": "h2b"}, dict(blocks1[2]), | |
| 53 | + {"key": "b4", "kind": "pricing_plan", "text": "New: Team plan", "path": "Pricing > Team", "hash": "h4", "simhash": 0, "weight": 1.5, "order": 3, "attrs": {}}] | |
| 54 | + blocks3 = blocks2 + [{"key": "b5", "kind": "footer", "text": "Footer updated", "path": "", "hash": "h5", "simhash": 0, "weight": 0.2, "order": 4, "attrs": {}}] | |
| 55 | + keys = {} | |
| 56 | + for name, payload in (("t1", text1), ("t2", text2), ("t3", text3), ("b1", json.dumps(blocks1)), ("b2", json.dumps(blocks2)), ("b3", json.dumps(blocks3))): | |
| 57 | + key, _size, _created = archive.put_text(payload) | |
| 58 | + keys[name] = key | |
| 59 | + d["objects"].append(key) | |
| 60 | + | |
| 61 | + async with transaction() as conn: | |
| 62 | + await execute(conn, "insert into countries (code, name, region, subregion, lat, lon) values ('ZZ', 'Ztestland', 'Test Region', 'Test Sub', 45.5, -73.6) " | |
| 63 | + "on conflict (code) do nothing") | |
| 64 | + await execute(conn, "insert into industries (slug, name, description, keywords, sort_order) values (:s, 'Ztest Industry', 'Synthetic test industry', " | |
| 65 | + "array['ztest'], 999) on conflict (slug) do nothing", s=d["industry"]) | |
| 66 | + await execute(conn, "insert into connectors (id, name, version, category) values (:id, 'Ztest generic', '1', 'homepage') on conflict (id) do nothing", id=ids["connector"]) | |
| 67 | + for key, slug, name, imp, indexed in ((ids["alpha"], d["alpha_slug"], f"Ztest Alpha {SUFFIX}", 0.9, True), (ids["beta"], d["beta_slug"], f"Ztest Beta {SUFFIX}", 0.5, False)): | |
| 68 | + await execute(conn, "insert into companies (id, slug, display_name, legal_name, canonical_domain, website, description, industries, industry_primary, country, " | |
| 69 | + "hq_city, public_company, importance, tier, indexed, status, onboarding_status, first_observed_at, last_observed_at, last_event_at) values " | |
| 70 | + "(:id, :slug, :name, :legal, :domain, :website, 'Synthetic company used by the API tests', cast(:inds as text[]), :ip, 'ZZ', 'Testville', " | |
| 71 | + "false, :imp, 2, :indexed, 'ACTIVE', 'active', :t0, :t1, :t1)", | |
| 72 | + id=key, slug=slug, name=name, legal=name + " Inc.", domain=f"{slug}.example", website=f"https://{slug}.example", inds=[d["industry"]], | |
| 73 | + ip=d["industry"], imp=imp, indexed=indexed, t0=now - timedelta(days=40), t1=now - timedelta(minutes=5)) | |
| 74 | + await execute(conn, "insert into company_aliases (company_id, alias, alias_norm, kind) values (:c, :a, :n, 'brand')", c=ids["alpha"], a=f"Alphaz {SUFFIX}", n=f"alphaz{SUFFIX}") | |
| 75 | + await execute(conn, "insert into domains (id, company_id, domain, kind) values (:id, :c, :d, 'primary')", id=f"dom_ztestapi{SUFFIX}1", c=ids["alpha"], d=f"{d['alpha_slug']}.example") | |
| 76 | + await execute(conn, "insert into company_relationships (id, from_company_id, to_company_id, kind, confidence) values (:id, :a, :b, 'PARTNER_OF', 0.8)", | |
| 77 | + id=f"rel_ztestapi{SUFFIX}1", a=ids["alpha"], b=ids["beta"]) | |
| 78 | + for sid, cid, surface, url, snaps, changes, events in ((ids["sensor_home"], ids["alpha"], "pricing", f"https://{d['alpha_slug']}.example/pricing", 3, 1, 2), | |
| 79 | + (ids["sensor_careers"], ids["alpha"], "careers", f"https://{d['alpha_slug']}.example/careers", 0, 0, 1), | |
| 80 | + (ids["sensor_beta"], ids["beta"], "homepage", f"https://{d['beta_slug']}.example/", 0, 0, 1)): | |
| 81 | + await execute(conn, "insert into sensors (id, company_id, surface, connector_id, url, canonical_url, domain, status, tier, quality_score, base_interval_s, " | |
| 82 | + "current_interval_s, next_run_at, last_run_at, last_success_at, last_change_at, last_status, observation_count, snapshot_count, change_count, " | |
| 83 | + "meaningful_change_count, event_count, last_snapshot_id) values (:id, :cid, :surface, :conn, :url, :url, :domain, 'active', 'C', 80, 21600, 21600, " | |
| 84 | + ":next, :last, :last, :last, 200, :obs, :snaps, :changes, :changes, :events, :lastsnap)", | |
| 85 | + id=sid, cid=cid, surface=surface, conn=ids["connector"], url=url, domain=url.split("/")[2], next=now + timedelta(hours=1), | |
| 86 | + last=now - timedelta(minutes=10), obs=snaps + 4, snaps=snaps, changes=changes, events=events, lastsnap=ids["snap3"] if snaps else None) | |
| 87 | + for sid_key, ver, tkey, bkey, prev, when in (("snap1", 1, "t1", "b1", None, now - timedelta(days=2)), ("snap2", 2, "t2", "b2", ids["snap1"], now - timedelta(days=1)), | |
| 88 | + ("snap3", 3, "t3", "b3", ids["snap2"], now - timedelta(hours=2))): | |
| 89 | + await execute(conn, "insert into snapshots (id, sensor_id, company_id, previous_snapshot_id, version_no, fetched_at, content_hash, normalized_hash, structural_hash, " | |
| 90 | + "text_key, blocks_key, extracted, extracted_summary, title, language, text_length, block_count) values (:id, :sid, :cid, :prev, :ver, :when, :h, :h, :h, " | |
| 91 | + ":tk, :bk, cast(:ex as jsonb), cast(:sum as jsonb), 'Ztest Alpha — pricing', 'en', :tl, :bc)", | |
| 92 | + id=ids[sid_key], sid=ids["sensor_home"], cid=ids["alpha"], prev=prev, ver=ver, when=when, h=keys[tkey], tk=keys[tkey], bk=keys[bkey], | |
| 93 | + ex=jsonb({"plans": [{"plan_name": "Starter"}]}), sum=jsonb({"plan_count": ver + 1}), tl=len(text1), bc=3) | |
| 94 | + diff = {"added": [{"key": "b4", "kind": "pricing_plan", "path": "Pricing > Team", "before": None, "after": "New: Team plan", "weight": 1.5, "similarity": None}], | |
| 95 | + "removed": [], "modified": [{"key": "b2", "kind": "pricing_plan", "path": "Pricing > Starter", "before": "Starter plan $10 per month", | |
| 96 | + "after": "Starter plan $12 per month", "weight": 1.5, "similarity": 0.9}], "moved": [], | |
| 97 | + "counts": {"added": 1, "removed": 0, "modified": 1, "moved": 0}, "text_delta_ratio": 0.18, "similarity": 0.82, "reasons": ["pricing_plan modified"]} | |
| 98 | + await execute(conn, "insert into changes (id, sensor_id, company_id, surface, snapshot_before, snapshot_after, detected_at, significance, kind, blocks_added, blocks_modified, " | |
| 99 | + "text_delta_ratio, similarity, diff, structured_delta, status) values (:id, :sid, :cid, 'pricing', :b, :a, :when, 0.72, 'major', 1, 1, 0.18, 0.82, " | |
| 100 | + "cast(:diff as jsonb), cast(:sd as jsonb), 'processed')", | |
| 101 | + id=ids["change"], sid=ids["sensor_home"], cid=ids["alpha"], b=ids["snap1"], a=ids["snap2"], when=now - timedelta(days=1), diff=jsonb(diff), | |
| 102 | + sd=jsonb({"plans": {"price_changed": [{"plan_name": "Starter", "before": 10, "after": 12, "currency": "USD", "billing_period": "month", "pct": 20}]}})) | |
| 103 | + events = ( | |
| 104 | + (ids["ev_pricing"], ids["alpha"], ids["sensor_home"], ids["change"], "pricing", "PRICING", "PRICE_INCREASE", 0.8, 0.92, "HIGH_CONFIDENCE", | |
| 105 | + "Starter plan price increased from $10 to $12 per month", "$10", "$12", now - timedelta(minutes=30), ["pricing"], "active"), | |
| 106 | + (ids["ev_hiring"], ids["alpha"], ids["sensor_careers"], None, "careers", "HIRING", "JOB_COUNT_INCREASE", 0.6, 0.85, "HIGH_CONFIDENCE", | |
| 107 | + "Open positions increased from 1 to 2", "1", "2", now - timedelta(hours=3), ["hiring", "ai"], "active"), | |
| 108 | + (ids["ev_product"], ids["alpha"], ids["sensor_home"], None, "pricing", "PRODUCT", "NEW_PRODUCT", 0.7, 0.75, "LIKELY", | |
| 109 | + "New plan listed: Team", None, "Team", now - timedelta(hours=6), ["product"], "active"), | |
| 110 | + (ids["ev_leader"], ids["beta"], ids["sensor_beta"], None, "homepage", "LEADERSHIP", "NEW_EXECUTIVE", 0.8, 0.7, "LIKELY", | |
| 111 | + "New executive listed: Jane Ztest (Chief Test Officer)", None, "Jane Ztest", now - timedelta(days=3), ["leadership"], "active"), | |
| 112 | + (ids["ev_retracted"], ids["beta"], ids["sensor_beta"], None, "homepage", "OTHER", "OTHER", 0.2, 0.5, "INFERRED", | |
| 113 | + "Retracted synthetic event", None, None, now - timedelta(days=4), [], "retracted"), | |
| 114 | + ) | |
| 115 | + for eid, cid, sid, chg, surface, et, est, imp, conf, label, title, old, new, when, tags, status in events: | |
| 116 | + await execute(conn, "insert into events (id, company_id, sensor_id, change_id, surface, event_type, event_subtype, importance, confidence, confidence_label, title, " | |
| 117 | + "summary, old_value, new_value, payload, entities, tags, detected_at, source_url, origin, status, dedupe_key) values (:id, :cid, :sid, :chg, :surface, " | |
| 118 | + ":et, :est, :imp, :conf, :label, :title, 'Detected on a monitored public page.', :old, :new, '{}'::jsonb, '{}'::jsonb, cast(:tags as text[]), :when, " | |
| 119 | + ":url, 'deterministic', :status, :dk)", | |
| 120 | + id=eid, cid=cid, sid=sid, chg=chg, surface=surface, et=et, est=est, imp=imp, conf=conf, label=label, title=title, old=old, new=new, tags=tags, | |
| 121 | + when=when, url=f"https://{d['alpha_slug']}.example/{surface}", status=status, dk=f"ztest:{eid}") | |
| 122 | + await execute(conn, "insert into event_sources (event_id, sensor_id, source_url, snapshot_id, surface, kind) values (:e, :s, :u, :snap, 'pricing', 'primary')", | |
| 123 | + e=ids["ev_pricing"], s=ids["sensor_home"], u=f"https://{d['alpha_slug']}.example/pricing", snap=ids["snap2"]) | |
| 124 | + for i, (title, status, ai, country) in enumerate((("Senior ML Engineer", "open", True, "ZZ"), ("Account Executive", "open", False, "ZZ"), | |
| 125 | + ("Office Manager", "no_longer_listed", False, None))): | |
| 126 | + await execute(conn, "insert into jobs (id, company_id, sensor_id, fingerprint, title, department, location_text, city, country, remote, employment_type, first_seen_at, " | |
| 127 | + "last_seen_at, removed_at, status, is_ai) values (:id, :cid, :sid, :fp, :title, 'Engineering', 'Testville', 'Testville', :country, :remote, 'full_time', " | |
| 128 | + ":seen, :last, :removed, :status, :ai)", | |
| 129 | + id=f"job_ztestapi{SUFFIX}{i}", cid=ids["alpha"], sid=ids["sensor_careers"], fp=f"ztest-{SUFFIX}-{i}", title=title, country=country, remote=(i == 0), | |
| 130 | + seen=now - timedelta(days=3 + i), last=now - timedelta(hours=1), removed=(now - timedelta(days=1)) if status != "open" else None, status=status, ai=ai) | |
| 131 | + await execute(conn, "insert into people (id, company_id, name, name_norm, title, role_category, is_executive, status) values (:id, :cid, 'Jane Ztest', :n, 'Chief Test Officer', 'other', true, 'listed')", | |
| 132 | + id=f"person_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"janeztest{SUFFIX}") | |
| 133 | + await execute(conn, "insert into people (id, company_id, name, name_norm, title, role_category, is_executive, status, removed_at) values (:id, :cid, 'John Former', :n, 'CFO', 'cfo', true, 'no_longer_listed', now())", | |
| 134 | + id=f"person_ztestapi{SUFFIX}2", cid=ids["alpha"], n=f"johnformer{SUFFIX}") | |
| 135 | + await execute(conn, "insert into products (id, company_id, name, name_norm, category, status) values (:id, :cid, 'Ztest Widget', :n, 'widgets', 'listed')", | |
| 136 | + id=f"prod_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"ztestwidget{SUFFIX}") | |
| 137 | + await execute(conn, "insert into pricing_plans (id, company_id, plan_name, plan_norm, currency, billing_period, price, price_text, features, status, version_no) values " | |
| 138 | + "(:id, :cid, 'Starter', 'starter', 'USD', 'month', 12, '$12 per month', '[\"1 seat\"]'::jsonb, 'current', 2)", id=f"plan_ztestapi{SUFFIX}2", cid=ids["alpha"]) | |
| 139 | + await execute(conn, "insert into pricing_plans (id, company_id, plan_name, plan_norm, currency, billing_period, price, price_text, status, version_no, valid_to) values " | |
| 140 | + "(:id, :cid, 'Starter', 'starter', 'USD', 'month', 10, '$10 per month', 'superseded', 1, now())", id=f"plan_ztestapi{SUFFIX}1", cid=ids["alpha"]) | |
| 141 | + await execute(conn, "insert into locations (id, company_id, kind, name, name_norm, city, country, lat, lon, status) values (:id, :cid, 'headquarters', 'Testville HQ', :n, 'Testville', 'ZZ', 45.5, -73.6, 'listed')", | |
| 142 | + id=f"loc_ztestapi{SUFFIX}1", cid=ids["alpha"], n=f"testvillehq{SUFFIX}") | |
| 143 | + await execute(conn, "insert into news_items (id, company_id, url, canonical_url, title, category, published_at) values (:id, :cid, :u, :u, 'Ztest Alpha announces Team plan', 'press', now())", | |
| 144 | + id=f"news_ztestapi{SUFFIX}1", cid=ids["alpha"], u=f"https://{d['alpha_slug']}.example/news/team-plan") | |
| 145 | + for cid, metric, value in ((ids["alpha"], "activity_score", 72.34), (ids["alpha"], "hiring_momentum_30d", 12.5), (ids["alpha"], "open_jobs", 2), (ids["alpha"], "ai_adoption", 40.0), | |
| 146 | + (ids["alpha"], "product_velocity", 55.0), (ids["beta"], "activity_score", 30.0), (ids["beta"], "hiring_momentum_30d", -8.0)): | |
| 147 | + await execute(conn, "insert into metrics_current (company_id, metric, value, confidence, inputs, formula_version) values (:cid, :m, :v, 0.8, '{\"events\": 3}'::jsonb, 'metrics-v1')", | |
| 148 | + cid=cid, m=metric, v=value) | |
| 149 | + for i in range(10): | |
| 150 | + await execute(conn, "insert into metric_series (company_id, metric, day, value, confidence, formula_version) values (:cid, 'activity_score', :day, :v, 0.8, 'metrics-v1')", | |
| 151 | + cid=ids["alpha"], day=(now - timedelta(days=9 - i)).date(), v=50 + i * 2) | |
| 152 | + await execute(conn, "insert into company_daily (company_id, day, observations, changes, events) values (:cid, :day, 5, 1, 2)", cid=ids["alpha"], day=(now - timedelta(days=1)).date()) | |
| 153 | + await execute(conn, "insert into signals (id, company_id, scope, scope_key, kind, strength, confidence, title, explanation, window_days) values (:id, :cid, 'company', :slug, " | |
| 154 | + "'hiring_surge', 0.7, 0.6, 'Hiring surge signal', 'Open positions doubled over 7 days.', 7)", id=ids["signal"], cid=ids["alpha"], slug=d["alpha_slug"]) | |
| 155 | + await execute(conn, "insert into queue_jobs (id, kind, key, payload, status, attempts, max_attempts, last_error) values (:id, 'run_sensor', :key, '{}'::jsonb, 'dead', 3, 3, 'boom')", | |
| 156 | + id=ids["queue_dead"], key=f"ztest:{SUFFIX}:dead") | |
| 157 | + await execute(conn, "insert into failures (id, sensor_id, company_id, failure_class, status_code, message, url) values (:id, :sid, :cid, 'HTTP_5XX', 503, 'synthetic', 'https://x.example')", | |
| 158 | + id=ids["failure"], sid=ids["sensor_beta"], cid=ids["beta"]) | |
| 159 | + await execute(conn, "insert into review_queue (id, kind, ref_id, company_id, payload) values (:id, 'major_event', :ref, :cid, '{}'::jsonb)", | |
| 160 | + id=ids["review"], ref=ids["ev_pricing"], cid=ids["alpha"]) | |
| 161 | + await execute(conn, "insert into api_keys (id, key_hash, prefix, name, tier) values (:id, :h, :p, 'ztest paid', 'paid')", | |
| 162 | + id=ids["api_key"], h=hashlib.sha256(RAW_API_KEY.encode()).hexdigest(), p=RAW_API_KEY[:12]) | |
| 163 | + return d | |
| 164 | + | |
| 165 | + | |
| 166 | +async def _cleanup(d: dict[str, Any]) -> None: | |
| 167 | + async with transaction() as conn: | |
| 168 | + await execute(conn, "delete from companies where slug like :p", p=f"ztest-api-%{d['suffix']}%") | |
| 169 | + await execute(conn, "delete from companies where slug like 'ztest-api-created-%'") | |
| 170 | + await execute(conn, "delete from companies where canonical_domain like 'ztest-created-%'") | |
| 171 | + await execute(conn, "delete from queue_jobs where key like :p", p=f"ztest:{d['suffix']}%") | |
| 172 | + await execute(conn, "delete from queue_jobs where key like 'discover:co_ztestapi%'") | |
| 173 | + await execute(conn, "delete from api_keys where id = :id", id=d["api_key"]) | |
| 174 | + await execute(conn, "delete from owners where token_hash = :h", h=hashlib.sha256(OWNER_TOKEN.encode()).hexdigest()) | |
| 175 | + await execute(conn, "delete from connectors where id = :id", id=d["connector"]) | |
| 176 | + await execute(conn, "delete from industries where slug = :s", s=d["industry"]) | |
| 177 | + left = await fetch_val(conn, "select count(*) from companies where country = 'ZZ'") | |
| 178 | + if not left: | |
| 179 | + await execute(conn, "delete from countries where code = 'ZZ'") | |
| 180 | + for key in d["objects"]: | |
| 181 | + try: | |
| 182 | + os.unlink(archive.object_path(key)) | |
| 183 | + except OSError: | |
| 184 | + pass | |
| 185 | + cache.clear() | |
| 186 | + | |
| 187 | + | |
| 188 | +_STATE: dict[str, Any] = {"refs": 0, "data": None} # one dataset per process even though each test module re-exports the fixture | |
| 189 | + | |
| 190 | + | |
| 191 | +@pytest_asyncio.fixture(loop_scope="session", scope="session") | |
| 192 | +async def fixture_data(): # type: ignore[no-untyped-def] | |
| 193 | + from companyatlas.api import ratelimit | |
| 194 | + | |
| 195 | + if _STATE["data"] is None: | |
| 196 | + cache.clear() | |
| 197 | + _STATE["anon_limit"] = ratelimit.limiter.limits["anonymous"] | |
| 198 | + ratelimit.limiter.limits["anonymous"] = 1_000_000 # the suite alone exceeds 120 req/min; the 429 path has its own test | |
| 199 | + _STATE["data"] = await _insert_fixture() | |
| 200 | + _STATE["refs"] += 1 | |
| 201 | + try: | |
| 202 | + yield _STATE["data"] | |
| 203 | + finally: | |
| 204 | + _STATE["refs"] -= 1 | |
| 205 | + if _STATE["refs"] == 0: | |
| 206 | + data, _STATE["data"] = _STATE["data"], None | |
| 207 | + ratelimit.limiter.limits["anonymous"] = _STATE["anon_limit"] | |
| 208 | + await _cleanup(data) | |
| 209 | + await dispose() | |
| 210 | + | |
| 211 | + | |
| 212 | +@pytest_asyncio.fixture(loop_scope="session") | |
| 213 | +async def client(fixture_data): # type: ignore[no-untyped-def] | |
| 214 | + from companyatlas.api.main import app | |
| 215 | + | |
| 216 | + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver", timeout=30) as c: | |
| 217 | + yield c | |
| 218 | + | |
| 219 | + | |
| 220 | +__all__ = ["ADMIN", "OWNER", "OWNER_TOKEN", "RAW_API_KEY", "client", "fixture_data"] | |
added
tests/test_ask_parser.py
+73 −0
@@ -0,0 +1,73 @@ | ||
| 1 | +"""Deterministic question parser for `/ask`: countries, industries, intents → event types, windows, answer style; never invents results.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from companyatlas.services.llm.ask import Interpretation, build_answer, parse_question, route_question | |
| 5 | + | |
| 6 | +INDUSTRIES = {"fintech": "Fintech", "artificial-intelligence": "Artificial intelligence", "retail": "Retail", "software": "Software"} | |
| 7 | +COUNTRIES = {"CA": "Canada", "US": "United States", "JP": "Japan", "DE": "Germany"} | |
| 8 | + | |
| 9 | + | |
| 10 | +def test_hiring_ai_in_canada_last_month(): | |
| 11 | + it = parse_question("Which companies are hiring AI engineers in Canada in the last month?", industries=INDUSTRIES, countries=COUNTRIES) | |
| 12 | + assert it.countries == ["CA"] | |
| 13 | + assert it.intent == "ai" and "hiring" in it.intents | |
| 14 | + assert "HIRING" in it.event_types and "AI_HIRING" in it.event_subtypes and "ai" in it.tags | |
| 15 | + assert it.window_days == 30 | |
| 16 | + assert "engineers" in it.keywords and "canada" not in it.keywords | |
| 17 | + assert it.answer_style == "list" and 0.5 <= it.confidence <= 0.9 | |
| 18 | + assert it.filters["countries"] == ["CA"] | |
| 19 | + | |
| 20 | + | |
| 21 | +def test_pricing_increase_fintech_this_week(): | |
| 22 | + it = parse_question("fintech companies that raised prices this week", industries=INDUSTRIES) | |
| 23 | + assert it.industries == ["fintech"] and it.intent == "pricing" | |
| 24 | + assert it.event_types == ["PRICING"] and "PRICE_INCREASE" in it.event_subtypes and "PRICE_DECREASE" not in it.event_subtypes | |
| 25 | + assert it.window_days == 7 | |
| 26 | + | |
| 27 | + | |
| 28 | +def test_launch_expansion_leadership_and_count(): | |
| 29 | + launch = parse_question("How many Japanese companies launched new products since 2026-01-01?", countries=COUNTRIES) | |
| 30 | + assert launch.countries == ["JP"] and launch.intent == "launch" and launch.event_types == ["PRODUCT"] | |
| 31 | + assert launch.answer_style == "count" and launch.window_days is not None and launch.window_days > 200 | |
| 32 | + exp = parse_question("companies expanding into new countries", countries=COUNTRIES) | |
| 33 | + assert exp.intent == "expansion" and "COUNTRY_EXPANSION" in exp.event_subtypes | |
| 34 | + lead = parse_question("new CEO appointed at German companies", countries=COUNTRIES) | |
| 35 | + assert lead.intent == "leadership" and lead.countries == ["DE"] | |
| 36 | + | |
| 37 | + | |
| 38 | +def test_compare_timeline_trend_and_quotes(): | |
| 39 | + cmp = parse_question('compare "Stripe" vs "Adyen" hiring over the past 90 days') | |
| 40 | + assert cmp.answer_style == "compare" and cmp.companies == ["Stripe", "Adyen"] and cmp.window_days == 90 and "HIRING" in cmp.event_types | |
| 41 | + tl = parse_question("when did Acme change its terms of service?") | |
| 42 | + assert tl.answer_style == "timeline" and tl.intent == "legal" and tl.event_types == ["LEGAL"] | |
| 43 | + tr = parse_question("which industries are trending in developer docs") | |
| 44 | + assert tr.answer_style == "trend" and "DEVELOPER" in tr.event_types | |
| 45 | + | |
| 46 | + | |
| 47 | +def test_plain_search_has_no_filters_and_low_confidence(): | |
| 48 | + it = parse_question("acme corporation") | |
| 49 | + assert it.intent == "search" and it.filters == {"keywords": ["acme", "corporation"]} | |
| 50 | + assert it.confidence < 0.5 | |
| 51 | + empty = parse_question("") | |
| 52 | + assert empty.filters == {} and empty.to_dict()["parser_version"].startswith("ask-parser") | |
| 53 | + | |
| 54 | + | |
| 55 | +def test_important_events_and_regions(): | |
| 56 | + it = parse_question("major pricing changes at European software companies", industries=INDUSTRIES) | |
| 57 | + assert it.min_importance == 0.6 and it.regions == ["europe"] and it.countries == [] and it.industries == ["software"] | |
| 58 | + | |
| 59 | + | |
| 60 | +async def test_route_question_without_llm_is_deterministic(monkeypatch): | |
| 61 | + from companyatlas.config import settings | |
| 62 | + | |
| 63 | + monkeypatch.setattr(settings, "llm_base_url", "") | |
| 64 | + it = await route_question("companies hiring in Canada", countries=COUNTRIES) | |
| 65 | + assert it.source == "deterministic" and it.countries == ["CA"] | |
| 66 | + | |
| 67 | + | |
| 68 | +def test_build_answer_only_phrases_measured_counts(): | |
| 69 | + it = Interpretation(question="q", intent="hiring", countries=["CA"], window_days=30) | |
| 70 | + assert build_answer(it, companies=0, events=0).startswith("No monitored evidence matches this question yet in CA over the last 30 days") | |
| 71 | + text = build_answer(it, companies=12, events=48) | |
| 72 | + assert text.startswith("48 hiring-related events in CA over the last 30 days across 12 monitored companies.") | |
| 73 | + assert "confidence" in text | |
added
tests/test_events_rules.py
+332 −0
@@ -0,0 +1,332 @@ | ||
| 1 | +"""Deterministic event rules: subtypes, wording, importance/confidence, idempotency and cross-surface clustering.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime, timedelta | |
| 5 | + | |
| 6 | +import pytest | |
| 7 | +from factories import intel_db # noqa: F401 — registers the fixture | |
| 8 | + | |
| 9 | +from companyatlas.services.events import derive_events, safe_wording, scale_importance | |
| 10 | +from companyatlas.taxonomy import FORBIDDEN_WORDING, ChangeKind | |
| 11 | + | |
| 12 | +COMPANY = {"id": "co_test", "slug": "ztest-acme", "display_name": "Acme", "country": "CA", "industries": []} | |
| 13 | + | |
| 14 | + | |
| 15 | +def _change(surface: str, *, significance: float = 0.6, delta: dict | None = None, diff: dict | None = None, kind: str | None = None) -> dict: | |
| 16 | + from companyatlas.taxonomy import change_kind | |
| 17 | + | |
| 18 | + return {"id": "chg_test", "sensor_id": "sen_test", "company_id": "co_test", "surface": surface, "significance": significance, | |
| 19 | + "kind": kind or str(change_kind(significance)), "structured_delta": delta or {}, "diff": diff or {}, "detected_at": datetime(2026, 9, 12, 12, tzinfo=UTC), | |
| 20 | + "snapshot_before": "snap_a", "snapshot_after": "snap_b", "blocks_added": 0, "blocks_removed": 0, "blocks_modified": 0, "text_delta_ratio": 0.2} | |
| 21 | + | |
| 22 | + | |
| 23 | +def _sensor(surface: str, connector: str = "generic-html-v1") -> dict: | |
| 24 | + return {"id": "sen_test", "surface": surface, "connector_id": connector, "url": f"https://acme.example/{surface}"} | |
| 25 | + | |
| 26 | + | |
| 27 | +def _by_subtype(derived): # type: ignore[no-untyped-def] | |
| 28 | + out: dict[str, list] = {} | |
| 29 | + for d in derived.events: | |
| 30 | + out.setdefault(d.subtype, []).append(d) | |
| 31 | + return out | |
| 32 | + | |
| 33 | + | |
| 34 | +def _assert_clean(derived) -> None: # type: ignore[no-untyped-def] | |
| 35 | + for d in derived.events: | |
| 36 | + low = (d.title + " " + (d.summary or "")).lower() | |
| 37 | + assert not any(bad in low for bad in FORBIDDEN_WORDING), d.title | |
| 38 | + | |
| 39 | + | |
| 40 | +# ------------------------------------------------------------------------------------------------------------ hiring | |
| 41 | + | |
| 42 | + | |
| 43 | +def test_job_count_increase_aggregate_and_ai(): | |
| 44 | + added = [{"title": f"Engineer {i}", "country": "CA", "department": "Engineering"} for i in range(12)] | |
| 45 | + added[0]["title"] = "Senior Machine Learning Engineer" | |
| 46 | + added[1]["title"] = "AI Product Manager" | |
| 47 | + delta = {"jobs": {"added": added, "removed": [], "open_before": 40, "open_after": 52}} | |
| 48 | + d = derive_events(_change("careers", delta=delta), COMPANY, _sensor("careers")) | |
| 49 | + by = _by_subtype(d) | |
| 50 | + assert "JOB_COUNT_INCREASE" in by and "NEW_JOB" not in by # > 5 added → aggregate only | |
| 51 | + ev = by["JOB_COUNT_INCREASE"][0] | |
| 52 | + assert ev.title == "12 new positions detected on careers page" | |
| 53 | + assert ev.old_value == "40" and ev.new_value == "52" | |
| 54 | + assert len(ev.entities["jobs"]) == 12 | |
| 55 | + assert "AI_HIRING" in by and by["AI_HIRING"][0].title == "2 AI-related positions detected on careers page" | |
| 56 | + assert ev.confidence == 0.8 # HTML extraction | |
| 57 | + assert 0 < ev.importance <= 1 | |
| 58 | + assert d.needs_classification is False | |
| 59 | + _assert_clean(d) | |
| 60 | + | |
| 61 | + | |
| 62 | +def test_per_job_events_when_few_added_and_ats_confidence(): | |
| 63 | + delta = {"jobs": {"added": [{"title": "Data Scientist", "location_text": "Toronto, CA"}, {"title": "Account Executive", "remote": True}], "removed": [], | |
| 64 | + "open_before": 10, "open_after": 12}} | |
| 65 | + d = derive_events(_change("jobs_board", delta=delta), COMPANY, _sensor("jobs_board", "greenhouse-v1")) | |
| 66 | + by = _by_subtype(d) | |
| 67 | + assert len(by["NEW_JOB"]) == 2 | |
| 68 | + assert by["NEW_JOB"][0].title == "New position listed: Data Scientist (Toronto, CA)" | |
| 69 | + assert by["NEW_JOB"][1].title == "New position listed: Account Executive (Remote)" | |
| 70 | + assert all(e.confidence == 0.95 for e in d.events) # structured ATS JSON | |
| 71 | + | |
| 72 | + | |
| 73 | +def test_job_count_decrease_wording_and_freeze_signal(): | |
| 74 | + removed = [{"title": f"Role {i}"} for i in range(30)] | |
| 75 | + delta = {"jobs": {"added": [], "removed": removed, "open_before": 40, "open_after": 10}} | |
| 76 | + d = derive_events(_change("careers", significance=0.7, delta=delta), COMPANY, _sensor("careers")) | |
| 77 | + by = _by_subtype(d) | |
| 78 | + dec = by["JOB_COUNT_DECREASE"][0] | |
| 79 | + assert dec.title == "30 monitored job listings no longer visible on careers page" | |
| 80 | + assert "HIRING_FREEZE_SIGNAL" in by | |
| 81 | + assert "no longer visible" in by["HIRING_FREEZE_SIGNAL"][0].title | |
| 82 | + assert by["HIRING_FREEZE_SIGNAL"][0].review == "unexpected_activity" | |
| 83 | + _assert_clean(d) | |
| 84 | + | |
| 85 | + | |
| 86 | +def test_hiring_surge_uses_company_baseline(): | |
| 87 | + added = [{"title": f"Role {i}"} for i in range(9)] | |
| 88 | + delta = {"jobs": {"added": added, "removed": [], "open_before": 100, "open_after": 109}} | |
| 89 | + no_baseline = derive_events(_change("careers", delta=delta), COMPANY, _sensor("careers")) | |
| 90 | + assert "HIRING_SURGE" not in _by_subtype(no_baseline) # 9 < fallback threshold of 10 | |
| 91 | + baseline = {"jobs_new_weekly": {"mean": 1.0, "stddev": 1.0, "samples": 8}} | |
| 92 | + with_baseline = derive_events(_change("careers", delta=delta), COMPANY, _sensor("careers"), baseline=baseline) | |
| 93 | + surge = _by_subtype(with_baseline)["HIRING_SURGE"][0] | |
| 94 | + assert "baseline ≈ 1.0 new/week" in surge.title | |
| 95 | + | |
| 96 | + | |
| 97 | +# ------------------------------------------------------------------------------------------------------------ pricing | |
| 98 | + | |
| 99 | + | |
| 100 | +def test_price_increase_and_tier_changes(): | |
| 101 | + delta = {"plans": {"price_changed": [{"plan_name": "Pro", "before": 49, "after": 59, "currency": "USD", "billing_period": "month"}], | |
| 102 | + "added": [{"plan_name": "Enterprise", "contact_sales": True}], "removed": [{"plan_name": "Starter", "price": 9, "currency": "USD"}]}} | |
| 103 | + d = derive_events(_change("pricing", significance=0.7, delta=delta), COMPANY, _sensor("pricing")) | |
| 104 | + by = _by_subtype(d) | |
| 105 | + inc = by["PRICE_INCREASE"][0] | |
| 106 | + assert inc.title == "Pro plan price observed at $59 (was $49)" | |
| 107 | + assert inc.old_value == "$49" and inc.new_value == "$59" | |
| 108 | + assert inc.payload["pct"] == pytest.approx(20.4, abs=0.1) | |
| 109 | + assert by["NEW_PRICING_TIER"][0].title == "New pricing tier listed: Enterprise (contact sales)" | |
| 110 | + assert "enterprise" in by["NEW_PRICING_TIER"][0].tags | |
| 111 | + assert by["PRICING_TIER_REMOVED"][0].title == "Pricing tier no longer listed: Starter" | |
| 112 | + assert inc.importance > by["PRICING_TIER_REMOVED"][0].importance | |
| 113 | + | |
| 114 | + | |
| 115 | +def test_price_decrease_eur(): | |
| 116 | + delta = {"plans": {"price_changed": [{"plan_name": "Team", "before": 30, "after": 24, "currency": "EUR", "billing_period": "month", "pct": -20}]}} | |
| 117 | + d = derive_events(_change("pricing", delta=delta), COMPANY, _sensor("pricing")) | |
| 118 | + ev = _by_subtype(d)["PRICE_DECREASE"][0] | |
| 119 | + assert ev.title == "Team plan price observed at €24 (was €30)" | |
| 120 | + assert ev.summary == "Decrease of 20.0%, billed per month." | |
| 121 | + | |
| 122 | + | |
| 123 | +def test_generic_pricing_change_from_text_diff(): | |
| 124 | + diff = {"modified": [{"path": "Pricing > Pro", "before": "a", "after": "b"}, {"path": "Pricing > FAQ", "before": "c", "after": "d"}], "added": [], "removed": [], | |
| 125 | + "counts": {"added": 0, "removed": 0, "modified": 2}, "text_delta_ratio": 0.3} | |
| 126 | + d = derive_events(_change("pricing", diff=diff), COMPANY, _sensor("pricing")) | |
| 127 | + ev = _by_subtype(d)["PRICING_CHANGE"][0] | |
| 128 | + assert ev.title == "Pricing page materially updated (2 blocks changed)" | |
| 129 | + assert ev.confidence == 0.7 # text-diff only | |
| 130 | + assert ev.payload["sections"] == ["Pro", "FAQ"] | |
| 131 | + | |
| 132 | + | |
| 133 | +# ------------------------------------------------------------------------------------------------------------ leadership | |
| 134 | + | |
| 135 | + | |
| 136 | +def test_leadership_events_wording(): | |
| 137 | + delta = {"people": {"added": [{"name": "Jane Doe", "title": "Chief Financial Officer", "role_category": "cfo", "is_executive": True}], | |
| 138 | + "removed": [{"name": "John Roe", "title": "Chief Technology Officer", "role_category": "cto", "is_executive": True}], | |
| 139 | + "title_changed": [{"name": "Ann Lee", "before": "VP Operations", "after": "COO"}]}} | |
| 140 | + d = derive_events(_change("leadership", significance=0.7, delta=delta), COMPANY, _sensor("leadership")) | |
| 141 | + by = _by_subtype(d) | |
| 142 | + assert by["NEW_EXECUTIVE"][0].title == "Jane Doe listed as Chief Financial Officer on leadership page" | |
| 143 | + assert by["EXECUTIVE_NO_LONGER_LISTED"][0].title == "John Roe no longer listed on leadership page" | |
| 144 | + assert by["EXECUTIVE_TITLE_CHANGE"][0].title == "Ann Lee now listed as COO (was VP Operations)" | |
| 145 | + assert by["LEADERSHIP_CHANGE"][0].title == "Leadership page updated: 1 added, 1 no longer listed, 1 title change" | |
| 146 | + _assert_clean(d) | |
| 147 | + | |
| 148 | + | |
| 149 | +# ------------------------------------------------------------------------------------------------------------ products / locations | |
| 150 | + | |
| 151 | + | |
| 152 | +def test_product_and_location_rules(): | |
| 153 | + delta = {"products": {"added": [{"name": "Atlas Copilot", "url": "https://acme.example/copilot"}], "removed": [{"name": "Atlas Lite"}]}, | |
| 154 | + "locations": {"added": [{"name": "Toronto office", "city": "Toronto", "country": "CA", "kind": "office"}, {"name": "Tokyo", "city": "Tokyo", "country": "JP", "kind": "office"}], | |
| 155 | + "removed": [{"name": "Berlin", "city": "Berlin", "country": "DE", "kind": "office"}], "new_countries": ["JP"]}} | |
| 156 | + d = derive_events(_change("locations", delta=delta), COMPANY, _sensor("locations"), country_names={"JP": "Japan"}) | |
| 157 | + by = _by_subtype(d) | |
| 158 | + assert by["NEW_PRODUCT"][0].title == "New product listed: Atlas Copilot" | |
| 159 | + assert "ai" in by["NEW_PRODUCT"][0].tags | |
| 160 | + assert by["PRODUCT_REMOVED"][0].title == "Product no longer listed: Atlas Lite" | |
| 161 | + assert by["NEW_LOCATION"][0].title == "New office listed: Toronto, CA" | |
| 162 | + assert by["OFFICE_REMOVED"][0].title == "Office no longer listed: Berlin, DE" | |
| 163 | + assert by["COUNTRY_EXPANSION"][0].title == "New country presence listed: Japan (Tokyo)" | |
| 164 | + assert by["COUNTRY_EXPANSION"][0].importance > by["NEW_LOCATION"][0].importance | |
| 165 | + | |
| 166 | + | |
| 167 | +# ------------------------------------------------------------------------------------------------------------ news / legal / homepage | |
| 168 | + | |
| 169 | + | |
| 170 | +def test_news_subtype_heuristics(): | |
| 171 | + delta = {"news": {"added": [{"title": "Acme announces Q3 2026 earnings results", "url": "https://acme.example/ir/q3", "category": "press"}, | |
| 172 | + {"title": "How we built our new AI search", "url": "https://acme.example/blog/ai-search", "category": "blog"}, | |
| 173 | + {"title": "v2.4 — webhooks and SDK updates", "url": "https://acme.example/changelog/2-4", "category": "changelog"}]}} | |
| 174 | + diff = {"added": [{"path": "Feed", "before": None, "after": "Acme announces Q3 2026 earnings results — revenue up…"}], "modified": [], "removed": [], "counts": {"added": 1}} | |
| 175 | + d = derive_events(_change("feed", delta=delta, diff=diff), COMPANY, _sensor("feed", "rss-feed-v1")) | |
| 176 | + by = _by_subtype(d) | |
| 177 | + assert "EARNINGS_RELEASE" in by and by["EARNINGS_RELEASE"][0].title.startswith("Earnings release: ") | |
| 178 | + assert "BLOG_POST" in by and "ai" in by["BLOG_POST"][0].tags | |
| 179 | + assert "CHANGELOG_ENTRY" in by | |
| 180 | + assert all(e.confidence == 0.9 for e in d.events) # feed / JSON-LD grade evidence | |
| 181 | + assert "EARNINGS_RELEASE" in d.summarize | |
| 182 | + | |
| 183 | + | |
| 184 | +def test_terms_change_sections_and_review(): | |
| 185 | + diff = {"modified": [{"path": "Terms > 7. Termination", "before": "x", "after": "y"}, {"path": "Terms > 12. Governing law", "before": "x", "after": "y"}], | |
| 186 | + "added": [{"path": "Terms > 14. Arbitration", "before": None, "after": "New section"}], "removed": [], "counts": {"added": 1, "removed": 0, "modified": 2}, | |
| 187 | + "text_delta_ratio": 0.12} | |
| 188 | + d = derive_events(_change("legal_terms", significance=0.55, diff=diff), COMPANY, _sensor("legal_terms")) | |
| 189 | + ev = _by_subtype(d)["TERMS_CHANGE"][0] | |
| 190 | + assert ev.title == "Terms of service page materially updated (3 sections changed)" | |
| 191 | + assert ev.payload["sections"] == ["7. Termination", "12. Governing law", "14. Arbitration"] | |
| 192 | + assert ev.review == "legal_sensitive" and ev.confidence == 0.7 | |
| 193 | + assert "TERMS_CHANGE" in d.summarize | |
| 194 | + | |
| 195 | + | |
| 196 | +def test_homepage_redesign_vs_change_and_messaging(): | |
| 197 | + diff = {"modified": [{"path": "Hero", "before": "Old", "after": "New"}], "added": [], "removed": [], "counts": {"added": 0, "removed": 0, "modified": 1}, "text_delta_ratio": 0.6} | |
| 198 | + meta = {"meta": {"title_changed": {"before": "Acme — Payments", "after": "Acme — AI Payments Platform"}}} | |
| 199 | + major = derive_events(_change("homepage", significance=0.7, diff=diff, delta=meta), COMPANY, _sensor("homepage")) | |
| 200 | + by = _by_subtype(major) | |
| 201 | + assert "HOMEPAGE_REDESIGN" in by and "MESSAGING_CHANGE" in by | |
| 202 | + assert by["MESSAGING_CHANGE"][0].old_value == "Acme — Payments" | |
| 203 | + meaningful = derive_events(_change("homepage", significance=0.5, diff=diff), COMPANY, _sensor("homepage")) | |
| 204 | + assert list(_by_subtype(meaningful)) == ["WEBSITE_CHANGE"] | |
| 205 | + assert meaningful.needs_classification is True and meaningful.classification_reason == "ambiguous_surface" | |
| 206 | + | |
| 207 | + | |
| 208 | +def test_docs_api_changelog_text_rules(): | |
| 209 | + diff = {"added": [{"path": "Changelog", "before": None, "after": "2026-09-10: Added bulk export endpoint\nDetails…"}], "modified": [], "removed": [], | |
| 210 | + "counts": {"added": 1, "removed": 0, "modified": 0}, "text_delta_ratio": 0.15} | |
| 211 | + assert _by_subtype(derive_events(_change("changelog", diff=diff), COMPANY, _sensor("changelog")))["CHANGELOG_ENTRY"][0].title == \ | |
| 212 | + "Changelog entry detected: 2026-09-10: Added bulk export endpoint" | |
| 213 | + assert "API_CHANGE" in _by_subtype(derive_events(_change("api", diff=diff), COMPANY, _sensor("api"))) | |
| 214 | + assert "DOC_CHANGE" in _by_subtype(derive_events(_change("docs", diff=diff), COMPANY, _sensor("docs"))) | |
| 215 | + | |
| 216 | + | |
| 217 | +# ------------------------------------------------------------------------------------------------------------ guards | |
| 218 | + | |
| 219 | + | |
| 220 | +@pytest.mark.parametrize("significance", [0.05, 0.3]) | |
| 221 | +def test_noise_and_minor_never_emit(significance): | |
| 222 | + delta = {"jobs": {"added": [{"title": "X"}] * 20, "removed": [], "open_before": 1, "open_after": 21}} | |
| 223 | + d = derive_events(_change("careers", significance=significance, delta=delta), COMPANY, _sensor("careers")) | |
| 224 | + assert d.events == [] and d.needs_classification is False | |
| 225 | + | |
| 226 | + | |
| 227 | +def test_unknown_surface_without_structure_requests_classification(): | |
| 228 | + diff = {"modified": [{"path": "Partners", "before": "a", "after": "b"}], "added": [], "removed": [], "counts": {"modified": 1}} | |
| 229 | + d = derive_events(_change("partners", diff=diff), COMPANY, _sensor("partners")) | |
| 230 | + assert d.events == [] and d.needs_classification and d.classification_reason == "no_deterministic_event" | |
| 231 | + | |
| 232 | + | |
| 233 | +def test_critical_changes_go_to_review_and_importance_scales(): | |
| 234 | + delta = {"plans": {"price_changed": [{"plan_name": "Pro", "before": 10, "after": 12, "currency": "USD"}]}} | |
| 235 | + crit = derive_events(_change("pricing", significance=0.9, delta=delta), COMPANY, _sensor("pricing")) | |
| 236 | + mid = derive_events(_change("pricing", significance=0.5, delta=delta), COMPANY, _sensor("pricing")) | |
| 237 | + assert crit.events[0].review == "major_event" and crit.events[0].importance > mid.events[0].importance | |
| 238 | + assert scale_importance(0.8, 0.5) == pytest.approx(0.8, abs=1e-6) | |
| 239 | + assert scale_importance(0.8, 1.0, 1.0) == 1.0 | |
| 240 | + | |
| 241 | + | |
| 242 | +def test_safe_wording_rewrites_forbidden_phrases(): | |
| 243 | + assert "no longer listed" in safe_wording("72 employees laid off") | |
| 244 | + assert "shut down" not in safe_wording("office shut down").lower() | |
| 245 | + | |
| 246 | + | |
| 247 | +# ============================================================================================================ database flow | |
| 248 | + | |
| 249 | + | |
| 250 | +@pytest.mark.usefixtures("intel_db") | |
| 251 | +async def test_process_pending_is_idempotent_and_clusters_across_surfaces(): | |
| 252 | + from factories import cleanup, make_change, make_company, make_sensor | |
| 253 | + | |
| 254 | + from companyatlas.db import fetch_all, fetch_one, fetch_val, transaction | |
| 255 | + from companyatlas.services.events import process_pending_changes, reprocess_events | |
| 256 | + | |
| 257 | + try: | |
| 258 | + async with transaction() as conn: | |
| 259 | + co = await make_company(conn) | |
| 260 | + newsroom = await make_sensor(conn, co, "newsroom") | |
| 261 | + feed = await make_sensor(conn, co, "feed", path="feed.xml") | |
| 262 | + careers = await make_sensor(conn, co, "careers") | |
| 263 | + item = {"title": "Acme launches Atlas Copilot for enterprises", "url": f"https://{co['canonical_domain']}/news/copilot", "category": "press"} | |
| 264 | + await make_change(conn, newsroom, significance=0.6, structured_delta={"news": {"added": [item]}}) | |
| 265 | + await make_change(conn, feed, significance=0.6, structured_delta={"news": {"added": [{**item, "url": item["url"] + "?utm=rss"}]}}) | |
| 266 | + await make_change(conn, careers, significance=0.6, structured_delta={"jobs": {"added": [{"title": "ML Engineer", "country": "CA"}], "removed": [], | |
| 267 | + "open_before": 5, "open_after": 6}}) | |
| 268 | + await make_change(conn, careers, significance=0.1, structured_delta={}) # noise → archived, no event | |
| 269 | + stats = await process_pending_changes(limit=50) | |
| 270 | + assert stats["changes"] == 3 and stats["archived"] == 1 | |
| 271 | + assert stats["events"] == 5 # NEWS×2 (one duplicate) + JOB_COUNT_INCREASE + AI_HIRING + NEW_JOB | |
| 272 | + assert stats["duplicates"] == 1 | |
| 273 | + async with transaction() as conn: | |
| 274 | + events = await fetch_all(conn, "select * from events where company_id = :c order by created_at", c=co["id"]) | |
| 275 | + news = [e for e in events if e["event_subtype"] == "NEWS_RELEASE"] | |
| 276 | + assert len(news) == 2 and {e["status"] for e in news} == {"active", "duplicate"} | |
| 277 | + canonical = next(e for e in news if e["status"] == "active") | |
| 278 | + dup = next(e for e in news if e["status"] == "duplicate") | |
| 279 | + assert canonical["cluster_id"] == dup["cluster_id"] | |
| 280 | + cluster = await fetch_one(conn, "select * from event_clusters where id = :id", id=canonical["cluster_id"]) | |
| 281 | + assert cluster["source_count"] == 2 and set(cluster["surfaces"]) == {"newsroom", "feed"} | |
| 282 | + assert canonical["confidence"] == pytest.approx(0.93, abs=0.011) # max(0.8 html, 0.9 feed) + 0.03 corroboration | |
| 283 | + assert canonical["payload"]["corroborations"] == 1 | |
| 284 | + sources = await fetch_all(conn, "select * from event_sources where event_id = :e", e=canonical["id"]) | |
| 285 | + assert {s["kind"] for s in sources} == {"primary", "corroboration"} | |
| 286 | + assert await fetch_val(conn, "select count(*) from changes where company_id = :c and status = 'pending'", c=co["id"]) == 0 | |
| 287 | + assert await fetch_val(conn, "select event_count from sensors where id = :s", s=careers["id"]) == 3 | |
| 288 | + assert await fetch_val(conn, "select last_event_at from companies where id = :c", c=co["id"]) is not None | |
| 289 | + assert await fetch_val(conn, "select count(*) from llm_jobs where company_id = :c", c=co["id"]) >= 0 | |
| 290 | + # re-run: nothing pending, nothing new | |
| 291 | + again = await process_pending_changes(limit=50) | |
| 292 | + assert again["events"] == 0 and again["changes"] == 0 | |
| 293 | + # reprocess over processed changes: dedupe keys make it a no-op | |
| 294 | + rep = await reprocess_events(datetime.now(UTC) - timedelta(days=1), company_id=co["id"]) | |
| 295 | + assert rep["changes"] == 3 and rep["events"] == 0 | |
| 296 | + async with transaction() as conn: | |
| 297 | + assert await fetch_val(conn, "select count(*) from events where company_id = :c", c=co["id"]) == 5 | |
| 298 | + finally: | |
| 299 | + await cleanup() | |
| 300 | + | |
| 301 | + | |
| 302 | +@pytest.mark.usefixtures("intel_db") | |
| 303 | +async def test_review_queue_and_llm_jobs_for_legal_change(monkeypatch): | |
| 304 | + from factories import cleanup, make_change, make_company, make_sensor | |
| 305 | + | |
| 306 | + from companyatlas.config import settings | |
| 307 | + from companyatlas.db import fetch_all, transaction | |
| 308 | + from companyatlas.services.events import process_pending_changes | |
| 309 | + | |
| 310 | + monkeypatch.setattr(settings, "llm_base_url", "https://llm.example/v1") | |
| 311 | + monkeypatch.setattr(settings, "llm_enabled", True) | |
| 312 | + try: | |
| 313 | + async with transaction() as conn: | |
| 314 | + co = await make_company(conn) | |
| 315 | + legal = await make_sensor(conn, co, "legal_terms") | |
| 316 | + other = await make_sensor(conn, co, "partners") | |
| 317 | + diff = {"modified": [{"path": "Terms > 7. Termination", "before": "old text", "after": "new text"}], "added": [], "removed": [], "counts": {"modified": 1}, "text_delta_ratio": 0.2} | |
| 318 | + await make_change(conn, legal, significance=0.55, diff=diff) | |
| 319 | + await make_change(conn, other, significance=0.5, diff=diff) | |
| 320 | + stats = await process_pending_changes() | |
| 321 | + assert stats["events"] == 1 and stats["llm_jobs"] == 2 | |
| 322 | + async with transaction() as conn: | |
| 323 | + reviews = await fetch_all(conn, "select kind from review_queue where company_id = :c", c=co["id"]) | |
| 324 | + assert {r["kind"] for r in reviews} == {"legal_sensitive"} | |
| 325 | + jobs = await fetch_all(conn, "select kind from llm_jobs where company_id = :c", c=co["id"]) | |
| 326 | + assert sorted(j["kind"] for j in jobs) == ["classify_change", "summarize_event"] | |
| 327 | + finally: | |
| 328 | + await cleanup() | |
| 329 | + | |
| 330 | + | |
| 331 | +def test_kind_enum_values_used_by_rules(): | |
| 332 | + assert ChangeKind.NOISE == "noise" and ChangeKind.CRITICAL == "critical" | |
added
tests/test_llm_gateway.py
+159 −0
@@ -0,0 +1,159 @@ | ||
| 1 | +"""LLM gateway: OpenAI-compatible JSON mode (respx-mocked), invalid-JSON repair path, 503/429 backoff, health, prompt loading.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import json | |
| 5 | + | |
| 6 | +import httpx | |
| 7 | +import pytest | |
| 8 | +import respx | |
| 9 | +from pydantic import BaseModel | |
| 10 | + | |
| 11 | +from companyatlas.services.llm.gateway import ( | |
| 12 | + LLMError, | |
| 13 | + LLMValidationError, | |
| 14 | + OpenAICompatibleProvider, | |
| 15 | + extract_json_text, | |
| 16 | + parse_json_object, | |
| 17 | +) | |
| 18 | +from companyatlas.services.llm.prompts import available_prompts, load_prompt | |
| 19 | +from companyatlas.services.llm.schemas import ChangeClassification, EventSummary, LegalDiffSummary | |
| 20 | + | |
| 21 | +BASE = "https://llm.test/v1" | |
| 22 | + | |
| 23 | + | |
| 24 | +class Probe(BaseModel): | |
| 25 | + ok: bool | |
| 26 | + n: int | |
| 27 | + | |
| 28 | + | |
| 29 | +def _completion(content: str, *, prompt_tokens: int = 50, completion_tokens: int = 20) -> dict: | |
| 30 | + return {"id": "x", "object": "chat.completion", "model": "qwen3-4b-instruct-2507-4bit", | |
| 31 | + "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], | |
| 32 | + "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens}} | |
| 33 | + | |
| 34 | + | |
| 35 | +def _provider(**kw) -> OpenAICompatibleProvider: # type: ignore[no-untyped-def] | |
| 36 | + sleeps: list[float] = [] | |
| 37 | + | |
| 38 | + async def fake_sleep(s: float) -> None: | |
| 39 | + sleeps.append(s) | |
| 40 | + | |
| 41 | + p = OpenAICompatibleProvider(base_url=BASE, api_key="k", timeout_s=5, sleep=fake_sleep, **kw) | |
| 42 | + p.sleeps = sleeps # type: ignore[attr-defined] | |
| 43 | + return p | |
| 44 | + | |
| 45 | + | |
| 46 | +@respx.mock | |
| 47 | +async def test_complete_json_happy_path(): | |
| 48 | + route = respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion('{"ok": true, "n": 3}'))) | |
| 49 | + p = _provider() | |
| 50 | + res = await p.complete_json("small", "sys", "user", Probe, max_tokens=50) | |
| 51 | + assert res.data == Probe(ok=True, n=3) | |
| 52 | + assert res.model == "qwen3-4b-instruct-2507-4bit" and res.request_tokens == 50 and res.response_tokens == 20 and res.attempts == 1 | |
| 53 | + sent = json.loads(route.calls[0].request.content) | |
| 54 | + assert sent["response_format"] == {"type": "json_object"} and sent["model"] == "qwen3-4b-instruct-2507-4bit" | |
| 55 | + assert "JSON schema" in sent["messages"][0]["content"] and sent["messages"][1]["content"] == "user" | |
| 56 | + assert route.calls[0].request.headers["Authorization"] == "Bearer k" | |
| 57 | + await p.close() | |
| 58 | + | |
| 59 | + | |
| 60 | +@respx.mock | |
| 61 | +async def test_invalid_json_triggers_one_repair_round_trip(): | |
| 62 | + route = respx.post(f"{BASE}/chat/completions") | |
| 63 | + route.side_effect = [httpx.Response(200, json=_completion("Sure! Here you go: {\"ok\": true, \"n\": \"three\"}")), | |
| 64 | + httpx.Response(200, json=_completion("```json\n{\"ok\": true, \"n\": 3}\n```"))] | |
| 65 | + p = _provider() | |
| 66 | + res = await p.complete_json("small", "sys", "user", Probe) | |
| 67 | + assert res.data.n == 3 and res.attempts == 2 and res.repaired is True | |
| 68 | + assert res.request_tokens == 100 # both calls accounted | |
| 69 | + repair_msgs = json.loads(route.calls[1].request.content)["messages"] | |
| 70 | + assert repair_msgs[-1]["role"] == "user" and "not valid" in repair_msgs[-1]["content"] | |
| 71 | + await p.close() | |
| 72 | + | |
| 73 | + | |
| 74 | +@respx.mock | |
| 75 | +async def test_still_invalid_after_repair_raises_validation_error(): | |
| 76 | + respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(200, json=_completion("no json at all"))) | |
| 77 | + p = _provider() | |
| 78 | + with pytest.raises(LLMValidationError): | |
| 79 | + await p.complete_json("small", "sys", "user", Probe) | |
| 80 | + await p.close() | |
| 81 | + | |
| 82 | + | |
| 83 | +@respx.mock | |
| 84 | +async def test_backoff_on_503_then_success(): | |
| 85 | + route = respx.post(f"{BASE}/chat/completions") | |
| 86 | + route.side_effect = [httpx.Response(503, text="loading model"), httpx.Response(429, text="busy", headers={"Retry-After": "7"}), | |
| 87 | + httpx.Response(200, json=_completion('{"ok": true, "n": 1}'))] | |
| 88 | + p = _provider(max_tries=6, backoff_initial_s=15, backoff_max_s=120) | |
| 89 | + res = await p.complete_json("small", "sys", "user", Probe) | |
| 90 | + assert res.data.n == 1 and route.call_count == 3 | |
| 91 | + assert p.sleeps == [15.0, 7.0] # exponential backoff, then Retry-After honoured | |
| 92 | + await p.close() | |
| 93 | + | |
| 94 | + | |
| 95 | +@respx.mock | |
| 96 | +async def test_backoff_exhausted_raises_retryable_error(): | |
| 97 | + respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(503, text="still loading")) | |
| 98 | + p = _provider(max_tries=3, backoff_initial_s=15, backoff_max_s=120) | |
| 99 | + with pytest.raises(LLMError) as exc: | |
| 100 | + await p.complete_json("small", "sys", "user", Probe) | |
| 101 | + assert exc.value.retryable and exc.value.status == 503 | |
| 102 | + assert p.sleeps == [15.0, 30.0] | |
| 103 | + await p.close() | |
| 104 | + | |
| 105 | + | |
| 106 | +@respx.mock | |
| 107 | +async def test_non_retryable_error_and_health(): | |
| 108 | + respx.post(f"{BASE}/chat/completions").mock(return_value=httpx.Response(401, text="bad key")) | |
| 109 | + respx.get(f"{BASE}/models").mock(return_value=httpx.Response(200, json={"data": [{"id": "qwen3-4b-instruct-2507-4bit"}, {"id": "qwen3.6-35b-a3b-4bit"}]})) | |
| 110 | + p = _provider() | |
| 111 | + with pytest.raises(LLMError) as exc: | |
| 112 | + await p.complete_json("small", "sys", "user", Probe) | |
| 113 | + assert not exc.value.retryable and exc.value.status == 401 and p.sleeps == [] | |
| 114 | + h = await p.health() | |
| 115 | + assert h.ok and "qwen3-4b-instruct-2507-4bit" in h.models | |
| 116 | + await p.close() | |
| 117 | + | |
| 118 | + | |
| 119 | +async def test_not_configured(): | |
| 120 | + p = OpenAICompatibleProvider(base_url="", api_key="") | |
| 121 | + from companyatlas.services.llm.gateway import LLMNotConfigured | |
| 122 | + | |
| 123 | + with pytest.raises(LLMNotConfigured): | |
| 124 | + await p.complete_json("small", "s", "u", Probe) | |
| 125 | + assert (await p.health()).ok is False | |
| 126 | + | |
| 127 | + | |
| 128 | +def test_json_extraction_tolerates_reasoning_and_fences(): | |
| 129 | + assert extract_json_text("<think>hmm</think>\n```json\n{\"a\": 1}\n```") == '{"a": 1}' | |
| 130 | + assert extract_json_text('prefix {"a": {"b": 2}} suffix') == '{"a": {"b": 2}}' | |
| 131 | + with pytest.raises(ValueError): | |
| 132 | + extract_json_text("nothing here") | |
| 133 | + with pytest.raises(ValueError): | |
| 134 | + parse_json_object("[1, 2]", Probe) | |
| 135 | + | |
| 136 | + | |
| 137 | +def test_schemas_enforce_taxonomy_and_wording(): | |
| 138 | + c = ChangeClassification(event_subtype="price increase", importance=0.7, confidence=0.8, title="Pro plan now $59", tags=["Pricing", "pricing", "x"]) | |
| 139 | + assert c.event_subtype == "PRICE_INCREASE" and c.tags == ["pricing", "x"] | |
| 140 | + assert ChangeClassification(event_subtype="MADE_UP", importance=0, confidence=0, title="abc").event_subtype == "OTHER" | |
| 141 | + with pytest.raises(ValueError): | |
| 142 | + EventSummary(summary="The company laid off 72 employees") | |
| 143 | + legal = LegalDiffSummary(summary="Terms: 2 sections changed.", materiality="huge", sections_changed=[{"section": "7", "change": "notice 30 days"}]) | |
| 144 | + assert legal.materiality == "unclear" and legal.sections_changed[0].section == "7" | |
| 145 | + with pytest.raises(ValueError): | |
| 146 | + ChangeClassification(event_subtype="OTHER", importance=1.5, confidence=0.5, title="x" * 10) | |
| 147 | + | |
| 148 | + | |
| 149 | +def test_prompt_files_load_by_name_and_version(): | |
| 150 | + avail = available_prompts() | |
| 151 | + for task in ("change-classifier", "event-summarizer", "legal-diff", "industry-tagger", "ask-router"): | |
| 152 | + assert "v1" in avail[task] | |
| 153 | + p = load_prompt(task, "v1") | |
| 154 | + assert p.ref == f"{task}/v1" and len(p.system) > 200 and not p.system.startswith("<!--") | |
| 155 | + if task in ("change-classifier", "event-summarizer", "legal-diff"): | |
| 156 | + assert "laid off" in p.system.lower() and "fired" in p.system.lower() # content prompts state the wording ban | |
| 157 | + assert load_prompt("legal-diff").version == "v1" | |
| 158 | + with pytest.raises(FileNotFoundError): | |
| 159 | + load_prompt("does-not-exist") | |
added
tests/test_metrics.py
+184 −0
@@ -0,0 +1,184 @@ | ||
| 1 | +"""Metrics: pure formulas (reproducible, coverage-normalised) and the DB flow (metrics_current / metric_series / company_daily / global_daily).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, date, datetime, timedelta | |
| 5 | + | |
| 6 | +import pytest | |
| 7 | +from factories import intel_db # noqa: F401 — registers the fixture | |
| 8 | + | |
| 9 | +from companyatlas.services.metrics import ( | |
| 10 | + activity_score, | |
| 11 | + ai_adoption, | |
| 12 | + anomaly_score, | |
| 13 | + corporate_change_index, | |
| 14 | + hiring_momentum, | |
| 15 | + historical_coverage, | |
| 16 | + saturate, | |
| 17 | +) | |
| 18 | +from companyatlas.taxonomy import CCI_FORMULA_VERSION, METRICS_FORMULA_VERSION, Metric | |
| 19 | + | |
| 20 | +# ------------------------------------------------------------------------------------------------------------ pure | |
| 21 | + | |
| 22 | + | |
| 23 | +def test_activity_score_coverage_normalisation(): | |
| 24 | + changes = [(1.0, "meaningful")] * 10 | |
| 25 | + small = activity_score(changes, [], active_sensors=2) | |
| 26 | + big = activity_score(changes, [], active_sensors=32) | |
| 27 | + assert small and big | |
| 28 | + assert small.value > big.value # same raw activity spread over 16× more sensors → lower density | |
| 29 | + assert small.inputs["active_sensors"] == 2 and small.formula_version == METRICS_FORMULA_VERSION | |
| 30 | + assert 0 < big.value < small.value <= 100 | |
| 31 | + | |
| 32 | + | |
| 33 | +def test_activity_score_decay_and_weights(): | |
| 34 | + fresh = activity_score([(0.0, "critical")], [], 1) | |
| 35 | + old = activity_score([(30.0, "critical")], [], 1) | |
| 36 | + minor_kind = activity_score([(0.0, "meaningful")], [], 1) | |
| 37 | + assert fresh.value > old.value > 0 | |
| 38 | + assert fresh.value > minor_kind.value | |
| 39 | + assert activity_score([], [], 0) is None # no sensors, no activity → no row | |
| 40 | + | |
| 41 | + | |
| 42 | +def test_hiring_momentum_requires_listings(): | |
| 43 | + assert hiring_momentum(10, 2, window=30) is None # < 3 listings at the reference point | |
| 44 | + m = hiring_momentum(12, 8, window=30, extra={"jobs_new_30d": 5}) | |
| 45 | + assert m.metric == Metric.HIRING_MOMENTUM_30D and m.value == 50.0 and m.inputs["jobs_new_30d"] == 5 | |
| 46 | + assert hiring_momentum(5, 10, window=7).value == -50.0 | |
| 47 | + | |
| 48 | + | |
| 49 | +def test_ai_adoption_renormalises_over_available_inputs(): | |
| 50 | + jobs_only = ai_adoption(ai_open=5, open_jobs=10, ai_events_90d=0, keyword_hits=0, has_text_inputs=False) | |
| 51 | + assert jobs_only.value == pytest.approx(100 * (0.5 * 1.0 + 0.25 * 0.0) / 0.75, abs=0.01) | |
| 52 | + none = ai_adoption(ai_open=None, open_jobs=None, ai_events_90d=0, keyword_hits=0, has_text_inputs=False) | |
| 53 | + assert none is None | |
| 54 | + text = ai_adoption(ai_open=None, open_jobs=None, ai_events_90d=3, keyword_hits=5, has_text_inputs=True) | |
| 55 | + assert 0 < text.value < 100 and set(text.inputs["weights_used"]) == {"events", "keywords"} | |
| 56 | + | |
| 57 | + | |
| 58 | +def test_cci_renormalises_and_maps_momentum(): | |
| 59 | + full = corporate_change_index({Metric.HIRING_MOMENTUM_30D: 50.0, Metric.PRODUCT_VELOCITY: 80.0, Metric.GEO_EXPANSION: 20.0, Metric.LEADERSHIP_ACTIVITY: 0.0, | |
| 60 | + Metric.DEVELOPER_MOMENTUM: 40.0, Metric.COMMUNICATION_ACTIVITY: 60.0, Metric.PRICING_ACTIVITY: 10.0}) | |
| 61 | + expected = 0.25 * 75 + 0.20 * 80 + 0.15 * 20 + 0.15 * 0 + 0.10 * 40 + 0.10 * 60 + 0.05 * 10 | |
| 62 | + assert full.value == pytest.approx(expected, abs=0.01) and full.formula_version == CCI_FORMULA_VERSION | |
| 63 | + partial = corporate_change_index({Metric.PRODUCT_VELOCITY: 80.0, Metric.PRICING_ACTIVITY: 20.0}) | |
| 64 | + assert partial.value == pytest.approx((0.20 * 80 + 0.05 * 20) / 0.25, abs=0.01) | |
| 65 | + assert partial.inputs["weight_coverage"] == 0.25 and partial.confidence < full.confidence | |
| 66 | + assert corporate_change_index({}) is None | |
| 67 | + | |
| 68 | + | |
| 69 | +def test_anomaly_and_coverage(): | |
| 70 | + assert anomaly_score(10, 2.0, 1.0, samples=2) is None | |
| 71 | + z = anomaly_score(10, 2.0, 1.0, samples=8) | |
| 72 | + assert z.value == 8.0 | |
| 73 | + flat = anomaly_score(3, 0.0, 0.0, samples=8) | |
| 74 | + assert flat.value == 6.0 # stddev floor 0.5 | |
| 75 | + cov = historical_coverage(observed=50, expected=100.0, days_with_obs=15, days_since_first=30, surfaces=4) | |
| 76 | + assert cov.value == pytest.approx(100 * (0.5 * 0.5 + 0.3 * 0.5 + 0.2 * 0.5), abs=0.01) | |
| 77 | + assert historical_coverage(observed=0, expected=0.0, days_with_obs=0, days_since_first=0, surfaces=0) is None | |
| 78 | + assert saturate(6.0, 6.0) == pytest.approx(0.632, abs=0.001) | |
| 79 | + | |
| 80 | + | |
| 81 | +# ------------------------------------------------------------------------------------------------------------ database | |
| 82 | + | |
| 83 | + | |
| 84 | +@pytest.mark.usefixtures("intel_db") | |
| 85 | +async def test_company_metrics_end_to_end(): | |
| 86 | + from factories import ( | |
| 87 | + CONNECTOR_ATS, | |
| 88 | + cleanup, | |
| 89 | + make_change, | |
| 90 | + make_company, | |
| 91 | + make_event, | |
| 92 | + make_job, | |
| 93 | + make_location, | |
| 94 | + make_observation, | |
| 95 | + make_sensor, | |
| 96 | + ) | |
| 97 | + | |
| 98 | + from companyatlas.db import fetch_all, transaction | |
| 99 | + from companyatlas.services.metrics import compute_company_metrics | |
| 100 | + | |
| 101 | + try: | |
| 102 | + async with transaction() as conn: | |
| 103 | + co = await make_company(conn, first_observed_days_ago=40) | |
| 104 | + board = await make_sensor(conn, co, "jobs_board", connector_id=CONNECTOR_ATS, created_days_ago=40) | |
| 105 | + await make_sensor(conn, co, "pricing", created_days_ago=40) | |
| 106 | + await make_sensor(conn, co, "docs", created_days_ago=40) | |
| 107 | + await make_sensor(conn, co, "locations", created_days_ago=40) | |
| 108 | + # jobs: 8 open now, 4 of them existed 30 days ago (+ 1 removed since) → momentum_30d = (8 − 5) / 5 = +60 % | |
| 109 | + for i in range(4): | |
| 110 | + await make_job(conn, co, title=f"Engineer {i}", first_seen_days_ago=45, country=None, sensor_id=board["id"]) | |
| 111 | + await make_job(conn, co, title="Old role", first_seen_days_ago=45, removed_days_ago=10, country=None, sensor_id=board["id"]) | |
| 112 | + for i in range(4): | |
| 113 | + await make_job(conn, co, title=f"ML Engineer {i}", first_seen_days_ago=5, is_ai=True, country="JP", sensor_id=board["id"]) | |
| 114 | + await make_location(conn, co, name="Tokyo office", country="JP", city="Tokyo", first_seen_days_ago=3) | |
| 115 | + await make_event(conn, co, subtype="NEW_PRODUCT", importance=0.7, days_ago=2) | |
| 116 | + await make_event(conn, co, subtype="DOC_CHANGE", importance=0.4, days_ago=3) | |
| 117 | + await make_event(conn, co, subtype="PRICE_INCREASE", importance=0.8, days_ago=4) | |
| 118 | + await make_event(conn, co, subtype="COUNTRY_EXPANSION", importance=0.8, days_ago=3) | |
| 119 | + await make_event(conn, co, subtype="AI_HIRING", importance=0.5, days_ago=5, tags=["hiring", "ai"]) | |
| 120 | + await make_change(conn, board, significance=0.6, detected_at=datetime.now(UTC) - timedelta(days=1), status="processed") | |
| 121 | + for d in range(0, 40, 2): | |
| 122 | + await make_observation(conn, board, days_ago=d) | |
| 123 | + stats = await compute_company_metrics([co["id"]]) | |
| 124 | + assert stats["companies"] == 1 | |
| 125 | + async with transaction() as conn: | |
| 126 | + rows = {r["metric"]: r for r in await fetch_all(conn, "select * from metrics_current where company_id = :c", c=co["id"])} | |
| 127 | + series = await fetch_all(conn, "select metric from metric_series where company_id = :c and day = current_date", c=co["id"]) | |
| 128 | + assert set(series and [r["metric"] for r in series]) == set(rows) | |
| 129 | + assert rows["open_jobs"]["value"] == 8 | |
| 130 | + assert rows["hiring_momentum_30d"]["value"] == pytest.approx(60.0) | |
| 131 | + assert rows["hiring_momentum_30d"]["inputs"]["jobs_new_30d"] == 4 and rows["hiring_momentum_30d"]["confidence"] == pytest.approx(0.9, abs=0.01) | |
| 132 | + assert rows["hiring_momentum_7d"]["value"] == pytest.approx(100.0) # 4 open 7 days ago (removed role already gone), 4 new since | |
| 133 | + assert "hiring_momentum_90d" not in rows # nothing was open 90 days ago → no fabricated momentum | |
| 134 | + assert rows["ai_adoption"]["value"] > 40 and rows["ai_adoption"]["inputs"]["ai_open"] == 4 | |
| 135 | + assert rows["geo_expansion"]["inputs"]["new_countries_90d"] == ["JP"] and rows["geo_expansion"]["value"] > 0 | |
| 136 | + assert rows["product_velocity"]["value"] > 0 and rows["pricing_activity"]["value"] > 0 and rows["developer_momentum"]["value"] > 0 | |
| 137 | + assert "leadership_activity" not in rows # no leadership sensor and no leadership events → no row | |
| 138 | + cci = rows["corporate_change_index"] | |
| 139 | + assert cci["formula_version"] == CCI_FORMULA_VERSION and 0 < cci["value"] <= 100 and "leadership_activity" not in cci["inputs"]["components"] | |
| 140 | + assert rows["activity_score"]["value"] > 0 and rows["activity_score"]["inputs"]["active_sensors"] == 4 | |
| 141 | + assert 0 < rows["historical_coverage"]["value"] <= 100 | |
| 142 | + assert "anomaly_score" not in rows # no baseline yet | |
| 143 | + finally: | |
| 144 | + await cleanup() | |
| 145 | + | |
| 146 | + | |
| 147 | +@pytest.mark.usefixtures("intel_db") | |
| 148 | +async def test_daily_aggregates_and_activity_index(): | |
| 149 | + from factories import cleanup, make_change, make_company, make_event, make_observation, make_sensor | |
| 150 | + | |
| 151 | + from companyatlas.db import fetch_all, fetch_one, transaction | |
| 152 | + from companyatlas.services.metrics import compute_daily | |
| 153 | + | |
| 154 | + day0 = date(2001, 1, 10) # far in the past: never collides with live data, cleaned by factories | |
| 155 | + try: | |
| 156 | + async with transaction() as conn: | |
| 157 | + a = await make_company(conn, country="CA", industries=["fintech"]) | |
| 158 | + b = await make_company(conn, country="US", industries=["retail"]) | |
| 159 | + sa = await make_sensor(conn, a, "homepage") | |
| 160 | + sb = await make_sensor(conn, b, "homepage") | |
| 161 | + for i in range(7): | |
| 162 | + d = datetime(2001, 1, 4 + i, 12, tzinfo=UTC) | |
| 163 | + for s in (sa, sb): | |
| 164 | + await make_observation(conn, s, days_ago=(datetime.now(UTC) - d).total_seconds() / 86400) | |
| 165 | + if i < 6: | |
| 166 | + await make_change(conn, sa, significance=0.6, detected_at=d, status="processed") | |
| 167 | + # day0: burst on company a | |
| 168 | + for _ in range(5): | |
| 169 | + await make_change(conn, sa, significance=0.7, detected_at=datetime(2001, 1, 10, 13, tzinfo=UTC), status="processed") | |
| 170 | + await make_event(conn, a, subtype="PRICE_INCREASE", days_ago=(datetime.now(UTC) - datetime(2001, 1, 10, 14, tzinfo=UTC)).total_seconds() / 86400) | |
| 171 | + for i in range(6): | |
| 172 | + await compute_daily(date(2001, 1, 4 + i)) | |
| 173 | + result = await compute_daily(day0) | |
| 174 | + assert result["companies"] == 2 and result["meaningful_changes"] == 5 and result["events"] == 1 | |
| 175 | + assert result["baseline_days"] == 6 | |
| 176 | + assert result["activity_index"] == pytest.approx(500.0) # 5 changes / 2 sensors vs baseline 1 change / 2 sensors → ×100 | |
| 177 | + async with transaction() as conn: | |
| 178 | + g = await fetch_one(conn, "select * from global_daily where day = :d", d=day0) | |
| 179 | + assert g["sensors_active"] == 2 and g["companies_active"] == 2 and g["events_by_type"] == {"PRICING": 1} | |
| 180 | + assert g["by_country"]["CA"]["meaningful_changes"] == 5 and g["by_industry"]["fintech"]["events"] == 1 | |
| 181 | + cd = {r["company_id"]: r for r in await fetch_all(conn, "select * from company_daily where day = :d", d=day0)} | |
| 182 | + assert cd[a["id"]]["meaningful_changes"] == 5 and cd[b["id"]]["meaningful_changes"] == 0 and cd[b["id"]]["observations"] == 1 | |
| 183 | + finally: | |
| 184 | + await cleanup() | |
added
tests/test_signals.py
+107 −0
@@ -0,0 +1,107 @@ | ||
| 1 | +"""Signals: pure detectors (wording, strength, evidence) and the DB upsert/expiry/scope aggregation flow.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime, timedelta | |
| 5 | + | |
| 6 | +import pytest | |
| 7 | +from factories import intel_db # noqa: F401 — registers the fixture | |
| 8 | + | |
| 9 | +from companyatlas.services.signals import SignalInputs, detect | |
| 10 | +from companyatlas.taxonomy import FORBIDDEN_WORDING, Metric | |
| 11 | + | |
| 12 | +NOW = datetime(2026, 9, 12, 12, tzinfo=UTC) | |
| 13 | + | |
| 14 | + | |
| 15 | +def _ev(subtype: str, *, days_ago: float = 2, event_type: str | None = None, tags: list[str] | None = None, title: str = "") -> dict: | |
| 16 | + from companyatlas.taxonomy import EVENT_SUBTYPES, EventType | |
| 17 | + | |
| 18 | + return {"id": f"evt_{subtype}_{days_ago}", "event_type": event_type or str(EVENT_SUBTYPES.get(subtype, (EventType.OTHER, 0))[0]), "event_subtype": subtype, | |
| 19 | + "importance": 0.6, "tags": tags or [], "detected_at": NOW - timedelta(days=days_ago), "title": title} | |
| 20 | + | |
| 21 | + | |
| 22 | +def _inputs(**kw) -> SignalInputs: # type: ignore[no-untyped-def] | |
| 23 | + base = {"metrics": {}, "events": [], "jobs": {"open_now": 0, "new_30d": 0, "ai_new_30d": 0, "new_countries_90d": []}, "now": NOW} | |
| 24 | + base.update(kw) | |
| 25 | + return SignalInputs(**base) | |
| 26 | + | |
| 27 | + | |
| 28 | +def _kinds(drafts): # type: ignore[no-untyped-def] | |
| 29 | + return {d.kind: d for d in drafts} | |
| 30 | + | |
| 31 | + | |
| 32 | +def test_hiring_surge_and_freeze_wording(): | |
| 33 | + surge = _kinds(detect(_inputs(metrics={Metric.HIRING_MOMENTUM_30D: 45.0}, jobs={"open_now": 40, "new_30d": 15, "ai_new_30d": 0, "new_countries_90d": []}))) | |
| 34 | + assert "hiring_surge" in surge and surge["hiring_surge"].evidence["hiring_momentum_30d"] == 45.0 and 0 < surge["hiring_surge"].strength <= 1 | |
| 35 | + freeze = _kinds(detect(_inputs(metrics={Metric.HIRING_MOMENTUM_30D: -55.0}, jobs={"open_now": 9, "new_30d": 0, "ai_new_30d": 0, "new_countries_90d": []}))) | |
| 36 | + assert "hiring_freeze" in freeze | |
| 37 | + text = (freeze["hiring_freeze"].title + freeze["hiring_freeze"].explanation).lower() | |
| 38 | + assert "no longer visible" in text and not any(b in text for b in FORBIDDEN_WORDING) | |
| 39 | + assert "hiring_surge" not in freeze | |
| 40 | + | |
| 41 | + | |
| 42 | +def test_launch_buildup_requires_multiple_categories_within_14_days(): | |
| 43 | + partial = detect(_inputs(events=[_ev("NEW_PRODUCT"), _ev("DOC_CHANGE")])) | |
| 44 | + assert "launch_buildup" not in _kinds(partial) | |
| 45 | + full = _kinds(detect(_inputs(events=[_ev("NEW_PRODUCT"), _ev("DOC_CHANGE"), _ev("CHANGELOG_ENTRY"), _ev("JOB_COUNT_INCREASE")]))) | |
| 46 | + assert "launch_buildup" in full | |
| 47 | + d = full["launch_buildup"] | |
| 48 | + assert d.window_days == 14 and d.title.startswith("Possible launch preparation signal") and "probabilistic" in d.explanation | |
| 49 | + assert set(d.evidence["categories"]) == {"product", "docs", "changelog", "careers"} and len(d.evidence["event_ids"]) == 4 | |
| 50 | + stale = detect(_inputs(events=[_ev("NEW_PRODUCT", days_ago=20), _ev("DOC_CHANGE", days_ago=20), _ev("CHANGELOG_ENTRY", days_ago=20), _ev("JOB_COUNT_INCREASE", days_ago=20)])) | |
| 51 | + assert "launch_buildup" not in _kinds(stale) | |
| 52 | + | |
| 53 | + | |
| 54 | +def test_expansion_pricing_developer_enterprise_ai_abnormal(): | |
| 55 | + drafts = _kinds(detect(_inputs( | |
| 56 | + metrics={Metric.DEVELOPER_MOMENTUM: 70.0, Metric.ANOMALY_SCORE: 3.1}, | |
| 57 | + events=[_ev("COUNTRY_EXPANSION"), _ev("NEW_PRICING_TIER"), _ev("PRICING_TIER_REMOVED"), _ev("API_CHANGE"), _ev("DOC_CHANGE"), _ev("CHANGELOG_ENTRY"), | |
| 58 | + _ev("AI_HIRING"), _ev("NEWS_RELEASE", tags=["ai"]), _ev("MESSAGING_CHANGE", title="Now with SSO and audit logs for enterprise")], | |
| 59 | + jobs={"open_now": 20, "new_30d": 8, "ai_new_30d": 4, "new_countries_90d": ["JP"]}, plans_contact_sales_new=1, locations_new_countries=["JP"]))) | |
| 60 | + assert {"expansion", "pricing_migration", "developer_push", "enterprise_repositioning", "ai_acceleration", "abnormal_activity"} <= set(drafts) | |
| 61 | + assert drafts["expansion"].evidence["new_countries"] == ["JP"] | |
| 62 | + assert drafts["pricing_migration"].evidence["subtypes"] == ["NEW_PRICING_TIER", "PRICING_TIER_REMOVED"] | |
| 63 | + assert drafts["ai_acceleration"].evidence["ai_new_30d"] == 4 and "public job titles" in drafts["ai_acceleration"].explanation | |
| 64 | + assert drafts["abnormal_activity"].window_days == 7 and drafts["abnormal_activity"].evidence["anomaly_z"] == 3.1 | |
| 65 | + for d in drafts.values(): | |
| 66 | + assert 0 < d.strength <= 1 and 0 < d.confidence <= 1 and "signal" in d.title.lower() | |
| 67 | + | |
| 68 | + | |
| 69 | +def test_sparse_inputs_produce_nothing(): | |
| 70 | + assert detect(_inputs()) == [] | |
| 71 | + assert detect(_inputs(metrics={Metric.ANOMALY_SCORE: 1.0}, events=[_ev("BLOG_POST")])) == [] | |
| 72 | + | |
| 73 | + | |
| 74 | +@pytest.mark.usefixtures("intel_db") | |
| 75 | +async def test_signal_upsert_expiry_and_scope_aggregate(): | |
| 76 | + from factories import cleanup, make_company, make_event | |
| 77 | + | |
| 78 | + from companyatlas.db import execute, fetch_all, transaction | |
| 79 | + from companyatlas.services.signals import compute_signals | |
| 80 | + | |
| 81 | + try: | |
| 82 | + async with transaction() as conn: | |
| 83 | + companies = [await make_company(conn, country="CA", industries=["fintech"]) for _ in range(3)] | |
| 84 | + for co in companies: | |
| 85 | + for st in ("NEW_PRICING_TIER", "PRICING_TIER_REMOVED"): | |
| 86 | + await make_event(conn, co, subtype=st, days_ago=1) | |
| 87 | + ids = [c["id"] for c in companies] | |
| 88 | + first = await compute_signals(ids) | |
| 89 | + assert first["inserted"] == 3 and first["scope"] >= 1 | |
| 90 | + second = await compute_signals(ids) | |
| 91 | + assert second["inserted"] == 0 and second["updated"] == 3 # idempotent: same signal updated, not duplicated | |
| 92 | + async with transaction() as conn: | |
| 93 | + rows = await fetch_all(conn, "select scope, scope_key, kind, status, evidence from signals where company_id = any(cast(:ids as text[])) or (scope <> 'company' and scope_key = 'CA' and kind = 'pricing_migration')", ids=ids) | |
| 94 | + company_rows = [r for r in rows if r["scope"] == "company"] | |
| 95 | + assert len(company_rows) == 3 and all(r["kind"] == "pricing_migration" and r["status"] == "active" for r in company_rows) | |
| 96 | + scope_rows = [r for r in rows if r["scope"] == "country"] | |
| 97 | + assert scope_rows and scope_rows[0]["evidence"]["count"] >= 3 | |
| 98 | + # evidence disappears → signal expires | |
| 99 | + await execute(conn, "update events set status = 'retracted' where company_id = :c", c=ids[0]) | |
| 100 | + third = await compute_signals([ids[0]]) | |
| 101 | + assert third["expired"] == 1 | |
| 102 | + async with transaction() as conn: | |
| 103 | + st = await fetch_all(conn, "select status from signals where company_id = :c", c=ids[0]) | |
| 104 | + assert [r["status"] for r in st] == ["expired"] | |
| 105 | + await execute(conn, "delete from signals where scope = 'country' and scope_key = 'CA' and kind = 'pricing_migration'") | |
| 106 | + finally: | |
| 107 | + await cleanup() | |
| 108 | ||