skillfactor-pipeline/pipeline/extract_local.py

"""Local-LLM extraction of job postings via Ollama (RTX-3090 box).

Replaces the in-session Claude extraction (see p3b docstring) so the full
catalog (~180k ads) can run outside the Claude loop. Claude only spot-checks
samples (see qa_sample.py).

- Prompt:  pipeline/prompts/extract_posting.txt  (strict JSON, 3 few-shots)
- Schema:  pipeline/prompts/extract_schema.json  (also passed to Ollama as
           structured-output format, which pins the keys/types server-side)
- Config:  .env  OLLAMA_URL / OLLAMA_MODEL
- Resume:  output file is JSONL keyed by job_id; already-extracted ids are
           skipped, so re-runs after an abort cost nothing.

Usage:
  python pipeline/extract_local.py --ads data/raw/jobs/recruiter_ads.json \
      --out data/evidence/recruitment-consultant.jsonl [--limit 20] [--offset 0]

Exit code 1 if >10% of processed ads failed validation (QUALITY_BAR).
"""
import argparse
import json
import os
import re
import sys
import time

import requests
from dotenv import load_dotenv

BASE = os.path.join(os.path.dirname(__file__), "..")
load_dotenv(os.path.join(BASE, ".env"))

PROMPT_FILE = os.path.join(os.path.dirname(__file__), "prompts", "extract_posting.txt")
SCHEMA_FILE = os.path.join(os.path.dirname(__file__), "prompts", "extract_schema.json")

MAX_POSTING_CHARS = 12000  # keeps prompt+ad inside num_ctx for 8-14B models
ARRAY_FIELDS = {
    "hard_skills": (12, 60), "tools": (12, 60), "methods": (10, 60),
    "responsibilities": (10, 80), "qualifications": (8, 80), "soft_skills": (8, 40),
}
SENIORITY = {"junior", "mid", "senior", "lead", "n/a"}
# gemma3 copies these prefixes into qualifications although they are skills,
# not formal requirements (testset 2026-07-07: 1/20 records) -> drop them there
NOT_A_QUALIFICATION = ("strong ", "understanding of", "knowledge of", "familiarity with")
ACRONYM_TOOLS = {"ats": "applicant tracking system (ATS)", "crm": "CRM",
                 "hris": "HRIS", "erp": "ERP", "cad": "CAD", "pos": "POS"}


def validate(rec):
    """Enforce schema + normalization; returns (clean_record, list_of_issues).

    Hard failures (missing keys, wrong types) -> issues, record unusable.
    Soft deviations (dupes, over-long lists, stray punctuation) are repaired
    silently -- the local model gets a strict prompt, the validator gets the
    final word.
    """
    issues = []
    if not isinstance(rec, dict):
        return None, ["not a JSON object"]
    clean = {}
    for field, (max_items, max_len) in ARRAY_FIELDS.items():
        val = rec.get(field)
        if val is None:
            issues.append(f"missing {field}")
            val = []
        if not isinstance(val, list):
            issues.append(f"{field} not a list")
            val = []
        out, seen = [], set()
        for item in val:
            if not isinstance(item, str):
                continue
            item = re.sub(r"\s+", " ", item.replace("<EFBFBD>", "'")).strip(" .,;:-")
            if field == "qualifications" and item.lower().startswith(NOT_A_QUALIFICATION):
                continue
            if field == "tools":
                item = ACRONYM_TOOLS.get(item.lower(), item)
            if not (2 <= len(item) <= max_len):
                continue
            key = item.lower()
            if key in seen:
                continue
            seen.add(key)
            out.append(item)
        clean[field] = out[:max_items]
    sen = rec.get("seniority")
    if sen not in SENIORITY:
        issues.append(f"seniority invalid: {sen!r}")
        sen = "n/a"
    clean["seniority"] = sen
    extra = set(rec) - set(ARRAY_FIELDS) - {"seniority"}
    if extra:
        issues.append(f"extra keys: {sorted(extra)}")
    hard_fail = [i for i in issues if i.startswith(("missing", "not a", "seniority invalid"))
                 or "not a list" in i]
    return clean, hard_fail


def call_ollama(prompt_tpl, schema, posting, url, model, retries=3):
    prompt = prompt_tpl.replace("{POSTING}", posting[:MAX_POSTING_CHARS])
    body = {
        "model": model, "prompt": prompt, "stream": False,
        "format": schema,  # Ollama structured output: server enforces shape
        "options": {"temperature": 0, "num_ctx": 8192, "num_predict": 1024},
    }
    for attempt in range(retries):
        try:
            r = requests.post(f"{url}/api/generate", json=body, timeout=300)
            r.raise_for_status()
            raw = r.json().get("response", "")
            return json.loads(raw), None
        except (requests.RequestException, json.JSONDecodeError) as exc:
            err = f"{type(exc).__name__}: {exc}"
            time.sleep(5 * (attempt + 1))
    return None, err


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ads", required=True, help="JSON list of ads with job_id+description")
    ap.add_argument("--out", required=True, help="JSONL output (resumable)")
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--offset", type=int, default=0)
    args = ap.parse_args()

    url = os.environ.get("OLLAMA_URL", "").rstrip("/")
    model = os.environ.get("OLLAMA_MODEL", "")
    if not url or not model:
        sys.exit("OLLAMA_URL / OLLAMA_MODEL missing in .env")

    prompt_tpl = open(PROMPT_FILE, encoding="utf-8").read()
    schema = json.load(open(SCHEMA_FILE, encoding="utf-8"))
    schema.pop("$schema", None)  # Ollama's validator chokes on meta keys

    ads = json.load(open(args.ads, encoding="utf-8"))
    done = set()
    if os.path.exists(args.out):
        with open(args.out, encoding="utf-8") as f:
            done = {json.loads(line)["job_id"] for line in f if line.strip()}

    todo = [a for a in ads if a.get("job_id") not in done and a.get("description")]
    todo = todo[args.offset:]
    if args.limit:
        todo = todo[:args.limit]
    print(f"{len(ads)} ads, {len(done)} already extracted, processing {len(todo)}")

    ok = fail = 0
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "a", encoding="utf-8") as out:
        for i, ad in enumerate(todo, 1):
            t0 = time.time()
            rec, err = call_ollama(prompt_tpl, schema, ad["description"], url, model)
            if rec is not None:
                clean, hard_fail = validate(rec)
                if not hard_fail:
                    clean["job_id"] = ad["job_id"]
                    out.write(json.dumps(clean, ensure_ascii=False) + "\n")
                    out.flush()
                    ok += 1
                    print(f"  [{i}/{len(todo)}] ok    {time.time()-t0:5.1f}s  {ad.get('title','')[:60]}")
                    continue
                err = "; ".join(hard_fail)
            fail += 1
            print(f"  [{i}/{len(todo)}] FAIL  {time.time()-t0:5.1f}s  {ad.get('title','')[:40]}  {err}")

    total = ok + fail
    print(f"done: {ok} ok, {fail} failed" + (f" ({100*fail/total:.0f}% fail)" if total else ""))
    if total and fail / total > 0.10:
        sys.exit(1)


if __name__ == "__main__":
    main()