242 lines
13 KiB
Markdown
242 lines
13 KiB
Markdown
# Refactor Plan: Prescreen Worker → Hybrid Extraction (Heuristic + LLM Fallback)
|
||
|
||
> Status: delivered through Phase 4; Phase 5 rollout pending merge
|
||
> Scope: `src/contract_check/worker_prescreen/`, `src/contract_check/core/llm/`, `src/contract_check/core/config.py`, `src/contract_check/core/metrics.py`
|
||
> Supersedes: the regex-only extractor described in `worker_prescreen/extractor.py` (regex-v1)
|
||
|
||
---
|
||
|
||
## 1. Goal
|
||
|
||
Replace the monolithic regex extractor (`extractor.py`, ~260 lines of Russian/BY legal
|
||
regexes) with a **hybrid two-stage extractor**:
|
||
|
||
1. **Stage 1 — Heuristic (deterministic, zero-cost):** keyword dictionaries,
|
||
positional windows, sentence scanning, per-field confidence weights.
|
||
No regex — plain string operations only.
|
||
2. **Stage 2 — LLM fallback (only when Stage 1 confidence is low):** reuse the
|
||
existing `core/llm` provider infrastructure (Ollama Cloud / YandexGPT) with a
|
||
JSON-schema-constrained extraction prompt.
|
||
|
||
Non-goals:
|
||
|
||
- No changes to routing semantics (`router.py` thresholds/decisions stay as-is).
|
||
- No DB schema migration (`prescreen_results.extractor_version` already exists,
|
||
`String(32)` fits `"heuristic-v2"` / `"llm-v1"`).
|
||
- No changes to MQ topology, messages (`PrescreenCompleted` shape unchanged),
|
||
or downstream `worker_analyze`.
|
||
- Auto-approve stays disabled by default.
|
||
|
||
## 2. Current State
|
||
|
||
```
|
||
worker_prescreen/
|
||
├── consumer.py # consumes prescreen.q → handler
|
||
├── handler.py # Stage 4: asyncio.to_thread(extract_contract_meta, text)
|
||
├── extractor.py # PrescreenContractMeta + 10 regex patterns ← TO REPLACE
|
||
├── router.py # decides auto_approve|manual_review|deep_analysis
|
||
└── config.py
|
||
```
|
||
|
||
Problems:
|
||
|
||
- Regexes are brittle (line-noise, OCR artifacts, whitespace variants) and hard
|
||
to extend (every new contract type = new regex).
|
||
- No recovery path: low-confidence extraction always ends in `manual_review`.
|
||
- `prescreen_results.extractor_version` is never written by the handler INSERT —
|
||
it silently relies on the DB default `'regex-v1'`.
|
||
|
||
## 3. Target Architecture
|
||
|
||
```
|
||
┌────────────────────────────────────────────┐
|
||
│ handler.py │
|
||
│ Stage 4: meta = await extractor.extract() │
|
||
└───────────────┬────────────────────────────┘
|
||
│
|
||
HybridMetaExtractor (orchestrator)
|
||
│
|
||
┌────────────────────┴──────────────────────┐
|
||
▼ ▼
|
||
HeuristicExtractor (sync, to_thread) LLMPrescreenExtractor
|
||
keyword + positional + sentence wraps core/llm provider
|
||
scanning, confidence-weighted extract_prescreen(text)
|
||
│ │
|
||
confidence ≥ threshold? ──── no ────► run LLM, validate JSON,
|
||
│ yes merge over heuristic meta
|
||
▼
|
||
PrescreenContractMeta → router.decide() → persist (with extractor_version)
|
||
```
|
||
|
||
### 3.1 New module layout (flat, matches existing package conventions)
|
||
|
||
```
|
||
worker_prescreen/
|
||
├── extractor.py # PrescreenContractMeta (model, unchanged shape)
|
||
│ # + MetadataExtractor protocol + re-export shim
|
||
├── extractor_heuristic.py# HeuristicExtractor (Stage 1)
|
||
├── extractor_llm.py # LLMPrescreenExtractor (Stage 2, dict → pydantic validation)
|
||
├── extractor_hybrid.py # HybridMetaExtractor (threshold + merge + kill-switch)
|
||
├── handler.py # Stage 4 rewired; INSERT gains extractor_version
|
||
└── router.py # unchanged
|
||
```
|
||
|
||
### 3.2 Heuristic stage design (regex-free)
|
||
|
||
| Field | Method | Notes |
|
||
|---|---|---|
|
||
| `contract_type` | Phrase dictionary match on normalized text | `{"supply": ["договор поставки", "договор купли-продажи", ...], ...}` — lowercase + whitespace-collapse once, then `in` checks |
|
||
| `party_a/b` | Positional: scan only the header window (text up to first `1.` / `ПРЕДМЕТ` heading, capped ~1500 chars); token-scan for entity-form tokens (ООО, ИП, АО, …) then capture until quote `»`/`"` close or line end | Replaces `_ENTITY_RE` |
|
||
| `total_amount` + `currency` | Trigger-word scan (`составляет`, `стоимость`, `цена`, …) then manual digit-window parse (`_scan_number` helper walking chars, handling spaces/commas); currency via token lookup in the following ~50 chars | Replaces `_AMOUNT_RE` |
|
||
| `start/end_date` | Trigger-word scan (`действует с`, `с … по …`) then `_parse_ddmmyyyy` manual splitter (already regex-free in v1 — keep) | Replaces `_DATE_RE`/`_END_DATE_RE` |
|
||
| boolean clauses | Sentence segmentation via `str.split` on `. ` / `\n`, keyword membership per sentence (`неустойка`, `штраф`, `расторгнуть`, `арбитражн`, …) | Replaces `_PENALTY_RE` etc. |
|
||
|
||
Confidence scoring becomes **weighted** instead of flat field coverage:
|
||
|
||
```python
|
||
FIELD_WEIGHTS = {
|
||
"contract_type": 0.25,
|
||
"party_a": 0.15,
|
||
"party_b": 0.15,
|
||
"total_amount": 0.15,
|
||
"currency": 0.05,
|
||
"start_date": 0.10,
|
||
"end_date": 0.05,
|
||
"has_penalty_clause": 0.05,
|
||
"has_termination_clause": 0.05,
|
||
"has_arbitration": 0.05,
|
||
} # sums to 1.0
|
||
```
|
||
|
||
Optional per-field method bonus (exact phrase match = full weight, positional
|
||
window hit = full weight, fuzzy tail hit = ×0.7) — start simple, weights
|
||
constant, tune later with real data.
|
||
|
||
### 3.3 LLM fallback design
|
||
|
||
- **Protocol extension** (`core/llm/port.py`):
|
||
`async def extract_prescreen(self, text: str) -> dict[str, Any]` — returns raw
|
||
JSON dict; **no import of worker_prescreen** (keeps layering clean).
|
||
Implemented by both `OllamaCloudProvider` and `YandexGPTProvider` as a thin
|
||
wrapper over their existing JSON-chat + repair-loop machinery
|
||
(`_chat_json` / `responseFormat=json_schema`) with a dedicated
|
||
`PRESCREEN_SYSTEM` prompt: extract the 10 meta fields, cite nothing, JSON only.
|
||
- **Worker-side wrapper** (`extractor_llm.py`):
|
||
`dict` → `PrescreenContractMeta` via pydantic (rejects hallucinated fields,
|
||
re-`None`s unknown enum values), then recomputes weighted confidence.
|
||
- **Input cap:** only the first `prescreen_llm_max_chars` (default 20 000) chars
|
||
are sent — meta lives in the header for templated contracts; cost control.
|
||
- **Merge rule:** LLM values override heuristic `None`s and low-confidence
|
||
fields; boolean clauses become OR(heuristic, llm) — both sources are
|
||
presence-checks, false positives are cheap, false negatives route wrong.
|
||
|
||
### 3.4 Failure semantics
|
||
|
||
- LLM fallback error (quota, timeout, invalid JSON after repair):
|
||
**do not fail the message.** Log + keep heuristic meta,
|
||
`extractor_version="heuristic-v2"`, record the error in
|
||
`prescreen_results.auto_findings` (`{"llm_fallback_error": "..."}`) and the
|
||
`prescreen_fallback_runs_total{outcome="failed"}` counter.
|
||
- Heuristic stage is pure string ops — its only failure mode is `None` fields,
|
||
which is already handled by low confidence → fallback / manual_review.
|
||
|
||
### 3.5 Config additions (`core/config.py`, `--- prescreen stage ---`)
|
||
|
||
```python
|
||
prescreen_llm_fallback_enabled: bool = Field(default=False, ...) # kill-switch; flip to True after burn-in
|
||
prescreen_llm_fallback_threshold: float = Field(default=0.75, ...) # ≤ router threshold
|
||
prescreen_llm_max_chars: int = Field(default=20_000, ...)
|
||
```
|
||
|
||
### 3.6 Metrics additions (`core/metrics.py`)
|
||
|
||
```python
|
||
prescreen_fallback_runs = Counter(
|
||
"contract_check_prescreen_fallback_runs_total",
|
||
"...", labelnames=["outcome"], # used | failed | skipped | disabled
|
||
)
|
||
prescreen_extraction_stage = Histogram(
|
||
"contract_check_prescreen_extraction_stage_seconds",
|
||
"...", labelnames=["stage"], # heuristic | llm
|
||
)
|
||
```
|
||
|
||
`extractor_version` values: `"regex-v1"` (legacy shim flag), `"heuristic-v2"`
|
||
(LLM not run / disabled / failed) and `"hybrid-llm-v1"` (LLM result merged).
|
||
|
||
## 4. Implementation Phases
|
||
|
||
Each phase is independently shippable and reverted by config/env flag.
|
||
|
||
### Phase 1 — Port + heuristic extractor (no behavior change for routing)
|
||
- [x] `extractor.py`: keep `PrescreenContractMeta` + weighted `_score_confidence`;
|
||
add `MetadataExtractor` protocol (`def extract(text: str) -> PrescreenContractMeta`);
|
||
keep `extract_contract_meta` as a delegating shim so
|
||
`tests/unit/test_prescreen_extractor.py` and `worker_prescreen/__init__.py`
|
||
keep importing it.
|
||
- [x] New `extractor_heuristic.py` implementing all helpers from §3.2
|
||
(`_scan_number`, `_find_after_trigger`, header-window splitter,
|
||
sentence splitter, dictionaries as module-level constants).
|
||
- [x] Port existing unit-test fixtures (`SIMPLE_SUPPLY`, `MINIMAL`, boolean-flag
|
||
parametrize) onto `HeuristicExtractor`; they must pass with identical
|
||
expected values.
|
||
- [x] `extractor_version` shim reports `"heuristic-v2"`.
|
||
|
||
### Phase 2 — LLM `extract_prescreen` on providers
|
||
- [x] Extend `LLMProvider` protocol + both adapters
|
||
(`ollama_cloud.py`, `yandex_gpt.py`): new `PRESCREEN_SYSTEM` prompt,
|
||
JSON schema for the 10 fields, reuse repair loop, truncate input to
|
||
`chunk_size`-independent small cap.
|
||
- [x] Unit tests with a fake transport (both adapters already have this pattern):
|
||
valid dict, invalid enum → `None`, malformed JSON → repair once → fail.
|
||
- [x] `extractor_llm.py`: dict → `PrescreenContractMeta` validation +
|
||
weighted confidence.
|
||
|
||
### Phase 3 — Orchestrator + handler wiring
|
||
- [x] `extractor_hybrid.py`: threshold check, LLM call (async), merge rules,
|
||
error swallowing, extractor_version selection, kill-switch.
|
||
- [x] `handler.py` Stage 4: replace `asyncio.to_thread(extract_contract_meta, …)`
|
||
with `await self._extractor.extract(contract_text)` where
|
||
`self._extractor` is injectable (constructor arg, defaults to hybrid) —
|
||
mirrors the existing `provider` injection pattern in `AnalyzeHandler`.
|
||
- [x] Handler Stage 6 INSERT: bind `extractor_version`.
|
||
- [x] Config + metrics from §3.5/§3.6.
|
||
|
||
### Phase 4 — Tests + verification
|
||
- [x] Unit: orchestrator matrix — high confidence skips LLM; low confidence
|
||
merges; LLM failure → heuristic result + `outcome="failed"`; disabled →
|
||
`outcome="disabled"`.
|
||
- [x] Integration (`tests/integration/test_prescreen_worker.py` pattern):
|
||
run handler in-process with a stub provider; assert
|
||
`prescreen_results.extractor_version` persisted, routing unchanged.
|
||
- [x] Backfill spot-check: replayed 8 RU/BY-style contract samples through
|
||
regex-v1 vs heuristic-v2. Confidence deltas: all neutral or better,
|
||
heuristic fixes regex under-parsing on NBSP/noise and captures fuller
|
||
party names. Record below in PR description.
|
||
- [x] `ruff check . && mypy src && pytest tests/unit tests/integration -k prescreen`.
|
||
(ruff/mypy findings outside the touched files are from pre-existing
|
||
uncommitted changes; prescreen tests pass.)
|
||
|
||
### Phase 5 — Rollout
|
||
- [ ] Merge with `prescreen_llm_fallback_enabled=false` (pure heuristic).
|
||
- [ ] Observe `prescreen_fallback_runs` / `prescreen_confidence` for a few days.
|
||
- [x] Document in `docs/ARCHITECTURE.md` prescreen section + `.env.example`
|
||
completed. The regex notes in `SPIKE_PHASE0.md` are no longer referenced
|
||
from the extractor docstring. Staging/prod enablement remains an
|
||
operational step after merge.
|
||
|
||
## 5. Risks & Mitigations
|
||
|
||
| Risk | Mitigation |
|
||
|---|---|
|
||
| Heuristic drops accuracy vs regex on edge templates | Port 100% of v1 unit tests before deleting regex; keep `extractor.py` regex code one release behind a flag (`PRESCREEN_KEEP_REGEX=true` env, temporary) |
|
||
| LLM hallucinates fields | pydantic validation whitelists contract_type enums; numeric/date parse checks; booleans only OR-merged |
|
||
| LLM latency blows up prescreen SLA | 20k char cap, existing provider timeouts, `prescreen_extraction_stage` histogram; fallback failures never block the pipeline |
|
||
| Cost creep on fallback rate | `prescreen_fallback_runs_total` alert; threshold tunable without deploy via env |
|
||
| Layering violation (core ← worker import) | Protocol returns plain `dict`; pydantic model stays in `worker_prescreen` |
|
||
|
||
## 6. Rollback
|
||
|
||
1. Config: `PRESCREEN_LLM_FALLBACK_ENABLED=false` → deterministic heuristic only.
|
||
2. Full: revert merge — no DB migration to undo; `extractor_version` strings are
|
||
informational only.
|