SPB Git
7.8 KB · 203 lines python
1# ============================================================2# SPID — LLM extractor (OpenAI-compatible; default gpt-4o-mini)3#4# Uses the chat-completions JSON mode to extract structured5# project records.  Cost-controlled: only project-relevant6# windows of each section are sent to the model.7# ============================================================8from __future__ import annotations910import json11import re12import time1314import httpx1516from spid import taxonomy as tx17from spid.config import get_logger, get_settings18from spid.extract.base import ProjectMention19from spid.extract.rules import _PROJECT_GATE2021log = get_logger("extract.llm")2223_SYSTEM = """You are a financial analyst extracting CORPORATE PROJECTS from SEC filing text.24A "project" is any concrete initiative, investment, program, expansion, construction,25product/R&D effort, acquisition/integration, technology deployment, partnership, or26strategic program the company is pursuing.2728Return STRICT JSON: {"projects": [ ... ]}. Each project object has:29  project_name  : short descriptive name (use the official name if one is stated)30  project_type  : EXACTLY ONE of the allowed types below31  description   : 1 sentence, grounded in the text (no speculation)32  objective     : the goal/why (or "")33  status        : one of planned | in_progress | completed | mentioned34  amount_usd    : numeric USD figure if stated (e.g. 1500000000), else null35  locations     : array of places (city/state/country), else []36  technologies  : array of technologies involved, else []37  partners      : array of named partners/JV counterparties, else []38  benefits      : array of expected benefits, else []39  risks         : array of stated risks, else []40  years         : array of 4-digit years mentioned, else []41  confidence    : 0.0-1.0 (how clearly this is a real, concrete project)4243Allowed project_type values:44%s4546Rules:47- Extract ONLY real initiatives, not generic risk-factor boilerplate or hypotheticals.48- If the text describes no concrete project, return {"projects": []}.49- Do not invent facts. Numbers/locations must appear in the text.50- Prefer fewer high-quality projects over many vague ones.""" % "\n".join(51    f"  - {k}: {v}" for k, v in tx.TYPE_LABELS.items()52)535455def _focus(text: str, budget: int = 9000) -> str:56    """Keep only windows around project-signal hits to bound token cost."""57    if len(text) <= budget:58        return text59    spans: list[tuple[int, int]] = []60    for m in _PROJECT_GATE.finditer(text):61        spans.append((max(0, m.start() - 220), min(len(text), m.end() + 380)))62        if sum(b - a for a, b in spans) > budget * 1.4:63            break64    if not spans:65        return text[:budget]66    # merge overlapping windows67    spans.sort()68    merged = [spans[0]]69    for a, b in spans[1:]:70        if a <= merged[-1][1]:71            merged[-1] = (merged[-1][0], max(merged[-1][1], b))72        else:73            merged.append((a, b))74    chunks, total = [], 075    for a, b in merged:76        chunks.append(text[a:b])77        total += b - a78        if total >= budget:79            break80    return "\n...\n".join(chunks)818283class LLMExtractor:84    name = "llm"8586    def __init__(self) -> None:87        self.s = get_settings()88        if not self.s.llm_api_key:89            raise RuntimeError("No LLM API key configured (OPENAI_API_KEY / MACLUSTR_API_KEY).")90        self._client = httpx.Client(91            base_url=self.s.llm_base_url,92            headers={"Authorization": f"Bearer {self.s.llm_api_key}"},93            timeout=self.s.llm_timeout,94        )9596    def _call(self, content: str, retries: int = 3) -> dict:97        payload = {98            "model": self.s.llm_model,99            "messages": [100                {"role": "system", "content": _SYSTEM},101                {"role": "user", "content": content},102            ],103            "temperature": 0,104            "response_format": {"type": "json_object"},105            "max_tokens": 2000,106        }107        last = None108        for attempt in range(retries):109            try:110                r = self._client.post("/chat/completions", json=payload)111                if r.status_code == 429 or r.status_code >= 500:112                    raise httpx.HTTPStatusError("retryable", request=r.request, response=r)113                r.raise_for_status()114                txt = r.json()["choices"][0]["message"]["content"]115                return json.loads(txt)116            except (httpx.HTTPError, json.JSONDecodeError, KeyError) as e:117                last = e118                time.sleep(1.5 * (attempt + 1))119        log.warning("LLM call failed after %d retries: %s", retries, last)120        return {"projects": []}121122    def extract(self, *, text: str, accession_number: str, cik: str,123                section_id: str, section_name: str) -> list[ProjectMention]:124        if not text or not _PROJECT_GATE.search(text):125            return []126        content = (f"SECTION: {section_name}\n\n" + _focus(text))[:14000]127        data = self._call(content)128        out: list[ProjectMention] = []129        for p in data.get("projects", []) or []:130            if not isinstance(p, dict):131                continue132            ptype = self._norm_type(p.get("project_type", ""))133            desc = str(p.get("description", "")).strip()[:600]134            if not desc:135                continue136            amount = p.get("amount_usd")137            try:138                amount = float(amount) if amount not in (None, "", "null") else None139            except (TypeError, ValueError):140                amount = None141            anchor = f"{ptype}:{str(p.get('project_name',''))[:40]}"142            conf = p.get("confidence", 0.6)143            try:144                conf = float(conf)145            except (TypeError, ValueError):146                conf = 0.6147            out.append(ProjectMention(148                mention_id=ProjectMention.make_id(accession_number, section_id, ptype, anchor),149                accession_number=accession_number, cik=cik,150                section_id=section_id, section_name=section_name,151                project_name=str(p.get("project_name") or tx.TYPE_LABELS.get(ptype, ptype))[:200],152                project_type=ptype, description=desc,153                objective=str(p.get("objective", ""))[:400],154                amount_usd=amount,155                amount_raw=(f"${amount:,.0f}" if amount else None),156                locations=_strlist(p.get("locations")),157                technologies=_strlist(p.get("technologies")) or tx.extract_technologies(desc),158                partners=_strlist(p.get("partners")),159                benefits=_strlist(p.get("benefits")),160                risks=_strlist(p.get("risks")),161                years=_intlist(p.get("years")) or tx.extract_years(desc),162                status=self._norm_status(p.get("status")),163                confidence=round(max(0.0, min(1.0, conf)), 3),164                backend=self.name,165            ))166        return out167168    @staticmethod169    def _norm_type(raw: str) -> str:170        raw = (raw or "").strip().lower().replace(" ", "_").replace("-", "_")171        if raw in tx.PROJECT_TYPES:172            return raw173        # fuzzy contains match against labels/keys174        for k, label in tx.TYPE_LABELS.items():175            if raw and (raw in k or raw in label.lower().replace(" ", "_")):176                return k177        return "technology_deployment"178179    @staticmethod180    def _norm_status(raw) -> str:181        raw = (str(raw or "")).strip().lower()182        return raw if raw in ("planned", "in_progress", "completed", "mentioned") else "mentioned"183184185def _strlist(v) -> list[str]:186    if not v:187        return []188    if isinstance(v, str):189        v = [v]190    return [str(x).strip()[:120] for x in v if str(x).strip()][:8]191192193def _intlist(v) -> list[int]:194    out = []195    for x in (v or []):196        try:197            y = int(str(x)[:4])198            if 1980 <= y <= 2040:199                out.append(y)200        except (TypeError, ValueError):201            continue202    return sorted(set(out))203
Files · main
Menu

Theme